Add support for dbt selector arg for DAG parsing

cosmos/config.py

@@ -40,6 +40,7 @@ class RenderConfig:
     :param load_method: The parsing method for loading the dbt model. Defaults to AUTOMATIC
     :param select: A list of dbt select arguments (e.g. 'config.materialized:incremental')
     :param exclude: A list of dbt exclude arguments (e.g. 'tag:nightly')
+    :param selector: Name of a dbt YAML selector to use for parsing. Only supported when using ``load_method=LoadMode.DBT_LS``.
     :param dbt_deps: Configure to run dbt deps when using dbt ls for dag parsing
     :param node_converters: a dictionary mapping a ``DbtResourceType`` into a callable. Users can control how to render dbt nodes in Airflow. Only supported when using ``load_method=LoadMode.DBT_MANIFEST`` or ``LoadMode.DBT_LS``.
     :param dbt_executable_path: The path to the dbt executable for dag generation. Defaults to dbt if available on the path.
@@ -52,6 +53,7 @@ class RenderConfig:
     load_method: LoadMode = LoadMode.AUTOMATIC
     select: list[str] = field(default_factory=list)
     exclude: list[str] = field(default_factory=list)
+    selector: str | None = None
     dbt_deps: bool = True
     node_converters: dict[DbtResourceType, Callable[..., Any]] | None = None
     dbt_executable_path: str | Path = get_system_dbt()

cosmos/dbt/graph.py

@@ -195,6 +195,9 @@ def run_dbt_ls(
         if self.project.dbt_vars:
             ls_command.extend(["--vars", yaml.dump(self.project.dbt_vars)])
 
+        if self.render_config.selector:
+            ls_command.extend(["--selector", self.render_config.selector])
+
         ls_command.extend(self.local_flags)
 
         stdout = run_command(ls_command, tmp_dir, env_vars)
@@ -291,6 +294,11 @@ def load_via_custom_parser(self) -> None:
         """
         logger.info("Trying to parse the dbt project `%s` using a custom Cosmos method...", self.project.project_name)
 
+        if self.render_config.selector:
+            raise CosmosLoadDbtException(
+                "RenderConfig.selector is not yet supported when loading dbt projects using the LoadMode.CUSTOM parser."
+            )
+
         if not self.render_config.project_path or not self.execution_config.project_path:
             raise CosmosLoadDbtException(
                 "Unable to load dbt project without RenderConfig.dbt_project_path and ExecutionConfig.dbt_project_path"
@@ -349,6 +357,11 @@ def load_from_dbt_manifest(self) -> None:
         """
         logger.info("Trying to parse the dbt project `%s` using a dbt manifest...", self.project.project_name)
 
+        if self.render_config.selector:
+            raise CosmosLoadDbtException(
+                "RenderConfig.selector is not yet supported when loading dbt projects using the LoadMode.DBT_MANIFEST parser."
+            )
+
         if not self.project.is_manifest_available():
             raise CosmosLoadDbtException(f"Unable to load manifest using {self.project.manifest_path}")
 

docs/configuration/render-config.rst

@@ -11,6 +11,7 @@ The ``RenderConfig`` class takes the following arguments:
 - ``test_behavior``: how to run tests. Defaults to running a model's tests immediately after the model is run. For more information, see the `Testing Behavior <testing-behavior.html>`_ section.
 - ``load_method``: how to load your dbt project. See `Parsing Methods <parsing-methods.html>`_ for more information.
 - ``select`` and ``exclude``: which models to include or exclude from your DAGs. See `Selecting & Excluding <selecting-excluding.html>`_ for more information.
+- ``selector``: (new in v1.3) name of a dbt YAML selector to use for DAG parsing. Only supported when using ``load_method=LoadMode.DBT_LS``. See `Selecting & Excluding <selecting-excluding.html>`_ for more information.
 - ``dbt_deps``: A Boolean to run dbt deps when using dbt ls for dag parsing. Default True
 - ``node_converters``: a dictionary mapping a ``DbtResourceType`` into a callable. Users can control how to render dbt nodes in Airflow. Only supported when using ``load_method=LoadMode.DBT_MANIFEST`` or ``LoadMode.DBT_LS``. Find more information below.
 - ``dbt_executable_path``: The path to the dbt executable for dag generation. Defaults to dbt if available on the path.

docs/configuration/selecting-excluding.rst

@@ -3,7 +3,13 @@
 Selecting & Excluding
 =======================
 
-Cosmos allows you to filter to a subset of your dbt project in each ``DbtDag`` / ``DbtTaskGroup`` using the ``select`` and ``exclude`` parameters in the ``RenderConfig`` class.
+Cosmos allows you to filter to a subset of your dbt project in each ``DbtDag`` / ``DbtTaskGroup`` using the ``select `` and ``exclude`` parameters in the ``RenderConfig`` class.
+
+ Since Cosmos 1.3, the ``selector`` parameter is also available in ``RenderConfig`` when using the ``LoadMode.DBT_LS`` to parse the dbt project into Airflow.
+
+
+Using ``select`` and ``exclude``
+--------------------------------
 
 The ``select`` and ``exclude`` parameters are lists, with values like the following:
 
@@ -84,3 +90,24 @@ Examples:
             exclude=["node_name+"],  # node_name and its children
         )
     )
+
+Using ``selector``
+--------------------------------
+.. note::
+    Only currently supported using the ``dbt_ls`` parsing method since Cosmos 1.3 where the selector is passed directly to the dbt CLI command. \
+    If  ``select`` and/or ``exclude`` are used with ``selector``, dbt will ignore the ``select`` and ``exclude`` parameters.
+
+The ``selector`` parameter is a string that references a `dbt YAML selector <https://docs.getdbt.com/reference/node-selection/yaml-selectors>`_ already defined in a dbt project.
+
+Examples:
+
+.. code-block:: python
+
+    from cosmos import DbtDag, RenderConfig, LoadMode
+
+    jaffle_shop = DbtDag(
+        render_config=RenderConfig(
+            selector="my_selector",  # this selector must be defined in your dbt project
+            load_method=LoadMode.DBT_LS,
+        )
+    )