From 618072069ef1d3dfab26060e711e0a6512745f61 Mon Sep 17 00:00:00 2001 From: Tyler Wilcox Date: Sat, 23 May 2026 21:18:05 -0400 Subject: [PATCH] AI: Brokerage Ingestion infra --- overview/admin.py | 4 +- overview/brokerage_types.py | 35 ++ overview/brokerages/__init__.py | 11 + overview/brokerages/base.py | 19 + overview/brokerages/random.py | 172 +++++++++ overview/brokerages/registry.py | 34 ++ overview/brokerages/rest.py | 131 +++++++ overview/management/__init__.py | 1 + overview/management/commands/__init__.py | 1 + .../management/commands/sync_brokerage.py | 70 ++++ ...gestionrun_account_external_id_and_more.py | 96 +++++ ...okerageconnection_access_token_and_more.py | 43 +++ overview/models.py | 101 +++++- overview/services.py | 111 ++++++ overview/tests.py | 332 +++++++++++++++++- 15 files changed, 1157 insertions(+), 4 deletions(-) create mode 100644 overview/brokerage_types.py create mode 100644 overview/brokerages/__init__.py create mode 100644 overview/brokerages/base.py create mode 100644 overview/brokerages/random.py create mode 100644 overview/brokerages/registry.py create mode 100644 overview/brokerages/rest.py create mode 100644 overview/management/__init__.py create mode 100644 overview/management/commands/__init__.py create mode 100644 overview/management/commands/sync_brokerage.py create mode 100644 overview/migrations/0004_brokerageconnection_ingestionrun_account_external_id_and_more.py create mode 100644 overview/migrations/0005_remove_brokerageconnection_access_token_and_more.py create mode 100644 overview/services.py diff --git a/overview/admin.py b/overview/admin.py index 770d475..3610578 100644 --- a/overview/admin.py +++ b/overview/admin.py @@ -1,8 +1,10 @@ from django.contrib import admin -from .models import Account, AccountHolding, Asset +from .models import Account, AccountHolding, Asset, BrokerageConnection, IngestionRun # Register your models here. admin.site.register(Account) admin.site.register(AccountHolding) admin.site.register(Asset) +admin.site.register(BrokerageConnection) +admin.site.register(IngestionRun) diff --git a/overview/brokerage_types.py b/overview/brokerage_types.py new file mode 100644 index 0000000..68cfe9e --- /dev/null +++ b/overview/brokerage_types.py @@ -0,0 +1,35 @@ +from dataclasses import dataclass, field +from datetime import datetime + + +@dataclass(frozen=True) +class BrokerageAccountPayload: + external_id: str + name: str + number: str + balance: int + + +@dataclass(frozen=True) +class BrokerageAssetPayload: + external_id: str + ticker: str + price: float + + +@dataclass(frozen=True) +class BrokerageHoldingPayload: + external_id: str + account_external_id: str + asset_external_id: str + ticker: str + amount: float + cost_basis: float + purchase_date: datetime + + +@dataclass(frozen=True) +class BrokerageSnapshot: + accounts: list[BrokerageAccountPayload] = field(default_factory=list) + assets: list[BrokerageAssetPayload] = field(default_factory=list) + holdings: list[BrokerageHoldingPayload] = field(default_factory=list) diff --git a/overview/brokerages/__init__.py b/overview/brokerages/__init__.py new file mode 100644 index 0000000..4d134e0 --- /dev/null +++ b/overview/brokerages/__init__.py @@ -0,0 +1,11 @@ +from .base import BaseBrokerageAdapter, BrokerageAdapterError +from .registry import get_brokerage_adapter, register_brokerage_adapter +from .random import RandomBrokerageAdapter + +__all__ = [ + "BaseBrokerageAdapter", + "BrokerageAdapterError", + "RandomBrokerageAdapter", + "get_brokerage_adapter", + "register_brokerage_adapter", +] diff --git a/overview/brokerages/base.py b/overview/brokerages/base.py new file mode 100644 index 0000000..f7ef37c --- /dev/null +++ b/overview/brokerages/base.py @@ -0,0 +1,19 @@ +from abc import ABC, abstractmethod + +from overview.brokerage_types import BrokerageSnapshot +from overview.models import BrokerageConnection + + +class BrokerageAdapterError(Exception): + pass + + +class BaseBrokerageAdapter(ABC): + provider = "" + + def __init__(self, connection: BrokerageConnection): + self.connection = connection + + @abstractmethod + def fetch_snapshot(self) -> BrokerageSnapshot: + raise NotImplementedError diff --git a/overview/brokerages/random.py b/overview/brokerages/random.py new file mode 100644 index 0000000..c02509b --- /dev/null +++ b/overview/brokerages/random.py @@ -0,0 +1,172 @@ +import random +from datetime import datetime, timedelta + +from overview.brokerage_types import ( + BrokerageAccountPayload, + BrokerageAssetPayload, + BrokerageHoldingPayload, + BrokerageSnapshot, +) + +from .base import BaseBrokerageAdapter + + +class RandomBrokerageAdapter(BaseBrokerageAdapter): + provider = "random-demo" + + DEFAULT_TICKERS = [ + "AAPL", + "MSFT", + "GOOGL", + "AMZN", + "NVDA", + "TSLA", + "META", + "BRK.B", + "JPM", + "V", + "NFLX", + "AMD", + ] + + def fetch_snapshot(self) -> BrokerageSnapshot: + generator = random.Random(self.connection.get_setting("seed")) + account_count = max(1, int(self.connection.get_setting("account_count", 2))) + asset_count = max(1, int(self.connection.get_setting("asset_count", 8))) + holdings_per_account = max( + 1, + int(self.connection.get_setting("holdings_per_account", 4)), + ) + reference_timestamp = self.connection.get_setting( + "reference_timestamp", + "2025-01-01T00:00:00+00:00", + ) + reference_time = datetime.fromisoformat( + str(reference_timestamp).replace("Z", "+00:00") + ) + tickers = self.connection.get_setting("tickers", self.DEFAULT_TICKERS) + + if not isinstance(tickers, list) or not tickers: + tickers = self.DEFAULT_TICKERS + + assets = self._build_assets(generator, asset_count, tickers) + accounts = self._build_accounts(generator, account_count) + holdings = self._build_holdings( + generator, + accounts=accounts, + assets=assets, + holdings_per_account=holdings_per_account, + reference_time=reference_time, + ) + + return BrokerageSnapshot( + accounts=accounts, + assets=assets, + holdings=holdings, + ) + + def _build_accounts( + self, + generator: random.Random, + account_count: int, + ) -> list[BrokerageAccountPayload]: + account_labels = [ + "Individual", + "Joint", + "Retirement", + "Roth IRA", + "Trust", + "Education", + ] + accounts = [] + + for index in range(account_count): + suffix = generator.randint(1000, 9999) + name = account_labels[index % len(account_labels)] + accounts.append( + BrokerageAccountPayload( + external_id=f"acct-{index + 1}", + name=f"{name} Account {index + 1}", + number=f"{suffix}{generator.randint(10, 99)}", + balance=0, + ) + ) + + return accounts + + def _build_assets( + self, + generator: random.Random, + asset_count: int, + tickers: list[str], + ) -> list[BrokerageAssetPayload]: + assets = [] + + for index in range(asset_count): + ticker = tickers[index % len(tickers)] + base_price = round(generator.uniform(25, 750), 2) + assets.append( + BrokerageAssetPayload( + external_id=f"asset-{index + 1}", + ticker=ticker, + price=base_price, + ) + ) + + return assets + + def _build_holdings( + self, + generator: random.Random, + accounts: list[BrokerageAccountPayload], + assets: list[BrokerageAssetPayload], + holdings_per_account: int, + reference_time: datetime, + ) -> list[BrokerageHoldingPayload]: + holdings = [] + asset_lookup = {asset.external_id: asset for asset in assets} + days_back = int(self.connection.get_setting("max_purchase_age_days", 3650)) + next_holding_id = 1 + refreshed_accounts = [] + + for account in accounts: + selected_assets = generator.sample( + assets, + k=min(holdings_per_account, len(assets)), + ) + balance = 0.0 + + for asset in selected_assets: + amount = round(generator.uniform(1, 50), 4) + current_value = amount * asset_lookup[asset.external_id].price + price_offset = generator.uniform(0.7, 1.15) + cost_basis = round(asset.price * price_offset, 2) + purchase_date = reference_time - timedelta( + days=generator.randint(1, max(1, days_back)) + ) + + holdings.append( + BrokerageHoldingPayload( + external_id=f"holding-{next_holding_id}", + account_external_id=account.external_id, + asset_external_id=asset.external_id, + ticker=asset.ticker, + amount=amount, + cost_basis=cost_basis, + purchase_date=purchase_date, + ) + ) + next_holding_id += 1 + balance += current_value + + refreshed_accounts.append( + BrokerageAccountPayload( + external_id=account.external_id, + name=account.name, + number=account.number, + balance=int(round(balance)), + ) + ) + + accounts[:] = refreshed_accounts + return holdings diff --git a/overview/brokerages/registry.py b/overview/brokerages/registry.py new file mode 100644 index 0000000..a709417 --- /dev/null +++ b/overview/brokerages/registry.py @@ -0,0 +1,34 @@ +from django.conf import settings +from django.utils.module_loading import import_string + +from overview.models import BrokerageConnection + +from .base import BaseBrokerageAdapter, BrokerageAdapterError +from .random import RandomBrokerageAdapter +from .rest import GenericRestBrokerageAdapter + + +_DEFAULT_ADAPTERS = { + GenericRestBrokerageAdapter.provider: GenericRestBrokerageAdapter, + RandomBrokerageAdapter.provider: RandomBrokerageAdapter, +} + + +def register_brokerage_adapter(provider: str, adapter_class: type[BaseBrokerageAdapter]): + _DEFAULT_ADAPTERS[provider] = adapter_class + + +def get_brokerage_adapter(connection: BrokerageConnection) -> BaseBrokerageAdapter: + adapter_class = _DEFAULT_ADAPTERS.get(connection.provider) + configured_paths = getattr(settings, "BROKERAGE_ADAPTERS", {}) + configured_path = configured_paths.get(connection.provider) + + if configured_path: + adapter_class = import_string(configured_path) + + if not adapter_class: + raise BrokerageAdapterError( + f"No brokerage adapter is registered for provider '{connection.provider}'." + ) + + return adapter_class(connection) diff --git a/overview/brokerages/rest.py b/overview/brokerages/rest.py new file mode 100644 index 0000000..fa58545 --- /dev/null +++ b/overview/brokerages/rest.py @@ -0,0 +1,131 @@ +import json +from datetime import datetime, timezone +from urllib.parse import urljoin +from urllib.request import Request, urlopen + +from overview.brokerage_types import ( + BrokerageAccountPayload, + BrokerageAssetPayload, + BrokerageHoldingPayload, + BrokerageSnapshot, +) + +from .base import BaseBrokerageAdapter, BrokerageAdapterError + + +class GenericRestBrokerageAdapter(BaseBrokerageAdapter): + provider = "generic-rest" + + def fetch_snapshot(self) -> BrokerageSnapshot: + accounts = self._fetch_collection( + self.connection.get_setting("account_endpoint", "/accounts") + ) + assets = self._fetch_collection( + self.connection.get_setting("asset_endpoint", "/assets") + ) + holdings = self._fetch_collection( + self.connection.get_setting("holding_endpoint", "/holdings") + ) + + return BrokerageSnapshot( + accounts=[self._build_account(item) for item in accounts], + assets=[self._build_asset(item) for item in assets], + holdings=[self._build_holding(item) for item in holdings], + ) + + def _fetch_collection(self, endpoint: str) -> list[dict]: + url = urljoin(self.connection.base_url, endpoint) + request = Request(url, headers=self._headers()) + + try: + with urlopen(request, timeout=30) as response: + payload = json.load(response) + except Exception as exc: + raise BrokerageAdapterError( + f"Unable to fetch brokerage data from {url}: {exc}" + ) from exc + + if isinstance(payload, dict): + for key in ("results", "data", "items"): + value = payload.get(key) + if isinstance(value, list): + return value + + if not isinstance(payload, list): + raise BrokerageAdapterError( + f"Expected a list payload from {url}, received {type(payload).__name__}" + ) + + return payload + + def _headers(self) -> dict[str, str]: + headers = {"Accept": "application/json"} + token = self.connection.get_credential("access_token") + api_key = self.connection.get_credential("api_key") + api_key_header = self.connection.get_setting("api_key_header", "X-API-Key") + + if token: + headers["Authorization"] = f"Bearer {token}" + if api_key: + headers[api_key_header] = api_key + + extra_headers = self.connection.get_setting("headers", {}) + if isinstance(extra_headers, dict): + headers.update({str(key): str(value) for key, value in extra_headers.items()}) + + return headers + + def _build_account(self, item: dict) -> BrokerageAccountPayload: + account_id_field = self.connection.get_setting("account_id_field", "id") + account_name_field = self.connection.get_setting("account_name_field", "name") + account_number_field = self.connection.get_setting("account_number_field", "number") + account_balance_field = self.connection.get_setting("account_balance_field", "balance") + + return BrokerageAccountPayload( + external_id=str(item[account_id_field]), + name=item[account_name_field], + number=item.get(account_number_field, ""), + balance=int(item.get(account_balance_field, 0)), + ) + + def _build_asset(self, item: dict) -> BrokerageAssetPayload: + asset_id_field = self.connection.get_setting("asset_id_field", "id") + asset_ticker_field = self.connection.get_setting("asset_ticker_field", "ticker") + asset_price_field = self.connection.get_setting("asset_price_field", "price") + + return BrokerageAssetPayload( + external_id=str(item[asset_id_field]), + ticker=item[asset_ticker_field], + price=float(item.get(asset_price_field, 0)), + ) + + def _build_holding(self, item: dict) -> BrokerageHoldingPayload: + holding_id_field = self.connection.get_setting("holding_id_field", "id") + holding_account_field = self.connection.get_setting("holding_account_field", "account_id") + holding_asset_field = self.connection.get_setting("holding_asset_field", "asset_id") + holding_ticker_field = self.connection.get_setting("holding_ticker_field", "ticker") + holding_amount_field = self.connection.get_setting("holding_amount_field", "amount") + holding_cost_basis_field = self.connection.get_setting( + "holding_cost_basis_field", "cost_basis" + ) + holding_purchase_date_field = self.connection.get_setting( + "holding_purchase_date_field", "purchase_date" + ) + + purchase_date = item.get(holding_purchase_date_field) + if purchase_date: + purchase_date = datetime.fromisoformat( + str(purchase_date).replace("Z", "+00:00") + ) + else: + purchase_date = datetime.now(timezone.utc) + + return BrokerageHoldingPayload( + external_id=str(item[holding_id_field]), + account_external_id=str(item[holding_account_field]), + asset_external_id=str(item[holding_asset_field]), + ticker=item[holding_ticker_field], + amount=float(item.get(holding_amount_field, 0)), + cost_basis=float(item.get(holding_cost_basis_field, 0)), + purchase_date=purchase_date, + ) diff --git a/overview/management/__init__.py b/overview/management/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/overview/management/__init__.py @@ -0,0 +1 @@ + diff --git a/overview/management/commands/__init__.py b/overview/management/commands/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/overview/management/commands/__init__.py @@ -0,0 +1 @@ + diff --git a/overview/management/commands/sync_brokerage.py b/overview/management/commands/sync_brokerage.py new file mode 100644 index 0000000..f588b78 --- /dev/null +++ b/overview/management/commands/sync_brokerage.py @@ -0,0 +1,70 @@ +from django.core.management.base import BaseCommand, CommandError + +from overview.models import BrokerageConnection +from overview.services import BrokerageIngestionService + + +class Command(BaseCommand): + help = "Sync investment accounts, assets, and holdings from a brokerage REST API." + + def add_arguments(self, parser): + parser.add_argument( + "--connection", + required=True, + help="Connection name or provider key to sync.", + ) + + def handle(self, *args, **options): + identifier = options["connection"] + connection = self._get_connection(identifier) + + run = BrokerageIngestionService(connection).ingest() + self.stdout.write( + self.style.SUCCESS( + "Sync finished: " + f"{run.accounts_processed} account(s), " + f"{run.assets_processed} asset(s), " + f"{run.holdings_processed} holding(s)." + ) + ) + + def _get_connection(self, identifier: str) -> BrokerageConnection: + named_connection = BrokerageConnection.objects.filter( + name=identifier, + is_active=True, + ).first() + if named_connection: + return named_connection + + provider_matches = list( + BrokerageConnection.objects.filter( + provider=identifier, + is_active=True, + ).order_by("name") + ) + if len(provider_matches) == 1: + return provider_matches[0] + if len(provider_matches) > 1: + matches = ", ".join(connection.name for connection in provider_matches) + raise CommandError( + f"Provider '{identifier}' matches multiple active connections: {matches}. " + "Re-run with --connection set to the connection name." + ) + + active_connections = list( + BrokerageConnection.objects.filter(is_active=True).order_by( + "provider", + "name", + ) + ) + available = ", ".join( + f"{connection.name} ({connection.provider})" + for connection in active_connections + ) + if not available: + available = "none" + + raise CommandError( + f"Active brokerage connection or provider '{identifier}' was not found. " + f"Available active connections: {available}." + ) diff --git a/overview/migrations/0004_brokerageconnection_ingestionrun_account_external_id_and_more.py b/overview/migrations/0004_brokerageconnection_ingestionrun_account_external_id_and_more.py new file mode 100644 index 0000000..44deaf2 --- /dev/null +++ b/overview/migrations/0004_brokerageconnection_ingestionrun_account_external_id_and_more.py @@ -0,0 +1,96 @@ +# Generated by Django 6.0.5 on 2026-05-24 00:51 + +import django.db.models.deletion +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('overview', '0003_accountholding_delete_assetholding_asset_price_and_more'), + ] + + operations = [ + migrations.CreateModel( + name='BrokerageConnection', + fields=[ + ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('provider', models.CharField(max_length=64)), + ('name', models.CharField(max_length=128)), + ('base_url', models.URLField()), + ('access_token', models.CharField(blank=True, max_length=512)), + ('api_key', models.CharField(blank=True, max_length=255)), + ('account_endpoint', models.CharField(default='/accounts', max_length=255)), + ('asset_endpoint', models.CharField(default='/assets', max_length=255)), + ('holding_endpoint', models.CharField(default='/holdings', max_length=255)), + ('is_active', models.BooleanField(default=True)), + ('created_at', models.DateTimeField(auto_now_add=True)), + ('updated_at', models.DateTimeField(auto_now=True)), + ], + ), + migrations.CreateModel( + name='IngestionRun', + fields=[ + ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('status', models.CharField(choices=[('pending', 'Pending'), ('running', 'Running'), ('succeeded', 'Succeeded'), ('failed', 'Failed')], default='pending', max_length=16)), + ('started_at', models.DateTimeField(blank=True, null=True)), + ('finished_at', models.DateTimeField(blank=True, null=True)), + ('accounts_processed', models.PositiveIntegerField(default=0)), + ('assets_processed', models.PositiveIntegerField(default=0)), + ('holdings_processed', models.PositiveIntegerField(default=0)), + ('error_message', models.TextField(blank=True)), + ('created_at', models.DateTimeField(auto_now_add=True)), + ], + ), + migrations.AddField( + model_name='account', + name='external_id', + field=models.CharField(blank=True, max_length=255), + ), + migrations.AddField( + model_name='account', + name='provider', + field=models.CharField(blank=True, max_length=64), + ), + migrations.AddField( + model_name='accountholding', + name='external_id', + field=models.CharField(blank=True, max_length=255), + ), + migrations.AddField( + model_name='accountholding', + name='provider', + field=models.CharField(blank=True, max_length=64), + ), + migrations.AddField( + model_name='asset', + name='external_id', + field=models.CharField(blank=True, max_length=255), + ), + migrations.AddField( + model_name='asset', + name='provider', + field=models.CharField(blank=True, max_length=64), + ), + migrations.AddConstraint( + model_name='account', + constraint=models.UniqueConstraint(fields=('provider', 'external_id'), name='unique_account_provider_external_id'), + ), + migrations.AddConstraint( + model_name='accountholding', + constraint=models.UniqueConstraint(fields=('provider', 'external_id'), name='unique_holding_provider_external_id'), + ), + migrations.AddConstraint( + model_name='asset', + constraint=models.UniqueConstraint(fields=('provider', 'external_id'), name='unique_asset_provider_external_id'), + ), + migrations.AddConstraint( + model_name='brokerageconnection', + constraint=models.UniqueConstraint(fields=('provider', 'name'), name='unique_brokerage_connection_provider_name'), + ), + migrations.AddField( + model_name='ingestionrun', + name='connection', + field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='overview.brokerageconnection'), + ), + ] diff --git a/overview/migrations/0005_remove_brokerageconnection_access_token_and_more.py b/overview/migrations/0005_remove_brokerageconnection_access_token_and_more.py new file mode 100644 index 0000000..be92f98 --- /dev/null +++ b/overview/migrations/0005_remove_brokerageconnection_access_token_and_more.py @@ -0,0 +1,43 @@ +# Generated by Django 6.0.5 on 2026-05-24 00:58 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('overview', '0004_brokerageconnection_ingestionrun_account_external_id_and_more'), + ] + + operations = [ + migrations.RemoveField( + model_name='brokerageconnection', + name='access_token', + ), + migrations.RemoveField( + model_name='brokerageconnection', + name='account_endpoint', + ), + migrations.RemoveField( + model_name='brokerageconnection', + name='api_key', + ), + migrations.RemoveField( + model_name='brokerageconnection', + name='asset_endpoint', + ), + migrations.RemoveField( + model_name='brokerageconnection', + name='holding_endpoint', + ), + migrations.AddField( + model_name='brokerageconnection', + name='credentials', + field=models.JSONField(blank=True, default=dict), + ), + migrations.AddField( + model_name='brokerageconnection', + name='settings', + field=models.JSONField(blank=True, default=dict), + ), + ] diff --git a/overview/models.py b/overview/models.py index b3d39a5..b090461 100644 --- a/overview/models.py +++ b/overview/models.py @@ -1,14 +1,40 @@ from django.db import models -# Create your models here. + class Account(models.Model): name = models.TextField(max_length=1024) number = models.TextField(max_length=1024) balance = models.IntegerField() + provider = models.CharField(max_length=64, blank=True) + external_id = models.CharField(max_length=255, blank=True) + + class Meta: + constraints = [ + models.UniqueConstraint( + fields=["provider", "external_id"], + name="unique_account_provider_external_id", + ), + ] + + def __str__(self): + return self.name class Asset(models.Model): ticker = models.TextField(max_length=20) price = models.FloatField() + provider = models.CharField(max_length=64, blank=True) + external_id = models.CharField(max_length=255, blank=True) + + class Meta: + constraints = [ + models.UniqueConstraint( + fields=["provider", "external_id"], + name="unique_asset_provider_external_id", + ), + ] + + def __str__(self): + return self.ticker class AccountHolding(models.Model): account = models.ForeignKey(Account, on_delete=models.CASCADE) @@ -17,6 +43,16 @@ class AccountHolding(models.Model): ammount = models.FloatField() cost_basis = models.FloatField() purchase_date = models.DateTimeField() + provider = models.CharField(max_length=64, blank=True) + external_id = models.CharField(max_length=255, blank=True) + + class Meta: + constraints = [ + models.UniqueConstraint( + fields=["provider", "external_id"], + name="unique_holding_provider_external_id", + ), + ] def holding_value(self): return self.ammount * self.asset.price @@ -26,4 +62,65 @@ class AccountHolding(models.Model): def holding_gain(self): return self.holding_value() - self.holding_principle() - + + def __str__(self): + return f"{self.account.name}: {self.ticker}" + + +class BrokerageConnection(models.Model): + provider = models.CharField(max_length=64) + name = models.CharField(max_length=128) + base_url = models.URLField() + credentials = models.JSONField(default=dict, blank=True) + settings = models.JSONField(default=dict, blank=True) + is_active = models.BooleanField(default=True) + created_at = models.DateTimeField(auto_now_add=True) + updated_at = models.DateTimeField(auto_now=True) + + class Meta: + constraints = [ + models.UniqueConstraint( + fields=["provider", "name"], + name="unique_brokerage_connection_provider_name", + ), + ] + + def __str__(self): + return f"{self.provider}: {self.name}" + + def get_credential(self, key: str, default=None): + return self.credentials.get(key, default) + + def get_setting(self, key: str, default=None): + return self.settings.get(key, default) + + +class IngestionRun(models.Model): + STATUS_PENDING = "pending" + STATUS_RUNNING = "running" + STATUS_SUCCEEDED = "succeeded" + STATUS_FAILED = "failed" + + STATUS_CHOICES = [ + (STATUS_PENDING, "Pending"), + (STATUS_RUNNING, "Running"), + (STATUS_SUCCEEDED, "Succeeded"), + (STATUS_FAILED, "Failed"), + ] + + connection = models.ForeignKey(BrokerageConnection, on_delete=models.CASCADE) + status = models.CharField( + max_length=16, + choices=STATUS_CHOICES, + default=STATUS_PENDING, + ) + started_at = models.DateTimeField(null=True, blank=True) + finished_at = models.DateTimeField(null=True, blank=True) + accounts_processed = models.PositiveIntegerField(default=0) + assets_processed = models.PositiveIntegerField(default=0) + holdings_processed = models.PositiveIntegerField(default=0) + error_message = models.TextField(blank=True) + created_at = models.DateTimeField(auto_now_add=True) + + def __str__(self): + return f"{self.connection} [{self.status}]" diff --git a/overview/services.py b/overview/services.py new file mode 100644 index 0000000..9d9970b --- /dev/null +++ b/overview/services.py @@ -0,0 +1,111 @@ +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), + ) diff --git a/overview/tests.py b/overview/tests.py index 7ce503c..2d3c8a3 100644 --- a/overview/tests.py +++ b/overview/tests.py @@ -1,3 +1,333 @@ +from datetime import datetime, timezone +from io import StringIO + +from django.core.management import call_command +from django.core.management.base import CommandError from django.test import TestCase -# Create your tests here. +from .brokerage_types import ( + BrokerageAccountPayload, + BrokerageAssetPayload, + BrokerageHoldingPayload, + BrokerageSnapshot, +) +from .brokerages.base import BaseBrokerageAdapter, BrokerageAdapterError +from .brokerages.random import RandomBrokerageAdapter +from .brokerages.registry import register_brokerage_adapter +from .models import Account, AccountHolding, Asset, BrokerageConnection, IngestionRun +from .services import BrokerageIngestionService + + +class FakeBrokerageAdapter(BaseBrokerageAdapter): + provider = "fake-broker" + + def __init__(self, connection, snapshot=None, error=None): + super().__init__(connection) + self.snapshot = snapshot + self.error = error + + def fetch_snapshot(self): + if self.error: + raise self.error + return self.snapshot + + +class AlternateFakeBrokerageAdapter(FakeBrokerageAdapter): + provider = "alternate-fake-broker" + + +class ExplodingBrokerageAdapter(BaseBrokerageAdapter): + provider = "error-broker" + + def fetch_snapshot(self): + raise RuntimeError("brokerage unavailable") + + +class BrokerageIngestionServiceTests(TestCase): + def setUp(self): + register_brokerage_adapter(FakeBrokerageAdapter.provider, FakeBrokerageAdapter) + register_brokerage_adapter( + AlternateFakeBrokerageAdapter.provider, + AlternateFakeBrokerageAdapter, + ) + register_brokerage_adapter( + ExplodingBrokerageAdapter.provider, + ExplodingBrokerageAdapter, + ) + + self.connection = BrokerageConnection.objects.create( + provider=FakeBrokerageAdapter.provider, + name="Primary", + base_url="https://broker.example/api/", + credentials={"access_token": "token-1"}, + settings={"account_endpoint": "/accounts"}, + ) + + def test_ingest_creates_and_updates_portfolio_records(self): + initial_snapshot = BrokerageSnapshot( + accounts=[ + BrokerageAccountPayload( + external_id="acct-1", + name="Brokerage Account", + number="1234", + balance=10000, + ) + ], + assets=[ + BrokerageAssetPayload( + external_id="asset-1", + ticker="AAPL", + price=195.12, + ) + ], + holdings=[ + BrokerageHoldingPayload( + external_id="holding-1", + account_external_id="acct-1", + asset_external_id="asset-1", + ticker="AAPL", + amount=10, + cost_basis=150.25, + purchase_date=datetime(2025, 1, 15, tzinfo=timezone.utc), + ) + ], + ) + + updated_snapshot = BrokerageSnapshot( + accounts=[ + BrokerageAccountPayload( + external_id="acct-1", + name="Brokerage Account", + number="1234", + balance=12500, + ) + ], + assets=[ + BrokerageAssetPayload( + external_id="asset-1", + ticker="AAPL", + price=205.55, + ) + ], + holdings=[ + BrokerageHoldingPayload( + external_id="holding-1", + account_external_id="acct-1", + asset_external_id="asset-1", + ticker="AAPL", + amount=12, + cost_basis=150.25, + purchase_date=datetime(2025, 1, 15, tzinfo=timezone.utc), + ) + ], + ) + + first_run = BrokerageIngestionService( + self.connection, + adapter=FakeBrokerageAdapter(self.connection, snapshot=initial_snapshot), + ).ingest() + second_run = BrokerageIngestionService( + self.connection, + adapter=FakeBrokerageAdapter(self.connection, snapshot=updated_snapshot), + ).ingest() + + self.assertEqual(first_run.status, IngestionRun.STATUS_SUCCEEDED) + self.assertEqual(second_run.status, IngestionRun.STATUS_SUCCEEDED) + self.assertEqual(Account.objects.count(), 1) + self.assertEqual(Asset.objects.count(), 1) + self.assertEqual(AccountHolding.objects.count(), 1) + + account = Account.objects.get() + asset = Asset.objects.get() + holding = AccountHolding.objects.get() + + self.assertEqual(account.balance, 12500) + self.assertEqual(asset.price, 205.55) + self.assertEqual(holding.ammount, 12) + self.assertEqual(second_run.accounts_processed, 1) + self.assertEqual(second_run.assets_processed, 1) + self.assertEqual(second_run.holdings_processed, 1) + + def test_ingest_marks_run_failed_when_client_errors(self): + failing_connection = BrokerageConnection.objects.create( + provider=ExplodingBrokerageAdapter.provider, + name="Failing", + base_url="https://broker.example/api/", + ) + + with self.assertRaises(RuntimeError): + BrokerageIngestionService(failing_connection).ingest() + + run = IngestionRun.objects.get() + self.assertEqual(run.status, IngestionRun.STATUS_FAILED) + self.assertIn("brokerage unavailable", run.error_message) + + def test_multiple_brokerages_can_ingest_same_external_ids(self): + alternate_connection = BrokerageConnection.objects.create( + provider=AlternateFakeBrokerageAdapter.provider, + name="Secondary", + base_url="https://broker-two.example/api/", + ) + + snapshot = BrokerageSnapshot( + accounts=[ + BrokerageAccountPayload( + external_id="acct-1", + name="Shared External Account", + number="9999", + balance=5000, + ) + ], + assets=[ + BrokerageAssetPayload( + external_id="asset-1", + ticker="MSFT", + price=410.01, + ) + ], + holdings=[ + BrokerageHoldingPayload( + external_id="holding-1", + account_external_id="acct-1", + asset_external_id="asset-1", + ticker="MSFT", + amount=3, + cost_basis=300.00, + purchase_date=datetime(2025, 2, 1, tzinfo=timezone.utc), + ) + ], + ) + + BrokerageIngestionService( + self.connection, + adapter=FakeBrokerageAdapter(self.connection, snapshot=snapshot), + ).ingest() + BrokerageIngestionService( + alternate_connection, + adapter=AlternateFakeBrokerageAdapter( + alternate_connection, + snapshot=snapshot, + ), + ).ingest() + + self.assertEqual(Account.objects.count(), 2) + self.assertEqual(Asset.objects.count(), 2) + self.assertEqual(AccountHolding.objects.count(), 2) + self.assertEqual( + set(Account.objects.values_list("provider", flat=True)), + {FakeBrokerageAdapter.provider, AlternateFakeBrokerageAdapter.provider}, + ) + + def test_missing_adapter_raises_clear_error(self): + unknown_connection = BrokerageConnection.objects.create( + provider="unknown-broker", + name="Unknown", + base_url="https://unknown.example/api/", + ) + + with self.assertRaises(BrokerageAdapterError): + BrokerageIngestionService(unknown_connection).ingest() + + def test_random_adapter_generates_seeded_snapshot(self): + random_connection = BrokerageConnection.objects.create( + provider=RandomBrokerageAdapter.provider, + name="Randomized", + base_url="https://local.random/", + settings={ + "seed": 17, + "account_count": 2, + "asset_count": 5, + "holdings_per_account": 3, + "tickers": ["AAPL", "MSFT", "TSLA", "NVDA", "AMD"], + }, + ) + + first_snapshot = RandomBrokerageAdapter(random_connection).fetch_snapshot() + second_snapshot = RandomBrokerageAdapter(random_connection).fetch_snapshot() + + self.assertEqual(first_snapshot.accounts, second_snapshot.accounts) + self.assertEqual(first_snapshot.assets, second_snapshot.assets) + self.assertEqual(first_snapshot.holdings, second_snapshot.holdings) + self.assertEqual(len(first_snapshot.accounts), 2) + self.assertEqual(len(first_snapshot.assets), 5) + self.assertEqual(len(first_snapshot.holdings), 6) + + def test_random_adapter_ingests_generated_data(self): + random_connection = BrokerageConnection.objects.create( + provider=RandomBrokerageAdapter.provider, + name="Random Import", + base_url="https://local.random/", + settings={ + "seed": 99, + "account_count": 3, + "asset_count": 6, + "holdings_per_account": 2, + }, + ) + + run = BrokerageIngestionService(random_connection).ingest() + + self.assertEqual(run.status, IngestionRun.STATUS_SUCCEEDED) + self.assertEqual(run.accounts_processed, 3) + self.assertEqual(run.assets_processed, 6) + self.assertEqual(run.holdings_processed, 6) + self.assertEqual( + Account.objects.filter(provider=RandomBrokerageAdapter.provider).count(), + 3, + ) + self.assertEqual( + Asset.objects.filter(provider=RandomBrokerageAdapter.provider).count(), + 6, + ) + self.assertEqual( + AccountHolding.objects.filter( + provider=RandomBrokerageAdapter.provider + ).count(), + 6, + ) + + def test_sync_command_accepts_provider_identifier(self): + random_connection = BrokerageConnection.objects.create( + provider=RandomBrokerageAdapter.provider, + name="Demo Random", + base_url="https://local.random/", + settings={ + "seed": 123, + "account_count": 1, + "asset_count": 2, + "holdings_per_account": 1, + }, + ) + + output = StringIO() + call_command( + "sync_brokerage", + connection=random_connection.provider, + stdout=output, + ) + + self.assertIn("Sync finished", output.getvalue()) + self.assertEqual( + Account.objects.filter(provider=RandomBrokerageAdapter.provider).count(), + 1, + ) + + def test_sync_command_requires_name_when_provider_has_multiple_connections(self): + BrokerageConnection.objects.create( + provider=RandomBrokerageAdapter.provider, + name="Random One", + base_url="https://local.random/", + ) + BrokerageConnection.objects.create( + provider=RandomBrokerageAdapter.provider, + name="Random Two", + base_url="https://local.random/", + ) + + with self.assertRaises(CommandError) as exc: + call_command( + "sync_brokerage", + connection=RandomBrokerageAdapter.provider, + ) + + self.assertIn("matches multiple active connections", str(exc.exception))