diff --git a/schedule/pyproject.toml b/schedule/pyproject.toml index dfdeb12..9d46bdf 100644 --- a/schedule/pyproject.toml +++ b/schedule/pyproject.toml @@ -31,7 +31,12 @@ build-backend = "hatchling.build" [dependency-groups] dev = [ "pytest>=9.1.1", + "pytest-asyncio>=1.4.0", ] +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["tests"] + [tool.hatch.build.targets.wheel] packages = ["src/schedule"] diff --git a/schedule/src/schedule/service.py b/schedule/src/schedule/service.py index 57f4920..86c2abb 100644 --- a/schedule/src/schedule/service.py +++ b/schedule/src/schedule/service.py @@ -186,16 +186,25 @@ class SchedulerService: return await self.orchestrator.dispatch_run(run_id) -def build_object_store() -> Any: +def build_object_store(bucket_name: str | None = None) -> Any: """Construct the AsyncStorageBackend the worker reads version artifacts from. Respects ``settings.storage_backend``: in ``local`` mode points at the local ``${local_storage_base_dir}/version`` directory (the same place - ``backend.storage_api`` writes to), in ``s3`` mode points at - ``settings.s3_version_bucket``. The worker only reads version artifacts - via this store (run logs / results still go through the backend's - HTTP storage API), so pointing at the version bucket/directory is the - correct resolution regardless of the artifact's workspace. + ``backend.storage_api`` writes to), in ``s3`` mode points at the + bucket named by ``bucket_name`` (default ``settings.s3_version_bucket``). + + The bucket override exists because some workspaces configure a custom + ``Workspaces.artifact_bucket``; their version artifacts are uploaded + into that bucket by the backend, so the worker must read from the + same place — not the global version bucket. Run logs / results still + go through the backend's HTTP storage API (which honors the override + on its own), so this factory only matters for the artifact-download + path. + + For local mode the override is ignored: custom bucket subdirectories + only make sense in S3 mode, and the version base_dir is the right + destination either way. """ if settings.storage_backend == "local": return create_storage( @@ -211,7 +220,7 @@ def build_object_store() -> Any: { "type": "s3", "mode": "async", - "bucket": settings.s3_version_bucket, + "bucket": bucket_name or settings.s3_version_bucket, "endpoint_url": settings.s3_endpoint, "aws_access_key_id": settings.s3_access_key, "aws_secret_access_key": settings.s3_secret_key, diff --git a/schedule/src/schedule/worker.py b/schedule/src/schedule/worker.py index 4fd9e71..759b7dd 100644 --- a/schedule/src/schedule/worker.py +++ b/schedule/src/schedule/worker.py @@ -58,6 +58,30 @@ class NodeExecutor: self.session_factory = session_factory self.object_store = object_store self.storage_client = storage_client + # Per-bucket object store cache. The injected ``object_store`` is + # the default version-bucket store; workspaces that override + # ``Workspaces.artifact_bucket`` need a store bound to that custom + # bucket (P0-3 fix). Build lazily so the common (no-override) path + # incurs no extra cost. + self._bucket_stores: dict[str, Any] = { + settings.s3_version_bucket: object_store, + } + + def _store_for(self, bucket_name: str) -> Any: + """Return the AsyncStorageBackend bound to ``bucket_name``. + + Caches per-bucket stores on first use; the default version bucket + always reuses the injected ``object_store`` so the common path + stays zero-allocation. + """ + store = self._bucket_stores.get(bucket_name) + if store is not None: + return store + from schedule.service import build_object_store + + store = build_object_store(bucket_name=bucket_name) + self._bucket_stores[bucket_name] = store + return store async def handle_node_execute( self, @@ -380,7 +404,11 @@ class NodeExecutor: object_key: str, content_hash: str, ) -> bytes: - content = await self.object_store.get(object_key) + # Honor the artifact's actual bucket (P0-3 fix): the artifact may + # live in ``Workspaces.artifact_bucket`` rather than the global + # version bucket the default ``object_store`` is bound to. + store = self._store_for(bucket_name) + content = await store.get(object_key) if hashlib.sha256(content).hexdigest() != content_hash: raise ValueError("stable version artifact hash mismatch") logger.debug( diff --git a/schedule/tests/__init__.py b/schedule/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/schedule/tests/test_worker.py b/schedule/tests/test_worker.py new file mode 100644 index 0000000..71c6f09 --- /dev/null +++ b/schedule/tests/test_worker.py @@ -0,0 +1,154 @@ +"""Tests for NodeExecutor bucket-aware artifact download (P0-3). + +Pre-fix the worker's ``object_store`` was bound to the global version +bucket at startup, so a workspace whose ``Workspaces.artifact_bucket`` +points at a custom S3 bucket would always 404 on download. The fix +introduces a per-bucket store cache so the worker reads from whatever +bucket the artifact actually lives in. + +These tests cover only the routing logic — no live S3 / MySQL. +""" + +from __future__ import annotations + +import hashlib +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +def _make_executor() -> tuple[SimpleNamespace, MagicMock, MagicMock]: + """Construct a NodeExecutor with the bucket-store cache pre-seeded. + + Returns ``(executor, default_store, storage_client)``. The default + store is whatever the worker would receive via ``object_store=`` at + construction time; it must be returned untouched by ``_store_for`` for + the global version bucket. + """ + from schedule.worker import NodeExecutor + + default_store = MagicMock(name="default_store") + storage_client = MagicMock(name="storage_client") + executor = NodeExecutor( + session_factory=MagicMock(), + object_store=default_store, + storage_client=storage_client, + ) + return executor, default_store, storage_client + + +def test_store_for_default_bucket_returns_injected_store() -> None: + """The default version bucket must reuse the injected ``object_store``. + + Otherwise the common (no-override) path would pay for an extra + ``create_storage`` call on every run — wasted work. + """ + from common.config import settings + + executor, default_store, _ = _make_executor() + store = executor._store_for(settings.s3_version_bucket) + assert store is default_store + + +def test_store_for_custom_bucket_uses_build_factory(monkeypatch: pytest.MonkeyPatch) -> None: + """A custom bucket_name must produce a store via ``build_object_store``.""" + custom_store = MagicMock(name="custom_store") + with patch( + "schedule.service.build_object_store", + return_value=custom_store, + ) as mock_build: + executor, _, _ = _make_executor() + store = executor._store_for("my-workspace-artifacts") + + assert store is custom_store + mock_build.assert_called_once_with(bucket_name="my-workspace-artifacts") + + +def test_store_for_custom_bucket_is_cached(monkeypatch: pytest.MonkeyPatch) -> None: + """Repeated lookups for the same custom bucket must hit the cache.""" + custom_store = MagicMock(name="custom_store") + with patch( + "schedule.service.build_object_store", + return_value=custom_store, + ) as mock_build: + executor, _, _ = _make_executor() + executor._store_for("my-workspace-artifacts") + executor._store_for("my-workspace-artifacts") + executor._store_for("my-workspace-artifacts") + + mock_build.assert_called_once_with(bucket_name="my-workspace-artifacts") + + +@pytest.mark.asyncio +async def test_download_artifact_uses_per_bucket_store() -> None: + """_download_artifact must call ``store.get(object_key)`` on the + bucket-bound store, NOT the default version-bucket store. + + Pre-fix this would route every workspace's reads through the global + version bucket, silently 404'ing for any workspace that overrode + ``artifact_bucket``. + """ + custom_store = MagicMock(name="custom_store") + custom_store.get = AsyncMock(return_value=b"hello world") + executor, default_store, _ = _make_executor() + # Pretend the cache already has the custom bucket wired up. + executor._bucket_stores["my-workspace-artifacts"] = custom_store + + payload = b"hello world" + expected_hash = hashlib.sha256(payload).hexdigest() + result = await executor._download_artifact( + bucket_name="my-workspace-artifacts", + object_key="ws-1/user-1/script.py", + content_hash=expected_hash, + ) + + assert result == payload + custom_store.get.assert_awaited_once_with("ws-1/user-1/script.py") + # The default store must NOT have been touched — that was the bug. + default_store.get.assert_not_called() + + +@pytest.mark.asyncio +async def test_download_artifact_uses_default_store_when_bucket_matches() -> None: + """Sanity: default bucket reads still flow through the injected store.""" + from common.config import settings + + default_store = MagicMock(name="default_store") + default_store.get = AsyncMock(return_value=b"payload") + executor = type(_make_executor()[0])( + session_factory=MagicMock(), + object_store=default_store, + storage_client=MagicMock(), + ) + payload = b"payload" + result = await executor._download_artifact( + bucket_name=settings.s3_version_bucket, + object_key="ws-1/user-1/script.py", + content_hash=hashlib.sha256(payload).hexdigest(), + ) + + assert result == payload + default_store.get.assert_awaited_once_with("ws-1/user-1/script.py") + + +@pytest.mark.asyncio +async def test_download_artifact_rejects_hash_mismatch() -> None: + """Hash mismatch must raise regardless of which store was used — + protects against corrupted bytes landing in any bucket. + """ + custom_store = MagicMock(name="custom_store") + custom_store.get = AsyncMock(return_value=b"corrupted") + executor, default_store, _ = _make_executor() + executor._bucket_stores["my-workspace-artifacts"] = custom_store + + with pytest.raises(ValueError, match="hash mismatch"): + await executor._download_artifact( + bucket_name="my-workspace-artifacts", + object_key="ws-1/user-1/script.py", + content_hash="0" * 64, # cannot match corrupted bytes + ) + + # Reads still went through the bucket-bound store, not the default. + custom_store.get.assert_awaited_once() + default_store.get.assert_not_called() diff --git a/uv.lock b/uv.lock index 282f753..766a613 100644 --- a/uv.lock +++ b/uv.lock @@ -2585,6 +2585,7 @@ dependencies = [ [package.dev-dependencies] dev = [ { name = "pytest" }, + { name = "pytest-asyncio" }, ] [package.metadata] @@ -2603,7 +2604,10 @@ requires-dist = [ ] [package.metadata.requires-dev] -dev = [{ name = "pytest", specifier = ">=9.1.1" }] +dev = [ + { name = "pytest", specifier = ">=9.1.1" }, + { name = "pytest-asyncio", specifier = ">=1.4.0" }, +] [[package]] name = "send2trash"