diff --git a/.changes/unreleased/Under the Hood-20241106-113249.yaml b/.changes/unreleased/Under the Hood-20241106-113249.yaml new file mode 100644 index 000000000..0437a8c88 --- /dev/null +++ b/.changes/unreleased/Under the Hood-20241106-113249.yaml @@ -0,0 +1,6 @@ +kind: Under the Hood +body: remove SnowflakeAdapterResponse in favor of updated AdapterResponse in base +time: 2024-11-06T11:32:49.503467-08:00 +custom: + Author: colin-rogers-dbt + Issue: "1233" diff --git a/dbt/adapters/snowflake/connections.py b/dbt/adapters/snowflake/connections.py index 10bee30f0..fc2c09c19 100644 --- a/dbt/adapters/snowflake/connections.py +++ b/dbt/adapters/snowflake/connections.py @@ -84,11 +84,6 @@ def snowflake_private_key(private_key: RSAPrivateKey) -> bytes: ) -@dataclass -class SnowflakeAdapterResponse(AdapterResponse): - query_id: str = "" - - @dataclass class SnowflakeCredentials(Credentials): account: str @@ -447,17 +442,17 @@ def cancel(self, connection): logger.debug("Cancel query '{}': {}".format(connection_name, res)) @classmethod - def get_response(cls, cursor) -> SnowflakeAdapterResponse: + def get_response(cls, cursor) -> AdapterResponse: code = cursor.sqlstate if code is None: code = "SUCCESS" - - return SnowflakeAdapterResponse( + query_id = str(cursor.sfqid) if cursor.sfqid is not None else None + return AdapterResponse( _message="{} {}".format(code, cursor.rowcount), rows_affected=cursor.rowcount, code=code, - query_id=cursor.sfqid, + query_id=query_id, ) # disable transactional logic by default on Snowflake diff --git a/tests/unit/test_snowflake_adapter.py b/tests/unit/test_snowflake_adapter.py index 32e73eb45..aa580aad2 100644 --- a/tests/unit/test_snowflake_adapter.py +++ b/tests/unit/test_snowflake_adapter.py @@ -60,8 +60,10 @@ def setUp(self): self.handle = mock.MagicMock(spec=snowflake_connector.SnowflakeConnection) self.cursor = self.handle.cursor.return_value self.mock_execute = self.cursor.execute + self.mock_execute.return_value = mock.MagicMock(sfqid="42") self.patcher = mock.patch("dbt.adapters.snowflake.connections.snowflake.connector.connect") self.snowflake = self.patcher.start() + self.snowflake.connect.cursor.return_value = mock.MagicMock(sfqid="42") # Create the Manifest.state_check patcher @mock.patch("dbt.parser.manifest.ManifestLoader.build_manifest_state_check") @@ -90,7 +92,6 @@ def _mock_state_check(self): self.qh_patch = mock.patch.object(self.adapter.connections.query_header, "add") self.mock_query_header_add = self.qh_patch.start() self.mock_query_header_add.side_effect = lambda q: "/* dbt */\n{}".format(q) - self.adapter.acquire_connection() inject_adapter(self.adapter, SnowflakePlugin)