WIP: add polls

tlapbot/__init__.py

@@ -25,6 +25,11 @@ def create_app(test_config: None = None) -> Flask:
     app.config.from_pyfile('config.py', silent=True)
     app.config.from_pyfile('redeems.py', silent=True)
 
+    # set up polls if they're enabled
+    if app.config['POLLS']:
+        app.config.from_object('tlapbot.defaults.default_polls')
+        app.config.from_pyfile('polls.py', silent=True)
+
     # Make logging work for gunicorn-ran instances of tlapbot.
     if app.config['GUNICORN']:
         gunicorn_logger = logging.getLogger('gunicorn.error')

tlapbot/db.py

@@ -40,6 +40,20 @@ def insert_counters(db: sqlite3.Connection) -> bool:
     return True
 
 
+def insert_polls(db: sqlite3.Connection) -> bool:
+    for poll, poll_info in current_app.config['POLLS'].items():
+        for option in poll_info["options"]:
+            try:
+                db.execute(
+                    "INSERT INTO polls(poll_name, option, points) VALUES(?, ?, 0)",
+                    (poll, option)
+                )
+            except sqlite3.Error as e:
+                print(f"Failed inserting poll option {poll} {option} to db: {e.args[0]}")
+                return False
+    return True
+
+
 def init_db() -> bool:
     db = get_db()
 
@@ -78,6 +92,32 @@ def refresh_counters() -> bool:
     return insert_counters(db)
 
 
+def refresh_polls() -> bool:
+    db = get_db()
+
+    try:
+        db.execute("DELETE FROM polls")
+        db.commit()
+    except sqlite3.Error as e:
+        print("Error occurred deleting old counters:", e.args[0])
+        return False
+    return insert_polls(db)
+
+
+def reset_poll(poll: str) -> bool:
+    try:
+        db = get_db()
+        db.execute(
+            "UPDATE polls SET points = 0 WHERE poll_name = ?",
+            (poll,)
+        )
+        db.commit()
+        return True
+    except sqlite3.Error as e:
+        current_app.logger.error(f"Error occurred adding a milestone: {e.args[0]}")
+        return False
+
+
 def refresh_milestones() -> bool:
     db = get_db()
     # delete old milestones
@@ -208,6 +248,25 @@ def hard_reset_milestone_command(milestone: str) -> None:
         click.echo(f"Hard reset milestone {milestone}.")
 
 
+@click.command('refresh-polls')
+@with_appcontext
+def refresh_polls_command() -> None:
+    """Initialize all polls from the polls file,
+    delete polls not in polls file."""
+    if refresh_polls():
+        click.echo('Refreshed polls.')
+
+
+@click.command('reset-polls')
+@click.argument('milestone')
+def reset_poll_command(poll: str) -> None:
+    """Resets polls progress back to zero."""
+    if reset_poll(poll):
+        click.echo(f"Reset poll {poll}.")
+    else:
+        click.echo(f"Could not reset poll {poll}.")
+
+
 def init_app(app: Flask) -> None:
     app.teardown_appcontext(close_db)
     app.cli.add_command(init_db_command)

tlapbot/defaults/default_config.py

@@ -7,4 +7,5 @@ PASSIVE=False
 LIST_REDEEMS=False
 ACTIVE_CATEGORIES=[]
 GUNICORN=False
-PREFIX='!'
+PREFIX='!'
+POLLS=True

tlapbot/defaults/default_polls.py

@@ -0,0 +1,3 @@
+POLLS={
+    "favourite_food": {"info": "Vote for your favourite food.", "options": ["pizza", "soup", "curry", "schnitzel"]}
+}

tlapbot/owncast_webhooks.py

@@ -7,6 +7,7 @@ from tlapbot.owncast_helpers import (add_user_to_database, change_display_name,
         read_users_points, remove_duplicate_usernames)
 from tlapbot.help_message import send_help
 from tlapbot.redeems_handler import handle_redeem
+from tlapbot.poll_handler import handle_poll_vote
 
 
 bp = Blueprint('owncast_webhooks', __name__)
@@ -57,6 +58,8 @@ def owncast_webhook() -> Any | None:
                 change_display_name(db, user_id, display_name)
                 if data["eventData"]["user"]["authenticated"]:
                     remove_duplicate_usernames(db, user_id, display_name)
+            elif data["eventData"]["rawBody"].startswith(f"{prefix}vote"):
+                handle_poll_vote(data["eventData"]["rawBody"], user_id)
             elif data["eventData"]["rawBody"].startswith(prefix):
                 handle_redeem(data["eventData"]["rawBody"], user_id)
     return data

tlapbot/poll_handler.py

@@ -0,0 +1,31 @@
+from sqlite3 import Connection
+from tlapbot.db import get_db
+from tlapbot.owncast_requests import send_chat
+from tlapbot.owncast_helpers import read_users_points, use_points
+from tlapbot.redeems.polls import poll_option_exists, vote_in_poll
+
+
+def handle_poll_vote(message: str, user_id: str) -> None:
+    split_message = message[5:].split(maxsplit=1)
+    if len(split_message < 3):
+        send_chat("Can't vote for poll, not enough arguments in message.")
+    poll_name = split_message[0]
+    option = split_message[1]
+    poll_points = split_message[2]
+
+    db: Connection = get_db()
+    if not poll_option_exists(db, poll_name, option):
+        send_chat("Can't vote for poll, poll and option combination not found.")
+        return
+    if not poll_points.isdigit():
+        send_chat(f"Can't vote in {poll_name} poll, donation amount is not a number.")
+        return
+    user_points = read_users_points(db, user_id)
+    if not user_points:
+        send_chat(f"Can't vote in {poll_name} poll, failed to read users' points.")
+        return
+    if user_points < poll_points:
+        send_chat(f"Can't vote in {poll_name} poll, you're voting with more points than you have.")
+    if (vote_in_poll(db, poll_name, option, poll_points) and 
+            use_points(db, user_id, poll_points)):
+        send_chat(f"{poll_points} donated to {option} in poll {poll_name}")

tlapbot/redeems/polls.py

@@ -0,0 +1,38 @@
+from flask import current_app
+from sqlite3 import Error, Connection
+
+
+def poll_option_exists(db: Connection, poll_name: str, option: str) -> bool | None:
+    """Returns None only if error was logged."""
+    try:
+        cursor = db.execute(
+            "SELECT poll_name, option FROM polls WHERE poll_name = ? and option = ?",
+            (poll_name, option)
+        )
+        counter = cursor.fetchone()
+        if counter is None:
+            current_app.logger.warning("Poll and option combination not found in database.")
+            current_app.logger.warning("Maybe you forgot to run the refresh-polls CLI command "
+                                       "after you added a new poll to the config?")
+            return False
+        return True
+    except Error as e:
+        current_app.logger.error(f"Error occurred checking if poll exists: {e.args[0]}")
+        current_app.logger.error(f"For poll&option: {poll_name}:{option}")
+
+
+def vote_in_poll(db: Connection, poll_name: str, option: str, points: int) -> bool | None:
+    if points <= 0:
+        current_app.logger.warning(f"Vote for {poll_name}:{option} was not a positive integer.")
+        return
+    if poll_option_exists(db, poll_name, option):
+        try:
+            cursor = db.execute(
+                "UPDATE polls SET progress = progress + ? WHERE poll_name = ? and option = ?",
+                (points, poll_name, option)
+                )
+            db.commit()
+            return True
+        except Error as e:
+            current_app.logger.error(f"Error occurred updating milestone: {e.args[0]}")
+        return False

tlapbot/schema.sql

@@ -28,4 +28,11 @@ CREATE TABLE redeem_queue (
   redeemer_id TEXT NOT NULL,
   note TEXT,
   FOREIGN KEY (redeemer_id) REFERENCES points (id)
-);
+);
+
+CREATE TABLE polls (
+  id INTEGER PRIMARY KEY AUTOINCREMENT,
+  points INTEGER NOT NULL,
+  option TEXT NOT NULL,
+  poll_name TEXT NOT NULL,
+)