gizmoduck/finance/exports.py

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,
    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"),
    ),
    (
        "categories.csv",
        Category,
        ("id", "name", "kind", "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")

    reservation_fd = None
    owns_reservation = False
    temporary_path = None
    published = False
    phase = "reservation"
    try:
        try:
            reservation_fd = os.open(
                destination,
                os.O_WRONLY | os.O_CREAT | os.O_EXCL,
                0o600,
            )
            owns_reservation = True
        except FileExistsError:
            raise ExportError("target_exists") from None
        finally:
            if reservation_fd is not None:
                os.close(reservation_fd)
                reservation_fd = None

        phase = "snapshot"
        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())
        os.replace(temporary_path, destination)
        temporary_path = None
        published = True
        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
        if owns_reservation and not published:
            try:
                destination.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)