91 lines
2.6 KiB
Python
91 lines
2.6 KiB
Python
import MySQLdb
|
|
import MySQLdb.cursors
|
|
|
|
from typing import Union, Optional
|
|
|
|
from .connection import Connection
|
|
from .logger import Logger
|
|
|
|
|
|
logger = Logger()
|
|
|
|
|
|
class Database:
|
|
"""Class wrapping MySQL database connection"""
|
|
|
|
def __init__(self, vessel):
|
|
"""Initialize a new Database object"""
|
|
self.vessel = vessel
|
|
self._con = None
|
|
self._ssh = None
|
|
|
|
def _execute(
|
|
self,
|
|
query: str,
|
|
parameters: Optional[tuple] = None,
|
|
ctype: Optional[MySQLdb.cursors.BaseCursor] = None,
|
|
retry: bool = True,
|
|
):
|
|
"""Execute a query on the database
|
|
|
|
Args:
|
|
query (str): SQL query to execute
|
|
parameters (tuple, optional): Parameters to use to replace
|
|
placeholders in the query, if any. Defaults to None.
|
|
"""
|
|
try:
|
|
cur = self.getCursor(ctype)
|
|
cur.execute(query, parameters)
|
|
self.commit() # Instantly commit after every (potential) write action
|
|
return cur.fetchall()
|
|
except (AttributeError, MySQLdb.OperationalError):
|
|
if retry:
|
|
self._connect()
|
|
return self._execute(query, parameters, ctype, False)
|
|
raise
|
|
except MySQLdb.ProgrammingError:
|
|
logger.error(f"Error in query: {query}")
|
|
logger.error(f"Parameters: {str(parameters)}")
|
|
raise
|
|
|
|
def _connect(self) -> None:
|
|
if self.vessel.ssh:
|
|
self._ssh = Connection(self.vessel)
|
|
port = self._ssh.forward_tcp(3306)
|
|
host = "127.0.0.1"
|
|
else:
|
|
port = 3306
|
|
host = self.vessel.host
|
|
|
|
self._con = MySQLdb.connect(
|
|
host=host,
|
|
user=self.vessel.username,
|
|
passwd=self.vessel.password,
|
|
db=self.vessel.database,
|
|
port=port,
|
|
)
|
|
|
|
def commit(self) -> None:
|
|
"""Commit the current database transaction
|
|
|
|
N.B.: Commit instantly after every write action to make the database
|
|
"thread-safe". Connections will time out if the database is locked for
|
|
more than five seconds.
|
|
"""
|
|
self._con.commit()
|
|
|
|
def getCursor(
|
|
self, ctype: Optional[MySQLdb.cursors.BaseCursor] = None
|
|
) -> MySQLdb.cursors.BaseCursor:
|
|
"""Return a cursor to operate on the MySQL database
|
|
|
|
Returns:
|
|
MySQLdb.Cursor: Cursor object to execute queries on
|
|
"""
|
|
return self._con.cursor(ctype)
|
|
|
|
def __del__(self):
|
|
"""Close database connection on removal of the Database object"""
|
|
if self._con:
|
|
self._con.close()
|