93 lines
3.3 KiB
Python
93 lines
3.3 KiB
Python
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
|
|
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)
|
|
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"
|
|
"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):
|
|
conn = psycopg2.connect("dbname={} user={}".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)
|