From 26a8b1f9ee557aeb506226c8b8aa835b78c4cd52 Mon Sep 17 00:00:00 2001 From: Emily Rockman Date: Tue, 28 May 2024 14:22:18 -0500 Subject: [PATCH] reflect dbt-common changes --- core/dbt/cli/requires.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/core/dbt/cli/requires.py b/core/dbt/cli/requires.py index a7a282f2891..fd3eb49a625 100644 --- a/core/dbt/cli/requires.py +++ b/core/dbt/cli/requires.py @@ -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 @@ -100,15 +105,14 @@ def wrapper(*args, **kwargs): def setup_record_replay(): rec_mode = get_record_mode_from_env() + rec_types = get_record_types_from_env() recorder: Optional[Recorder] = None if rec_mode == RecorderMode.REPLAY: - recording_path = os.environ["DBT_REPLAY"] - recorder = Recorder(RecorderMode.REPLAY, recording_path) + recording_path = os.environ.get("DBT_RECORDER_REPLAY_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 @@ -116,7 +120,7 @@ def setup_record_replay(): 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")