Stonks/overview/services.py

from dataclasses import dataclass

from django.db import transaction
from django.utils import timezone

from .brokerage_types import BrokerageSnapshot
from .brokerages.base import BaseBrokerageAdapter
from .brokerages.registry import get_brokerage_adapter
from .models import (
    Account,
    AccountHolding,
    AccountSnapshot,
    AccountTransaction,
    Asset,
    AssetPriceSnapshot,
    BrokerageConnection,
    HoldingSnapshot,
    IngestionRun,
)


@dataclass(frozen=True)
class IngestionResult:
    accounts_processed: int
    assets_processed: int
    holdings_processed: int
    transactions_processed: int


class BrokerageIngestionService:
    def __init__(
        self,
        connection: BrokerageConnection,
        adapter: BaseBrokerageAdapter | None = None,
    ):
        self.connection = connection
        self.adapter = adapter or get_brokerage_adapter(connection)

    def ingest(self) -> IngestionRun:
        run = IngestionRun.objects.create(
            connection=self.connection,
            status=IngestionRun.STATUS_RUNNING,
            started_at=timezone.now(),
        )

        try:
            snapshot = self.adapter.fetch_snapshot()
            result = self._persist_snapshot(snapshot, run)
        except Exception as exc:
            run.status = IngestionRun.STATUS_FAILED
            run.error_message = str(exc)
            run.finished_at = timezone.now()
            run.save(
                update_fields=["status", "error_message", "finished_at"]
            )
            raise

        run.status = IngestionRun.STATUS_SUCCEEDED
        run.finished_at = timezone.now()
        run.accounts_processed = result.accounts_processed
        run.assets_processed = result.assets_processed
        run.holdings_processed = result.holdings_processed
        run.transactions_processed = result.transactions_processed
        run.save(
            update_fields=[
                "status",
                "finished_at",
                "accounts_processed",
                "assets_processed",
                "holdings_processed",
                "transactions_processed",
            ]
        )
        return run

    @transaction.atomic
    def _persist_snapshot(
        self,
        snapshot: BrokerageSnapshot,
        run: IngestionRun,
    ) -> IngestionResult:
        accounts_by_external_id = {}
        assets_by_external_id = {}
        snapshot_at = run.started_at or timezone.now()

        for payload in snapshot.accounts:
            account, _ = Account.objects.update_or_create(
                provider=self.connection.provider,
                external_id=payload.external_id,
                defaults={
                    "name": payload.name,
                    "number": payload.number,
                    "balance": payload.balance,
                },
            )
            accounts_by_external_id[payload.external_id] = account
            AccountSnapshot.objects.update_or_create(
                ingestion_run=run,
                account=account,
                defaults={
                    "balance": payload.balance,
                    "snapshot_at": snapshot_at,
                },
            )

        for payload in snapshot.assets:
            asset, _ = Asset.objects.update_or_create(
                provider=self.connection.provider,
                external_id=payload.external_id,
                defaults={
                    "ticker": payload.ticker,
                    "price": payload.price,
                },
            )
            assets_by_external_id[payload.external_id] = asset
            AssetPriceSnapshot.objects.update_or_create(
                ingestion_run=run,
                asset=asset,
                defaults={
                    "price": payload.price,
                    "snapshot_at": snapshot_at,
                },
            )

        for payload in snapshot.holdings:
            account = accounts_by_external_id[payload.account_external_id]
            asset = assets_by_external_id[payload.asset_external_id]
            holding, _ = AccountHolding.objects.update_or_create(
                provider=self.connection.provider,
                external_id=payload.external_id,
                defaults={
                    "account": account,
                    "asset": asset,
                    "ticker": payload.ticker,
                    "ammount": payload.amount,
                    "cost_basis": payload.cost_basis,
                    "purchase_date": payload.purchase_date,
                },
            )
            HoldingSnapshot.objects.update_or_create(
                ingestion_run=run,
                holding=holding,
                defaults={
                    "account": account,
                    "asset": asset,
                    "amount": payload.amount,
                    "cost_basis": payload.cost_basis,
                    "market_price": asset.price,
                    "value": payload.amount * asset.price,
                    "snapshot_at": snapshot_at,
                },
            )

        for payload in snapshot.transactions:
            account = accounts_by_external_id[payload.account_external_id]
            asset = None
            if payload.asset_external_id:
                asset = assets_by_external_id[payload.asset_external_id]

            AccountTransaction.objects.update_or_create(
                provider=self.connection.provider,
                external_id=payload.external_id,
                defaults={
                    "account": account,
                    "asset": asset,
                    "transaction_type": payload.transaction_type,
                    "date": payload.date,
                    "value": payload.value,
                    "name": payload.name,
                    "quantity": payload.quantity,
                    "realized_gain": payload.realized_gain,
                },
            )

        return IngestionResult(
            accounts_processed=len(snapshot.accounts),
            assets_processed=len(snapshot.assets),
            holdings_processed=len(snapshot.holdings),
            transactions_processed=len(snapshot.transactions),
        )