Skip to content

Commit

Permalink
fix vars per common changes
Browse files Browse the repository at this point in the history
  • Loading branch information
emmyoop committed May 23, 2024
1 parent 7e0238c commit ea76a40
Showing 1 changed file with 13 additions and 8 deletions.
21 changes: 13 additions & 8 deletions core/dbt/cli/requires.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,12 @@
from dbt_common.events.helpers import get_json_string_utcnow
from dbt_common.exceptions import DbtBaseException as DbtException
from dbt_common.invocation import reset_invocation_id
from dbt_common.record import Recorder, RecorderMode, get_record_mode_from_env
from dbt_common.record import (
Recorder,
RecorderMode,
get_record_mode_from_env,
get_record_types_from_env,
)
from dbt_common.utils import cast_dict_to_dict_of_strings


Expand Down Expand Up @@ -99,26 +104,26 @@ def wrapper(*args, **kwargs):


def setup_record_replay():
rec_mode, rec_type = get_record_mode_from_env()
rec_mode = get_record_mode_from_env()
rec_types = get_record_types_from_env()

recorder: Optional[Recorder] = None
if rec_mode == RecorderMode.REPLAY:
# TODO: handle the case where the recording path is not set
recording_path = os.environ[
"DBT_REPLAY"
"DBT_RECORDER_REPLAY_PATH"
] # TODO: need to add into new logic getting the path as well
recorder = Recorder(RecorderMode.REPLAY, recording_path)
recorder = Recorder(RecorderMode.REPLAY, types=rec_types, recording_path=recording_path)
elif rec_mode == RecorderMode.RECORD:
recorder = Recorder(RecorderMode.RECORD)
elif rec_mode == RecorderMode.RECORD_QUERIES:
recorder = Recorder(RecorderMode.RECORD_QUERIES)
recorder = Recorder(RecorderMode.RECORD, types=rec_types)

get_invocation_context().recorder = recorder


def tear_down_record_replay():
recorder = get_invocation_context().recorder
if recorder is not None:
if recorder.mode == RecorderMode.RECORD or recorder.mode == RecorderMode.RECORD_QUERIES:
if recorder.mode == RecorderMode.RECORD:
recorder.write("recording.json")
elif recorder.mode == RecorderMode.REPLAY:
recorder.write_diffs("replay_diffs.json")
Expand Down

0 comments on commit ea76a40

Please sign in to comment.