fisk/core/models.py

import stripe
from alpaca.common.exceptions import APIError
from alpaca.trading.client import TradingClient
from alpaca.trading.requests import GetAssetsRequest
from django.conf import settings
from django.contrib.auth.models import AbstractUser
from django.db import models

from core.lib import trades
from core.lib.customers import get_or_create, update_customer_fields
from core.util import logs

log = logs.get_logger(__name__)


class Plan(models.Model):
    name = models.CharField(max_length=255, unique=True)
    description = models.CharField(max_length=1024, null=True, blank=True)
    cost = models.IntegerField()
    product_id = models.CharField(max_length=255, unique=True, null=True, blank=True)
    image = models.CharField(max_length=1024, null=True, blank=True)

    def __str__(self):
        return f"{self.name} (£{self.cost})"


class User(AbstractUser):
    # Stripe customer ID
    stripe_id = models.CharField(max_length=255, null=True, blank=True)
    last_payment = models.DateTimeField(null=True, blank=True)
    plans = models.ManyToManyField(Plan, blank=True)
    email = models.EmailField(unique=True)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._original = self

    def save(self, *args, **kwargs):
        """
        Override the save function to create a Stripe customer.
        """
        if settings.STRIPE_ENABLED:
            if not self.stripe_id:  # stripe ID not stored
                self.stripe_id = get_or_create(
                    self.email, self.first_name, self.last_name
                )

            to_update = {}
            if self.email != self._original.email:
                to_update["email"] = self.email
            if self.first_name != self._original.first_name:
                to_update["first_name"] = self.first_name
            if self.last_name != self._original.last_name:
                to_update["last_name"] = self.last_name

            update_customer_fields(self.stripe_id, **to_update)

        super().save(*args, **kwargs)

    def delete(self, *args, **kwargs):
        if settings.STRIPE_ENABLED:
            if self.stripe_id:
                stripe.Customer.delete(self.stripe_id)
                log.info(f"Deleted Stripe customer {self.stripe_id}")
        super().delete(*args, **kwargs)

    def has_plan(self, plan):
        plan_list = [plan.name for plan in self.plans.all()]
        return plan in plan_list


class Account(models.Model):
    EXCHANGE_CHOICES = (("alpaca", "Alpaca"),)
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=255)
    exchange = models.CharField(choices=EXCHANGE_CHOICES, max_length=255)
    api_key = models.CharField(max_length=255)
    api_secret = models.CharField(max_length=255)
    sandbox = models.BooleanField(default=False)
    supported_symbols = models.JSONField(default=list)

    def __str__(self):
        name = f"{self.name} ({self.exchange})"
        if self.sandbox:
            name += " (sandbox)"
        return name

    def save(self, *args, **kwargs):
        """
        Override the save function to update supported symbols.
        """
        try:
            request = GetAssetsRequest(status="active", asset_class="crypto")
            assets = self.client.get_all_assets(filter=request)
            asset_list = [x["symbol"] for x in assets if "symbol" in x]
            self.supported_symbols = asset_list
            print("Supported symbols", self.supported_symbols)
        except APIError as e:
            log.error(f"Could not get asset list: {e}")
            # return False

        super().save(*args, **kwargs)

    def get_client(self):
        trading_client = TradingClient(
            self.api_key, self.api_secret, paper=self.sandbox, raw_data=True
        )
        return trading_client

    @property
    def client(self):
        """
        Convenience property for one-off API calls.
        """
        return self.get_client()

    @classmethod
    def get_by_id(cls, account_id, user):
        return cls.objects.get(id=account_id, user=user)


class Session(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    request = models.CharField(max_length=255, null=True, blank=True)
    session = models.CharField(max_length=255, null=True, blank=True)
    subscription_id = models.CharField(max_length=255, null=True, blank=True)
    plan = models.ForeignKey(Plan, null=True, blank=True, on_delete=models.CASCADE)


class Hook(models.Model):
    DIRECTION_CHOICES = (
        ("buy", "Buy"),
        ("sell", "Sell"),
    )
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=1024, null=True, blank=True, unique=True)
    hook = models.CharField(max_length=255, unique=True)
    direction = models.CharField(choices=DIRECTION_CHOICES, max_length=255)
    received = models.IntegerField(default=0)

    def __str__(self):
        return f"{self.name} ({self.hook})"


class Trade(models.Model):
    TYPE_CHOICES = (
        ("market", "Market"),
        ("limit", "Limit"),
    )
    DIRECTION_CHOICES = (
        ("buy", "Buy"),
        ("sell", "Sell"),
    )
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    account = models.ForeignKey(Account, on_delete=models.CASCADE)
    hook = models.ForeignKey(Hook, on_delete=models.CASCADE, null=True, blank=True)
    symbol = models.CharField(max_length=255)
    type = models.CharField(choices=TYPE_CHOICES, max_length=255)
    amount = models.FloatField(null=True, blank=True)
    amount_usd = models.FloatField(null=True, blank=True)
    price = models.FloatField(null=True, blank=True)
    stop_loss = models.FloatField(null=True, blank=True)
    take_profit = models.FloatField(null=True, blank=True)
    status = models.CharField(max_length=255, null=True, blank=True)
    direction = models.CharField(choices=DIRECTION_CHOICES, max_length=255)

    # To populate from the trade
    order_id = models.CharField(max_length=255, null=True, blank=True)
    client_order_id = models.CharField(max_length=255, null=True, blank=True)
    response = models.JSONField(null=True, blank=True)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._original = self

    def post(self):
        return trades.post_trade(self)

    def delete(self, *args, **kwargs):
        # close the trade
        super().delete(*args, **kwargs)


class Callback(models.Model):
    hook = models.ForeignKey(Hook, on_delete=models.CASCADE)
    title = models.CharField(max_length=1024, null=True, blank=True)
    message = models.CharField(max_length=1024, null=True, blank=True)
    period = models.CharField(max_length=255, null=True, blank=True)
    sent = models.BigIntegerField(null=True, blank=True)
    trade = models.BigIntegerField(null=True, blank=True)
    exchange = models.CharField(max_length=255, null=True, blank=True)
    base = models.CharField(max_length=255, null=True, blank=True)
    quote = models.CharField(max_length=255, null=True, blank=True)
    contract = models.CharField(max_length=255, null=True, blank=True)
    price = models.FloatField(null=True, blank=True)
    symbol = models.CharField(max_length=255)


class Strategy(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=255)
    description = models.TextField(null=True, blank=True)
    account = models.ForeignKey(Account, on_delete=models.CASCADE)
    hooks = models.ManyToManyField(Hook)
    enabled = models.BooleanField(default=False)
    take_profit_percent = models.FloatField(default=3.0)
    stop_loss_percent = models.FloatField(default=1.0)
    price_slippage_percent = models.FloatField(default=2.5)
    trade_size_percent = models.FloatField(default=2.5)

    def __str__(self):
        return self.name


# class Perms(models.Model):
#     class Meta:
#         permissions = (
#             ("bypass_hashing", "Can bypass field hashing"),  #
#             ("bypass_blacklist", "Can bypass the blacklist"),  #
#             ("bypass_encryption", "Can bypass field encryption"),  #
#             ("bypass_obfuscation", "Can bypass field obfuscation"),  #
#             ("bypass_delay", "Can bypass data delay"),  #
#             ("bypass_randomisation", "Can bypass data randomisation"),  #
#             ("post_irc", "Can post to IRC"),
#             ("post_discord", "Can post to Discord"),
#             ("query_search", "Can search with query strings"),  #
#             ("use_insights", "Can use the Insights page"),
#             ("index_int", "Can use the internal index"),
#             ("index_meta", "Can use the meta index"),
#             ("restricted_sources", "Can access restricted sources"),
#         )