112 lines
3.8 KiB
Python
112 lines
3.8 KiB
Python
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, Asset, BrokerageConnection, IngestionRun
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class IngestionResult:
|
|
accounts_processed: int
|
|
assets_processed: int
|
|
holdings_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)
|
|
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.save(
|
|
update_fields=[
|
|
"status",
|
|
"finished_at",
|
|
"accounts_processed",
|
|
"assets_processed",
|
|
"holdings_processed",
|
|
]
|
|
)
|
|
return run
|
|
|
|
@transaction.atomic
|
|
def _persist_snapshot(self, snapshot: BrokerageSnapshot) -> IngestionResult:
|
|
accounts_by_external_id = {}
|
|
assets_by_external_id = {}
|
|
|
|
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
|
|
|
|
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
|
|
|
|
for payload in snapshot.holdings:
|
|
account = accounts_by_external_id[payload.account_external_id]
|
|
asset = assets_by_external_id[payload.asset_external_id]
|
|
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,
|
|
},
|
|
)
|
|
|
|
return IngestionResult(
|
|
accounts_processed=len(snapshot.accounts),
|
|
assets_processed=len(snapshot.assets),
|
|
holdings_processed=len(snapshot.holdings),
|
|
)
|