Pass full column definition instead of just typename

Author: eaglesemanation

internal/cmd/shim.go

@@ -6,7 +6,6 @@ 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"
 )
 
@@ -60,16 +59,38 @@ 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,
+func columnSlice(cols []*catalog.Column, table *catalog.Table) []*plugin.Column {
+	var columns []*plugin.Column
+	for _, c := range cols {
+		l := -1
+		if c.Length != nil {
+			l = *c.Length
+		}
+		var tableId *plugin.Identifier = nil
+		if table != nil {
+			tableId = &plugin.Identifier{
+				Catalog: table.Rel.Catalog,
+				Schema:  table.Rel.Schema,
+				Name:    table.Rel.Name,
+			}
+		}
+		columns = append(columns, &plugin.Column{
+			Name: c.Name,
+			Type: &plugin.Identifier{
+				Catalog: c.Type.Catalog,
+				Schema:  c.Type.Schema,
+				Name:    c.Type.Name,
+			},
+			Comment:   c.Comment,
+			NotNull:   c.IsNotNull,
+			Unsigned:  c.IsUnsigned,
+			IsArray:   c.IsArray,
+			ArrayDims: int32(c.ArrayDims),
+			Length:    int32(l),
+			Table:     tableId,
 		})
 	}
-	return ids
+	return columns
 }
 
 func pluginCatalog(c *catalog.Catalog) *plugin.Catalog {
@@ -87,47 +108,21 @@ func pluginCatalog(c *catalog.Catalog) *plugin.Catalog {
 				})
 			case *catalog.CompositeType:
 				cts = append(cts, &plugin.CompositeType{
-					Name:         typ.Name,
-					Comment:      typ.Comment,
-					ColTypeNames: identifierSlice(typ.ColTypeNames),
+					Name:    typ.Name,
+					Comment: typ.Comment,
+					Columns: columnSlice(typ.Columns, nil),
 				})
 			}
 		}
 		var tables []*plugin.Table
 		for _, t := range s.Tables {
-			var columns []*plugin.Column
-			for _, c := range t.Columns {
-				l := -1
-				if c.Length != nil {
-					l = *c.Length
-				}
-				columns = append(columns, &plugin.Column{
-					Name: c.Name,
-					Type: &plugin.Identifier{
-						Catalog: c.Type.Catalog,
-						Schema:  c.Type.Schema,
-						Name:    c.Type.Name,
-					},
-					Comment:   c.Comment,
-					NotNull:   c.IsNotNull,
-					Unsigned:  c.IsUnsigned,
-					IsArray:   c.IsArray,
-					ArrayDims: int32(c.ArrayDims),
-					Length:    int32(l),
-					Table: &plugin.Identifier{
-						Catalog: t.Rel.Catalog,
-						Schema:  t.Rel.Schema,
-						Name:    t.Rel.Name,
-					},
-				})
-			}
 			tables = append(tables, &plugin.Table{
 				Rel: &plugin.Identifier{
 					Catalog: t.Rel.Catalog,
 					Schema:  t.Rel.Schema,
 					Name:    t.Rel.Name,
 				},
-				Columns: columns,
+				Columns: columnSlice(t.Columns, t),
 				Comment: t.Comment,
 			})
 		}

internal/plugin/codegen.pb.go

@@ -28,9 +28,9 @@ func (e *Enum) isType() {
 }
 
 type CompositeType struct {
-	Name         string
-	ColTypeNames []*ast.TypeName
-	Comment      string
+	Name    string
+	Columns []*Column
+	Comment string
 }
 
 func (ct *CompositeType) isType() {
@@ -102,11 +102,11 @@ func stringSlice(list *ast.List) []string {
 	return items
 }
 
-func columnTypeNamesSlice(list *ast.List) []*ast.TypeName {
-	items := []*ast.TypeName{}
+func columnDefsSlice(list *ast.List) []*ast.ColumnDef {
+	items := []*ast.ColumnDef{}
 	for _, item := range list.Items {
 		if n, ok := item.(*ast.ColumnDef); ok {
-			items = append(items, n.TypeName)
+			items = append(items, n)
 		}
 	}
 	return items
@@ -146,10 +146,22 @@ func (c *Catalog) createCompositeType(stmt *ast.CompositeTypeStmt) error {
 	if _, _, err := schema.getType(stmt.TypeName); err == nil {
 		return sqlerr.TypeExists(tbl.Name)
 	}
-	schema.Types = append(schema.Types, &CompositeType{
-		Name:         stmt.TypeName.Name,
-		ColTypeNames: columnTypeNamesSlice(stmt.ColDefList),
-	})
+	ct := &CompositeType{
+		Name: stmt.TypeName.Name,
+	}
+	for _, col := range columnDefsSlice(stmt.ColDefList) {
+		ct.Columns = append(ct.Columns, &Column{
+			Name:       col.Colname,
+			Type:       *col.TypeName,
+			IsNotNull:  col.IsNotNull,
+			IsUnsigned: col.IsUnsigned,
+			IsArray:    col.IsArray,
+			ArrayDims:  col.ArrayDims,
+			Comment:    col.Comment,
+			Length:     col.Length,
+		})
+	}
+	schema.Types = append(schema.Types, ct)
 	return nil
 }
 
@@ -350,9 +362,9 @@ func (c *Catalog) renameType(stmt *ast.RenameTypeStmt) error {
 
 	case *CompositeType:
 		schema.Types[idx] = &CompositeType{
-			Name:         newName,
-			ColTypeNames: typ.ColTypeNames,
-			Comment:      typ.Comment,
+			Name:    newName,
+			Columns: typ.Columns,
+			Comment: typ.Comment,
 		}
 
 	case *Enum:
@@ -390,8 +402,8 @@ func (c *Catalog) updateTypeNames(typeUpdater func(t *ast.TypeName)) error {
 			if !ok {
 				continue
 			}
-			for _, fieldType := range composite.ColTypeNames {
-				typeUpdater(fieldType)
+			for _, fieldType := range composite.Columns {
+				typeUpdater(&fieldType.Type)
 			}
 		}
 	}

internal/sql/catalog/types.go

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