Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Address review comments on PR 1347 regarding profile_config for ExecutionMode.Docker #1413

Merged
merged 7 commits into from
Dec 20, 2024
Merged
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions cosmos/operators/docker.py
Original file line number Diff line number Diff line change
@@ -4,6 +4,8 @@

from airflow.utils.context import Context

from cosmos.config import ProfileConfig
from cosmos.exceptions import CosmosValueError
from cosmos.operators.base import (
AbstractDbtBaseOperator,
DbtBuildMixin,
@@ -42,8 +44,17 @@ class DbtDockerBaseOperator(AbstractDbtBaseOperator, DockerOperator): # type: i
def __init__(
self,
image: str, # Make image a required argument since it's required by DockerOperator
profile_config: ProfileConfig | None = None,
**kwargs: Any,
) -> None:
self.profile_config = profile_config
if self.profile_config and not self.profile_config.profiles_yml_filepath:
raise CosmosValueError(
"For ExecutionMode.DOCKER, specifying ProfileConfig only works with profiles_yml_filepath method and "
"it must be specified. ProfileConfig with ProfileMapping method is not supported as the underlying "
"Airflow connections are not available in the Docker container for the mapping to work."
)

super().__init__(image=image, **kwargs)

def build_and_run_cmd(self, context: Context, cmd_flags: list[str] | None = None) -> Any:
12 changes: 12 additions & 0 deletions docs/getting_started/docker.rst
Original file line number Diff line number Diff line change
@@ -97,3 +97,15 @@ Enable and trigger a run of the `jaffle_shop_docker <https://github.com/astronom

.. figure:: https://github.com/astronomer/astronomer-cosmos/raw/main/docs/_static/jaffle_shop_docker_dag_run.png
:width: 800


Specifying ProfileConfig
+++++++++++++++++++++++++

Starting with Cosmos 1.8.0, you can use the ``profile_config`` argument in your Dbt DAG Docker operators to reference
profiles for your dbt project defined in a profiles.yml file. To do so, provide the file’s path via the
``profiles_yml_path`` parameter in ``profile_config``.

Note that in ``ExecutionMode.DOCKER``, the ``profile_config`` is only compatible with the ``profiles_yml_path``
approach. The ``profile_mapping`` method will not work because the required Airflow connections cannot be accessed
within the Docker container to map them to the dbt profile.
27 changes: 27 additions & 0 deletions tests/operators/test_docker.py
Original file line number Diff line number Diff line change
@@ -5,6 +5,8 @@
from airflow.utils.context import Context
from pendulum import datetime

from cosmos import ProfileConfig
from cosmos.exceptions import CosmosValueError
from cosmos.operators.docker import (
DbtBuildDockerOperator,
DbtCloneDockerOperator,
@@ -13,6 +15,7 @@
DbtSeedDockerOperator,
DbtTestDockerOperator,
)
from cosmos.profiles import PostgresUserPasswordProfileMapping


@pytest.fixture()
@@ -133,3 +136,27 @@ def test_dbt_docker_build_command():
"start_time: '{{ data_interval_start.strftime(''%Y%m%d%H%M%S'') }}'\n",
"--no-version-check",
]


def test_profile_config_without_profiles_yml_raises_error(base_operator):
with pytest.raises(CosmosValueError) as err:
base_operator(
conn_id="my_airflow_connection",
task_id="my-task",
image="my_image",
project_dir="my/dir",
append_env=False,
profile_config=ProfileConfig(
profile_name="profile_name",
target_name="target_name",
profile_mapping=PostgresUserPasswordProfileMapping(
conn_id="example_conn",
profile_args={"schema": "public"},
),
),
)

error_message = str(err.value)
expected_err_msg = "For ExecutionMode.DOCKER, specifying ProfileConfig only works with profiles_yml_filepath method"

assert expected_err_msg in error_message
11 changes: 6 additions & 5 deletions tests/test_converter.py
Original file line number Diff line number Diff line change
@@ -134,7 +134,7 @@ def test_validate_arguments_schema_in_task_args():
"execution_mode,operator_args",
[
(ExecutionMode.KUBERNETES, {}),
# (ExecutionMode.DOCKER, {"image": "sample-image"}),
(ExecutionMode.DOCKER, {"image": "sample-image"}),
],
)
@patch("cosmos.converter.DbtGraph.filtered_nodes", nodes)
@@ -168,6 +168,7 @@ def test_converter_creates_dag_with_seed(mock_load_dbt_graph, execution_mode, op
"execution_mode,operator_args",
[
(ExecutionMode.KUBERNETES, {}),
(ExecutionMode.DOCKER, {"image": "sample-image"}),
],
)
@patch("cosmos.converter.DbtGraph.filtered_nodes", nodes)
@@ -201,7 +202,7 @@ def test_converter_creates_dag_with_project_path_str(mock_load_dbt_graph, execut
"execution_mode,virtualenv_dir,operator_args",
[
(ExecutionMode.KUBERNETES, Path("/some/virtualenv/dir"), {}),
# (ExecutionMode.DOCKER, {"image": "sample-image"}),
(ExecutionMode.DOCKER, Path("/some/virtualenv/dir"), {"image": "sample-image"}),
],
)
@patch("cosmos.converter.DbtGraph.filtered_nodes", nodes)
@@ -241,7 +242,7 @@ def test_converter_raises_warning(mock_load_dbt_graph, execution_mode, virtualen
"execution_mode,operator_args",
[
(ExecutionMode.KUBERNETES, {}),
# (ExecutionMode.DOCKER, {"image": "sample-image"}),
(ExecutionMode.DOCKER, {"image": "sample-image"}),
],
)
@patch("cosmos.converter.DbtGraph.filtered_nodes", nodes)
@@ -278,7 +279,7 @@ def test_converter_fails_execution_config_no_project_dir(mock_load_dbt_graph, ex
"execution_mode,operator_args",
[
(ExecutionMode.KUBERNETES, {}),
# (ExecutionMode.DOCKER, {"image": "sample-image"}),
(ExecutionMode.DOCKER, {"image": "sample-image"}),
],
)
@patch("cosmos.converter.DbtGraph.filtered_nodes", nodes)
@@ -317,7 +318,7 @@ def test_converter_fails_project_config_path_and_execution_config_path(
"execution_mode,operator_args",
[
(ExecutionMode.KUBERNETES, {}),
# (ExecutionMode.DOCKER, {"image": "sample-image"}),
(ExecutionMode.DOCKER, {"image": "sample-image"}),
],
)
@patch("cosmos.converter.DbtGraph.filtered_nodes", nodes)