Py-Hanabi/hanabi/database/database.py

96 lines
3.4 KiB
Python
Raw Normal View History

2023-07-04 18:53:18 +02:00
from typing import Optional
from pathlib import Path
import yaml
import psycopg2
import platformdirs
from hanabi import constants
from hanabi import logger
class LazyDBCursor:
def __init__(self):
self.__cur: Optional[psycopg2.cursor] = None
def __getattr__(self, item):
if self.__cur is None:
raise ValueError(
"DB cursor used in uninitialized state. Did you forget to initialize the DB connection?"
)
return getattr(self.__cur, item)
def set_cur(self, cur):
self.__cur = cur
class LazyDBConnection:
def __init__(self):
self.__conn: Optional[psycopg2.connection] = None
def __getattr__(self, item):
if self.__conn is None:
raise ValueError(
"DB connection used in uninitialized state. Did you forget to initialize the DB connection?"
)
return getattr(self.__conn, item)
def set_conn(self, conn):
self.__conn = conn
class DBConnectionManager:
def __init__(self):
self.lazy_conn: LazyDBConnection = LazyDBConnection()
self.lazy_cur: LazyDBCursor = LazyDBCursor()
self.config_file = Path(platformdirs.user_config_dir(constants.APP_NAME, ensure_exists=True)) / 'config.yaml'
self.db_name: str = constants.DEFAULT_DB_NAME
self.db_user: str = constants.DEFAULT_DB_USER
2023-07-05 20:59:20 +02:00
self.db_pass: Optional[str] = None
def read_config(self):
logger.debug("DB connection configuration read from {}".format(self.config_file))
if self.config_file.exists():
with open(self.config_file, "r") as f:
config = yaml.safe_load(f)
self.db_name = config.get('dbname', None)
self.db_user = config.get('dbuser', None)
2023-07-05 20:59:20 +02:00
self.db_pass = config.get('dbpass', None)
if self.db_name is None:
logger.verbose("Falling back to default database name {}".format(constants.DEFAULT_DB_NAME))
self.db_name = constants.DEFAULT_DB_NAME
if self.db_user is None:
logger.verbose("Falling back to default database user {}".format(constants.DEFAULT_DB_USER))
self.db_user = constants.DEFAULT_DB_USER
else:
logger.info(
"No configuration file for database connection found, falling back to default values "
"(dbname={}, dbuser={}).".format(
constants.DEFAULT_DB_NAME, constants.DEFAULT_DB_USER
)
)
logger.info(
"Note: To turn off this message, create a config file at {}".format(self.config_file)
)
def create_config_file(self):
if self.config_file.exists():
raise FileExistsError("Configuration file already exists, not overriding.")
self.config_file.write_text(
"dbname: {}\n"
2023-07-05 20:59:20 +02:00
"dbuser: {}\n"
"dbpass: null".format(
constants.DEFAULT_DB_NAME,
constants.DEFAULT_DB_USER
)
)
logger.info("Initialised default config file {}".format(self.config_file))
def connect(self):
2023-07-05 22:53:35 +02:00
conn = psycopg2.connect("dbname='{}' user='{}' password='{}' host='localhost'".format(
self.db_name, self.db_user, self.db_pass)
)
cur = conn.cursor()
self.lazy_conn.set_conn(conn)
self.lazy_cur.set_cur(cur)