Repository: TobikoData/sqlmesh
Pull request #2491 (2024-04-22 23:01:13)

fix: use pg_attribute for postgres to handle nested data types closes…
… #2449

the old method for postgres didn't return the types for nested columns, text[] was returned as array.

this new method returns text[] properly. i couldn't do this for redshift because late binding views are not accounted for in this table. it shouldn't be a problem though because redshift doesn't have nested types like postgres


Patch:
diff --git a/sqlmesh/core/engine_adapter/base_postgres.py b/sqlmesh/core/engine_adapter/base_postgres.py
index 4ca5627f..9ad49c4f 100644
--- a/sqlmesh/core/engine_adapter/base_postgres.py
+++ b/sqlmesh/core/engine_adapter/base_postgres.py
@@ -22,25 +22,37 @@ if t.TYPE_CHECKING:
 
 class BasePostgresEngineAdapter(EngineAdapter):
     DEFAULT_BATCH_SIZE = 400
-    COLUMNS_TABLE = "information_schema.columns"
     CATALOG_SUPPORT = CatalogSupport.SINGLE_CATALOG_ONLY
     COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY
     COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY
 
+    def _columns_query(self, table: exp.Table) -> exp.Select:
+        sql = (
+            exp.select(
+                "attname AS column_name",
+                "pg_catalog.format_type(atttypid, atttypmod) AS data_type",
+            )
+            .from_("pg_catalog.pg_attribute")
+            .join("pg_catalog.pg_class", on="oid = attrelid")
+            .join("pg_catalog.pg_namespace", on="oid = relnamespace")
+            .where(
+                exp.and_(
+                    "attnum > 0",
+                    "NOT attisdropped",
+                    exp.column("nspname").eq(table.alias_or_name),
+                )
+            )
+        )
+        if table.args.get("db"):
+            sql = sql.where(exp.column("relname").eq(table.args["db"].name))
+        return sql
+
     def columns(
         self, table_name: TableName, include_pseudo_columns: bool = False
     ) -> t.Dict[str, exp.DataType]:
         """Fetches column names and types for the target table."""
         table = exp.to_table(table_name)
-        sql = (
-            exp.select("column_name", "data_type")
-            .from_(self.COLUMNS_TABLE)
-            .where(f"table_name = '{table.alias_or_name}'")
-        )
-        if table.args.get("db"):
-            sql = sql.where(f"table_schema = '{table.args['db'].name}'")
-
-        self.execute(sql)
+        self.execute(self._columns_query(table))
         resp = self.cursor.fetchall()
         if not resp:
             raise SQLMeshError("Could not get columns for table '%s'. Table not found.", table_name)
diff --git a/sqlmesh/core/engine_adapter/redshift.py b/sqlmesh/core/engine_adapter/redshift.py
index 981f9c78..d8bb59c5 100644
--- a/sqlmesh/core/engine_adapter/redshift.py
+++ b/sqlmesh/core/engine_adapter/redshift.py
@@ -34,12 +34,21 @@ class RedshiftEngineAdapter(
     NonTransactionalTruncateMixin,
 ):
     DIALECT = "redshift"
-    COLUMNS_TABLE = "svv_columns"  # Includes late-binding views
     CURRENT_CATALOG_EXPRESSION = exp.func("current_database")
     # Redshift doesn't support comments for VIEWs WITH NO SCHEMA BINDING (which we always use)
     COMMENT_CREATION_VIEW = CommentCreationView.UNSUPPORTED
     SUPPORTS_REPLACE_TABLE = False
 
+    def _columns_query(self, table: exp.Table) -> exp.Select:
+        sql = (
+            exp.select("column_name", "data_type")
+            .from_("svv_columns")  # Includes late-binding views
+            .where(exp.column("table_name").eq(table.alias_or_name))
+        )
+        if table.args.get("db"):
+            sql = sql.where(exp.column("table_schema").eq(table.args["db"].name))
+        return sql
+
     @property
     def cursor(self) -> t.Any:
         # Redshift by default uses a `format` paramstyle that has issues when we try to write our snapshot


Test patch:
diff --git a/tests/core/engine_adapter/test_base_postgres.py b/tests/core/engine_adapter/test_base_postgres.py
index a7a9369f..2e8809b1 100644
--- a/tests/core/engine_adapter/test_base_postgres.py
+++ b/tests/core/engine_adapter/test_base_postgres.py
@@ -16,7 +16,12 @@ def test_columns(make_mocked_engine_adapter: t.Callable):
 
     resp = adapter.columns("db.table")
     adapter.cursor.execute.assert_called_once_with(
-        """SELECT "column_name", "data_type" FROM "information_schema"."columns" WHERE "table_name" = 'table' AND "table_schema" = 'db'"""
+        'SELECT "attname" AS "column_name", '
+        '"pg_catalog".FORMAT_TYPE("atttypid", "atttypmod") AS "data_type" '
+        'FROM "pg_catalog"."pg_attribute" '
+        'JOIN "pg_catalog"."pg_class" ON "oid" = "attrelid" '
+        'JOIN "pg_catalog"."pg_namespace" ON "oid" = "relnamespace" '
+        """WHERE ("attnum" > 0 AND NOT "attisdropped" AND "nspname" = 'table') AND "relname" = 'db'"""
     )
     assert resp == {"col": exp.DataType.build("INT")}