import sys
import six
from sqlalchemy.orm import Session


class SessionUow(object):
    def __init__(self, session):
        # type: (Session) -> None
        self.session = session

    def execute_with_transaction(self, func):
        try:
            ret = func()
            self.session.commit()
            return ret
        except:
            self.session.rollback()
            six.reraise(*sys.exc_info())

    def execute_method_with_transaction(self, obj, method_name, *args, **kwargs):
        def inner_func():
            method = getattr(obj, method_name)
            return method(*args, **kwargs)

        return self.execute_with_transaction(inner_func)
