futool-db-lite/session/session.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/6/22 19:53
# @Author  : old tom
# @File    : session.py
# @Project : futool-db-lite
# @Desc    :
from sqlalchemy import Engine
from executor.sql_executor import SQLExecutor
from transaction.connect_transaction import TransactionFactory


class CreateSessionError(Exception):
    def __init__(self, msg):
        Exception.__init__(self, msg)


class SqlsessionFactory(object):
    """
    工厂模式创建sqlSession
    """

    def __init__(self, engine: Engine):
        self._engine = engine

    def open_session(self):
        try:
            conn = self._engine.connect()
            tx_factory = TransactionFactory(conn)
            tx = tx_factory.create_transaction()
            executor = SQLExecutor(tx)
            return Sqlsession(executor)
        except Exception as e:
            raise CreateSessionError(msg=f'创建sqlSession异常,e={e}')


class Sqlsession(object):
    def __init__(self, executor: SQLExecutor):
        self.executor = executor

    def select_one(self, sql):
        return self.executor.query(sql).fetchone()

    def select_many(self, sql, size):
        return self.executor.query(sql).fetchmany(size)

    def select_all(self, sql):
        return self.executor.query(sql).fetchall()

    def insert(self, sql):
        return self.executor.execute_update(sql)

    def update(self, sql):
        return self.executor.execute_update(sql)

    def delete(self, sql):
        return self.executor.execute_update(sql)

    def begin_transaction(self):
        """
        开启事务
        :return:
        """
        self.executor.get_transaction().begin_transaction()

    def commit(self):
        """
        提交
        :return:
        """
        self.executor.get_transaction().commit()

    def rollback(self):
        """
        回滚
        :return:
        """
        self.executor.get_transaction().rollback()

    def close(self):
        """
        关闭连接
        :return:
        """
        self.executor.get_connection().close()