From 7d879f0de92e7e2eccdb541369098e6783c2a294 Mon Sep 17 00:00:00 2001 From: Peter Allen Webb Date: Tue, 25 Jun 2024 13:50:06 -0400 Subject: [PATCH 1/5] Add record/replay support. --- dbt/adapters/snowflake/connections.py | 40 ++++++++++++------ dbt/adapters/snowflake/record.py | 59 +++++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 13 deletions(-) create mode 100644 dbt/adapters/snowflake/record.py diff --git a/dbt/adapters/snowflake/connections.py b/dbt/adapters/snowflake/connections.py index 4db007f19..c786167db 100644 --- a/dbt/adapters/snowflake/connections.py +++ b/dbt/adapters/snowflake/connections.py @@ -36,6 +36,7 @@ DbtConfigError, ) from dbt_common.exceptions import DbtDatabaseError +from dbt_common.record import get_record_mode_from_env, RecorderMode from dbt.adapters.exceptions.connection import FailedToConnectError from dbt.adapters.contracts.connection import AdapterResponse, Connection, Credentials from dbt.adapters.sql import SQLConnectionManager @@ -43,6 +44,7 @@ from dbt_common.events.functions import warn_or_error from dbt.adapters.events.types import AdapterEventWarning, AdapterEventError from dbt_common.ui import line_wrap_message, warning_tag +from dbt.adapters.snowflake.record import SnowflakeRecordReplayHandle if TYPE_CHECKING: import agate @@ -370,20 +372,32 @@ def connect(): if creds.query_tag: session_parameters.update({"QUERY_TAG": creds.query_tag}) + handle = None + + # In replay mode, we won't connect to a real database at all, while + # in record and diff modes we do, but insert an intermediate handle + # object which monitors native connection activity. + rec_mode = get_record_mode_from_env() + handle = None + if rec_mode != RecorderMode.REPLAY: + handle = snowflake.connector.connect( + account=creds.account, + database=creds.database, + schema=creds.schema, + warehouse=creds.warehouse, + role=creds.role, + autocommit=True, + client_session_keep_alive=creds.client_session_keep_alive, + application="dbt", + insecure_mode=creds.insecure_mode, + session_parameters=session_parameters, + **creds.auth_args(), + ) - handle = snowflake.connector.connect( - account=creds.account, - database=creds.database, - schema=creds.schema, - warehouse=creds.warehouse, - role=creds.role, - autocommit=True, - client_session_keep_alive=creds.client_session_keep_alive, - application="dbt", - insecure_mode=creds.insecure_mode, - session_parameters=session_parameters, - **creds.auth_args(), - ) + if rec_mode is not None: + # If using the record/replay mechanism, regardless of mode, we + # use a wrapper. + handle = SnowflakeRecordReplayHandle(handle, connection) return handle diff --git a/dbt/adapters/snowflake/record.py b/dbt/adapters/snowflake/record.py new file mode 100644 index 000000000..6d475b3e8 --- /dev/null +++ b/dbt/adapters/snowflake/record.py @@ -0,0 +1,59 @@ +import dataclasses +from typing import Optional + +from dbt.adapters.record import RecordReplayHandle, RecordReplayCursor +from dbt_common.record import record_function, Record, Recorder + + +class SnowflakeRecordReplayHandle(RecordReplayHandle): + def cursor(self): + cursor = None if self.native_handle is None else self.native_handle.cursor() + return SnowflakeRecordReplayCursor(cursor, self.connection) + + +@dataclasses.dataclass +class CursorGetSqlStateParams: + connection_name: str + + +@dataclasses.dataclass +class CursorGetSqlStateResult: + msg: Optional[str] + + +class CursorGetSqlStateRecord(Record): + params_cls = CursorGetSqlStateParams + result_cls = CursorGetSqlStateResult + + +Recorder.register_record_type(CursorGetSqlStateRecord) + + +@dataclasses.dataclass +class CursorGetSqfidParams: + connection_name: str + + +@dataclasses.dataclass +class CursorGetSqfidResult: + msg: Optional[str] + + +class CursorGetSqfidRecord(Record): + params_cls = CursorGetSqfidParams + result_cls = CursorGetSqfidResult + + +Recorder.register_record_type(CursorGetSqfidRecord) + + +class SnowflakeRecordReplayCursor(RecordReplayCursor): + @property + @record_function(CursorGetSqlStateRecord, method=True, id_field_name="connection_name") + def sqlstate(self): + return self.native_cursor.sqlstate + + @property + @record_function(CursorGetSqfidRecord, method=True, id_field_name="connection_name") + def sfqid(self): + return self.native_cursor.sfqid From 247a593f88d788214c6ac4eb5b120103bf9353bf Mon Sep 17 00:00:00 2001 From: Peter Allen Webb Date: Sun, 7 Jul 2024 15:08:06 -0400 Subject: [PATCH 2/5] Add group to record types. --- dbt/adapters/snowflake/record.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/dbt/adapters/snowflake/record.py b/dbt/adapters/snowflake/record.py index 6d475b3e8..92671b1ef 100644 --- a/dbt/adapters/snowflake/record.py +++ b/dbt/adapters/snowflake/record.py @@ -21,12 +21,11 @@ class CursorGetSqlStateResult: msg: Optional[str] +@Recorder.register_record_type class CursorGetSqlStateRecord(Record): params_cls = CursorGetSqlStateParams result_cls = CursorGetSqlStateResult - - -Recorder.register_record_type(CursorGetSqlStateRecord) + group = "Database" @dataclasses.dataclass @@ -39,12 +38,11 @@ class CursorGetSqfidResult: msg: Optional[str] +@Recorder.register_record_type class CursorGetSqfidRecord(Record): params_cls = CursorGetSqfidParams result_cls = CursorGetSqfidResult - - -Recorder.register_record_type(CursorGetSqfidRecord) + group = "Database" class SnowflakeRecordReplayCursor(RecordReplayCursor): From 8e0bf4c42c98b52868af2c0ca0a0ef2e8b69621c Mon Sep 17 00:00:00 2001 From: Peter Allen Webb Date: Tue, 16 Jul 2024 17:45:09 -0400 Subject: [PATCH 3/5] Re-organize record/replay code to match dbt-adapters --- dbt/adapters/snowflake/record.py | 57 ------------------- dbt/adapters/snowflake/record/__init__.py | 2 + .../snowflake/record/cursor/cursor.py | 21 +++++++ dbt/adapters/snowflake/record/cursor/sfqid.py | 21 +++++++ .../snowflake/record/cursor/sqlstate.py | 21 +++++++ dbt/adapters/snowflake/record/handle.py | 12 ++++ 6 files changed, 77 insertions(+), 57 deletions(-) delete mode 100644 dbt/adapters/snowflake/record.py create mode 100644 dbt/adapters/snowflake/record/__init__.py create mode 100644 dbt/adapters/snowflake/record/cursor/cursor.py create mode 100644 dbt/adapters/snowflake/record/cursor/sfqid.py create mode 100644 dbt/adapters/snowflake/record/cursor/sqlstate.py create mode 100644 dbt/adapters/snowflake/record/handle.py diff --git a/dbt/adapters/snowflake/record.py b/dbt/adapters/snowflake/record.py deleted file mode 100644 index 92671b1ef..000000000 --- a/dbt/adapters/snowflake/record.py +++ /dev/null @@ -1,57 +0,0 @@ -import dataclasses -from typing import Optional - -from dbt.adapters.record import RecordReplayHandle, RecordReplayCursor -from dbt_common.record import record_function, Record, Recorder - - -class SnowflakeRecordReplayHandle(RecordReplayHandle): - def cursor(self): - cursor = None if self.native_handle is None else self.native_handle.cursor() - return SnowflakeRecordReplayCursor(cursor, self.connection) - - -@dataclasses.dataclass -class CursorGetSqlStateParams: - connection_name: str - - -@dataclasses.dataclass -class CursorGetSqlStateResult: - msg: Optional[str] - - -@Recorder.register_record_type -class CursorGetSqlStateRecord(Record): - params_cls = CursorGetSqlStateParams - result_cls = CursorGetSqlStateResult - group = "Database" - - -@dataclasses.dataclass -class CursorGetSqfidParams: - connection_name: str - - -@dataclasses.dataclass -class CursorGetSqfidResult: - msg: Optional[str] - - -@Recorder.register_record_type -class CursorGetSqfidRecord(Record): - params_cls = CursorGetSqfidParams - result_cls = CursorGetSqfidResult - group = "Database" - - -class SnowflakeRecordReplayCursor(RecordReplayCursor): - @property - @record_function(CursorGetSqlStateRecord, method=True, id_field_name="connection_name") - def sqlstate(self): - return self.native_cursor.sqlstate - - @property - @record_function(CursorGetSqfidRecord, method=True, id_field_name="connection_name") - def sfqid(self): - return self.native_cursor.sfqid diff --git a/dbt/adapters/snowflake/record/__init__.py b/dbt/adapters/snowflake/record/__init__.py new file mode 100644 index 000000000..f763dc3a4 --- /dev/null +++ b/dbt/adapters/snowflake/record/__init__.py @@ -0,0 +1,2 @@ +from dbt.adapters.snowflake.record.cursor.cursor import SnowflakeRecordReplayCursor +from dbt.adapters.snowflake.record.handle import SnowflakeRecordReplayHandle diff --git a/dbt/adapters/snowflake/record/cursor/cursor.py b/dbt/adapters/snowflake/record/cursor/cursor.py new file mode 100644 index 000000000..a07468867 --- /dev/null +++ b/dbt/adapters/snowflake/record/cursor/cursor.py @@ -0,0 +1,21 @@ +from dbt_common.record import record_function + +from dbt.adapters.record import RecordReplayCursor +from dbt.adapters.snowflake.record.cursor.sfqid import CursorGetSfqidRecord +from dbt.adapters.snowflake.record.cursor.sqlstate import CursorGetSqlStateRecord + + +class SnowflakeRecordReplayCursor(RecordReplayCursor): + """A custom extension of RecordReplayCursor that adds the sqlstate + and sfqid properties which are specific to snowflake-connector.""" + + @property + @property + @record_function(CursorGetSqlStateRecord, method=True, id_field_name="connection_name") + def sqlstate(self): + return self.native_cursor.sqlstate + + @property + @record_function(CursorGetSfqidRecord, method=True, id_field_name="connection_name") + def sfqid(self): + return self.native_cursor.sfqid diff --git a/dbt/adapters/snowflake/record/cursor/sfqid.py b/dbt/adapters/snowflake/record/cursor/sfqid.py new file mode 100644 index 000000000..e39c857d3 --- /dev/null +++ b/dbt/adapters/snowflake/record/cursor/sfqid.py @@ -0,0 +1,21 @@ +import dataclasses +from typing import Optional + +from dbt_common.record import Record, Recorder + + +@dataclasses.dataclass +class CursorGetSfqidParams: + connection_name: str + + +@dataclasses.dataclass +class CursorGetSfqidResult: + msg: Optional[str] + + +@Recorder.register_record_type +class CursorGetSfqidRecord(Record): + params_cls = CursorGetSfqidParams + result_cls = CursorGetSfqidResult + group = "Database" diff --git a/dbt/adapters/snowflake/record/cursor/sqlstate.py b/dbt/adapters/snowflake/record/cursor/sqlstate.py new file mode 100644 index 000000000..5619058fd --- /dev/null +++ b/dbt/adapters/snowflake/record/cursor/sqlstate.py @@ -0,0 +1,21 @@ +import dataclasses +from typing import Optional + +from dbt_common.record import Record, Recorder + + +@dataclasses.dataclass +class CursorGetSqlStateParams: + connection_name: str + + +@dataclasses.dataclass +class CursorGetSqlStateResult: + msg: Optional[str] + + +@Recorder.register_record_type +class CursorGetSqlStateRecord(Record): + params_cls = CursorGetSqlStateParams + result_cls = CursorGetSqlStateResult + group = "Database" diff --git a/dbt/adapters/snowflake/record/handle.py b/dbt/adapters/snowflake/record/handle.py new file mode 100644 index 000000000..046bb911b --- /dev/null +++ b/dbt/adapters/snowflake/record/handle.py @@ -0,0 +1,12 @@ +from dbt.adapters.record import RecordReplayHandle + +from dbt.adapters.snowflake.record.cursor.cursor import SnowflakeRecordReplayCursor + + +class SnowflakeRecordReplayHandle(RecordReplayHandle): + """A custom extension of RecordReplayHandle that returns a + snowflake-connector-specific SnowflakeRecordReplayCursor object.""" + + def cursor(self): + cursor = None if self.native_handle is None else self.native_handle.cursor() + return SnowflakeRecordReplayCursor(cursor, self.connection) From 8d31e978841b8db5f2cda2a90d10c2f971644d63 Mon Sep 17 00:00:00 2001 From: Peter Allen Webb Date: Tue, 16 Jul 2024 17:47:25 -0400 Subject: [PATCH 4/5] Add changelog entry. --- .changes/unreleased/Under the Hood-20240716-174655.yaml | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 .changes/unreleased/Under the Hood-20240716-174655.yaml diff --git a/.changes/unreleased/Under the Hood-20240716-174655.yaml b/.changes/unreleased/Under the Hood-20240716-174655.yaml new file mode 100644 index 000000000..61174fc64 --- /dev/null +++ b/.changes/unreleased/Under the Hood-20240716-174655.yaml @@ -0,0 +1,6 @@ +kind: Under the Hood +body: Add record/replay support. +time: 2024-07-16T17:46:55.11204-04:00 +custom: + Author: peterallenwebb + Issue: "1106" From 0de8a6e02df676ea80417cf7e499c64e69330a48 Mon Sep 17 00:00:00 2001 From: Peter Webb Date: Tue, 16 Jul 2024 17:51:16 -0400 Subject: [PATCH 5/5] Update .changes/unreleased/Under the Hood-20240716-174655.yaml Co-authored-by: Colin Rogers <111200756+colin-rogers-dbt@users.noreply.github.com> --- .changes/unreleased/Under the Hood-20240716-174655.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.changes/unreleased/Under the Hood-20240716-174655.yaml b/.changes/unreleased/Under the Hood-20240716-174655.yaml index 61174fc64..14c3c8d76 100644 --- a/.changes/unreleased/Under the Hood-20240716-174655.yaml +++ b/.changes/unreleased/Under the Hood-20240716-174655.yaml @@ -1,5 +1,5 @@ kind: Under the Hood -body: Add record/replay support. +body: Add support for experimental record/replay testing. time: 2024-07-16T17:46:55.11204-04:00 custom: Author: peterallenwebb