gear-orders/db/models.py

import pytz
import datetime
from peewee import *
from settings import SQLITE_DB, MASTODON_INSTANCE
from util import sqlite_time

database = SqliteDatabase(SQLITE_DB)

class BaseModel(Model):
    class Meta:
        database = database

class MastodonServer(BaseModel):
    name = TextField(unique=True, null=False)

    client_id = TextField(null=False)
    client_secret = TextField(null=False)

    class Meta:
        table_name = 'mastodon_server'

class User(BaseModel):
    telegram_username = TextField(unique=True)
    telegram_chat_id = IntegerField()
    telegram_photo_url = TextField(null=True)

    mastodon_server = ForeignKeyField(
        MastodonServer,
        column_name='mastodon_server_id',
        field='id',
        null=True,
    )
    mastodon_username = TextField(null=True)
    mastodon_access_token = TextField(null=True)

    mastodon_attn_list = TextField(null=True)
    mastodon_post_public = BooleanField(null=True, default=False)

    verify_mastodon_alt_text = BooleanField(null=False, default=False)
    verify_mastodon_favorite = BooleanField(null=False, default=False)
    verify_delay = IntegerField(null=True)

    permission_orders_pools_view = BooleanField(null=False, default=True)
    permission_orders_pools_details = BooleanField(null=False, default=True)
    permission_orders_pools_edit = BooleanField(null=False, default=True)

    def mastodon_account(self):
        if self.mastodon_server is None or self.mastodon_username is None:
            return
        elif self.mastodon_server.name == MASTODON_INSTANCE:
            return self.mastodon_username
        else:
            return f"{self.mastodon_username}@{self.mastodon_server}"
        
    def mastodon_visibility(self):
        return 'public' if self.mastodon_post_public else 'direct'

    def __str__(self):
        return self.telegram_username

    class Meta:
        table_name = 'user'

class DomSubUsers(BaseModel):
    dom = ForeignKeyField(column_name='dom_id', field='id', model=User)
    sub = ForeignKeyField(column_name='sub_id', field='id', model=User)

    class Meta:
        table_name = 'dom_sub'
        indexes = (
            (('dom_id', 'sub_id'), True),
        )

#
# Order Configuration
#
class OrdersPool(BaseModel):
    name = TextField()

    user = ForeignKeyField(column_name='user_id', field='id', model=User, backref='orders_pools')

    scheduled = BooleanField(default=False)
    probability = FloatField(null=True)
    weekdays = BooleanField(null=True)
    weekends = BooleanField(null=True)
    time = TextField(null=True)
    confirm_delay = IntegerField(null=True)

    punishment_pool = ForeignKeyField(
        'self',
        column_name='punishment_pool_id',
        field='id',
        backref='punishments_for',
        null=True
    )

    updated_at = DateTimeField(null=False)

    def to_dict(self):
        return {
            'id': self.id,
            'name': self.name,
            'scheduled': self.scheduled,
            'probability': self.probability,
            'weekdays': self.weekdays,
            'weekends': self.weekends,
            'time': self.time,
            'confirm_delay': self.confirm_delay,
            'punishment_pool_id': self.punishment_pool_id,
            'orders': [{
                'id': order.id,
                'name': order.name,
                'weight': order.weight,
                'repeat': order.repeat,
                'add_ons': [
                {
                    'id': add_on.id,
                    'name': add_on.name,
                    'probability': add_on.probability 
                } for add_on in order.add_ons
                ]
            } for order in self.orders]
        }

    def save(self, *args, **kwargs):
        self.updated_at = sqlite_time(datetime.datetime.now(datetime.UTC))

        return super(OrdersPool, self).save(*args, **kwargs)

    def __str__(self):
        return f"{self.name}[{self.user.telegram_username}]"

    class Meta:
        table_name = 'orders_pool'

class Order(BaseModel):
    name = TextField()
    weight = IntegerField()
    repeat = FloatField()

    pool = ForeignKeyField(column_name='orders_pool_id', field='id', model=OrdersPool, backref='orders')

    class Meta:
        table_name = 'order'

class OrderAddOn(BaseModel):
    name = TextField()
    probability = FloatField()

    order = ForeignKeyField(column_name='order_id', field='id', model=Order, backref='add_ons')

    class Meta:
        table_name = 'order_add_on'

#
# Order State
#
class OrderStatus(BaseModel):
    confirmed_at = DateTimeField(null=True)
    created_at = DateTimeField()
    due_at = DateTimeField(null=True)
    verify_at = DateTimeField(null=True)
    mastodon_id = TextField()
    text = TextField()

    pool = ForeignKeyField(
        column_name='orders_pool_id',
        field='id',
        model=OrdersPool,
        backref='statuses',
        null=True
        )
    
    user = ForeignKeyField(
        column_name='user_id',
        field='id',
        model=User,
        backref='statuses',
        null=False
        )
    
    punishment_for = ForeignKeyField(
        'self',
        column_name='punishment_for_id',
        field='id',
        backref='punishment',
        null=True
    )

    def __str__(self):
        return f"{self.pool} {self.id}"
    
    class Meta:
        table_name = 'order_status'

class Repeat(BaseModel):
    count = IntegerField(default=0)
    orders = TextField()
    probability = FloatField()

    orders_pool = ForeignKeyField(
        column_name='orders_pool_id',
        field='id',
        model=OrdersPool,
        null=False,
        unique=True,
        backref='repeat'
        )

    class Meta:
        table_name = 'repeat'

class SkipDay(BaseModel):
    date = DateField(unique=True)

    user = ForeignKeyField(
        User,
        column_name='user_id',
        field='id',
        null=False,
        backref='skip_days'
    )

    pool = ForeignKeyField(
        column_name='orders_pool_id',
        field='id',
        model=OrdersPool,
        backref='skip_days',
        null=True
    )

    class Meta:
        table_name = 'skip_day'

class TimelineEvent(BaseModel):
    updated_at = DateTimeField(null=False)
    type = TextField(null=False)
    text = TextField(null=False)
    extra = TextField(null=True)

    user = ForeignKeyField(
        User,
        column_name='user_id',
        field='id',
        null=False,
        backref="timeline_events"
    )

    actor = ForeignKeyField(
        User,
        column_name='actor_user_id',
        field='id',
        null=True,
        backref="timeline_events_actor"
    )

    orders_pool = ForeignKeyField(
        OrdersPool,
        column_name='orders_pool_id',
        field='id',
        null=True,
        backref="timeline_events"
    )

    order_status = ForeignKeyField(
        OrderStatus,
        column_name='order_status_id',
        field='id',
        null=True,
        backref="timeline_events"
    )

    class Meta:
        table_name = 'timeline_event'