diff --git a/dbt/adapters/postgres/record.py b/dbt/adapters/postgres/record.py deleted file mode 100644 index 52248a7..0000000 --- a/dbt/adapters/postgres/record.py +++ /dev/null @@ -1,35 +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 PostgresRecordReplayHandle(RecordReplayHandle): - def cursor(self): - cursor = None if self.native_handle is None else self.native_handle.cursor() - return PostgresRecordReplayCursor(cursor, self.connection) - - -@dataclasses.dataclass -class CursorGetStatusMessageParams: - connection_name: str - - -@dataclasses.dataclass -class CursorGetStatusMessageResult: - msg: Optional[str] - - -@Recorder.register_record_type -class CursorGetStatusMessageRecord(Record): - params_cls = CursorGetStatusMessageParams - result_cls = CursorGetStatusMessageResult - group = "Database" - - -class PostgresRecordReplayCursor(RecordReplayCursor): - @property - @record_function(CursorGetStatusMessageRecord, method=True, id_field_name="connection_name") - def statusmessage(self): - return self.native_cursor.statusmessage diff --git a/dbt/adapters/postgres/record/__init__.py b/dbt/adapters/postgres/record/__init__.py new file mode 100644 index 0000000..9b8b9b3 --- /dev/null +++ b/dbt/adapters/postgres/record/__init__.py @@ -0,0 +1,2 @@ +from dbt.adapters.postgres.record.cursor.cursor import PostgresRecordReplayCursor +from dbt.adapters.postgres.record.handle import PostgresRecordReplayHandle diff --git a/dbt/adapters/postgres/record/cursor/cursor.py b/dbt/adapters/postgres/record/cursor/cursor.py new file mode 100644 index 0000000..fc13e3c --- /dev/null +++ b/dbt/adapters/postgres/record/cursor/cursor.py @@ -0,0 +1,13 @@ +from dbt_common.record import record_function + +from dbt.adapters.record import RecordReplayCursor + +from dbt.adapters.postgres.record.cursor.status import CursorGetStatusMessageRecord + + +class PostgresRecordReplayCursor(RecordReplayCursor): + @property + @record_function(CursorGetStatusMessageRecord, method=True, id_field_name="connection_name") + def statusmessage(self): + return self.native_cursor.statusmessage + diff --git a/dbt/adapters/postgres/record/cursor/status.py b/dbt/adapters/postgres/record/cursor/status.py new file mode 100644 index 0000000..55b4075 --- /dev/null +++ b/dbt/adapters/postgres/record/cursor/status.py @@ -0,0 +1,20 @@ +import dataclasses +from typing import Optional + +from dbt_common.record import Record, Recorder + +@dataclasses.dataclass +class CursorGetStatusMessageParams: + connection_name: str + + +@dataclasses.dataclass +class CursorGetStatusMessageResult: + msg: Optional[str] + + +@Recorder.register_record_type +class CursorGetStatusMessageRecord(Record): + params_cls = CursorGetStatusMessageParams + result_cls = CursorGetStatusMessageResult + group = "Database" diff --git a/dbt/adapters/postgres/record/handle.py b/dbt/adapters/postgres/record/handle.py new file mode 100644 index 0000000..c6a589a --- /dev/null +++ b/dbt/adapters/postgres/record/handle.py @@ -0,0 +1,8 @@ +from dbt.adapters.record import RecordReplayHandle + +from dbt.adapters.postgres.record.cursor.cursor import PostgresRecordReplayCursor + +class PostgresRecordReplayHandle(RecordReplayHandle): + def cursor(self): + cursor = None if self.native_handle is None else self.native_handle.cursor() + return PostgresRecordReplayCursor(cursor, self.connection)