diff --git a/cosmos/operators/docker.py b/cosmos/operators/docker.py index 05671b4d0..81b7196b3 100644 --- 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 index 9da31f00d..4883605dc 100644 --- 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)