AI: Brokerage Ingestion infra
This commit is contained in:
Tyler Wilcox
parent
2958a4b459
commit
618072069e
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
@ -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,
|
||||
)
|
||||
|
|
@ -0,0 +1 @@
|
|||
|
||||
|
|
@ -0,0 +1 @@
|
|||
|
||||
|
|
@ -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}."
|
||||
)
|
||||
|
|
@ -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'),
|
||||
),
|
||||
]
|
||||
|
|
@ -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),
|
||||
),
|
||||
]
|
||||
|
|
@ -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
|
||||
|
|
@ -27,3 +63,64 @@ 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}]"
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
@ -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))
|
||||
|
|
|
|||
Loading…
Reference in New Issue