[Bugfix] Dbt docs operator should not look for graph.gpickle file w…

…hen `--no-write-json` is passed. (#883) Fixes #882 --------- Co-authored-by: Tatiana Al-Chueyr <[email protected]>
astronomer · Mar 21, 2024 · 4e3a3c8 · 4e3a3c8
1 parent ddb25cb
commit 4e3a3c8
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 deletions.
diff --git a/cosmos/operators/local.py b/cosmos/operators/local.py
@@ -570,19 +570,21 @@ class DbtDocsLocalOperator(DbtLocalBaseOperator):
     """
 
     ui_color = "#8194E0"
-    required_files = ["index.html", "manifest.json", "graph.gpickle", "catalog.json"]
+    required_files = ["index.html", "manifest.json", "catalog.json"]
     base_cmd = ["docs", "generate"]
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.check_static_flag()
 
     def check_static_flag(self) -> None:
-        flag = "--static"
         if self.dbt_cmd_flags:
-            if flag in self.dbt_cmd_flags:
+            if "--static" in self.dbt_cmd_flags:
                 # For the --static flag we only upload the generated static_index.html file
                 self.required_files = ["static_index.html"]
+        if self.dbt_cmd_global_flags:
+            if "--no-write-json" in self.dbt_cmd_global_flags and "graph.gpickle" in self.required_files:
+                self.required_files.remove("graph.gpickle")
 
 
 class DbtDocsCloudLocalOperator(DbtDocsLocalOperator, ABC):

diff --git a/tests/operators/test_local.py b/tests/operators/test_local.py
@@ -759,6 +759,21 @@ def test_dbt_docs_local_operator_with_static_flag():
     assert operator.required_files == ["static_index.html"]
 
 
+def test_dbt_docs_local_operator_ignores_graph_gpickle():
+    # Check when --no-write-json is passed, graph.gpickle is removed.
+    # This is only currently relevant for subclasses, but will become more generally relevant in the future.
+    class CustomDbtDocsLocalOperator(DbtDocsLocalOperator):
+        required_files = ["index.html", "manifest.json", "graph.gpickle", "catalog.json"]
+
+    operator = CustomDbtDocsLocalOperator(
+        task_id="fake-task",
+        project_dir="fake-dir",
+        profile_config=profile_config,
+        dbt_cmd_global_flags=["--no-write-json"],
+    )
+    assert operator.required_files == ["index.html", "manifest.json", "catalog.json"]
+
+
 @patch("cosmos.hooks.subprocess.FullOutputSubprocessHook.send_sigint")
 def test_dbt_local_operator_on_kill_sigint(mock_send_sigint) -> None: