|
|
|
|
@@ -21,18 +21,24 @@ from uuid import uuid4
|
|
|
|
|
|
|
|
|
|
from django.conf import settings
|
|
|
|
|
from django.core.cache import cache
|
|
|
|
|
from django.db.models import QuerySet, prefetch_related_objects
|
|
|
|
|
from django.db import connection
|
|
|
|
|
from django.db.models import Prefetch, Q, QuerySet, prefetch_related_objects
|
|
|
|
|
from django.db.models.expressions import RawSQL
|
|
|
|
|
from django.utils import timezone
|
|
|
|
|
from openpyxl import Workbook
|
|
|
|
|
|
|
|
|
|
from organizations.models import OrganizationSourceRecord, SourceGroup
|
|
|
|
|
from organizations.models import (
|
|
|
|
|
OrganizationSourceFinancialLine,
|
|
|
|
|
OrganizationSourceRecord,
|
|
|
|
|
SourceGroup,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
EXPORT_FORMAT_CSV = "csv"
|
|
|
|
|
EXPORT_FORMAT_XLSX = "xlsx"
|
|
|
|
|
EXPORT_FORMAT_JSON = "json"
|
|
|
|
|
EXPORT_FORMATS = (EXPORT_FORMAT_CSV, EXPORT_FORMAT_XLSX, EXPORT_FORMAT_JSON)
|
|
|
|
|
FINANCIAL_SOURCE_GROUP = SourceGroup.FINANCIAL_INDICATORS.value
|
|
|
|
|
EXPORT_MANIFEST_VERSION = 1
|
|
|
|
|
EXPORT_MANIFEST_VERSION = 2
|
|
|
|
|
CURRENT_EXPORT_MANIFEST_FILE_NAME = "current.json"
|
|
|
|
|
GENERATION_MANIFEST_FILE_NAME = "manifest.json"
|
|
|
|
|
GENERATION_DIRECTORY_NAME = "generations"
|
|
|
|
|
@@ -120,6 +126,7 @@ class SourceRecordExportGeneration:
|
|
|
|
|
|
|
|
|
|
generation_id: str
|
|
|
|
|
generated_at: str
|
|
|
|
|
export_year: int
|
|
|
|
|
artifacts: tuple[SourceRecordExportArtifact, ...]
|
|
|
|
|
records_count: int
|
|
|
|
|
|
|
|
|
|
@@ -198,6 +205,7 @@ def build_source_record_export_artifacts(
|
|
|
|
|
generations_directory.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
|
|
|
|
generated_at_datetime = _normalize_generation_datetime(now or timezone.now())
|
|
|
|
|
export_year = _export_year(generated_at_datetime)
|
|
|
|
|
generation_id = (
|
|
|
|
|
f"{generated_at_datetime.strftime('%Y%m%dT%H%M%SZ')}-{uuid4().hex[:8]}"
|
|
|
|
|
)
|
|
|
|
|
@@ -214,6 +222,7 @@ def build_source_record_export_artifacts(
|
|
|
|
|
headers, records_count = _spool_source_group_rows(
|
|
|
|
|
source_group=source_group,
|
|
|
|
|
output_path=row_spool_path,
|
|
|
|
|
export_year=export_year,
|
|
|
|
|
)
|
|
|
|
|
source_record_counts[source_group] = records_count
|
|
|
|
|
|
|
|
|
|
@@ -253,6 +262,7 @@ def build_source_record_export_artifacts(
|
|
|
|
|
generation = SourceRecordExportGeneration(
|
|
|
|
|
generation_id=generation_id,
|
|
|
|
|
generated_at=generated_at_datetime.isoformat(),
|
|
|
|
|
export_year=export_year,
|
|
|
|
|
artifacts=tuple(artifacts),
|
|
|
|
|
records_count=sum(source_record_counts.values()),
|
|
|
|
|
)
|
|
|
|
|
@@ -306,6 +316,7 @@ def build_source_records_export_archive(
|
|
|
|
|
source_groups: Sequence[str],
|
|
|
|
|
export_format: str,
|
|
|
|
|
export_directory: str | Path | None = None,
|
|
|
|
|
requested_at: datetime | None = None,
|
|
|
|
|
) -> SourceRecordExportArchive:
|
|
|
|
|
"""Package selected prepared files without querying source-record tables."""
|
|
|
|
|
|
|
|
|
|
@@ -313,6 +324,13 @@ def build_source_records_export_archive(
|
|
|
|
|
generation = load_current_source_record_export_generation(
|
|
|
|
|
export_directory=root_directory,
|
|
|
|
|
)
|
|
|
|
|
requested_at_datetime = _normalize_generation_datetime(
|
|
|
|
|
requested_at or timezone.now()
|
|
|
|
|
)
|
|
|
|
|
if generation.export_year != _export_year(requested_at_datetime):
|
|
|
|
|
raise SourceRecordExportArtifactsUnavailable(
|
|
|
|
|
"Prepared source-record export belongs to a different calendar year."
|
|
|
|
|
)
|
|
|
|
|
artifacts_by_key: dict[
|
|
|
|
|
tuple[str, str],
|
|
|
|
|
list[SourceRecordExportArtifact],
|
|
|
|
|
@@ -338,10 +356,11 @@ def build_source_records_export_archive(
|
|
|
|
|
sorted(artifacts, key=lambda artifact: artifact.part_number)
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
generated_at = datetime.fromisoformat(generation.generated_at)
|
|
|
|
|
timestamp = generated_at.strftime("%Y%m%d_%H%M%S")
|
|
|
|
|
return SourceRecordExportArchive(
|
|
|
|
|
archive_name=f"organization_source_records_export_{timestamp}.zip",
|
|
|
|
|
archive_name=_build_source_records_archive_name(
|
|
|
|
|
source_groups=source_groups,
|
|
|
|
|
requested_at=requested_at_datetime,
|
|
|
|
|
),
|
|
|
|
|
archive_chunks=_stream_zip_archive(selected_artifacts),
|
|
|
|
|
files_count=len(selected_artifacts),
|
|
|
|
|
generated_at=generation.generated_at,
|
|
|
|
|
@@ -355,9 +374,11 @@ def create_source_record_export_download_ticket(
|
|
|
|
|
) -> SourceRecordExportDownloadTicket:
|
|
|
|
|
"""Validate prepared files and cache a short-lived download capability."""
|
|
|
|
|
|
|
|
|
|
requested_at = _normalize_generation_datetime(timezone.now())
|
|
|
|
|
package = build_source_records_export_archive(
|
|
|
|
|
source_groups=source_groups,
|
|
|
|
|
export_format=export_format,
|
|
|
|
|
requested_at=requested_at,
|
|
|
|
|
)
|
|
|
|
|
expires_in = max(
|
|
|
|
|
1,
|
|
|
|
|
@@ -372,6 +393,7 @@ def create_source_record_export_download_ticket(
|
|
|
|
|
payload = {
|
|
|
|
|
"sources": list(source_groups),
|
|
|
|
|
"format": export_format,
|
|
|
|
|
"requested_at": requested_at.isoformat(),
|
|
|
|
|
}
|
|
|
|
|
for _attempt in range(3):
|
|
|
|
|
ticket = secrets.token_urlsafe(32)
|
|
|
|
|
@@ -405,6 +427,7 @@ def consume_source_record_export_download_ticket(
|
|
|
|
|
try:
|
|
|
|
|
source_groups = payload["sources"]
|
|
|
|
|
export_format = payload["format"]
|
|
|
|
|
requested_at_value = payload["requested_at"]
|
|
|
|
|
if (
|
|
|
|
|
not isinstance(source_groups, list)
|
|
|
|
|
or not source_groups
|
|
|
|
|
@@ -415,14 +438,19 @@ def consume_source_record_export_download_ticket(
|
|
|
|
|
)
|
|
|
|
|
or len(source_groups) != len(set(source_groups))
|
|
|
|
|
or export_format not in EXPORT_FORMATS
|
|
|
|
|
or not isinstance(requested_at_value, str)
|
|
|
|
|
):
|
|
|
|
|
raise ValueError
|
|
|
|
|
requested_at = datetime.fromisoformat(requested_at_value)
|
|
|
|
|
if timezone.is_naive(requested_at):
|
|
|
|
|
raise ValueError
|
|
|
|
|
except (KeyError, TypeError, ValueError):
|
|
|
|
|
raise SourceRecordExportTicketInvalid from None
|
|
|
|
|
|
|
|
|
|
return build_source_records_export_archive(
|
|
|
|
|
source_groups=source_groups,
|
|
|
|
|
export_format=export_format,
|
|
|
|
|
requested_at=requested_at,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@@ -475,6 +503,23 @@ def _normalize_generation_datetime(value: datetime) -> datetime:
|
|
|
|
|
return value.astimezone(UTC)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _export_year(value: datetime) -> int:
|
|
|
|
|
"""Return the calendar year in the configured service timezone."""
|
|
|
|
|
return timezone.localtime(value).year
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _build_source_records_archive_name(
|
|
|
|
|
*,
|
|
|
|
|
source_groups: Sequence[str],
|
|
|
|
|
requested_at: datetime,
|
|
|
|
|
) -> str:
|
|
|
|
|
source_name = "__".join(
|
|
|
|
|
SOURCE_GROUP_EXPORT_FILE_STEMS[source_group] for source_group in source_groups
|
|
|
|
|
)
|
|
|
|
|
timestamp = timezone.localtime(requested_at).strftime("%Y%m%d_%H%M%S")
|
|
|
|
|
return f"{source_name}_{timestamp}.zip"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _source_group_export_formats(source_group: str) -> tuple[str, ...]:
|
|
|
|
|
if source_group == FINANCIAL_SOURCE_GROUP:
|
|
|
|
|
return (EXPORT_FORMAT_JSON,)
|
|
|
|
|
@@ -495,8 +540,12 @@ def _build_source_group_file_name(*, source_group: str, file_format: str) -> str
|
|
|
|
|
return f"{SOURCE_GROUP_EXPORT_FILE_STEMS[source_group]}.{file_format}"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _source_group_queryset(source_group: str) -> QuerySet[OrganizationSourceRecord]:
|
|
|
|
|
return (
|
|
|
|
|
def _source_group_queryset(
|
|
|
|
|
source_group: str,
|
|
|
|
|
*,
|
|
|
|
|
export_year: int,
|
|
|
|
|
) -> QuerySet[OrganizationSourceRecord]:
|
|
|
|
|
queryset = (
|
|
|
|
|
OrganizationSourceRecord.objects.filter(extension__source_group=source_group)
|
|
|
|
|
.select_related("extension", "extension__organization")
|
|
|
|
|
# Export order is not part of the file contract. Clearing the model's
|
|
|
|
|
@@ -504,22 +553,52 @@ def _source_group_queryset(source_group: str) -> QuerySet[OrganizationSourceReco
|
|
|
|
|
# source groups with millions of rows.
|
|
|
|
|
.order_by()
|
|
|
|
|
)
|
|
|
|
|
if source_group == FINANCIAL_SOURCE_GROUP:
|
|
|
|
|
return queryset.filter(financial_lines__year=export_year).distinct()
|
|
|
|
|
|
|
|
|
|
if connection.vendor == "postgresql":
|
|
|
|
|
queryset = queryset.annotate(
|
|
|
|
|
export_record_year=RawSQL(
|
|
|
|
|
"substring(record_date FROM '([0-9]{4})')::integer",
|
|
|
|
|
(),
|
|
|
|
|
)
|
|
|
|
|
)
|
|
|
|
|
return queryset.filter(
|
|
|
|
|
Q(export_record_year=export_year)
|
|
|
|
|
| Q(record_date="", created_at__year=export_year)
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
year_pattern = rf"(^|[^0-9]){export_year}([^0-9]|$)"
|
|
|
|
|
return queryset.filter(
|
|
|
|
|
Q(record_date__regex=year_pattern)
|
|
|
|
|
| Q(record_date="", created_at__year=export_year)
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _iter_source_records(
|
|
|
|
|
*,
|
|
|
|
|
source_group: str,
|
|
|
|
|
include_financial_lines: bool,
|
|
|
|
|
export_year: int,
|
|
|
|
|
) -> Iterator[OrganizationSourceRecord]:
|
|
|
|
|
iterator = _source_group_queryset(source_group).iterator(
|
|
|
|
|
chunk_size=SOURCE_RECORD_EXPORT_ITERATOR_CHUNK_SIZE
|
|
|
|
|
)
|
|
|
|
|
iterator = _source_group_queryset(
|
|
|
|
|
source_group,
|
|
|
|
|
export_year=export_year,
|
|
|
|
|
).iterator(chunk_size=SOURCE_RECORD_EXPORT_ITERATOR_CHUNK_SIZE)
|
|
|
|
|
while True:
|
|
|
|
|
batch = list(islice(iterator, SOURCE_RECORD_EXPORT_ITERATOR_CHUNK_SIZE))
|
|
|
|
|
if not batch:
|
|
|
|
|
return
|
|
|
|
|
if include_financial_lines:
|
|
|
|
|
prefetch_related_objects(batch, "financial_lines")
|
|
|
|
|
prefetch_related_objects(
|
|
|
|
|
batch,
|
|
|
|
|
Prefetch(
|
|
|
|
|
"financial_lines",
|
|
|
|
|
queryset=OrganizationSourceFinancialLine.objects.filter(
|
|
|
|
|
year=export_year
|
|
|
|
|
).order_by(),
|
|
|
|
|
),
|
|
|
|
|
)
|
|
|
|
|
yield from batch
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@@ -527,6 +606,7 @@ def _spool_source_group_rows(
|
|
|
|
|
*,
|
|
|
|
|
source_group: str,
|
|
|
|
|
output_path: Path,
|
|
|
|
|
export_year: int,
|
|
|
|
|
) -> tuple[list[str], int]:
|
|
|
|
|
include_financial_lines = source_group == FINANCIAL_SOURCE_GROUP
|
|
|
|
|
payload_headers: set[str] = set()
|
|
|
|
|
@@ -538,6 +618,7 @@ def _spool_source_group_rows(
|
|
|
|
|
for record in _iter_source_records(
|
|
|
|
|
source_group=source_group,
|
|
|
|
|
include_financial_lines=include_financial_lines,
|
|
|
|
|
export_year=export_year,
|
|
|
|
|
):
|
|
|
|
|
row = _build_record_row(
|
|
|
|
|
record,
|
|
|
|
|
@@ -799,6 +880,7 @@ def _generation_manifest_payload(
|
|
|
|
|
"version": EXPORT_MANIFEST_VERSION,
|
|
|
|
|
"generation_id": generation.generation_id,
|
|
|
|
|
"generated_at": generation.generated_at,
|
|
|
|
|
"export_year": generation.export_year,
|
|
|
|
|
"records_count": generation.records_count,
|
|
|
|
|
"artifacts_count": generation.artifacts_count,
|
|
|
|
|
"files_count": generation.files_count,
|
|
|
|
|
@@ -830,6 +912,9 @@ def _generation_from_manifest(
|
|
|
|
|
generation_id = str(payload["generation_id"])
|
|
|
|
|
generated_at = str(payload["generated_at"])
|
|
|
|
|
datetime.fromisoformat(generated_at)
|
|
|
|
|
export_year = int(payload["export_year"])
|
|
|
|
|
if not 1 <= export_year <= 9999:
|
|
|
|
|
raise ValueError("Source-record export year is invalid.")
|
|
|
|
|
records_count = int(payload["records_count"])
|
|
|
|
|
artifact_payloads = payload["artifacts"]
|
|
|
|
|
if not isinstance(artifact_payloads, list):
|
|
|
|
|
@@ -884,6 +969,7 @@ def _generation_from_manifest(
|
|
|
|
|
return SourceRecordExportGeneration(
|
|
|
|
|
generation_id=generation_id,
|
|
|
|
|
generated_at=generated_at,
|
|
|
|
|
export_year=export_year,
|
|
|
|
|
artifacts=tuple(artifacts),
|
|
|
|
|
records_count=records_count,
|
|
|
|
|
)
|
|
|
|
|
|