Skip to content

Commit

Permalink
Reorganize record/replay code to match dbt-adapters.
Browse files Browse the repository at this point in the history
  • Loading branch information
peterallenwebb committed Jul 16, 2024
1 parent 64e2ea6 commit bdbfc3b
Show file tree
Hide file tree
Showing 5 changed files with 43 additions and 35 deletions.
35 changes: 0 additions & 35 deletions dbt/adapters/postgres/record.py

This file was deleted.

2 changes: 2 additions & 0 deletions dbt/adapters/postgres/record/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
from dbt.adapters.postgres.record.cursor.cursor import PostgresRecordReplayCursor
from dbt.adapters.postgres.record.handle import PostgresRecordReplayHandle
13 changes: 13 additions & 0 deletions dbt/adapters/postgres/record/cursor/cursor.py
Original file line number Diff line number Diff line change
@@ -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

20 changes: 20 additions & 0 deletions dbt/adapters/postgres/record/cursor/status.py
Original file line number Diff line number Diff line change
@@ -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"
8 changes: 8 additions & 0 deletions dbt/adapters/postgres/record/handle.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit bdbfc3b

Please sign in to comment.