feat: export organization source records

src/organizations/serializers.py

@@ -8,7 +8,9 @@ from organizations.models import (
     OrganizationSourceExtension,
     OrganizationSourceFinancialLine,
     OrganizationSourceRecord,
+    SourceGroup,
 )
+from organizations.source_record_export import EXPORT_FORMATS
 
 
 class OrganizationSourceFinancialLineSerializer(serializers.ModelSerializer):
@@ -127,6 +129,25 @@ class OrganizationSourceRecordListResponseSerializer(serializers.Serializer):
     meta = serializers.JSONField(read_only=True, allow_null=True)
 
 
+class OrganizationSourceRecordExportRequestSerializer(serializers.Serializer):
+    """Request for source-record ZIP export."""
+
+    sources = serializers.ListField(
+        child=serializers.ChoiceField(choices=SourceGroup.choices),
+        allow_empty=False,
+    )
+    format = serializers.ChoiceField(
+        choices=[
+            (export_format, export_format.upper()) for export_format in EXPORT_FORMATS
+        ],
+    )
+
+    def validate_sources(self, value: list[str]) -> list[str]:
+        if len(value) != len(set(value)):
+            raise serializers.ValidationError("Источники не должны повторяться.")
+        return value
+
+
 class OrganizationSourceExtensionSerializer(serializers.ModelSerializer):
     """Compact source extension representation."""
 

src/organizations/source_record_export.py

@@ -0,0 +1,316 @@
+"""ZIP export for organization source records."""
+
+from __future__ import annotations
+
+import csv
+import json
+import zipfile
+from collections.abc import Sequence
+from dataclasses import dataclass
+from datetime import date, datetime
+from decimal import Decimal
+from io import BytesIO, StringIO
+from typing import Any
+
+from django.db.models import QuerySet
+from django.utils import timezone
+from openpyxl import Workbook
+
+from organizations.models import 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
+
+SOURCE_GROUP_EXPORT_FILE_STEMS: dict[str, str] = {
+    SourceGroup.FINANCIAL_INDICATORS.value: "financial-indicators",
+    SourceGroup.GOVERNMENT_PROCUREMENTS.value: "public-procurements",
+    SourceGroup.INDUSTRIAL_PRODUCTION.value: "manufacturers-and-products",
+    SourceGroup.PLANNED_INSPECTIONS.value: "planned-inspections",
+    SourceGroup.BANKRUPTCY.value: "bankruptcy-procedures",
+    SourceGroup.DEFENSE_SUPPLIERS.value: "defense-unreliable-suppliers",
+    SourceGroup.ARBITRATION.value: "arbitration-cases",
+    SourceGroup.SECURITY_REGISTRIES.value: "information-security-registries",
+    SourceGroup.VACANCIES.value: "labor-vacancies",
+}
+
+ORGANIZATION_EXPORT_FIELDS = ["Наименование", "ИНН", "ОГРН", "КПП"]
+SOURCE_RECORD_EXPORT_FIELDS = [
+    "uid",
+    "source_group",
+    "source",
+    "record_type",
+    "external_id",
+    "title",
+    "record_date",
+    "amount",
+    "status",
+    "url",
+    "load_batch",
+    "created_at",
+    "updated_at",
+]
+FINANCIAL_LINE_FIELDS = [
+    "id",
+    "form_code",
+    "line_code",
+    "line_name",
+    "year",
+    "period_start",
+    "period_end",
+]
+
+
+@dataclass(frozen=True)
+class SourceRecordExportArchive:
+    """In-memory source records export archive."""
+
+    archive_name: str
+    archive_bytes: bytes
+    files_count: int
+
+
+def build_source_records_export_archive(
+    *,
+    source_groups: Sequence[str],
+    export_format: str,
+    now: datetime | None = None,
+) -> SourceRecordExportArchive:
+    """Build a ZIP archive with one source-record file per source group."""
+
+    timestamp = (now or timezone.now()).strftime("%Y%m%d_%H%M%S")
+    archive_buffer = BytesIO()
+
+    with zipfile.ZipFile(
+        archive_buffer,
+        mode="w",
+        compression=zipfile.ZIP_DEFLATED,
+    ) as archive:
+        for source_group in source_groups:
+            file_format = _resolve_source_group_export_format(
+                source_group=source_group,
+                requested_format=export_format,
+            )
+            file_name = _build_source_group_file_name(
+                source_group=source_group,
+                file_format=file_format,
+            )
+            archive.writestr(
+                file_name,
+                _render_source_group_file(
+                    source_group=source_group,
+                    file_format=file_format,
+                ),
+            )
+
+    return SourceRecordExportArchive(
+        archive_name=f"organization_source_records_export_{timestamp}.zip",
+        archive_bytes=archive_buffer.getvalue(),
+        files_count=len(source_groups),
+    )
+
+
+def _resolve_source_group_export_format(
+    *,
+    source_group: str,
+    requested_format: str,
+) -> str:
+    if source_group == FINANCIAL_SOURCE_GROUP:
+        return EXPORT_FORMAT_JSON
+    return requested_format
+
+
+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 (
+        OrganizationSourceRecord.objects.filter(extension__source_group=source_group)
+        .select_related("extension", "extension__organization")
+        .prefetch_related("financial_lines")
+        .order_by(
+            "extension__organization__name",
+            "source",
+            "record_date",
+            "uid",
+        )
+    )
+
+
+def _render_source_group_file(*, source_group: str, file_format: str) -> bytes:
+    queryset = _source_group_queryset(source_group)
+    payload_headers = _collect_payload_headers(queryset)
+    headers = [
+        *ORGANIZATION_EXPORT_FIELDS,
+        *SOURCE_RECORD_EXPORT_FIELDS,
+        *payload_headers,
+    ]
+    include_financial_lines = source_group == FINANCIAL_SOURCE_GROUP
+
+    if file_format == EXPORT_FORMAT_CSV:
+        return _render_csv_file(
+            queryset=_source_group_queryset(source_group),
+            headers=headers,
+            payload_headers=payload_headers,
+        )
+    if file_format == EXPORT_FORMAT_XLSX:
+        return _render_xlsx_file(
+            queryset=_source_group_queryset(source_group),
+            headers=headers,
+            payload_headers=payload_headers,
+        )
+    return _render_json_file(
+        queryset=_source_group_queryset(source_group),
+        headers=headers,
+        payload_headers=payload_headers,
+        include_financial_lines=include_financial_lines,
+    )
+
+
+def _collect_payload_headers(
+    queryset: QuerySet[OrganizationSourceRecord],
+) -> list[str]:
+    payload_headers: set[str] = set()
+
+    for record in queryset.iterator(chunk_size=1000):
+        payload_headers.update(_flatten_payload(record.payload).keys())
+
+    return sorted(payload_headers)
+
+
+def _render_csv_file(
+    *,
+    queryset: QuerySet[OrganizationSourceRecord],
+    headers: Sequence[str],
+    payload_headers: Sequence[str],
+) -> bytes:
+    output = StringIO()
+    writer = csv.DictWriter(output, fieldnames=list(headers), lineterminator="\n")
+    writer.writeheader()
+
+    for record in queryset.iterator(chunk_size=1000):
+        row = _build_record_row(record, payload_headers=payload_headers)
+        writer.writerow({key: _serialize_flat_value(row.get(key)) for key in headers})
+
+    return ("\ufeff" + output.getvalue()).encode("utf-8")
+
+
+def _render_xlsx_file(
+    *,
+    queryset: QuerySet[OrganizationSourceRecord],
+    headers: Sequence[str],
+    payload_headers: Sequence[str],
+) -> bytes:
+    workbook = Workbook(write_only=True)
+    worksheet = workbook.create_sheet(title="data")
+    worksheet.append(list(headers))
+
+    for record in queryset.iterator(chunk_size=1000):
+        row = _build_record_row(record, payload_headers=payload_headers)
+        worksheet.append([_serialize_flat_value(row.get(key)) for key in headers])
+
+    output = BytesIO()
+    workbook.save(output)
+    return output.getvalue()
+
+
+def _render_json_file(
+    *,
+    queryset: QuerySet[OrganizationSourceRecord],
+    headers: Sequence[str],
+    payload_headers: Sequence[str],
+    include_financial_lines: bool,
+) -> bytes:
+    rows = []
+
+    for record in queryset.iterator(chunk_size=1000):
+        row = _build_record_row(record, payload_headers=payload_headers)
+        json_row = {key: _serialize_json_value(row.get(key)) for key in headers}
+        if include_financial_lines:
+            json_row["financial_lines"] = [
+                {
+                    field_name: _serialize_json_value(
+                        getattr(financial_line, field_name)
+                    )
+                    for field_name in FINANCIAL_LINE_FIELDS
+                }
+                for financial_line in record.financial_lines.all()
+            ]
+        rows.append(json_row)
+
+    return json.dumps(rows, ensure_ascii=False, indent=2).encode("utf-8")
+
+
+def _build_record_row(
+    record: OrganizationSourceRecord,
+    *,
+    payload_headers: Sequence[str],
+) -> dict[str, Any]:
+    organization = record.extension.organization
+    payload = _flatten_payload(record.payload)
+
+    row: dict[str, Any] = {
+        "Наименование": organization.full_name
+        or organization.short_name
+        or organization.name,
+        "ИНН": organization.inn,
+        "ОГРН": organization.ogrn,
+        "КПП": organization.kpp,
+        "uid": str(record.uid),
+        "source_group": record.extension.source_group,
+        "source": record.source,
+        "record_type": record.record_type,
+        "external_id": record.external_id,
+        "title": record.title,
+        "record_date": record.record_date,
+        "amount": record.amount,
+        "status": record.status,
+        "url": record.url,
+        "load_batch": record.load_batch,
+        "created_at": record.created_at,
+        "updated_at": record.updated_at,
+    }
+
+    for header in payload_headers:
+        row[header] = payload.get(header)
+
+    return row
+
+
+def _flatten_payload(value: Any, *, prefix: str = "payload") -> dict[str, Any]:
+    if not isinstance(value, dict):
+        return {prefix: value} if value not in (None, "") else {}
+
+    flattened: dict[str, Any] = {}
+    for key, item in value.items():
+        item_key = f"{prefix}.{key}"
+        if isinstance(item, dict) and item:
+            flattened.update(_flatten_payload(item, prefix=item_key))
+        else:
+            flattened[item_key] = item
+    return flattened
+
+
+def _serialize_flat_value(value: Any) -> str | int | float | bool:
+    if value is None:
+        return ""
+    if isinstance(value, dict | list | tuple):
+        return json.dumps(value, ensure_ascii=False, sort_keys=True)
+    if isinstance(value, Decimal):
+        return str(value)
+    if isinstance(value, datetime | date):
+        return value.isoformat()
+    return value
+
+
+def _serialize_json_value(value: Any) -> Any:
+    if isinstance(value, dict | list | tuple):
+        return json.dumps(value, ensure_ascii=False, sort_keys=True)
+    if isinstance(value, Decimal):
+        return str(value)
+    if isinstance(value, datetime | date):
+        return value.isoformat()
+    return value

src/organizations/views.py

@@ -11,12 +11,14 @@ from django.conf import settings
 from django.core.cache import cache
 from django.db.models import CharField, Q
 from django.db.models.functions import Cast
+from django.http import HttpResponse
 from django_filters import rest_framework as filters
 from drf_yasg import openapi
 from drf_yasg.utils import swagger_auto_schema
+from rest_framework import status
 from rest_framework.decorators import action
 from rest_framework.filters import OrderingFilter, SearchFilter
-from rest_framework.permissions import AllowAny, IsAuthenticated
+from rest_framework.permissions import AllowAny, IsAdminUser, IsAuthenticated
 from rest_framework.response import Response
 from rest_framework.viewsets import ReadOnlyModelViewSet
 
@@ -35,9 +37,11 @@ from organizations.models import (
 from organizations.serializers import (
     OrganizationSerializer,
     OrganizationSourceExtensionSerializer,
+    OrganizationSourceRecordExportRequestSerializer,
     OrganizationSourceRecordListResponseSerializer,
     OrganizationSourceRecordSerializer,
 )
+from organizations.source_record_export import build_source_records_export_archive
 
 ORGANIZATIONS_TAG = swagger_tag("Организации", "Organizations")
 FALSE_QUERY_VALUES = {"0", "false", "no", "off"}
@@ -449,6 +453,8 @@ class OrganizationSourceRecordViewSet(ReadOnlyModelViewSet):
     ordering = ["-created_at", "-uid"]
 
     def get_permissions(self):
+        if self.action == "export":
+            return [IsAdminUser()]
         if getattr(settings, "ORGANIZATIONS_V2_ALLOW_ANONYMOUS", False):
             return [AllowAny()]
         return super().get_permissions()
@@ -553,3 +559,39 @@ class OrganizationSourceRecordViewSet(ReadOnlyModelViewSet):
     )
     def list(self, request, *args: Any, **kwargs: Any) -> Response:
         return super().list(request, *args, **kwargs)
+
+    @swagger_auto_schema(
+        tags=[ORGANIZATIONS_TAG],
+        operation_id="v2_organization_source_records_export",
+        operation_summary="Выгрузить записи источников",
+        operation_description=(
+            "Формирует ZIP-архив с одним файлом на выбранную группу источников. "
+            "Финансово-экономические показатели всегда выгружаются в JSON."
+        ),
+        request_body=OrganizationSourceRecordExportRequestSerializer,
+        responses={
+            200: openapi.Response(
+                description="ZIP-архив записей источников.",
+                schema=openapi.Schema(type=openapi.TYPE_FILE),
+            ),
+            400: "Некорректные параметры выгрузки.",
+            403: "Доступ разрешён только администраторам.",
+        },
+    )
+    @action(detail=False, methods=["post"], url_path="export")
+    def export(self, request, *args: Any, **kwargs: Any) -> HttpResponse:
+        serializer = OrganizationSourceRecordExportRequestSerializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+
+        package = build_source_records_export_archive(
+            source_groups=serializer.validated_data["sources"],
+            export_format=serializer.validated_data["format"],
+        )
+        response = HttpResponse(package.archive_bytes, content_type="application/zip")
+        response.status_code = status.HTTP_200_OK
+        response[
+            "Content-Disposition"
+        ] = f'attachment; filename="{package.archive_name}"'
+        response["Content-Length"] = str(len(package.archive_bytes))
+        response["X-Source-Export-Files"] = str(package.files_count)
+        return response