reportmonster/src/reportmonster/classes/database.py

69 lines
2.1 KiB
Python

import MySQLdb
import MySQLdb.cursors
from typing import Union, Optional
from classes.connection import Connection
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
self._connect()
def _execute(self, query: str, parameters: Optional[tuple] = None, ctype: Optional[MySQLdb.cursors.BaseCursor] = None) -> None:
"""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.
"""
cur = self.getCursor(ctype)
cur.execute(query, parameters)
self.commit() # Instantly commit after every (potential) write action
return cur.fetchall()
def _connect(self):
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
"""
self._con.close()