5938fd95de
* Added exception in pool * Solved issue with message codec * Changed pg pool to normal pg connection * fix(redshift-connector): Fixed close connection when exception
66 lines
2 KiB
Python
66 lines
2 KiB
Python
import logging
|
|
import time
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy import MetaData
|
|
from sqlalchemy.orm import sessionmaker, session
|
|
from contextlib import contextmanager
|
|
import logging
|
|
from decouple import config as _config
|
|
from decouple import Choices
|
|
from contextlib import contextmanager
|
|
from decouple import config
|
|
|
|
logging.basicConfig(level=config("LOGLEVEL", default=logging.INFO))
|
|
logging.getLogger('apscheduler').setLevel(config("LOGLEVEL", default=logging.INFO))
|
|
|
|
sslmode = _config('DB_SSLMODE',
|
|
cast=Choices(['disable', 'allow', 'prefer', 'require', 'verify-ca', 'verify-full']),
|
|
default='allow'
|
|
)
|
|
|
|
conn_str = config('string_connection', default='')
|
|
if conn_str == '':
|
|
pg_host = config("pg_host")
|
|
pg_dbname = config("pg_dbname")
|
|
pg_user = config("pg_user")
|
|
pg_password = config("pg_password")
|
|
pg_port = config("pg_port", cast=int)
|
|
else:
|
|
import urllib.parse
|
|
conn_str = urllib.parse.unquote(conn_str)
|
|
usr_info, host_info = conn_str.split('@')
|
|
i = usr_info.find('://')
|
|
pg_user, pg_password = usr_info[i+3:].split(':')
|
|
host_info, pg_dbname = host_info.split('/')
|
|
i = host_info.find(':')
|
|
if i == -1:
|
|
pg_host = host_info
|
|
pg_port = 5432
|
|
else:
|
|
pg_host, pg_port = host_info.split(':')
|
|
pg_port = int(pg_port)
|
|
|
|
conn_str = f"postgresql://{pg_user}:{pg_password}@{pg_host}:{pg_port}/{pg_dbname}"
|
|
|
|
class PostgresClient:
|
|
CONNECTION_STRING: str = conn_str
|
|
_sessions = sessionmaker()
|
|
|
|
def __init__(self):
|
|
self.engine = create_engine(self.CONNECTION_STRING, connect_args={'sslmode': sslmode})
|
|
|
|
@contextmanager
|
|
def get_live_session(self) -> session:
|
|
"""
|
|
This is a session that can be committed.
|
|
Changes will be reflected in the database.
|
|
"""
|
|
# Automatic transaction and connection handling in session
|
|
connection = self.engine.connect()
|
|
my_session = type(self)._sessions(bind=connection)
|
|
|
|
yield my_session
|
|
|
|
my_session.close()
|
|
connection.close()
|
|
|