Skip to content

Commit

Permalink
Fix ACI integration (astronomer#872)
Browse files Browse the repository at this point in the history
The inheritance order of the `DbtAzureContainerInstanceBaseOperator` was
incorrect, as well as the wrong execute method being called. This led to
no dbt command actually being built nor run. I noticed this when testing
the 1.4.0a1 release.

Example log before:
```
[2024-03-01, 20:27:29 UTC] {container_instances.py:366} INFO - Usage: dbt [OPTIONS] COMMAND [ARGS]...
[2024-03-01, 20:27:29 UTC] {container_instances.py:366} INFO - 
```

After this fix:
```
[2024-03-01, 20:28:44 UTC] {container_instances.py:366} INFO - �[0m20:28:44  Found 5 models, 3 seeds, 20 tests, 0 sources, 0 exposures, 0 metrics, 403 macros, 0 groups, 0 semantic models
[2024-03-01, 20:28:44 UTC] {container_instances.py:366} INFO - �[0m20:28:44
```


## Checklist

- [x] I have made corresponding changes to the documentation (if
required)
- [x] I have added tests that prove my fix is effective or that my
feature works
  • Loading branch information
danielvdende authored and arojasb3 committed Jul 14, 2024
1 parent cef5bec commit e5b3039
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 8 deletions.
16 changes: 8 additions & 8 deletions cosmos/operators/azure_container_instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
)


class DbtAzureContainerInstanceBaseOperator(AzureContainerInstancesOperator, AbstractDbtBaseOperator): # type: ignore
class DbtAzureContainerInstanceBaseOperator(AbstractDbtBaseOperator, AzureContainerInstancesOperator): # type: ignore
"""
Executes a dbt core cli command in an Azure Container Instance
"""
Expand Down Expand Up @@ -66,7 +66,7 @@ def __init__(
def build_and_run_cmd(self, context: Context, cmd_flags: list[str] | None = None) -> None:
self.build_command(context, cmd_flags)
self.log.info(f"Running command: {self.command}")
result = super().execute(context)
result = AzureContainerInstancesOperator.execute(self, context)
logger.info(result)

def build_command(self, context: Context, cmd_flags: list[str] | None = None) -> None:
Expand All @@ -79,13 +79,13 @@ def build_command(self, context: Context, cmd_flags: list[str] | None = None) ->
self.command: list[str] = dbt_cmd


class DbtLSAzureContainerInstanceOperator(DbtLSMixin, DbtAzureContainerInstanceBaseOperator):
class DbtLSAzureContainerInstanceOperator(DbtLSMixin, DbtAzureContainerInstanceBaseOperator): # type: ignore
"""
Executes a dbt core ls command.
"""


class DbtSeedAzureContainerInstanceOperator(DbtSeedMixin, DbtAzureContainerInstanceBaseOperator):
class DbtSeedAzureContainerInstanceOperator(DbtSeedMixin, DbtAzureContainerInstanceBaseOperator): # type: ignore
"""
Executes a dbt core seed command.
Expand All @@ -95,22 +95,22 @@ class DbtSeedAzureContainerInstanceOperator(DbtSeedMixin, DbtAzureContainerInsta
template_fields: Sequence[str] = DbtAzureContainerInstanceBaseOperator.template_fields + DbtRunMixin.template_fields # type: ignore[operator]


class DbtSnapshotAzureContainerInstanceOperator(DbtSnapshotMixin, DbtAzureContainerInstanceBaseOperator):
class DbtSnapshotAzureContainerInstanceOperator(DbtSnapshotMixin, DbtAzureContainerInstanceBaseOperator): # type: ignore
"""
Executes a dbt core snapshot command.
"""


class DbtRunAzureContainerInstanceOperator(DbtRunMixin, DbtAzureContainerInstanceBaseOperator):
class DbtRunAzureContainerInstanceOperator(DbtRunMixin, DbtAzureContainerInstanceBaseOperator): # type: ignore
"""
Executes a dbt core run command.
"""

template_fields: Sequence[str] = DbtAzureContainerInstanceBaseOperator.template_fields + DbtRunMixin.template_fields # type: ignore[operator]


class DbtTestAzureContainerInstanceOperator(DbtTestMixin, DbtAzureContainerInstanceBaseOperator):
class DbtTestAzureContainerInstanceOperator(DbtTestMixin, DbtAzureContainerInstanceBaseOperator): # type: ignore
"""
Executes a dbt core test command.
"""
Expand All @@ -121,7 +121,7 @@ def __init__(self, on_warning_callback: Callable[..., Any] | None = None, **kwar
self.on_warning_callback = on_warning_callback


class DbtRunOperationAzureContainerInstanceOperator(DbtRunOperationMixin, DbtAzureContainerInstanceBaseOperator):
class DbtRunOperationAzureContainerInstanceOperator(DbtRunOperationMixin, DbtAzureContainerInstanceBaseOperator): # type: ignore
"""
Executes a dbt core run-operation command.
Expand Down
22 changes: 22 additions & 0 deletions tests/operators/test_azure_container_instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -145,3 +145,25 @@ def test_dbt_azure_container_instance_build_command():
"start_time: '{{ data_interval_start.strftime(''%Y%m%d%H%M%S'') }}'\n",
"--no-version-check",
]


@patch("cosmos.operators.azure_container_instance.AzureContainerInstancesOperator.execute")
def test_dbt_azure_container_instance_build_and_run_cmd(mock_execute):
dbt_base_operator = ConcreteDbtAzureContainerInstanceOperator(
ci_conn_id="my_airflow_connection",
task_id="my-task",
image="my_image",
region="Mordor",
name="my-aci",
resource_group="my-rg",
project_dir="my/dir",
environment_variables={"FOO": "BAR"},
)
mock_build_command = MagicMock()
dbt_base_operator.build_command = mock_build_command

mock_context = MagicMock()
dbt_base_operator.build_and_run_cmd(context=mock_context)

mock_build_command.assert_called_with(mock_context, None)
mock_execute.assert_called_once_with(dbt_base_operator, mock_context)

0 comments on commit e5b3039

Please sign in to comment.