parley/storage.py

import sqlite3 as sql

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

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

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

    def get_connections(self, game: str):
        return self.connections.get(game, [])

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

    def get_backlog(self, game: str, c=None) -> [(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=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=None):
        if c is None:
            c = self.db.cursor()

        game_id = 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()