add dns_config and diagnostics parameters to AzureContainerInstancesOperator

airflow/providers/microsoft/azure/operators/container_instances.py

@@ -25,8 +25,10 @@
 from azure.mgmt.containerinstance.models import (
     Container,
     ContainerGroup,
+    ContainerGroupDiagnostics,
     ContainerGroupSubnetId,
     ContainerPort,
+    DnsConfiguration,
     EnvironmentVariable,
     IpAddress,
     ResourceRequests,
@@ -90,6 +92,8 @@ class AzureContainerInstancesOperator(BaseOperator):
         Possible values include: 'Always', 'OnFailure', 'Never'
     :param ip_address: The IP address type of the container group.
     :param subnet_ids: The subnet resource IDs for a container group
+    :param dns_config: The DNS configuration for a container group.
+    :param diagnostics: Container group diagnostic information (Log Analytics).
 
     **Example**::
 
@@ -113,6 +117,13 @@ class AzureContainerInstancesOperator(BaseOperator):
             memory_in_gb=14.0,
             cpu=4.0,
             gpu=GpuResource(count=1, sku="K80"),
+            dns_config=["10.0.0.10", "10.0.0.11"],
+            diagnostics={
+                "log_analytics": {
+                    "workspaceId": "workspaceid",
+                    "workspaceKey": "workspaceKey",
+                }
+            },
             command=["/bin/echo", "world"],
             task_id="start_container",
         )
@@ -145,6 +156,8 @@ def __init__(
         ip_address: IpAddress | None = None,
         ports: list[ContainerPort] | None = None,
         subnet_ids: list[ContainerGroupSubnetId] | None = None,
+        dns_config: DnsConfiguration | None = None,
+        diagnostics: ContainerGroupDiagnostics | None = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -183,6 +196,8 @@ def __init__(
         self.ip_address = ip_address
         self.ports = ports
         self.subnet_ids = subnet_ids
+        self.dns_config = dns_config
+        self.diagnostics = diagnostics
 
     def execute(self, context: Context) -> int:
         # Check name again in case it was templated.
@@ -256,6 +271,8 @@ def execute(self, context: Context) -> int:
                 tags=self.tags,
                 ip_address=self.ip_address,
                 subnet_ids=self.subnet_ids,
+                dns_config=self.dns_config,
+                diagnostics=self.diagnostics,
             )
 
             self._ci_hook.create_or_update(self.resource_group, self.name, container_group)

tests/providers/microsoft/azure/operators/test_container_instances.py

@@ -410,3 +410,53 @@ def test_execute_with_missing_events(self, log_mock, aci_mock):
         aci.execute(None)
 
         assert log_mock.call_count == 0
+
+    @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook")
+    def test_execute_with_dnsconfig(self, aci_mock):
+        dns_config = MagicMock()
+
+        aci_mock.return_value.get_state.return_value = make_mock_container(
+            state="Terminated", exit_code=0, detail_status="test"
+        )
+        aci_mock.return_value.exists.return_value = False
+
+        aci = AzureContainerInstancesOperator(
+            ci_conn_id=None,
+            registry_conn_id=None,
+            resource_group="resource-group",
+            name="container-name",
+            image="container-image",
+            region="region",
+            task_id="task",
+            dns_config=dns_config,
+        )
+        aci.execute(None)
+        assert aci_mock.return_value.create_or_update.call_count == 1
+        (_, _, called_cg), _ = aci_mock.return_value.create_or_update.call_args
+
+        assert called_cg.dns_config == dns_config
+
+    @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook")
+    def test_execute_with_diagnostics(self, aci_mock):
+        diagnostics = MagicMock()
+
+        aci_mock.return_value.get_state.return_value = make_mock_container(
+            state="Terminated", exit_code=0, detail_status="test"
+        )
+        aci_mock.return_value.exists.return_value = False
+
+        aci = AzureContainerInstancesOperator(
+            ci_conn_id=None,
+            registry_conn_id=None,
+            resource_group="resource-group",
+            name="container-name",
+            image="container-image",
+            region="region",
+            task_id="task",
+            diagnostics=diagnostics,
+        )
+        aci.execute(None)
+        assert aci_mock.return_value.create_or_update.call_count == 1
+        (_, _, called_cg), _ = aci_mock.return_value.create_or_update.call_args
+
+        assert called_cg.diagnostics == diagnostics