catalog: support functions with table parameters (#541)

Co-authored-by: Kyle Conroy <kyle@conroy.org>

Author: maxhawkins

internal/endtoend/testdata/func_args/endtoend.json

@@ -0,0 +1,4 @@
+{
+    "experimental_parser_only": true
+}
+  

internal/endtoend/testdata/func_args/go/query.sql.go

@@ -4,6 +4,8 @@ CREATE FUNCTION plus(a integer, b integer) RETURNS integer AS $$
     END;
 $$ LANGUAGE plpgsql;
 
+CREATE FUNCTION table_args(x INT) RETURNS TABLE (y INT) AS 'SELECT x' LANGUAGE sql;
+
 -- name: Plus :one
 SELECT plus(b => $2, a => $1);
 
@@ -16,4 +18,5 @@ SELECT make_interval(days => $1::int);
 -- name: MakeIntervalMonths :one
 SELECT make_interval(months => sqlc.arg('months')::int);
 
-
+-- name: TableArgs :one
+SELECT table_args(x => $1);

internal/endtoend/testdata/func_args/query.sql

@@ -42,6 +42,23 @@ func parseFuncName(node nodes.Node) (*ast.FuncName, error) {
 	}, nil
 }
 
+func parseFuncParamMode(m nodes.FunctionParameterMode) (ast.FuncParamMode, error) {
+	switch m {
+	case 'i':
+		return ast.FuncParamIn, nil
+	case 'o':
+		return ast.FuncParamOut, nil
+	case 'b':
+		return ast.FuncParamInOut, nil
+	case 'v':
+		return ast.FuncParamVariadic, nil
+	case 't':
+		return ast.FuncParamTable, nil
+	default:
+		return -1, fmt.Errorf("parse func param: invalid mode %v", m)
+	}
+}
+
 func parseTypeName(node nodes.Node) (*ast.TypeName, error) {
 	rel, err := parseRelation(node)
 	if err != nil {
@@ -434,9 +451,14 @@ func translate(node nodes.Node) (ast.Node, error) {
 			if err != nil {
 				return nil, err
 			}
+			mode, err := parseFuncParamMode(arg.Mode)
+			if err != nil {
+				return nil, err
+			}
 			fp := &ast.FuncParam{
 				Name: arg.Name,
 				Type: tn,
+				Mode: mode,
 			}
 			if arg.Defexpr != nil {
 				fp.DefExpr = &ast.TODO{}

internal/postgresql/parse.go

@@ -1,9 +1,20 @@
 package ast
 
+type FuncParamMode int
+
+const (
+	FuncParamIn FuncParamMode = iota
+	FuncParamOut
+	FuncParamInOut
+	FuncParamVariadic
+	FuncParamTable
+)
+
 type FuncParam struct {
 	Name    *string
 	Type    *TypeName
 	DefExpr Node // Will always be &ast.TODO
+	Mode    FuncParamMode
 }
 
 func (n *FuncParam) Pos() int {

internal/sql/ast/func_param.go

@@ -106,11 +106,13 @@ func (s *Schema) getFunc(rel *ast.FuncName, tns []*ast.TypeName) (*Function, int
 		if s.Funcs[i].Name != rel.Name {
 			continue
 		}
-		if len(s.Funcs[i].Args) != len(tns) {
+
+		args := s.Funcs[i].InArgs()
+		if len(args) != len(tns) {
 			continue
 		}
 		found := true
-		for j := range s.Funcs[i].Args {
+		for j := range args {
 			if !sameType(s.Funcs[i].Args[j].Type, tns[j]) {
 				found = false
 				break
@@ -215,10 +217,24 @@ type Function struct {
 	Desc       string
 }
 
+func (f *Function) InArgs() []*Argument {
+	var args []*Argument
+	for _, a := range f.Args {
+		switch a.Mode {
+		case ast.FuncParamTable, ast.FuncParamOut:
+			continue
+		default:
+			args = append(args, a)
+		}
+	}
+	return args
+}
+
 type Argument struct {
 	Name       string
 	Type       *ast.TypeName
 	HasDefault bool
+	Mode       ast.FuncParamMode
 }
 
 func New(def string) *Catalog {

internal/sql/catalog/catalog.go

@@ -30,6 +30,7 @@ func (c *Catalog) createFunction(stmt *ast.CreateFunctionStmt) error {
 		fn.Args[i] = &Argument{
 			Name:       name,
 			Type:       arg.Type,
+			Mode:       arg.Mode,
 			HasDefault: arg.DefExpr != nil,
 		}
 		types[i] = arg.Type

internal/sql/catalog/func.go

@@ -38,7 +38,9 @@ func (c *Catalog) GetFuncN(rel *ast.FuncName, n int) (Function, error) {
 			if s.Funcs[i].Name != rel.Name {
 				continue
 			}
-			if len(s.Funcs[i].Args) == n {
+
+			args := s.Funcs[i].InArgs()
+			if len(args) == n {
 				return *s.Funcs[i], nil
 			}
 		}

internal/sql/catalog/public.go

@@ -39,7 +39,7 @@ func (v *funcCallVisitor) Visit(node ast.Node) astutils.Visitor {
 		args = len(call.Args.Items)
 	}
 	for _, fun := range funs {
-		if len(fun.Args) == args {
+		if len(fun.InArgs()) == args {
 			return v
 		}
 	}