From a6d8be1eb5db2dcae2527b8e07cc71e993d8c7ec Mon Sep 17 00:00:00 2001
From: William Chu <william.chu@uptickhq.com>
Date: Tue, 22 Oct 2024 16:25:06 +1100
Subject: [PATCH] feat(cli): efficiently find latest image using prefix search
 for latest image tag

---
 gitops/common/app.py   |  5 +++
 gitops/main.py         |  4 +++
 gitops/utils/images.py | 77 ++++++++++++++++++++++++++++++++----------
 3 files changed, 69 insertions(+), 17 deletions(-)

diff --git a/gitops/common/app.py b/gitops/common/app.py
index df95fb8..eb5ea98 100644
--- a/gitops/common/app.py
+++ b/gitops/common/app.py
@@ -113,6 +113,11 @@ def image_repository_name(self) -> str:
         """305686791668.dkr.ecr.ap-southeast-2.amazonaws.com/[uptick]:yoink-9f03ac80f3"""
         return self.image.split(":")[0].split("/")[-1]
 
+    @property
+    def image_repository(self) -> str:
+        """[305686791668.dkr.ecr.ap-southeast-2.amazonaws.com]/uptick:yoink-9f03ac80f3"""
+        return self.image.split("/")[0]
+
     @property
     def image_tag(self) -> str:
         """305686791668.dkr.ecr.ap-southeast-2.amazonaws.com/uptick:[yoink-9f03ac80f3]"""
diff --git a/gitops/main.py b/gitops/main.py
index fda4a7e..84f1f92 100644
--- a/gitops/main.py
+++ b/gitops/main.py
@@ -1,7 +1,11 @@
+import logging
+
 from invoke import Collection, Program, Task
 
 from . import __version__, core, shorthands
 
+logging.basicConfig(level=logging.INFO, format="%(message)s")
+
 version = __version__
 
 namespace = Collection()
diff --git a/gitops/utils/images.py b/gitops/utils/images.py
index f7d137c..ee48e9c 100644
--- a/gitops/utils/images.py
+++ b/gitops/utils/images.py
@@ -1,13 +1,22 @@
+import datetime
+import logging
 from functools import lru_cache
 from hashlib import md5
+from typing import TYPE_CHECKING
 
 import boto3
+import botocore
+import botocore.exceptions
 from colorama import Fore
 
 from .cli import colourise
 
+if TYPE_CHECKING:
+    from mypy_boto3_ecr.type_defs import ImageDetailTypeDef  # type: ignore
 BATCH_SIZE = 100
 
+logger = logging.getLogger(__name__)
+
 
 def get_image(tag: str) -> str:
     """Finds a specific image in ECR."""
@@ -16,30 +25,64 @@ def get_image(tag: str) -> str:
 
 
 @lru_cache
-def get_latest_image(repository_name: str, prefix: str) -> str | None:
-    """Finds latest image in ECR with the given prefix and returns the image tag"""
-    ecr_client = boto3.client("ecr")
+def get_latest_image(repository_name: str, prefix: str, ecr_repository: str | None = None) -> str | None:  # noqa:C901
+    """Finds latest image in ECR with the given prefix and returns the image tag
+
+    param ecr_repository is expected in this format: 305686791668.dkr.ecr.ap-southeast-2.amazonaws.com
+    """
+    describe_image_args = {}
+    region_name = None
+    if ecr_repository:
+        account_id = ecr_repository.split(".")[0]
+        describe_image_args["registryId"] = account_id
+        region_name = ecr_repository.split(".")[3]
+
+    ecr_client = boto3.client("ecr", region_name=region_name)
     client_paginator = ecr_client.get_paginator("describe_images")
 
     results = []
-    for ecr_response in client_paginator.paginate(
-        repositoryName=repository_name,
-        filter={"tagStatus": "TAGGED"},
-        maxResults=BATCH_SIZE,
-    ):
-        for image in ecr_response["imageDetails"]:
-            if prefix != "":
-                if prefix_tags := [tag for tag in image["imageTags"] if tag.startswith(prefix + "-")]:
-                    results.append((prefix_tags[0], image["imagePushedAt"]))
-            else:
-                if prefix_tags := [tag for tag in image["imageTags"] if "-" not in tag]:
-                    results.append((prefix_tags[0], image["imagePushedAt"]))
 
+    # First we try to find the image with `*latest`
+    image_tag = f"{prefix}-latest" if prefix else "latest"
+
+    def add_image_to_results(image: "ImageDetailTypeDef") -> None:
+        if prefix != "":
+            if prefix_tags := [
+                tag for tag in image["imageTags"] if tag.startswith(prefix + "-") and "latest" not in tag
+            ]:
+                results.append((prefix_tags[0], image["imagePushedAt"]))
+        else:
+            if prefix_tags := [tag for tag in image["imageTags"] if "-" not in tag and "latest" not in tag]:
+                results.append((prefix_tags[0], image["imagePushedAt"]))
+
+    try:
+        image = ecr_client.describe_images(
+            repositoryName=repository_name,
+            imageIds=[{"imageTag": image_tag}],
+            **describe_image_args,  # type: ignore
+        )["imageDetails"][0]
+        # TODO: Remove this check after we've fully migrated to the new image tagging scheme
+        if image["imagePushedAt"] < datetime.datetime(2024, 9, 1, 0, 0, 0, tzinfo=datetime.timezone.utc):
+            logger.warning(f"Image {image_tag} is too old to be considered latest.")
+            raise ValueError("Image is too old")
+
+        add_image_to_results(image)
+    except (botocore.exceptions.ClientError, ValueError):
+        # Ok we couldn't find the -latest image; lets scan images manually
+        for ecr_response in client_paginator.paginate(
+            repositoryName=repository_name,
+            filter={
+                "tagStatus": "TAGGED",
+            },
+            **describe_image_args,  # type: ignore
+        ):
+            for image in ecr_response["imageDetails"]:
+                add_image_to_results(image)
     if not results:
         if prefix:
-            print(f'No images found in repository: {repository_name} with tag "{prefix}-*".')
+            logger.info(f'No images found in repository: {repository_name} with tag "{prefix}-*".')
         else:
-            print(f"No images found in repository: {repository_name}")
+            logger.info(f"No images found in repository: {repository_name}")
         return None
 
     latest_image_tag = sorted(results, key=lambda image: image[1], reverse=True)[0][0]