Skip to content

Commit

Permalink
remove SnowflakeAdapterResponse in favor of updated AdapterResponse i…
Browse files Browse the repository at this point in the history
…n base
  • Loading branch information
colin-rogers-dbt committed Nov 6, 2024
1 parent a16654b commit 0e17068
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 10 deletions.
13 changes: 4 additions & 9 deletions dbt/adapters/snowflake/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion tests/unit/test_snowflake_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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)

Expand Down

0 comments on commit 0e17068

Please sign in to comment.