List enum types when reading schema (#443)

Part of #376
xataio · Oct 31, 2024 · db74779 · db74779
1 parent b02c324
commit db74779
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 15 deletions.
diff --git a/pkg/schema/schema.go b/pkg/schema/schema.go
@@ -73,6 +73,9 @@ type Column struct {
 
 	// Optional comment for the column
 	Comment string `json:"comment"`
+
+	// Will contain possible enum values if the type is an enum
+	EnumValues []string `json:"enumValues"`
 }
 
 // Index represents an index on a table

diff --git a/pkg/state/init.sql b/pkg/state/init.sql
@@ -89,13 +89,13 @@ BEGIN
                     'oid', t.oid,
                     'comment', descr.description,
                     'columns', (SELECT COALESCE(json_object_agg(name, c), '{}'::json)
-                                FROM (SELECT attr.attname                                                                                        AS name,
-                                             pg_get_expr(def.adbin, def.adrelid)                                                                 AS default,
+                                FROM (SELECT attr.attname                                                                                     AS name,
+                                             pg_get_expr(def.adbin, def.adrelid)                                                              AS default,
                                              NOT (
                                                  attr.attnotnull
                                                      OR tp.typtype = 'd'
                                                      AND tp.typnotnull
-                                                 )                                                                                               AS nullable,
+                                                 )                                                                                            AS nullable,
                                              CASE
                                                  WHEN 'character varying' :: regtype = ANY
                                                       (ARRAY [attr.atttypid, tp.typelem]) THEN REPLACE(
@@ -110,18 +110,24 @@ BEGIN
                                                          'timestamptz'
                                                                                                )
                                                  ELSE format_type(attr.atttypid, attr.atttypmod)
-                                                 END                                                                                             AS type,
-                                             descr.description                                                                                   AS comment,
-                                             (EXISTS (SELECT 1
-                                                      FROM pg_constraint
-                                                      WHERE conrelid = attr.attrelid
-                                                        AND ARRAY [attr.attnum::int] @> conkey::int[]
-                                                        AND contype = 'u') OR EXISTS (SELECT 1
-                                                                                      FROM pg_index
-                                                                                               JOIN pg_class ON pg_class.oid = pg_index.indexrelid
-                                                                                      WHERE indrelid = attr.attrelid
-                                                                                        AND indisunique
-                                                                                        AND ARRAY [attr.attnum::int] @> pg_index.indkey::int[])) AS unique
+                                                 END                                                                                          AS type,
+                                             descr.description                                                                                AS comment,
+                                             (EXISTS
+
+
+                                                  (SELECT 1
+                                                   FROM pg_constraint
+                                                   WHERE conrelid = attr.attrelid
+                                                     AND ARRAY [attr.attnum::int] @> conkey::int[]
+                                                     AND contype = 'u') OR EXISTS (SELECT 1
+                                                                                   FROM pg_index
+                                                                                            JOIN pg_class ON pg_class.oid = pg_index.indexrelid
+                                                                                   WHERE indrelid = attr.attrelid
+                                                                                     AND indisunique
+                                                                                     AND ARRAY [attr.attnum::int] @> pg_index.indkey::int[])) AS unique,
+                                             (SELECT array_agg(e.enumlabel ORDER BY e.enumsortorder)
+                                              from pg_enum as e
+                                              WHERE e.enumtypid = tp.oid)                                                                     AS enumValues
                                       FROM pg_attribute AS attr
                                                INNER JOIN pg_type AS tp ON attr.atttypid = tp.oid
                                                LEFT JOIN pg_attrdef AS def ON attr.attrelid = def.adrelid

diff --git a/pkg/state/state_test.go b/pkg/state/state_test.go
@@ -816,6 +816,36 @@ func TestReadSchema(t *testing.T) {
 					},
 				},
 			},
+			{
+				name:       "custom enum types",
+				createStmt: "CREATE TYPE review AS ENUM ('good', 'bad', 'ugly'); CREATE TABLE public.table1 (name text, review review);",
+				wantSchema: &schema.Schema{
+					Name: "public",
+					Tables: map[string]schema.Table{
+						"table1": {
+							Name: "table1",
+							Columns: map[string]schema.Column{
+								"name": {
+									Name:     "name",
+									Type:     "text",
+									Nullable: true,
+								},
+								"review": {
+									Name:       "review",
+									Type:       "public.review",
+									Nullable:   true,
+									EnumValues: []string{"good", "bad", "ugly"},
+								},
+							},
+							PrimaryKey:        []string{},
+							Indexes:           map[string]schema.Index{},
+							ForeignKeys:       map[string]schema.ForeignKey{},
+							CheckConstraints:  map[string]schema.CheckConstraint{},
+							UniqueConstraints: map[string]schema.UniqueConstraint{},
+						},
+					},
+				},
+			},
 		}
 
 		for _, tt := range tests {