futool-db-lite/session/engine/dbengine.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/6/22 19:54
# @Author  : old tom
# @File    : dbengine.py
# @Project : futool-db-lite
# @Desc    :
from urllib.parse import quote_plus as urlquote
from sqlalchemy import create_engine
from config.db_config import DbConfigLoader, DEFAULT_CONF_PATH
from session.session import SqlsessionFactory


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


class DatabaseEngine(object):
    """
    数据库连接
    dialect+driver://username:password@host:port/database
    mysql: mysql+pymysql://scott:tiger@localhost/foo
    sqlserver: mssql+pymssql://scott:tiger@hostname:port/dbname
    """
    DB_URL = {
        'postgresql': 'postgresql+psycopg2://{0}:{1}@{2}:{3}/{4}',
        'oracle': 'oracle+cx_oracle://{0}:{1}@{2}:{3}/?service_name={4}'
    }

    def __init__(self, db_name, conf_path=DEFAULT_CONF_PATH):
        """
        pool_size:连接池大小
        pool_recycle:连接回收（秒）
        pool_pre_ping:测试连接，执行select 1
        参考 https://docs.sqlalchemy.org/en/20/core/pooling.html#connection-pool-configuration
        """
        loader = DbConfigLoader(db_name, conf_path)
        # 数据库配置
        conf = loader.db_conf
        if conf.dialect not in self.DB_URL:
            raise DatabaseEngineError(msg='不支持的数据库类型')
        # urlquote 处理密码中的特殊字符
        url = self.DB_URL[conf.dialect].format(conf.user, urlquote(conf.passwd), conf.host, conf.port, conf.database)
        self.engine = create_engine(url, pool_size=conf.pool_size, pool_recycle=conf.pool_recycle, pool_pre_ping=True,
                                    echo=conf.show_sql)

    def create_session_factory(self):
        """
        创建session工厂
        :return:
        """
        return SqlsessionFactory(self.engine)