Pavel Chernikov
Pavel Chernikov

Reputation: 2266

Can I use pymysql.connect() with "with" statement?

The following is listed as example in pymysql:

conn = pymysql.connect(...)
with conn.cursor() as cursor:
    cursor.execute(...)
    ...
conn.close()

Can I use the following instead, or will this leave a lingering connection? (it executes successfully)

import pymysql
with pymysql.connect(...) as cursor:
    cursor.execute('show tables')

(python 3, latest pymysql)

Upvotes: 17

Views: 10103

Answers (5)

Eugene Yarmash
Eugene Yarmash

Reputation: 150225

As of PyMySQL v.1.1.1, the connection object is a context manager and therefore can be used in a with statement:

with pymysql.connect(**connect_args) as conn:
    with conn.cursor() as cursor:
        # do something with the cursor

On leaving the with code block the connection will be closed, so don't do this if you want to use the same connection in multiple contexts. This behavior differs from how the connection context manager works in other database access libraries (e.g. psycopg2 or sqlite), where exiting the with block just commits or rolls back the transaction. As it stands, the Python DBAPI doesn't require the connection class to implement __enter__ and __exit__ methods at all, which doesn't promote consistency in implementations.

Upvotes: 1

sberfiacreix
sberfiacreix

Reputation: 61

A recent update to Pymysql (https://github.com/PyMySQL/PyMySQL/pull/886/files) now calls close() on exit, so using with is supported and reflected in their documentation.

import pymysql.cursors

# Connect to the database
connection = pymysql.connect(host='localhost',
                             user='user',
                             password='passwd',
                             database='db',
                             cursorclass=pymysql.cursors.DictCursor)

with connection:
.
.

Upvotes: 2

bitsplit
bitsplit

Reputation: 1060

As an alternative to this, since I wanted to support the context manager pattern for a connection, I implemented it with a monkey patch. Not the best approach, but it's something.

import pymysql


MONKEYPATCH_PYMYSQL_CONNECTION = True


def monkeypatch_pymysql_connection():
    Connection = pymysql.connections.Connection

    def enter_patch(self):
        return self

    def exit_patch(self, exc, value, traceback):
        try:
            self.rollback()  # Implicit rollback when connection closed per PEP-249
        finally:
            self.close()

    Connection.__enter__ = enter_patch
    Connection.__exit__ = exit_patch


if MONKEYPATCH_PYMYSQL_CONNECTION:
    monkeypatch_pymysql_connection()
    MONKEYPATCH_PYMYSQL_CONNECTION = False  # Prevent patching more than once

This approach worked for my use case. I would prefer to have __enter__ and __exit__ methods in the Connection class. That approach, however, was rejected by the developers when they addressed the issue in late 2018.

Upvotes: 1

Maxim Yanchenko
Maxim Yanchenko

Reputation: 156

As it was pointed out, the Cursor takes care of itself, but all the Connection's support for context manager was removed completely just a few days ago, so the only option now is to write yours:

https://github.com/PyMySQL/PyMySQL/pull/763

https://github.com/PyMySQL/PyMySQL/issues/446

Upvotes: 8

Mike
Mike

Reputation: 7203

This does not look safe, if you look here, the __enter__ and __exit__ functions are what are called in a with clause. For the pymysql connection they look like this:

def __enter__(self):
    """Context manager that returns a Cursor"""
    return self.cursor()

def __exit__(self, exc, value, traceback):
    """On successful exit, commit. On exception, rollback"""
    if exc:
        self.rollback()
    else:
        self.commit()

So it doesn't look like the exit clause closes the connection, which means it would be lingering. I'm not sure why they did it this way. You could make your own wrappers that do this though.

You could recycle a connection by creating multiple cursors with it (the source for cursors is here) the cursor methods look like this:

def __enter__(self):
    return self

def __exit__(self, *exc_info):
    del exc_info
    self.close()

So they do close themselves. You could create a single connection and reuse it with multiple cursors in with clauses.

If you want to hide the logic of closing connections behind a with clause, e.g. a context manager, a simple way to do it would be like this:

from contextlib import contextmanager
import pymysql


@contextmanager
def get_connection(*args, **kwargs):
    connection = pymysql.connect(*args, **kwargs)
    try:
        yield connection
    finally:
        connection.close()

You could then use that context manager like this:

with get_connection(...) as con:
    with con.cursor() as cursor:
        cursor.execute(...)

Upvotes: 20

Related Questions