reflect dbt-common changes

dbt-labs · May 28, 2024 · 26a8b1f · 26a8b1f
1 parent 36d4f99
commit 26a8b1f
Showing 1 changed file with 11 additions and 7 deletions.
diff --git 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,23 +105,22 @@ 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
 
 
 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")