Propagate composite type field types from parser to plugin proto

Author: eaglesemanation

internal/cmd/shim.go

@@ -6,6 +6,7 @@ import (
 	"github.com/sqlc-dev/sqlc/internal/config/convert"
 	"github.com/sqlc-dev/sqlc/internal/info"
 	"github.com/sqlc-dev/sqlc/internal/plugin"
+	"github.com/sqlc-dev/sqlc/internal/sql/ast"
 	"github.com/sqlc-dev/sqlc/internal/sql/catalog"
 )
 
@@ -59,6 +60,18 @@ func pluginWASM(p config.Plugin) *plugin.Codegen_WASM {
 	return nil
 }
 
+func identifierSlice(types []*ast.TypeName) []*plugin.Identifier {
+	ids := []*plugin.Identifier{}
+	for _, typ := range types {
+		ids = append(ids, &plugin.Identifier{
+			Catalog: typ.Catalog,
+			Schema:  typ.Schema,
+			Name:    typ.Name,
+		})
+	}
+	return ids
+}
+
 func pluginCatalog(c *catalog.Catalog) *plugin.Catalog {
 	var schemas []*plugin.Schema
 	for _, s := range c.Schemas {
@@ -74,8 +87,9 @@ func pluginCatalog(c *catalog.Catalog) *plugin.Catalog {
 				})
 			case *catalog.CompositeType:
 				cts = append(cts, &plugin.CompositeType{
-					Name:    typ.Name,
-					Comment: typ.Comment,
+					Name:         typ.Name,
+					Comment:      typ.Comment,
+					ColTypeNames: identifierSlice(typ.ColTypeNames),
 				})
 			}
 		}

internal/engine/postgresql/convert.go

@@ -889,7 +889,8 @@ func convertCompositeTypeStmt(n *pg.CompositeTypeStmt) *ast.CompositeTypeStmt {
 	}
 	rel := parseRelationFromRangeVar(n.Typevar)
 	return &ast.CompositeTypeStmt{
-		TypeName: rel.TypeName(),
+		TypeName:   rel.TypeName(),
+		ColDefList: convertSlice(n.GetColdeflist()),
 	}
 }
 

internal/engine/postgresql/parse.go

@@ -392,9 +392,26 @@ func translate(node *nodes.Node) (ast.Node, error) {
 	case *nodes.Node_CompositeTypeStmt:
 		n := inner.CompositeTypeStmt
 		rel := parseRelationFromRangeVar(n.Typevar)
-		return &ast.CompositeTypeStmt{
-			TypeName: rel.TypeName(),
-		}, nil
+		stmt := &ast.CompositeTypeStmt{
+			TypeName:   rel.TypeName(),
+			ColDefList: &ast.List{},
+		}
+		for _, val := range n.GetColdeflist() {
+			switch item := val.Node.(type) {
+			case *nodes.Node_ColumnDef:
+				rel, err := parseRelationFromNodes(item.ColumnDef.TypeName.Names)
+				if err != nil {
+					return nil, err
+				}
+				stmt.ColDefList.Items = append(stmt.ColDefList.Items, &ast.ColumnDef{
+					Colname:    item.ColumnDef.GetColname(),
+					TypeName:   rel.TypeName(),
+					IsLocal:    true,
+					CollClause: convertCollateClause(item.ColumnDef.GetCollClause()),
+				})
+			}
+		}
+		return stmt, nil
 
 	case *nodes.Node_CreateStmt:
 		n := inner.CreateStmt

internal/plugin/codegen.pb.go

@@ -1,7 +1,8 @@
 package ast
 
 type CompositeTypeStmt struct {
-	TypeName *TypeName
+	TypeName   *TypeName
+	ColDefList *List
 }
 
 func (n *CompositeTypeStmt) Pos() int {

internal/sql/ast/composite_type_stmt.go

@@ -495,6 +495,7 @@ func (a *application) apply(parent ast.Node, name string, iter *iterator, n ast.
 
 	case *ast.CompositeTypeStmt:
 		a.apply(n, "TypeName", nil, n.TypeName)
+		a.apply(n, "ColDefList", nil, n.ColDefList)
 
 	case *ast.Const:
 		a.apply(n, "Xpr", nil, n.Xpr)

internal/sql/astutils/rewrite.go

@@ -28,8 +28,9 @@ func (e *Enum) isType() {
 }
 
 type CompositeType struct {
-	Name    string
-	Comment string
+	Name         string
+	ColTypeNames []*ast.TypeName
+	Comment      string
 }
 
 func (ct *CompositeType) isType() {
@@ -101,6 +102,16 @@ func stringSlice(list *ast.List) []string {
 	return items
 }
 
+func columnTypeNamesSlice(list *ast.List) []*ast.TypeName {
+	items := []*ast.TypeName{}
+	for _, item := range list.Items {
+		if n, ok := item.(*ast.ColumnDef); ok {
+			items = append(items, n.TypeName)
+		}
+	}
+	return items
+}
+
 func (c *Catalog) getType(rel *ast.TypeName) (Type, int, error) {
 	ns := rel.Schema
 	if ns == "" {
@@ -136,7 +147,8 @@ func (c *Catalog) createCompositeType(stmt *ast.CompositeTypeStmt) error {
 		return sqlerr.TypeExists(tbl.Name)
 	}
 	schema.Types = append(schema.Types, &CompositeType{
-		Name: stmt.TypeName.Name,
+		Name:         stmt.TypeName.Name,
+		ColTypeNames: columnTypeNamesSlice(stmt.ColDefList),
 	})
 	return nil
 }
@@ -277,16 +289,11 @@ func (c *Catalog) alterTypeSetSchema(stmt *ast.AlterTypeSetSchemaStmt) error {
 	oldSchema.Types = append(oldSchema.Types[:idx], oldSchema.Types[idx+1:]...)
 	newSchema.Types = append(newSchema.Types, typ)
 
-	// Update all the table columns with the new type
-	for _, schema := range c.Schemas {
-		for _, table := range schema.Tables {
-			for _, column := range table.Columns {
-				if column.Type == oldType {
-					column.Type.Schema = *stmt.NewSchema
-				}
-			}
+	c.updateTypeNames(func(t *ast.TypeName) {
+		if *t == oldType {
+			t.Schema = *stmt.NewSchema
 		}
-	}
+	})
 	return nil
 }
 
@@ -343,8 +350,9 @@ func (c *Catalog) renameType(stmt *ast.RenameTypeStmt) error {
 
 	case *CompositeType:
 		schema.Types[idx] = &CompositeType{
-			Name:    newName,
-			Comment: typ.Comment,
+			Name:         newName,
+			ColTypeNames: typ.ColTypeNames,
+			Comment:      typ.Comment,
 		}
 
 	case *Enum:
@@ -359,16 +367,33 @@ func (c *Catalog) renameType(stmt *ast.RenameTypeStmt) error {
 
 	}
 
-	// Update all the table columns with the new type
+	c.updateTypeNames(func(t *ast.TypeName) {
+		if *t == *stmt.Type {
+			t.Name = newName
+		}
+	})
+
+	return nil
+}
+
+func (c *Catalog) updateTypeNames(typeUpdater func(t *ast.TypeName)) error {
 	for _, schema := range c.Schemas {
+		// Update all the table columns with the new type
 		for _, table := range schema.Tables {
 			for _, column := range table.Columns {
-				if column.Type == *stmt.Type {
-					column.Type.Name = newName
-				}
+				typeUpdater(&column.Type)
+			}
+		}
+		// Update all the composite fields with the new type
+		for _, typ := range schema.Types {
+			composite, ok := typ.(*CompositeType)
+			if !ok {
+				continue
+			}
+			for _, fieldType := range composite.ColTypeNames {
+				typeUpdater(fieldType)
 			}
 		}
 	}
-
 	return nil
 }

internal/sql/catalog/types.go

@@ -61,6 +61,7 @@ message Schema {
 message CompositeType {
   string name = 1;
   string comment = 2;
+  repeated Identifier col_type_names = 3;
 }
 
 message Enum {