Merge pull request #1174 from dolthub/daylon/create-function-conversion

Parse CREATE FUNCTION into runnable functions

go.mod

@@ -26,6 +26,7 @@ require (
 	github.com/madflojo/testcerts v1.1.1
 	github.com/mitchellh/go-ps v1.0.0
 	github.com/mitchellh/go-wordwrap v1.0.1
+	github.com/pganalyze/pg_query_go/v6 v6.0.0
 	github.com/pierrre/geohash v1.0.0
 	github.com/pkg/profile v1.5.0
 	github.com/sergi/go-diff v1.1.0

go.sum

@@ -729,6 +729,8 @@ github.com/pborman/getopt v0.0.0-20180729010549-6fdd0a2c7117/go.mod h1:85jBQOZwp
 github.com/pborman/uuid v1.2.0/go.mod h1:X/NO0urCmaxf9VXbdlT7C2Yzkj2IKimNn4k+gtPdI/k=
 github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic=
 github.com/performancecopilot/speed v3.0.0+incompatible/go.mod h1:/CLtqpZ5gBg1M9iaPbIdPPGyKcA8hKdoy6hAWba7Yac=
+github.com/pganalyze/pg_query_go/v6 v6.0.0 h1:in6RkR/apfqlAtvqgDxd4Y4o87a5Pr8fkKDB4DrDo2c=
+github.com/pganalyze/pg_query_go/v6 v6.0.0/go.mod h1:nvTHIuoud6e1SfrUaFwHqT0i4b5Nr+1rPWVds3B5+50=
 github.com/pierrec/lz4 v1.0.2-0.20190131084431-473cd7ce01a1/go.mod h1:3/3N9NVKO0jef7pBehbT1qWhCMrIgbYNnFAZCqQ5LRc=
 github.com/pierrec/lz4 v2.0.5+incompatible/go.mod h1:pdkljMzZIN41W+lC3N2tnIh5sFi+IEE17M5jbnwPHcY=
 github.com/pierrec/lz4/v4 v4.1.6 h1:ueMTcBBFrbT8K4uGDNNZPa8Z7LtPV7Cl0TDjaeHxP44=

server/analyzer/resolve_type.go

@@ -20,6 +20,8 @@ import (
 	"github.com/dolthub/go-mysql-server/sql/plan"
 	"github.com/dolthub/go-mysql-server/sql/transform"
 
+	pgnodes "github.com/dolthub/doltgresql/server/node"
+
 	"github.com/dolthub/doltgresql/core"
 	"github.com/dolthub/doltgresql/core/id"
 	"github.com/dolthub/doltgresql/server/expression"
@@ -47,6 +49,32 @@ func ResolveTypeForNodes(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node,
 	return transform.Node(node, func(node sql.Node) (sql.Node, transform.TreeIdentity, error) {
 		var same = transform.SameTree
 		switch n := node.(type) {
+		case *plan.AddColumn:
+			col := n.Column()
+			if rt, ok := col.Type.(*pgtypes.DoltgresType); ok && !rt.IsResolvedType() {
+				dt, err := resolveType(ctx, rt)
+				if err != nil {
+					return nil, transform.NewTree, err
+				}
+				same = transform.NewTree
+				col.Type = dt
+			}
+			return node, same, nil
+		case *pgnodes.CreateFunction:
+			retType, err := resolveType(ctx, n.ReturnType)
+			if err != nil {
+				return nil, transform.NewTree, err
+			}
+			paramTypes := make([]*pgtypes.DoltgresType, len(n.ParameterTypes))
+			for i := range n.ParameterTypes {
+				paramTypes[i], err = resolveType(ctx, n.ParameterTypes[i])
+				if err != nil {
+					return nil, transform.NewTree, err
+				}
+			}
+			n.ReturnType = retType
+			n.ParameterTypes = paramTypes
+			return node, transform.NewTree, nil
 		case *plan.CreateTable:
 			for _, col := range n.TargetSchema() {
 				if rt, ok := col.Type.(*pgtypes.DoltgresType); ok && !rt.IsResolvedType() {
@@ -59,17 +87,6 @@ func ResolveTypeForNodes(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node,
 				}
 			}
 			return node, same, nil
-		case *plan.AddColumn:
-			col := n.Column()
-			if rt, ok := col.Type.(*pgtypes.DoltgresType); ok && !rt.IsResolvedType() {
-				dt, err := resolveType(ctx, rt)
-				if err != nil {
-					return nil, transform.NewTree, err
-				}
-				same = transform.NewTree
-				col.Type = dt
-			}
-			return node, same, nil
 		case *plan.ModifyColumn:
 			col := n.NewColumn()
 			if rt, ok := col.Type.(*pgtypes.DoltgresType); ok && !rt.IsResolvedType() {
@@ -116,6 +133,9 @@ func ResolveTypeForExprs(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node,
 
 // resolveType resolves any type that is unresolved yet. (e.g.: domain types, built-in types that schema specified, etc.)
 func resolveType(ctx *sql.Context, typ *pgtypes.DoltgresType) (*pgtypes.DoltgresType, error) {
+	if typ.IsResolvedType() {
+		return typ, nil
+	}
 	schema, err := core.GetSchemaName(ctx, nil, typ.Schema())
 	if err != nil {
 		return nil, err
@@ -126,6 +146,13 @@ func resolveType(ctx *sql.Context, typ *pgtypes.DoltgresType) (*pgtypes.Doltgres
 	}
 	resolvedTyp, exists := typs.GetType(id.NewType(schema, typ.Name()))
 	if !exists {
+		// If a blank schema is provided, then we'll also try the pg_catalog, since a type is most likely to be there
+		if typ.Schema() == "" {
+			resolvedTyp, exists = typs.GetType(id.NewType("pg_catalog", typ.Name()))
+			if exists {
+				return resolvedTyp, nil
+			}
+		}
 		return nil, pgtypes.ErrTypeDoesNotExist.New(typ.Name())
 	}
 	return resolvedTyp, nil

server/ast/alter_function.go

@@ -22,7 +22,7 @@ import (
 
 // nodeAlterFunction handles *tree.AlterFunction nodes.
 func nodeAlterFunction(ctx *Context, node *tree.AlterFunction) (vitess.Statement, error) {
-	err := verifyRedundantRoutineOption(ctx, node.Options)
+	_, err := validateRoutineOptions(ctx, node.Options)
 	if err != nil {
 		return nil, err
 	}

server/ast/alter_procedure.go

@@ -22,7 +22,7 @@ import (
 
 // nodeAlterProcedure handles *tree.AlterProcedure nodes.
 func nodeAlterProcedure(ctx *Context, node *tree.AlterProcedure) (vitess.Statement, error) {
-	err := verifyRedundantRoutineOption(ctx, node.Options)
+	_, err := validateRoutineOptions(ctx, node.Options)
 	if err != nil {
 		return nil, err
 	}

server/ast/context.go

@@ -14,19 +14,24 @@
 
 package ast
 
-import "github.com/dolthub/doltgresql/server/auth"
+import (
+	"github.com/dolthub/doltgresql/postgres/parser/parser"
+	"github.com/dolthub/doltgresql/server/auth"
+)
 
 // Context contains any relevant context for the AST conversion. For example, the auth system uses the context to
 // determine which larger statement an expression exists in, which may influence how the expression should handle
 // authorization.
 type Context struct {
-	authContext *auth.AuthContext
+	authContext   *auth.AuthContext
+	originalQuery string
 }
 
 // NewContext returns a new *Context.
-func NewContext() *Context {
+func NewContext(postgresStmt parser.Statement) *Context {
 	return &Context{
-		authContext: auth.NewAuthContext(),
+		authContext:   auth.NewAuthContext(),
+		originalQuery: postgresStmt.SQL,
 	}
 }
 

server/ast/convert.go

@@ -25,7 +25,7 @@ import (
 
 // Convert converts a Postgres AST into a Vitess AST.
 func Convert(postgresStmt parser.Statement) (vitess.Statement, error) {
-	ctx := NewContext()
+	ctx := NewContext(postgresStmt)
 	switch stmt := postgresStmt.AST.(type) {
 	case *tree.AlterAggregate:
 		return nodeAlterAggregate(ctx, stmt)

server/ast/create_function.go

@@ -15,33 +15,65 @@
 package ast
 
 import (
-	"github.com/cockroachdb/errors"
+	"strings"
 
+	"github.com/cockroachdb/errors"
 	vitess "github.com/dolthub/vitess/go/vt/sqlparser"
 
 	"github.com/dolthub/doltgresql/postgres/parser/sem/tree"
+	pgnodes "github.com/dolthub/doltgresql/server/node"
+	"github.com/dolthub/doltgresql/server/plpgsql"
+	pgtypes "github.com/dolthub/doltgresql/server/types"
 )
 
 // nodeCreateFunction handles *tree.CreateFunction nodes.
 func nodeCreateFunction(ctx *Context, node *tree.CreateFunction) (vitess.Statement, error) {
-	err := verifyRedundantRoutineOption(ctx, node.Options)
+	options, err := validateRoutineOptions(ctx, node.Options)
 	if err != nil {
 		return nil, err
 	}
-
-	return NotYetSupportedError("CREATE FUNCTION statement is not yet supported")
-}
-
-// verifyRedundantRoutineOption checks for each option defined only once.
-// If there is multiple definition of the same option, it returns an error.
-func verifyRedundantRoutineOption(ctx *Context, options []tree.RoutineOption) error {
-	var optDefined = make(map[tree.FunctionOption]struct{})
-	for _, opt := range options {
-		if _, ok := optDefined[opt.OptionType]; ok {
-			return errors.Errorf("ERROR:  conflicting or redundant options")
-		} else {
-			optDefined[opt.OptionType] = struct{}{}
+	// We only support PL/pgSQL for now, so we'll verify that first
+	if languageOption, ok := options[tree.OptionLanguage]; ok {
+		if strings.ToLower(languageOption.Language) != "plpgsql" {
+			return nil, errors.Errorf("CREATE FUNCTION only supports PL/pgSQL for now")
 		}
+	} else {
+		return nil, errors.Errorf("CREATE FUNCTION does not define an input language")
+	}
+	// PL/pgSQL is different from standard Postgres SQL, so we have to use a special parser to handle it.
+	// This parser also requires the full `CREATE FUNCTION` string, so we'll pass that.
+	parsedBody, err := plpgsql.Parse(ctx.originalQuery)
+	if err != nil {
+		return nil, err
+	}
+	// Grab the rest of the information that we'll need to create the function
+	tableName := node.Name.ToTableName()
+	schemaName := tableName.Schema()
+	if len(schemaName) == 0 {
+		// TODO: fix function finder such that it doesn't always assume pg_catalog
+		schemaName = "pg_catalog"
+	}
+	retType := pgtypes.Void
+	if len(node.RetType) == 1 {
+		retType = pgtypes.NewUnresolvedDoltgresType("", strings.ToLower(node.RetType[0].Type.SQLString()))
+	}
+	paramNames := make([]string, len(node.Args))
+	paramTypes := make([]*pgtypes.DoltgresType, len(node.Args))
+	for i, arg := range node.Args {
+		paramNames[i] = arg.Name.String()
+		paramTypes[i] = pgtypes.NewUnresolvedDoltgresType("", strings.ToLower(arg.Type.SQLString()))
 	}
-	return nil
+	// Returns the stored procedure call with all options
+	return vitess.InjectedStatement{
+		Statement: pgnodes.NewCreateFunction(
+			tableName.Table(),
+			schemaName,
+			retType,
+			paramNames,
+			paramTypes,
+			true, // TODO: implement strict check
+			parsedBody,
+		),
+		Children: nil,
+	}, nil
 }

server/ast/create_procedure.go

@@ -15,17 +15,33 @@
 package ast
 
 import (
+	"github.com/cockroachdb/errors"
+
 	vitess "github.com/dolthub/vitess/go/vt/sqlparser"
 
 	"github.com/dolthub/doltgresql/postgres/parser/sem/tree"
 )
 
 // nodeCreateProcedure handles *tree.CreateProcedure nodes.
 func nodeCreateProcedure(ctx *Context, node *tree.CreateProcedure) (vitess.Statement, error) {
-	err := verifyRedundantRoutineOption(ctx, node.Options)
+	_, err := validateRoutineOptions(ctx, node.Options)
 	if err != nil {
 		return nil, err
 	}
 
 	return NotYetSupportedError("CREATE PROCEDURE statement is not yet supported")
 }
+
+// validateRoutineOptions ensures that each option is defined only once. Returns a map containing all options, or an
+// error if an option is invalid or is defined multiple times.
+func validateRoutineOptions(ctx *Context, options []tree.RoutineOption) (map[tree.FunctionOption]tree.RoutineOption, error) {
+	var optDefined = make(map[tree.FunctionOption]tree.RoutineOption)
+	for _, opt := range options {
+		if _, ok := optDefined[opt.OptionType]; ok {
+			return nil, errors.Errorf("ERROR:  conflicting or redundant options")
+		} else {
+			optDefined[opt.OptionType] = opt
+		}
+	}
+	return optDefined, nil
+}