SQLAlchemy 事务管理

sa_trsanction.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# date:        2019/11/16
# author:      he.zhiming
#

from __future__ import (absolute_import, unicode_literals)

import sqlalchemy
import functools


class SessionContext:
    """
    使用方式:
        with SessoinContext() as session:
            # code

    """
    def __init__(self, session_obj=None, using=None):
        nested = True
        if session_obj is None:
            session_obj = get_new_session()
            nested = False

        self._session = session_obj
        self._nested = nested


    def __enter__(self):
        return self._session

    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            if exc_type is None:
                try:
                    self._session.commit()
                except:
                    self._session.rollback()
                    raise
                else:
                    self.exit()
            else:
                self._session.rollback()
        finally:
            if not self._nested:
                self._session.close()

    def exit(self):
        transaction_commit_signal.signal()


def transactional(using=None):
    """

    使用方式:
        @transactional()
        def get_user(id_):
            session只能add,不能做commit操作

    :param using:
    :return:
    """
    def _decorator(service_func):
        @functools.wraps(service_func)
        def __wrapper(*args, **kwargs):
            has_classarg = False
            if args:
                funcargs = getattr(args[0], service_func.__name__, None)
                if funcargs is not None and hasattr(funcargs, '__self__'):
                    has_classarg = True


            with SessionContext(using=using) as session:
                if has_classarg:
                    classarg = args[0]
                    args = args[1:]
                    return service_func(classarg, session, *args, **kwargs)
                else:
                    return service_func(session, *args, **kwargs)

        return __wrapper

    return _decorator

class C:
    def __enter__(self):
        return "123"

    def __exit__(self, exc_type, exc_val, exc_tb):

        print(1)

        raise exc_type(exc_val)



def test_session():
    # with SessionContext()

    with SessionContext() as session: