Stonks/overview/management/commands/sync_brokerage.py

72 lines
2.4 KiB
Python

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), "
f"{run.transactions_processed} transaction(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}."
)