Repository: TobikoData/sqlmesh
Pull request #2791 (2024-06-18 22:29:35)

Feat: increase default type precision in redshift/mssql view workaround
Redshift and MSSQL use the `VarcharSizeWorkaroundMixin` to infer the size of a varchar or other "max" column when it cannot be resolved from model code.

Data may be truncated if it is inferred to the default length (256 in Redshift, 1 in MSSQL) and that was not intended by the user, so this PR increases the length to "max" if a column with the default length is detected. 


Patch:
diff --git a/sqlmesh/core/engine_adapter/mixins.py b/sqlmesh/core/engine_adapter/mixins.py
index acd34481..171d453c 100644
--- a/sqlmesh/core/engine_adapter/mixins.py
+++ b/sqlmesh/core/engine_adapter/mixins.py
@@ -228,6 +228,29 @@ class NonTransactionalTruncateMixin(EngineAdapter):
 
 
 class VarcharSizeWorkaroundMixin(EngineAdapter):
+    def _default_precision_to_max(
+        self, columns_to_types: t.Dict[str, exp.DataType]
+    ) -> t.Dict[str, exp.DataType]:
+        # get default lengths for types that support "max" length
+        types_with_max_default_param = {
+            k: [self.SCHEMA_DIFFER.parameterized_type_defaults[k][0][0]]
+            for k in self.SCHEMA_DIFFER.max_parameter_length
+            if k in self.SCHEMA_DIFFER.parameterized_type_defaults
+        }
+
+        # Redshift and MSSQL have a bug where CTAS statements have non-deterministic types. If a LIMIT
+        # is applied to a CTAS statement, VARCHAR (and possibly other) types sometimes revert to their
+        # default length of 256 (Redshift) or 1 (MSSQL). If we detect that a type has its default length
+        # and supports "max" length, we convert it to "max" length to prevent inadvertent data truncation.
+        for col_name, col_type in columns_to_types.items():
+            if col_type.this in types_with_max_default_param and col_type.expressions:
+                parameter = self.SCHEMA_DIFFER.get_type_parameters(col_type)
+                type_default = types_with_max_default_param[col_type.this]
+                if parameter == type_default:
+                    col_type.set("expressions", [exp.DataTypeParam(this=exp.var("max"))])
+
+        return columns_to_types
+
     def _build_create_table_exp(
         self,
         table_name_or_schema: t.Union[exp.Schema, TableName],
@@ -268,7 +291,9 @@ class VarcharSizeWorkaroundMixin(EngineAdapter):
                 temp_view_name, select_statement, replace=False, no_schema_binding=False
             )
             try:
-                columns_to_types_from_view = self.columns(temp_view_name)
+                columns_to_types_from_view = self._default_precision_to_max(
+                    self.columns(temp_view_name)
+                )
 
                 schema = self._build_schema_exp(
                     exp.to_table(table_name_or_schema),
diff --git a/sqlmesh/core/schema_diff.py b/sqlmesh/core/schema_diff.py
index 56566acb..a049c250 100644
--- a/sqlmesh/core/schema_diff.py
+++ b/sqlmesh/core/schema_diff.py
@@ -372,8 +372,8 @@ class SchemaDiffer(PydanticModel):
         if current_type.this == new_type.this and not current_type.is_type(
             *exp.DataType.NESTED_TYPES
         ):
-            current_params = self._get_type_parameters(current_type)
-            new_params = self._get_type_parameters(new_type)
+            current_params = self.get_type_parameters(current_type)
+            new_params = self.get_type_parameters(new_type)
 
             if len(current_params) != len(new_params):
                 return False
@@ -381,7 +381,7 @@ class SchemaDiffer(PydanticModel):
             return all(new >= current for current, new in zip(current_params, new_params))
         return False
 
-    def _get_type_parameters(self, type: exp.DataType) -> t.List[t.Union[int, float]]:
+    def get_type_parameters(self, type: exp.DataType) -> t.List[t.Union[int, float]]:
         def _str_to_number(string: str, allows_max_param: bool) -> t.Union[int, float]:
             try:
                 return int(string)


Test patch:
diff --git a/tests/core/engine_adapter/test_mssql.py b/tests/core/engine_adapter/test_mssql.py
index b051bc21..47d08475 100644
--- a/tests/core/engine_adapter/test_mssql.py
+++ b/tests/core/engine_adapter/test_mssql.py
@@ -77,6 +77,65 @@ def test_columns(make_mocked_engine_adapter: t.Callable):
     )
 
 
+def test_varchar_workaround_to_max(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
+    adapter = make_mocked_engine_adapter(MSSQLEngineAdapter)
+
+    columns = {
+        "binary1": exp.DataType.build("BINARY(1)", dialect=adapter.dialect),
+        "varbinary": exp.DataType.build("VARBINARY", dialect=adapter.dialect),
+        "varbinary1": exp.DataType.build("VARBINARY(1)", dialect=adapter.dialect),
+        "varbinary2": exp.DataType.build("VARBINARY(2)", dialect=adapter.dialect),
+        "varchar": exp.DataType.build("VARCHAR", dialect=adapter.dialect),
+        "varchar1": exp.DataType.build("VARCHAR(1)", dialect=adapter.dialect),
+        "varchar2": exp.DataType.build("VARCHAR(2)", dialect=adapter.dialect),
+        "nvarchar": exp.DataType.build("NVARCHAR", dialect=adapter.dialect),
+        "nvarchar1": exp.DataType.build("NVARCHAR(1)", dialect=adapter.dialect),
+        "nvarchar2": exp.DataType.build("NVARCHAR(2)", dialect=adapter.dialect),
+    }
+
+    assert adapter._default_precision_to_max(columns) == {
+        "binary1": exp.DataType.build("BINARY(1)", dialect=adapter.dialect),
+        "varbinary": exp.DataType.build("VARBINARY", dialect=adapter.dialect),
+        "varbinary1": exp.DataType.build("VARBINARY(max)", dialect=adapter.dialect),
+        "varbinary2": exp.DataType.build("VARBINARY(2)", dialect=adapter.dialect),
+        "varchar": exp.DataType.build("VARCHAR", dialect=adapter.dialect),
+        "varchar1": exp.DataType.build("VARCHAR(max)", dialect=adapter.dialect),
+        "varchar2": exp.DataType.build("VARCHAR(2)", dialect=adapter.dialect),
+        "nvarchar": exp.DataType.build("NVARCHAR", dialect=adapter.dialect),
+        "nvarchar1": exp.DataType.build("NVARCHAR(max)", dialect=adapter.dialect),
+        "nvarchar2": exp.DataType.build("NVARCHAR(2)", dialect=adapter.dialect),
+    }
+
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.base.random_id",
+        return_value="test_random_id",
+    )
+
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.mssql.MSSQLEngineAdapter.table_exists",
+        return_value=True,
+    )
+
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.mssql.MSSQLEngineAdapter.columns",
+        return_value=columns,
+    )
+
+    adapter.ctas(
+        table_name="test_schema.test_table",
+        query_or_df=parse_one(
+            "SELECT binary, varbinary1 + 1 AS varbinary1, varbinary2 AS varbinary2, varchar, varchar1, varchar2, nvarchar, nvarchar1, nvarchar2 FROM (SELECT * FROM table WHERE FALSE LIMIT 0) WHERE d > 0 AND FALSE LIMIT 0"
+        ),
+        exists=False,
+    )
+
+    assert to_sql_calls(adapter) == [
+        "CREATE VIEW [__temp_ctas_test_random_id] AS SELECT [binary], [varbinary1] + 1 AS [varbinary1], [varbinary2] AS [varbinary2], [varchar], [varchar1], [varchar2], [nvarchar], [nvarchar1], [nvarchar2] FROM (SELECT * FROM [table]);",
+        "DROP VIEW IF EXISTS [__temp_ctas_test_random_id];",
+        "CREATE TABLE [test_schema].[test_table] ([binary1] BINARY(1), [varbinary] VARBINARY, [varbinary1] VARBINARY(max), [varbinary2] VARBINARY(2), [varchar] VARCHAR, [varchar1] VARCHAR(max), [varchar2] VARCHAR(2), [nvarchar] NVARCHAR, [nvarchar1] NVARCHAR(max), [nvarchar2] NVARCHAR(2));",
+    ]
+
+
 def test_table_exists(make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(MSSQLEngineAdapter)
     adapter.cursor.fetchone.return_value = (1,)
diff --git a/tests/core/engine_adapter/test_redshift.py b/tests/core/engine_adapter/test_redshift.py
index 057e9956..b7990be0 100644
--- a/tests/core/engine_adapter/test_redshift.py
+++ b/tests/core/engine_adapter/test_redshift.py
@@ -29,6 +29,57 @@ def test_columns(adapter: t.Callable):
     assert resp == {"col": exp.DataType.build("INT")}
 
 
+def test_varchar_size_workaround(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
+    adapter = make_mocked_engine_adapter(RedshiftEngineAdapter)
+
+    columns = {
+        "char": exp.DataType.build("CHAR", dialect=adapter.dialect),
+        "char1": exp.DataType.build("CHAR(1)", dialect=adapter.dialect),
+        "char2": exp.DataType.build("CHAR(2)", dialect=adapter.dialect),
+        "varchar": exp.DataType.build("VARCHAR", dialect=adapter.dialect),
+        "varchar256": exp.DataType.build("VARCHAR(256)", dialect=adapter.dialect),
+        "varchar2": exp.DataType.build("VARCHAR(2)", dialect=adapter.dialect),
+    }
+
+    assert adapter._default_precision_to_max(columns) == {
+        "char": exp.DataType.build("CHAR", dialect=adapter.dialect),
+        "char1": exp.DataType.build("CHAR(max)", dialect=adapter.dialect),
+        "char2": exp.DataType.build("CHAR(2)", dialect=adapter.dialect),
+        "varchar": exp.DataType.build("VARCHAR", dialect=adapter.dialect),
+        "varchar256": exp.DataType.build("VARCHAR(max)", dialect=adapter.dialect),
+        "varchar2": exp.DataType.build("VARCHAR(2)", dialect=adapter.dialect),
+    }
+
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.base.random_id",
+        return_value="test_random_id",
+    )
+
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.redshift.RedshiftEngineAdapter.table_exists",
+        return_value=True,
+    )
+
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.redshift.RedshiftEngineAdapter.columns",
+        return_value=columns,
+    )
+
+    adapter.ctas(
+        table_name="test_schema.test_table",
+        query_or_df=parse_one(
+            "SELECT char, char1 + 1 AS char1, char2 AS char2, varchar, varchar256, varchar2 FROM (SELECT * FROM table WHERE FALSE LIMIT 0) WHERE d > 0 AND FALSE LIMIT 0"
+        ),
+        exists=False,
+    )
+
+    assert to_sql_calls(adapter) == [
+        'CREATE VIEW "__temp_ctas_test_random_id" AS SELECT "char", "char1" + 1 AS "char1", "char2" AS "char2", "varchar", "varchar256", "varchar2" FROM (SELECT * FROM "table")',
+        'DROP VIEW IF EXISTS "__temp_ctas_test_random_id" CASCADE',
+        'CREATE TABLE "test_schema"."test_table" ("char" CHAR, "char1" CHAR(max), "char2" CHAR(2), "varchar" VARCHAR, "varchar256" VARCHAR(max), "varchar2" VARCHAR(2))',
+    ]
+
+
 def test_create_table_from_query_exists_no_if_not_exists(
     adapter: t.Callable, mocker: MockerFixture
 ):