hanabi-league/database.py

121 lines
4.1 KiB
Python
Raw Normal View History

2023-11-22 15:31:36 +01:00
import psycopg2
import psycopg2.extensions
2023-11-22 18:22:53 +01:00
import psycopg2.errors
import unidecode
2023-11-22 15:31:36 +01:00
import constants
2023-11-22 15:31:36 +01:00
from config import read_db_config
from log_setup import logger
class DBConnectionManager:
def __init__(self):
self._conn = None
def connect(self):
config = read_db_config()
2023-11-22 16:27:04 +01:00
logger.debug("Establishing database connection with dbname={}, user={}, password={}".format(
config.db_name, config.db_user, config.db_pass
))
2023-11-22 15:31:36 +01:00
self._conn = psycopg2.connect("dbname='{}' user='{}' password='{}' host='localhost'".format(
config.db_name, config.db_user, config.db_pass
))
logger.debug("Established database connection.")
def get_connection(self) -> psycopg2.extensions.connection:
2023-11-22 15:31:36 +01:00
"""
Get the database connection.
If not already connected, this reads the database config file and connects to the DB.
Otherwise, the already active connection is returned.
"""
if self._conn is None:
self.connect()
return self._conn
2023-11-22 18:22:53 +01:00
def get_new_cursor(self) -> psycopg2.extensions.cursor:
if self._conn is None:
self.connect()
return self._conn.cursor()
2023-11-22 15:31:36 +01:00
# Global instance that will hold our DB connection
conn_manager = DBConnectionManager()
def get_existing_tables():
conn = conn_manager.get_connection()
cur = conn.cursor()
table_names = ", ".join("'{}'".format(tablename) for tablename in constants.DB_TABLE_NAMES)
cur.execute(
"SELECT tablename FROM pg_tables WHERE"
" schemaname = 'public' AND"
" tablename IN ({})".format(
table_names
)
)
return [table for (table,) in cur.fetchall()]
2023-11-22 16:21:30 +01:00
def init_database():
"""
Warning: This drops all existing tables from the database
"""
conn = conn_manager.get_connection()
cur = conn.cursor()
with open(constants.DATABASE_SCHEMA_PATH, "r") as f:
cur.execute(f.read())
conn.commit()
logger.verbose("Initialized DB tables.")
2023-11-22 18:22:53 +01:00
def add_player_name(player_name: str):
conn = conn_manager.get_connection()
cur = conn.cursor()
try:
cur.execute("INSERT INTO users (player_name) VALUES (%s)", (player_name,))
conn.commit()
except psycopg2.errors.UniqueViolation:
logger.warn("Player name {} already exists in the database, aborting insertion.".format(player_name))
conn.rollback()
def add_user_name_to_player(hanabi_username: str, player_name: str):
normalized_username = unidecode.unidecode(hanabi_username).lower()
cur = conn_manager.get_new_cursor()
cur.execute("SELECT id FROM users WHERE player_name = (%s)", (player_name,))
user_id = cur.fetchone()
if user_id is None:
logger.error("Display name {} not found in database, cannot add username to it.".format(player_name))
return
else:
cur.execute("SELECT username, player_name from user_accounts "
"INNER JOIN users"
" ON user_accounts.user_id = users.id "
"WHERE "
" normalized_username = (%s)",
(normalized_username,)
)
res = cur.fetchone()
if res is not None:
existing_username, existing_player_name = res
if existing_player_name == player_name:
logger.warn("Hanabi username {} is already registered to player {}, attempted to re-register it.".format(
existing_username, existing_player_name
))
else:
logger.error("Hanabi username {} is already associated to player {}, cannot register it to player {}.".format(
res[0], res[1], player_name
))
return
cur.execute(
"INSERT INTO user_accounts (username, normalized_username, user_id) VALUES (%s, %s, %s)",
(hanabi_username, normalized_username, user_id)
)
conn_manager.get_connection().commit()
def add_player(player_name: str, user_name: str):
add_player_name(player_name)
add_user_name_to_player(user_name, player_name)