Skip to content

Commit 7cdbf81

Browse files
fresioASfresioAS
committed
add new column, copy to new column, drop original column, rename new column to old
1 parent f0acd3b commit 7cdbf81

File tree

2 files changed

+144
-33
lines changed

2 files changed

+144
-33
lines changed

‎sqlmesh/core/engine_adapter/fabric.py‎

Lines changed: 103 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
from sqlmesh.utils.errors import SQLMeshError
1616
from sqlmesh.utils.connection_pool import ConnectionPool
1717
from sqlmesh.core.schema_diff import TableAlterOperation
18+
from sqlmesh.utils import random_id
1819

1920

2021
logger = logging.getLogger(__name__)
@@ -156,15 +157,111 @@ def set_current_catalog(self, catalog_name: str) -> None:
156157
)
157158

158159
def alter_table(
159-
self,
160-
alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]],
160+
self, alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]]
161161
) -> None:
162162
"""
163-
Disables ALTER TABLE for Fabric since it has limited support.
164-
By making this a no-op, we signal to the caller to fall back to a
165-
more reliable drop/add strategy for columns to apply schema changes.
163+
Applies alter expressions to a table. Fabric has limited support for ALTER TABLE,
164+
so this method implements a workaround for column type changes.
165+
This method is self-contained and sets its own catalog context.
166166
"""
167-
return
167+
if not alter_expressions:
168+
return
169+
170+
# Get the target table from the first expression to determine the correct catalog.
171+
first_op = alter_expressions[0]
172+
expression = first_op.expression if isinstance(first_op, TableAlterOperation) else first_op
173+
if not isinstance(expression, exp.Alter) or not expression.this.catalog:
174+
# Fallback for unexpected scenarios
175+
logger.warning(
176+
"Could not determine catalog from alter expression, executing with current context."
177+
)
178+
super().alter_table(alter_expressions)
179+
return
180+
181+
target_catalog = expression.this.catalog
182+
self.set_current_catalog(target_catalog)
183+
184+
with self.transaction():
185+
for op in alter_expressions:
186+
expression = op.expression if isinstance(op, TableAlterOperation) else op
187+
188+
if not isinstance(expression, exp.Alter):
189+
self.execute(expression)
190+
continue
191+
192+
for action in expression.actions:
193+
table_name = expression.this
194+
195+
table_name_without_catalog = table_name.copy()
196+
table_name_without_catalog.set("catalog", None)
197+
198+
is_type_change = isinstance(action, exp.AlterColumn) and action.args.get(
199+
"dtype"
200+
)
201+
202+
if is_type_change:
203+
column_to_alter = action.this
204+
new_type = action.args["dtype"]
205+
temp_column_name_str = f"{column_to_alter.name}__{random_id(short=True)}"
206+
temp_column_name = exp.to_identifier(temp_column_name_str)
207+
208+
logger.info(
209+
"Applying workaround for column '%s' on table '%s' to change type to '%s'.",
210+
column_to_alter.sql(),
211+
table_name.sql(),
212+
new_type.sql(),
213+
)
214+
215+
# Step 1: Add a temporary column.
216+
add_column_expr = exp.Alter(
217+
this=table_name_without_catalog.copy(),
218+
kind="TABLE",
219+
actions=[
220+
exp.ColumnDef(this=temp_column_name.copy(), kind=new_type.copy())
221+
],
222+
)
223+
add_sql = self._to_sql(add_column_expr)
224+
self.execute(add_sql)
225+
226+
# Step 2: Copy and cast data.
227+
update_sql = self._to_sql(
228+
exp.Update(
229+
this=table_name_without_catalog.copy(),
230+
expressions=[
231+
exp.EQ(
232+
this=temp_column_name.copy(),
233+
expression=exp.Cast(
234+
this=column_to_alter.copy(), to=new_type.copy()
235+
),
236+
)
237+
],
238+
)
239+
)
240+
self.execute(update_sql)
241+
242+
# Step 3: Drop the original column.
243+
drop_sql = self._to_sql(
244+
exp.Alter(
245+
this=table_name_without_catalog.copy(),
246+
kind="TABLE",
247+
actions=[exp.Drop(this=column_to_alter.copy(), kind="COLUMN")],
248+
)
249+
)
250+
self.execute(drop_sql)
251+
252+
# Step 4: Rename the temporary column.
253+
old_name_qualified = f"{table_name_without_catalog.sql(dialect=self.dialect)}.{temp_column_name.sql(dialect=self.dialect)}"
254+
new_name_unquoted = column_to_alter.sql(
255+
dialect=self.dialect, identify=False
256+
)
257+
rename_sql = f"EXEC sp_rename '{old_name_qualified}', '{new_name_unquoted}', 'COLUMN'"
258+
self.execute(rename_sql)
259+
else:
260+
# For other alterations, execute directly.
261+
direct_alter_expr = exp.Alter(
262+
this=table_name_without_catalog.copy(), kind="TABLE", actions=[action]
263+
)
264+
self.execute(direct_alter_expr)
168265

169266

170267
class FabricHttpClient:

‎tests/core/engine_adapter/test_fabric.py‎

Lines changed: 41 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88

99
from sqlmesh.core.engine_adapter import FabricEngineAdapter
1010
from tests.core.engine_adapter import to_sql_calls
11-
from sqlmesh.core.engine_adapter.shared import DataObject
1211

1312
pytestmark = [pytest.mark.engine, pytest.mark.fabric]
1413

@@ -73,41 +72,56 @@ def test_insert_overwrite_by_time_partition(adapter: FabricEngineAdapter):
7372
]
7473

7574

76-
def test_replace_query(adapter: FabricEngineAdapter, mocker: MockerFixture):
77-
mocker.patch.object(
78-
adapter,
79-
"_get_data_objects",
80-
return_value=[DataObject(schema="", name="test_table", type="table")],
81-
)
82-
adapter.replace_query(
83-
"test_table", parse_one("SELECT a FROM tbl"), {"a": exp.DataType.build("int")}
75+
def test_alter_table_column_type_workaround(adapter: FabricEngineAdapter, mocker: MockerFixture):
76+
"""
77+
Tests the alter_table method's workaround for changing a column's data type.
78+
"""
79+
# Mock set_current_catalog to avoid connection pool side effects
80+
set_catalog_mock = mocker.patch.object(adapter, "set_current_catalog")
81+
# Mock random_id to have a predictable temporary column name
82+
mocker.patch("sqlmesh.core.engine_adapter.fabric.random_id", return_value="abcdef")
83+
84+
alter_expression = exp.Alter(
85+
this=exp.to_table("my_db.my_schema.my_table"),
86+
actions=[
87+
exp.AlterColumn(
88+
this=exp.to_column("col_a"),
89+
dtype=exp.DataType.build("BIGINT"),
90+
)
91+
],
8492
)
8593

86-
# This behavior is inherited from MSSQLEngineAdapter and should be TRUNCATE + INSERT
87-
assert to_sql_calls(adapter) == [
88-
"TRUNCATE TABLE [test_table];",
89-
"INSERT INTO [test_table] ([a]) SELECT [a] FROM [tbl];",
94+
adapter.alter_table([alter_expression])
95+
96+
set_catalog_mock.assert_called_once_with("my_db")
97+
98+
expected_calls = [
99+
"ALTER TABLE [my_schema].[my_table] ADD [col_a__abcdef] BIGINT;",
100+
"UPDATE [my_schema].[my_table] SET [col_a__abcdef] = CAST([col_a] AS BIGINT);",
101+
"ALTER TABLE [my_schema].[my_table] DROP COLUMN [col_a];",
102+
"EXEC sp_rename 'my_schema.my_table.col_a__abcdef', 'col_a', 'COLUMN'",
90103
]
91104

105+
assert to_sql_calls(adapter) == expected_calls
92106

93-
def test_alter_table_is_noop(adapter: FabricEngineAdapter):
94-
"""
95-
Tests that alter_table is a no-op for Fabric.
96107

97-
The adapter should not execute any SQL, signaling to the caller
98-
that it should use the fallback strategy (drop/add)
99-
to apply schema changes.
108+
def test_alter_table_direct_alteration(adapter: FabricEngineAdapter, mocker: MockerFixture):
109+
"""
110+
Tests the alter_table method for direct alterations like adding a column.
100111
"""
112+
set_catalog_mock = mocker.patch.object(adapter, "set_current_catalog")
113+
101114
alter_expression = exp.Alter(
102-
this=exp.to_table("test_table"),
103-
actions=[
104-
exp.ColumnDef(
105-
this=exp.to_column("new_col"),
106-
kind=exp.DataType(this=exp.DataType.Type.INT),
107-
)
108-
],
115+
this=exp.to_table("my_db.my_schema.my_table"),
116+
actions=[exp.ColumnDef(this=exp.to_column("new_col"), kind=exp.DataType.build("INT"))],
109117
)
110118

111119
adapter.alter_table([alter_expression])
112120

113-
adapter.cursor.execute.assert_not_called()
121+
set_catalog_mock.assert_called_once_with("my_db")
122+
123+
expected_calls = [
124+
"ALTER TABLE [my_schema].[my_table] ADD [new_col] INT;",
125+
]
126+
127+
assert to_sql_calls(adapter) == expected_calls

0 commit comments

Comments
 (0)