Context manager for Python's MySQLdb

问题:

I am used to (spoiled by?) python's SQLite interface to deal with SQL databases. One nice feature in python's SQLite's API the "context manager," i.e., python's with statement. I usually execute queries in the following way:

import as sqlite

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

With the code above, if my query modifies the database and I forget to run conn.commit(),the context manager runs it for me automatically upon exiting the with statement. It also handles exceptions nicely: if an exception occurs before I commit anything, then the database is rolled back.

I am now using the MySQLdb interface, which doesn't seem to support a similar context manager out of the box. How do I create my own? There is a related question here, but it doesn't offer a complete solution.

回答1:

Previously, MySQLdb connections were context managers. As of this commit on 2018-12-04, however, MySQLdb connections are no longer context managers, and users must explicitly call conn.commit() or conn.rollback(), or write their own context manager, such as the one below.

---

You could use something like this:

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


class Cursor(object):
    def __init__(self,
                 cursorclass=Cursor,
                 host=config.HOST, user=config.USER,
                 passwd=config.PASS, dbname=config.MYDB,
                 driver=MySQLdb,
                 ):
        self.cursorclass = cursorclass
        self.host = host
        self.user = user
        self.passwd = passwd
        self.dbname = dbname
        self.driver = driver
        self.connection = self.driver.connect(
            host=host, user=user, passwd=passwd, db=dbname,
            cursorclass=cursorclass)
        self.cursor = self.connection.cursor()

    def __iter__(self):
        for item in self.cursor:
            yield item

    def __enter__(self):
        return self.cursor

    def __exit__(self, ext_type, exc_value, traceback):
        self.cursor.close()
        if isinstance(exc_value, Exception):
            self.connection.rollback()
        else:
            self.connection.commit()
        self.connection.close()

with Cursor() as cursor:
    print(cursor)
    connection = (cursor.connection)
    print(connection)

To use it you would place config.py in your PYTHONPATH and define the HOST, USER, PASS, MYDB variables there.

回答2:

Think things have changed since this question was originally asked. Somewhat confusingly (from my point of view at least), for recent versions of MySQLdb, if you use a connection in a context you get a cursor (as per the oursql example), not something that closes automatically (as you would if you opened a file for instance).

Here's what I do:

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')