gear-orders/db.py

import os
import sqlite3
import datetime

SQLITE_DB = os.environ.get('SQLITE_DB', 'db.sqlite3')

TABLE_REPEATS = '''
    CREATE TABLE IF NOT EXISTS repeats (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        probability FLOAT NOT NULL,
        orders TEXT NOT NULL,
        count INTEGER DEFAULT 0
    );
'''

TABLE_ORDER_STATUS = '''
    CREATE TABLE IF NOT EXISTS order_status (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        mastodon_id TEXT NOT NULL,
        created_at TIMESTAMP NOT NULL,
        text TEXT NOT NULL,
        confirmed_at TIMESTAMP,
        punishment_id INTEGER,
        FOREIGN KEY(punishment_id) REFERENCES punishment_status(id)
    );
'''

TABLE_PUNISHMENT_STATUS = '''
    CREATE TABLE IF NOT EXISTS punishment_status (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        mastodon_id TEXT NOT NULL,
        created_at TIMESTAMP NOT NULL,
        text TEXT NOT NULL,
        confirmed_at TIMESTAMP
    );
'''

class Database:
    def __init__(self):
        self.conn = sqlite3.connect(SQLITE_DB)
        self.conn.row_factory = sqlite3.Row
        self.table_init(TABLE_REPEATS)
        self.table_init(TABLE_PUNISHMENT_STATUS)
        self.table_init(TABLE_ORDER_STATUS)
        
    def table_init(self, table_sql):
        c = self.conn.cursor()
        c.execute(table_sql)
        self.conn.commit()

    def update(self, sql, args=[]):
        c = self.conn.cursor()
        c.execute(sql, args)
        self.conn.commit()
        return c.lastrowid

    def repeat_get(self):
        c = self.conn.cursor()
        sql = 'SELECT id, probability, orders, count FROM repeats LIMIT 1'
        c.execute(sql)
        return c.fetchone()

    def repeat_increment(self):
        self.update('UPDATE repeats SET count = count + 1')

    def repeat_put(self, probability, orders):
        self.update(
            'INSERT INTO repeats (probability, orders) VALUES (?, ?);',
            [probability, orders]
        )

    def repeat_clear(self):
        self.update('DELETE FROM repeats')

    def order_status_put(self, mastodon_id, created_at, text):
        self.update(
            '''
                INSERT INTO order_status
                (mastodon_id, created_at, text)
                VALUES (?, ?, ?);
            ''',
            [
                mastodon_id,
                created_at,
                text
            ]
        )

    def order_status_outstanding(self):
        c = self.conn.cursor()
        sql = '''
            SELECT id, mastodon_id, created_at, confirmed_at
            FROM order_status
            WHERE confirmed_at IS NULL AND punishment_id IS NULL
        '''
        c.execute(sql)
        return c.fetchall()

    def order_status_confirm(self, id, confirmed_at):
        self.update(
            '''
                UPDATE order_status
                SET confirmed_at=?
                WHERE id=?;
            ''',
            [
                confirmed_at,
                id
            ]
        )
    
    def punishment_status_put(self, order_status_id, mastodon_id, created_at, text):
        punishment_status_id = self.update(
            '''
                INSERT INTO punishment_status
                (mastodon_id, created_at, text)
                VALUES (?, ?, ?);
            ''',
            [
                mastodon_id,
                created_at,
                text
            ]
        )
        self.update(
            '''
                UPDATE order_status
                SET punishment_id=?
                WHERE id=?
            ''',
            [punishment_status_id, order_status_id]
        )