Skip to content

Commit

Permalink
Fix docker worker volumes (#14516)
Browse files Browse the repository at this point in the history
  • Loading branch information
zzstoatzz authored Jul 8, 2024
1 parent 7bf536f commit baeeed6
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 1 deletion.
27 changes: 26 additions & 1 deletion src/integrations/prefect-docker/prefect_docker/worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,32 @@ class ImagePullPolicy(enum.Enum):
NEVER = "Never"


VolumeStr = Annotated[str, AfterValidator(lambda v: ":" in v)]
def assert_volume_str(v: str) -> str:
"""Assert that a string is a valid Docker volume string."""
if not isinstance(v, str):
raise ValueError("Volume must be a string")

# Regex pattern for valid Docker volume strings, including Windows paths
pattern = r"^([a-zA-Z]:\\|/)?[^:]+:(/)?[^:]+(:ro|:rw)?$"

match = re.match(pattern, v)
if not match:
raise ValueError(f"Invalid volume string: {v!r}")

_, _, mode = match.groups()

# Check for empty parts
if ":" not in v or v.startswith(":") or v.endswith(":"):
raise ValueError(f"Volume string contains empty part: {v!r}")

# If there's a mode, it must be 'ro' or 'rw'
if mode and mode not in (":ro", ":rw"):
raise ValueError(f"Invalid volume mode: {mode!r}. Must be ':ro' or ':rw'")

return v


VolumeStr = Annotated[str, AfterValidator(assert_volume_str)]


class DockerWorkerJobConfiguration(BaseJobConfiguration):
Expand Down
37 changes: 37 additions & 0 deletions src/integrations/prefect-docker/tests/test_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,9 @@
CONTAINER_LABELS,
DockerWorker,
DockerWorkerJobConfiguration,
VolumeStr,
)
from pydantic import TypeAdapter, ValidationError

from prefect.client.schemas import FlowRun
from prefect.events import RelatedResource
Expand Down Expand Up @@ -262,6 +264,41 @@ async def test_uses_volumes_setting(
assert "c:d" in call_volumes


@pytest.mark.parametrize(
"volume_str",
[
"a:b",
"/host/path:/container/path",
"named_volume:/app/data",
"/home/user:/home/docker:ro",
"/home/user:/home/docker:rw",
"C:\\path\\on\\windows:/path/in/container",
"\\\\host\\share:/path/in/container",
],
)
def test_valid_volume_strings(volume_str):
assert TypeAdapter(VolumeStr).validate_python(volume_str) == volume_str


@pytest.mark.parametrize(
"volume_str",
[
"invalid_volume",
":missing_host",
"missing_container:",
"/double:/colon:/path",
"/path:/path:invalid_mode",
":/:",
" : : ",
"/host:/container:rw:extra",
"", # empty string
],
)
def test_invalid_volume_strings(volume_str):
with pytest.raises(ValidationError, match="Invalid volume"):
TypeAdapter(VolumeStr).validate_python(volume_str)


async def test_uses_privileged_setting(
mock_docker_client, flow_run, default_docker_worker_job_configuration
):
Expand Down

0 comments on commit baeeed6

Please sign in to comment.