57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
"""Quota-safe monthly claims for organization lookups in Checko."""
|
|
|
|
from datetime import date, datetime
|
|
|
|
from apps.parsers.models import CheckoCollectionAttempt
|
|
from django.utils import timezone
|
|
|
|
|
|
def collection_period_month(at: date | datetime | None = None) -> date:
|
|
if at is None:
|
|
local_date = timezone.localdate()
|
|
elif isinstance(at, datetime):
|
|
local_date = timezone.localtime(at).date() if timezone.is_aware(at) else at.date()
|
|
else:
|
|
local_date = at
|
|
return local_date.replace(day=1)
|
|
|
|
|
|
def claim_monthly_collection(
|
|
*,
|
|
organization_id: str,
|
|
source: str,
|
|
at: date | datetime | None = None,
|
|
) -> CheckoCollectionAttempt | None:
|
|
"""Atomically claim this month's only allowed Checko collection attempt."""
|
|
attempt, created = CheckoCollectionAttempt.objects.get_or_create(
|
|
organization_id=organization_id,
|
|
source=source,
|
|
period_month=collection_period_month(at),
|
|
defaults={"status": CheckoCollectionAttempt.Status.IN_PROGRESS},
|
|
)
|
|
return attempt if created else None
|
|
|
|
|
|
def finish_collection(
|
|
attempt: CheckoCollectionAttempt,
|
|
*,
|
|
records_count: int,
|
|
error: Exception | str | None = None,
|
|
) -> None:
|
|
"""Persist the outcome without releasing the monthly claim."""
|
|
attempt.records_count = max(int(records_count), 0)
|
|
attempt.error_message = str(error or "")
|
|
attempt.status = (
|
|
CheckoCollectionAttempt.Status.FAILED
|
|
if error is not None
|
|
else CheckoCollectionAttempt.Status.SUCCESS
|
|
)
|
|
attempt.save(
|
|
update_fields=[
|
|
"records_count",
|
|
"error_message",
|
|
"status",
|
|
"updated_at",
|
|
]
|
|
)
|