Merge branch 'main' into snowflake-encrypted-pk-path

cosmos/airflow/graph.py

@@ -56,7 +56,7 @@ def create_test_task_metadata(
  execution_mode: ExecutionMode,
  task_args: dict[str, Any],
  on_warning_callback: Callable[..., Any] | None = None,
- model_name: str | None = None,
+ node: DbtNode | None = None,
 ) -> TaskMetadata:
  """
  Create the metadata that will be used to instantiate the Airflow Task that will be used to run the Dbt test node.
@@ -66,13 +66,18 @@ def create_test_task_metadata(
  :param task_args: Arguments to be used to instantiate an Airflow Task
  :param on_warning_callback: A callback function called on warnings with additional Context variables “test_names”
  and “test_results” of type List.
- :param model_name: If the test relates to a specific model, the name of the model it relates to
+ :param node: If the test relates to a specific node, the node reference
  :returns: The metadata necessary to instantiate the source dbt node as an Airflow task.
  """
  task_args = dict(task_args)
  task_args["on_warning_callback"] = on_warning_callback
- if model_name is not None:
- task_args["models"] = model_name
+ if node is not None:
+ if node.resource_type == DbtResourceType.MODEL:
+ task_args["models"] = node.name
+ elif node.resource_type == DbtResourceType.SOURCE:
+ task_args["select"] = f"source:{node.unique_id[len('source.'):]}"
+ else: # tested with node.resource_type == DbtResourceType.SEED or DbtResourceType.SNAPSHOT
+ task_args["select"] = node.name
  return TaskMetadata(
  id=test_task_name,
  operator_class=calculate_operator_class(
@@ -112,6 +117,8 @@ def create_task_metadata(
  task_id = "run"
  else:
  task_id = f"{node.name}_{node.resource_type.value}"
+ if use_task_group is True:
+ task_id = node.resource_type.value
 
  task_metadata = TaskMetadata(
  id=task_id,
@@ -163,7 +170,7 @@ def generate_task_or_group(
  "test",
  execution_mode,
  task_args=task_args,
- model_name=node.name,
+ node=node,
  on_warning_callback=on_warning_callback,
  )
  test_task = create_airflow_task(test_meta, dag, task_group=model_task_group)

cosmos/constants.py

@@ -70,7 +70,7 @@ def _missing_value_(cls, value): # type: ignore
 
 DEFAULT_DBT_RESOURCES = DbtResourceType.__members__.values()
 
-
-TESTABLE_DBT_RESOURCES = {
- DbtResourceType.MODEL
-} # TODO: extend with DbtResourceType.SOURCE, DbtResourceType.SNAPSHOT, DbtResourceType.SEED)
+# dbt test runs tests defined on models, sources, snapshots, and seeds.
+# It expects that you have already created those resources through the appropriate commands.
+# https://docs.getdbt.com/reference/commands/test
+TESTABLE_DBT_RESOURCES = {DbtResourceType.MODEL, DbtResourceType.SOURCE, DbtResourceType.SNAPSHOT, DbtResourceType.SEED}

cosmos/profiles/__init__.py

@@ -21,6 +21,7 @@
 from .trino.certificate import TrinoCertificateProfileMapping
 from .trino.jwt import TrinoJWTProfileMapping
 from .trino.ldap import TrinoLDAPProfileMapping
+from .vertica.user_pass import VerticaUserPasswordProfileMapping
 
 profile_mappings: list[Type[BaseProfileMapping]] = [
  AthenaAccessKeyProfileMapping,
@@ -38,6 +39,7 @@
  TrinoLDAPProfileMapping,
  TrinoCertificateProfileMapping,
  TrinoJWTProfileMapping,
+ VerticaUserPasswordProfileMapping,
 ]
 
 
@@ -74,4 +76,5 @@ def get_automatic_profile_mapping(
  "TrinoLDAPProfileMapping",
  "TrinoCertificateProfileMapping",
  "TrinoJWTProfileMapping",
+ "VerticaUserPasswordProfileMapping",
 ]

cosmos/profiles/vertica/__init__.py

@@ -0,0 +1,5 @@
+"Vertica Airflow connection -> dbt profile mappings"
+
+from .user_pass import VerticaUserPasswordProfileMapping
+
+__all__ = ["VerticaUserPasswordProfileMapping"]

cosmos/profiles/vertica/user_pass.py

@@ -0,0 +1,76 @@
+"Maps Airflow Vertica connections using user + password authentication to dbt profiles."
+from __future__ import annotations
+
+from typing import Any
+
+from ..base import BaseProfileMapping
+
+
+class VerticaUserPasswordProfileMapping(BaseProfileMapping):
+ """
+ Maps Airflow Vertica connections using user + password authentication to dbt profiles.
+ https://docs.getdbt.com/reference/warehouse-setups/vertica-setup
+ https://airflow.apache.org/docs/apache-airflow-providers-vertica/stable/connections/vertica.html
+ """
+
+ airflow_connection_type: str = "vertica"
+ dbt_profile_type: str = "vertica"
+
+ required_fields = [
+ "host",
+ "user",
+ "password",
+ "database",
+ "schema",
+ ]
+ secret_fields = [
+ "password",
+ ]
+ airflow_param_mapping = {
+ "host": "host",
+ "user": "login",
+ "password": "password",
+ "port": "port",
+ "schema": "schema",
+ "database": "extra.database",
+ "autocommit": "extra.autocommit",
+ "backup_server_node": "extra.backup_server_node",
+ "binary_transfer": "extra.binary_transfer",
+ "connection_load_balance": "extra.connection_load_balance",
+ "connection_timeout": "extra.connection_timeout",
+ "disable_copy_local": "extra.disable_copy_local",
+ "kerberos_host_name": "extra.kerberos_host_name",
+ "kerberos_service_name": "extra.kerberos_service_name",
+ "log_level": "extra.log_level",
+ "log_path": "extra.log_path",
+ "oauth_access_token": "extra.oauth_access_token",
+ "request_complex_types": "extra.request_complex_types",
+ "session_label": "extra.session_label",
+ "ssl": "extra.ssl",
+ "unicode_error": "extra.unicode_error",
+ "use_prepared_statements": "extra.use_prepared_statements",
+ "workload": "extra.workload",
+ }
+
+ @property
+ def profile(self) -> dict[str, Any | None]:
+ "Gets profile. The password is stored in an environment variable."
+ profile = {
+ "port": 5433,
+ **self.mapped_params,
+ **self.profile_args,
+ # password should always get set as env var
+ "password": self.get_env_var_format("password"),
+ }
+
+ return self.filter_null(profile)
+
+ @property
+ def mock_profile(self) -> dict[str, Any | None]:
+ "Gets mock profile. Defaults port to 5433."
+ parent_mock = super().mock_profile
+
+ return {
+ "port": 5433,
+ **parent_mock,
+ }

pyproject.toml

@@ -55,6 +55,7 @@ dbt-all = [
  "dbt-redshift",
  "dbt-snowflake",
  "dbt-spark",
+ "dbt-vertica",
 ]
 dbt-athena = [
  "dbt-athena-community",
@@ -80,6 +81,9 @@ dbt-snowflake = [
 dbt-spark = [
  "dbt-spark",
 ]
+dbt-vertica = [
+ "dbt-vertica<=1.5.4",
+]
 openlineage = [
  "openlineage-integration-common",
  "openlineage-airflow",
@@ -165,18 +169,18 @@ test = 'pytest -vv --durations=0 . -m "not integration" --ignore=tests/test_exam
 test-cov = """pytest -vv --cov=cosmos --cov-report=term-missing --cov-report=xml --durations=0 -m "not integration" --ignore=tests/test_example_dags.py --ignore=tests/test_example_dags_no_connections.py"""
 # we install using the following workaround to overcome installation conflicts, such as:
 # apache-airflow 2.3.0 and dbt-core [0.13.0 - 1.5.2] and jinja2>=3.0.0 because these package versions have conflicting dependencies
-test-integration-setup = """pip uninstall -y dbt-core dbt-databricks dbt-sqlite dbt-postgres dbt-sqlite; \
+test-integration-setup = """pip uninstall dbt-postgres dbt-databricks dbt-vertica; \
 rm -rf airflow.*; \
 airflow db init; \
-pip install 'dbt-core' 'dbt-databricks' 'dbt-postgres' 'openlineage-airflow'"""
-test-integration = """pytest -vv \
+pip install 'dbt-core' 'dbt-databricks' 'dbt-postgres' 'dbt-vertica' 'openlineage-airflow'"""
+test-integration = """rm -rf dbt/jaffle_shop/dbt_packages;
+pytest -vv \
 --cov=cosmos \
 --cov-report=term-missing \
 --cov-report=xml \
 --durations=0 \
 -m integration \
--k 'not (sqlite or example_cosmos_sources or example_cosmos_python_models or example_virtualenv or cosmos_manifest_example)'
-"""
+-k 'not (sqlite or example_cosmos_sources or example_cosmos_python_models or example_virtualenv or cosmos_manifest_example)'"""
 test-integration-expensive = """pytest -vv \
 --cov=cosmos \
 --cov-report=term-missing \

tests/airflow/test_graph.py

@@ -5,6 +5,7 @@
 import pytest
 from airflow import __version__ as airflow_version
 from airflow.models import DAG
+from airflow.utils.task_group import TaskGroup
 from packaging import version
 
 from cosmos.airflow.graph import (
@@ -13,6 +14,7 @@
  calculate_operator_class,
  create_task_metadata,
  create_test_task_metadata,
+ generate_task_or_group,
 )
 from cosmos.config import ProfileConfig
 from cosmos.constants import DbtResourceType, ExecutionMode, TestBehavior
@@ -101,6 +103,50 @@ def test_build_airflow_graph_with_after_each():
  assert dag.leaves[0].task_id == "child_run"
 
 
+@pytest.mark.parametrize(
+ "node_type,task_suffix",
+ [(DbtResourceType.MODEL, "run"), (DbtResourceType.SEED, "seed"), (DbtResourceType.SNAPSHOT, "snapshot")],
+)
+def test_create_task_group_for_after_each_supported_nodes(node_type, task_suffix):
+ """
+ dbt test runs tests defined on models, sources, snapshots, and seeds.
+ It expects that you have already created those resources through the appropriate commands.
+ https://docs.getdbt.com/reference/commands/test
+ """
+ with DAG("test-task-group-after-each", start_date=datetime(2022, 1, 1)) as dag:
+ node = DbtNode(
+ name="dbt_node",
+ unique_id="dbt_node",
+ resource_type=node_type,
+ file_path=SAMPLE_PROJ_PATH / "gen2/models/parent.sql",
+ tags=["has_child"],
+ config={"materialized": "view"},
+ depends_on=[],
+ has_test=True,
+ )
+ output = generate_task_or_group(
+ dag=dag,
+ task_group=None,
+ node=node,
+ execution_mode=ExecutionMode.LOCAL,
+ task_args={
+ "project_dir": SAMPLE_PROJ_PATH,
+ "profile_config": ProfileConfig(
+ profile_name="default",
+ target_name="default",
+ profile_mapping=PostgresUserPasswordProfileMapping(
+ conn_id="fake_conn",
+ profile_args={"schema": "public"},
+ ),
+ ),
+ },
+ test_behavior=TestBehavior.AFTER_EACH,
+ on_warning_callback=None,
+ )
+ assert isinstance(output, TaskGroup)
+ assert list(output.children.keys()) == [f"dbt_node.{task_suffix}", "dbt_node.test"]
+
+
 @pytest.mark.skipif(
  version.parse(airflow_version) < version.parse("2.4"),
  reason="Airflow DAG did not have task_group_dict until the 2.4 release",
@@ -259,7 +305,12 @@ def test_create_task_metadata_seed(caplog, use_task_group):
  args={},
  use_task_group=use_task_group,
  )
- assert metadata.id == "my_seed_seed"
+
+ if not use_task_group:
+ assert metadata.id == "my_seed_seed"
+ else:
+ assert metadata.id == "seed"
+
  assert metadata.operator_class == "cosmos.operators.docker.DbtSeedDockerOperator"
  assert metadata.arguments == {"models": "my_seed"}
 
@@ -280,14 +331,32 @@ def test_create_task_metadata_snapshot(caplog):
  assert metadata.arguments == {"models": "my_snapshot"}
 
 
-def test_create_test_task_metadata():
+@pytest.mark.parametrize(
+ "node_type,node_unique_id,selector_key,selector_value",
+ [
+ (DbtResourceType.MODEL, "node_name", "models", "node_name"),
+ (DbtResourceType.SEED, "node_name", "select", "node_name"),
+ (DbtResourceType.SOURCE, "source.node_name", "select", "source:node_name"),
+ (DbtResourceType.SNAPSHOT, "node_name", "select", "node_name"),
+ ],
+)
+def test_create_test_task_metadata(node_type, node_unique_id, selector_key, selector_value):
+ sample_node = DbtNode(
+ name="node_name",
+ unique_id=node_unique_id,
+ resource_type=node_type,
+ depends_on=[],
+ file_path="",
+ tags=[],
+ config={},
+ )
  metadata = create_test_task_metadata(
  test_task_name="test_no_nulls",
  execution_mode=ExecutionMode.LOCAL,
  task_args={"task_arg": "value"},
  on_warning_callback=True,
- model_name="my_model",
+ node=sample_node,
  )
  assert metadata.id == "test_no_nulls"
  assert metadata.operator_class == "cosmos.operators.local.DbtTestLocalOperator"
- assert metadata.arguments == {"task_arg": "value", "on_warning_callback": True, "models": "my_model"}
+ assert metadata.arguments == {"task_arg": "value", "on_warning_callback": True, selector_key: selector_value}