From eddbb21e3fddd8183ab3cf3797bd5ed791fcf5ce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=98yvind=20Sigerstad?= <newsflash@gmail.com>
Date: Wed, 18 Dec 2024 04:53:39 +0100
Subject: [PATCH] feat: add KeepLastestNImages rule (#153)

* feat: add KeepLastestNImages rule

* fix: remove unwanted diff

* fix: remove unused import
---
 README.md                           |  8 ++++
 artifactory_cleanup/rules/docker.py | 27 +++++++++++
 tests/data/all-built-in-rules.yaml  |  2 +
 tests/test_rules_docker.py          | 69 +++++++++++++++++++++++++++++
 4 files changed, 106 insertions(+)

diff --git a/README.md b/README.md
index 37d3f85..c9d7210 100644
--- a/README.md
+++ b/README.md
@@ -350,6 +350,14 @@ policies:
   custom_regexp: "[^\\d][\\._]((\\d+\\.)+\\d+)"
 ```
 
+- `KeepLatestNDockerImages(count=N)` - Leaves N
+  most recently updated Docker image digests. This ensures all tags matching the same digest is kept.
+
+```yaml
+- rule: KeepLatestNDockerImages
+  count: 1
+```
+
 - `DeleteDockerImageIfNotContainedInProperties(docker_repo='docker-local', properties_prefix='my-prop', image_prefix=None, full_docker_repo_name=None)`
   \- Remove Docker image, if it is not found in the properties of the artifact repository.
 
diff --git a/artifactory_cleanup/rules/docker.py b/artifactory_cleanup/rules/docker.py
index 40c54c0..4bf2682 100644
--- a/artifactory_cleanup/rules/docker.py
+++ b/artifactory_cleanup/rules/docker.py
@@ -191,6 +191,33 @@ class ExcludeDockerImages(FilterDockerImages):
     boolean_operator = "$and"
 
 
+class KeepLatestNDockerImages(RuleForDocker):
+    """
+    Leaves ``count`` Docker image digests for each image. This allows tags that have the same digest to be kept.
+    """
+
+    def __init__(self, count: int):
+        self.count = count
+
+    def filter(self, artifacts):
+        artifacts = self._manifest_to_docker_images(artifacts)
+        artifacts_by_path = defaultdict(list)
+
+        for artifact in artifacts:
+            path = artifact["path"]
+            artifacts_by_path[path].append(artifact)
+
+        for path, _artifacts in artifacts_by_path.items():
+            sha256s_to_keep = set()
+            _artifacts.sort(reverse=True, key=lambda x: x['updated'])
+            for artifact in _artifacts:
+                if len(sha256s_to_keep) < self.count:
+                    sha256s_to_keep.add(artifact['sha256'])
+                if artifact['sha256'] in sha256s_to_keep:
+                    artifacts.keep(artifact)
+
+        return artifacts
+
 class KeepLatestNVersionImagesByProperty(RuleForDocker):
     r"""
     Leaves ``count`` Docker images with the same major.
diff --git a/tests/data/all-built-in-rules.yaml b/tests/data/all-built-in-rules.yaml
index 5a7151e..fd8eaf8 100644
--- a/tests/data/all-built-in-rules.yaml
+++ b/tests/data/all-built-in-rules.yaml
@@ -97,6 +97,8 @@ artifactory-cleanup:
           masks:
             - "*production*"
             - "*release*"
+        - rule: KeepLatestNDockerImages
+          count: 1
         - rule: KeepLatestNVersionImagesByProperty
           count: 1
           custom_regexp: "[^\\d][\\._]((\\d+\\.)+\\d+)"
diff --git a/tests/test_rules_docker.py b/tests/test_rules_docker.py
index 0bf1666..e590fe2 100644
--- a/tests/test_rules_docker.py
+++ b/tests/test_rules_docker.py
@@ -1,5 +1,6 @@
 from artifactory_cleanup import CleanupPolicy
 from artifactory_cleanup.rules import (
+    KeepLatestNDockerImages,
     KeepLatestNVersionImagesByProperty,
     ArtifactsList,
     RuleForDocker,
@@ -7,6 +8,74 @@
 )
 
 
+class TestKeepLatestNDockerImages:
+    def test_filter(self):
+        data = [
+            {
+                "path": "foobar/latest",
+                "sha256": "9908cd35ccb5774c0da1c8bae657c48fe42f865a62fded44043fcfe2a09f3e31",
+                "name": "manifest.json",
+                "updated": "2021-03-20T13:54:52.383+02:00",
+                "properties": {"docker.manifest": "v0.1.99"},
+            },
+            {
+                "path": "foobar/2021-03-19T13-53-52.383",
+                "sha256": "d26a1bf07bd081a5b389d8402df8fec682267d0b02276768b8a9b2bb6bd149a6",
+                "name": "manifest.json",
+                "updated": "2021-03-19T13:53:52.383+02:00",
+                "properties": {"docker.manifest": "v0.1.99"},
+            },
+            {
+                "path": "foobar/sha256__9908cd35ccb5774c0da1c8bae657c48fe42f865a62fded44043fcfe2a09f3e31",
+                "sha256": "9908cd35ccb5774c0da1c8bae657c48fe42f865a62fded44043fcfe2a09f3e31",
+                "name": "manifest.json",
+                "updated": "2021-03-19T13:52:52.383+02:00",
+                "properties": {"docker.manifest": "v0.1.99"},
+            },
+            {
+                "path": "baz/latest",
+                "sha256": "42988e2a52ad999d0038b46a7528f6526e4b9e2093f0bf7522eb61ac316715d3",
+                "name": "manifest.json",
+                "updated": "2021-03-20T13:54:52.383+02:00",
+                "properties": {"docker.manifest": "v0.1.99"},
+            },
+            {
+                "path": "qux/0.0.1",
+                "sha256": "1fbffb7bb96039fae4a89ddd2cbac16b285fac333bc928d6464665e953828054",
+                "name": "manifest.json",
+                "updated": "2021-03-20T13:54:52.383+02:00",
+                "properties": {"docker.manifest": "v0.1.99"},
+            },
+            {
+                "path": "qux/0.0.2",
+                "sha256": "46c469dbbff12818441d22aa8fed36869e71f3f9a7ec317d57219744e1688e25",
+                "name": "manifest.json",
+                "updated": "2021-03-20T13:53:52.383+02:00",
+                "properties": {"docker.manifest": "v0.1.99"},
+            },
+        ]   
+
+        artifacts = ArtifactsList.from_response(data)
+        policy = CleanupPolicy("test", KeepLatestNDockerImages(count=1))
+        assert policy.filter(artifacts) == [
+            {
+                "path": "foobar",
+                "sha256": "d26a1bf07bd081a5b389d8402df8fec682267d0b02276768b8a9b2bb6bd149a6",
+                "name": "2021-03-19T13-53-52.383",
+                "updated": "2021-03-19T13:53:52.383+02:00",
+                "properties": {"docker.manifest": "v0.1.99"},
+                "stats": {},
+            },
+            {
+                "path": "qux",
+                "sha256": "46c469dbbff12818441d22aa8fed36869e71f3f9a7ec317d57219744e1688e25",
+                "name": "0.0.2",
+                "updated": "2021-03-20T13:53:52.383+02:00",
+                "properties": {"docker.manifest": "v0.1.99"},
+                "stats": {},
+            },
+        ]
+
 class TestKeepLatestNVersionImagesByProperty:
     def test_filter(self):
         # Skip collecting docker size