From 6e9c63adfaf86a5c34a786644e1fbcfd36cd2adc Mon Sep 17 00:00:00 2001 From: Ryan Hatter <25823361+RNHTTR@users.noreply.github.com> Date: Fri, 5 Jan 2024 06:20:05 -0500 Subject: [PATCH 1/9] Update examples to use the latest astro-runtime (10.0.0) (#777) --- dev/Dockerfile | 2 +- docs/getting_started/astro.rst | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/dev/Dockerfile b/dev/Dockerfile index 90c49ed6c..b929be8b1 100644 --- a/dev/Dockerfile +++ b/dev/Dockerfile @@ -1,4 +1,4 @@ -FROM quay.io/astronomer/astro-runtime:7.3.0-base +FROM quay.io/astronomer/astro-runtime:10.0.0-base USER root diff --git a/docs/getting_started/astro.rst b/docs/getting_started/astro.rst index c0bedc7e6..8aaa194e5 100644 --- a/docs/getting_started/astro.rst +++ b/docs/getting_started/astro.rst @@ -20,7 +20,7 @@ Create a virtual environment in your ``Dockerfile`` using the sample below. Be s .. code-block:: docker - FROM quay.io/astronomer/astro-runtime:8.8.0 + FROM quay.io/astronomer/astro-runtime:10.0.0 # install dbt into a virtual environment RUN python -m venv dbt_venv && source dbt_venv/bin/activate && \ From bf5182c74d9852df43063fb412b3ac5c1a0e93d2 Mon Sep 17 00:00:00 2001 From: yuriy Date: Thu, 11 Jan 2024 01:13:01 +0100 Subject: [PATCH 2/9] move from config to base_profile --- cosmos/config.py | 49 +------------ cosmos/profiles/__init__.py | 3 +- cosmos/profiles/base.py | 82 ++++++++++++++++++---- dev/dags/basic_cosmos_dag.py | 4 +- docs/templates/index.rst.jinja2 | 36 +--------- tests/profiles/test_base_profile.py | 84 ++++++++++++++++++----- tests/test_config.py | 103 ---------------------------- 7 files changed, 141 insertions(+), 220 deletions(-) diff --git a/cosmos/config.py b/cosmos/config.py index 3a7f5b315..46e3f1915 100644 --- a/cosmos/config.py +++ b/cosmos/config.py @@ -212,8 +212,6 @@ class ProfileConfig: :param target_name: The name of the dbt target to use. :param profiles_yml_filepath: The path to a profiles.yml file to use. :param profile_mapping: A mapping of Airflow connections to dbt profiles. - :param dbt_config_vars: Dictionary of dbt configs for the project with profile mapping. Not active if you use profiles_yml_filepath. This argument overrides configs defined in your profiles.yml - file. The dictionary is dumped to a yaml string. Details https://docs.getdbt.com/docs/core/connect-data-platform/profiles.yml """ # should always be set to be explicit @@ -225,11 +223,9 @@ class ProfileConfig: # should be set if using cosmos to map Airflow connections to dbt profiles profile_mapping: BaseProfileMapping | None = None - dbt_config_vars: dict[str, Any] | None = None def __post_init__(self) -> None: self.validate_profile() - self.validate_dbt_config_vars() def validate_profile(self) -> None: "Validates that we have enough information to render a profile." @@ -240,46 +236,6 @@ def validate_profile(self) -> None: "Both profiles_yml_filepath and profile_mapping are defined and are mutually exclusive. Ensure only one of these is defined." ) - def validate_dbt_config_vars(self) -> None: - "Validates config vars for profile." - - vars_checks: dict[str, dict[str, Any]] = { - "send_anonymous_usage_stats": {"var_type": bool}, - "partial_parse": {"var_type": bool}, - "use_experimental_parser": {"var_type": bool}, - "static_parser": {"var_type": bool}, - "printer_width": {"var_type": int}, - "write_json": {"var_type": bool}, - "warn_error": {"var_type": str}, - "warn_error_options": {"var_type": dict, "accepted_values": {"include", "exclude"}}, - "log_format": {"var_type": str, "accepted_values": {"text", "json", "default"}}, - "debug": {"var_type": bool}, - "version_check": {"var_type": bool}, - } - - if self.dbt_config_vars: - if self.profiles_yml_filepath: - raise CosmosValueError( - "Both profiles_yml_filepath and dbt_config_vars are defined and are mutually exclusive. Ensure only one of these is defined." - ) - - for var_key, var_value in self.dbt_config_vars.items(): - if var_key not in list(vars_checks): - raise CosmosValueError(f"dbt config var {var_key}: {var_value} is not supported") - - vars_check = vars_checks.get(var_key, {}) - - var_type = vars_check.get("var_type", Any) - if not isinstance(var_value, var_type): - raise CosmosValueError(f"dbt config var {var_key}: {var_value} must be a {var_type}") - - accepted_values = vars_check.get("accepted_values") - if accepted_values: - if var_value not in accepted_values: - raise CosmosValueError( - f"dbt config var {var_key}: {var_value} must be one of {accepted_values}" - ) - def validate_profiles_yml(self) -> None: "Validates a user-supplied profiles.yml is present" if self.profiles_yml_filepath and not Path(self.profiles_yml_filepath).exists(): @@ -296,10 +252,7 @@ def ensure_profile( elif self.profile_mapping: profile_contents = self.profile_mapping.get_profile_file_contents( - profile_name=self.profile_name, - target_name=self.target_name, - use_mock_values=use_mock_values, - dbt_config_vars=self.dbt_config_vars, + profile_name=self.profile_name, target_name=self.target_name, use_mock_values=use_mock_values ) if use_mock_values: diff --git a/cosmos/profiles/__init__.py b/cosmos/profiles/__init__.py index 1f39a91a0..e7eae57a1 100644 --- a/cosmos/profiles/__init__.py +++ b/cosmos/profiles/__init__.py @@ -6,7 +6,7 @@ from .athena import AthenaAccessKeyProfileMapping -from .base import BaseProfileMapping +from .base import BaseProfileMapping, DbtConfigVars from .bigquery.service_account_file import GoogleCloudServiceAccountFileProfileMapping from .bigquery.service_account_keyfile_dict import GoogleCloudServiceAccountDictProfileMapping from .bigquery.oauth import GoogleCloudOauthProfileMapping @@ -81,4 +81,5 @@ def get_automatic_profile_mapping( "TrinoCertificateProfileMapping", "TrinoJWTProfileMapping", "VerticaUserPasswordProfileMapping", + "DbtConfigVars", ] diff --git a/cosmos/profiles/base.py b/cosmos/profiles/base.py index f10ad596b..c1b8f93ac 100644 --- a/cosmos/profiles/base.py +++ b/cosmos/profiles/base.py @@ -5,7 +5,8 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import Any +from typing import Any, Optional +import dataclasses from typing import TYPE_CHECKING import yaml @@ -24,6 +25,65 @@ logger = get_logger(__name__) +@dataclasses.dataclass +class DbtConfigVars: + send_anonymous_usage_stats: Optional[bool] = False + partial_parse: Optional[bool] = None + use_experimental_parser: Optional[bool] = None + static_parser: Optional[bool] = None + printer_width: Optional[int] = None + write_json: Optional[bool] = None + warn_error: Optional[str] = None + warn_error_options: Optional[dict[str, Any]] = None + log_format: Optional[str] = None + debug: Optional[bool] = None + version_check: Optional[bool] = None + + def _validate_data(self) -> None: + checks: dict[str, dict[str, Any]] = { + "send_anonymous_usage_stats": {"var_type": bool}, + "partial_parse": {"var_type": bool}, + "use_experimental_parser": {"var_type": bool}, + "static_parser": {"var_type": bool}, + "printer_width": {"var_type": int}, + "write_json": {"var_type": bool}, + "warn_error": {"var_type": str}, + "warn_error_options": {"var_type": dict, "accepted_values": {"include", "exclude"}}, + "log_format": {"var_type": str, "accepted_values": {"text", "json", "default"}}, + "debug": {"var_type": bool}, + "version_check": {"var_type": bool}, + } + for field_name, field_def in self.__dataclass_fields__.items(): + field_value = getattr(self, field_name) + + if not field_value is None: + vars_check = checks.get(field_name, {}) + accepted_values = vars_check.get("accepted_values") + var_type = vars_check.get("var_type", Any) + + if not isinstance(field_value, var_type): + raise CosmosValueError(f"dbt config var {field_name}: {field_value} must be a {var_type}") + + if accepted_values: + if field_value not in accepted_values: + raise CosmosValueError( + f"dbt config var {field_name}: {field_value} must be one of {accepted_values}" + ) + + def __post_init__(self) -> None: + self._validate_data() + + def as_dict(self) -> Optional[dict[str, Any]]: + result = { + field.name: getattr(self, field.name) + for field in dataclasses.fields(self) + if getattr(self, field.name) is not None + } + if result != {}: + return result + return None + + class BaseProfileMapping(ABC): """ A base class that other profile mappings should inherit from to ensure consistency. @@ -41,11 +101,13 @@ class BaseProfileMapping(ABC): _conn: Connection | None = None - def __init__(self, conn_id: str, profile_args: dict[str, Any] | None = None, disable_event_tracking: bool = False): + def __init__( + self, conn_id: str, profile_args: dict[str, Any] | None = None, dbt_config_vars: DbtConfigVars | None = None + ): self.conn_id = conn_id self.profile_args = profile_args or {} self._validate_profile_args() - self.disable_event_tracking = disable_event_tracking + self.dbt_config_vars = dbt_config_vars or DbtConfigVars() def _validate_profile_args(self) -> None: """ @@ -158,11 +220,7 @@ def env_vars(self) -> dict[str, str]: return env_vars def get_profile_file_contents( - self, - profile_name: str, - target_name: str = "cosmos_target", - use_mock_values: bool = False, - dbt_config_vars: dict[str, Any] | None = None, + self, profile_name: str, target_name: str = "cosmos_target", use_mock_values: bool = False ) -> str: """ Translates the profile into a string that can be written to a profiles.yml file. @@ -184,11 +242,9 @@ def get_profile_file_contents( } } - if not dbt_config_vars is None: - profile_contents["config"] = dbt_config_vars - - if self.disable_event_tracking: - profile_contents["config"] = {"send_anonymous_usage_stats": "False"} + congig_vars = self.dbt_config_vars.as_dict() + if congig_vars: + profile_contents["config"] = congig_vars return str(yaml.dump(profile_contents, indent=4)) diff --git a/dev/dags/basic_cosmos_dag.py b/dev/dags/basic_cosmos_dag.py index 79dfa4a4e..485d767a2 100755 --- a/dev/dags/basic_cosmos_dag.py +++ b/dev/dags/basic_cosmos_dag.py @@ -7,7 +7,7 @@ from pathlib import Path from cosmos import DbtDag, ProjectConfig, ProfileConfig -from cosmos.profiles import PostgresUserPasswordProfileMapping +from cosmos.profiles import PostgresUserPasswordProfileMapping, DbtConfigVars DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt" DBT_ROOT_PATH = Path(os.getenv("DBT_ROOT_PATH", DEFAULT_DBT_ROOT_PATH)) @@ -18,8 +18,8 @@ profile_mapping=PostgresUserPasswordProfileMapping( conn_id="airflow_db", profile_args={"schema": "public"}, + dbt_config_vars=DbtConfigVars(send_anonymous_usage_stats=True), ), - dbt_config_vars={"send_anonymous_usage_stats": False}, ) # [START local_example] diff --git a/docs/templates/index.rst.jinja2 b/docs/templates/index.rst.jinja2 index 80a757a1f..0ca25bf4f 100644 --- a/docs/templates/index.rst.jinja2 +++ b/docs/templates/index.rst.jinja2 @@ -86,36 +86,6 @@ but override the ``database`` and ``schema`` values: Note that when using a profile mapping, the profiles.yml file gets generated with the profile name and target name you specify in ``ProfileConfig``. -Disabling dbt event tracking --------------------------------- -.. versionadded:: 1.3 - -By default `dbt will track events `_ by sending anonymous usage data -when dbt commands are invoked. Users have an option to opt out of event tracking by updating their ``profiles.yml`` file. - -If you'd like to disable this behavior in the Cosmos generated profile, you can pass ``disable_event_tracking=True`` to the profile mapping like in -the example below: - -.. code-block:: python - - from cosmos.profiles import SnowflakeUserPasswordProfileMapping - - profile_config = ProfileConfig( - profile_name="my_profile_name", - target_name="my_target_name", - profile_mapping=SnowflakeUserPasswordProfileMapping( - conn_id="my_snowflake_conn_id", - profile_args={ - "database": "my_snowflake_database", - "schema": "my_snowflake_schema", - }, - disable_event_tracking=True, - ), - ) - - dag = DbtDag(profile_config=profile_config, ...) - - Dbt config vars -------------------------------- .. versionadded:: 1.3 @@ -124,7 +94,7 @@ The parts of ``profiles.yml``, which aren't specific to a particular data platfo .. code-block:: python - from cosmos.profiles import SnowflakeUserPasswordProfileMapping + from cosmos.profiles import SnowflakeUserPasswordProfileMapping, DbtConfigVars profile_config = ProfileConfig( profile_name="my_profile_name", @@ -135,10 +105,8 @@ The parts of ``profiles.yml``, which aren't specific to a particular data platfo "database": "my_snowflake_database", "schema": "my_snowflake_schema", }, + dbt_config_vars=DbtConfigVars(send_anonymous_usage_stats=True), ), - dbt_config_vars={ - "debug": True, - }, ) dag = DbtDag(profile_config=profile_config, ...) diff --git a/tests/profiles/test_base_profile.py b/tests/profiles/test_base_profile.py index 6f0859ca2..d391b7689 100644 --- a/tests/profiles/test_base_profile.py +++ b/tests/profiles/test_base_profile.py @@ -5,7 +5,7 @@ import pytest import yaml -from cosmos.profiles.base import BaseProfileMapping +from cosmos.profiles.base import BaseProfileMapping, DbtConfigVars from cosmos.exceptions import CosmosValueError @@ -38,33 +38,79 @@ def test_validate_profile_args(profile_arg: str): ) -@pytest.mark.parametrize("disable_event_tracking", [True, False]) -def test_disable_event_tracking(disable_event_tracking: bool): +@pytest.mark.parametrize("dbt_config_var,dbt_config_value", [("debug", True), ("debug", False)]) +def test_validate_dbt_config_vars(dbt_config_var: str, dbt_config_value: Any): """ - Tests the config block in the profile is set correctly if disable_event_tracking is set. + Tests the config block in the profile is set correctly. """ + dbt_config_vars = {dbt_config_var: dbt_config_value} test_profile = TestProfileMapping( conn_id="fake_conn_id", - disable_event_tracking=disable_event_tracking, + dbt_config_vars=DbtConfigVars(**dbt_config_vars), ) profile_contents = yaml.safe_load(test_profile.get_profile_file_contents(profile_name="fake-profile-name")) - assert ("config" in profile_contents) == disable_event_tracking - if disable_event_tracking: - assert profile_contents["config"]["send_anonymous_usage_stats"] == "False" + assert "config" in profile_contents + assert profile_contents["config"] == dbt_config_vars -@pytest.mark.parametrize("dbt_config_var,dbt_config_value", [("debug", True), ("debug", False)]) -def test_validate_dbt_config_vars(dbt_config_var: str, dbt_config_value: Any): - """ - Tests the config block in the profile is set correctly. - """ - test_profile = TestProfileMapping(conn_id="fake_conn_id") - profile_contents = yaml.safe_load( - test_profile.get_profile_file_contents( - profile_name="fake-profile-name", dbt_config_vars={dbt_config_var: dbt_config_value} +def test_profile_config_validate_dbt_config_vars_empty(): + test_profile = TestProfileMapping( + conn_id="fake_conn_id", + dbt_config_vars=None, + ) + assert test_profile.dbt_config_vars is None + + +def test_profile_config_validate_dbt_config_vars_check_unexpected_var(): + dbt_config_var = "unexpected_var" + dbt_config_value = True + dbt_config_vars = {dbt_config_var: dbt_config_value} + + with pytest.raises(CosmosValueError) as err_info: + TestProfileMapping( + conn_id="fake_conn_id", + dbt_config_vars=DbtConfigVars(**dbt_config_vars), ) + assert err_info.value.args[0] == f"dbt config var {dbt_config_var}: {dbt_config_value} is not supported" + + +@pytest.mark.parametrize( + "dbt_config_var,dbt_config_value", + [("send_anonymous_usage_stats", 2), ("send_anonymous_usage_stats", None)], +) +def test_profile_config_validate_dbt_config_vars_check_unexpected_types(dbt_config_var: str, dbt_config_value: Any): + dbt_config_vars = {dbt_config_var: dbt_config_value} + + with pytest.raises(CosmosValueError) as err_info: + TestProfileMapping( + conn_id="fake_conn_id", + dbt_config_vars=DbtConfigVars(**dbt_config_vars), + ) + assert err_info.value.args[0].startswith(f"dbt config var {dbt_config_var}: {dbt_config_value} must be a ") + + +@pytest.mark.parametrize("dbt_config_var,dbt_config_value", [("send_anonymous_usage_stats", True)]) +def test_profile_config_validate_dbt_config_vars_check_expected_types(dbt_config_var: str, dbt_config_value: Any): + dbt_config_vars = {dbt_config_var: dbt_config_value} + + profile_config = TestProfileMapping( + conn_id="fake_conn_id", + dbt_config_vars=DbtConfigVars(**dbt_config_vars), ) + assert profile_config.dbt_config_vars == dbt_config_vars - assert "config" in profile_contents - assert profile_contents["config"] == {dbt_config_var: dbt_config_value} + +@pytest.mark.parametrize( + "dbt_config_var,dbt_config_value", + [("log_format", "text2")], +) +def test_profile_config_validate_dbt_config_vars_check_values(dbt_config_var: str, dbt_config_value: Any): + dbt_config_vars = {dbt_config_var: dbt_config_value} + + with pytest.raises(CosmosValueError) as err_info: + TestProfileMapping( + conn_id="fake_conn_id", + dbt_config_vars=DbtConfigVars(**dbt_config_vars), + ) + assert err_info.value.args[0].startswith(f"dbt config var {dbt_config_var}: {dbt_config_value} must be one of ") diff --git a/tests/test_config.py b/tests/test_config.py index dbef30b50..6fa53b10c 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,5 +1,4 @@ from pathlib import Path -from typing import Any from unittest.mock import patch from cosmos.profiles.postgres.user_pass import PostgresUserPasswordProfileMapping @@ -134,108 +133,6 @@ def test_profile_config_validate_both(): ) -def test_profile_config_validate_dbt_config_vars_with_profiles_yml_filepath(): - dbt_config_vars = dict(send_anonymous_usage_stats=True) - - with pytest.raises(CosmosValueError) as err_info: - ProfileConfig( - profile_name="test", - target_name="test", - profiles_yml_filepath=SAMPLE_PROFILE_YML, - dbt_config_vars=dbt_config_vars, - ) - assert ( - err_info.value.args[0] - == "Both profiles_yml_filepath and dbt_config_vars are defined and are mutually exclusive. Ensure only one of these is defined." - ) - - -def test_profile_config_validate_dbt_config_vars_empty(): - profile_config = ProfileConfig( - profile_name="test", - target_name="test", - profile_mapping=PostgresUserPasswordProfileMapping(conn_id="test", profile_args={}), - dbt_config_vars=None, - ) - assert profile_config.dbt_config_vars is None - - -def test_profile_config_validate_dbt_config_vars_check_unexpected_var(): - dbt_config_var = "unexpected_var" - dbt_config_value = True - dbt_config_vars = {dbt_config_var: dbt_config_value} - - with pytest.raises(CosmosValueError) as err_info: - ProfileConfig( - profile_name="test", - target_name="test", - profile_mapping=PostgresUserPasswordProfileMapping(conn_id="test", profile_args={}), - dbt_config_vars=dbt_config_vars, - ) - assert err_info.value.args[0] == f"dbt config var {dbt_config_var}: {dbt_config_value} is not supported" - - -def test_profile_config_validate_dbt_config_vars_check_expected_var(): - dbt_config_var = "send_anonymous_usage_stats" - dbt_config_value = True - dbt_config_vars = {dbt_config_var: dbt_config_value} - - profile_config = ProfileConfig( - profile_name="test", - target_name="test", - profile_mapping=PostgresUserPasswordProfileMapping(conn_id="test", profile_args={}), - dbt_config_vars=dbt_config_vars, - ) - assert profile_config.dbt_config_vars == dbt_config_vars - - -@pytest.mark.parametrize( - "dbt_config_var,dbt_config_value", - [("send_anonymous_usage_stats", 2), ("send_anonymous_usage_stats", None)], -) -def test_profile_config_validate_dbt_config_vars_check_unexpected_types(dbt_config_var: str, dbt_config_value: Any): - dbt_config_vars = {dbt_config_var: dbt_config_value} - - with pytest.raises(CosmosValueError) as err_info: - ProfileConfig( - profile_name="test", - target_name="test", - profile_mapping=PostgresUserPasswordProfileMapping(conn_id="test", profile_args={}), - dbt_config_vars=dbt_config_vars, - ) - assert err_info.value.args[0].startswith(f"dbt config var {dbt_config_var}: {dbt_config_value} must be a ") - - -@pytest.mark.parametrize("dbt_config_var,dbt_config_value", [("send_anonymous_usage_stats", True)]) -def test_profile_config_validate_dbt_config_vars_check_expected_types(dbt_config_var: str, dbt_config_value: Any): - dbt_config_vars = {dbt_config_var: dbt_config_value} - - profile_config = ProfileConfig( - profile_name="test", - target_name="test", - profile_mapping=PostgresUserPasswordProfileMapping(conn_id="test", profile_args={}), - dbt_config_vars=dbt_config_vars, - ) - assert profile_config.dbt_config_vars == dbt_config_vars - - -@pytest.mark.parametrize( - "dbt_config_var,dbt_config_value", - [("log_format", "text2")], -) -def test_profile_config_validate_dbt_config_vars_check_values(dbt_config_var: str, dbt_config_value: Any): - dbt_config_vars = {dbt_config_var: dbt_config_value} - - with pytest.raises(CosmosValueError) as err_info: - ProfileConfig( - profile_name="test", - target_name="test", - profile_mapping=PostgresUserPasswordProfileMapping(conn_id="test", profile_args={}), - dbt_config_vars=dbt_config_vars, - ) - assert err_info.value.args[0].startswith(f"dbt config var {dbt_config_var}: {dbt_config_value} must be one of ") - - def test_profile_config_validate_profiles_yml(): profile_config = ProfileConfig(profile_name="test", target_name="test", profiles_yml_filepath="/tmp/no-exists") with pytest.raises(CosmosValueError) as err_info: From a0e9d285fb5f35eca856c88a377c5b66c137a82e Mon Sep 17 00:00:00 2001 From: yuriy Date: Fri, 15 Dec 2023 14:40:35 +0100 Subject: [PATCH 3/9] add-configs-in-profile --- cosmos/config.py | 47 +++++++++++++++++++++++++++++++++++- dev/dags/basic_cosmos_dag.py | 1 + 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/cosmos/config.py b/cosmos/config.py index 46e3f1915..069821621 100644 --- a/cosmos/config.py +++ b/cosmos/config.py @@ -212,6 +212,8 @@ class ProfileConfig: :param target_name: The name of the dbt target to use. :param profiles_yml_filepath: The path to a profiles.yml file to use. :param profile_mapping: A mapping of Airflow connections to dbt profiles. + :param dbt_config_vars: Dictionary of dbt configs for the project. This argument overrides configs defined in your profiles.yml + file. The dictionary is dumped to a yaml string. Details https://docs.getdbt.com/docs/core/connect-data-platform/profiles.yml """ # should always be set to be explicit @@ -223,9 +225,11 @@ class ProfileConfig: # should be set if using cosmos to map Airflow connections to dbt profiles profile_mapping: BaseProfileMapping | None = None + dbt_config_vars: dict[str, Any] | None = None def __post_init__(self) -> None: self.validate_profile() + self.validate_dbt_config_vars() def validate_profile(self) -> None: "Validates that we have enough information to render a profile." @@ -236,6 +240,44 @@ def validate_profile(self) -> None: "Both profiles_yml_filepath and profile_mapping are defined and are mutually exclusive. Ensure only one of these is defined." ) + def validate_dbt_config_vars(self) -> None: + "Validates config vars for profile." + + vars_checks = { + "send_anonymous_usage_stats": {'var_type': bool}, + "use_colors": {'var_type': bool}, + "partial_parse": {'var_type': bool}, + "printer_width": {'var_type': int}, + "write_json": {'var_type': bool}, + "warn_error": {'var_type': str}, + "log_format": {'var_type': str, 'accepted_values': {"text", "json", "default"}}, + "debug": {'var_type': bool}, + "version_check": {'var_type': bool}, + "fail_fast": {'var_type': bool}, + "use_experimental_parser": {'var_type': bool}, + "static_parser": {'var_type': bool}, + } + + if self.dbt_config_vars: + for var_key, var_value in self.dbt_config_vars.items(): + + vars_check = vars_checks.get(var_key) + if vars_check: + + var_type = vars_check.get('var_type') + if var_type: + if not isinstance(var_value, var_type): + raise CosmosValueError(f"dbt config var {var_key}: {var_value} must be a {var_type}") + + accepted_values = vars_check.get('accepted_values') + if accepted_values: + if var_value not in accepted_values: + raise CosmosValueError( + f"dbt config var {var_key}: {var_value} must be one of {accepted_values}" + ) + else: + warnings.warn(f"dbt config var {var_key}: {var_value} is not supported") + def validate_profiles_yml(self) -> None: "Validates a user-supplied profiles.yml is present" if self.profiles_yml_filepath and not Path(self.profiles_yml_filepath).exists(): @@ -252,7 +294,10 @@ def ensure_profile( elif self.profile_mapping: profile_contents = self.profile_mapping.get_profile_file_contents( - profile_name=self.profile_name, target_name=self.target_name, use_mock_values=use_mock_values + profile_name=self.profile_name, + target_name=self.target_name, + use_mock_values=use_mock_values, + dbt_config_vars=self.dbt_config_vars, ) if use_mock_values: diff --git a/dev/dags/basic_cosmos_dag.py b/dev/dags/basic_cosmos_dag.py index 485d767a2..f9a61ed13 100755 --- a/dev/dags/basic_cosmos_dag.py +++ b/dev/dags/basic_cosmos_dag.py @@ -20,6 +20,7 @@ profile_args={"schema": "public"}, dbt_config_vars=DbtConfigVars(send_anonymous_usage_stats=True), ), + dbt_config_vars={'send_anonymous_usage_stats': False}, ) # [START local_example] From 22b61414a06e3b9735b1735bd1368a7409d2d939 Mon Sep 17 00:00:00 2001 From: yuriy Date: Sat, 30 Dec 2023 12:59:56 +0100 Subject: [PATCH 4/9] add type check --- cosmos/config.py | 32 +++++++++++++++----------------- dev/dags/basic_cosmos_dag.py | 2 +- 2 files changed, 16 insertions(+), 18 deletions(-) diff --git a/cosmos/config.py b/cosmos/config.py index 069821621..014a77208 100644 --- a/cosmos/config.py +++ b/cosmos/config.py @@ -243,33 +243,31 @@ def validate_profile(self) -> None: def validate_dbt_config_vars(self) -> None: "Validates config vars for profile." - vars_checks = { - "send_anonymous_usage_stats": {'var_type': bool}, - "use_colors": {'var_type': bool}, - "partial_parse": {'var_type': bool}, - "printer_width": {'var_type': int}, - "write_json": {'var_type': bool}, - "warn_error": {'var_type': str}, - "log_format": {'var_type': str, 'accepted_values': {"text", "json", "default"}}, - "debug": {'var_type': bool}, - "version_check": {'var_type': bool}, - "fail_fast": {'var_type': bool}, - "use_experimental_parser": {'var_type': bool}, - "static_parser": {'var_type': bool}, + vars_checks: dict[str, dict[str, Any]] = { + "send_anonymous_usage_stats": {"var_type": bool}, + "use_colors": {"var_type": bool}, + "partial_parse": {"var_type": bool}, + "printer_width": {"var_type": int}, + "write_json": {"var_type": bool}, + "warn_error": {"var_type": str}, + "log_format": {"var_type": str, "accepted_values": {"text", "json", "default"}}, + "debug": {"var_type": bool}, + "version_check": {"var_type": bool}, + "fail_fast": {"var_type": bool}, + "use_experimental_parser": {"var_type": bool}, + "static_parser": {"var_type": bool}, } if self.dbt_config_vars: for var_key, var_value in self.dbt_config_vars.items(): - vars_check = vars_checks.get(var_key) if vars_check: - - var_type = vars_check.get('var_type') + var_type = vars_check.get("var_type") if var_type: if not isinstance(var_value, var_type): raise CosmosValueError(f"dbt config var {var_key}: {var_value} must be a {var_type}") - accepted_values = vars_check.get('accepted_values') + accepted_values = vars_check.get("accepted_values") if accepted_values: if var_value not in accepted_values: raise CosmosValueError( diff --git a/dev/dags/basic_cosmos_dag.py b/dev/dags/basic_cosmos_dag.py index f9a61ed13..46bdf6f93 100755 --- a/dev/dags/basic_cosmos_dag.py +++ b/dev/dags/basic_cosmos_dag.py @@ -20,7 +20,7 @@ profile_args={"schema": "public"}, dbt_config_vars=DbtConfigVars(send_anonymous_usage_stats=True), ), - dbt_config_vars={'send_anonymous_usage_stats': False}, + dbt_config_vars={"send_anonymous_usage_stats": False}, ) # [START local_example] From bafe365efd8564dcfd7474c58736229ea02d6ab1 Mon Sep 17 00:00:00 2001 From: yuriy Date: Thu, 11 Jan 2024 01:34:37 +0100 Subject: [PATCH 5/9] update tests --- docs/templates/index.rst.jinja2 | 2 +- tests/profiles/test_base_profile.py | 27 +++------------------------ 2 files changed, 4 insertions(+), 25 deletions(-) diff --git a/docs/templates/index.rst.jinja2 b/docs/templates/index.rst.jinja2 index 0ca25bf4f..a41a63c37 100644 --- a/docs/templates/index.rst.jinja2 +++ b/docs/templates/index.rst.jinja2 @@ -88,7 +88,7 @@ you specify in ``ProfileConfig``. Dbt config vars -------------------------------- -.. versionadded:: 1.3 +.. versionadded:: 1.3.2 The parts of ``profiles.yml``, which aren't specific to a particular data platform `dbt docs `_ diff --git a/tests/profiles/test_base_profile.py b/tests/profiles/test_base_profile.py index d391b7689..b81ce4e7d 100644 --- a/tests/profiles/test_base_profile.py +++ b/tests/profiles/test_base_profile.py @@ -51,33 +51,12 @@ def test_validate_dbt_config_vars(dbt_config_var: str, dbt_config_value: Any): profile_contents = yaml.safe_load(test_profile.get_profile_file_contents(profile_name="fake-profile-name")) assert "config" in profile_contents - assert profile_contents["config"] == dbt_config_vars - - -def test_profile_config_validate_dbt_config_vars_empty(): - test_profile = TestProfileMapping( - conn_id="fake_conn_id", - dbt_config_vars=None, - ) - assert test_profile.dbt_config_vars is None - - -def test_profile_config_validate_dbt_config_vars_check_unexpected_var(): - dbt_config_var = "unexpected_var" - dbt_config_value = True - dbt_config_vars = {dbt_config_var: dbt_config_value} - - with pytest.raises(CosmosValueError) as err_info: - TestProfileMapping( - conn_id="fake_conn_id", - dbt_config_vars=DbtConfigVars(**dbt_config_vars), - ) - assert err_info.value.args[0] == f"dbt config var {dbt_config_var}: {dbt_config_value} is not supported" + assert profile_contents["config"][dbt_config_var] == dbt_config_value @pytest.mark.parametrize( "dbt_config_var,dbt_config_value", - [("send_anonymous_usage_stats", 2), ("send_anonymous_usage_stats", None)], + [("send_anonymous_usage_stats", 2), ("send_anonymous_usage_stats", "aaa")], ) def test_profile_config_validate_dbt_config_vars_check_unexpected_types(dbt_config_var: str, dbt_config_value: Any): dbt_config_vars = {dbt_config_var: dbt_config_value} @@ -98,7 +77,7 @@ def test_profile_config_validate_dbt_config_vars_check_expected_types(dbt_config conn_id="fake_conn_id", dbt_config_vars=DbtConfigVars(**dbt_config_vars), ) - assert profile_config.dbt_config_vars == dbt_config_vars + assert profile_config.dbt_config_vars.as_dict() == dbt_config_vars @pytest.mark.parametrize( From 98d9fc64d98ffe5e5bee90378d12ad68054073f6 Mon Sep 17 00:00:00 2001 From: yuriy Date: Thu, 11 Jan 2024 01:42:34 +0100 Subject: [PATCH 6/9] rebase --- cosmos/config.py | 45 +-------------------------------------------- 1 file changed, 1 insertion(+), 44 deletions(-) diff --git a/cosmos/config.py b/cosmos/config.py index 014a77208..46e3f1915 100644 --- a/cosmos/config.py +++ b/cosmos/config.py @@ -212,8 +212,6 @@ class ProfileConfig: :param target_name: The name of the dbt target to use. :param profiles_yml_filepath: The path to a profiles.yml file to use. :param profile_mapping: A mapping of Airflow connections to dbt profiles. - :param dbt_config_vars: Dictionary of dbt configs for the project. This argument overrides configs defined in your profiles.yml - file. The dictionary is dumped to a yaml string. Details https://docs.getdbt.com/docs/core/connect-data-platform/profiles.yml """ # should always be set to be explicit @@ -225,11 +223,9 @@ class ProfileConfig: # should be set if using cosmos to map Airflow connections to dbt profiles profile_mapping: BaseProfileMapping | None = None - dbt_config_vars: dict[str, Any] | None = None def __post_init__(self) -> None: self.validate_profile() - self.validate_dbt_config_vars() def validate_profile(self) -> None: "Validates that we have enough information to render a profile." @@ -240,42 +236,6 @@ def validate_profile(self) -> None: "Both profiles_yml_filepath and profile_mapping are defined and are mutually exclusive. Ensure only one of these is defined." ) - def validate_dbt_config_vars(self) -> None: - "Validates config vars for profile." - - vars_checks: dict[str, dict[str, Any]] = { - "send_anonymous_usage_stats": {"var_type": bool}, - "use_colors": {"var_type": bool}, - "partial_parse": {"var_type": bool}, - "printer_width": {"var_type": int}, - "write_json": {"var_type": bool}, - "warn_error": {"var_type": str}, - "log_format": {"var_type": str, "accepted_values": {"text", "json", "default"}}, - "debug": {"var_type": bool}, - "version_check": {"var_type": bool}, - "fail_fast": {"var_type": bool}, - "use_experimental_parser": {"var_type": bool}, - "static_parser": {"var_type": bool}, - } - - if self.dbt_config_vars: - for var_key, var_value in self.dbt_config_vars.items(): - vars_check = vars_checks.get(var_key) - if vars_check: - var_type = vars_check.get("var_type") - if var_type: - if not isinstance(var_value, var_type): - raise CosmosValueError(f"dbt config var {var_key}: {var_value} must be a {var_type}") - - accepted_values = vars_check.get("accepted_values") - if accepted_values: - if var_value not in accepted_values: - raise CosmosValueError( - f"dbt config var {var_key}: {var_value} must be one of {accepted_values}" - ) - else: - warnings.warn(f"dbt config var {var_key}: {var_value} is not supported") - def validate_profiles_yml(self) -> None: "Validates a user-supplied profiles.yml is present" if self.profiles_yml_filepath and not Path(self.profiles_yml_filepath).exists(): @@ -292,10 +252,7 @@ def ensure_profile( elif self.profile_mapping: profile_contents = self.profile_mapping.get_profile_file_contents( - profile_name=self.profile_name, - target_name=self.target_name, - use_mock_values=use_mock_values, - dbt_config_vars=self.dbt_config_vars, + profile_name=self.profile_name, target_name=self.target_name, use_mock_values=use_mock_values ) if use_mock_values: From 5339fbb1ec673ee29b502d94cb82caab2c8451a9 Mon Sep 17 00:00:00 2001 From: yuriy Date: Thu, 11 Jan 2024 01:56:22 +0100 Subject: [PATCH 7/9] update dev dag --- dev/dags/basic_cosmos_dag.py | 1 - 1 file changed, 1 deletion(-) diff --git a/dev/dags/basic_cosmos_dag.py b/dev/dags/basic_cosmos_dag.py index 46bdf6f93..485d767a2 100755 --- a/dev/dags/basic_cosmos_dag.py +++ b/dev/dags/basic_cosmos_dag.py @@ -20,7 +20,6 @@ profile_args={"schema": "public"}, dbt_config_vars=DbtConfigVars(send_anonymous_usage_stats=True), ), - dbt_config_vars={"send_anonymous_usage_stats": False}, ) # [START local_example] From 8a0a30a011ca923fdf1721c12993fbed0c92d002 Mon Sep 17 00:00:00 2001 From: yuriy Date: Thu, 11 Jan 2024 02:31:36 +0100 Subject: [PATCH 8/9] rebase --- cosmos/profiles/base.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cosmos/profiles/base.py b/cosmos/profiles/base.py index c1b8f93ac..15558c93b 100644 --- a/cosmos/profiles/base.py +++ b/cosmos/profiles/base.py @@ -242,9 +242,9 @@ def get_profile_file_contents( } } - congig_vars = self.dbt_config_vars.as_dict() - if congig_vars: - profile_contents["config"] = congig_vars + config_vars = self.dbt_config_vars.as_dict() + if config_vars: + profile_contents["config"] = config_vars return str(yaml.dump(profile_contents, indent=4)) From 73c173393d45dcea968de3f9a67b44eb7d3f6c21 Mon Sep 17 00:00:00 2001 From: yuriy Date: Thu, 11 Jan 2024 02:42:39 +0100 Subject: [PATCH 9/9] rebase --- cosmos/profiles/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/cosmos/profiles/base.py b/cosmos/profiles/base.py index 15558c93b..4f357d070 100644 --- a/cosmos/profiles/base.py +++ b/cosmos/profiles/base.py @@ -53,6 +53,7 @@ def _validate_data(self) -> None: "debug": {"var_type": bool}, "version_check": {"var_type": bool}, } + for field_name, field_def in self.__dataclass_fields__.items(): field_value = getattr(self, field_name)