refactor

schedule/src/schedule/main.py

@@ -26,10 +26,7 @@ async def lifespan(app: Any) -> AsyncIterator[None]:
         session_factory=session_factory,
         redis=redis,
         object_store=build_object_store(),
-        storage_client=SchedulerStorageClient(
-            storage_http_client,
-            os.environ["INTERNAL_SERVICE_TOKEN"],
-        ),
+        storage_client=SchedulerStorageClient(storage_http_client),
         workspace_root=Path(
             os.getenv("WORKSPACE_ROOT", "/workspace/workspaces")
         ),

schedule/src/schedule/storage_client.py

@@ -1,16 +1,13 @@
+"""Schedule-specific storage client built on the shared ``StorageClient``."""
+
 from __future__ import annotations
 
-import base64
 from typing import Any
 
-import httpx
+from common.storage.client import StorageClient
 
 
-class SchedulerStorageClient:
-    def __init__(self, client: httpx.AsyncClient, service_token: str) -> None:
-        self.client = client
-        self.headers = {"X-Service-Token": service_token}
-
+class SchedulerStorageClient(StorageClient):
     async def create_object(
         self,
         *,
@@ -22,20 +19,17 @@ class SchedulerStorageClient:
         content: bytes,
         idempotency_key: str,
     ) -> dict[str, Any]:
-        response = await self.client.post(
-            "/internal/v1/objects",
-            headers=self.headers,
-            json={
-                "workspace_id": workspace_id,
-                "user_id": user_id,
-                "usage_type": usage_type,
-                "file_name": file_name,
-                "content_type": content_type,
-                "content_base64": base64.b64encode(content).decode("ascii"),
-                "visibility": "workspace",
-                "is_immutable": True,
-                "idempotency_key": idempotency_key,
-            },
+        return await self.create_server_object(
+            workspace_id=workspace_id,
+            user_id=user_id,
+            usage_type=usage_type,
+            file_name=file_name,
+            content_type=content_type,
+            content=content,
+            visibility="workspace",
+            is_immutable=True,
+            idempotency_key=idempotency_key,
         )
-        response.raise_for_status()
-        return response.json()["data"]
+
+
+__all__ = ["SchedulerStorageClient"]