<reponame>sachinraina112/flask_app
import logging
from contextlib import contextmanager

import sqlalchemy as db
import sqlalchemy.orm as orm
from AppConfig import db_config
from db.Security import Security
from sqlalchemy.exc import DBAPIError, SQLAlchemyError

logger = logging.getLogger(__name__)



class DBEngine:

    _engine = None
    _Session = None

    @staticmethod
    def _attach_pool_events():
        # Called when a connection returns to the pool.
        @db.event.listens_for(DBEngine._engine, 'checkin')
        def receive_checkin(dbapi_connection, connection_record):
            logger.debug(f"Connection returned to the pool with id {id(dbapi_connection)}")

        # Called when a connection is retrieved from the Pool.
        @db.event.listens_for(DBEngine._engine, 'checkout')
        def receive_checkout(dbapi_connection, connection_record, connection_proxy):
            logger.debug(f"Connection retrieved from the pool with id {id(dbapi_connection)}")

        # Called when a DBAPI connection is closed.
        @db.event.listens_for(DBEngine._engine, 'close')
        def receive_close(dbapi_connection, connection_record):
            logger.debug(f"Connection to the DB is closed with id {id(dbapi_connection)}")

        # Called at the moment a particular DBAPI connection is first created for a given Pool.
        @db.event.listens_for(DBEngine._engine, 'connect')
        def receive_connect(dbapi_connection, connection_record):
            logger.debug(f"New DB connection is created with id {id(dbapi_connection)}")

        # Called when a DBAPI connection is to be “invalidated”.
        @db.event.listens_for(DBEngine._engine, 'invalidate')
        def receive_invalidate(dbapi_connection, connection_record, exception):
            logger.debug(f"DB connection is invalidated with id {id(dbapi_connection)}")

    @staticmethod
    def _attach_connection_events():

        @db.event.listens_for(DBEngine._engine, 'before_execute')
        def receive_before_execute(conn, clauseelement, multiparams, params):
            logger.info(f"Executing SQL statement {clauseelement} with parameter {params} on connection {id(conn.connection.connection)}")

        @db.event.listens_for(DBEngine._engine, 'after_execute')
        def receive_after_execute(conn, clauseelement, multiparams, params, result):
            logger.info(f"Query executed successfully on connection {id(conn.connection.connection)}")

        @db.event.listens_for(DBEngine._engine, "handle_error")
        def handle_exception(context):
            logger.error(f"Exception occurred while performing DB operation {context}")

    @staticmethod
    def create_db_engine():

        db_user = db_config['user']
        sec = Security()

        password = sec.decrypt(sec.string_to_bytes(db_config["password"]))
        database = db_config['database']
        host = db_config['host']
        port = db_config['port']
        database_url = f"mysql+pymysql://{db_user}:{password}@{host}:{port}/{database}"
        DBEngine._engine = db.create_engine(database_url, pool_size=5, max_overflow=10, pool_recycle=3600, pool_pre_ping=True, logging_name="sql-pool")
        DBEngine._attach_pool_events()
        DBEngine._attach_connection_events()
        return DBEngine._engine

    @staticmethod
    def get_engine():
        if DBEngine._engine is None:
            DBEngine._engine = DBEngine.create_db_engine()
            logger.info("SQL engine created successfully")
        return DBEngine._engine

    @staticmethod
    @contextmanager
    def get_session():
        engine = DBEngine.get_engine()
        if DBEngine._Session is None:
            session_factory = orm.sessionmaker(bind=engine)
            DBEngine._Session = orm.scoped_session(session_factory)
            logger.info("Session factory created successfully")
        session = DBEngine._Session()
        try:
            yield session
            session.commit()
        except (DBAPIError, SQLAlchemyError) as e:
            logger.error(f"Error occurred while doing db operation {e}")
            logger.exception(e)
            session.rollback()
            raise
        finally:
            # TODO: check this is required
            session.close()

    @staticmethod
    def get_object_session(obj):
        return DBEngine._Session.object_session(obj)

import logging
from contextlib import contextmanager

import sqlalchemy as db
import sqlalchemy.orm as orm
from AppConfig import db_config
from db.Security import Security
from sqlalchemy.exc import DBAPIError, SQLAlchemyError

logger = logging.getLogger(__name__)



class DBEngine:

    _engine = None
    _Session = None

    @staticmethod
    def _attach_pool_events():
        # Called when a connection returns to the pool.
        @db.event.listens_for(DBEngine._engine, 'checkin')
        def receive_checkin(dbapi_connection, connection_record):
            logger.debug(f"Connection returned to the pool with id {id(dbapi_connection)}")

        # Called when a connection is retrieved from the Pool.
        @db.event.listens_for(DBEngine._engine, 'checkout')
        def receive_checkout(dbapi_connection, connection_record, connection_proxy):
            logger.debug(f"Connection retrieved from the pool with id {id(dbapi_connection)}")

        # Called when a DBAPI connection is closed.
        @db.event.listens_for(DBEngine._engine, 'close')
        def receive_close(dbapi_connection, connection_record):
            logger.debug(f"Connection to the DB is closed with id {id(dbapi_connection)}")

        # Called at the moment a particular DBAPI connection is first created for a given Pool.
        @db.event.listens_for(DBEngine._engine, 'connect')
        def receive_connect(dbapi_connection, connection_record):
            logger.debug(f"New DB connection is created with id {id(dbapi_connection)}")

        # Called when a DBAPI connection is to be “invalidated”.
        @db.event.listens_for(DBEngine._engine, 'invalidate')
        def receive_invalidate(dbapi_connection, connection_record, exception):
            logger.debug(f"DB connection is invalidated with id {id(dbapi_connection)}")

    @staticmethod
    def _attach_connection_events():

        @db.event.listens_for(DBEngine._engine, 'before_execute')
        def receive_before_execute(conn, clauseelement, multiparams, params):
            logger.info(f"Executing SQL statement {clauseelement} with parameter {params} on connection {id(conn.connection.connection)}")

        @db.event.listens_for(DBEngine._engine, 'after_execute')
        def receive_after_execute(conn, clauseelement, multiparams, params, result):
            logger.info(f"Query executed successfully on connection {id(conn.connection.connection)}")

        @db.event.listens_for(DBEngine._engine, "handle_error")
        def handle_exception(context):
            logger.error(f"Exception occurred while performing DB operation {context}")

    @staticmethod
    def create_db_engine():

        db_user = db_config['user']
        sec = Security()

        password = sec.decrypt(sec.string_to_bytes(db_config["password"]))
        database = db_config['database']
        host = db_config['host']
        port = db_config['port']
        database_url = f"mysql+pymysql://{db_user}:{password}@{host}:{port}/{database}"
        DBEngine._engine = db.create_engine(database_url, pool_size=5, max_overflow=10, pool_recycle=3600, pool_pre_ping=True, logging_name="sql-pool")
        DBEngine._attach_pool_events()
        DBEngine._attach_connection_events()
        return DBEngine._engine

    @staticmethod
    def get_engine():
        if DBEngine._engine is None:
            DBEngine._engine = DBEngine.create_db_engine()
            logger.info("SQL engine created successfully")
        return DBEngine._engine

    @staticmethod
    @contextmanager
    def get_session():
        engine = DBEngine.get_engine()
        if DBEngine._Session is None:
            session_factory = orm.sessionmaker(bind=engine)
            DBEngine._Session = orm.scoped_session(session_factory)
            logger.info("Session factory created successfully")
        session = DBEngine._Session()
        try:
            yield session
            session.commit()
        except (DBAPIError, SQLAlchemyError) as e:
            logger.error(f"Error occurred while doing db operation {e}")
            logger.exception(e)
            session.rollback()
            raise
        finally:
            # TODO: check this is required
            session.close()

    @staticmethod
    def get_object_session(obj):
        return DBEngine._Session.object_session(obj)
# Called when a connection returns to the pool. # Called when a connection is retrieved from the Pool. # Called when a DBAPI connection is closed. # Called at the moment a particular DBAPI connection is first created for a given Pool. # Called when a DBAPI connection is to be “invalidated”. # TODO: check this is required