From aa1e16d60e757101f3072a434cb5fdffe62c1fcf Mon Sep 17 00:00:00 2001
From: Alisha <alisha@nucleuscloud.com>
Date: Fri, 17 Jan 2025 16:50:18 -0800
Subject: [PATCH] add all required schemas for datatypes

---
 .../go/db/dbschemas/postgresql/system.sql.go  |  2 +-
 .../sql/postgresql/queries/system.sql         |  2 +-
 .../sqlmanager/postgres/postgres-manager.go   | 37 +++++++++++++++++++
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/backend/gen/go/db/dbschemas/postgresql/system.sql.go b/backend/gen/go/db/dbschemas/postgresql/system.sql.go
index a092b4c2bc..8944430143 100644
--- a/backend/gen/go/db/dbschemas/postgresql/system.sql.go
+++ b/backend/gen/go/db/dbschemas/postgresql/system.sql.go
@@ -275,7 +275,7 @@ WITH custom_types AS (
     JOIN
         pg_catalog.pg_namespace n ON n.oid = t.typnamespace
     WHERE
-        n.nspname = $1
+        n.nspname NOT IN('pg_catalog', 'pg_toast', 'information_schema')
         AND t.typtype IN ('d', 'e', 'c')
 ),
 table_columns AS (
diff --git a/backend/pkg/dbschemas/sql/postgresql/queries/system.sql b/backend/pkg/dbschemas/sql/postgresql/queries/system.sql
index 9c9279e38e..bf1b0ea0b5 100644
--- a/backend/pkg/dbschemas/sql/postgresql/queries/system.sql
+++ b/backend/pkg/dbschemas/sql/postgresql/queries/system.sql
@@ -436,7 +436,7 @@ WITH custom_types AS (
     JOIN
         pg_catalog.pg_namespace n ON n.oid = t.typnamespace
     WHERE
-        n.nspname = sqlc.arg('schema')
+        n.nspname NOT IN('pg_catalog', 'pg_toast', 'information_schema')
         AND t.typtype IN ('d', 'e', 'c')
 ),
 table_columns AS (
diff --git a/backend/pkg/sqlmanager/postgres/postgres-manager.go b/backend/pkg/sqlmanager/postgres/postgres-manager.go
index 28a064ec99..1daa77b3ee 100644
--- a/backend/pkg/sqlmanager/postgres/postgres-manager.go
+++ b/backend/pkg/sqlmanager/postgres/postgres-manager.go
@@ -529,6 +529,7 @@ func (p *PostgresManager) GetSchemaInitStatements(
 		return nil
 	})
 
+	datatypes := &sqlmanager_shared.SchemaTableDataTypeResponse{}
 	dataTypeStmts := []string{}
 	errgrp.Go(func() error {
 		datatypeCfg, err := p.GetSchemaTableDataTypes(errctx, tables)
@@ -536,6 +537,7 @@ func (p *PostgresManager) GetSchemaInitStatements(
 			return fmt.Errorf("unable to retrieve postgres schema table data types: %w", err)
 		}
 		dataTypeStmts = datatypeCfg.GetStatements()
+		datatypes = datatypeCfg
 		return nil
 	})
 
@@ -578,6 +580,9 @@ func (p *PostgresManager) GetSchemaInitStatements(
 		return nil, err
 	}
 
+	additionalSchemaStmts := getSchemaCreationStatementsFromDataTypes(tables, datatypes)
+	schemaStmts = append(schemaStmts, additionalSchemaStmts...)
+
 	return []*sqlmanager_shared.InitSchemaStatements{
 		{Label: SchemasLabel, Statements: schemaStmts},
 		{Label: "data types", Statements: dataTypeStmts},
@@ -589,6 +594,38 @@ func (p *PostgresManager) GetSchemaInitStatements(
 	}, nil
 }
 
+// Finds any schemas referenced in datatypes that don't exist in tables and returns the statements to create them
+func getSchemaCreationStatementsFromDataTypes(tables []*sqlmanager_shared.SchemaTable, datatypes *sqlmanager_shared.SchemaTableDataTypeResponse) []string {
+	schemaStmts := []string{}
+	schemaSet := map[string]struct{}{}
+	for _, table := range tables {
+		schemaSet[table.Schema] = struct{}{}
+	}
+
+	// Check each datatype schema against the table schemas
+	for _, composite := range datatypes.Composites {
+		if _, exists := schemaSet[composite.Schema]; !exists {
+			schemaStmts = append(schemaStmts, fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %q;", composite.Schema))
+			schemaSet[composite.Schema] = struct{}{}
+		}
+	}
+
+	for _, enum := range datatypes.Enums {
+		if _, exists := schemaSet[enum.Schema]; !exists {
+			schemaStmts = append(schemaStmts, fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %q;", enum.Schema))
+			schemaSet[enum.Schema] = struct{}{}
+		}
+	}
+
+	for _, domain := range datatypes.Domains {
+		if _, exists := schemaSet[domain.Schema]; !exists {
+			schemaStmts = append(schemaStmts, fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %q;", domain.Schema))
+			schemaSet[domain.Schema] = struct{}{}
+		}
+	}
+	return schemaStmts
+}
+
 func wrapPgIdempotentIndex(
 	schema,
 	constraintname,