From bc92c53f4b17e524bbb5c1c52d0a7fd994ba19b6 Mon Sep 17 00:00:00 2001 From: Ben Knight Date: Wed, 5 Aug 2026 17:37:32 +0000 Subject: [PATCH 1/3] fix: drain and close probe cursors so an attention cannot roll back the build get_column_schema_from_query wants a query's column shape and gets it by executing the query, reading cursor.description, then returning without fetching the rows or closing the cursor. Closing a cursor whose result set the server is still producing makes the driver cancel the request, and the cancel arrives as an attention. Every connection runs SET XACT_ABORT ON (#718), under which SQL Server answers an attention by rolling back the open transaction. None of that raises, so the damage surfaces later and elsewhere: a snapshot loses the staging table it built moments earlier in the same transaction and fails with `Invalid object name '..._dbt_tmp'`, and a contract-enforced model silently loses its in-transaction pre_hook writes. Only queries opening with a CTE were exposed. Anything else is wrapped as `select * from (...) where 1 = 0` by sqlserver__get_empty_subquery_sql and returns nothing, which is why snapshot staging queries (`with snapshot_query as ...`, both check and timestamp strategies) and CTE-headed contract models were the ones that broke, while everything else stayed quiet. The size of the result set is not the trigger; the timing is. The cancel only raises an attention while the request is still in flight, so a few hundred rows at zero client delay is enough, while ~10ms of client work before the close makes even 20MB safe. Existing coverage exercised these paths with tens of narrow rows and so never tripped it -- the snapshot fixtures in test_transactions.py are now sized past the boundary, and the previously untested contract + CTE combination is covered too. _discard_pending_results lives next to _try_drain_nextset rather than introducing a second draining idiom, and is applied to _get_row_count as well, which leaked its cursor on every expand_column_types call. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 1 + dbt/adapters/sqlserver/sqlserver_adapter.py | 43 ++++-- .../sqlserver/sqlserver_connections.py | 48 ++++++ .../adapter/dbt/test_constraints.py | 85 +++++++++++ .../adapter/dbt/test_transactions.py | 88 +++++++++-- .../functional/adapter/dbt/test_unit_tests.py | 52 +++++++ .../mssql/test_expand_column_types.py | 42 ++++++ .../test_get_column_schema_from_query.py | 138 ++++++++++++++++++ 8 files changed, 475 insertions(+), 22 deletions(-) create mode 100644 tests/unit/adapters/mssql/test_get_column_schema_from_query.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e658860b..a28169a2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ #### Bugfixes +- Fix snapshots failing on their second and later runs with `Invalid object name '..._dbt_tmp'`, and contract-enforced models silently losing their in-transaction `pre_hook` writes. `get_column_schema_from_query` reads a query's column shape by executing it, then returned without fetching the rows or closing the cursor. Closing a cursor whose result set the server is still producing makes the driver cancel the request, and SQL Server answers that cancel by rolling back the open transaction, since every connection runs `SET XACT_ABORT ON` (#718). Nothing is raised for any of it, so the snapshot lost the staging table it had just built and failed against it a statement later. The probe now drains and closes its cursor, as does the row-count probe in `expand_column_types`. Only queries opening with a CTE were affected - anything else is wrapped as `select * from (...) where 1 = 0` by `sqlserver__get_empty_subquery_sql` and returns no rows - which is why snapshot staging queries (`with snapshot_query as ...`, both `check` and `timestamp` strategies) and CTE-headed contract models were the ones that broke. - Fix models failing with `Incorrect syntax near '\'` when the schema name needs delimiters, such as a domain-qualified `domain\user`. The clustered columnstore index name embeds the schema and was emitted as a bare identifier, so the generated DDL did not parse. [#409](https://github.com/dbt-msft/dbt-sqlserver/issues/409) - Fix identifiers built inside string literals not being quoted, which broke schema names containing a `.` or a `"`. `OBJECT_ID('schema.table')` returns `NULL` rather than erroring for such a name, so the failures were silent: the drop-before-create guards in `create_table_as` treated an existing table as absent (then hit `Msg 2714`), and the mask introspection in `apply_masks` found no columns, so configured masks were never applied. `sp_rename` was affected too, failing the table rename-swap with `No item by the name of ...`. All now pass quoted, qualified names. [#785](https://github.com/dbt-msft/dbt-sqlserver/issues/785) diff --git a/dbt/adapters/sqlserver/sqlserver_adapter.py b/dbt/adapters/sqlserver/sqlserver_adapter.py index a5a63d9d..fc891540 100644 --- a/dbt/adapters/sqlserver/sqlserver_adapter.py +++ b/dbt/adapters/sqlserver/sqlserver_adapter.py @@ -29,7 +29,10 @@ ) from dbt.adapters.sqlserver.sqlserver_column import SQLServerColumn, SQLServerColumnNative from dbt.adapters.sqlserver.sqlserver_configs import SQLServerConfigs -from dbt.adapters.sqlserver.sqlserver_connections import SQLServerConnectionManager +from dbt.adapters.sqlserver.sqlserver_connections import ( + SQLServerConnectionManager, + _discard_pending_results, +) from dbt.adapters.sqlserver.sqlserver_mask import ColumnMask from dbt.adapters.sqlserver.sqlserver_mask import mask_changes as _mask_changes from dbt.adapters.sqlserver.sqlserver_mask import resolve_masks as _resolve_masks @@ -176,16 +179,33 @@ def _behavior_flags(self) -> List[BehaviorFlag]: @available.parse(lambda *a, **k: []) def get_column_schema_from_query(self, sql: str) -> List[BaseColumn]: - """Get a list of the Columns with names and data types from the given sql.""" + """Get a list of the Columns with names and data types from the given sql. + + Only the result *shape* is wanted, but the query still runs, so the + cursor comes back holding the whole result set. Usually that set is + empty: dbt-core's ``get_column_schema_from_query`` macro wraps the + query first, and ``sqlserver__get_empty_subquery_sql`` renders that as + ``select * from (...) where 1 = 0``. A query that opens with a CTE + cannot be wrapped that way, though, and is passed through untouched + (dbt/include/sqlserver/macros/adapters/columns.sql), so snapshot + staging queries and CTE-headed contract models arrive here in full. + + Either way the cursor must not be abandoned holding rows -- see + ``_discard_pending_results`` for what that costs. + """ _, cursor = self.connections.add_select_query(sql) - columns = [ - self.Column.create( - column_name, self.connections.data_type_code_to_name(column_type_code) - ) - # https://peps.python.org/pep-0249/#description - for column_name, column_type_code, *_ in cursor.description - ] + try: + columns = [ + self.Column.create( + column_name, self.connections.data_type_code_to_name(column_type_code) + ) + # https://peps.python.org/pep-0249/#description + for column_name, column_type_code, *_ in cursor.description + ] + finally: + _discard_pending_results(cursor) + return columns @classmethod @@ -381,7 +401,10 @@ def _get_row_count(self, relation) -> int: """Return the number of rows in the given relation.""" sql = f"SELECT COUNT_BIG(*) FROM {relation}" _, cursor = self.connections.add_select_query(sql) - row = cursor.fetchone() + try: + row = cursor.fetchone() + finally: + _discard_pending_results(cursor) return int(row[0]) if row else 0 def expand_column_types(self, goal, current, max_rows: int = 1000000): diff --git a/dbt/adapters/sqlserver/sqlserver_connections.py b/dbt/adapters/sqlserver/sqlserver_connections.py index 5bd501e8..f1b88184 100644 --- a/dbt/adapters/sqlserver/sqlserver_connections.py +++ b/dbt/adapters/sqlserver/sqlserver_connections.py @@ -234,6 +234,54 @@ def _try_drain_nextset(cursor: Any) -> bool: raise +# Rows per round trip when shedding a result set nobody asked for. Large +# enough that discarding even a big one costs a handful of round trips. +_DISCARD_CHUNK_SIZE = 10000 + + +def _discard_pending_results(cursor: Any) -> None: + """Consume and close *cursor*, leaving nothing for the driver to cancel. + + Closing a cursor whose result set the server is still producing makes the + driver cancel the request, and the cancel reaches SQL Server as an + *attention*. Every connection opened here runs ``SET XACT_ABORT ON`` + (``_apply_session_settings``, for dbt-msft/dbt-sqlserver#718), and SQL + Server answers an attention under ``XACT_ABORT ON`` by rolling back the + open transaction. + + None of that raises -- an attention is not an error -- so a materialization + that abandons a cursor part-way through a build silently loses whatever it + had already created inside that transaction, then fails further on against + relations that no longer exist. Fetching the rows first makes the close an + ordinary one. + + Failures while discarding are logged and swallowed: the caller already has + what it came for, and the connection is about to be reused for real work, + so trouble shedding rows nobody wanted must not become the error the user + sees. + """ + try: + while True: + # ``description`` is None for statements that return no rows at + # all, where fetching would raise rather than yield nothing. + if cursor.description is not None: + while cursor.fetchmany(_DISCARD_CHUNK_SIZE): + pass + # nextset() only advances; whatever rows the set it lands on holds + # still have to be fetched, hence the outer loop. + if not _try_drain_nextset(cursor): + break + except Exception as e: + # AdapterLogger cannot serialize an exception as a log argument, so + # interpolate rather than passing ``e`` through. + logger.debug(f"Discarding a pending result set failed: {e}") + + try: + cursor.close() + except Exception as e: + logger.debug(f"Closing a cursor failed: {e}") + + # Mapping of Apache Arrow type codes (integers) to SQL Server type names. # ADBC cursors report column types as Arrow type codes; this map translates # them for use in get_column_schema_from_query() and related column expansion. diff --git a/tests/functional/adapter/dbt/test_constraints.py b/tests/functional/adapter/dbt/test_constraints.py index 34459093..b763e6ea 100644 --- a/tests/functional/adapter/dbt/test_constraints.py +++ b/tests/functional/adapter/dbt/test_constraints.py @@ -750,3 +750,88 @@ def test__constraints_enforcement_rollback( # Its result includes the expected error messages self.assert_expected_error_messages(failing_results[0].message, expected_error_messages) + + +# --------------------------------------------------------------------------- +# Contract enforcement on a model whose SQL opens with a CTE +# --------------------------------------------------------------------------- +# +# Every fixture above renders to a plain `select`, so +# sqlserver__get_empty_subquery_sql wraps it as `select * from (...) where 1 = 0` +# and the contract probe in columns_spec_ddl.sql costs nothing. That wrapper +# cannot wrap a query that already starts with a CTE, and passes it through +# untouched instead (dbt/include/sqlserver/macros/adapters/columns.sql), so a +# CTE-headed model runs in full just to have its column shape read -- and the +# cursor holding the result is abandoned. +# +# Abandoning it while the server is still working on the request makes the +# driver cancel it, and the attention that sends rolls back the open +# transaction under `SET XACT_ABORT ON` without raising anything. The model's +# in-transaction pre-hook is what that rollback destroys here. Sized megabytes +# past the point where the server has finished streaming; see the note in +# tests/functional/adapter/dbt/test_transactions.py for why there is no +# threshold constant to use instead. +_CTE_CONTRACT_ROWS = 5000 + +cte_contract_model_sql = ( + """ +{{ config( + materialized='table', + contract={'enforced': true}, + pre_hook="INSERT INTO {{ this.schema }}.contract_audit_log (msg) VALUES ('before_main')" +) }} +with source_data as ( + select top (%d) + row_number() over (order by (select null)) as id, + cast(replicate('x', 500) as varchar(8000)) as payload + from sys.all_objects a cross join sys.all_objects b +) +select id, payload from source_data +""" + % _CTE_CONTRACT_ROWS +) + +cte_contract_schema_yml = """ +version: 2 +models: + - name: cte_contract_model + config: + contract: + enforced: true + columns: + - name: id + data_type: bigint + - name: payload + data_type: varchar(8000) +""" + + +class TestCteModelConstraintsColumnsEqual: + """The contract probe must not disturb the transaction it runs inside.""" + + @pytest.fixture(scope="class") + def models(self): + return { + "cte_contract_model.sql": cte_contract_model_sql, + "constraints_schema.yml": cte_contract_schema_yml, + } + + def test_contract_probe_leaves_the_transaction_intact(self, project): + project.run_sql( + "CREATE TABLE {schema}.contract_audit_log (msg varchar(100))", + ) + + results = run_dbt(["run", "-s", "cte_contract_model"]) + assert len(results) == 1 + assert results[0].status == "success" + + rows = project.run_sql("select count(*) from {schema}.contract_audit_log", fetch="one") + assert rows[0] == 1, ( + "the pre-hook ran inside the model's transaction and its row is gone: " + "the contract probe abandoned its cursor, the driver cancelled the " + "request, and XACT_ABORT rolled the transaction back" + ) + + relation = relation_from_name(project.adapter, "cte_contract_model") + built = project.run_sql(f"select count(*) from {relation}", fetch="one") + assert built[0] == _CTE_CONTRACT_ROWS diff --git a/tests/functional/adapter/dbt/test_transactions.py b/tests/functional/adapter/dbt/test_transactions.py index b764f1cd..bde53ec9 100644 --- a/tests/functional/adapter/dbt/test_transactions.py +++ b/tests/functional/adapter/dbt/test_transactions.py @@ -29,10 +29,33 @@ def project_config_update(self): select 1/0 as boom """ -_snapshot_seed_csv = """id,name,updated_at -1,alice,2024-01-01 00:00:00 -2,bob,2024-01-01 00:00:00 -""" +# The snapshot source is sized deliberately. A snapshot's second run probes the +# shape of its staging query through get_column_schema_from_query +# (check_time_data_types -> get_updated_at_column_data_type), and that query +# starts with a CTE, so sqlserver__get_empty_subquery_sql cannot wrap it in +# `where 1 = 0` and it runs in full. If the cursor holding that result set is +# abandoned while the server is still working on the request, the driver +# cancels it; the attention that sends rolls back the open transaction under +# `SET XACT_ABORT ON`, silently, taking the staging table with it. +# +# The failure is a race rather than a size threshold -- measured against SQL +# Server 2022, a few hundred rows is already enough at zero client delay, while +# ~10ms of client-side work before the close makes even 20MB safe -- so there is +# no constant to encode. These fixtures sit megabytes past the boundary so the +# server is unambiguously still streaming, whatever the runner's speed. A +# handful of narrow rows (which is what a seed gives you) never trips it, which +# is why this went unnoticed. +_SNAPSHOT_ROWS = 5000 +_SNAPSHOT_PAYLOAD_WIDTH = 500 + +_snapshot_source_sql = """ +{{ config(materialized='table') }} +select top (%d) + row_number() over (order by (select null)) as id, + cast(replicate('{{ var("payload_char", "x") }}', %d) as varchar(8000)) as payload, + cast('{{ var("snap_updated_at", "2024-01-01") }}' as datetime2) as updated_at +from sys.all_objects a cross join sys.all_objects b +""" % (_SNAPSHOT_ROWS, _SNAPSHOT_PAYLOAD_WIDTH) _snapshot_sql = """ {% snapshot snap %} @@ -42,7 +65,19 @@ def project_config_update(self): strategy='timestamp', updated_at='updated_at', ) }} -select * from {{ ref('snap_seed') }} +select * from {{ ref('snap_source') }} +{% endsnapshot %} +""" + +_snapshot_check_sql = """ +{% snapshot snap_check %} +{{ config( + target_schema=schema, + unique_key='id', + strategy='check', + check_cols=['payload'], +) }} +select * from {{ ref('snap_source') }} {% endsnapshot %} """ @@ -73,15 +108,12 @@ def models(self): select 1/0 as boom """, "failing_model.sql": _failing_model_sql, + "snap_source.sql": _snapshot_source_sql, } - @pytest.fixture(scope="class") - def seeds(self): - return {"snap_seed.csv": _snapshot_seed_csv} - @pytest.fixture(scope="class") def snapshots(self): - return {"snap.sql": _snapshot_sql} + return {"snap.sql": _snapshot_sql, "snap_check.sql": _snapshot_check_sql} def test_table_materialization(self, project): results = run_dbt(["run", "--models", "table_model"]) @@ -137,18 +169,50 @@ def test_side_effect_rolled_back(self, project): assert rows[0] == 0 def test_snapshot_create_and_merge(self, project): - run_dbt(["seed"]) + """Timestamp strategy, with a second run that has changes to write. + + The merge reads the staging table built earlier in the same + transaction, so a probe that silently rolls that transaction back + surfaces here as `Invalid object name '..._dbt_tmp'`. + """ + run_dbt(["run", "--models", "snap_source"]) results = run_dbt(["snapshot", "--select", "snap"]) assert len(results) == 1 assert results[0].status == "success" rows = project.run_sql("select count(*) from {schema}.snap", fetch="one") - assert rows[0] == 2 + assert rows[0] == _SNAPSHOT_ROWS + # move every row's updated_at forward, so the second run has a full + # changeset to stage and merge rather than converging to zero rows + run_dbt(["run", "--models", "snap_source", "--vars", "snap_updated_at: '2024-06-01'"]) results = run_dbt(["snapshot", "--select", "snap"]) assert len(results) == 1 assert results[0].status == "success" + rows = project.run_sql("select count(*) from {schema}.snap", fetch="one") + assert rows[0] == _SNAPSHOT_ROWS * 2 + + def test_snapshot_check_strategy_create_and_merge(self, project): + """Same path as above via the check strategy: both strategies build + their staging query through sqlserver__snapshot_staging_table, so both + hand the probe a CTE-headed query that runs in full.""" + run_dbt(["run", "--models", "snap_source"]) + results = run_dbt(["snapshot", "--select", "snap_check"]) + assert len(results) == 1 + assert results[0].status == "success" + + rows = project.run_sql("select count(*) from {schema}.snap_check", fetch="one") + assert rows[0] == _SNAPSHOT_ROWS + + run_dbt(["run", "--models", "snap_source", "--vars", "payload_char: y"]) + results = run_dbt(["snapshot", "--select", "snap_check"]) + assert len(results) == 1 + assert results[0].status == "success" + + rows = project.run_sql("select count(*) from {schema}.snap_check", fetch="one") + assert rows[0] == _SNAPSHOT_ROWS * 2 + class BaseFailingModelWithSideEffect: @pytest.fixture(scope="class") diff --git a/tests/functional/adapter/dbt/test_unit_tests.py b/tests/functional/adapter/dbt/test_unit_tests.py index 8567c800..f6d902eb 100644 --- a/tests/functional/adapter/dbt/test_unit_tests.py +++ b/tests/functional/adapter/dbt/test_unit_tests.py @@ -127,3 +127,55 @@ def test_unit_test_data_type(self, project, data_types): run_dbt(["test", "--select", "my_model"]) except Exception: raise AssertionError(f"unit test failed when testing model with {sql_value}") + + +# The contract branch of sqlserver__unit_test_create_table_as had no coverage. +# It reaches get_assert_columns_equivalent, and so the same probe that runs a +# CTE-headed query in full for snapshots and contract-enforced models (see +# tests/functional/adapter/dbt/test_constraints.py). It is safe here only +# because a unit test's inputs are replaced by its fixture rows, so the query +# being probed returns a handful of rows however large the real model is -- +# this pins that branch so a change to it does not go unnoticed. +contract_unit_test_model_sql = """ +select tested_column from {{ ref('my_upstream_model') }} +""" + +contract_unit_test_yml = """ +version: 2 +models: + - name: my_contract_model + config: + contract: + enforced: true + columns: + - name: tested_column + data_type: int +unit_tests: + - name: test_my_contract_model + model: my_contract_model + given: + - input: ref('my_upstream_model') + rows: + - {tested_column: 1} + expect: + rows: + - {tested_column: 1} +""" + + +class TestUnitTestWithContract: + @pytest.fixture(scope="class") + def models(self): + return { + "my_upstream_model.sql": upstream_model_sql, + "my_contract_model.sql": contract_unit_test_model_sql, + "schema.yml": contract_unit_test_yml, + } + + def test_unit_test_runs_under_an_enforced_contract(self, project): + results = run_dbt(["run"]) + assert len(results) == 2 + + results = run_dbt(["test", "--select", "my_contract_model"]) + assert len(results) == 1 + assert results[0].status == "pass" diff --git a/tests/unit/adapters/mssql/test_expand_column_types.py b/tests/unit/adapters/mssql/test_expand_column_types.py index 3aafc53b..4cd33439 100644 --- a/tests/unit/adapters/mssql/test_expand_column_types.py +++ b/tests/unit/adapters/mssql/test_expand_column_types.py @@ -166,3 +166,45 @@ def test_varchar_max_to_bounded_does_not_expand(self, adapter): adapter.expand_column_types(goal, current, max_rows=-1) adapter.alter_column_type.assert_not_called() + + +class TestGetRowCount: + """expand_column_types' row-count probe owns the cursor it is handed by + add_select_query, and has to release it.""" + + @pytest.fixture + def raw_adapter(self): + config = MagicMock() + config.flags = {} + config.project_name = "test" + config.credentials.type = "sqlserver" + return SQLServerAdapter(config, MagicMock()) + + @staticmethod + def _cursor(count=42): + """A single-row COUNT_BIG result. Explicit about being exhausted -- + a MagicMock's fetchmany/nextset are truthy forever, which would hang + any caller that drains before closing.""" + cursor = MagicMock() + cursor.fetchone.return_value = (count,) + cursor.fetchmany.return_value = [] + cursor.nextset.return_value = False + return cursor + + @staticmethod + def _attach(adapter, cursor): + adapter.connections.add_select_query = MagicMock(return_value=(MagicMock(), cursor)) + return cursor + + def test_returns_the_count(self, raw_adapter): + self._attach(raw_adapter, self._cursor()) + + assert raw_adapter._get_row_count(make_rel()) == 42 + + def test_closes_the_cursor(self, raw_adapter): + cursor = self._cursor() + self._attach(raw_adapter, cursor) + + raw_adapter._get_row_count(make_rel()) + + cursor.close.assert_called_once() diff --git a/tests/unit/adapters/mssql/test_get_column_schema_from_query.py b/tests/unit/adapters/mssql/test_get_column_schema_from_query.py new file mode 100644 index 00000000..05bdd26a --- /dev/null +++ b/tests/unit/adapters/mssql/test_get_column_schema_from_query.py @@ -0,0 +1,138 @@ +"""A cursor taken from ``add_select_query`` must be drained and closed. + +Abandoning a cursor while the server is still streaming its result makes the +driver cancel the request. The cancel arrives as an *attention*, and every +connection this adapter opens runs under ``SET XACT_ABORT ON`` +(``SQLServerConnectionManager._set_session_options``, #718), so SQL Server +answers an attention by rolling back the open transaction. An attention is not +an error, so nothing is raised and nothing is logged: the caller carries on and +fails later against relations the rollback removed. + +These tests pin the invariant without a warehouse -- the cursor must be closed, +and it must have nothing left outstanding at the moment it is closed. The +end-to-end versions live with the callers that trip it: the snapshot staging +probe in tests/functional/adapter/dbt/test_transactions.py and the contract +probe in tests/functional/adapter/dbt/test_constraints.py. +""" + +from unittest.mock import MagicMock + +import pytest + +from dbt.adapters.sqlserver.sqlserver_adapter import SQLServerAdapter + + +class FakeCursor: + """Records whether anything was still outstanding when it was closed.""" + + def __init__(self, rows, description=None, extra_result_sets=0): + self._rows = list(rows) + self._extra_result_sets = extra_result_sets + # PEP 249: (name, type_code, display_size, internal_size, precision, + # scale, null_ok) + self.description = description or [ + ("id", 4, None, None, None, None, None), + ("payload", 12, None, None, None, None, None), + ] + self.closed = False + self.rows_pending_at_close = None + self.result_sets_pending_at_close = None + self.fetchmany_calls = 0 + + def fetchmany(self, size): + self.fetchmany_calls += 1 + batch, self._rows = self._rows[:size], self._rows[size:] + return batch + + def fetchone(self): + return self._rows.pop(0) if self._rows else None + + def nextset(self): + if self._extra_result_sets: + self._extra_result_sets -= 1 + self._rows = [("row",)] * 25 + return True + return False + + def close(self): + self.closed = True + self.rows_pending_at_close = len(self._rows) + self.result_sets_pending_at_close = self._extra_result_sets + + +@pytest.fixture +def adapter(): + config = MagicMock() + config.flags = {} + config.project_name = "test" + config.credentials.type = "sqlserver" + return SQLServerAdapter(config, MagicMock()) + + +def attach(adapter, cursor): + adapter.connections.add_select_query = MagicMock(return_value=(MagicMock(), cursor)) + adapter.connections.data_type_code_to_name = MagicMock(return_value="varchar") + return cursor + + +class TestGetColumnSchemaFromQuery: + def test_returns_the_column_schema(self, adapter): + attach(adapter, FakeCursor(rows=[("a", "b")] * 100)) + + columns = adapter.get_column_schema_from_query("select 1") + + assert [c.column for c in columns] == ["id", "payload"] + + def test_closes_the_cursor(self, adapter): + cursor = attach(adapter, FakeCursor(rows=[("a", "b")] * 100)) + + adapter.get_column_schema_from_query("select 1") + + assert cursor.closed, "cursor was abandoned instead of closed" + + def test_leaves_no_rows_outstanding_at_close(self, adapter): + cursor = attach(adapter, FakeCursor(rows=[("a", "b")] * 25_000)) + + adapter.get_column_schema_from_query("select 1") + + assert cursor.rows_pending_at_close == 0, ( + "closing with rows still pending makes the driver cancel the request, " + "which rolls back the open transaction under XACT_ABORT ON" + ) + + def test_leaves_no_result_sets_outstanding_at_close(self, adapter): + """Draining must re-fetch after each ``nextset()``, not just advance past it.""" + cursor = attach(adapter, FakeCursor(rows=[("a", "b")] * 50, extra_result_sets=2)) + + adapter.get_column_schema_from_query("select 1") + + assert cursor.result_sets_pending_at_close == 0 + assert cursor.rows_pending_at_close == 0 + + def test_fetches_in_batches_rather_than_one_row_at_a_time(self, adapter): + cursor = attach(adapter, FakeCursor(rows=[("a", "b")] * 100_000)) + + adapter.get_column_schema_from_query("select 1") + + # 100k rows in a bounded number of round trips, not 100k of them + assert 0 < cursor.fetchmany_calls <= 50 + + def test_closes_the_cursor_when_column_building_raises(self, adapter): + cursor = attach(adapter, FakeCursor(rows=[("a", "b")] * 100)) + adapter.connections.data_type_code_to_name = MagicMock(side_effect=ValueError("boom")) + + with pytest.raises(ValueError): + adapter.get_column_schema_from_query("select 1") + + assert cursor.closed, "cursor leaked when column building failed" + + def test_a_failure_while_discarding_is_not_raised_to_the_caller(self, adapter): + """The caller already has its metadata; a discard problem must not become + the error the user sees.""" + cursor = attach(adapter, FakeCursor(rows=[("a", "b")] * 100)) + cursor.fetchmany = MagicMock(side_effect=RuntimeError("driver went away")) + + columns = adapter.get_column_schema_from_query("select 1") + + assert [c.column for c in columns] == ["id", "payload"] + assert cursor.closed From 7cee2570f0b1a3c3da7ac22dd1346d7d9a0eba01 Mon Sep 17 00:00:00 2001 From: Ben Knight Date: Wed, 5 Aug 2026 17:38:04 +0000 Subject: [PATCH 2/3] perf: describe CTE-headed probe queries instead of executing them sqlserver__get_empty_subquery_sql neuters a probe query as `select * from (...) where 1 = 0`, but a query opening with a CTE cannot be wrapped that way and is passed through untouched, so it ran in full purely to have its column names read. A snapshot executed its whole staging query once for the probe and again to build the staging table; a contract-enforced model ran twice per build. Draining that result set (previous commit) makes it safe but not cheap, and the cost tracks data volume. sp_describe_first_result_set compiles the query and reports its shape without scanning, which is all this method ever wanted. It is already how sqlserver__get_columns_in_query handles CTEs (#698). The macro cannot be the fix point: its output is also embedded in a CREATE VIEW body by dbt-core's unit-test materialization, where only a bare SELECT is legal, and reading cursor.description of an `exec sp_describe...` batch would describe the procedure's own result shape rather than the query's. Reported names and types are unchanged. Reading cursor.description reports Python classes, which collapse whole type families -- every integer width arrives as int, every string type as varchar -- so the describe path is mapped back onto exactly those names and fed through Column.create as before, leaving the dbt_sqlserver_use_native_string_types flag working without the mapping knowing it exists. Rather than guess, the describe path declines and lets the caller execute whenever it cannot guarantee agreement: the adbc backend (whose type names derive from Arrow codes and would disagree), a type outside the mapping, a query sp_describe_first_result_set refuses to describe such as one reading a #temp table, or a describe that returns nothing. Every gap degrades to slower, never to silently different. TestCteProbeAvoidsExecution pins both halves: that a CTE-headed probe no longer executes its query, and that the two branches agree across 22 types. The second is a guard rather than a symptom and passes either way -- it is what makes the mapping safe to maintain. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 1 + dbt/adapters/sqlserver/sqlserver_adapter.py | 133 ++++++++++++++++++ .../adapter/dbt/test_constraints.py | 68 +++++++++ 3 files changed, 202 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a28169a2..93a86c6b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ #### Under the hood +- `get_column_schema_from_query` now reads a CTE-headed query's shape with `sp_describe_first_result_set` instead of executing it. Such queries cannot be wrapped as `select * from (...) where 1 = 0`, so they previously ran in full - a snapshot executed its whole staging query once for the probe and again to build the staging table, and a contract-enforced model ran twice per build. Describing compiles without scanning, so the cost no longer tracks data volume. Reported column names and types are unchanged: the describe path is mapped back onto the coarser names reading `cursor.description` produces, and falls back to executing whenever it cannot guarantee that (the `adbc` backend, a type outside the mapping, or a query `sp_describe_first_result_set` declines to describe, such as one reading a `#temp` table). - **Behavior change (generated SQL only):** the few macros that hand-formatted `[bracket]` identifiers now use `adapter.quote()`, matching the `"double quoted"` identifiers `{{ relation }}` already rendered, so one statement no longer mixes both styles. Affects `USE` (now via the existing `get_use_database_sql()` helper), `CREATE SCHEMA`, contract column lists, grantees, index/constraint names and generated test view names. Server-side `QUOTENAME()` keeps brackets by design. Visible only if you parse dbt's generated SQL or have custom macros assuming brackets. [#785](https://github.com/dbt-msft/dbt-sqlserver/issues/785) - Identifier quoting escapes an embedded `"` by doubling it (`ab"cd` → `"ab""cd"`), in both `adapter.quote()` and relation rendering (`SQLServerRelation.quoted`). Needed twice over: a `"` requires no escaping inside `[brackets]` but does inside double quotes, so the delimiter change above would otherwise have rejected names the adapter previously accepted; and the string-literal fix above renders identifiers through the relation, so a schema containing a `"` depends on it. [#785](https://github.com/dbt-msft/dbt-sqlserver/issues/785) diff --git a/dbt/adapters/sqlserver/sqlserver_adapter.py b/dbt/adapters/sqlserver/sqlserver_adapter.py index fc891540..9760b92a 100644 --- a/dbt/adapters/sqlserver/sqlserver_adapter.py +++ b/dbt/adapters/sqlserver/sqlserver_adapter.py @@ -1,4 +1,5 @@ import datetime as _dt +import re from typing import Any, Dict, List, Optional, Tuple, Union import agate @@ -27,6 +28,7 @@ index_config_changes, normalize_drop_unmanaged, ) +from dbt.adapters.sqlserver.sqlserver_auth import is_adbc_backend from dbt.adapters.sqlserver.sqlserver_column import SQLServerColumn, SQLServerColumnNative from dbt.adapters.sqlserver.sqlserver_configs import SQLServerConfigs from dbt.adapters.sqlserver.sqlserver_connections import ( @@ -40,6 +42,58 @@ logger = AdapterLogger("SQLServer") +# Mirrors sqlserver__select_starts_with_cte +# (dbt/include/sqlserver/macros/adapters/columns.sql): a query opening with a +# CTE cannot be neutered as ``select * from (...) where 1 = 0``, so it reaches +# get_column_schema_from_query unwrapped and would otherwise be executed in +# full just to read its column names. +_SQL_COMMENT = re.compile(r"(?s)/\*.*?\*/|--[^\n]*\n") + +# sp_describe_first_result_set reports true SQL Server types; reading +# ``cursor.description`` reports Python classes, which collapse whole families +# (every integer width arrives as ``int``, every string type as ``varchar``). +# Contract comparison comes through this method either way, so the describe +# path is mapped back onto exactly the names the execute path yields via +# ``data_type_code_to_name``. TestCteProbeAvoidsExecution pins the two +# together; a type missing from this map falls back to executing rather than +# guessing. +_SYSTEM_TYPE_TO_EXECUTED_NAME = { + "bigint": "int", + "int": "int", + "smallint": "int", + "tinyint": "int", + "bit": "bit", + "decimal": "decimal", + "numeric": "decimal", + "money": "decimal", + "smallmoney": "decimal", + "float": "float", + "real": "float", + "date": "date", + "time": "time", + "datetime": "datetime2(6)", + "smalldatetime": "datetime2(6)", + "datetime2": "datetime2(6)", + "char": "varchar", + "nchar": "varchar", + "varchar": "varchar", + "nvarchar": "varchar", + "text": "varchar", + "ntext": "varchar", + "xml": "varchar", + "uniqueidentifier": "varchar", + "binary": "varbinary", + "varbinary": "varbinary", + "image": "varbinary", + "timestamp": "varbinary", + "rowversion": "varbinary", + "datetimeoffset": "varbinary", + "sql_variant": "varbinary", + "hierarchyid": "varbinary", + "geography": "varbinary", + "geometry": "varbinary", +} + def _normalize_result_datetimes( result: Union[Tuple, List[Tuple], None], @@ -193,6 +247,11 @@ def get_column_schema_from_query(self, sql: str) -> List[BaseColumn]: Either way the cursor must not be abandoned holding rows -- see ``_discard_pending_results`` for what that costs. """ + if _SQL_COMMENT.sub("", sql).strip().lower().startswith("with"): + described = self._describe_result_set(sql) + if described is not None: + return described + _, cursor = self.connections.add_select_query(sql) try: @@ -208,6 +267,80 @@ def get_column_schema_from_query(self, sql: str) -> List[BaseColumn]: return columns + def _describe_result_set(self, sql: str) -> Optional[List[BaseColumn]]: + """Read a query's column shape without running it, or None to fall back. + + ``sp_describe_first_result_set`` compiles the query and reports its + result shape, which is all this method ever wanted. It is already how + ``sqlserver__get_columns_in_query`` handles CTEs (#698). + + Returns None -- deliberately, rather than raising -- whenever the + describe cannot be trusted to match what executing would have reported: + an unsupported backend, a query it refuses to describe (it cannot see + through ``#temp`` tables, where executing works), or a type absent from + ``_SYSTEM_TYPE_TO_EXECUTED_NAME``. The caller then executes as before, + which is slower but never disagrees with itself. + """ + credentials = self.connections.profile.credentials + if is_adbc_backend(credentials.backend): + # ADBC derives its type names from Arrow codes (int64 -> bigint, + # large_string -> varchar(max)), so the map above -- built for the + # pyodbc/mssql-python collapse -- would make the two branches + # disagree on that backend. + return None + + # Inline rather than bound: mssql-python binds str as varchar and the + # procedure demands nvarchar(max). columns.sql:24 escapes it the same + # way for the same reason. + describe_sql = "exec sp_describe_first_result_set @tsql = N'{}'".format( + sql.replace("'", "''") + ) + + try: + _, cursor = self.connections.add_select_query(describe_sql) + except Exception as e: + logger.debug(f"Could not describe a CTE query, falling back to executing it: {e}") + return None + + try: + fields = [description[0].lower() for description in cursor.description] + rows = cursor.fetchall() + except Exception as e: + logger.debug(f"Could not read a described result set, executing the query: {e}") + return None + finally: + _discard_pending_results(cursor) + + try: + hidden, name, type_name = ( + fields.index("is_hidden"), + fields.index("name"), + fields.index("system_type_name"), + ) + except ValueError: # pragma: no cover - shape is fixed by SQL Server + return None + + columns = [] + for row in rows: + if row[hidden]: + continue + # "varchar(10)" / "decimal(10,2)" -> "varchar" / "decimal" + base_type = str(row[type_name]).split("(")[0].strip().lower() + executed_name = _SYSTEM_TYPE_TO_EXECUTED_NAME.get(base_type) + if executed_name is None or row[name] is None: + logger.debug( + f"Describing a CTE query reported {base_type!r}, which has no " + "equivalent in the executed path; executing it instead" + ) + return None + columns.append(self.Column.create(row[name], executed_name)) + + # Every select has at least one column, so nothing described means + # sp_describe_first_result_set could not work the shape out. Returning + # an empty list would read as "this query has no columns" and surface + # as a baffling contract mismatch; execute instead. + return columns or None + @classmethod def quote(cls, identifier: str) -> str: """Double-quote an identifier, doubling any embedded double quote. diff --git a/tests/functional/adapter/dbt/test_constraints.py b/tests/functional/adapter/dbt/test_constraints.py index b763e6ea..5b2c14af 100644 --- a/tests/functional/adapter/dbt/test_constraints.py +++ b/tests/functional/adapter/dbt/test_constraints.py @@ -835,3 +835,71 @@ def test_contract_probe_leaves_the_transaction_intact(self, project): relation = relation_from_name(project.adapter, "cte_contract_model") built = project.run_sql(f"select count(*) from {relation}", fetch="one") assert built[0] == _CTE_CONTRACT_ROWS + + +# The type set the probe has to report consistently. Contract comparison reads +# these, so whichever way the probe learns a query's shape must agree with the +# other -- see TestCteProbeAvoidsExecution. +_PROBE_TYPE_MATRIX = [ + "bigint", + "int", + "smallint", + "tinyint", + "bit", + "decimal(10,2)", + "numeric(5,1)", + "money", + "float", + "real", + "date", + "time", + "datetime", + "datetime2(3)", + "char(5)", + "nchar(5)", + "varchar(10)", + "nvarchar(20)", + "varchar(max)", + "nvarchar(max)", + "uniqueidentifier", + "varbinary(10)", +] + + +class TestCteProbeAvoidsExecution: + """A CTE-headed query reaches get_column_schema_from_query unwrapped, + because sqlserver__get_empty_subquery_sql cannot neuter it with + `where 1 = 0`. Reading its shape should not mean running it.""" + + def test_cte_probe_does_not_execute_the_query(self, project): + # 1/0 compiles cleanly and only fails when the query actually runs, so + # getting columns back is proof the probe did not execute it. + sql = "with q as (select 1/0 as boom, cast('x' as varchar(10)) as t) select * from q" + + with project.adapter.connection_named("_probe"): + columns = project.adapter.get_column_schema_from_query(sql) + + assert [c.column for c in columns] == ["boom", "t"] + + def test_cte_probe_reports_the_same_types_as_the_wrapped_probe(self, project): + """Guard, not a symptom: this has to hold before and after any change + to how the CTE branch reads metadata, or contract comparisons shift + under models that merely happen to open with a CTE.""" + mismatches = [] + + with project.adapter.connection_named("_probe"): + for type_sql in _PROBE_TYPE_MATRIX: + plain = f"select cast(null as {type_sql}) as c" + wrapped = project.adapter.get_column_schema_from_query( + f"select * from ({plain}) dbt_sbq_tmp where 1 = 0" + ) + cte = project.adapter.get_column_schema_from_query( + f"with q as ({plain}) select * from q" + ) + if [(c.column, c.dtype) for c in cte] != [(c.column, c.dtype) for c in wrapped]: + mismatches.append( + f"{type_sql}: cte={[(c.column, c.dtype) for c in cte]} " + f"wrapped={[(c.column, c.dtype) for c in wrapped]}" + ) + + assert not mismatches, "probe branches disagree on:\n" + "\n".join(mismatches) From 33544280b68ade23f0eb212c00f8897dc096ada8 Mon Sep 17 00:00:00 2001 From: Ben Knight Date: Thu, 6 Aug 2026 06:07:20 +0000 Subject: [PATCH 3/3] fix: resolve probe column types per backend and complete the Arrow maps Two ways the column probe reported a type the driver would not have. The describe path (previous commit) maps sp_describe_first_result_set's system_type_name back onto the name reading cursor.description yields, but that name is the *driver's* choice of Python class, not SQL Server's, and the two ODBC backends do not agree on three of them. mssql-python decodes uniqueidentifier as uuid.UUID, datetimeoffset as datetime and sql_variant as str where pyodbc hands back str, bytearray and bytearray, so one fixed table could not serve both: a CTE-headed model with a uniqueidentifier column described as varchar and executed as uniqueidentifier, and its contract comparison shifted under it. The lookup now follows the backend in use. Two entries resist tabulation and say so rather than guess. datetimeoffset on pyodbc depends on whether add_query has registered the -155 output converter yet -- bytearray on a connection's first query, str on every one after -- so it is absent and the query is executed, which agrees with itself by construction. uniqueidentifier on pyodbc depends on the process-global pyodbc.native_uuid flag, which is readable, so it is read at probe time. The adbc backend, which skips the describe path entirely, had gaps of its own in the Arrow maps. tinyint is SQL Server's only unsigned integer and arrives as uint8, uniqueidentifier as the arrow.uuid extension type; neither was mapped, so the probe raised "Unsupported SQL Server type code" and the model failed outright. real arrives as Arrow "float", which is the 4-byte float -- it was mapped to SQL Server's 8-byte float, so a contract declaring real compared against FLOAT and failed. The integer-keyed map is keyed on Arrow type ids but held values matching no Arrow enum: id 3 (int8) said varchar, 7 (int32) said float, 12 (double) said decimal. Nothing reaches it today -- ADBC reports pyarrow DataType objects, which the name-keyed map handles -- so it was silently wrong rather than broken. Both maps are now derived from the same set and pinned against pyarrow itself instead of transcribed ids. Types the ADBC driver cannot represent at all (datetimeoffset, and sql_variant, which arrives as an Arrow struct) fail before the probe is reached, so the functional matrix covers those two on the ODBC backends only. --- CHANGELOG.md | 2 +- dbt/adapters/sqlserver/sqlserver_adapter.py | 58 +++++++++++-- .../sqlserver/sqlserver_connections.py | 69 +++++++++++---- dbt/adapters/sqlserver/sqlserver_runtime.py | 5 ++ .../adapter/dbt/test_constraints.py | 16 +++- .../test_get_column_schema_from_query.py | 67 ++++++++++++++- .../test_sqlserver_connection_manager.py | 86 +++++++++++++++++-- 7 files changed, 267 insertions(+), 36 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 93a86c6b..9f6f9149 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,7 +18,7 @@ #### Under the hood -- `get_column_schema_from_query` now reads a CTE-headed query's shape with `sp_describe_first_result_set` instead of executing it. Such queries cannot be wrapped as `select * from (...) where 1 = 0`, so they previously ran in full - a snapshot executed its whole staging query once for the probe and again to build the staging table, and a contract-enforced model ran twice per build. Describing compiles without scanning, so the cost no longer tracks data volume. Reported column names and types are unchanged: the describe path is mapped back onto the coarser names reading `cursor.description` produces, and falls back to executing whenever it cannot guarantee that (the `adbc` backend, a type outside the mapping, or a query `sp_describe_first_result_set` declines to describe, such as one reading a `#temp` table). +- `get_column_schema_from_query` now reads a CTE-headed query's shape with `sp_describe_first_result_set` instead of executing it. Such queries cannot be wrapped as `select * from (...) where 1 = 0`, so they previously ran in full - a snapshot executed its whole staging query once for the probe and again to build the staging table, and a contract-enforced model ran twice per build. Describing compiles without scanning, so the cost no longer tracks data volume. Reported column names and types are unchanged: the describe path is mapped back onto the coarser names reading `cursor.description` produces, and falls back to executing whenever it cannot guarantee that (the `adbc` backend, a type outside the mapping, or a query `sp_describe_first_result_set` declines to describe, such as one reading a `#temp` table). The Python class behind a column is the driver's choice rather than SQL Server's, and the backends do not agree on all of them - `mssql-python` decodes `uniqueidentifier`, `datetimeoffset` and `sql_variant` into richer types than `pyodbc` does - so the mapping follows the backend in use, and `datetimeoffset` on `pyodbc` (whose class depends on when the `-155` output converter was registered) falls back to executing. - **Behavior change (generated SQL only):** the few macros that hand-formatted `[bracket]` identifiers now use `adapter.quote()`, matching the `"double quoted"` identifiers `{{ relation }}` already rendered, so one statement no longer mixes both styles. Affects `USE` (now via the existing `get_use_database_sql()` helper), `CREATE SCHEMA`, contract column lists, grantees, index/constraint names and generated test view names. Server-side `QUOTENAME()` keeps brackets by design. Visible only if you parse dbt's generated SQL or have custom macros assuming brackets. [#785](https://github.com/dbt-msft/dbt-sqlserver/issues/785) - Identifier quoting escapes an embedded `"` by doubling it (`ab"cd` → `"ab""cd"`), in both `adapter.quote()` and relation rendering (`SQLServerRelation.quoted`). Needed twice over: a `"` requires no escaping inside `[brackets]` but does inside double quotes, so the delimiter change above would otherwise have rejected names the adapter previously accepted; and the string-literal fix above renders identifiers through the relation, so a schema containing a `"` depends on it. [#785](https://github.com/dbt-msft/dbt-sqlserver/issues/785) diff --git a/dbt/adapters/sqlserver/sqlserver_adapter.py b/dbt/adapters/sqlserver/sqlserver_adapter.py index 9760b92a..fbd02285 100644 --- a/dbt/adapters/sqlserver/sqlserver_adapter.py +++ b/dbt/adapters/sqlserver/sqlserver_adapter.py @@ -28,7 +28,7 @@ index_config_changes, normalize_drop_unmanaged, ) -from dbt.adapters.sqlserver.sqlserver_auth import is_adbc_backend +from dbt.adapters.sqlserver.sqlserver_auth import is_adbc_backend, is_mssql_python_backend from dbt.adapters.sqlserver.sqlserver_column import SQLServerColumn, SQLServerColumnNative from dbt.adapters.sqlserver.sqlserver_configs import SQLServerConfigs from dbt.adapters.sqlserver.sqlserver_connections import ( @@ -39,6 +39,7 @@ from dbt.adapters.sqlserver.sqlserver_mask import mask_changes as _mask_changes from dbt.adapters.sqlserver.sqlserver_mask import resolve_masks as _resolve_masks from dbt.adapters.sqlserver.sqlserver_relation import SQLServerRelation +from dbt.adapters.sqlserver.sqlserver_runtime import _get_pyodbc logger = AdapterLogger("SQLServer") @@ -57,6 +58,10 @@ # ``data_type_code_to_name``. TestCteProbeAvoidsExecution pins the two # together; a type missing from this map falls back to executing rather than # guessing. +# +# The names below are pyodbc's. The class a driver picks for a column is its +# own choice, not SQL Server's, and mssql-python decodes three of these +# differently -- see ``_MSSQL_PYTHON_TYPE_OVERRIDES``. _SYSTEM_TYPE_TO_EXECUTED_NAME = { "bigint": "int", "int": "int", @@ -87,13 +92,54 @@ "image": "varbinary", "timestamp": "varbinary", "rowversion": "varbinary", - "datetimeoffset": "varbinary", + # datetimeoffset is deliberately absent *from this map*: add_query + # registers the -155 output converter after the execute that needed it, + # so pyodbc reports the column as bytearray on a connection's first query + # and as str on every one after. No fixed name mirrors that, so on pyodbc + # describing gives up and the query is executed -- which agrees with + # itself by construction. mssql-python needs no converter and is pinned + # in _MSSQL_PYTHON_TYPE_OVERRIDES. "sql_variant": "varbinary", "hierarchyid": "varbinary", "geography": "varbinary", "geometry": "varbinary", } +# mssql-python decodes three types into richer Python objects than pyodbc +# does: uniqueidentifier as ``uuid.UUID`` (pyodbc: ``str``), datetimeoffset +# as ``datetime`` and sql_variant as ``str`` (pyodbc: ``bytearray`` for +# both). Executing reports those classes, so describing has to agree. +_MSSQL_PYTHON_TYPE_OVERRIDES = { + "uniqueidentifier": "uniqueidentifier", + "datetimeoffset": "datetime2(6)", + "sql_variant": "varchar", +} + + +def _executed_name_for_system_type(base_type: str, backend: Any) -> Optional[str]: + """The name the executed probe would report for a described column type. + + None means "no confident answer" -- the caller then executes the query, + which is slower but cannot disagree with itself. + """ + if is_mssql_python_backend(backend): + overridden = _MSSQL_PYTHON_TYPE_OVERRIDES.get(base_type) + if overridden is not None: + return overridden + + elif base_type == "uniqueidentifier": + # pyodbc yields uuid.UUID or str for a GUID depending on its + # module-level ``native_uuid`` flag -- process-global state anything + # in the process can flip, so read it rather than assume a default. + try: + native_uuid = bool(_get_pyodbc().native_uuid) + except Exception as e: # pragma: no cover - pyodbc is present if in use + logger.debug(f"Could not read pyodbc.native_uuid, executing the query: {e}") + return None + return "uniqueidentifier" if native_uuid else "varchar" + + return _SYSTEM_TYPE_TO_EXECUTED_NAME.get(base_type) + def _normalize_result_datetimes( result: Union[Tuple, List[Tuple], None], @@ -277,9 +323,9 @@ def _describe_result_set(self, sql: str) -> Optional[List[BaseColumn]]: Returns None -- deliberately, rather than raising -- whenever the describe cannot be trusted to match what executing would have reported: an unsupported backend, a query it refuses to describe (it cannot see - through ``#temp`` tables, where executing works), or a type absent from - ``_SYSTEM_TYPE_TO_EXECUTED_NAME``. The caller then executes as before, - which is slower but never disagrees with itself. + through ``#temp`` tables, where executing works), or a type this + backend's driver has no known executed name for. The caller then + executes as before, which is slower but never disagrees with itself. """ credentials = self.connections.profile.credentials if is_adbc_backend(credentials.backend): @@ -326,7 +372,7 @@ def _describe_result_set(self, sql: str) -> Optional[List[BaseColumn]]: continue # "varchar(10)" / "decimal(10,2)" -> "varchar" / "decimal" base_type = str(row[type_name]).split("(")[0].strip().lower() - executed_name = _SYSTEM_TYPE_TO_EXECUTED_NAME.get(base_type) + executed_name = _executed_name_for_system_type(base_type, credentials.backend) if executed_name is None or row[name] is None: logger.debug( f"Describing a CTE query reported {base_type!r}, which has no " diff --git a/dbt/adapters/sqlserver/sqlserver_connections.py b/dbt/adapters/sqlserver/sqlserver_connections.py index f1b88184..5dd0e2a7 100644 --- a/dbt/adapters/sqlserver/sqlserver_connections.py +++ b/dbt/adapters/sqlserver/sqlserver_connections.py @@ -283,27 +283,40 @@ def _discard_pending_results(cursor: Any) -> None: # Mapping of Apache Arrow type codes (integers) to SQL Server type names. -# ADBC cursors report column types as Arrow type codes; this map translates -# them for use in get_column_schema_from_query() and related column expansion. -# Reference: pyarrow type enum values (pa.int32().__class__.__name__ yields the -# Arrow type name string, and the type's ``id`` property is the integer code). +# A cursor reporting a column type as a bare integer means the Arrow type id +# -- ``pa.int32().id`` -- so the keys are exactly those ids, and the names +# must agree with ARROW_STRING_TYPE_TO_NAME below, which is the same mapping +# reached by the printed type name. +# +# The ADBC mssql driver does not take this path: it hands over pyarrow +# ``DataType`` objects, which the string map handles. This is here for a +# cursor implementation that reports ids instead. +# +# ``pa.uuid()`` is deliberately absent. Its id is 31, Arrow's generic +# EXTENSION id, shared by every extension type -- an id alone cannot say +# which one, so the name is the only reliable signal for it. ARROW_TYPE_CODE_TO_NAME: dict[int, str] = { 1: "bit", # pa.bool_() - 3: "varchar", # pa.string() / pa.utf8() - 4: "varbinary", # pa.binary() - 5: "varchar(max)", # pa.large_string() / pa.large_utf8() - 6: "real", # pa.float32() - 7: "float", # pa.float64() - 8: "int", # pa.int32() + 2: "tinyint", # pa.uint8() + 3: "smallint", # pa.int8() + 4: "int", # pa.uint16() + 5: "smallint", # pa.int16() + 6: "bigint", # pa.uint32() + 7: "int", # pa.int32() 9: "bigint", # pa.int64() - 10: "smallint", # pa.int8() - 11: "smallint", # pa.int16() - 12: "decimal", # pa.decimal128() - 14: "date", # pa.date32() - 16: "date", # pa.date64() - 17: "datetime2(6)", # pa.timestamp() - 18: "time", # pa.time32() - 19: "time", # pa.time64() + 11: "real", # pa.float32() + 12: "float", # pa.float64() + 13: "varchar", # pa.string() / pa.utf8() + 14: "varbinary", # pa.binary() + 16: "date", # pa.date32() + 17: "date", # pa.date64() + 18: "datetime2(6)", # pa.timestamp() + 19: "time", # pa.time32() + 20: "time", # pa.time64() + 23: "decimal", # pa.decimal128() + 24: "decimal", # pa.decimal256() + 34: "varchar(max)", # pa.large_string() / pa.large_utf8() + 35: "varbinary", # pa.large_binary() } # Some ADBC drivers / cursor implementations report the Arrow type name as a @@ -313,9 +326,21 @@ def _discard_pending_results(cursor: Any) -> None: "int16": "smallint", "int32": "int", "int64": "bigint", + # tinyint is SQL Server's only unsigned integer, and the ADBC mssql + # driver reports it as uint8 -- the one unsigned Arrow type reachable + # from a SQL Server result set. The wider unsigned types have no T-SQL + # equivalent and are widened to a signed type that holds their range; + # uint64 has none, so it is left to raise rather than silently truncate. + "uint8": "tinyint", + "uint16": "int", + "uint32": "bigint", "float32": "real", "float64": "float", - "float": "float", + # pyarrow prints float32 as "float" and float64 as "double", so a bare + # "float" here is Arrow's 4-byte float -- SQL Server's real, not its + # float. The pyodbc / mssql-python backends never reach this map: they + # report the Python class, which arrives as "". + "float": "real", "double": "float", "string": "varchar", "utf8": "varchar", @@ -324,6 +349,7 @@ def _discard_pending_results(cursor: Any) -> None: "bool": "bit", "boolean": "bit", "decimal128": "decimal", + "decimal256": "decimal", "decimal": "decimal", "date32": "date", "date64": "date", @@ -334,6 +360,11 @@ def _discard_pending_results(cursor: Any) -> None: "timestamp": "datetime2(6)", "binary": "varbinary", "large_binary": "varbinary", + # uniqueidentifier arrives as the canonical Arrow UUID extension type, + # whose printed form carries no "[" or "(" for the base-type split to + # trim, so the full name is the key. + "extension": "uniqueidentifier", + "uuid": "uniqueidentifier", } # Attribute used to stash the in-flight pyodbc / mssql-python cursor on a diff --git a/dbt/adapters/sqlserver/sqlserver_runtime.py b/dbt/adapters/sqlserver/sqlserver_runtime.py index ab4e0527..d2a46a13 100644 --- a/dbt/adapters/sqlserver/sqlserver_runtime.py +++ b/dbt/adapters/sqlserver/sqlserver_runtime.py @@ -45,6 +45,11 @@ class PyodbcModuleProtocol(Protocol): InterfaceError: type[Exception] DatabaseError: type[Exception] pooling: bool + # Whether pyodbc decodes a uniqueidentifier column as uuid.UUID rather + # than str. Read by the CTE probe, which has to report the same type + # executing the query would have. Module-global, so it can change under + # us; nothing here sets it. + native_uuid: bool def connect(self, *args: Any, **kwargs: Any) -> Any: ... diff --git a/tests/functional/adapter/dbt/test_constraints.py b/tests/functional/adapter/dbt/test_constraints.py index 5b2c14af..e628805e 100644 --- a/tests/functional/adapter/dbt/test_constraints.py +++ b/tests/functional/adapter/dbt/test_constraints.py @@ -2,6 +2,7 @@ import pytest +from dbt.adapters.sqlserver.sqlserver_auth import is_adbc_backend from dbt.tests.adapter.constraints.fixtures import ( model_data_type_schema_yml, my_incremental_model_sql, @@ -865,6 +866,16 @@ def test_contract_probe_leaves_the_transaction_intact(self, project): "varbinary(10)", ] +# The ADBC mssql driver cannot represent these at all ("Unknown type +# DATETIMEOFFSET", and sql_variant arrives as an Arrow struct), so selecting +# one fails before the probe is reached. The ODBC-based backends read both, +# and disagree with each other on what Python type they are -- exactly what +# the describe branch has to follow -- so they are still worth pinning there. +_PROBE_TYPE_MATRIX_NON_ADBC = [ + "datetimeoffset", + "sql_variant", +] + class TestCteProbeAvoidsExecution: """A CTE-headed query reaches get_column_schema_from_query unwrapped, @@ -886,9 +897,12 @@ def test_cte_probe_reports_the_same_types_as_the_wrapped_probe(self, project): to how the CTE branch reads metadata, or contract comparisons shift under models that merely happen to open with a CTE.""" mismatches = [] + matrix = list(_PROBE_TYPE_MATRIX) + if not is_adbc_backend(project.adapter.config.credentials.backend): + matrix += _PROBE_TYPE_MATRIX_NON_ADBC with project.adapter.connection_named("_probe"): - for type_sql in _PROBE_TYPE_MATRIX: + for type_sql in matrix: plain = f"select cast(null as {type_sql}) as c" wrapped = project.adapter.get_column_schema_from_query( f"select * from ({plain}) dbt_sbq_tmp where 1 = 0" diff --git a/tests/unit/adapters/mssql/test_get_column_schema_from_query.py b/tests/unit/adapters/mssql/test_get_column_schema_from_query.py index 05bdd26a..5a697145 100644 --- a/tests/unit/adapters/mssql/test_get_column_schema_from_query.py +++ b/tests/unit/adapters/mssql/test_get_column_schema_from_query.py @@ -15,11 +15,16 @@ probe in tests/functional/adapter/dbt/test_constraints.py. """ -from unittest.mock import MagicMock +from types import SimpleNamespace +from unittest.mock import MagicMock, patch import pytest -from dbt.adapters.sqlserver.sqlserver_adapter import SQLServerAdapter +from dbt.adapters.sqlserver.sqlserver_adapter import ( + SQLServerAdapter, + _executed_name_for_system_type, +) +from dbt.adapters.sqlserver.sqlserver_credentials import SQLServerBackend class FakeCursor: @@ -136,3 +141,61 @@ def test_a_failure_while_discarding_is_not_raised_to_the_caller(self, adapter): assert [c.column for c in columns] == ["id", "payload"] assert cursor.closed + + +class TestDescribedTypeFollowsTheBackend: + """``_describe_result_set`` reports the name the *executed* probe would + have reported, and that is the driver's choice of Python class, not SQL + Server's type. The two ODBC backends disagree on three types, so one + fixed table cannot serve both -- a CTE-headed model with a + uniqueidentifier column shifted its contract type when it did.""" + + @pytest.mark.parametrize( + "backend, system_type, expected", + [ + # pyodbc: bytearray for sql_variant (uniqueidentifier is not here + # -- it depends on a live flag, see the next test). + (SQLServerBackend.pyodbc, "sql_variant", "varbinary"), + # ... except datetimeoffset, whose class depends on whether the + # -155 output converter was registered yet. No answer is better + # than a guess: None sends the caller back to executing. + (SQLServerBackend.pyodbc, "datetimeoffset", None), + # mssql-python: uuid.UUID, str, datetime + (SQLServerBackend.mssql_python, "uniqueidentifier", "uniqueidentifier"), + (SQLServerBackend.mssql_python, "sql_variant", "varchar"), + (SQLServerBackend.mssql_python, "datetimeoffset", "datetime2(6)"), + # Types both drivers collapse the same way are backend-independent. + (SQLServerBackend.pyodbc, "bigint", "int"), + (SQLServerBackend.mssql_python, "bigint", "int"), + (SQLServerBackend.pyodbc, "nvarchar", "varchar"), + (SQLServerBackend.mssql_python, "nvarchar", "varchar"), + # Unknown types fall back to executing on either backend. + (SQLServerBackend.pyodbc, "some_future_type", None), + (SQLServerBackend.mssql_python, "some_future_type", None), + ], + ) + def test_executed_name_for_system_type(self, backend, system_type, expected): + assert _executed_name_for_system_type(system_type, backend) == expected + + def test_pyodbc_guid_follows_the_native_uuid_flag(self): + """pyodbc hands back uuid.UUID instead of str when native_uuid is on. + It is process-global and anything in the process can set it, so the + describe branch reads it rather than assuming a default.""" + for native_uuid, expected in ((True, "uniqueidentifier"), (False, "varchar")): + with patch( + "dbt.adapters.sqlserver.sqlserver_adapter._get_pyodbc", + return_value=SimpleNamespace(native_uuid=native_uuid), + ): + actual = _executed_name_for_system_type( + "uniqueidentifier", SQLServerBackend.pyodbc + ) + assert actual == expected + + def test_an_unreadable_pyodbc_falls_back_to_executing(self): + with patch( + "dbt.adapters.sqlserver.sqlserver_adapter._get_pyodbc", + side_effect=RuntimeError("pyodbc is not importable"), + ): + actual = _executed_name_for_system_type("uniqueidentifier", SQLServerBackend.pyodbc) + + assert actual is None diff --git a/tests/unit/adapters/mssql/test_sqlserver_connection_manager.py b/tests/unit/adapters/mssql/test_sqlserver_connection_manager.py index 206c31e4..f9ea6f41 100644 --- a/tests/unit/adapters/mssql/test_sqlserver_connection_manager.py +++ b/tests/unit/adapters/mssql/test_sqlserver_connection_manager.py @@ -30,6 +30,7 @@ ) from dbt.adapters.sqlserver.sqlserver_backend import is_pyodbc_handle as _is_pyodbc_handle from dbt.adapters.sqlserver.sqlserver_connections import ( + ARROW_TYPE_CODE_TO_NAME, SQLServerConnectionManager, ) from dbt.adapters.sqlserver.sqlserver_credentials import ( @@ -690,14 +691,15 @@ def test_data_type_code_to_name_handles_repr_and_arrow_codes() -> None: assert SQLServerConnectionManager.data_type_code_to_name("") == "varchar" assert SQLServerConnectionManager.data_type_code_to_name("int") == "int" - # Arrow integer type codes (ADBC path). + # Arrow integer type codes (ADBC path). The keys are Arrow type ids, so + # these are ``pa.().id`` -- see the pyarrow cross-check below. assert SQLServerConnectionManager.data_type_code_to_name(1) == "bit" # bool_ - assert SQLServerConnectionManager.data_type_code_to_name(3) == "varchar" # string / utf8 - assert SQLServerConnectionManager.data_type_code_to_name(8) == "int" # int32 - assert SQLServerConnectionManager.data_type_code_to_name(7) == "float" # float64 - assert SQLServerConnectionManager.data_type_code_to_name(17) == "datetime2(6)" # timestamp - assert SQLServerConnectionManager.data_type_code_to_name(10) == "smallint" # int8 - assert SQLServerConnectionManager.data_type_code_to_name(5) == "varchar(max)" # large_string + assert SQLServerConnectionManager.data_type_code_to_name(13) == "varchar" # string / utf8 + assert SQLServerConnectionManager.data_type_code_to_name(7) == "int" # int32 + assert SQLServerConnectionManager.data_type_code_to_name(12) == "float" # float64 + assert SQLServerConnectionManager.data_type_code_to_name(18) == "datetime2(6)" # timestamp + assert SQLServerConnectionManager.data_type_code_to_name(3) == "smallint" # int8 + assert SQLServerConnectionManager.data_type_code_to_name(34) == "varchar(max)" # large_string # Unrecognised Arrow integer code raises rather than silently # mis-reporting the column type. @@ -714,6 +716,76 @@ def test_data_type_code_to_name_handles_repr_and_arrow_codes() -> None: SQLServerConnectionManager.data_type_code_to_name("nonexistent_type") +def test_data_type_code_to_name_covers_every_arrow_type_sql_server_emits() -> None: + """The ADBC mssql driver reports column types as pyarrow ``DataType`` + objects, not codes or names. A type missing from the map is not a wrong + name, it is a hard failure the moment a model selects that column, so + every Arrow type a SQL Server result set can produce is pinned here.""" + pa = pytest.importorskip("pyarrow") + + # source SQL Server type -> Arrow type the driver reports -> mapped name + expected = [ + ("bigint", pa.int64(), "bigint"), + ("int", pa.int32(), "int"), + ("smallint", pa.int16(), "smallint"), + ("tinyint", pa.uint8(), "tinyint"), # the only unsigned type in T-SQL + ("bit", pa.bool_(), "bit"), + ("decimal(10,2)", pa.decimal128(10, 2), "decimal"), + ("float", pa.float64(), "float"), # T-SQL float is 8-byte ... + ("real", pa.float32(), "real"), # ... and real is 4-byte + ("date", pa.date32(), "date"), + ("time", pa.time64("ns"), "time"), + ("datetime2", pa.timestamp("us", tz="UTC"), "datetime2(6)"), + ("varchar", pa.string(), "varchar"), + ("varbinary", pa.binary(), "varbinary"), + ("uniqueidentifier", pa.uuid(), "uniqueidentifier"), + ] + + for sql_type, arrow_type, name in expected: + assert SQLServerConnectionManager.data_type_code_to_name(arrow_type) == name, sql_type + + +def test_arrow_type_code_map_keys_are_real_arrow_type_ids() -> None: + """``ARROW_TYPE_CODE_TO_NAME`` is keyed on ``pa.().id``, and the two + Arrow maps have to answer alike -- a cursor reporting ids and one + reporting names must not disagree about the same column. Both are easy to + hand-write wrong, so pyarrow is asked rather than trusted from memory.""" + pa = pytest.importorskip("pyarrow") + + by_id = { + pa.bool_(): "bit", + pa.uint8(): "tinyint", + pa.int8(): "smallint", + pa.uint16(): "int", + pa.int16(): "smallint", + pa.uint32(): "bigint", + pa.int32(): "int", + pa.int64(): "bigint", + pa.float32(): "real", + pa.float64(): "float", + pa.string(): "varchar", + pa.binary(): "varbinary", + pa.date32(): "date", + pa.date64(): "date", + pa.timestamp("us"): "datetime2(6)", + pa.time32("s"): "time", + pa.time64("us"): "time", + pa.decimal128(5, 2): "decimal", + pa.decimal256(40, 2): "decimal", + pa.large_string(): "varchar(max)", + pa.large_binary(): "varbinary", + } + + for arrow_type, name in by_id.items(): + assert ARROW_TYPE_CODE_TO_NAME.get(arrow_type.id) == name, str(arrow_type) + # ... and the name-keyed map agrees with the id-keyed one. + assert SQLServerConnectionManager.data_type_code_to_name(arrow_type) == name + + # Every id in the map is accounted for above: no leftover entry keyed on + # an id no pyarrow type actually has. + assert set(ARROW_TYPE_CODE_TO_NAME) == {t.id for t in by_id} + + def test_mssql_python_active_directory_default_passes() -> None: credentials = SQLServerCredentials( driver=None,