[prefect-docker] add cached docker build and push steps

src/integrations/prefect-docker/prefect_docker/deployments/steps.py

@@ -24,8 +24,10 @@
     ```
 """
 
+import json
 import os
 import sys
+from functools import wraps
 from pathlib import Path
 from typing import Dict, List, Optional
 
@@ -59,7 +61,7 @@ class BuildDockerImageResult(TypedDict):
     tag: str
     image: str
     image_id: str
-    additional_tags: Optional[str]
+    additional_tags: Optional[List[str]]
 
 
 class PushDockerImageResult(TypedDict):
@@ -74,9 +76,9 @@ class PushDockerImageResult(TypedDict):
     """
 
     image_name: str
-    tag: str
+    tag: Optional[str]
     image: str
-    additional_tags: Optional[str]
+    additional_tags: Optional[List[str]]
 
 
 def build_docker_image(
@@ -328,3 +330,115 @@ def push_docker_image(
         "image": f"{image_name}:{tag}",
         "additional_tags": additional_tags,
     }
+
+
+def _make_hashable(obj):
+    if isinstance(obj, dict):
+        return json.dumps(obj, sort_keys=True)
+    elif isinstance(obj, list):
+        return tuple(_make_hashable(v) for v in obj)
+    return obj
+
+
+def cacheable(func):
+    cache = {}
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        key = (
+            tuple(_make_hashable(arg) for arg in args),
+            tuple((k, _make_hashable(v)) for k, v in sorted(kwargs.items())),
+        )
+        if key not in cache:
+            cache[key] = func(*args, **kwargs)
+        return cache[key]
+
+    return wrapper
+
+
+@cacheable
+def cached_build_docker_image(
+    image_name: str,
+    dockerfile: str,
+    tag: Optional[str] = None,
+    additional_tags: Optional[List[str]] = None,
+    **build_kwargs,
+) -> BuildDockerImageResult:
+    """Cached version of `prefect_docker.deployments.steps.build_docker_image`.
+
+    Args:
+        image_name: The name of the Docker image to build, including the registry and
+            repository.
+        dockerfile: The path to the Dockerfile used to build the image. If "auto" is
+            passed, a temporary Dockerfile will be created to build the image.
+        tag: The tag to apply to the built image.
+        additional_tags: Additional tags on the image, in addition to `tag`, to apply to the built image.
+        **build_kwargs: Additional keyword arguments to pass to Docker when building
+            the image. Available options can be found in the [`docker-py`](https://docker-py.readthedocs.io/en/stable/images.html#docker.models.images.ImageCollection.build)
+            documentation.
+
+    Returns:
+        A dictionary containing the image name and tag of the built image.
+
+    Example:
+        Build a Docker image that is cached if the same image is built again in the same `prefect deploy` run:
+        ```yaml
+        build:
+            - prefect_docker.deployments.steps.cached_build_docker_image:
+                id: build-image
+                requires: prefect-docker
+                image_name: repo-name/image-name
+                tag: dev
+        ```
+    """
+    return build_docker_image(
+        image_name=image_name,
+        dockerfile=dockerfile,
+        tag=tag,
+        additional_tags=additional_tags,
+        **build_kwargs,
+    )
+
+
+@cacheable
+def cached_push_docker_image(
+    image_name: str,
+    tag: str,
+    credentials: Optional[Dict] = None,
+    additional_tags: Optional[List[str]] = None,
+    **push_kwargs,
+) -> PushDockerImageResult:
+    """Cached version of `prefect_docker.deployments.steps.push_docker_image`.
+
+    Args:
+        image_name: The name of the Docker image to push, including the registry and
+            repository.
+        tag: The tag of the Docker image to push.
+        credentials: A dictionary containing the username, password, and URL for the
+            registry to push the image to.
+        additional_tags: Additional tags on the image, in addition to `tag`, to apply to the built image.
+        **push_kwargs: Additional keyword arguments to pass to Docker when pushing
+            the image. Available options can be found in the [`docker-py`](https://docker-py.readthedocs.io/en/stable/images.html#docker.models.images.ImageCollection.push)
+            documentation.
+
+    Returns:
+        A dictionary containing the image name and tag of the pushed image.
+
+    Example:
+        Push a Docker image that is cached if the same image is pushed again in the same `prefect deploy` run:
+        ```yaml
+        push:
+            - prefect_docker.deployments.steps.cached_push_docker_image:
+                requires: prefect-docker
+                image_name: repo-name/image-name
+                tag: dev
+                credentials: "{{ prefect.blocks.docker-registry-credentials.dev-registry }}"
+        ```
+    """
+    return push_docker_image(
+        image_name=image_name,
+        tag=tag,
+        credentials=credentials,
+        additional_tags=additional_tags,
+        **push_kwargs,
+    )

src/integrations/prefect-docker/tests/deployments/test_steps.py

@@ -2,15 +2,20 @@
 import sys
 from pathlib import Path
 from unittest import mock
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import docker
 import docker.errors
 import docker.models.containers
 import docker.models.images
 import pendulum
 import pytest
-from prefect_docker.deployments.steps import build_docker_image, push_docker_image
+from prefect_docker.deployments.steps import (
+    build_docker_image,
+    cached_build_docker_image,
+    cached_push_docker_image,
+    push_docker_image,
+)
 
 import prefect
 import prefect.utilities.dockerutils
@@ -364,3 +369,74 @@ def test_push_docker_image_raises_on_event_error(mock_docker_client):
         push_docker_image(
             image_name=FAKE_IMAGE_NAME, tag=FAKE_TAG, credentials=FAKE_CREDENTIALS
         )
+
+
+def test_cached_build_docker_image(mock_docker_client):
+    image_name = "registry/repo"
+    dockerfile = "Dockerfile"
+    tag = "mytag"
+    additional_tags = ["tag1", "tag2"]
+    expected_result = {
+        "image": f"{image_name}:{tag}",
+        "tag": tag,
+        "image_name": image_name,
+        "image_id": FAKE_CONTAINER_ID,
+        "additional_tags": additional_tags,
+    }
+
+    with patch(
+        "prefect_docker.deployments.steps.build_docker_image"
+    ) as mock_build_docker_image:
+        mock_build_docker_image.return_value = expected_result
+
+        # Call the cached function multiple times with the same arguments
+        for _ in range(3):
+            result = cached_build_docker_image(
+                image_name=image_name,
+                dockerfile=dockerfile,
+                tag=tag,
+                additional_tags=additional_tags,
+            )
+            assert result == expected_result
+
+        mock_build_docker_image.assert_called_once_with(
+            image_name=image_name,
+            dockerfile=dockerfile,
+            tag=tag,
+            additional_tags=additional_tags,
+        )
+
+
+def test_cached_push_docker_image(mock_docker_client):
+    image_name = FAKE_IMAGE_NAME
+    tag = FAKE_TAG
+    credentials = FAKE_CREDENTIALS
+    additional_tags = FAKE_ADDITIONAL_TAGS
+    expected_result = {
+        "image_name": image_name,
+        "tag": tag,
+        "image": f"{image_name}:{tag}",
+        "additional_tags": additional_tags,
+    }
+
+    with patch(
+        "prefect_docker.deployments.steps.push_docker_image"
+    ) as mock_push_docker_image:
+        mock_push_docker_image.return_value = expected_result
+
+        # Call the cached function multiple times with the same arguments
+        for _ in range(3):
+            result = cached_push_docker_image(
+                image_name=image_name,
+                tag=tag,
+                credentials=credentials,
+                additional_tags=additional_tags,
+            )
+            assert result == expected_result
+
+        mock_push_docker_image.assert_called_once_with(
+            image_name=image_name,
+            tag=tag,
+            credentials=credentials,
+            additional_tags=additional_tags,
+        )