OwnchatBot/ownchatbot/db.py

import sqlite3
from flask import current_app, g
from ownchatbot.reward_handlers import goal_reached, get_queue
from flask.cli import with_appcontext
import click
from time import time
import os
import logging


def rem_from_queue(reward_name):  # Remove a reward from the queue
    db = get_db()

    try:
        cursor = db.execute(
            "DELETE FROM reward_queue WHERE reward = ?", (reward_name,)
        )
        db.commit()
    except sqlite3.Error as rfqerror:
        current_app.logger.error(f'Couldn\'t remove {reward_name} from reward queue: {rfqerror.args[0]}')
        return False
    return True


def init_app(app):
    app.teardown_appcontext(close_db)
    app.cli.add_command(init_db_command)
    app.cli.add_command(update_db_command)
    

def get_db():
    if 'db' not in g:
        g.db = sqlite3.connect(
            current_app.config['DATABASE'],
            detect_types=sqlite3.PARSE_DECLTYPES
        )
        g.db.row_factory = sqlite3.Row

    return g.db


def close_db(e=None):
    db = g.pop('db', None)

    if db is not None:
        db.close()


def init_db():  # Setup the database. If it already exists, apply new schema
    db = get_db()

    with current_app.open_resource('schema.sql') as f:  # Set up database
        db.executescript(f.read().decode('utf8'))

    reread_votes()
    reread_goals()
    if get_queue(db):  # Make sure database is initialized
        return True


def update_db():
    db = get_db()

    update_db_path = os.path.join(current_app.root_path, 'update_db.sql')  # Check if the update_db.sql file exists
    if os.path.exists(update_db_path):  # Update it if it does
        with current_app.open_resource('update_db.sql') as f:
            db.executescript(f.read().decode('utf8'))

    if get_queue(db):  # Make sure database is initialized
        return True


@click.command('init-db')
@with_appcontext
def init_db_command():  # If there is already a database, clear it and do setup
    if init_db():
        click.echo('Database initialized.')


@click.command('update-db')
@with_appcontext
def update_db_command():  # If there is already a database, clear it and do setup
    if update_db():
        click.echo('Database initialized.')


def clear_reward_queue():  # Completely clear the reward queue
    db = get_db()

    try:
        cursor = db.execute(
            "DELETE FROM reward_queue"
        )
        cursor.execute(
            "UPDATE votes SET count = 0"
        )
        db.commit()
    except sqlite3.Error as serror:
        current_app.logger.error(f'Couldn\'t clear reward queue: {serror.args[0]}')
        return False
    return True


def clear_fulfilled_rewards():  # Clears only fulfilled rewards from the queue
    db = get_db()

    try:
        cursor = db.execute(
            "DELETE FROM reward_queue WHERE fulfilled = 1"
        )
        db.commit()
    except sqlite3.Error as serror:
        current_app.logger.error(f'Couldn\'t clear fulfilled rewards: {serror.args[0]}')
        return False
    return True


def rem_all_votes():  # USED TO BE "clear_votes" Clear all votes from the database
    db = get_db()

    try:
        db.execute("DELETE FROM votes")
        db.commit()
    except sqlite3.Error as cverror:
        current_app.logger.error(f'Couldn\'t clear all votes: {cverror.args[0]}')
        return False
    if put_votes(db):
        return True


def rem_vote():  # Remove a single vote from the database
    db = get_db(vote)

    try:
        db.execute("DELETE FROM votes WHERE name = ?", ('vote',))
        db.commit()
    except sqlite3.Error as rverror:
        current_app.logger.error(f'Couldn\'t remove \"{vote}\" from database: {rverror.args[0]}')
        return False
    if put_votes(db):
        return True


def is_cool(reward_name):  # Check if a reward is cooling down.
    current_time = time()  # In seconds
    rewards = current_app.config['REWARDS']
    this_cool_down = rewards[reward_name]['cooldown']
    if this_cool_down > 0:  # If this reward has a cooldown
        try:  # Get the cooldown time for it
            db = get_db()
            cursor = db.execute(
                "SELECT created FROM cool_downs WHERE reward = ?", (reward_name,)
                )
            current_cds = cursor.fetchall()
        except sqlite3.Error as icerror:
            current_app.logger.error(f'Couldn\'t get \"{reward_name}\" from database: {icerror.args[0]}')
        if current_cds:
            last_time = current_cds[0][0]
            hot_time = current_time - last_time
            if hot_time < this_cool_down * 60:  # If it is in the database and still within the cooldown time, return False
                hot_time = this_cool_down - (hot_time / 60)  # Get remaining minutes in cooldown
                return False, hot_time  # Return False, and the time left on the cooldown
            else:  # If it is in the database and the cooldown has expired, update the cooldown time, and return True
                try:
                    db.execute(
                        "UPDATE cool_downs SET created = ? WHERE reward = ?",
                        (current_time, reward_name)
                        )
                    db.commit()
                    return True, 0
                except sqlite3.Error as scerror:
                    current_app.logger.error(f'Couldn\'t update \"{reward_name}\"\'s cooldown time in the database: {scerror.args[0]}')
        else:  # If it is not in the database, add it and return True
            try:
                db.execute(
                        "INSERT INTO cool_downs(created, reward) VALUES(?, ?)",
                        (current_time, reward_name)
                    )
                db.commit()
                return True, 0
            except sqlite3.Error as scerror:
                current_app.logger.error(f'Couldn\'t add \"{reward_name}\" to database: {scerror.args[0]}')
    else:  # If the redeem has no cooldown
        return True, 0


def rem_cool(reward_name):  # Remove a reward from the database
    try:
        db = get_db()
        cursor = db.execute(
            "DELETE FROM cool_downs WHERE reward = ?", (reward_name,)
            )
        current_cds = cursor.fetchall()
    except sqlite3.Error as icerror:
        current_app.logger.error(f'Couldn\'t remove \"{reward_name}\" from database: {icerror.args[0]}')
        return False
    return True


def put_votes(db):  # Reread votes from rewards.py, and sync with database
    info = details['info']
    for reward, details in current_app.config['REWARDS'].items():
        if details['type'] == 'vote':
            try:
                db.execute(
                    "INSERT INTO votes(name, count, info) VALUES(?, 0, ?)",
                    (reward, info)
                )
                db.commit()
            except sqlite3.Error as serror:
                current_app.logger.error(f'Couldn\'t insert \"{vote}\" into database: {serror.args[0]}')
                return False
    return True


def reread_votes():  # Reread votes from rewards.py, and sync with database
    db = get_db()

    try:  # Remove votes not in rewards.py
        cursor = db.execute("SELECT name FROM votes")
        votes = cursor.fetchall()
        chopping_block = []
        for vote in votes:
            vote = vote[0]
            if vote not in current_app.config['REWARDS'].keys():
                chopping_block.append(vote)
            elif current_app.config['REWARDS'][vote]['type'] != 'vote':
                chopping_block.append(vote)
        for vote in chopping_block:
            cursor.execute("DELETE FROM votes WHERE name = ?", (vote,))
        db.commit()
    except sqlite3.Error as serror:
        current_app.logger.error(f'Couldn\'t clear deleted votes from database: {serror.args[0]}')
        return False

    try:  # Add new votes found in rewards.py
        for reward, details in current_app.config['REWARDS'].items():
            if details['type'] == 'vote':  # Check if the vote is already in the database
                cursor = db.execute(
                    "SELECT count FROM votes WHERE name = ?",
                    (reward,)
                )
                result = cursor.fetchone()
                if result is None:  # If it isn't, add it
                    info = details['info']
                    if details['type'] == 'vote':
                        db.execute(
                            "INSERT INTO votes(name, count, info) VALUES(?, 0, ?)",
                            (reward, details["info"])
                        )
                else:  # If it is, update the info
                    cursor.execute(
                        "UPDATE votes SET info = ? WHERE name = ?",
                        (details['info'], reward)
                    )
                db.commit()
    except sqlite3.Error as serror:
        current_app.logger.error(f'Couldn\'t insert \"{vote}\" into database: {serror.args[0]}')
        return False
    return True


def reread_goals():  # Reread goals from rewards.py, and sync with database
    db = get_db()

    try:  # Remove goals not in rewards.py
        cursor = db.execute("SELECT name FROM goals")
        goals = cursor.fetchall()
        chopping_block = []
        for goal in goals:
            goal = goal[0]
            if goal not in current_app.config['REWARDS'].keys():
                chopping_block.append(goal)
            elif current_app.config['REWARDS'][goal]['type'] != 'goal':
                chopping_block.append(goal)
        for goal in chopping_block:
            cursor.execute("DELETE FROM goals WHERE name = ?", (goal,))
        db.commit()
    except sqlite3.Error as serror:
        current_app.logger.error(f'Couldn\'t clear removed goals from database: {serror.args[0]}')
        return False

    try:  # Add new goals found in rewards.py
        for reward, details in current_app.config['REWARDS'].items():
            if details['type'] == 'goal':  # Check if the goal is already in the database
                cursor = db.execute(
                    "SELECT target FROM goals WHERE name = ?",
                    (reward,)
                )
                result = cursor.fetchone()
                if result is None:  # If it isn't, add it
                    cursor.execute(
                        "INSERT INTO goals(name, progress, target, info, complete, milestones) VALUES(?, 0, ?, ?, FALSE, 0)",
                        (reward, details['target'], details['info'])
                    )
                elif result != details['target']:  # If it is, update the target
                    cursor.execute(
                        "UPDATE goals SET target = ?, info = ? WHERE name = ?",
                        (details["target"], details["info"], reward)
                    )
        db.commit()
    except sqlite3.Error as serror:
        current_app.logger.error(f'Couldn\'t insert \"{reward}\" into database: {serror.args[0]}')
        return False
    return True


def reset_goal(goal):  # Set goal progress back to zero
    if goal not in current_app.config['REWARDS']:  # If it doesn't exist in rewards.py
        current_app.logger.error(f'Couldn\'t reset goal, {goal} not in rewards file.')
        return False
    try:
        db = get_db()
        target = current_app.config['REWARDS'][goal]['target']
        info = current_app.config['REWARDS'][goal]['info']
        cursor = db.cursor()
        cursor.execute(
                        "UPDATE goals SET progress = 0, target = ?, complete = FALSE, info = ?, milestones = 0 WHERE name = ?",
                        (target, info, goal)
                    )
        db.commit()
        return True
    except sqlite3.Error as rgerror:
        current_app.logger.error(f'Couldn\'t reset {goal} in the goals table: {rgerror.args[0]}')
        return False


def reset_vote(vote):
    if vote not in current_app.config['REWARDS']:  # Check if it exists in rewards.py
        current_app.logger.error(f'Couldn\'t reset vote, {vote} not in rewards file.')
        return False
    else:
        try:
            db = get_db()
            cursor = db.cursor()
            cursor.execute(
                        "UPDATE votes SET count = 0, voters = '' WHERE name = ?",
                        (vote,)
                    )
            db.commit()
        except sqlite3.Error as rverror:
            current_app.logger.error(f'Couldn\'t reset {vote} in the votes table: {rverror.args[0]}')
            return False
    return True