From 54101cbd22bec0cede7fe3c6c02d6f7aaec5992e Mon Sep 17 00:00:00 2001 From: "John W. Carbone" Date: Mon, 10 Jul 2023 14:47:37 -0400 Subject: [PATCH] [DRAFT] Support for PostgreSQL "RECORD" return types In this draft change, we've added support for PostgreSQL functions that return the "RECORD" data type. --- internal/compiler/output_columns.go | 21 ++++++++++++++-- internal/compiler/query.go | 6 +++-- internal/compiler/query_catalog.go | 17 +++++++++++-- internal/engine/postgresql/parse.go | 25 +++++++++++++++++++ internal/sql/ast/create_function_stmt.go | 5 ++-- internal/sql/catalog/func.go | 31 ++++++++++++++++++++++++ 6 files changed, 97 insertions(+), 8 deletions(-) diff --git a/internal/compiler/output_columns.go b/internal/compiler/output_columns.go index 26a421f025..f03ebf8dd1 100644 --- a/internal/compiler/output_columns.go +++ b/internal/compiler/output_columns.go @@ -515,13 +515,30 @@ func (c *Compiler) sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, erro } fn, err := qc.GetFunc(funcCall.Func) - if err != nil { + if err == nil { + table := &Table{ + Rel: &ast.TableName{ + Catalog: fn.ReturnTable.Rel.Catalog, + Schema: fn.ReturnTable.Rel.Schema, + //Name: fn.ReturnType.Name, + Name: fn.ReturnTable.Rel.Name, + }, + } + + for _, c := range fn.Columns { + table.Columns = append(table.Columns, &Column{ + Name: c.Name, + DataType: c.DataType, + }) + } + tables = append(tables, table) continue } table, err := qc.GetTable(&ast.TableName{ Catalog: fn.ReturnType.Catalog, Schema: fn.ReturnType.Schema, - Name: fn.ReturnType.Name, + //Name: fn.ReturnType.Name, + Name: fn.Rel.Name, }) if err != nil { if n.Alias == nil || len(n.Alias.Colnames.Items) == 0 { diff --git a/internal/compiler/query.go b/internal/compiler/query.go index e77f555dbd..65b693ef19 100644 --- a/internal/compiler/query.go +++ b/internal/compiler/query.go @@ -5,8 +5,10 @@ import ( ) type Function struct { - Rel *ast.FuncName - ReturnType *ast.TypeName + Rel *ast.FuncName + ReturnType *ast.TypeName + Columns []*Column + ReturnTable *Table } type Table struct { diff --git a/internal/compiler/query_catalog.go b/internal/compiler/query_catalog.go index 2b6577c2e9..1b0846fbdf 100644 --- a/internal/compiler/query_catalog.go +++ b/internal/compiler/query_catalog.go @@ -87,8 +87,21 @@ func (qc QueryCatalog) GetFunc(rel *ast.FuncName) (*Function, error) { if len(funcs) == 0 { return nil, fmt.Errorf("function not found: %s", rel.Name) } + + var funcRecord *Table + var cols []*Column + rt := funcs[0].ReturnTable + if rt != nil { + for _, c := range rt.Columns { + cols = append(cols, ConvertColumn(rt.Rel, c)) + } + funcRecord = &Table{Rel: rt.Rel, Columns: cols} + } + return &Function{ - Rel: rel, - ReturnType: funcs[0].ReturnType, + Rel: rel, + ReturnType: funcs[0].ReturnType, + ReturnTable: funcRecord, + Columns: cols, }, nil } diff --git a/internal/engine/postgresql/parse.go b/internal/engine/postgresql/parse.go index fa1a54a911..23db91ce0c 100644 --- a/internal/engine/postgresql/parse.go +++ b/internal/engine/postgresql/parse.go @@ -504,6 +504,9 @@ func translate(node *nodes.Node) (ast.Node, error) { } stmt.Params.Items = append(stmt.Params.Items, fp) } + if stmt.ReturnType != nil && stmt.ReturnType.Name == "record" { + stmt.ReturnTable = funcReturnTable(stmt) + } return stmt, nil case *nodes.Node_CreateSchemaStmt: @@ -629,3 +632,25 @@ func translate(node *nodes.Node) (ast.Node, error) { return convert(node) } } + +// makes a pseudo table from a function for table return type +func funcReturnTable(stmt *ast.CreateFunctionStmt) *ast.CreateTableStmt { + name := ast.TableName{ + Name: stmt.Func.Name, + Catalog: stmt.Func.Catalog, + Schema: stmt.Func.Schema, + } + table := &ast.CreateTableStmt{ + Name: &name, + } + for _, param := range stmt.Params.Items { + funcParam := param.(*ast.FuncParam) + if funcParam.Mode == ast.FuncParamTable { + table.Cols = append(table.Cols, &ast.ColumnDef{ + Colname: *funcParam.Name, + TypeName: funcParam.Type, + }) + } + } + return table +} diff --git a/internal/sql/ast/create_function_stmt.go b/internal/sql/ast/create_function_stmt.go index a1e0b8f728..e95c4d4a3b 100644 --- a/internal/sql/ast/create_function_stmt.go +++ b/internal/sql/ast/create_function_stmt.go @@ -6,8 +6,9 @@ type CreateFunctionStmt struct { ReturnType *TypeName Func *FuncName // TODO: Undertand these two fields - Options *List - WithClause *List + Options *List + WithClause *List + ReturnTable *CreateTableStmt } func (n *CreateFunctionStmt) Pos() int { diff --git a/internal/sql/catalog/func.go b/internal/sql/catalog/func.go index 80acd1da50..effa51c371 100644 --- a/internal/sql/catalog/func.go +++ b/internal/sql/catalog/func.go @@ -2,6 +2,7 @@ package catalog import ( "errors" + "fmt" "github.com/kyleconroy/sqlc/internal/sql/ast" "github.com/kyleconroy/sqlc/internal/sql/sqlerr" @@ -17,6 +18,7 @@ type Function struct { Comment string Desc string ReturnTypeNullable bool + ReturnTable *Table } type Argument struct { @@ -53,6 +55,9 @@ func (c *Catalog) createFunction(stmt *ast.CreateFunctionStmt) error { Args: make([]*Argument, len(stmt.Params.Items)), ReturnType: stmt.ReturnType, } + if stmt.ReturnType.Name == "record" { + fn.ReturnTable = createFuncReturnTable(stmt.ReturnTable) + } types := make([]*ast.TypeName, len(stmt.Params.Items)) for i, item := range stmt.Params.Items { arg := item.(*ast.FuncParam) @@ -82,6 +87,32 @@ func (c *Catalog) createFunction(stmt *ast.CreateFunctionStmt) error { return nil } +func createFuncReturnTable(stmt *ast.CreateTableStmt) *Table { + tbl := Table{ + Rel: stmt.Name, + Comment: stmt.Comment, + } + for _, col := range stmt.Cols { + tc := &Column{ + Name: col.Colname, + Type: *col.TypeName, + IsNotNull: col.IsNotNull, + IsUnsigned: col.IsUnsigned, + IsArray: col.IsArray, + Comment: col.Comment, + Length: col.Length, + } + if col.Vals != nil { + typeName := ast.TypeName{ + Name: fmt.Sprintf("%s_%s", stmt.Name.Name, col.Colname), + } + tc.Type = typeName + } + tbl.Columns = append(tbl.Columns, tc) + } + return &tbl +} + func (c *Catalog) dropFunction(stmt *ast.DropFunctionStmt) error { for _, spec := range stmt.Funcs { ns := spec.Name.Schema