tao.chen/model-platform
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'optimize-by-taochen' into develop

Browse Source 
# Conflicts:
#	.env.example
#	backend/src/backend/main.py
#	common/src/common/db/models.py
#	common/src/common/eventing.py
#	docker-compose.yml
#	runtime/Dockerfile
#	runtime/pyproject.toml
#	runtime/src/runtime/main.py
#	schedule/src/schedule/main.py
#	schedule/src/schedule/service.py
This commit is contained in:
tao.chen
2026-07-31 09:29:03 +08:00
parent c32e6e4cae 91767461d6
commit 62c7d7465e
59 changed files with 3016 additions and 4423 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/src/backend/dependencies.py
+4 -8
View File
@@ -3,10 +3,11 @@ from __future__ import annotations
from dataclasses import dataclass
from typing import AsyncIterator
from fastapi import Header, HTTPException, Request, status
from fastapi import Depends, Header, HTTPException, Request, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from common.db import session_scope
from common.db.models import (
Roles,
Users,
@@ -29,13 +30,8 @@ class RequestContext:
async def database_session(request: Request) -> AsyncIterator[AsyncSession]:
async with request.app.state.session_factory() as session:
try:
yield session
await session.commit()
except Exception:
await session.rollback()
raise
async with session_scope(request.app.state.session_factory) as session:
yield session
async def request_context(
backend/src/backend/file_locks.py
-110
View File
@@ -1,110 +0,0 @@
from __future__ import annotations
from typing import Any, Awaitable, Callable
from fastapi import APIRouter, Depends, Request, status
from fastapi.responses import JSONResponse
from backend.dependencies import (
RequestContext,
request_context,
)
from backend.runtime_client import RuntimeClientError
from backend.schemas import FileLockTokenRequest
router = APIRouter(tags=["file-locks"])
async def runtime_response(
context: RequestContext,
operation: Callable[[], Awaitable[dict[str, Any]]],
*,
success_status: int = status.HTTP_200_OK,
) -> JSONResponse:
try:
data = await operation()
except RuntimeClientError as exc:
error = exc.detail
if not isinstance(error, dict) or "code" not in error:
error = {
"code": "RUNTIME_REQUEST_FAILED",
"message": str(error),
"retryable": exc.status_code >= 500,
"details": {},
}
return JSONResponse(
status_code=exc.status_code,
content={"request_id": context.request_id, "error": error},
)
return JSONResponse(
status_code=success_status,
content={
"request_id": context.request_id,
"data": data,
"meta": {},
},
)
@router.post(
"/api/v1/files/{storage_object_id}/lock",
status_code=status.HTTP_201_CREATED,
)
async def acquire_file_lock(
storage_object_id: str,
request: Request,
context: RequestContext = Depends(request_context),
) -> JSONResponse:
return await runtime_response(
context,
lambda: request.app.state.runtime_client.acquire_file_lock(
{
"workspace_id": context.workspace.workspace_id,
"storage_object_id": storage_object_id,
"user_id": context.user.user_id,
"request_id": context.request_id,
}
),
success_status=status.HTTP_201_CREATED,
)
@router.post("/api/v1/file-locks/{edit_session_id}/heartbeat")
async def heartbeat_file_lock(
edit_session_id: str,
payload: FileLockTokenRequest,
request: Request,
context: RequestContext = Depends(request_context),
) -> JSONResponse:
return await runtime_response(
context,
lambda: request.app.state.runtime_client.heartbeat_file_lock(
edit_session_id,
{
"workspace_id": context.workspace.workspace_id,
"user_id": context.user.user_id,
"lock_token": payload.lock_token,
},
),
)
@router.delete("/api/v1/file-locks/{edit_session_id}")
async def release_file_lock(
edit_session_id: str,
payload: FileLockTokenRequest,
request: Request,
context: RequestContext = Depends(request_context),
) -> JSONResponse:
return await runtime_response(
context,
lambda: request.app.state.runtime_client.release_file_lock(
edit_session_id,
{
"workspace_id": context.workspace.workspace_id,
"user_id": context.user.user_id,
"lock_token": payload.lock_token,
},
),
)
backend/src/backend/jupyter.py
+221 -66
View File
@@ -1,84 +1,239 @@
from __future__ import annotations
import base64
import hashlib
import hmac
import json
import os
import re
import time
from typing import Optional
from fastapi import APIRouter, Depends, Request, status
from fastapi.responses import JSONResponse
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from sqlalchemy import select
from backend.dependencies import (
RequestContext,
request_context,
)
from backend.dependencies import database_session
from backend.runtime_client import RuntimeClientError
from backend.schemas import (
CreateJupyterAccessTicketRequest,
)
from common.db.models import Scripts, Users, WorkspaceMembers, Workspaces
from sqlalchemy.ext.asyncio import AsyncSession
router = APIRouter(tags=["jupyter"])
security = HTTPBearer(auto_error=False)
def cookie_secure() -> bool:
return os.getenv("COOKIE_SECURE", "false").strip().lower() in {
"1",
"true",
"yes",
"on",
}
JWT_SECRET = os.environ.get("JWT_SECRET", "dev-only-not-for-production")
JWT_ALGORITHM = "HS256"
@router.post(
"/api/v1/jupyter/access-tickets",
status_code=status.HTTP_201_CREATED,
)
async def create_jupyter_access_ticket(
payload: CreateJupyterAccessTicketRequest,
request: Request,
context: RequestContext = Depends(request_context),
) -> JSONResponse:
def _b64decode(value: str) -> bytes:
padding = "=" * (-len(value) % 4)
return base64.urlsafe_b64decode(value + padding)
def verify_jwt_token(token: str) -> dict:
"""Parse a signed JWT from the access_token cookie.
Returns a payload dict containing at least ``sub`` (user_id) and
``exp``; raises 401 on missing, malformed, or expired tokens. The
current implementation is intentionally minimal — replace with a
real verification once a public-key/HSM source is wired in.
"""
if not token:
raise HTTPException(
status_code=401,
detail="Missing Authentication Token",
)
try:
ticket_data = (
await request.app.state.runtime_client
.create_jupyter_access_ticket(
payload.edit_session_id,
{
"workspace_id": context.workspace.workspace_id,
"user_id": context.user.user_id,
"lock_token": payload.lock_token,
},
)
)
except RuntimeClientError as exc:
error = exc.detail
if not isinstance(error, dict) or "code" not in error:
error = {
"code": "JUPYTER_TICKET_FAILED",
"message": str(error),
"retryable": exc.status_code >= 500,
"details": {},
}
return JSONResponse(
status_code=exc.status_code,
content={"request_id": context.request_id, "error": error},
header_b64, payload_b64, signature_b64 = token.split(".", 2)
except ValueError as exc:
raise HTTPException(
status_code=401,
detail="Invalid Authentication Token",
) from exc
signing_input = f"{header_b64}.{payload_b64}".encode()
expected = hmac.new(
JWT_SECRET.encode(),
signing_input,
hashlib.sha256,
).digest()
try:
signature = _b64decode(signature_b64)
except Exception as exc: # pragma: no cover - malformed b64
raise HTTPException(
status_code=401,
detail="Invalid Authentication Token",
) from exc
if not hmac.compare_digest(expected, signature):
raise HTTPException(
status_code=401,
detail="Invalid Authentication Token",
)
raw_ticket = ticket_data.pop("ticket")
max_age = int(ticket_data.pop("expires_in_seconds"))
response = JSONResponse(
status_code=status.HTTP_201_CREATED,
content={
"request_id": context.request_id,
"data": ticket_data,
"meta": {},
},
try:
payload = json.loads(_b64decode(payload_b64))
except Exception as exc:
raise HTTPException(
status_code=401,
detail="Invalid Authentication Token",
) from exc
exp = payload.get("exp")
if not isinstance(exp, (int, float)) or exp < time.time():
raise HTTPException(
status_code=401,
detail="Token expired",
)
return payload
def extract_notebook_path(
uri: str,
workspace_id: str,
) -> Optional[str]:
"""Pull the relative notebook path out of the original request URI.
Only ``/jupyter/{workspace_id}/notebooks/*.ipynb`` requests are
subject to file-level lock checks; everything else (tree views,
``/api/contents`` and WebSocket upgrades) bypasses the lock.
"""
pattern = rf"^/jupyter/{re.escape(workspace_id)}/notebooks/(.+\.ipynb)"
match = re.match(pattern, uri)
if match:
return match.group(1)
return None
async def check_notebook_is_locked(
session: AsyncSession,
workspace_id: str,
notebook_path: str,
user_id: str,
) -> bool:
"""Decide whether the request is allowed to read the notebook.
Returns ``True`` only when the notebook exists, is owned by a
different user, and is currently locked. The owner is always let
through; a missing row is treated as "not owned yet" and allowed.
"""
statement = select(Scripts.owner_user_id, Scripts.is_locked).where(
Scripts.workspace_id == workspace_id,
Scripts.script_name == notebook_path,
Scripts.script_type == "notebook",
Scripts.is_deleted == 0,
)
response.set_cookie(
key="jupyter_access",
value=raw_ticket,
max_age=max_age,
path="/jupyter/",
secure=cookie_secure(),
httponly=True,
samesite="lax",
row = (await session.execute(statement)).one_or_none()
if row is None:
return False
owner_user_id, is_locked = row
if owner_user_id == user_id:
return False
return bool(is_locked)
async def load_active_membership(
session: AsyncSession,
user_id: str,
workspace_id: str,
) -> tuple[Users, Workspaces]:
"""Resolve the user's active membership in the workspace."""
statement = (
select(Users, Workspaces)
.join(
WorkspaceMembers,
WorkspaceMembers.user_id == Users.user_id,
)
.join(
Workspaces,
Workspaces.workspace_id == WorkspaceMembers.workspace_id,
)
.where(
Users.user_id == user_id,
Users.status == "active",
WorkspaceMembers.workspace_id == workspace_id,
WorkspaceMembers.member_status == "active",
Workspaces.status == "active",
)
)
return response
row = (await session.execute(statement)).one_or_none()
if row is None:
raise HTTPException(
status_code=403,
detail="active workspace membership is required",
)
return row
@router.get("/api/v1/auth/jupyter")
async def verify_jupyter_access(
request: Request,
response: Response,
auth: Optional[HTTPAuthorizationCredentials] = Depends(security),
session: AsyncSession = Depends(database_session),
) -> dict:
"""Nginx auth_request subrequest handler.
Returns 200 with two response headers so Nginx can proxy the
request to the user's Jupyter instance:
* ``x-upstream-addr`` = ``{base_url}:{port}``
* ``x-jupyter-internal-token`` = the runtime-issued token
"""
workspace_id = request.headers.get("X-Original-Workspace-Id")
original_uri = request.headers.get("X-Original-URI", "")
if not workspace_id:
raise HTTPException(
status_code=400,
detail="Missing Workspace ID",
)
cookie_token = request.cookies.get("access_token")
bearer_token = auth.credentials if auth else None
token = cookie_token or bearer_token
payload = verify_jwt_token(token)
user_id = payload.get("sub")
if not user_id:
raise HTTPException(
status_code=401,
detail="Invalid Authentication Token",
)
await load_active_membership(session, user_id, workspace_id)
notebook_path = extract_notebook_path(original_uri, workspace_id)
if notebook_path and await check_notebook_is_locked(
session, workspace_id, notebook_path, user_id,
):
raise HTTPException(
status_code=403,
detail=f"Notebook '{notebook_path}' is currently locked",
)
runtime_client = request.app.state.runtime_client
ws_info = await runtime_client.get_workspace(workspace_id)
if not ws_info or ws_info.get("status") != "running":
try:
ws_info = await runtime_client.start_workspace(workspace_id)
except RuntimeClientError as exc:
raise HTTPException(
status_code=exc.status_code,
detail=exc.detail,
) from exc
target_port = ws_info.get("port")
jupyter_token = ws_info.get("token")
jupyter_base_url = ws_info.get("base_url")
if not target_port:
raise HTTPException(
status_code=500,
detail="Jupyter instance returned no port",
)
response.headers["x-upstream-addr"] = f"{jupyter_base_url}:{target_port}"
response.headers["x-jupyter-internal-token"] = jupyter_token or ""
return {"status": "ok"}
backend/src/backend/legacy_jupyter_auth.py
-147
View File
@@ -1,147 +0,0 @@
# coding=utf-8
"""
@Time :2026/7/27
@Author :tao.chen
"""
import os
import re
import httpx
from fastapi import FastAPI, Request, Response, HTTPException, Depends, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from typing import Optional
from loguru import logger
app = FastAPI(title="Jupyter Auth & Router Backend")
RUNTIME_BASE_URL = os.getenv("RUNTIME_BASE_URL", "http://127.0.0.1:8001")
security = HTTPBearer(auto_error=False)
# ------------------------------------------------------------------
# 1. Runtime 交互 Client
# ------------------------------------------------------------------
class RuntimeClient:
"""与 Runtime 进程管理器服务交互"""
@staticmethod
async def get_workspace(workspace_id: str) -> Optional[dict]:
"""按需查询单个 workspace 进程"""
async with httpx.AsyncClient(base_url=RUNTIME_BASE_URL) as client:
try:
resp = await client.post(
"/api/v1/jupyter",
json={"action": "get", "workspace_id": workspace_id},
timeout=3.0,
)
if resp.status_code == 200:
return resp.json()
return None
except httpx.RequestError:
return None
@staticmethod
async def start_workspace(workspace_id: str) -> dict:
"""进程未运行时主动触发启动"""
async with httpx.AsyncClient(base_url=RUNTIME_BASE_URL) as client:
resp = await client.post(
"/api/v1/jupyter",
json={"action": "start", "workspace_id": workspace_id},
timeout=10.0,
)
if resp.status_code == 200:
return resp.json()
raise HTTPException(
status_code=500, detail="Failed to start Jupyter instance"
)
# ------------------------------------------------------------------
# 2. 数据库与权限模拟 (请根据实际 MySQL ORM 修改)
# ------------------------------------------------------------------
async def check_notebook_is_locked(workspace_id: str, notebook_path: str) -> bool:
"""
查数据库：判断特定 Notebook 文件是否被锁定
:param workspace_id: 工作区 ID
:param notebook_path: 相对路径，如 "test.ipynb" 或 "folder/demo.ipynb"
"""
# 模拟锁定数据库：假定 test_locked.ipynb 被锁定
locked_notebooks = {
("test1234", "test_locked.ipynb"): True,
}
return locked_notebooks.get((workspace_id, notebook_path), False)
def verify_jwt_token(token: str) -> str:
"""校验 JWT 令牌"""
if token == "invalid-token":
raise HTTPException(status_code=401, detail="Invalid Authentication Token")
return "user_001"
def extract_notebook_path(uri: str, workspace_id: str) -> Optional[str]:
"""
从原始请求 URI 中提取请求的 .ipynb 文件相对路径
例如: /jupyter/test1234/notebooks/folder/test.ipynb -> folder/test.ipynb
"""
pattern = rf"^/jupyter/{re.escape(workspace_id)}/notebooks/(.+\.ipynb)"
match = re.match(pattern, uri)
if match:
return match.group(1)
return None
# ------------------------------------------------------------------
# 3. 核心 Auth 接口 (针对 Nginx auth_request)
# ------------------------------------------------------------------
@app.get("/api/v1/auth/jupyter")
async def verify_jupyter_access(
request: Request,
response: Response,
auth: Optional[HTTPAuthorizationCredentials] = Depends(security),
):
# 获取 Nginx 传入的元数据
workspace_id = request.headers.get("X-Original-Workspace-Id")
original_uri = request.headers.get("X-Original-URI", "")
cookie_token = request.cookies.get("access_token")
bearer_token = auth.credentials if auth else None
token = bearer_token or cookie_token
# if not token:
# raise HTTPException(status_code=401, detail="Missing Authentication Token")
if not workspace_id:
raise HTTPException(status_code=400, detail="Missing Workspace ID")
# 基础身份认证
# current_user_id = verify_jwt_token(token)
# 精准锁校验：只有在访问 .ipynb 文件时才检查 is_locked
notebook_path = extract_notebook_path(original_uri, workspace_id)
if notebook_path:
is_locked = await check_notebook_is_locked(workspace_id, notebook_path)
if is_locked:
raise HTTPException(
status_code=403,
detail=f"Notebook '{notebook_path}' is currently locked",
)
# 获取或启动 Jupyter 子进程
ws_info = await RuntimeClient.get_workspace(workspace_id)
if not ws_info or ws_info.get("status") != "running":
ws_info = await RuntimeClient.start_workspace(workspace_id)
target_port = ws_info.get("port")
jupyter_token = ws_info.get("token")
jupyter_base_url = ws_info.get("base_url")
if not target_port:
raise HTTPException(
status_code=500, detail="Jupyter instance returned no port"
)
# 通过 Response Header 返回 Upstream 地址与 Token 给 Nginx
response.headers["x-upstream-addr"] = f"{jupyter_base_url}:{target_port}"
response.headers["x-jupyter-internal-token"] = jupyter_token or ""
return {"status": "ok"}
backend/src/backend/main.py
+2 -24
View File
@@ -13,12 +13,10 @@ from common.db import create_database_engine, create_session_factory
from common.service_app import create_service_app
from common.storage import RustFSObjectStore
from backend.admin import router as admin_router
from backend.file_locks import router as file_locks_router
from backend.jupyter import router as jupyter_router
from backend.resources import router as resources_router
from backend.runtime_client import RuntimeClient
from backend.schedule_runs import router as schedule_runs_router
from backend.schedule_client import ScheduleExecutorClient
from backend.schedules import router as schedules_router
from backend.scripts import router as scripts_router
from backend.storage_api import app as storage_app
@@ -41,10 +39,6 @@ async def lifespan(app: Any) -> AsyncIterator[None]:
"RUSTFS_INTERNAL_ENDPOINT",
"http://rustfs:9000",
),
public_endpoint=os.getenv(
"RUSTFS_PUBLIC_ENDPOINT",
"http://localhost:9000",
),
access_key=os.environ["RUSTFS_ACCESS_KEY"],
secret_key=os.environ["RUSTFS_SECRET_KEY"],
)
@@ -61,31 +55,16 @@ async def lifespan(app: Any) -> AsyncIterator[None]:
base_url="http://backend.internal",
timeout=httpx.Timeout(30.0),
)
app.state.storage_client = StorageClient(
storage_http_client,
os.environ["INTERNAL_SERVICE_TOKEN"],
)
app.state.storage_client = StorageClient(storage_http_client)
runtime_http_client = httpx.AsyncClient(
base_url=os.getenv("RUNTIME_API_URL", "http://runtime:8000"),
timeout=httpx.Timeout(30.0),
)
app.state.runtime_client = RuntimeClient(
runtime_http_client,
os.environ["INTERNAL_SERVICE_TOKEN"],
)
schedule_http_client = httpx.AsyncClient(
base_url=os.getenv("SCHEDULE_API_URL", "http://schedule:8000"),
timeout=httpx.Timeout(10.0),
)
app.state.schedule_client = ScheduleExecutorClient(
schedule_http_client,
os.environ["INTERNAL_SERVICE_TOKEN"],
)
app.state.runtime_client = RuntimeClient(runtime_http_client)
try:
yield
finally:
await runtime_http_client.aclose()
await schedule_http_client.aclose()
await storage_http_client.aclose()
await engine.dispose()
@@ -94,7 +73,6 @@ app = create_service_app(
os.getenv("SERVICE_NAME", "backend"),
lifespan=lifespan,
)
app.include_router(file_locks_router)
app.include_router(jupyter_router)
app.include_router(resources_router)
app.include_router(schedule_runs_router)
backend/src/backend/resources.py
-1
View File
@@ -78,7 +78,6 @@ async def create_resource_upload(
"expected_size_bytes": payload.expected_size_bytes,
"expected_hash": payload.expected_hash,
"idempotency_key": idempotency_key,
"url_scope": "public",
}
)
return {"request_id": context.request_id, "data": data, "meta": {}}
backend/src/backend/runtime_client.py
+40 -63
View File
@@ -12,14 +12,20 @@ class RuntimeClientError(Exception):
detail: Any
_RUNTIME_TRANSPORT_ERROR = RuntimeClientError(
503,
{
"code": "RUNTIME_UNAVAILABLE",
"message": "Runtime Manager 暂时不可用",
"retryable": True,
"details": {},
},
)
class RuntimeClient:
def __init__(
self,
client: httpx.AsyncClient,
service_token: str,
) -> None:
def __init__(self, client: httpx.AsyncClient) -> None:
self.client = client
self.headers = {"X-Service-Token": service_token}
async def _request(
self,
@@ -28,22 +34,9 @@ class RuntimeClient:
payload: dict[str, Any],
) -> dict[str, Any]:
try:
response = await self.client.request(
method,
path,
json=payload,
headers=self.headers,
)
response = await self.client.request(method, path, json=payload)
except httpx.RequestError as exc:
raise RuntimeClientError(
503,
{
"code": "RUNTIME_UNAVAILABLE",
"message": "Runtime Manager 暂时不可用",
"retryable": True,
"details": {},
},
) from exc
raise _RUNTIME_TRANSPORT_ERROR from exc
if response.is_error:
try:
detail = response.json().get("detail", response.text)
@@ -52,53 +45,37 @@ class RuntimeClient:
raise RuntimeClientError(response.status_code, detail)
return response.json()
async def acquire_file_lock(
async def get_workspace(
self,
payload: dict[str, Any],
) -> dict[str, Any]:
return (
await self._request(
"POST",
"/internal/v1/file-locks/acquire",
payload,
)
)["data"]
workspace_id: str,
) -> dict[str, Any] | None:
"""Ask Runtime for a live workspace.
async def heartbeat_file_lock(
self,
edit_session_id: str,
payload: dict[str, Any],
) -> dict[str, Any]:
return (
await self._request(
Returns the workspace descriptor when the Jupyter process is
``running``; returns ``None`` when the workspace is not yet
started or has been torn down so callers can fall through to
:meth:`start_workspace`.
"""
try:
return await self._request(
"POST",
f"/internal/v1/file-locks/{edit_session_id}/heartbeat",
payload,
"/api/v1/jupyter",
{"action": "get", "workspace_id": workspace_id},
)
)["data"]
except RuntimeClientError as exc:
if exc.status_code == 404:
return None
raise
async def release_file_lock(
async def start_workspace(
self,
edit_session_id: str,
payload: dict[str, Any],
workspace_id: str,
) -> dict[str, Any]:
return (
await self._request(
"DELETE",
f"/internal/v1/file-locks/{edit_session_id}",
payload,
)
)["data"]
return await self._request(
"POST",
"/api/v1/jupyter",
{"action": "start", "workspace_id": workspace_id},
)
async def create_jupyter_access_ticket(
self,
edit_session_id: str,
payload: dict[str, Any],
) -> dict[str, Any]:
return (
await self._request(
"POST",
f"/internal/v1/jupyter/access-tickets/{edit_session_id}",
payload,
)
)["data"]
__all__ = ["RuntimeClient", "RuntimeClientError"]
backend/src/backend/schedule_runs.py
+1 -1
View File
@@ -20,7 +20,7 @@ from backend.dependencies import (
database_session,
request_context,
)
from backend.schedule_schemas import StrictModel
from common.schemas import StrictModel
from backend.schedules import (
graph_rows,
schedule_row,
backend/src/backend/schedule_schemas.py
+3 -6
View File
@@ -1,15 +1,12 @@
from __future__ import annotations
from datetime import datetime
from typing import Any, Literal
from typing import Any
from pydantic import Field, field_validator, model_validator
from backend.schemas import StrictModel
TriggerType = Literal["manual", "cron", "api"]
FailurePolicy = Literal["stop", "continue"]
from common.db.models.schedules import FailurePolicy, TriggerType
from common.schemas import StrictModel
def _required_text(value: str) -> str:
backend/src/backend/schemas.py
+2 -13
View File
@@ -2,11 +2,9 @@ from __future__ import annotations
from typing import Literal
from pydantic import BaseModel, ConfigDict, Field, field_validator
from pydantic import Field, field_validator
class StrictModel(BaseModel):
model_config = ConfigDict(extra="forbid")
from common.schemas import StrictModel
class CreateResourceUploadRequest(StrictModel):
@@ -61,12 +59,3 @@ class PublishVersionRequest(StrictModel):
class DownloadUrlRequest(StrictModel):
expires_seconds: int = Field(default=300, ge=30, le=3600)
class FileLockTokenRequest(StrictModel):
lock_token: str = Field(min_length=32, max_length=256)
class CreateJupyterAccessTicketRequest(StrictModel):
edit_session_id: str = Field(min_length=26, max_length=26)
lock_token: str = Field(min_length=32, max_length=256)
backend/src/backend/storage_api.py
+102 -166
View File
@@ -12,28 +12,26 @@ from datetime import UTC, datetime, timedelta
from pathlib import Path, PurePosixPath
from typing import Any, AsyncIterator
from fastapi import Depends, Header, HTTPException, Request, status
from fastapi import Depends, HTTPException, Request, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from common.db import create_database_engine, create_session_factory
from common.db import create_database_engine, create_session_factory, session_scope
from common.db.models import (
StorageObjects,
UploadSessions,
Users,
WorkspaceMembers,
Workspaces,
)
Workspaces)
from common.ids import new_ulid
from common.service_app import create_service_app
from common.storage import RustFSObjectStore
from backend.storage_schemas import (
from common.storage.schemas import (
CompleteUploadRequest,
CreateUploadRequest,
DownloadUrlRequest,
RegisterWorkspaceObjectRequest,
ServerObjectRequest,
)
ServerObjectRequest)
def utcnow() -> datetime:
@@ -47,8 +45,7 @@ def hash_bytes(value: str) -> bytes:
def normalized_idempotency_key(
workspace_id: str,
user_id: str,
value: str,
) -> str:
value: str) -> str:
digest = hashlib.sha256(
f"{workspace_id}:{user_id}:{value}".encode("utf-8")
).hexdigest()
@@ -92,23 +89,15 @@ async def lifespan(app: Any) -> AsyncIterator[None]:
app.state.object_store = RustFSObjectStore(
internal_endpoint=os.getenv(
"RUSTFS_INTERNAL_ENDPOINT",
"http://rustfs:9000",
),
public_endpoint=os.getenv(
"RUSTFS_PUBLIC_ENDPOINT",
"http://localhost:9000",
),
"http://rustfs:9000"),
access_key=os.environ["RUSTFS_ACCESS_KEY"],
secret_key=os.environ["RUSTFS_SECRET_KEY"],
)
secret_key=os.environ["RUSTFS_SECRET_KEY"])
app.state.default_bucket = os.getenv(
"RUSTFS_DEFAULT_BUCKET",
"model-platform",
)
"model-platform")
await asyncio.to_thread(
app.state.object_store.ensure_bucket,
app.state.default_bucket,
)
app.state.default_bucket)
try:
yield
finally:
@@ -117,75 +106,51 @@ async def lifespan(app: Any) -> AsyncIterator[None]:
app = create_service_app(
os.getenv("SERVICE_NAME", "storage-api"),
lifespan=lifespan,
)
lifespan=lifespan)
async def database_session(request: Request) -> AsyncIterator[AsyncSession]:
async with request.app.state.session_factory() as session:
try:
yield session
await session.commit()
except Exception:
await session.rollback()
raise
def verify_internal_service(
x_service_token: str = Header(alias="X-Service-Token"),
) -> None:
expected = os.environ.get("INTERNAL_SERVICE_TOKEN", "")
if not expected or not secrets.compare_digest(expected, x_service_token):
raise HTTPException(
status.HTTP_401_UNAUTHORIZED,
"invalid internal service identity",
)
async with session_scope(request.app.state.session_factory) as session:
yield session
async def require_workspace_member(
session: AsyncSession,
workspace_id: str,
user_id: str,
) -> Workspaces:
user_id: str) -> Workspaces:
statement = (
select(Workspaces)
.join(
WorkspaceMembers,
WorkspaceMembers.workspace_id == Workspaces.workspace_id,
)
WorkspaceMembers.workspace_id == Workspaces.workspace_id)
.join(Users, Users.user_id == WorkspaceMembers.user_id)
.where(
Workspaces.workspace_id == workspace_id,
Workspaces.status == "active",
WorkspaceMembers.user_id == user_id,
WorkspaceMembers.member_status == "active",
Users.status == "active",
)
Users.status == "active")
)
workspace = await session.scalar(statement)
if workspace is None:
raise HTTPException(
status.HTTP_403_FORBIDDEN,
"user is not an active workspace member",
)
"user is not an active workspace member")
return workspace
async def create_upload_record(
payload: CreateUploadRequest,
session: AsyncSession,
request: Request,
) -> dict[str, Any]:
request: Request) -> dict[str, Any]:
workspace = await require_workspace_member(
session,
payload.workspace_id,
payload.user_id,
)
payload.user_id)
stored_key = normalized_idempotency_key(
payload.workspace_id,
payload.user_id,
payload.idempotency_key,
)
payload.idempotency_key)
existing = await session.scalar(
select(UploadSessions).where(
UploadSessions.idempotency_key == stored_key
@@ -201,8 +166,7 @@ async def create_upload_record(
):
raise HTTPException(
status.HTTP_409_CONFLICT,
"idempotency key was used with different upload metadata",
)
"idempotency key was used with different upload metadata")
upload = existing
else:
upload_id = new_ulid()
@@ -226,16 +190,14 @@ async def create_upload_record(
expires_at=utcnow() + timedelta(minutes=15),
expected_size_bytes=payload.expected_size_bytes,
expected_hash=payload.expected_hash,
content_type=payload.content_type,
)
content_type=payload.content_type)
session.add(upload)
await session.flush()
if upload.upload_status == "completed" and upload.storage_object_id:
storage_object = await session.get(
StorageObjects,
upload.storage_object_id,
)
upload.storage_object_id)
return {
"upload_id": upload.upload_id,
"status": upload.upload_status,
@@ -246,33 +208,55 @@ async def create_upload_record(
if upload.upload_status not in {"created", "uploading"}:
raise HTTPException(
status.HTTP_409_CONFLICT,
f"upload cannot continue from status {upload.upload_status}",
)
f"upload cannot continue from status {upload.upload_status}")
url, headers = request.app.state.object_store.presign_put(
bucket_name=upload.bucket_name,
object_key=upload.object_key,
content_type=upload.content_type or "application/octet-stream",
expected_hash=upload.expected_hash,
expires_seconds=900,
public=payload.url_scope == "public",
expires_seconds=900)
presigned_url = request.app.state.object_store.rewrite_to_public_path(
url,
public_base_url=_public_base_url(request),
)
return {
"upload_id": upload.upload_id,
"status": upload.upload_status,
"method": "PUT",
"presigned_url": url,
"presigned_url": presigned_url,
"required_headers": headers,
"expires_at": upload.expires_at.isoformat(),
}
def _public_base_url(request: Request) -> str:
"""Return the public base URL the client should use.
Falls back to the inbound request's ``Host`` header and the scheme
Nginx forwards via ``X-Forwarded-Proto`` so the resulting
presigned URL always points at the public edge rather than the
in-cluster RustFS endpoint.
"""
forwarded_proto = request.headers.get("x-forwarded-proto", "").strip()
scheme = forwarded_proto or request.url.scheme or "http"
host = (
request.headers.get("x-forwarded-host", "").strip()
or request.headers.get("host", "").strip()
)
if not host:
raise HTTPException(
status.HTTP_400_BAD_REQUEST,
"cannot determine public host for presigned URL",
)
return f"{scheme}://{host}"
async def complete_upload_record(
upload_id: str,
payload: CompleteUploadRequest,
session: AsyncSession,
request: Request,
) -> StorageObjects:
request: Request) -> StorageObjects:
upload = await session.scalar(
select(UploadSessions)
.where(UploadSessions.upload_id == upload_id)
@@ -285,14 +269,12 @@ async def complete_upload_record(
if item is None:
raise HTTPException(
status.HTTP_409_CONFLICT,
"completed upload has no storage object",
)
"completed upload has no storage object")
return item
if upload.upload_status not in {"created", "uploading"}:
raise HTTPException(
status.HTTP_409_CONFLICT,
f"upload cannot be completed from status {upload.upload_status}",
)
f"upload cannot be completed from status {upload.upload_status}")
if upload.expires_at < utcnow():
upload.upload_status = "expired"
raise HTTPException(status.HTTP_409_CONFLICT, "upload expired")
@@ -301,13 +283,11 @@ async def complete_upload_record(
head = await asyncio.to_thread(
request.app.state.object_store.head,
bucket_name=upload.bucket_name,
object_key=upload.object_key,
)
object_key=upload.object_key)
except Exception as exc:
raise HTTPException(
status.HTTP_409_CONFLICT,
"uploaded object is not available",
) from exc
"uploaded object is not available") from exc
actual_size = int(head.get("ContentLength", 0))
if (
@@ -317,8 +297,7 @@ async def complete_upload_record(
upload.upload_status = "failed"
raise HTTPException(
status.HTTP_409_CONFLICT,
"uploaded object size does not match expected_size_bytes",
)
"uploaded object size does not match expected_size_bytes")
actual_content_type = str(
head.get("ContentType") or "application/octet-stream"
)
@@ -326,8 +305,7 @@ async def complete_upload_record(
upload.upload_status = "failed"
raise HTTPException(
status.HTTP_409_CONFLICT,
"uploaded object content type does not match",
)
"uploaded object content type does not match")
metadata = {
str(key).lower(): str(value).lower()
for key, value in dict(head.get("Metadata") or {}).items()
@@ -337,14 +315,12 @@ async def complete_upload_record(
actual_hash = await asyncio.to_thread(
request.app.state.object_store.sha256,
bucket_name=upload.bucket_name,
object_key=upload.object_key,
)
object_key=upload.object_key)
if upload.expected_hash and actual_hash != upload.expected_hash:
upload.upload_status = "failed"
raise HTTPException(
status.HTTP_409_CONFLICT,
"uploaded object hash does not match expected_hash",
)
"uploaded object hash does not match expected_hash")
file_name = upload.object_key.rsplit("/", 1)[-1]
item = StorageObjects(
@@ -367,8 +343,7 @@ async def complete_upload_record(
visibility=payload.visibility,
is_immutable=int(payload.is_immutable),
object_status="available",
created_by=upload.user_id,
)
created_by=upload.user_id)
session.add(item)
await session.flush()
await session.refresh(item)
@@ -378,45 +353,37 @@ async def complete_upload_record(
return item
@app.post("/internal/v1/uploads", dependencies=[Depends(verify_internal_service)])
@app.post("/internal/v1/uploads")
async def create_upload(
payload: CreateUploadRequest,
request: Request,
session: AsyncSession = Depends(database_session),
) -> dict[str, Any]:
session: AsyncSession = Depends(database_session)) -> dict[str, Any]:
return {
"data": await create_upload_record(payload, session, request),
}
@app.post(
"/internal/v1/uploads/{upload_id}/complete",
dependencies=[Depends(verify_internal_service)],
)
"/internal/v1/uploads/{upload_id}/complete")
async def complete_upload(
upload_id: str,
payload: CompleteUploadRequest,
request: Request,
session: AsyncSession = Depends(database_session),
) -> dict[str, Any]:
session: AsyncSession = Depends(database_session)) -> dict[str, Any]:
item = await complete_upload_record(
upload_id,
payload,
session,
request,
)
request)
return {"data": storage_payload(item)}
@app.post(
"/internal/v1/uploads/{upload_id}/abort",
dependencies=[Depends(verify_internal_service)],
)
"/internal/v1/uploads/{upload_id}/abort")
async def abort_upload(
upload_id: str,
request: Request,
session: AsyncSession = Depends(database_session),
) -> dict[str, Any]:
session: AsyncSession = Depends(database_session)) -> dict[str, Any]:
upload = await session.scalar(
select(UploadSessions)
.where(UploadSessions.upload_id == upload_id)
@@ -427,39 +394,32 @@ async def abort_upload(
if upload.upload_status == "completed":
raise HTTPException(
status.HTTP_409_CONFLICT,
"completed upload cannot be aborted",
)
"completed upload cannot be aborted")
if upload.upload_status != "aborted":
await asyncio.to_thread(
request.app.state.object_store.delete,
bucket_name=upload.bucket_name,
object_key=upload.object_key,
)
object_key=upload.object_key)
upload.upload_status = "aborted"
return {"data": {"upload_id": upload_id, "status": "aborted"}}
@app.post(
"/internal/v1/objects",
dependencies=[Depends(verify_internal_service)],
)
"/internal/v1/objects")
async def create_server_object(
payload: ServerObjectRequest,
request: Request,
session: AsyncSession = Depends(database_session),
) -> dict[str, Any]:
session: AsyncSession = Depends(database_session)) -> dict[str, Any]:
try:
content = base64.b64decode(payload.content_base64, validate=True)
except (binascii.Error, ValueError) as exc:
raise HTTPException(
status.HTTP_422_UNPROCESSABLE_ENTITY,
"content_base64 is invalid",
) from exc
"content_base64 is invalid") from exc
if len(content) > 100 * 1024 * 1024:
raise HTTPException(
status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
"object exceeds 100 MiB server-side upload limit",
)
"object exceeds 100 MiB server-side upload limit")
content_hash = hashlib.sha256(content).hexdigest()
upload_result = await create_upload_record(
CreateUploadRequest(
@@ -470,12 +430,9 @@ async def create_server_object(
content_type=payload.content_type,
expected_size_bytes=len(content),
expected_hash=content_hash,
idempotency_key=payload.idempotency_key,
url_scope="internal",
),
idempotency_key=payload.idempotency_key),
session,
request,
)
request)
if upload_result.get("status") == "completed":
return {"data": upload_result["storage_object"], "meta": {"reused": True}}
@@ -483,42 +440,34 @@ async def create_server_object(
if upload is None:
raise HTTPException(
status.HTTP_500_INTERNAL_SERVER_ERROR,
"upload record disappeared",
)
"upload record disappeared")
await asyncio.to_thread(
request.app.state.object_store.put_bytes,
bucket_name=upload.bucket_name,
object_key=upload.object_key,
content=content,
content_type=payload.content_type,
content_hash=content_hash,
)
content_hash=content_hash)
item = await complete_upload_record(
upload.upload_id,
CompleteUploadRequest(
usage_type=payload.usage_type,
visibility=payload.visibility,
is_immutable=payload.is_immutable,
),
is_immutable=payload.is_immutable),
session,
request,
)
request)
return {"data": storage_payload(item), "meta": {"reused": False}}
@app.post(
"/internal/v1/workspace-objects",
dependencies=[Depends(verify_internal_service)],
)
"/internal/v1/workspace-objects")
async def register_workspace_object(
payload: RegisterWorkspaceObjectRequest,
session: AsyncSession = Depends(database_session),
) -> dict[str, Any]:
session: AsyncSession = Depends(database_session)) -> dict[str, Any]:
workspace = await require_workspace_member(
session,
payload.workspace_id,
payload.user_id,
)
payload.user_id)
pure_path = PurePosixPath(payload.relative_path.replace("\\", "/"))
if (
pure_path.is_absolute()
@@ -527,8 +476,7 @@ async def register_workspace_object(
):
raise HTTPException(
status.HTTP_422_UNPROCESSABLE_ENTITY,
"invalid workspace relative_path",
)
"invalid workspace relative_path")
relative_path = pure_path.as_posix()
workspace_root = Path(
os.getenv("WORKSPACE_ROOT", "/workspace/workspaces")
@@ -538,13 +486,11 @@ async def register_workspace_object(
if scoped_root != target and scoped_root not in target.parents:
raise HTTPException(
status.HTTP_422_UNPROCESSABLE_ENTITY,
"workspace path escapes its root",
)
"workspace path escapes its root")
if not target.is_file():
raise HTTPException(
status.HTTP_404_NOT_FOUND,
"workspace file does not exist",
)
"workspace file does not exist")
content = await asyncio.to_thread(target.read_bytes)
content_hash = hashlib.sha256(content).hexdigest()
stat_result = target.stat()
@@ -553,8 +499,7 @@ async def register_workspace_object(
select(StorageObjects).where(
StorageObjects.workspace_id == payload.workspace_id,
StorageObjects.storage_backend == "workspace_fs",
StorageObjects.path_hash == path_digest,
)
StorageObjects.path_hash == path_digest)
)
reused = item is not None
if item is None:
@@ -572,19 +517,16 @@ async def register_workspace_object(
visibility=payload.visibility,
is_immutable=0,
object_status="available",
created_by=payload.user_id,
)
created_by=payload.user_id)
session.add(item)
elif item.is_immutable:
raise HTTPException(
status.HTTP_409_CONFLICT,
"immutable workspace object cannot be updated",
)
"immutable workspace object cannot be updated")
elif item.owner_user_id != payload.user_id:
raise HTTPException(
status.HTTP_403_FORBIDDEN,
"workspace object belongs to another user",
)
"workspace object belongs to another user")
item.usage_type = payload.usage_type
item.file_name = target.name
item.file_extension = target.suffix.lower() or None
@@ -601,15 +543,12 @@ async def register_workspace_object(
@app.post(
"/internal/v1/objects/{storage_object_id}/download-url",
dependencies=[Depends(verify_internal_service)],
)
"/internal/v1/objects/{storage_object_id}/download-url")
async def create_download_url(
storage_object_id: str,
payload: DownloadUrlRequest,
request: Request,
session: AsyncSession = Depends(database_session),
) -> dict[str, Any]:
session: AsyncSession = Depends(database_session)) -> dict[str, Any]:
item = await session.get(StorageObjects, storage_object_id)
if item is None or item.object_status != "available":
raise HTTPException(status.HTTP_404_NOT_FOUND, "object not found")
@@ -620,18 +559,20 @@ async def create_download_url(
):
raise HTTPException(
status.HTTP_409_CONFLICT,
"object does not support a presigned URL",
)
"object does not support a presigned URL")
url = request.app.state.object_store.presign_get(
bucket_name=item.bucket_name,
object_key=item.object_key,
file_name=item.file_name,
expires_seconds=payload.expires_seconds,
expires_seconds=payload.expires_seconds)
presigned_url = request.app.state.object_store.rewrite_to_public_path(
url,
public_base_url=_public_base_url(request),
)
return {
"data": {
"storage_object_id": item.storage_object_id,
"presigned_url": url,
"presigned_url": presigned_url,
"method": "GET",
"expires_in_seconds": payload.expires_seconds,
}
@@ -639,14 +580,11 @@ async def create_download_url(
@app.delete(
"/internal/v1/objects/{storage_object_id}",
dependencies=[Depends(verify_internal_service)],
)
"/internal/v1/objects/{storage_object_id}")
async def delete_object(
storage_object_id: str,
request: Request,
session: AsyncSession = Depends(database_session),
) -> dict[str, Any]:
session: AsyncSession = Depends(database_session)) -> dict[str, Any]:
item = await session.scalar(
select(StorageObjects)
.where(StorageObjects.storage_object_id == storage_object_id)
@@ -657,15 +595,13 @@ async def delete_object(
if item.is_immutable:
raise HTTPException(
status.HTTP_409_CONFLICT,
"immutable object cannot be deleted",
)
"immutable object cannot be deleted")
if item.object_status != "deleted":
if item.storage_backend == "rustfs" and item.bucket_name and item.object_key:
await asyncio.to_thread(
request.app.state.object_store.delete,
bucket_name=item.bucket_name,
object_key=item.object_key,
)
object_key=item.object_key)
item.object_status = "deleted"
item.deleted_at = utcnow()
return {
backend/src/backend/storage_client.py
+3 -114
View File
@@ -1,116 +1,5 @@
from __future__ import annotations
"""Backward-compatible re-export of the shared storage client."""
import base64
from typing import Any
from common.storage.client import StorageClient
import httpx
from fastapi import HTTPException
class StorageClient:
def __init__(
self,
client: httpx.AsyncClient,
service_token: str,
) -> None:
self.client = client
self.headers = {"X-Service-Token": service_token}
async def _request(
self,
method: str,
path: str,
*,
payload: dict[str, Any] | None = None,
) -> dict[str, Any]:
response = await self.client.request(
method,
path,
json=payload,
headers=self.headers,
)
if response.is_error:
try:
detail = response.json().get("detail", response.text)
except ValueError:
detail = response.text
raise HTTPException(response.status_code, detail)
return response.json()
async def create_upload(
self,
payload: dict[str, Any],
) -> dict[str, Any]:
return (await self._request(
"POST",
"/internal/v1/uploads",
payload=payload,
))["data"]
async def complete_upload(
self,
upload_id: str,
payload: dict[str, Any],
) -> dict[str, Any]:
return (await self._request(
"POST",
f"/internal/v1/uploads/{upload_id}/complete",
payload=payload,
))["data"]
async def register_workspace_object(
self,
payload: dict[str, Any],
) -> dict[str, Any]:
return (await self._request(
"POST",
"/internal/v1/workspace-objects",
payload=payload,
))["data"]
async def create_server_object(
self,
*,
workspace_id: str,
user_id: str,
usage_type: str,
file_name: str,
content_type: str,
content: bytes,
visibility: str,
is_immutable: bool,
idempotency_key: str,
) -> dict[str, Any]:
result = await self._request(
"POST",
"/internal/v1/objects",
payload={
"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": visibility,
"is_immutable": is_immutable,
"idempotency_key": idempotency_key,
},
)
return result["data"]
async def create_download_url(
self,
storage_object_id: str,
expires_seconds: int,
) -> dict[str, Any]:
return (await self._request(
"POST",
f"/internal/v1/objects/{storage_object_id}/download-url",
payload={"expires_seconds": expires_seconds},
))["data"]
async def delete_object(self, storage_object_id: str) -> dict[str, Any]:
return (await self._request(
"DELETE",
f"/internal/v1/objects/{storage_object_id}",
))["data"]
__all__ = ["StorageClient"]
backend/src/backend/storage_schemas.py
-79
View File
@@ -1,79 +0,0 @@
from __future__ import annotations
from typing import Literal
from pydantic import BaseModel, ConfigDict, Field, field_validator
class StrictModel(BaseModel):
model_config = ConfigDict(extra="forbid")
class CreateUploadRequest(StrictModel):
workspace_id: str = Field(min_length=26, max_length=26)
user_id: str = Field(min_length=26, max_length=26)
usage_type: Literal[
"data_resource",
"version_artifact",
"snapshot",
"run_log",
"run_result",
]
file_name: str = Field(min_length=1, max_length=255)
content_type: str = Field(min_length=1, max_length=255)
expected_size_bytes: int = Field(ge=0, le=100 * 1024 * 1024)
expected_hash: str | None = Field(default=None, min_length=64, max_length=64)
idempotency_key: str = Field(min_length=8, max_length=128)
url_scope: Literal["public", "internal"] = "public"
@field_validator("expected_hash")
@classmethod
def validate_hash(cls, value: str | None) -> str | None:
if value is None:
return None
normalized = value.lower()
if any(character not in "0123456789abcdef" for character in normalized):
raise ValueError("expected_hash must be lowercase SHA-256 hex")
return normalized
class CompleteUploadRequest(StrictModel):
usage_type: Literal[
"data_resource",
"version_artifact",
"snapshot",
"run_log",
"run_result",
]
visibility: Literal["private", "workspace", "public"] = "private"
is_immutable: bool = False
class ServerObjectRequest(StrictModel):
workspace_id: str = Field(min_length=26, max_length=26)
user_id: str = Field(min_length=26, max_length=26)
usage_type: Literal[
"data_resource",
"version_artifact",
"snapshot",
"run_log",
"run_result",
]
file_name: str = Field(min_length=1, max_length=255)
content_type: str = Field(min_length=1, max_length=255)
content_base64: str = Field(min_length=1)
visibility: Literal["private", "workspace", "public"] = "private"
is_immutable: bool = False
idempotency_key: str = Field(min_length=8, max_length=128)
class RegisterWorkspaceObjectRequest(StrictModel):
workspace_id: str = Field(min_length=26, max_length=26)
user_id: str = Field(min_length=26, max_length=26)
relative_path: str = Field(min_length=1, max_length=1024)
usage_type: Literal["working_copy", "public_script"]
visibility: Literal["private", "workspace", "public"] = "private"
class DownloadUrlRequest(StrictModel):
expires_seconds: int = Field(default=300, ge=30, le=3600)