85 lines
3.0 KiB
Python
85 lines
3.0 KiB
Python
"""SQLAlchemy-backed upload-policy store."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from copy import deepcopy
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from deerflow.persistence.upload_policy.base import UploadPolicyStore
|
|
from deerflow.persistence.upload_policy.model import UploadPolicyRow
|
|
|
|
_GLOBAL_SCOPE = "global"
|
|
|
|
|
|
def _to_dict(row: UploadPolicyRow) -> dict[str, Any]:
|
|
updated_at: datetime | str | None = row.updated_at
|
|
if isinstance(updated_at, datetime):
|
|
updated_at = updated_at.isoformat()
|
|
return {
|
|
"policy": deepcopy(row.policy_json or {}),
|
|
"updated_by": row.updated_by,
|
|
"updated_at": updated_at,
|
|
"revision": int(row.revision or 1),
|
|
}
|
|
|
|
|
|
class UploadPolicyRepository(UploadPolicyStore):
|
|
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
|
self._sf = session_factory
|
|
|
|
async def ensure_seed(self, policy: dict[str, Any]) -> dict[str, Any]:
|
|
existing = await self.get_policy()
|
|
if existing is not None:
|
|
return existing
|
|
async with self._sf() as session:
|
|
row = UploadPolicyRow(
|
|
scope=_GLOBAL_SCOPE,
|
|
policy_json=deepcopy(policy),
|
|
revision=1,
|
|
updated_by=None,
|
|
)
|
|
session.add(row)
|
|
try:
|
|
await session.commit()
|
|
except IntegrityError:
|
|
# Another worker seeded the singleton concurrently.
|
|
await session.rollback()
|
|
winner = await self.get_policy()
|
|
if winner is None:
|
|
raise
|
|
return winner
|
|
await session.refresh(row)
|
|
return _to_dict(row)
|
|
|
|
async def get_policy(self) -> dict[str, Any] | None:
|
|
async with self._sf() as session:
|
|
row = await session.get(UploadPolicyRow, _GLOBAL_SCOPE)
|
|
return _to_dict(row) if row is not None else None
|
|
|
|
async def set_policy(self, policy: dict[str, Any], *, updated_by: str | None = None) -> dict[str, Any]:
|
|
async with self._sf() as session:
|
|
result = await session.execute(
|
|
select(UploadPolicyRow).where(UploadPolicyRow.scope == _GLOBAL_SCOPE).with_for_update()
|
|
)
|
|
row = result.scalar_one_or_none()
|
|
if row is None:
|
|
row = UploadPolicyRow(
|
|
scope=_GLOBAL_SCOPE,
|
|
policy_json=deepcopy(policy),
|
|
revision=1,
|
|
updated_by=updated_by,
|
|
)
|
|
session.add(row)
|
|
else:
|
|
row.policy_json = deepcopy(policy)
|
|
row.updated_by = updated_by
|
|
row.revision = int(row.revision or 0) + 1
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
return _to_dict(row)
|