SQLSource.py

# coding: utf-8

from contextlib import contextmanager
import logging

import pandas as pd
from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker


class SQLSource:
    """
    sql_source = SQLSource()
    with sql_source() as session:
        with SQLSource.transaction_scope(session) as session:
            session.query(...)
    """

    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.connection_string = ...
        self.db_name = self.connection_string.split('@')[1].split(':')[0]
        self.engine = create_engine(self.connection_string)
        self.Session = scoped_session(sessionmaker(bind=self.engine))


    @contextmanager
    def __call__(self):
        try:
            self.logger.debug('Scoped session to database ' + self.db_name)
            yield self.Session()
        finally:
            self.Session.remove()
        return 


    @staticmethod
    @contextmanager
    def transaction_scope(session):
        try:
            yield session
            lself.logger.debug('Session commit')
            session.commit()
        except:
            self.logger.debug('Session rollback')
            session.rollback()
            raise
        finally:
            self.logger.debug('Session close')
            session.close()