support explicit cast to domain type

core/typecollection/serialization.go

@@ -32,6 +32,9 @@ func (pgs *TypeCollection) Serialize(ctx context.Context) ([]byte, error) {
 	pgs.mutex.Lock()
 	defer pgs.mutex.Unlock()
 
+	// TODO: technically, can create type in pg_catalog schema
+	delete(pgs.schemaMap, "pg_catalog")
+
 	// Write all the types to the writer
 	writer := utils.NewWriter(256)
 	writer.VariableUint(0) // Version

core/typecollection/typecollection.go

@@ -27,40 +27,74 @@ type TypeCollection struct {
 	mutex     *sync.RWMutex
 }
 
-// GetType returns the type with the given schema and name.
-// Returns nil if the type cannot be found.
-func (pgs *TypeCollection) GetType(schName, typName string) (*types.DoltgresType, bool) {
-	pgs.mutex.RLock()
-	defer pgs.mutex.RUnlock()
+// Clone returns a new *TypeCollection with the same contents as the original.
+func (pgs *TypeCollection) Clone() *TypeCollection {
+	pgs.mutex.Lock()
+	defer pgs.mutex.Unlock()
 
-	if nameMap, ok := pgs.schemaMap[schName]; ok {
-		if typ, ok := nameMap[typName]; ok {
-			return typ, true
+	newCollection := &TypeCollection{
+		schemaMap: make(map[string]map[string]*types.DoltgresType),
+		mutex:     &sync.RWMutex{},
+	}
+	for schema, nameMap := range pgs.schemaMap {
+		if len(nameMap) == 0 {
+			continue
+		} else if schema == "pg_catalog" {
+			// TODO: technically, can create type in pg_catalog schema
+			continue
+		}
+		clonedNameMap := make(map[string]*types.DoltgresType)
+		for key, typ := range nameMap {
+			clonedNameMap[key] = typ
 		}
+		newCollection.schemaMap[schema] = clonedNameMap
 	}
-	return nil, false
+	return newCollection
 }
 
-// GetDomainType returns a domain type with the given schema and name.
-// Returns nil if the type cannot be found. It checks for domain type.
-func (pgs *TypeCollection) GetDomainType(schName, typName string) (*types.DoltgresType, bool) {
-	pgs.mutex.RLock()
-	defer pgs.mutex.RUnlock()
+// CreateType creates a new type.
+func (pgs *TypeCollection) CreateType(schName string, typ *types.DoltgresType) error {
+	pgs.mutex.Lock()
+	defer pgs.mutex.Unlock()
+
+	nameMap, ok := pgs.schemaMap[schName]
+	if !ok {
+		nameMap = make(map[string]*types.DoltgresType)
+		pgs.schemaMap[schName] = nameMap
+	}
+	if _, ok = nameMap[typ.Name]; ok {
+		return types.ErrTypeAlreadyExists.New(typ.Name)
+	}
+	nameMap[typ.Name] = typ
+	return nil
+}
+
+// DropType drops an existing type.
+func (pgs *TypeCollection) DropType(schName, typName string) error {
+	pgs.mutex.Lock()
+	defer pgs.mutex.Unlock()
+
+	if schName == "pg_catalog" {
+		// TODO: check if it's built-in type (technically, can create type in pg_catalog schema)
+		return types.ErrCannotDropType.New(typName)
+	}
 
 	if nameMap, ok := pgs.schemaMap[schName]; ok {
-		if typ, ok := nameMap[typName]; ok && typ.TypType == types.TypeType_Domain {
-			return typ, true
+		if _, ok = nameMap[typName]; ok {
+			delete(nameMap, typName)
+			return nil
 		}
 	}
-	return nil, false
+	return types.ErrTypeDoesNotExist.New(typName)
 }
 
 // GetAllTypes returns a map containing all types in the collection, grouped by the schema they're contained in.
-// Each type array is also sorted by the type name.
+// Each type array is also sorted by the type name. It includes built-in types.
 func (pgs *TypeCollection) GetAllTypes() (typesMap map[string][]*types.DoltgresType, schemaNames []string, totalCount int) {
 	pgs.mutex.RLock()
 	defer pgs.mutex.RUnlock()
 
+	pgs.addSupportedBuiltInTypes()
 	typesMap = make(map[string][]*types.DoltgresType)
 	for schemaName, nameMap := range pgs.schemaMap {
 		schemaNames = append(schemaNames, schemaName)
@@ -75,49 +109,78 @@ func (pgs *TypeCollection) GetAllTypes() (typesMap map[string][]*types.DoltgresT
 		typesMap[schemaName] = typs
 	}
 
-	// TODO: add built-in types
 	sort.Slice(schemaNames, func(i, j int) bool {
 		return schemaNames[i] < schemaNames[j]
 	})
 	return
 }
 
-// CreateType creates a new type.
-func (pgs *TypeCollection) CreateType(schema string, typ *types.DoltgresType) error {
-	pgs.mutex.Lock()
-	defer pgs.mutex.Unlock()
+// GetDomainType returns a domain type with the given schema and name.
+// Returns nil if the type cannot be found. It checks for domain type.
+func (pgs *TypeCollection) GetDomainType(schName, typName string) (*types.DoltgresType, bool) {
+	t, exists := pgs.GetType(schName, typName)
+	if !exists {
+		return nil, exists
+	}
+	if t.TypType == types.TypeType_Domain {
+		return t, exists
+	}
+	return nil, false
+}
 
-	nameMap, ok := pgs.schemaMap[schema]
-	if !ok {
-		nameMap = make(map[string]*types.DoltgresType)
-		pgs.schemaMap[schema] = nameMap
+// GetType returns the type with the given schema and name.
+// Returns nil if the type cannot be found.
+func (pgs *TypeCollection) GetType(schName, typName string) (*types.DoltgresType, bool) {
+	pgs.mutex.RLock()
+	defer pgs.mutex.RUnlock()
+
+	pgs.addSupportedBuiltInTypes()
+	if nameMap, ok := pgs.schemaMap[schName]; ok {
+		if typ, ok := nameMap[typName]; ok {
+			return typ, true
+		}
 	}
-	if _, ok = nameMap[typ.Name]; ok {
-		return types.ErrTypeAlreadyExists.New(typ.Name)
+
+	return nil, false
+}
+
+// GetTypeByOID returns the type matching given OID.
+func (pgs *TypeCollection) GetTypeByOID(oid uint32) (*types.DoltgresType, bool) {
+	pgs.mutex.RLock()
+	defer pgs.mutex.RUnlock()
+
+	pgs.addSupportedBuiltInTypes()
+	for _, nameMap := range pgs.schemaMap {
+		for _, typ := range nameMap {
+			if typ.OID == oid {
+				return typ, true
+			}
+		}
 	}
-	nameMap[typ.Name] = typ
-	return nil
+	return nil, false
 }
 
-// DropType drops an existing type.
-func (pgs *TypeCollection) DropType(schName, typName string) error {
+// HasType checks if a type exists with given schema and type name.
+func (pgs *TypeCollection) HasType(schema string, typName string) bool {
 	pgs.mutex.Lock()
 	defer pgs.mutex.Unlock()
 
-	if nameMap, ok := pgs.schemaMap[schName]; ok {
-		if _, ok = nameMap[typName]; ok {
-			delete(nameMap, typName)
-			return nil
-		}
+	pgs.addSupportedBuiltInTypes()
+	nameMap, ok := pgs.schemaMap[schema]
+	if !ok {
+		nameMap = make(map[string]*types.DoltgresType)
+		pgs.schemaMap[schema] = nameMap
 	}
-	return types.ErrTypeDoesNotExist.New(typName)
+	_, ok = nameMap[typName]
+	return ok
 }
 
 // IterateTypes iterates over all types in the collection.
 func (pgs *TypeCollection) IterateTypes(f func(schema string, typ *types.DoltgresType) error) error {
 	pgs.mutex.Lock()
 	defer pgs.mutex.Unlock()
 
+	pgs.addSupportedBuiltInTypes()
 	for schema, nameMap := range pgs.schemaMap {
 		for _, t := range nameMap {
 			if err := f(schema, t); err != nil {
@@ -128,24 +191,15 @@ func (pgs *TypeCollection) IterateTypes(f func(schema string, typ *types.Doltgre
 	return nil
 }
 
-// Clone returns a new *TypeCollection with the same contents as the original.
-func (pgs *TypeCollection) Clone() *TypeCollection {
-	pgs.mutex.Lock()
-	defer pgs.mutex.Unlock()
-
-	newCollection := &TypeCollection{
-		schemaMap: make(map[string]map[string]*types.DoltgresType),
-		mutex:     &sync.RWMutex{},
-	}
-	for schema, nameMap := range pgs.schemaMap {
-		if len(nameMap) == 0 {
-			continue
+// addSupportedBuiltInTypes adds supported built-in types in the type collection map
+// with 'pg_catalog' schema as key. It doesn't add if 'pg_catalog' entry exists in the map.
+func (pgs *TypeCollection) addSupportedBuiltInTypes() {
+	if _, ok := pgs.schemaMap["pg_catalog"]; !ok {
+		// add built-in types
+		pgCatTypeMap := make(map[string]*types.DoltgresType)
+		for _, t := range types.GetAllTypes() {
+			pgCatTypeMap[t.Name] = t
 		}
-		clonedNameMap := make(map[string]*types.DoltgresType)
-		for key, typ := range nameMap {
-			clonedNameMap[key] = typ
-		}
-		newCollection.schemaMap[schema] = clonedNameMap
+		pgs.schemaMap["pg_catalog"] = pgCatTypeMap
 	}
-	return newCollection
 }