Add default source nodes rendering

cosmos/airflow/graph.py

@@ -142,6 +142,7 @@ def create_task_metadata(
         DbtResourceType.SNAPSHOT: "DbtSnapshot",
         DbtResourceType.SEED: "DbtSeed",
         DbtResourceType.TEST: "DbtTest",
+        DbtResourceType.SOURCE: "DbtSource",
     }
     args = {**args, **{"models": node.resource_name}}
 
@@ -151,6 +152,15 @@ def create_task_metadata(
             task_id = f"{node.name}_run"
             if use_task_group is True:
                 task_id = "run"
+        elif node.resource_type == DbtResourceType.SOURCE:
+            task_id = f"{node.name}_source"
+            args["select"] = f"source:{node.resource_name}"
+            args["models"] = None
+            if use_task_group is True:
+                task_id = node.resource_type.value
+            if node.has_freshness is False:
+                # render sources without freshness as empty operators
+                return TaskMetadata(id=task_id, operator_class="airflow.operators.empty.EmptyOperator")
         else:
             task_id = f"{node.name}_{node.resource_type.value}"
             if use_task_group is True:

cosmos/core/airflow.py

@@ -29,7 +29,7 @@ def get_airflow_task(task: Task, dag: DAG, task_group: "TaskGroup | None" = None
         task_id=task.id,
         dag=dag,
         task_group=task_group,
-        extra_context=task.extra_context,
+        **({} if class_name == "EmptyOperator" else {"extra_context": task.extra_context}),
         **task.arguments,
     )
 

cosmos/dbt/graph.py

@@ -48,7 +48,7 @@ class CosmosLoadDbtException(Exception):
 @dataclass
 class DbtNode:
     """
-    Metadata related to a dbt node (e.g. model, seed, snapshot).
+    Metadata related to a dbt node (e.g. model, seed, snapshot, source).
     """
 
     unique_id: str
@@ -57,6 +57,7 @@ class DbtNode:
     file_path: Path
     tags: list[str] = field(default_factory=lambda: [])
     config: dict[str, Any] = field(default_factory=lambda: {})
+    has_freshness: bool = False
     has_test: bool = False
 
     @property
@@ -95,6 +96,30 @@ def context_dict(self) -> dict[str, Any]:
         }
 
 
+def is_freshness_effective(freshness: dict[str, Any]) -> bool:
+    """Function to find if a source has null freshness. Scenarios where freshness
+    looks like:
+    "freshness": {
+                "warn_after": {
+                    "count": null,
+                    "period": null
+                },
+                "error_after": {
+                    "count": null,
+                    "period": null
+                },
+                "filter": null
+            }
+    should be considered as null, this function ensures that."""
+    for _, value in freshness.items():
+        if isinstance(value, dict):
+            if any(subvalue is not None for subvalue in value.values()):
+                return True
+        elif value is not None:
+            return True
+    return False
+
+
 def run_command(command: list[str], tmp_dir: Path, env_vars: dict[str, str]) -> str:
     """Run a command in a subprocess, returning the stdout."""
     logger.info("Running command: `%s`", " ".join(command))
@@ -138,6 +163,11 @@ def parse_dbt_ls_output(project_path: Path | None, ls_stdout: str) -> dict[str,
                 file_path=project_path / node_dict["original_file_path"],
                 tags=node_dict.get("tags", []),
                 config=node_dict.get("config", {}),
+                has_freshness=(
+                    is_freshness_effective(node_dict.get("freshness", False))
+                    if node_dict["resource_type"] == DbtResourceType.SOURCE
+                    else False
+                ),
             )
             nodes[node.unique_id] = node
             logger.debug("Parsed dbt resource `%s` of type `%s`", node.unique_id, node.resource_type)
@@ -339,7 +369,14 @@ def run_dbt_ls(
         self, dbt_cmd: str, project_path: Path, tmp_dir: Path, env_vars: dict[str, str]
     ) -> dict[str, DbtNode]:
         """Runs dbt ls command and returns the parsed nodes."""
-        ls_command = [dbt_cmd, "ls", "--output", "json"]
+        ls_command = [
+            dbt_cmd,
+            "ls",
+            "--output",
+            "json",
+            "--output-keys",
+            "name alias unique_id resource_type depends_on original_file_path tags config freshness",
+        ]
 
         ls_args = self.dbt_ls_args
         ls_command.extend(self.local_flags)
@@ -605,6 +642,11 @@ def load_from_dbt_manifest(self) -> None:
                     file_path=self.execution_config.project_path / Path(node_dict["original_file_path"]),
                     tags=node_dict["tags"],
                     config=node_dict["config"],
+                    has_freshness=(
+                        node_dict["freshness"] is not None
+                        if node_dict["resource_type"] == "source" and "freshness" in node_dict
+                        else False
+                    ),
                 )
 
                 nodes[node.unique_id] = node

cosmos/operators/local.py

@@ -816,3 +816,12 @@ def __init__(self, **kwargs: str) -> None:
         raise DeprecationWarning(
             "The DbtDepsOperator has been deprecated. " "Please use the `install_deps` flag in dbt_args instead."
         )
+
+
+class DbtSourceLocalOperator(DbtLocalBaseOperator):
+    """
+    Executes a dbt source freshness command.
+    """
+
+    ui_color = "#34CCEB"
+    base_cmd = ["source", "freshness"]

cosmos/operators/virtualenv.py

@@ -18,6 +18,7 @@
     DbtRunOperationLocalOperator,
     DbtSeedLocalOperator,
     DbtSnapshotLocalOperator,
+    DbtSourceLocalOperator,
     DbtTestLocalOperator,
 )
 
@@ -162,3 +163,10 @@ class DbtDocsVirtualenvOperator(DbtVirtualenvBaseOperator, DbtDocsLocalOperator)
     Executes `dbt docs generate` command within a Python Virtual Environment, that is created before running the dbt
     command and deleted just after.
     """
+
+
+class DbtSourceVirtualenvOperator(DbtVirtualenvBaseOperator, DbtSourceLocalOperator):
+    """
+    Executes `dbt source freshness` command within a Python Virtual Environment, that is created before running the dbt
+    command and deleted just after.
+    """

dev/dags/dbt/jaffle_shop/models/staging/sources.yml

@@ -0,0 +1,26 @@
+
+version: 2
+
+sources:
+  - name: postgres_db
+    database: "{{ env_var('POSTGRES_DB') }}"
+    schema: "{{ env_var('POSTGRES_SCHEMA') }}"
+    tables:
+      - name: raw_customers
+        columns:
+          - name: id
+            tests:
+              - unique
+              - not_null
+      - name: raw_payments
+      - name: raw_orders
+        columns:
+          - name: id
+            tests:
+              - unique
+              - not_null
+        freshness:
+          warn_after:
+            count: 3650
+            period: day
+        loaded_at_field: CAST(order_date AS TIMESTAMP)

dev/dags/dbt/jaffle_shop/models/staging/stg_customers.sql

@@ -4,7 +4,7 @@ with source as (
     Normally we would select from the table here, but we are using seeds to load
     our data in this project
     #}
-    select * from {{ ref('raw_customers') }}
+    select * from {{ source('postgres_db', 'raw_customers') }}
 
 ),
 

dev/dags/dbt/jaffle_shop/models/staging/stg_orders.sql

@@ -4,7 +4,7 @@ with source as (
     Normally we would select from the table here, but we are using seeds to load
     our data in this project
     #}
-    select * from {{ ref('raw_orders') }}
+    select * from {{ source('postgres_db', 'raw_orders') }}
 
 ),
 

dev/dags/dbt/jaffle_shop/models/staging/stg_payments.sql

@@ -4,7 +4,7 @@ with source as (
     Normally we would select from the table here, but we are using seeds to load
     our data in this project
     #}
-    select * from {{ ref('raw_payments') }}
+    select * from {{ source('postgres_db', 'raw_payments') }}
 
 ),
 

tests/dbt/test_graph.py

@@ -453,13 +453,15 @@ def test_load_via_dbt_ls_with_exclude(postgres_profile_config):
     dbt_graph.load_via_dbt_ls()
     assert dbt_graph.nodes == dbt_graph.filtered_nodes
     # This test is dependent upon dbt >= 1.5.4
-    assert len(dbt_graph.nodes) == 7
+    assert len(dbt_graph.nodes) == 9
     expected_keys = [
         "model.jaffle_shop.customers",
         "model.jaffle_shop.stg_customers",
         "seed.jaffle_shop.raw_customers",
         "test.jaffle_shop.not_null_customers_customer_id.5c9bf9911d",
         "test.jaffle_shop.not_null_stg_customers_customer_id.e2cfb1f9aa",
+        "test.jaffle_shop.source_not_null_postgres_db_raw_customers_id.de3e9fff76",
+        "test.jaffle_shop.source_unique_postgres_db_raw_customers_id.6e5ad1d707",
         "test.jaffle_shop.unique_customers_customer_id.c5af1ff4b1",
         "test.jaffle_shop.unique_stg_customers_customer_id.c7614daada",
     ]
@@ -492,7 +494,7 @@ def test_load_via_dbt_ls_without_exclude(project_name, postgres_profile_config):
     dbt_graph.load_via_dbt_ls()
 
     assert dbt_graph.nodes == dbt_graph.filtered_nodes
-    assert len(dbt_graph.nodes) == 28
+    assert len(dbt_graph.nodes) == 35
 
 
 def test_load_via_custom_without_project_path():
@@ -1276,11 +1278,11 @@ def test_load_via_dbt_ls_with_selector_arg(tmp_dbt_project_dir, postgres_profile
     dbt_graph.load_via_dbt_ls()
 
     filtered_nodes = dbt_graph.filtered_nodes.keys()
-    assert len(filtered_nodes) == 4
+    assert len(filtered_nodes) == 6
     assert "model.jaffle_shop.stg_customers" in filtered_nodes
-    assert "seed.jaffle_shop.raw_customers" in filtered_nodes
-    # Two tests should be filtered
-    assert sum(node.startswith("test.jaffle_shop") for node in filtered_nodes) == 2
+    assert "source.jaffle_shop.postgres_db.raw_customers" in filtered_nodes
+    # Four tests should be filtered
+    assert sum(node.startswith("test.jaffle_shop") for node in filtered_nodes) == 4
 
 
 @pytest.mark.parametrize(

tests/operators/test_local.py

@@ -444,7 +444,7 @@ def test_run_operator_dataset_inlets_and_outlets(caplog):
 
     run_test_dag(dag)
 
-    assert run_operator.inlets == []
+    assert run_operator.inlets == [Dataset(uri="postgres://0.0.0.0:5432/postgres.public.raw_customers", extra=None)]
     assert run_operator.outlets == [Dataset(uri="postgres://0.0.0.0:5432/postgres.public.stg_customers", extra=None)]
     assert test_operator.inlets == [Dataset(uri="postgres://0.0.0.0:5432/postgres.public.stg_customers", extra=None)]
     assert test_operator.outlets == []