如何在DB类的不同实例间高效共享数据库引擎？

问题描述 #

开发CLI应用的DB封装类时，遇到一个问题：每次为不同表实例化DB类时，即使是同一个数据库URL，也会创建多个独立的引擎实例。这既不合理也可能带来资源浪费，希望能实现同一数据库URL仅创建一个引擎，同时保持原有的类使用方式，还要支持连接多个不同数据库，适配测试场景（测试时用不同临时数据库）。

当前DB类实现：

from sqlmodel import Session, SQLModel, create_engine, select

class DB:
    def __init__(self, url: str, table: SQLModel, *, echo=False):
        """Database wrapper specific to the supplied database table.

        url: URL of the database file.
        table: Model of the table to which the operations should refer.
                Must be a Subclass of SQLModel.
        """
        self.url = url
        self.table = table
        self.engine = create_engine(url, echo=echo)

    def create_metadata(self):
        """Creates metadata, call only once per database connection."""
        SQLModel.metadata.create_all(self.engine)

    def read_all(self):
        """Returns all rows of the table."""
        with Session(self.engine) as session:
            entries = session.exec(select(self.table)).all()
            return entries

    def read(self, _id):
        """Returns a row of the table."""
        with Session(self.engine) as session:
            entry = session.get(self.table, _id)
            return entry

    def add(self, **fields):
        """Adds a row to the table. Fields must map to the table definition."""
        with Session(self.engine) as session:
            entry = self.table(**fields)
            session.add(entry)
            session.commit()

    def update(self, _id, **updates):
        """Updates a row of the table. Updates must map to the table definition."""
        with Session(self.engine) as session:
            entry = self.read(_id)
            for key, val in updates.items():
                setattr(entry, key, val)
            session.add(entry)
            session.commit()

    def delete(self, _id):
        """Delete a row of the table."""
        with Session(self.engine) as session:
            entry = self.read(_id)
            session.delete(entry)
            session.commit()

原使用方式：

from db import DB
from models import Project, Account

URL = "sqlite:///database.db"

projects = DB(url=URL, table=Project)
accounts = DB(url=URL, table=Account)

projects.read_all()
accounts.read(4)

此时projects和accounts各自持有独立的引擎实例，不符合预期。

现有思路的局限性 #

1. 类级单引擎方案：仅支持全局单一数据库URL，无法适配多数据库和测试场景（测试需要不同临时库）。

class DB:
    engine = None

    def __init__(self, url: str, table: SQLModel, *, echo=False):
        self.url = url
        self.table = table
        if DB.engine is None:
            DB.engine = create_engine(url, echo=echo)

    def create_metadata(self):
        SQLModel.metadata.create_all(self.engine)

2. 分离Engine类方案：解决了多数据库问题，但增加了使用复杂度，需要同时管理Engine和DB实例。

class Engine:
    def __init__(self, url, *, echo=False):
        self.url = url
        self.engine = create_engine(url, echo=echo)
    
    def create_metadata(self):
        """Creates metadata, call only once per database connection."""
        SQLModel.metadata.create_all(self.engine)

class DB:
    def __init__(self, table: SQLModel, engine):
        self.table = table
        self.engine = engine.engine

最优解决方案：引擎缓存机制 #

在DB类内部维护一个类级别的引擎缓存字典，以(url, echo)作为键（确保相同URL不同echo也能创建独立引擎），值为对应的引擎实例。这样既保持了原有的使用方式，又能实现同一URL复用引擎，同时支持多数据库和测试场景。

修改后的DB类实现：

from sqlmodel import Session, SQLModel, create_engine, select

class DB:
    # 类级缓存：键为(url, echo)，值为对应的引擎实例
    _engine_cache = {}

    def __init__(self, url: str, table: SQLModel, *, echo=False):
        """Database wrapper specific to the supplied database table.

        url: URL of the database file.
        table: Model of the table to which the operations should refer.
                Must be a Subclass of SQLModel.
        """
        self.url = url
        self.table = table
        self.echo = echo
        # 从缓存获取引擎，不存在则创建
        cache_key = (url, echo)
        if cache_key not in DB._engine_cache:
            DB._engine_cache[cache_key] = create_engine(url, echo=echo)
        self.engine = DB._engine_cache[cache_key]

    def create_metadata(self):
        """Creates metadata, call only once per database connection."""
        SQLModel.metadata.create_all(self.engine)

    # 以下方法保持原实现不变
    def read_all(self):
        with Session(self.engine) as session:
            entries = session.exec(select(self.table)).all()
            return entries

    def read(self, _id):
        with Session(self.engine) as session:
            entry = session.get(self.table, _id)
            return entry

    def add(self, **fields):
        with Session(self.engine) as session:
            entry = self.table(**fields)
            session.add(entry)
            session.commit()

    def update(self, _id, **updates):
        with Session(self.engine) as session:
            entry = self.read(_id)
            for key, val in updates.items():
                setattr(entry, key, val)
            session.add(entry)
            session.commit()

    def delete(self, _id):
        with Session(self.engine) as session:
            entry = self.read(_id)
            session.delete(entry)
            session.commit()

使用效果 #

原有的调用代码完全不需要修改：

from db import DB
from models import Project, Account

URL = "sqlite:///database.db"

projects = DB(url=URL, table=Project)
accounts = DB(url=URL, table=Account)

# projects和accounts使用同一个引擎实例
assert projects.engine is accounts.engine

同时支持多数据库场景：

db1_url = "sqlite:///db1.db"
db2_url = "sqlite:///db2.db"

project_db = DB(url=db1_url, table=Project)
account_db = DB(url=db2_url, table=Account)

# 不同URL对应不同引擎
assert project_db.engine is not account_db.engine

测试场景中，每个测试用不同临时URL，也会创建独立引擎，互不干扰。

内容的提问来源于stack exchange，提问作者Jan