import csv import json import os import re import tempfile from dataclasses import dataclass from datetime import date, datetime, timezone as datetime_timezone from decimal import Decimal from hashlib import sha256 from io import BytesIO, StringIO from pathlib import Path from uuid import UUID from zipfile import ZIP_DEFLATED, ZipFile, ZipInfo from django.conf import settings from django.db import connection, transaction from django.utils import timezone from finance.models import ( Account, CardStatement, Category, CategoryGroup, InstallmentPlan, Payee, SourceRecord, Transaction, Transfer, ) EXPORT_FORMAT = "gizmoduck-export" EXPORT_VERSION = 1 APPLICATION_VERSION = "0.1.0" @dataclass(frozen=True) class ExportArtifact: filename: str content: bytes version: int row_counts: dict[str, int] @dataclass(frozen=True) class ExportResult: name: str path: Path version: int row_counts: dict[str, int] class ExportError(Exception): def __init__(self, code): self.code = code super().__init__(code) _PLAIN_ZIP_NAME = re.compile(r"\A[A-Za-z0-9][A-Za-z0-9._-]{0,198}\.zip\Z") SCHEMAS = ( ( "accounts.csv", Account, ("id", "name", "kind", "currency", "archived", "created_at", "updated_at"), ), ( "payees.csv", Payee, ("id", "name", "archived", "created_at", "updated_at"), ), ( "category_groups.csv", CategoryGroup, ("id", "name", "archived", "created_at", "updated_at"), ), ( "categories.csv", Category, ("id", "name", "kind", "group_id", "archived", "created_at", "updated_at"), ), ( "statements.csv", CardStatement, ( "id", "account_id", "period_start", "period_end", "due_date", "statement_balance", "created_at", "updated_at", ), ), ( "installment_plans.csv", InstallmentPlan, ("id", "account_id", "description", "total_installments", "created_at", "updated_at"), ), ( "transactions.csv", Transaction, ( "id", "account_id", "amount", "posting_date", "occurrence_date", "description", "kind", "payee_id", "category_id", "statement_id", "refund_of_id", "installment_plan_id", "installment_number", "created_at", "updated_at", ), ), ( "transfers.csv", Transfer, ("id", "outgoing_id", "incoming_id", "created_at", "updated_at"), ), ( "source_records.csv", SourceRecord, ( "id", "transaction_id", "document_hash", "source_format", "parser_version", "item_locator", "created_at", ), ), ) def build_export() -> ExportArtifact: with transaction.atomic(): with connection.cursor() as cursor: cursor.execute("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ") created_at = timezone.now() csv_files = { filename: _csv_content(model, fields) for filename, model, fields in SCHEMAS } file_manifest = { filename: { "rows": row_count, "sha256": sha256(content).hexdigest(), } for (filename, _model, _fields), (content, row_count) in zip( SCHEMAS, csv_files.values(), strict=True ) } manifest = { "format": EXPORT_FORMAT, "version": EXPORT_VERSION, "created_at": _value(created_at), "application_version": APPLICATION_VERSION, "files": file_manifest, } package = BytesIO() with ZipFile(package, "w") as archive: for filename, (content, _row_count) in csv_files.items(): _write_member(archive, filename, content) manifest_content = ( json.dumps(manifest, ensure_ascii=False, separators=(",", ":")) + "\n" ).encode("utf-8") _write_member(archive, "manifest.json", manifest_content) return ExportArtifact( filename=f"gizmoduck-export-{created_at.astimezone(datetime_timezone.utc):%Y%m%dT%H%M%SZ}.zip", content=package.getvalue(), version=EXPORT_VERSION, row_counts={name: details["rows"] for name, details in file_manifest.items()}, ) def write_export(name: str | None = None) -> ExportResult: if name is None: name = f"gizmoduck-export-{timezone.now().astimezone(datetime_timezone.utc):%Y%m%dT%H%M%SZ}.zip" if not isinstance(name, str) or not _PLAIN_ZIP_NAME.fullmatch(name): raise ExportError("invalid_name") try: root = Path(settings.EXPORT_ROOT).resolve(strict=True) except (OSError, RuntimeError, TypeError): raise ExportError("invalid_root") from None if not root.is_dir(): raise ExportError("invalid_root") destination = root / name if destination.is_symlink(): raise ExportError("invalid_name") if destination.exists(): raise ExportError("existing_destination") temporary_path = None phase = "snapshot" try: artifact = build_export() phase = "publication" temporary_fd, temporary_name = tempfile.mkstemp( prefix=".gizmoduck-export-", suffix=".tmp", dir=root ) temporary_path = Path(temporary_name) with os.fdopen(temporary_fd, "wb") as stream: stream.write(artifact.content) stream.flush() os.fsync(stream.fileno()) try: os.link(temporary_path, destination) except FileExistsError: raise ExportError("existing_destination") from None temporary_path.unlink() temporary_path = None return ExportResult( name=name, path=destination, version=artifact.version, row_counts=artifact.row_counts, ) except ExportError: raise except Exception: code = "export_failed" if phase == "snapshot" else "publication_failed" raise ExportError(code) from None finally: if temporary_path is not None: try: temporary_path.unlink(missing_ok=True) except OSError: pass def _csv_content(model, fields): stream = StringIO(newline="") writer = csv.writer(stream, lineterminator="\n") writer.writerow(fields) row_count = 0 for row in model.objects.order_by("id").values_list(*fields): writer.writerow(_value(value) for value in row) row_count += 1 return stream.getvalue().encode("utf-8"), row_count def _value(value): if value is None: return "" if isinstance(value, bool): return str(value).lower() if isinstance(value, Decimal): return format(value, "f") if isinstance(value, UUID): return str(value) if isinstance(value, (date, datetime)): return value.isoformat() return str(value) def _write_member(archive, filename, content): info = ZipInfo(filename, date_time=(1980, 1, 1, 0, 0, 0)) info.compress_type = ZIP_DEFLATED info.create_system = 3 info.external_attr = 0o600 << 16 archive.writestr(info, content, compress_type=ZIP_DEFLATED, compresslevel=9)