connectors/docker: add platform and architecture option via connector data

Signed-off-by: JP-Ellis <[email protected]>
Co-authored-by: Nick Mills-Barrett <[email protected]>

Author: JP-Ellis

src/pyinfra/connectors/docker.py

@@ -26,10 +26,14 @@
 
 class ConnectorData(TypedDict):
     docker_identifier: str
+    docker_platform: str
+    docker_architecture: str
 
 
 connector_data_meta: dict[str, DataMeta] = {
     "docker_identifier": DataMeta("ID of container or image to start from"),
+    "docker_platform": DataMeta("Platform to use for Docker image (e.g., linux/amd64)"),
+    "docker_architecture": DataMeta("Architecture to use for Docker image (e.g., amd64, arm64)"),
 }
 
 
@@ -108,9 +112,29 @@ def _find_start_docker_container(self, container_id) -> tuple[str, bool]:
         return container_id, True
 
     def _start_docker_image(self, image_name):
+        docker_cmd_parts = [
+            self.docker_cmd,
+            "run",
+            "-d",
+        ]
+
+        if self.data.get("docker_platform"):
+            docker_cmd_parts.extend(["--platform", self.data["docker_platform"]])
+        if self.data.get("docker_architecture"):
+            docker_cmd_parts.extend(["--arch", self.data["docker_architecture"]])
+
+        docker_cmd_parts.extend(
+            [
+                image_name,
+                "tail",
+                "-f",
+                "/dev/null",
+            ]
+        )
+
         try:
             return local.shell(
-                f"{self.docker_cmd} run -d {image_name} tail -f /dev/null",
+                " ".join(docker_cmd_parts),
                 splitlines=True,
             )[-1]  # last line is the container ID
         except PyinfraError as e:

tests/test_connectors/test_docker.py

@@ -180,6 +180,44 @@ def test_get_file_error(self):
         with self.assertRaises(IOError):
             host.get_file("not-a-file", "not-another-file", print_output=True)
 
+    def test_docker_platform(self):
+        inventory = make_inventory(
+            hosts=((f"@{self.connector_name}/not-an-image", {"docker_platform": "linux/amd64"}),),
+        )
+        State(inventory, Config())
+
+        host = inventory.get_host(f"@{self.connector_name}/not-an-image")
+        assert host.data.docker_platform == "linux/amd64"
+
+    def test_docker_architecture(self):
+        inventory = make_inventory(
+            hosts=((f"@{self.connector_name}/not-an-image", {"docker_architecture": "arm64"}),),
+        )
+        State(inventory, Config())
+
+        host = inventory.get_host(f"@{self.connector_name}/not-an-image")
+        assert host.data.docker_architecture == "arm64"
+
+    def test_connect_with_docker_platform(self):
+        inventory = make_inventory(
+            hosts=((f"@{self.connector_name}/not-an-image", {"docker_platform": "linux/amd64"}),),
+        )
+        state = State(inventory, Config())
+        host = inventory.get_host(f"@{self.connector_name}/not-an-image")
+        host.connect(reason=True)
+        assert len(state.active_hosts) == 0
+        host.disconnect()
+
+    def test_connect_with_docker_architecture(self):
+        inventory = make_inventory(
+            hosts=((f"@{self.connector_name}/not-an-image", {"docker_architecture": "arm64"}),),
+        )
+        state = State(inventory, Config())
+        host = inventory.get_host(f"@{self.connector_name}/not-an-image")
+        host.connect(reason=True)
+        assert len(state.active_hosts) == 0
+        host.disconnect()
+
 
 # Reuse the container testing code for docker and podman
 
@@ -188,6 +226,12 @@ def fake_docker_shell(command, splitlines=None):
     if command == "docker run -d not-an-image tail -f /dev/null":
         return ["containerid"]
 
+    if command == "docker run -d --platform linux/amd64 not-an-image tail -f /dev/null":
+        return ["containerid"]
+
+    if command == "docker run -d --arch arm64 not-an-image tail -f /dev/null":
+        return ["containerid"]
+
     if command == "docker commit containerid":
         return ["sha256:blahsomerandomstringdata"]
 
@@ -213,6 +257,12 @@ def fake_podman_shell(command, splitlines=None):
     if command == "podman run -d not-an-image tail -f /dev/null":
         return ["containerid"]
 
+    if command == "podman run -d --platform linux/amd64 not-an-image tail -f /dev/null":
+        return ["containerid"]
+
+    if command == "podman run -d --arch arm64 not-an-image tail -f /dev/null":
+        return ["containerid"]
+
     if command == "podman commit containerid":
         return ["sha256:blahsomerandomstringdata"]