diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py index 67918b6d14..63386d176c 100644 --- a/sqlmesh/core/dialect.py +++ b/sqlmesh/core/dialect.py @@ -551,6 +551,36 @@ def _parse_table_parts( return table +# Only needed for T-SQL: it spells a column's nullability right after its type, e.g. +# ALTER TABLE t ALTER COLUMN c INT NOT NULL. Without this the trailing clause is left +# over, so the whole statement falls back to a Command and any macros it contains (such as +# @this_model) are no longer resolved, which means they reach the engine verbatim. +# +# See: https://learn.microsoft.com/en-us/sql/t-sql/statements/alter-table-transact-sql +def _parse_alter_table_alter(self: Parser) -> t.Optional[exp.Expr]: + alter_column = self.__parse_alter_table_alter() # type: ignore + + if isinstance(alter_column, exp.AlterColumn) and alter_column.args.get("dtype"): + if self._match_pair(TokenType.NOT, TokenType.NULL): + alter_column.set("allow_null", False) + elif self._match(TokenType.NULL): + alter_column.set("allow_null", True) + + return alter_column + + +def altercolumn_sql(self: Generator, expression: exp.AlterColumn) -> str: + sql = self._altercolumn_sql(expression) # type: ignore + + # sqlglot's generator returns as soon as it renders the type, so the nullability parsed + # above has to be appended here + allow_null = expression.args.get("allow_null") + if expression.args.get("dtype") and allow_null is not None: + sql = f"{sql} NULL" if allow_null else f"{sql} NOT NULL" + + return sql + + def _parse_if(self: Parser) -> t.Optional[exp.Expr]: # If we fail to parse an IF function with expressions as arguments, we then try # to parse a statement / command to support the macro @IF(condition, statement) @@ -780,7 +810,7 @@ def _parse_interval_span(self: Parser, this: exp.Expr) -> exp.Interval: return interval -def _override(klass: t.Type[Tokenizer | Parser], func: t.Callable) -> None: +def _override(klass: t.Type[Tokenizer | Parser | Generator], func: t.Callable) -> None: name = func.__name__ setattr(klass, f"_{name}", getattr(klass, name)) setattr(klass, name, func) @@ -1194,6 +1224,8 @@ def extend_sqlglot() -> None: _override(Parser, _parse_interval_span) _override(Parser, _warn_unsupported) _override(Snowflake.Parser, _parse_table_parts) + _override(TSQL.Parser, _parse_alter_table_alter) + _override(TSQL.Generator, altercolumn_sql) # DuckDB's prefix absolute power operator `@` clashes with the macro syntax DuckDB.Parser.NO_PAREN_FUNCTION_PARSERS.pop("@", None) diff --git a/tests/core/test_dialect.py b/tests/core/test_dialect.py index e2f1daba3d..fb967a7f5d 100644 --- a/tests/core/test_dialect.py +++ b/tests/core/test_dialect.py @@ -995,6 +995,38 @@ def test_conditional_statement(): assert q.sql(dialect="tsql") == "@IF(@runtime_stage = 'evaluating', SELECT 1)" +def test_tsql_alter_column_nullability(): + # Issue #5932: T-SQL spells nullability right after the type in ALTER COLUMN. Without support + # for it the statement falls back to a Command, so any macros it contains go unresolved. + for sql, expected in [ + ( + "ALTER TABLE x ALTER COLUMN y INT NOT NULL", + "ALTER TABLE x ALTER COLUMN y INTEGER NOT NULL", + ), + ("ALTER TABLE x ALTER COLUMN y INT NULL", "ALTER TABLE x ALTER COLUMN y INTEGER NULL"), + ("ALTER TABLE x ALTER COLUMN y INT", "ALTER TABLE x ALTER COLUMN y INTEGER"), + ]: + e = parse_one(sql, read="tsql") + assert isinstance(e, exp.Alter) + assert e.sql(dialect="tsql") == expected + + # The macro must survive parsing so that it can be resolved later + e = parse_one( + "@IF(@runtime_stage = 'creating', ALTER TABLE @SQL('@this_model') ALTER COLUMN id INT NOT NULL);", + read="tsql", + ) + assert ( + e.sql(dialect="tsql") + == "@IF(@runtime_stage = 'creating', ALTER TABLE @SQL('@this_model') ALTER COLUMN id INTEGER NOT NULL)" + ) + + # Statements that don't carry a type are unaffected + assert ( + parse_one("ALTER TABLE x ALTER COLUMN y DROP NOT NULL", read="tsql").sql(dialect="tsql") + == "ALTER TABLE x ALTER COLUMN y DROP NOT NULL" + ) + + def test_model_name_cannot_be_string(): with pytest.raises(ParseError) as parse_error: parse( diff --git a/tests/core/test_model.py b/tests/core/test_model.py index 761f5703df..1f3cde265b 100644 --- a/tests/core/test_model.py +++ b/tests/core/test_model.py @@ -1958,6 +1958,45 @@ def test_render_definition(): assert "def test_macro(evaluator, v):" in d.format_model_expressions(model.render_definition()) +def test_tsql_alter_column_post_statement(make_snapshot: t.Callable) -> None: + # Issue #5932: the trailing NOT NULL made this parse as a Command, which left @this_model + # unresolved and sent the macro to the engine verbatim. + expressions = d.parse( + """ + MODEL ( + name test.test_model, + dialect tsql, + ); + + SELECT 1 AS id; + + @IF(@runtime_stage = 'creating', ALTER TABLE @SQL('@this_model') ALTER COLUMN id INT NOT NULL); + """ + ) + + model = load_sql_based_model(expressions, default_catalog="catalog") + + snapshot = make_snapshot(model) + snapshot.categorize_as(SnapshotChangeCategory.BREAKING) + + post_statements = model.render_post_statements( + snapshots={model.fqn: snapshot}, + runtime_stage=RuntimeStage.CREATING, + ) + + assert len(post_statements) == 1 + assert ( + post_statements[0].sql(dialect="tsql") + == f"ALTER TABLE [catalog].[sqlmesh__test].[test__test_model__{snapshot.version}] /* catalog.test.test_model */ ALTER COLUMN [id] INTEGER NOT NULL" + ) + + # The statement is skipped outside of the creating stage + assert not model.render_post_statements( + snapshots={model.fqn: snapshot}, + runtime_stage=RuntimeStage.EVALUATING, + ) + + def test_render_definition_with_defaults(): query = """ SELECT