Skip to content

Commit

Permalink
add query recording options
Browse files Browse the repository at this point in the history
  • Loading branch information
emmyoop committed May 23, 2024
1 parent a28ff8a commit 622188c
Showing 1 changed file with 11 additions and 0 deletions.
11 changes: 11 additions & 0 deletions dbt_common/record.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ class Diff:
class RecorderMode(Enum):
RECORD = 1
REPLAY = 2
RECORD_QUERIES = 3


class Recorder:
Expand Down Expand Up @@ -155,6 +156,10 @@ def get_record_mode_from_env() -> Optional[RecorderMode]:
if record_val is not None and record_val != "0" and record_val.lower() != "false":
return RecorderMode.RECORD

record_val = os.environ.get("DBT_RECORD_QUERIES")
if record_val is not None and record_val != "0" and record_val.lower() != "false":
return RecorderMode.RECORD_QUERIES

return None


Expand All @@ -176,6 +181,12 @@ def record_replay_wrapper(*args, **kwargs):
if recorder is None:
return func_to_record(*args, **kwargs)

if (
recorder.mode == RecorderMode.RECORD_QUERIES
and record_type.__name__ != "QueryRecord"
):
return func_to_record(*args, **kwargs)

# For methods, peel off the 'self' argument before calling the
# params constructor.
param_args = args[1:] if method else args
Expand Down

0 comments on commit 622188c

Please sign in to comment.