80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
import logging
|
|
|
|
import sqlalchemy
|
|
from sshtunnel import SSHTunnelForwarder
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
from config import SSH_HOST, SSH_USERNAME, SSH_PASSWORD, DB_NAME, DB_USER, DB_PASSWORD, DB_HOST, DB_PORT, SSH_PORT
|
|
|
|
|
|
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger('db_connector.py')
|
|
stream_handler = logging.StreamHandler()
|
|
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
stream_handler.setFormatter(formatter)
|
|
logger.addHandler(stream_handler)
|
|
|
|
|
|
class RemoteDB:
|
|
def __init__(self):
|
|
self.ssh_host = SSH_HOST
|
|
self.ssh_username = SSH_USERNAME
|
|
self.ssh_password = SSH_PASSWORD
|
|
self.db_name = DB_NAME
|
|
self.db_user = DB_USER
|
|
self.db_password = DB_PASSWORD
|
|
self.db_host = DB_HOST
|
|
self.db_port = DB_PORT
|
|
self.ssh_port = SSH_PORT
|
|
self.tunnel = None
|
|
self.engine = None
|
|
self.Session = None
|
|
self._connect()
|
|
|
|
def _connect(self):
|
|
try:
|
|
self.tunnel = SSHTunnelForwarder(
|
|
(self.ssh_host, self.ssh_port),
|
|
ssh_username=self.ssh_username,
|
|
ssh_password=self.ssh_password,
|
|
remote_bind_address=(self.db_host, self.db_port)
|
|
)
|
|
self.tunnel.start()
|
|
|
|
local_port = self.tunnel.local_bind_port
|
|
db_url = f"postgresql://{self.db_user}:{self.db_password}@localhost:{local_port}/{self.db_name}"
|
|
self.engine = create_engine(db_url)
|
|
self.Session = sessionmaker(bind=self.engine)
|
|
except Exception as e:
|
|
logger.exception(f"Connection failed: {e}")
|
|
|
|
def execute_query(self, query: str):
|
|
session = self.Session()
|
|
try:
|
|
result = session.execute(sqlalchemy.text(query))
|
|
session.commit()
|
|
return result.fetchall()
|
|
except Exception as e:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|
|
|
|
def execute_command(self, cmd: str):
|
|
session = self.Session()
|
|
try:
|
|
result = session.execute(sqlalchemy.text(cmd))
|
|
session.commit()
|
|
logger.debug(f"Command {cmd} committed.")
|
|
except Exception as e:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|
|
|
|
def close(self):
|
|
if self.engine:
|
|
self.engine.dispose()
|
|
if self.tunnel:
|
|
self.tunnel.stop()
|