Python MySQLdb的上下文管理器

Question

我习惯使用 Python 的 SQLite 接口来处理 SQL 数据库。Python 的 SQLite API 有一个很不错的功能，就是“上下文管理器”，也就是 Python 的 with 语句。我通常是这样执行查询的：

import as sqlite

with sqlite.connect(db_filename) as conn:
    query = "INSERT OR IGNORE INTO shapes VALUES (?,?);"
    results = conn.execute(query, ("ID1","triangle"))

使用上面的代码，如果我的查询修改了数据库，而我忘记运行 conn.commit()，那么上下文管理器会在我退出 with 语句时自动为我执行这个操作。它还很好地处理异常：如果在我提交任何内容之前发生了异常，那么数据库会被回滚到之前的状态。

现在我在使用 MySQLdb 接口，但它似乎没有直接支持类似的上下文管理器。我该如何创建自己的上下文管理器呢？这里有一个相关的问题 在这里，但它没有提供完整的解决方案。

Answer 1

我觉得自从这个问题最初被提出来后，情况有了一些变化。从我的角度来看，这有点让人困惑。对于最近版本的 MySQLdb 来说，如果你在一个特定的环境中使用连接，你会得到一个游标（就像 oursql 的例子那样），而不是一些会自动关闭的东西（就像你打开一个文件时那样）。

这是我通常的做法：

from contextlib import closing
with closing(getConnection()) as conn: #ensure that the connection is closed
    with conn as cursor:               #cursor will now auto-commit
        cursor.execute('SELECT * FROM tablename')

Answer 2

之前，MySQLdb 的连接是上下文管理器，也就是说你可以用它来自动管理数据库连接的开启和关闭。不过，从2018年12月4日的这个更新开始，MySQLdb 的连接就不再是上下文管理器了。现在，用户需要手动调用 conn.commit() 或 conn.rollback()，或者自己写一个上下文管理器，比如下面这个。

---

你可以使用类似这样的代码：

import config
import MySQLdb
import MySQLdb.cursors as mc
import _mysql_exceptions
import contextlib
DictCursor = mc.DictCursor
SSCursor = mc.SSCursor
SSDictCursor = mc.SSDictCursor
Cursor = mc.Cursor

@contextlib.contextmanager
def connection(cursorclass=Cursor,
               host=config.HOST, user=config.USER,
               passwd=config.PASS, dbname=config.MYDB,
               driver=MySQLdb):
    connection = driver.connect(
            host=host, user=user, passwd=passwd, db=dbname,
            cursorclass=cursorclass)
    try:
        yield connection
    except Exception:
        connection.rollback()
        raise
    else:
        connection.commit()
    finally:
        connection.close()

@contextlib.contextmanager
def cursor(cursorclass=Cursor, host=config.HOST, user=config.USER,
           passwd=config.PASS, dbname=config.MYDB):
    with connection(cursorclass, host, user, passwd, dbname) as conn:
        cursor = conn.cursor()
        try:
            yield cursor
        finally:
            cursor.close()


with cursor(SSDictCursor) as cur:
    print(cur)
    connection = cur.connection
    print(connection)
    sql = 'select * from table'
    cur.execute(sql)
    for row in cur:
        print(row)

要使用它，你需要把 config.py 放在你的 PYTHONPATH 中，并在里面定义 HOST、USER、PASS 和 MYDB 这些变量。