Stonks/overview/models.py

from django.db import models


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)
    asset = models.ForeignKey(Asset, on_delete=models.CASCADE)
    ticker = models.TextField(max_length=20)
    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
    
    def holding_principle(self):
        return self.ammount * self.cost_basis

    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}]"


class AccountSnapshot(models.Model):
    ingestion_run = models.ForeignKey(IngestionRun, on_delete=models.CASCADE)
    account = models.ForeignKey(Account, on_delete=models.CASCADE)
    balance = models.IntegerField()
    snapshot_at = models.DateTimeField()

    class Meta:
        constraints = [
            models.UniqueConstraint(
                fields=["ingestion_run", "account"],
                name="unique_account_snapshot_per_run",
            ),
        ]


class AssetPriceSnapshot(models.Model):
    ingestion_run = models.ForeignKey(IngestionRun, on_delete=models.CASCADE)
    asset = models.ForeignKey(Asset, on_delete=models.CASCADE)
    price = models.FloatField()
    snapshot_at = models.DateTimeField()

    class Meta:
        constraints = [
            models.UniqueConstraint(
                fields=["ingestion_run", "asset"],
                name="unique_asset_snapshot_per_run",
            ),
        ]


class HoldingSnapshot(models.Model):
    ingestion_run = models.ForeignKey(IngestionRun, on_delete=models.CASCADE)
    holding = models.ForeignKey(AccountHolding, on_delete=models.CASCADE)
    account = models.ForeignKey(Account, on_delete=models.CASCADE)
    asset = models.ForeignKey(Asset, on_delete=models.CASCADE)
    amount = models.FloatField()
    cost_basis = models.FloatField()
    market_price = models.FloatField()
    value = models.FloatField()
    snapshot_at = models.DateTimeField()

    class Meta:
        constraints = [
            models.UniqueConstraint(
                fields=["ingestion_run", "holding"],
                name="unique_holding_snapshot_per_run",
            ),
        ]