Python上下文管理器处理数据库连接的最佳实践

在Python中操作数据库时，确保连接在使用完毕后被正确关闭是避免资源泄漏的关键。手动调用 close() 方法容易因异常而遗漏，而上下文管理器（Context Manager）能自动处理资源的获取与释放。使用 with 语句配合上下文管理器是管理数据库连接最安全、最简洁的方式。

1. 理解上下文管理器的核心机制

上下文管理器是一个实现了 __enter__ 和 __exit__ 方法的对象。当你将它用于 with 语句时：

• 进入 with 块时，自动调用 __enter__ 方法，通常返回资源（如数据库连接）。
• 退出 with 块时（无论是否发生异常），自动调用 __exit__ 方法，执行清理逻辑（如关闭连接）。

编写自定义上下文管理器时，必须确保 __exit__ 能处理所有可能的异常路径。

2. 使用 contextlib.contextmanager 快速创建轻量级管理器

对于简单场景，无需定义完整类，可借助 contextlib 模块的 @contextmanager 装饰器。

安装并导入所需模块：

pip install psycopg2-binary  # 以PostgreSQL为例

定义一个生成数据库连接的上下文管理器函数：

from contextlib import contextmanager
import psycopg2

@contextmanager
def get_db_connection():
    conn = None
    try:
        conn = psycopg2.connect(
            host="localhost",
            database="mydb",
            user="user",
            password="password"
        )
        yield conn
    except Exception as e:
        if conn:
            conn.rollback()
        raise e
    finally:
        if conn and not conn.closed:
            conn.close()

使用该管理器执行查询：

with get_db_connection() as conn:
    with conn.cursor() as cur:
        cur.execute("SELECT * FROM users WHERE id = %s", (1,))
        result = cur.fetchone()
        print(result)

注意：psycopg2 的 Connection 对象本身已支持上下文管理协议，但其默认行为不自动提交事务，需手动处理。上述封装显式加入了 rollback() 逻辑以增强安全性。

3. 为常用数据库驱动封装健壮的连接管理器

不同数据库驱动对上下文管理的支持程度不同。以下是针对主流库的最佳实践封装。

SQLite（内置支持，但仍建议封装）

SQLite 的 sqlite3.Connection 支持 with，但仅管理事务（自动提交或回滚），不自动关闭连接。因此仍需手动关闭，或自行封装：

import sqlite3
from contextlib import contextmanager

@contextmanager
def get_sqlite_connection(db_path: str):
    conn = sqlite3.connect(db_path)
    try:
        yield conn
    except Exception:
        conn.rollback()
        raise
    else:
        conn.commit()
    finally:
        conn.close()

MySQL（PyMySQL 或 mysql-connector-python）

以 PyMySQL 为例，其连接对象不原生支持上下文管理，必须封装：

import pymysql
from contextlib import contextmanager

@contextmanager
def get_mysql_connection():
    conn = pymysql.connect(
        host='localhost',
        user='user',
        password='password',
        database='mydb',
        charset='utf8mb4'
    )
    try:
        yield conn
    except Exception:
        conn.rollback()
        raise
    else:
        conn.commit()
    finally:
        conn.close()

4. 避免常见陷阱

错误：依赖驱动自带的上下文管理而不验证行为

某些驱动（如 psycopg2）的 with conn: 仅控制事务，不会关闭连接。若误以为会自动关闭，将导致连接池耗尽。

验证方法：查阅官方文档中关于 __enter__ / __exit__ 的说明。

错误：在 __exit__ 中吞掉异常

不要在 __exit__ 中无条件返回 True，否则会隐藏错误：

# 错误示例
def __exit__(self, exc_type, exc_val, exc_tb):
    self.conn.close()
    return True  # ← 这会抑制所有异常！

正确做法：让 __exit__ 默认返回 None（即 False），使异常正常传播。

错误：未处理事务边界

在 __exit__ 中，若发生异常应执行 rollback()；若无异常则应 commit()。遗漏任一操作都可能导致数据不一致或锁等待。

5. 推荐：使用现成的高级库简化管理

对于生产环境，建议使用 SQLAlchemy 或 Django ORM，它们内置了更完善的连接和事务管理。

但若需直接操作原生连接，可采用以下通用模板：

from contextlib import contextmanager
import your_db_driver

@contextmanager
def managed_db_connection(**kwargs):
    conn = your_db_driver.connect(**kwargs)
    try:
        yield conn
    except Exception:
        conn.rollback()
        raise
    else:
        conn.commit()
    finally:
        conn.close()

调用时传入连接参数即可复用。

6. 性能与连接池考量

频繁创建/关闭连接开销大。在高并发场景中，应结合连接池使用上下文管理器。

例如，使用 psycopg2.pool：

from psycopg2 import pool
from contextlib import contextmanager

connection_pool = psycopg2.pool.SimpleConnectionPool(
    1, 20,
    host="localhost",
    database="mydb",
    user="user",
    password="password"
)

@contextmanager
def get_pooled_connection():
    conn = connection_pool.getconn()
    try:
        yield conn
    except Exception:
        conn.rollback()
        raise
    else:
        conn.commit()
    finally:
        connection_pool.putconn(conn)

注意：连接池的 getconn() 和 putconn() 必须成对出现，上下文管理器能完美保证这一点。

7. 最终检查清单

在部署前，请确认你的数据库上下文管理器满足以下条件：

编写测试用例验证异常路径：

def test_connection_rollback_on_error():
    with pytest.raises(ValueError):
        with get_db_connection() as conn:
            with conn.cursor() as cur:
                cur.execute("INSERT INTO logs (msg) VALUES ('test')")
                raise ValueError("模拟失败")
    # 验证 'test' 未写入数据库

始终通过 with 语句获取数据库连接。