Add more template fields to DbtBaseOperator (#786)

This partially addresses #754 via allowing for built-in templating
support for the `DbtBaseOperator`.

I also noticed `--full-refresh` was not documented so I added that in.

## Still missing

Manual run pattern is not documented; the fact that these fields are
templated is not documented. I don't really know where in the docs to
put this. The docs are very API-focused more than narrative-based or
suggestive, and Cosmos's maintainers prefer this style of documentation
so it's hard to find a spot for this. It's possible that that's fine and
we just keep this as a feature for more advanced users who dig into
source code to discover for themselves? 🤷

Co-authored-by: Tatiana Al-Chueyr <[email protected]>

cosmos/operators/base.py

@@ -8,6 +8,7 @@
 from airflow.models.baseoperator import BaseOperator
 from airflow.utils.context import Context
 from airflow.utils.operator_helpers import context_to_airflow_vars
+from airflow.utils.strings import to_boolean
 
 from cosmos.dbt.executable import get_system_dbt
 from cosmos.log import get_logger
@@ -61,7 +62,7 @@ class AbstractDbtBaseOperator(BaseOperator, metaclass=ABCMeta):
     :param dbt_cmd_global_flags: List of dbt global flags to be passed to the dbt command
     """
 
-    template_fields: Sequence[str] = ("env", "vars")
+    template_fields: Sequence[str] = ("env", "select", "exclude", "selector", "vars", "models")
     global_flags = (
         "project_dir",
         "select",
@@ -253,6 +254,26 @@ class DbtBuildMixin:
     base_cmd = ["build"]
     ui_color = "#8194E0"
 
+    template_fields: Sequence[str] = ("full_refresh",)
+
+    def __init__(self, full_refresh: bool | str = False, **kwargs: Any) -> None:
+        self.full_refresh = full_refresh
+        super().__init__(**kwargs)
+
+    def add_cmd_flags(self) -> list[str]:
+        flags = []
+
+        if isinstance(self.full_refresh, str):
+            # Handle template fields when render_template_as_native_obj=False
+            full_refresh = to_boolean(self.full_refresh)
+        else:
+            full_refresh = self.full_refresh
+
+        if full_refresh is True:
+            flags.append("--full-refresh")
+
+        return flags
+
 
 class DbtLSMixin:
     """
@@ -275,13 +296,20 @@ class DbtSeedMixin:
 
     template_fields: Sequence[str] = ("full_refresh",)
 
-    def __init__(self, full_refresh: bool = False, **kwargs: Any) -> None:
+    def __init__(self, full_refresh: bool | str = False, **kwargs: Any) -> None:
         self.full_refresh = full_refresh
         super().__init__(**kwargs)
 
     def add_cmd_flags(self) -> list[str]:
         flags = []
-        if self.full_refresh is True:
+
+        if isinstance(self.full_refresh, str):
+            # Handle template fields when render_template_as_native_obj=False
+            full_refresh = to_boolean(self.full_refresh)
+        else:
+            full_refresh = self.full_refresh
+
+        if full_refresh is True:
             flags.append("--full-refresh")
 
         return flags
@@ -307,13 +335,20 @@ class DbtRunMixin:
 
     template_fields: Sequence[str] = ("full_refresh",)
 
-    def __init__(self, full_refresh: bool = False, **kwargs: Any) -> None:
+    def __init__(self, full_refresh: bool | str = False, **kwargs: Any) -> None:
         self.full_refresh = full_refresh
         super().__init__(**kwargs)
 
     def add_cmd_flags(self) -> list[str]:
         flags = []
-        if self.full_refresh is True:
+
+        if isinstance(self.full_refresh, str):
+            # Handle template fields when render_template_as_native_obj=False
+            full_refresh = to_boolean(self.full_refresh)
+        else:
+            full_refresh = self.full_refresh
+
+        if full_refresh is True:
             flags.append("--full-refresh")
 
         return flags

cosmos/operators/docker.py

@@ -69,6 +69,8 @@ class DbtBuildDockerOperator(DbtBuildMixin, DbtDockerBaseOperator):
     Executes a dbt core build command.
     """
 
+    template_fields: Sequence[str] = DbtDockerBaseOperator.template_fields + DbtBuildMixin.template_fields  # type: ignore[operator]
+
 
 class DbtLSDockerOperator(DbtLSMixin, DbtDockerBaseOperator):
     """

cosmos/operators/kubernetes.py

@@ -102,6 +102,8 @@ class DbtBuildKubernetesOperator(DbtBuildMixin, DbtKubernetesBaseOperator):
     Executes a dbt core build command.
     """
 
+    template_fields: Sequence[str] = DbtKubernetesBaseOperator.template_fields + DbtBuildMixin.template_fields  # type: ignore[operator]
+
 
 class DbtLSKubernetesOperator(DbtLSMixin, DbtKubernetesBaseOperator):
     """

cosmos/operators/local.py

@@ -468,6 +468,8 @@ class DbtBuildLocalOperator(DbtBuildMixin, DbtLocalBaseOperator):
     Executes a dbt core build command.
     """
 
+    template_fields: Sequence[str] = DbtLocalBaseOperator.template_fields + DbtBuildMixin.template_fields  # type: ignore[operator]
+
 
 class DbtLSLocalOperator(DbtLSMixin, DbtLocalBaseOperator):
     """

cosmos/operators/virtualenv.py

@@ -104,7 +104,7 @@ def execute(self, context: Context) -> None:
         logger.info(output)
 
 
-class DbtBuildVirtualenvOperator(DbtVirtualenvBaseOperator, DbtBuildLocalOperator):
+class DbtBuildVirtualenvOperator(DbtVirtualenvBaseOperator, DbtBuildLocalOperator):  # type: ignore[misc]
     """
     Executes a dbt core build command within a Python Virtual Environment, that is created before running the dbt command
     and deleted just after.

docs/configuration/operator-args.rst

@@ -54,6 +54,7 @@ dbt-related
 - ``quiet``: run ``dbt`` in silent mode, only displaying its error logs.
 - ``vars``: (Deprecated since Cosmos 1.3 use ``ProjectConfig.dbt_vars`` instead) Supply variables to the project. This argument overrides variables defined in the ``dbt_project.yml``.
 - ``warn_error``: convert ``dbt`` warnings into errors.
+- ``full_refresh``: If True, then full refresh the node. This only applies to model and seed nodes.
 
 Airflow-related
 ...............
@@ -88,3 +89,33 @@ Sample usage
             "skip_exit_code": 1,
         }
     )
+
+
+Template fields
+---------------
+
+Some of the operator args are `template fields <https://airflow.apache.org/docs/apache-airflow/stable/howto/custom-operator.html#templating>`_ for your convenience.
+
+These template fields can be useful for hooking into Airflow `Params <https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/params.html>`_, or for more advanced customization with `XComs <https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/xcoms.html>`_.
+
+The following operator args support templating, and are accessible both through the  ``DbtDag`` and ``DbtTaskGroup`` constructors in addition to being accessible standalone:
+
+- ``env``
+- ``vars``
+- ``full_refresh`` (for the ``build``, ``seed``, and ``run`` operators since Cosmos 1.4.)
+
+.. note::
+    Using Jinja templating for ``env`` and ``vars`` may cause problems when using ``LoadMode.DBT_LS`` to render your DAG.
+
+The following template fields are only selectable when using the operators in a standalone context (starting in Cosmos 1.4):
+
+- ``select``
+- ``exclude``
+- ``selector``
+- ``models``
+
+Since Airflow resolves template fields during Airflow DAG execution and not DAG parsing,  the args above cannot be templated via ``DbtDag`` and ``DbtTaskGroup`` because both need to select dbt nodes during DAG parsing.
+
+Additionally, the SQL for compiled dbt models is stored in the template fields, which is viewable in the Airflow UI for each task run.
+This is provided for telemetry on task execution, and is not an operator arg.
+For more information about this, see the `Compiled SQL <compiled-sql.html>`_ docs.

tests/operators/test_base.py

@@ -51,8 +51,10 @@ def test_dbt_mixin_base_cmd(dbt_command, dbt_operator_class):
     assert [dbt_command] == dbt_operator_class.base_cmd
 
 
-@pytest.mark.parametrize("dbt_operator_class", [DbtSeedMixin, DbtRunMixin])
-@pytest.mark.parametrize("full_refresh, expected_flags", [(True, ["--full-refresh"]), (False, [])])
+@pytest.mark.parametrize("dbt_operator_class", [DbtSeedMixin, DbtRunMixin, DbtBuildMixin])
+@pytest.mark.parametrize(
+    "full_refresh, expected_flags", [("True", ["--full-refresh"]), (True, ["--full-refresh"]), (False, [])]
+)
 def test_dbt_mixin_add_cmd_flags_full_refresh(full_refresh, expected_flags, dbt_operator_class):
     dbt_mixin = dbt_operator_class(full_refresh=full_refresh)
     flags = dbt_mixin.add_cmd_flags()

tests/operators/test_local.py

@@ -567,6 +567,7 @@ def test_store_compiled_sql() -> None:
     "operator_class,kwargs,expected_call_kwargs",
     [
         (DbtSeedLocalOperator, {"full_refresh": True}, {"context": {}, "cmd_flags": ["--full-refresh"]}),
+        (DbtBuildLocalOperator, {"full_refresh": True}, {"context": {}, "cmd_flags": ["--full-refresh"]}),
         (DbtRunLocalOperator, {"full_refresh": True}, {"context": {}, "cmd_flags": ["--full-refresh"]}),
         (
             DbtTestLocalOperator,
@@ -650,8 +651,18 @@ def test_calculate_openlineage_events_completes_openlineage_errors(mock_processo
 @pytest.mark.parametrize(
     "operator_class,expected_template",
     [
-        (DbtSeedLocalOperator, ("env", "vars", "compiled_sql", "full_refresh")),
-        (DbtRunLocalOperator, ("env", "vars", "compiled_sql", "full_refresh")),
+        (
+            DbtSeedLocalOperator,
+            ("env", "select", "exclude", "selector", "vars", "models", "compiled_sql", "full_refresh"),
+        ),
+        (
+            DbtRunLocalOperator,
+            ("env", "select", "exclude", "selector", "vars", "models", "compiled_sql", "full_refresh"),
+        ),
+        (
+            DbtBuildLocalOperator,
+            ("env", "select", "exclude", "selector", "vars", "models", "compiled_sql", "full_refresh"),
+        ),
     ],
 )
 def test_dbt_base_operator_template_fields(operator_class, expected_template):