parley/storage.py

from sanic.websocket import WebSocketConnection # type: ignore
import sqlite3 as sql
from typing import Dict, List, Optional, Tuple

class Cache:
    connections: Dict[str, WebSocketConnection]

    def __init__(self):
        self.connections = {}

    def add_connection(self, game: str, ws: WebSocketConnection):
        if game not in self.connections:
            self.connections[game] = set()
        self.connections[game].add(ws)

    def remove_connection(self, game: str, ws: WebSocketConnection):
        if game in self.connections:
            self.connections[game].remove(ws)

    def get_connections(self, game: str) -> List[WebSocketConnection]:
        return self.connections.get(game, [])

class Storage:
    db: sql.Connection

    def __init__(self):
        self.db = sql.connect('samp.db')

    def get_backlog(self, game: str, c: Optional[sql.Cursor]=None) -> List[Tuple[str, str]]:
        if c is None:
            c = self.db.cursor()

        c.execute('SELECT c.user, c.content '
                  'FROM chat c, games g '
                  'WHERE c.game = g.id AND g.name = ?', [game])
        return list(c)

    def get_game_id(self, game: str, c: Optional[sql.Cursor]=None) -> int:
        if c is None:
            c = self.db.cursor()

        c.execute('SELECT id FROM games WHERE name = ?', [game])
        res = c.fetchone()
        if res is not None:
            return res[0]
        else:
            c.execute('INSERT INTO games (name) VALUES (?)', [game])
            self.db.commit()
            c.execute('SELECT id FROM games WHERE name = ?', [game])
            return c.fetchone()[0]

    def add_msg(self, user: str, content: str, game: str, c: Optional[sql.Cursor]=None):
        if c is None:
            c = self.db.cursor()

        game_id: int = self.get_game_id(game, c)
        c.execute('INSERT INTO chat (game, user, content) '
                  'VALUES (?, ?, ?)', [game_id, user, content])
        self.db.commit()

    def close(self):
        self.db.close()