Code - Python 上下文管理

发表于   |   分类于 技术 , Python   |   暂无评论

核心知识点

- 上下文管理器：实现__enter__和__exit__方法的对象，用于自动管理资源（打开→使用→关闭）

- 语法：with 上下文管理器 as 变量: （代码块执行完自动调用__exit__关闭资源）

- 企业常用场景：文件操作、数据库连接、网络请求（避免资源泄露）

- 简化实现：用@contextmanager装饰器快速创建上下文管理器（无需手动实现两个魔法方法）

实战案例：MySQL连接上下文管理器（AI接口数据持久化场景）

案例需求

开发一个MySQL连接的上下文管理器，实现：自动连接数据库→执行SQL→自动关闭连接（即使执行过程中报错也能正常关闭），适配济南企业常用的MySQL 8.0版本，用于AI接口的日志数据持久化（如保存用户调用记录）。

实现步骤

1. 安装依赖：pymysql（Python连接MySQL的库）
2. 准备MySQL环境：创建数据库（ai_dev）和表（api_call_log）
3. 用@contextmanager装饰器编写上下文管理器，封装连接逻辑
4. 测试：用with语句调用上下文管理器，执行插入日志的SQL，验证连接是否自动关闭

完整代码


    import pymysql
    from contextlib import contextmanager
    from datetime import datetime
    
    
    # ------------------- 1. 先执行SQL创建数据库和表（手动在MySQL客户端执行） -------------------
    # CREATE DATABASE IF NOT EXISTS ai_dev DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
    # USE ai_dev;
    # CREATE TABLE IF NOT EXISTS api_call_log (
    #     id INT AUTO_INCREMENT PRIMARY KEY,
    #     func_name VARCHAR(50) NOT NULL COMMENT '接口函数名',
    #     call_time DATETIME NOT NULL COMMENT '调用时间',
    #     user_id VARCHAR(32) NOT NULL COMMENT '调用用户ID',
    #     status TINYINT NOT NULL COMMENT '状态：1-成功，0-失败',
    #     created_at DATETIME DEFAULT CURRENT_TIMESTAMP
    # ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='AI接口调用日志表';
    @contextmanager
    def mysql_connection(host: str, port: int, user: str, password: str, db: str):
        """创建MySQL数据库连接上下文管理器"""
        try:
            conn = pymysql.connect(
                host=host,
                port=port,
                user=user,
                password=password,
                database=db,
                charset="utf8mb4",
            )
            print("数据库连接成功")
            yield conn
        except pymysql.MySQLError as e:
            print(f"数据库连接失败: {e}")
            raise
        finally:
            if conn is not None:  # 安全关闭
                conn.close()
                print("数据库连接已关闭")
    
    
    # ------------------- 2. 插入ai接口调用日志  -------------------
    def insert_api_call_log(
        conn: pymysql.Connection,
        func_name: str,
        user_id: str,
        status: int,
        call_time: datetime = None,
    ):
        if call_time is None:
            call_time = datetime.now()
        """插入AI接口调用日志"""
        sql = "INSERT INTO api_call_log (func_name, call_time, user_id, status) VALUES (%s, %s, %s, %s)"
        with conn.cursor() as cursor:
            cursor.execute(sql, (func_name, call_time, user_id, status))
            conn.commit()
            print("AI接口调用日志插入成功")
    
    
    if __name__ == "__main__":
        # 测试数据库连接和插入日志
        mysql_config = {
            "host": "localhost",
            "port": 3306,
            "user": "root",
            "password": "root",
            "db": "ai_dev",
        }
    
        with mysql_connection(**mysql_config) as conn:
            # 插入一条成功的调用日志（模拟文本摘要接口调用）
            insert_api_call_log(conn, "text_summary", "user_1001", 1)
            # 插入一条失败的调用日志（模拟参数错误）
            insert_api_call_log(conn, "text_summary", "user_1002", 0)
    
    print("日志插入完成，请查看数据库验证")

总结

该上下文管理器等价于PHP中用try-catch-finally封装PDO连接（try中连接→执行SQL，finally中关闭连接），但Python的with语句更简洁，无需手动编写finally块，通过上下文管理器自动完成资源的“打开-关闭”生命周期管理。
添加新评论