Repository: TobikoData/sqlmesh
Pull request #3225 (2024-10-04 15:51:53)

Fix(redshift): regression in varchar length workaround
***NOTE: redshift integration tests won't pass until sqlglot is bumped > 25.24.4***

#2791 introduced a regression in the varchar length workaround because it assumed the `columns()` method returned data types including their length/precision. That was true for MSSQL, but wasn't for Redshift.

This PR updates Redshift's `columns()` method to attach length/precision to data types.

Note: the `_columns_query()` method allowed postgres and redshift to use different columns queries inside postgres' `columns()` method. Redshift now needs to override the entire `columns()` method, so `_columns_query()` has been removed from postgres.

treysp: Feedback addressed, also updated the MSSQL implementation to align
georgesittas: > NOTE: redshift integration tests won't pass until sqlglot is bumped > 25.24.4

Released `v25.24.5`.

Patch:
diff --git a/sqlmesh/core/engine_adapter/base_postgres.py b/sqlmesh/core/engine_adapter/base_postgres.py
index 97b2523c..51840d75 100644
--- a/sqlmesh/core/engine_adapter/base_postgres.py
+++ b/sqlmesh/core/engine_adapter/base_postgres.py
@@ -26,7 +26,12 @@ class BasePostgresEngineAdapter(EngineAdapter):
     COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY
     COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY
 
-    def _columns_query(self, table: exp.Table) -> exp.Select:
+    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(
                 "attname AS column_name",
@@ -45,14 +50,8 @@ class BasePostgresEngineAdapter(EngineAdapter):
         )
         if table.args.get("db"):
             sql = sql.where(exp.column("nspname").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)
-        self.execute(self._columns_query(table))
+        self.execute(sql)
         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/mssql.py b/sqlmesh/core/engine_adapter/mssql.py
index 3044c001..c7fb39d9 100644
--- a/sqlmesh/core/engine_adapter/mssql.py
+++ b/sqlmesh/core/engine_adapter/mssql.py
@@ -68,6 +68,7 @@ class MSSQLEngineAdapter(
             exp.DataType.build("NVARCHAR", dialect=DIALECT).this: 2147483647,
         },
     )
+    VARIABLE_LENGTH_DATA_TYPES = {"binary", "varbinary", "char", "varchar", "nchar", "nvarchar"}
 
     def columns(
         self,
@@ -95,23 +96,34 @@ class MSSQLEngineAdapter(
 
         columns_raw = self.fetchall(sql, quote_identifiers=True)
 
-        def build_var_length_col(row: tuple) -> tuple:
-            var_len_chars = ("binary", "varbinary", "char", "varchar", "nchar", "nvarchar")
-            if row[1] in var_len_chars and row[2] > 0:
-                return (row[0], f"{row[1]}({row[2]})")
-            if row[1] in ("varbinary", "varchar", "nvarchar") and row[2] == -1:
-                return (row[0], f"{row[1]}(max)")
-            if row[1] in (
-                "decimal",
-                "numeric",
+        def build_var_length_col(
+            column_name: str,
+            data_type: str,
+            character_maximum_length: t.Optional[int] = None,
+            numeric_precision: t.Optional[int] = None,
+            numeric_scale: t.Optional[int] = None,
+        ) -> tuple:
+            data_type = data_type.lower()
+            if (
+                data_type in self.VARIABLE_LENGTH_DATA_TYPES
+                and character_maximum_length is not None
+                and character_maximum_length > 0
             ):
-                return (row[0], f"{row[1]}({row[3]}, {row[4]})")
-            if row[1] == "float":
-                return (row[0], f"{row[1]}({row[3]})")
+                return (column_name, f"{data_type}({character_maximum_length})")
+            if (
+                data_type in ("varbinary", "varchar", "nvarchar")
+                and character_maximum_length is not None
+                and character_maximum_length == -1
+            ):
+                return (column_name, f"{data_type}(max)")
+            if data_type in ("decimal", "numeric"):
+                return (column_name, f"{data_type}({numeric_precision}, {numeric_scale})")
+            if data_type == "float":
+                return (column_name, f"{data_type}({numeric_precision})")
 
-            return (row[0], row[1])
+            return (column_name, data_type)
 
-        columns = [build_var_length_col(col) for col in columns_raw]
+        columns = [build_var_length_col(*row) for row in columns_raw]
 
         return {
             column_name: exp.DataType.build(data_type, dialect=self.dialect)
diff --git a/sqlmesh/core/engine_adapter/redshift.py b/sqlmesh/core/engine_adapter/redshift.py
index 4f9a8009..86438f3b 100644
--- a/sqlmesh/core/engine_adapter/redshift.py
+++ b/sqlmesh/core/engine_adapter/redshift.py
@@ -54,16 +54,65 @@ class RedshiftEngineAdapter(
             exp.DataType.build("VARCHAR", dialect=DIALECT).this: 65535,
         },
     )
+    VARIABLE_LENGTH_DATA_TYPES = {
+        "char",
+        "character",
+        "nchar",
+        "varchar",
+        "character varying",
+        "nvarchar",
+        "varbyte",
+        "varbinary",
+        "binary varying",
+    }
+
+    def columns(
+        self,
+        table_name: TableName,
+        include_pseudo_columns: bool = True,
+    ) -> t.Dict[str, exp.DataType]:
+        table = exp.to_table(table_name)
 
-    def _columns_query(self, table: exp.Table) -> exp.Select:
         sql = (
-            exp.select("column_name", "data_type")
+            exp.select(
+                "column_name",
+                "data_type",
+                "character_maximum_length",
+                "numeric_precision",
+                "numeric_scale",
+            )
             .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
+
+        columns_raw = self.fetchall(sql, quote_identifiers=True)
+
+        def build_var_length_col(
+            column_name: str,
+            data_type: str,
+            character_maximum_length: t.Optional[int] = None,
+            numeric_precision: t.Optional[int] = None,
+            numeric_scale: t.Optional[int] = None,
+        ) -> tuple:
+            data_type = data_type.lower()
+            if (
+                data_type in self.VARIABLE_LENGTH_DATA_TYPES
+                and character_maximum_length is not None
+            ):
+                return (column_name, f"{data_type}({character_maximum_length})")
+            if data_type in ("decimal", "numeric"):
+                return (column_name, f"{data_type}({numeric_precision}, {numeric_scale})")
+
+            return (column_name, data_type)
+
+        columns = [build_var_length_col(*row) for row in columns_raw]
+
+        return {
+            column_name: exp.DataType.build(data_type, dialect=self.dialect)
+            for column_name, data_type in columns
+        }
 
     @property
     def cursor(self) -> t.Any:


Test patch:
diff --git a/tests/core/engine_adapter/integration/test_integration_redshift.py b/tests/core/engine_adapter/integration/test_integration_redshift.py
new file mode 100644
index 00000000..c3c438e6
--- /dev/null
+++ b/tests/core/engine_adapter/integration/test_integration_redshift.py
@@ -0,0 +1,87 @@
+import typing as t
+import pytest
+from tests.core.engine_adapter.integration import TestContext
+from sqlglot import exp
+
+pytestmark = [pytest.mark.remote, pytest.mark.engine, pytest.mark.redshift]
+
+
+@pytest.fixture
+def mark_gateway() -> t.Tuple[str, str]:
+    return "redshift", "inttest_redshift"
+
+
+@pytest.fixture
+def test_type() -> str:
+    return "query"
+
+
+def test_columns(ctx: TestContext):
+    ctx.init()
+
+    table = ctx.table("column_types")
+    col_strings = {
+        "char": ["char", "character", "nchar"],
+        "varchar": ["varchar", "character varying", "nvarchar"],
+        "varbinary": ["varbyte", "varbinary", "binary varying"],
+        "decimal": ["decimal", "numeric"],
+    }
+
+    # raw ddl
+    sql = f"CREATE TABLE {table} ("
+    sql += (
+        ", ".join(
+            f"{col.replace(' ', '_')}10 {col}(10)"
+            for col in [*col_strings["char"], *col_strings["varchar"], *col_strings["varbinary"]]
+        )
+        + ", "
+    )
+    # bare types that should have their default lengths of 1 added by columns()
+    sql += ", ".join(f"{col.replace(' ', '_')}1 {col}" for col in col_strings["char"]) + ", "
+    # bare types that should have their default lengths of 256 added by columns()
+    sql += ", ".join(f"{col.replace(' ', '_')}256 {col}" for col in col_strings["varchar"]) + ", "
+    sql += (
+        ", ".join(f"{col.replace(' ', '_')}172 {col}(17, 2)" for col in col_strings["decimal"])
+        + ")"
+    )
+
+    ctx.engine_adapter.cursor.execute(sql)
+    columns = ctx.engine_adapter.columns(table)
+
+    # columns to types
+    cols_to_types = {
+        f"{col.replace(' ', '_')}10": exp.DataType.build(f"{col}(10)", dialect=ctx.dialect)
+        for col in [*col_strings["char"], *col_strings["varchar"], *col_strings["varbinary"]]
+    }
+    cols_to_types.update(
+        {
+            f"{col.replace(' ', '_')}1": exp.DataType.build(f"{col}(1)", dialect=ctx.dialect)
+            for col in col_strings["char"]
+        }
+    )
+    cols_to_types.update(
+        {
+            f"{col.replace(' ', '_')}256": exp.DataType.build(f"{col}(256)", dialect=ctx.dialect)
+            for col in col_strings["varchar"]
+        }
+    )
+    cols_to_types.update(
+        {
+            f"{col.replace(' ', '_')}172": exp.DataType.build(f"{col}(17, 2)", dialect=ctx.dialect)
+            for col in col_strings["decimal"]
+        }
+    )
+
+    # did we convert the types from redshift correctly?
+    assert [col.sql(ctx.dialect) for col in columns.values()] == [
+        col.sql(ctx.dialect) for col in cols_to_types.values()
+    ]
+
+    # did we replace default char/varchar lengths with MAX correctly?
+    max_cols = [col for col in columns if col.endswith("1") or col.endswith("256")]
+    assert [
+        col.sql(ctx.dialect)
+        for col in ctx.engine_adapter._default_precision_to_max(  # type: ignore
+            {k: columns[k] for k in max_cols}
+        ).values()
+    ] == ["CHAR(max)", "CHAR(max)", "CHAR(max)", "VARCHAR(max)", "VARCHAR(max)", "VARCHAR(max)"]
diff --git a/tests/core/engine_adapter/test_redshift.py b/tests/core/engine_adapter/test_redshift.py
index b7990be0..f5102339 100644
--- a/tests/core/engine_adapter/test_redshift.py
+++ b/tests/core/engine_adapter/test_redshift.py
@@ -24,7 +24,7 @@ def test_columns(adapter: t.Callable):
     adapter.cursor.fetchall.return_value = [("col", "INT")]
     resp = adapter.columns("db.table")
     adapter.cursor.execute.assert_called_once_with(
-        """SELECT "column_name", "data_type" FROM "svv_columns" WHERE "table_name" = 'table' AND "table_schema" = 'db'"""
+        """SELECT "column_name", "data_type", "character_maximum_length", "numeric_precision", "numeric_scale" FROM "svv_columns" WHERE "table_name" = 'table' AND "table_schema" = 'db'"""
     )
     assert resp == {"col": exp.DataType.build("INT")}