PooledDB connection does not work with `with` statement

[hzhu212]

How to reproduce:

from dbutils.pooled_db import PooledDB

pool = PooledDB(some_creator)

conn = pool.connection()
with conn:
    print(1)

Will throw an AttributeError in with statement:

AttributeError                            Traceback (most recent call last)
<ipython-input-33-1657fbc224a7> in <module>
----> 1 with conn:
      2     print(1)

AttributeError: __enter__

What's wrong:

It seems conn (PooledDedicatedDBConnection object) has no attribute __enter__.

But when I type

conn.__enter__

it gives

<bound method SteadyDBConnection.__enter__ of <dbutils.steady_db.SteadyDBConnection object at 0x7fafbc4dbe10>>

So conn.__enter__ do exist, but it comes from internal SteadyDBConnection object, not conn itself.

After digging into source code, I find the __enter__ attribute is proxied from conn._con object with __getattr__:

def __getattr__(self, name):
        """Proxy all members of the class."""
        if self._con:
            return getattr(self._con, name)
        else:
            raise InvalidConnection

So the reason of the AttributeError is that __getattr__ won't take effect in with statement. I Googled this problem and find this: Why doesn't getattr work with exit?. It seems like a designed feature that we can't proxy __enter__ and __exit__ with __getattr__ in with statement.

Furthermore, even though __getattr__ works fine with __enter__ and __exit__, PooledDedicatedDBConnection shouldn't proxy them from SteadyDBConnection, because they act the wrong way.

The right effect of PooledDedicatedDBConnection.__exit__ should be return current connection to pool, like its close method. But what SteadyDBConnection's __enter__ and __exit__ methods do are very different.

A simple way to fix

A simple way to enable with statement on PooledDedicatedDBConnection object is putting __enter__ and __exit__ method into its definition, like this:

class PooledDedicatedDBConnection:
    ...

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()

Or, we can use contextlib.closing as an alternative solution, like this:

from contextlib import closing

with closing(conn) as conn:
    print(1)

This works fine, but not as convenient as other DB API 2 connections, because of the import thing.

[Cito]

Hi @hzhu212 - thanks a lot for your suggestion. I have now implemented this in the main branch. If it looks good to you, I can create a new release with this change.

[BrunoMarengo]

Hello @Cito ! hope you are fine! I was reading the source code from the main branch and tried to use the connection in a context manager. After some really confused faces I figured it out that the branch wasn't released yet! Are you planning to release it soon?

[Cito]

@BrunoMarengo I have released a new version 2.0.2 with this change now.