79 lines
3.2 KiB
Python
79 lines
3.2 KiB
Python
import sqlalchemy as sq
|
|
import sqlparse
|
|
import pandas as pd
|
|
|
|
from config import DatabaseConfig, DatabaseType
|
|
from keepass import KeePass
|
|
|
|
|
|
def _read_and_sql_file_and_strip_for_comments(filename: str):
|
|
query: str
|
|
with open(filename, 'r') as f:
|
|
query = f.read()
|
|
query = sqlparse.format(query, strip_comments=True)
|
|
return query
|
|
|
|
|
|
class DBAdapter:
|
|
_engine: sq.Engine
|
|
_database_config: DatabaseConfig
|
|
def __init__(self, keepass: KeePass, database_config: DatabaseConfig):
|
|
self._database_config = database_config
|
|
connection_string: str
|
|
keepass_entry = keepass.get_db_credentials()
|
|
|
|
match self._database_config.type:
|
|
case DatabaseType.PSQL:
|
|
connection_string = f'postgresql+pg8000://{keepass_entry.name}:{keepass_entry.password}@{self._database_config.host}/{self._database_config.name}'
|
|
case DatabaseType.ORCL:
|
|
connection_string = f'oracle:{keepass_entry.name}:{keepass_entry.password}@{self._database_config.host}:{self._database_config.port}:{self._database_config.ssid}'
|
|
case _:
|
|
raise Exception(f'Database type {database_config.type} is not supported')
|
|
|
|
self._engine: sq.Engine = sq.create_engine(connection_string)
|
|
|
|
def _set_transaction_readonly(self, conn: sq.Connection):
|
|
if not conn.in_transaction():
|
|
raise Exception('Connection is not in a transaction')
|
|
|
|
match self._database_config.type:
|
|
case DatabaseType.PSQL:
|
|
conn.execute(sq.text('SET TRANSACTION READ ONLY'))
|
|
case DatabaseType.ORCL:
|
|
conn.execute(sq.text('SET TRANSACTION READ ONLY'))
|
|
case _:
|
|
raise Exception(f'Database type {self.database_config.type} is not supported for readonly transactions')
|
|
def run_query(self, query: str, read_only = True) -> sq.CursorResult:
|
|
result: sq.CursorResult
|
|
with self._engine.connect() as conn:
|
|
with conn.begin():
|
|
try:
|
|
if read_only:
|
|
self._set_transaction_readonly(conn)
|
|
result = conn.execute(sq.text(query))
|
|
conn.commit()
|
|
except Exception as e:
|
|
conn.rollback()
|
|
raise e
|
|
return result
|
|
|
|
def run_sql_file_one_statement(self, filename: str = "query.sql", read_only = True) -> sq.CursorResult:
|
|
query = _read_and_sql_file_and_strip_for_comments(filename)
|
|
return self.run_query(query, read_only)
|
|
|
|
def run_sql_file_export_to_csv(self, input_name: str = "query.sql", output_name: str = "export.csv", read_only = True):
|
|
result: pd.DataFrame
|
|
query = _read_and_sql_file_and_strip_for_comments(input_name)
|
|
|
|
with self._engine.connect() as conn:
|
|
with conn.begin():
|
|
try:
|
|
if read_only:
|
|
self._set_transaction_readonly(conn)
|
|
result = pd.read_sql(query, conn)
|
|
conn.commit()
|
|
except Exception as e:
|
|
conn.rollback()
|
|
raise e
|
|
result.to_csv(output_name, index=False, sep=';')
|