Lookup composite types when selecting from function

Author: wickwirew

internal/compiler/output_columns.go

@@ -544,6 +544,11 @@ func (c *Compiler) sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, erro
 					Schema:  fn.ReturnType.Schema,
 					Name:    fn.ReturnType.Name,
 				})
+
+				// Successfully found the table
+				if err != nil {
+					table, err = qc.GetCompositeType(fn.ReturnType)
+				}
 			}
 			if table == nil || err != nil {
 				if n.Alias != nil && len(n.Alias.Colnames.Items) > 0 {

internal/compiler/query_catalog.go

@@ -107,3 +107,23 @@ func (qc QueryCatalog) GetFunc(rel *ast.FuncName) (*Function, error) {
 		ReturnType: funcs[0].ReturnType,
 	}, nil
 }
+
+func (qc QueryCatalog) GetCompositeType(rel *ast.TypeName) (*Table, error) {
+	ty, err := qc.catalog.GetCompostiteType(rel)
+	if err != nil {
+		return &Table{}, err
+	}
+
+	tblName := &ast.TableName{
+		Catalog: rel.Catalog,
+		Schema:  rel.Schema,
+		Name:    rel.Name,
+	}
+
+	var cols []*Column
+	for _, tyCol := range ty.Columns {
+		cols = append(cols, ConvertColumn(tblName, tyCol))
+	}
+
+	return &Table{Rel: tblName, Columns: cols}, nil
+}

internal/engine/postgresql/parse.go

@@ -392,9 +392,31 @@ func translate(node *nodes.Node) (ast.Node, error) {
 	case *nodes.Node_CompositeTypeStmt:
 		n := inner.CompositeTypeStmt
 		rel := parseRelationFromRangeVar(n.Typevar)
-		return &ast.CompositeTypeStmt{
+
+		stmt := &ast.CompositeTypeStmt{
 			TypeName: rel.TypeName(),
-		}, nil
+		}
+
+		for _, elt := range n.GetColdeflist() {
+			switch item := elt.Node.(type) {
+			case *nodes.Node_ColumnDef:
+				rel, err := parseRelationFromNodes(item.ColumnDef.TypeName.Names)
+				if err != nil {
+					return nil, err
+				}
+
+				stmt.Cols = append(stmt.Cols, &ast.ColumnDef{
+					Colname:    item.ColumnDef.Colname,
+					TypeName:   rel.TypeName(),
+					IsNotNull:  false, // Composite types cannot have constraints
+					IsArray:    isArray(item.ColumnDef.TypeName),
+					ArrayDims:  len(item.ColumnDef.TypeName.ArrayBounds),
+					PrimaryKey: false,
+				})
+			}
+		}
+
+		return stmt, nil
 
 	case *nodes.Node_CreateStmt:
 		n := inner.CreateStmt

internal/sql/ast/composite_type_stmt.go

@@ -2,6 +2,7 @@ package ast
 
 type CompositeTypeStmt struct {
 	TypeName *TypeName
+	Cols     []*ColumnDef
 }
 
 func (n *CompositeTypeStmt) Pos() int {

internal/sql/catalog/public.go

@@ -131,3 +131,21 @@ func (c *Catalog) GetTable(rel *ast.TableName) (Table, error) {
 		return *table, err
 	}
 }
+
+func (c *Catalog) GetCompostiteType(rel *ast.TypeName) (CompositeType, error) {
+	ty, _, err := c.getType(rel)
+	if err != nil {
+		return CompositeType{}, err
+	}
+
+	cTy, ok := ty.(*CompositeType)
+	if !ok {
+		return CompositeType{}, fmt.Errorf("Type %s is not a composite type", rel.Name)
+	}
+
+	if ty == nil {
+		return CompositeType{}, err
+	} else {
+		return *cTy, err
+	}
+}

internal/sql/catalog/types.go

@@ -30,6 +30,7 @@ func (e *Enum) isType() {
 type CompositeType struct {
 	Name    string
 	Comment string
+	Columns []*Column
 }
 
 func (ct *CompositeType) isType() {
@@ -135,8 +136,24 @@ func (c *Catalog) createCompositeType(stmt *ast.CompositeTypeStmt) error {
 	if _, _, err := schema.getType(stmt.TypeName); err == nil {
 		return sqlerr.TypeExists(tbl.Name)
 	}
+
+	var cols []*Column
+	for _, cDef := range stmt.Cols {
+		cols = append(cols, &Column{
+			Name:       cDef.Colname,
+			Type:       *cDef.TypeName,
+			IsNotNull:  cDef.IsNotNull,
+			IsUnsigned: cDef.IsUnsigned,
+			IsArray:    cDef.IsArray,
+			ArrayDims:  cDef.ArrayDims,
+			Comment:    cDef.Comment,
+			Length:     cDef.Length,
+		})
+	}
+
 	schema.Types = append(schema.Types, &CompositeType{
-		Name: stmt.TypeName.Name,
+		Name:    stmt.TypeName.Name,
+		Columns: cols,
 	})
 	return nil
 }