File size: 4,597 Bytes
30ea0e9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | from __future__ import annotations
from fastapi import APIRouter, Depends, Request
from app.audit import AuditLogger
from app.config import Settings
from app.deps import (
get_audit,
get_bucket_write_limiter,
get_hub,
get_settings_dep,
)
from app.errors import NotRegistered, RateLimited, SyncTooLarge
from app.hub import HubClient
from app.models import (
ArtifactSyncRequest,
SharedResourceSyncRequest,
SyncFile,
SyncResponse,
)
from app.naming import artifact_dest_dir, registration_path
from app.rate_limit import CompoundLimiter
from app.validation import (
check_dest_not_blocked,
resolve_source,
validate_shared_dest_path,
validate_slug,
)
router = APIRouter()
def _require_registered(hub: HubClient, agent_id: str) -> None:
try:
hub.read_central_text(registration_path(agent_id))
except Exception:
raise NotRegistered(agent_id)
def _check_sync_caps(settings: Settings, files: list) -> int:
total_bytes = sum(f.size for f in files)
if len(files) > settings.sync_max_files:
raise SyncTooLarge(f"{len(files)} files exceeds cap of {settings.sync_max_files}")
if total_bytes > settings.sync_max_bytes:
raise SyncTooLarge(
f"{total_bytes} bytes exceeds cap of {settings.sync_max_bytes}"
)
return total_bytes
def _execute_sync(
hub: HubClient,
src_bucket: str,
src_prefix: str,
dest_prefix: str,
) -> list[SyncFile]:
out: list[SyncFile] = []
for src_path, dest_path, size in hub.copy_tree_to_central(src_bucket, src_prefix, dest_prefix):
out.append(SyncFile(src_path=src_path, dest_path=dest_path, bytes=size))
return out
@router.post("/v1/artifacts:sync", response_model=SyncResponse)
def artifacts_sync(
req: ArtifactSyncRequest,
request: Request,
settings: Settings = Depends(get_settings_dep),
hub: HubClient = Depends(get_hub),
audit: AuditLogger = Depends(get_audit),
limiter: CompoundLimiter = Depends(get_bucket_write_limiter),
) -> SyncResponse:
validate_slug(req.dest_slug)
parsed, agent_id = resolve_source(settings, req.source)
_require_registered(hub, agent_id)
allowed, retry = limiter.try_consume(parsed.bucket)
if not allowed:
raise RateLimited(retry)
src_bucket = f"{parsed.org}/{parsed.bucket}"
src_prefix = parsed.path
files = hub.list_bucket_dir(src_bucket, src_prefix)
_check_sync_caps(settings, files)
dest_prefix = artifact_dest_dir(req.dest_slug, agent_id)
check_dest_not_blocked(dest_prefix)
copied = _execute_sync(hub, src_bucket, src_prefix, dest_prefix)
total = sum(f.bytes for f in copied)
audit.write(
agent_id=agent_id,
route="/v1/artifacts:sync",
via="bucket",
source=str(parsed),
target_path=dest_prefix,
bytes_count=total,
status_code=200,
caller_ip=request.client.host if request.client else None,
user_agent=request.headers.get("user-agent"),
extra={"file_count": len(copied)},
)
return SyncResponse(dest=dest_prefix, files=copied, bytes_copied=total)
@router.post("/v1/shared-resources:sync", response_model=SyncResponse)
def shared_resources_sync(
req: SharedResourceSyncRequest,
request: Request,
settings: Settings = Depends(get_settings_dep),
hub: HubClient = Depends(get_hub),
audit: AuditLogger = Depends(get_audit),
limiter: CompoundLimiter = Depends(get_bucket_write_limiter),
) -> SyncResponse:
parsed, agent_id = resolve_source(settings, req.source)
validate_shared_dest_path(req.dest_path, agent_id)
_require_registered(hub, agent_id)
allowed, retry = limiter.try_consume(parsed.bucket)
if not allowed:
raise RateLimited(retry)
src_bucket = f"{parsed.org}/{parsed.bucket}"
src_prefix = parsed.path
files = hub.list_bucket_dir(src_bucket, src_prefix)
_check_sync_caps(settings, files)
dest_prefix = f"shared_resources/{req.dest_path}"
copied = _execute_sync(hub, src_bucket, src_prefix, dest_prefix)
total = sum(f.bytes for f in copied)
audit.write(
agent_id=agent_id,
route="/v1/shared-resources:sync",
via="bucket",
source=str(parsed),
target_path=dest_prefix,
bytes_count=total,
status_code=200,
caller_ip=request.client.host if request.client else None,
user_agent=request.headers.get("user-agent"),
extra={"file_count": len(copied)},
)
return SyncResponse(dest=dest_prefix, files=copied, bytes_copied=total)
|