add docker image uri validated str

src/aibs_informatics_aws_lambda/handlers/batch/create.py

@@ -1,5 +1,6 @@
 import re
 from dataclasses import dataclass
+from typing import ClassVar, Optional
 
 from aibs_informatics_aws_utils.batch import (
     BatchJobBuilder,
@@ -16,11 +17,30 @@
     CreateDefinitionAndPrepareArgsResponse,
 )
 
+REGISTRY = r"(docker\.io|public\.ecr\.aws|ghcr\.io|(?:[\d]{12}).dkr.ecr.(?:[\w-]*).amazonaws.com)"
+REPO_NAME = r"([a-z0-9]+(?:[-_\.\/][a-z0-9]+)*)"
+IMAGE_TAG_OR_SHA = r"(?::([\w.\-_]{1,127})|@sha256:([A-Fa-f0-9]{64}))"
 
-def is_valid_docker_uri(uri):
-    pattern = r"^(docker\.io|public\.ecr\.aws)\/([a-z0-9]+([-_\.\/][a-z0-9]+)*)(:[a-z0-9]+([-_\.][a-z0-9]+)*)?$"
-    match = re.match(pattern, uri)
-    return match is not None
+
+# TODO: move to aibs-informatics-core once this is stable
+class DockerImageUri(ValidatedStr):
+    regex_pattern: ClassVar[re.Pattern] = re.compile(f"^{REGISTRY}/{REPO_NAME}{IMAGE_TAG_OR_SHA}?")
+
+    @property
+    def registry(self) -> str:
+        return self.get_match_groups()[0]
+
+    @property
+    def repo_name(self) -> str:
+        return self.get_match_groups()[1]
+
+    @property
+    def tag(self) -> Optional[str]:
+        return self.get_match_groups()[2]
+
+    @property
+    def sha256(self) -> Optional[str]:
+        return self.get_match_groups()[3]
 
 
 @dataclass
@@ -32,7 +52,7 @@ def handle(
     ) -> CreateDefinitionAndPrepareArgsResponse:
         job_def_builder = BatchJobBuilder(
             image=request.image
-            if is_valid_docker_uri(request.image)
+            if DockerImageUri.is_valid(request.image)
             else resolve_image_uri(request.image),
             job_definition_name=request.job_definition_name,
             job_name=request.job_name or f"{request.job_definition_name}-{uuid_str()}",

test/aibs_informatics_aws_lambda/handlers/batch/test_create.py

@@ -1,15 +1,34 @@
 from test.aibs_informatics_aws_lambda.base import LambdaHandlerTestCase
 
 from aibs_informatics_core.utils.hashing import uuid_str
+from pytest import mark
 
 from aibs_informatics_aws_lambda.common.api.handler import LambdaHandlerType
-from aibs_informatics_aws_lambda.handlers.batch.create import CreateDefinitionAndPrepareArgsHandler
+from aibs_informatics_aws_lambda.handlers.batch.create import (
+    CreateDefinitionAndPrepareArgsHandler,
+    DockerImageUri,
+)
 from aibs_informatics_aws_lambda.handlers.batch.model import (
     CreateDefinitionAndPrepareArgsRequest,
     CreateDefinitionAndPrepareArgsResponse,
 )
 
 
+@mark.parametrize(
+    "image, , is_valid",
+    [
+        ("docker.io/my-image:latest", True),
+        (f"docker.io/my-image@sha256:{'a' * 64}", True),
+        ("my-image:latest", False),
+        ("public.ecr.aws/my-image:latest", True),
+        ("ghcr.io/my-image:latest", True),
+        ("123456789012.dkr.ecr.us-west-2.amazonaws.com/my-image:latest", True),
+    ],
+)
+def test__DockerImageUri__is_valid(image, is_valid):
+    assert DockerImageUri.is_valid(image) == is_valid
+
+
 class CreateDefinitionAndPrepareArgsHandlerTests(LambdaHandlerTestCase):
     def setUp(self) -> None:
         super().setUp()