Add managed identity assignment support

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

@@ -27,6 +27,7 @@
     Container,
     ContainerGroup,
     ContainerGroupDiagnostics,
+    ContainerGroupIdentity,
     ContainerGroupSubnetId,
     ContainerPort,
     DnsConfiguration,
@@ -102,6 +103,7 @@ class AzureContainerInstancesOperator(BaseOperator):
     :param dns_config: The DNS configuration for a container group.
     :param diagnostics: Container group diagnostic information (Log Analytics).
     :param priority: Container group priority, Possible values include: 'Regular', 'Spot'
+    :param identity: List of User/System assigned identities for the container group.
 
     **Example**::
 
@@ -144,6 +146,12 @@ class AzureContainerInstancesOperator(BaseOperator):
                 }
             },
             priority="Regular",
+            identity = {
+                {
+                    "type": "UserAssigned",
+                    "resource_ids": ["/subscriptions/00000000-0000-0000-0000-00000000000/resourceGroups/my_rg/providers/Microsoft.ManagedIdentity/userAssignedIdentities/my_identity"],
+                },
+            }
             command=["/bin/echo", "world"],
             task_id="start_container",
         )
@@ -180,6 +188,7 @@ def __init__(
         dns_config: DnsConfiguration | None = None,
         diagnostics: ContainerGroupDiagnostics | None = None,
         priority: str | None = "Regular",
+        identity: ContainerGroupIdentity | None = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -222,6 +231,7 @@ def __init__(
         self.dns_config = dns_config
         self.diagnostics = diagnostics
         self.priority = priority
+        self.identity = identity
         if self.priority not in ["Regular", "Spot"]:
             raise AirflowException(
                 "Invalid value for the priority argument. "
@@ -304,6 +314,7 @@ def execute(self, context: Context) -> int:
                 dns_config=self.dns_config,
                 diagnostics=self.diagnostics,
                 priority=self.priority,
+                identity=self.identity,
             )
 
             self._ci_hook.create_or_update(self.resource_group, self.name, container_group)

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

@@ -586,6 +586,31 @@ def test_execute_with_spot_discount(self, aci_mock):
 
         assert aci_mock.return_value.delete.call_count == 1
 
+    @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook")
+    def test_execute_with_identity(self, aci_mock):
+        identity = 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",
+            identity=identity,
+        )
+        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.identity == identity
+
 
 class XcomMock:
     def __init__(self) -> None: