reduce how often list_containers is called (#196)

Author: bmc-msft

src/api-service/__app__/containers/__init__.py

@@ -15,6 +15,7 @@
 from ..onefuzzlib.azure.containers import (
     create_container,
     delete_container,
+    get_container_metadata,
     get_container_sas_url,
     get_containers,
 )
@@ -28,28 +29,29 @@ def get(req: func.HttpRequest) -> func.HttpResponse:
 
     if isinstance(request, Error):
         return not_ok(request, context="container get")
-
-    containers = get_containers()
-
     if request is not None:
-        if request.name in containers:
-            info = ContainerInfo(
-                name=request.name,
-                sas_url=get_container_sas_url(
-                    request.name,
-                    read=True,
-                    write=True,
-                    create=True,
-                    delete=True,
-                    list=True,
-                ),
-                metadata=containers[request.name],
+        metadata = get_container_metadata(request.name)
+        if metadata is None:
+            return not_ok(
+                Error(code=ErrorCode.INVALID_REQUEST, errors=["invalid container"]),
+                context=request.name,
             )
-            return ok(info)
-        return not_ok(
-            Error(code=ErrorCode.INVALID_REQUEST, errors=["invalid container"]),
-            context=request.name,
+
+        info = ContainerInfo(
+            name=request.name,
+            sas_url=get_container_sas_url(
+                request.name,
+                read=True,
+                write=True,
+                create=True,
+                delete=True,
+                list=True,
+            ),
+            metadata=metadata,
         )
+        return ok(info)
+
+    containers = get_containers()
 
     container_info = []
     for name in containers:

src/api-service/__app__/download/__init__.py

@@ -7,7 +7,11 @@
 from onefuzztypes.enums import ErrorCode
 from onefuzztypes.models import Error, FileEntry
 
-from ..onefuzzlib.azure.containers import blob_exists, get_containers, get_file_sas_url
+from ..onefuzzlib.azure.containers import (
+    blob_exists,
+    container_exists,
+    get_file_sas_url,
+)
 from ..onefuzzlib.request import not_ok, parse_uri, redirect
 
 
@@ -16,7 +20,7 @@ def get(req: func.HttpRequest) -> func.HttpResponse:
     if isinstance(request, Error):
         return not_ok(request, context="download")
 
-    if request.container not in get_containers():
+    if not container_exists(request.container):
         return not_ok(
             Error(code=ErrorCode.INVALID_REQUEST, errors=["invalid container"]),
             context=request.container,

src/api-service/__app__/notifications/__init__.py

@@ -10,7 +10,7 @@
 from onefuzztypes.models import Error
 from onefuzztypes.requests import NotificationCreate, NotificationGet
 
-from ..onefuzzlib.azure.containers import get_containers
+from ..onefuzzlib.azure.containers import container_exists
 from ..onefuzzlib.notifications.main import Notification
 from ..onefuzzlib.request import not_ok, ok, parse_request
 
@@ -29,8 +29,7 @@ def post(req: func.HttpRequest) -> func.HttpResponse:
     if isinstance(request, Error):
         return not_ok(request, context="notification create")
 
-    containers = get_containers()
-    if request.container not in containers:
+    if not container_exists(request.container):
         return not_ok(
             Error(code=ErrorCode.INVALID_REQUEST, errors=["invalid container"]),
             context=request.container,

src/api-service/__app__/onefuzzlib/azure/containers.py

@@ -10,10 +10,20 @@
 
 from azure.common import AzureHttpError, AzureMissingResourceHttpError
 from azure.storage.blob import BlobPermissions, ContainerPermissions
+from memoization import cached
 
 from .creds import get_blob_service
 
 
+@cached(ttl=5)
+def container_exists(name: str, account_id: Optional[str] = None) -> bool:
+    try:
+        get_blob_service(account_id).get_container_properties(name)
+        return True
+    except AzureHttpError:
+        return False
+
+
 def get_containers(account_id: Optional[str] = None) -> Dict[str, Dict[str, str]]:
     return {
         x.name: x.metadata
@@ -27,8 +37,7 @@ def get_container_metadata(
 ) -> Optional[Dict[str, str]]:
     try:
         result = get_blob_service(account_id).get_container_metadata(name)
-        if result:
-            return cast(Dict[str, str], result)
+        return cast(Dict[str, str], result)
     except AzureHttpError:
         pass
     return None

src/api-service/__app__/onefuzzlib/tasks/config.py

@@ -11,7 +11,7 @@
 from onefuzztypes.enums import Compare, ContainerPermission, ContainerType, TaskFeature
 from onefuzztypes.models import TaskConfig, TaskDefinition, TaskUnitConfig
 
-from ..azure.containers import blob_exists, get_container_sas_url, get_containers
+from ..azure.containers import blob_exists, container_exists, get_container_sas_url
 from ..azure.creds import get_fuzz_storage, get_instance_url
 from ..azure.queue import get_queue_sas
 from .defs import TASK_DEFINITIONS
@@ -58,12 +58,14 @@ def check_container(
 
 
 def check_containers(definition: TaskDefinition, config: TaskConfig) -> None:
-    all_containers = set(get_containers().keys())
+    checked = set()
 
     containers: Dict[ContainerType, List[str]] = {}
     for container in config.containers:
-        if container.name not in all_containers:
-            raise TaskConfigError("missing container: %s" % container.name)
+        if container.name not in checked:
+            if not container_exists(container.name):
+                raise TaskConfigError("missing container: %s" % container.name)
+            checked.add(container.name)
 
         if container.type not in containers:
             containers[container.type] = []