adds profile_config for Docker (#1347)

Adds ProfileConfig to Docker operators so that people can use TaskGroups. Closes #1346
astronomer · Dec 20, 2024 · 3b92421 · 3b92421
1 parent c21f8c9
commit 3b92421
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/cosmos/operators/docker.py b/cosmos/operators/docker.py
@@ -4,6 +4,7 @@
 
 from airflow.utils.context import Context
 
+from cosmos.config import ProfileConfig
 from cosmos.operators.base import (
     AbstractDbtBaseOperator,
     DbtBuildMixin,
@@ -42,8 +43,10 @@ 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
         super().__init__(image=image, **kwargs)
 
     def build_and_run_cmd(self, context: Context, cmd_flags: list[str] | None = None) -> Any:

diff --git a/tests/test_converter.py b/tests/test_converter.py
@@ -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)