Allow users to set the flag full-refresh when running models

astronomer · Sep 12, 2023 · fe36958 · fe36958
1 parent d09abc5
commit fe36958
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/cosmos/operators/local.py b/cosmos/operators/local.py
@@ -405,10 +405,21 @@ class DbtRunLocalOperator(DbtLocalBaseOperator):
     ui_color = "#7352BA"
     ui_fgcolor = "#F4F2FC"
 
-    def __init__(self, **kwargs: Any) -> None:
+    def __init__(self, full_refresh: bool = False, **kwargs: Any) -> None:
+        self.full_refresh = full_refresh
         super().__init__(**kwargs)
         self.base_cmd = ["run"]
 
+    def add_cmd_flags(self) -> list[str]:
+        flags = []
+        if self.full_refresh is True:
+            flags.append("--full-refresh")
+        return flags
+
+    def execute(self, context: Context) -> None:
+        cmd_flags = self.add_cmd_flags()
+        self.build_and_run_cmd(context=context, cmd_flags=cmd_flags)
+
 
 class DbtTestLocalOperator(DbtLocalBaseOperator):
     """

diff --git a/tests/operators/test_local.py b/tests/operators/test_local.py
@@ -245,6 +245,7 @@ def test_store_compiled_sql() -> None:
     "operator_class,kwargs,expected_call_kwargs",
     [
         (DbtSeedLocalOperator, {"full_refresh": True}, {"context": {}, "cmd_flags": ["--full-refresh"]}),
+        (DbtRunLocalOperator, {"full_refresh": True}, {"context": {}, "cmd_flags": ["--full-refresh"]}),
         (
             DbtRunOperationLocalOperator,
             {"args": {"days": 7, "dry_run": True}, "macro_name": "bla"},
@@ -264,7 +265,6 @@ def test_operator_execute_with_flags(mock_build_and_run_cmd, operator_class, kwa
     (
         DbtLSLocalOperator,
         DbtSnapshotLocalOperator,
-        DbtRunLocalOperator,
         DbtTestLocalOperator,
         DbtDocsLocalOperator,
         DbtDocsS3LocalOperator,