@@ -25,6 +25,7 @@ import (
2525 "github.com/dolthub/dolt/go/store/hash"
2626 "github.com/dolthub/dolt/go/store/prolly"
2727 "github.com/dolthub/dolt/go/store/prolly/tree"
28+ "github.com/dolthub/go-mysql-server/sql"
2829
2930 "github.com/dolthub/doltgresql/core/id"
3031 "github.com/dolthub/doltgresql/core/rootobject/objinterface"
@@ -101,8 +102,10 @@ func (pgs *TypeCollection) DropType(ctx context.Context, names ...id.Type) (err
101102}
102103
103104// GetAllTypes returns a map containing all types in the collection, grouped by the schema they're contained in.
104- // Each type array is also sorted by the type name. It includes built-in types.
105+ // Each type array is also sorted by the type name. It includes built-in types, but does not include types referring to
106+ // a table's row type.
105107func (pgs * TypeCollection ) GetAllTypes (ctx context.Context ) (typeMap map [string ][]* pgtypes.DoltgresType , schemaNames []string , totalCount int , err error ) {
108+ // TODO: this should probably get tables as well since tables create composite types matching their rows
106109 schemaNamesMap := make (map [string ]struct {})
107110 typeMap = make (map [string ][]* pgtypes.DoltgresType )
108111 err = pgs .IterateTypes (ctx , func (t * pgtypes.DoltgresType ) (stop bool , err error ) {
@@ -158,9 +161,21 @@ func (pgs *TypeCollection) GetType(ctx context.Context, name id.Type) (*pgtypes.
158161 }
159162 // The initial load is from the internal map
160163 h , err := pgs .underlyingMap .Get (ctx , string (name ))
161- if err != nil || h . IsEmpty () {
164+ if err != nil {
162165 return nil , err
163166 }
167+ if h .IsEmpty () {
168+ // If it's not a built-in type or created type, then check if it's a composite table row type
169+ sqlCtx , ok := ctx .(* sql.Context )
170+ if ! ok {
171+ return nil , nil
172+ }
173+ tbl , schema , err := pgs .getTable (sqlCtx , name .SchemaName (), name .TypeName ())
174+ if err != nil || tbl == nil {
175+ return nil , err
176+ }
177+ return pgs .tableToType (sqlCtx , tbl , schema )
178+ }
164179 data , err := pgs .ns .ReadBytes (ctx , h )
165180 if err != nil {
166181 return nil , err
@@ -180,19 +195,26 @@ func (pgs *TypeCollection) HasType(ctx context.Context, name id.Type) bool {
180195 if _ , ok := pgtypes .IDToBuiltInDoltgresType [name ]; ok {
181196 return true
182197 }
183-
198+ // Now we'll check our created types
184199 if _ , ok := pgs .accessedMap [name ]; ok {
185200 return true
186201 }
187202 ok , err := pgs .underlyingMap .Has (ctx , string (name ))
188203 if err == nil && ok {
189204 return true
190205 }
191- return false
206+ // If it's not a built-in type or created type, then check if it's a composite table row type
207+ sqlCtx , ok := ctx .(* sql.Context )
208+ if ! ok {
209+ return false
210+ }
211+ tbl , _ , err := pgs .getTable (sqlCtx , name .SchemaName (), name .TypeName ())
212+ return err == nil && tbl != nil
192213}
193214
194215// resolveName returns the fully resolved name of the given type. Returns an error if the name is ambiguous.
195216func (pgs * TypeCollection ) resolveName (ctx context.Context , schemaName string , typeName string ) (id.Type , error ) {
217+ // TODO: this should probably check table names as well since tables create composite types matching their rows
196218 // First check for an exact match in the built-in types
197219 inputID := id .NewType (schemaName , typeName )
198220 if _ , ok := pgtypes .IDToBuiltInDoltgresType [inputID ]; ok {
@@ -251,6 +273,7 @@ func (pgs *TypeCollection) resolveName(ctx context.Context, schemaName string, t
251273
252274// IterateTypes iterates over all types in the collection.
253275func (pgs * TypeCollection ) IterateTypes (ctx context.Context , f func (typ * pgtypes.DoltgresType ) (stop bool , err error )) error {
276+ // TODO: this should probably iterate tables as well since tables create composite types matching their rows
254277 // We can iterate the built-in types first
255278 for _ , t := range pgtypes .GetAllBuitInTypes () {
256279 stop , err := f (t )
@@ -368,3 +391,50 @@ func (pgs *TypeCollection) writeCache(ctx context.Context) (err error) {
368391 clear (pgs .accessedMap )
369392 return nil
370393}
394+
395+ // getTable returns the SQL table that matches the given schema and table name. Returns a nil table if one is not found.
396+ // This is intended for use with tableToType.
397+ func (* TypeCollection ) getTable (ctx * sql.Context , schema string , tblName string ) (tbl sql.Table , actualSchema string , err error ) {
398+ actualSchema , err = GetSchemaName (ctx , nil , schema )
399+ if err != nil {
400+ return nil , "" , err
401+ }
402+ tbl , err = GetSqlTableFromContext (ctx , "" , doltdb.TableName {
403+ Name : tblName ,
404+ Schema : actualSchema ,
405+ })
406+ if err != nil || tbl == nil {
407+ return nil , "" , err
408+ }
409+ if schTbl , ok := tbl .(sql.DatabaseSchemaTable ); ok {
410+ actualSchema = schTbl .DatabaseSchema ().SchemaName ()
411+ }
412+ return tbl , actualSchema , nil
413+ }
414+
415+ // tableToType handles type creation related to a table's composite row type.
416+ // https://www.postgresql.org/docs/15/sql-createtable.html
417+ func (* TypeCollection ) tableToType (ctx * sql.Context , tbl sql.Table , schema string ) (* pgtypes.DoltgresType , error ) {
418+ tblName := tbl .Name ()
419+ tblSch := tbl .Schema ()
420+ typeID := id .NewType (schema , tblName )
421+ relID := id .NewTable (schema , tblName ).AsId ()
422+ arrayID := id .NewType (schema , "_" + tblName )
423+ attrs := make ([]pgtypes.CompositeAttribute , len (tblSch ))
424+ for i , col := range tblSch {
425+ collation := "" // TODO: what should we use for the collation?
426+ colType , ok := col .Type .(* pgtypes.DoltgresType )
427+ if ! ok {
428+ // TODO: perhaps we should use a better error message stating that it uses a non-Doltgres type?
429+ return nil , pgtypes .ErrTypeDoesNotExist .New (tblName )
430+ }
431+ attrs [i ] = pgtypes .NewCompositeAttribute (ctx , relID , col .Name , colType .ID , int16 (i + 1 ), collation )
432+ }
433+ return pgtypes .NewCompositeType (ctx , relID , arrayID , typeID , attrs ), nil
434+ }
435+
436+ // GetSqlTableFromContext is a forward declaration to get around import cycles
437+ var GetSqlTableFromContext func (ctx * sql.Context , databaseName string , tableName doltdb.TableName ) (sql.Table , error )
438+
439+ // GetSchemaName is a forward declaration to get around import cycles
440+ var GetSchemaName func (ctx * sql.Context , db sql.Database , schemaName string ) (string , error )
0 commit comments