121 lines
4.6 KiB
Python
121 lines
4.6 KiB
Python
"""Atomic idempotency claims with a lease and fenced completion."""
|
|
|
|
import hashlib
|
|
import json
|
|
import uuid
|
|
from datetime import UTC, date, datetime, timedelta
|
|
|
|
from sqlalchemy import or_, select, update
|
|
from sqlalchemy.dialects.postgresql import insert
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.exceptions import AppException
|
|
from app.models.entities import IntegrationAssessment
|
|
from app.schemas.integration_risk import IntegrationRiskRequest
|
|
|
|
|
|
class IntegrationRiskRepository:
|
|
def __init__(self, session: AsyncSession):
|
|
self.session = session
|
|
|
|
async def get(self, client_id: str, request_id: str) -> IntegrationAssessment | None:
|
|
return await self.session.scalar(
|
|
select(IntegrationAssessment)
|
|
.where(
|
|
IntegrationAssessment.client_id == client_id,
|
|
IntegrationAssessment.request_id == request_id,
|
|
)
|
|
.execution_options(populate_existing=True)
|
|
)
|
|
|
|
async def claim(
|
|
self,
|
|
payload: IntegrationRiskRequest,
|
|
user_id: int,
|
|
default_date: date,
|
|
) -> tuple[IntegrationAssessment, bool]:
|
|
data = payload.model_dump(mode="json", by_alias=True)
|
|
# Preserve the digest of DingTalk requests created before externalUserId was introduced.
|
|
if data["operator"]["externalUserId"] is None:
|
|
data["operator"].pop("externalUserId")
|
|
digest = hashlib.sha256(
|
|
json.dumps(data, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode()
|
|
).hexdigest()
|
|
now = datetime.now(UTC)
|
|
token = uuid.uuid4().hex
|
|
lease_until = now + timedelta(seconds=150)
|
|
inserted = await self.session.scalar(
|
|
insert(IntegrationAssessment)
|
|
.values(
|
|
client_id=payload.source_system,
|
|
request_id=payload.request_id,
|
|
user_id=user_id,
|
|
request_hash=digest,
|
|
request_data=data,
|
|
as_of_date=payload.assessment.as_of_date or default_date,
|
|
status="processing",
|
|
owner_token=token,
|
|
lease_until=lease_until,
|
|
created_at=now,
|
|
updated_at=now,
|
|
)
|
|
.on_conflict_do_nothing(index_elements=["client_id", "request_id"])
|
|
.returning(IntegrationAssessment.id)
|
|
)
|
|
record = await self.get(payload.source_system, payload.request_id)
|
|
if record is None:
|
|
raise AppException("评估请求记录失败", code=503, status_code=503)
|
|
if record.request_hash != digest:
|
|
raise AppException("requestId 已被不同请求内容使用", code=409, status_code=409)
|
|
if inserted is not None:
|
|
await self.session.commit()
|
|
return record, True
|
|
if record.status == "completed":
|
|
return record, False
|
|
reclaimed = await self.session.scalar(
|
|
update(IntegrationAssessment)
|
|
.where(
|
|
IntegrationAssessment.id == record.id,
|
|
or_(IntegrationAssessment.status == "failed", IntegrationAssessment.lease_until <= now),
|
|
IntegrationAssessment.status != "completed",
|
|
)
|
|
.values(status="processing", owner_token=token, lease_until=lease_until, error_message=None, updated_at=now)
|
|
.execution_options(synchronize_session=False)
|
|
.returning(IntegrationAssessment.id)
|
|
)
|
|
await self.session.commit()
|
|
record = await self.get(payload.source_system, payload.request_id)
|
|
return record, reclaimed is not None
|
|
|
|
async def finish(
|
|
self,
|
|
record_id: int,
|
|
owner_token: str,
|
|
*,
|
|
result: dict | None = None,
|
|
markdown: str | None = None,
|
|
error: str | None = None,
|
|
) -> bool:
|
|
updated = await self.session.scalar(
|
|
update(IntegrationAssessment)
|
|
.where(
|
|
IntegrationAssessment.id == record_id,
|
|
IntegrationAssessment.owner_token == owner_token,
|
|
IntegrationAssessment.status == "processing",
|
|
)
|
|
.values(
|
|
status="failed" if error else "completed",
|
|
result_data=result,
|
|
report_markdown=markdown,
|
|
error_message=error,
|
|
updated_at=datetime.now(UTC),
|
|
)
|
|
.execution_options(synchronize_session=False)
|
|
.returning(IntegrationAssessment.id)
|
|
)
|
|
if updated is None:
|
|
await self.session.rollback()
|
|
return False
|
|
await self.session.commit()
|
|
return True
|