Improve performance of AzureContainerInstanceHook.exists using direct container lookup

providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py

@@ -21,6 +21,7 @@
 from typing import TYPE_CHECKING, Any, cast
 
 from azure.common.client_factory import get_client_from_auth_file, get_client_from_json_dict
+from azure.core.exceptions import ResourceNotFoundError
 from azure.identity import ClientSecretCredential, DefaultAzureCredential
 from azure.mgmt.containerinstance import ContainerInstanceManagementClient
 
@@ -154,10 +155,11 @@ def exists(self, resource_group: str, name: str) -> bool:
         :param resource_group: the name of the resource group
         :param name: the name of the container group
         """
-        for container in self.connection.container_groups.list_by_resource_group(resource_group):
-            if container.name == name:
-                return True
-        return False
+        try:
+            self.connection.container_groups.get(resource_group, name)
+            return True
+        except ResourceNotFoundError:
+            return False
 
     def test_connection(self):
         """Test a configured Azure Container Instance connection."""

providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py

@@ -20,9 +20,8 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
+from azure.core.exceptions import ResourceNotFoundError
 from azure.mgmt.containerinstance.models import (
-    Container,
-    ContainerGroup,
     Logs,
     ResourceRequests,
     ResourceRequirements,
@@ -96,25 +95,26 @@ def test_delete(self, delete_mock):
         self.hook.delete("resource_group", "aci-test")
         delete_mock.assert_called_once_with("resource_group", "aci-test")
 
-    @patch("azure.mgmt.containerinstance.operations.ContainerGroupsOperations.list_by_resource_group")
-    def test_exists_with_existing(self, list_mock):
-        list_mock.return_value = [
-            ContainerGroup(
-                os_type="Linux",
-                containers=[Container(name="test1", image="hello-world", resources=self.resources)],
-            )
-        ]
-        assert not self.hook.exists("test", "test1")
-
-    @patch("azure.mgmt.containerinstance.operations.ContainerGroupsOperations.list_by_resource_group")
-    def test_exists_with_not_existing(self, list_mock):
-        list_mock.return_value = [
-            ContainerGroup(
-                os_type="Linux",
-                containers=[Container(name="test1", image="hello-world", resources=self.resources)],
-            )
-        ]
+    @patch("azure.mgmt.containerinstance.operations.ContainerGroupsOperations.get")
+    def test_exists_with_existing(self, get_mock):
+        get_mock.return_value = MagicMock()
+
+        assert self.hook.exists("test", "test1")
+        get_mock.assert_called_once_with("test", "test1")
+
+    @patch("azure.mgmt.containerinstance.operations.ContainerGroupsOperations.get")
+    def test_exists_with_not_existing(self, get_mock):
+        get_mock.side_effect = ResourceNotFoundError("not found")
+
         assert not self.hook.exists("test", "not found")
+        get_mock.assert_called_once_with("test", "not found")
+
+    @patch("azure.mgmt.containerinstance.operations.ContainerGroupsOperations.get")
+    def test_exists_unexpected_exception(self, get_mock):
+        get_mock.side_effect = RuntimeError("Unexpected Exception")
+
+        with pytest.raises(RuntimeError):
+            self.hook.exists("test", "test")
 
     @patch("azure.mgmt.containerinstance.operations.ContainerGroupsOperations.list")
     def test_connection_success(self, mock_container_groups_list):