What is the pythonic way to use try/except/finally when having to release connections?

Viewed 65

I have a function that connects to a database and fetches some data and returns this data. This entire process is surrounded by a try except finally block. In the finally block, I release the resources if they exist even in case of an error. So what is the pythonic way to do the following:

def fetch_data():
    db_conn = None
    try:
        db_conn = DBConnection()
        data = db_conn.get_data()
    except Exception as ex:
        print(f"Exception occurred when fetching data: {ex}")
    finally:
        if db_conn:
            db_conn.close()

The initialization of db_conn = None and the finally block don't look elegant or pythonic enough and I was wondering if there was a better way of doing this?

1 Answers

You want to use a context manager. Context managers and the with statement were added to the language specifically to handle this pattern.

The code would become:

with DBConnection() as db_conn:
    data = db_conn.get_data()

See With Statement Context Managers. The implementation of DBConnection would need to supply __enter__ and __exit__ methods to handle this:

class DBConnection:
    def __init__(self):
        # .. initialisation of the instance. **Not** in the context manager

    def close(self):
        # ...

    def get_data(self):
        # ...
        return data_loaded

    def __enter__(self):
        # context is entered, `try:` 'opens' *after* this point.
        # Perhaps you want to actually connect to the database here
        # whatever is returned here is assignable via `with ... as name`
        # this can be a new object or self
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        # The context is exiting, equivalent of except ... and finally ...
        self.close()
        if exc_type:
            print(f"Exception occurred when fetching data: {exc_value}")
        # returning None will allow exceptions to propagate, returning
        # True (or other true value) will clear the exception.
        return True  # exception is cleared

I used return self in the __enter__ handler because that's a good pattern to use for context managers that fits your specific example well, but you could also return something else. Some database adapters return a transaction object or cursor at that point, for example.

Note that exceptions raised in __enter__ are not part of the context and are not handled there! If you need to handle exceptions in opening the database that you don't want to propagate, you'd have to postpone connecting until the first get_data() call.

Context managers encapsulate the try: ... except ...: finally: ... pattern. Also see the PEP 343 -- The *with* statement proposal that added this to Python:

This PEP adds a new statement "with" to the Python language to make it possible to factor out standard uses of try/finally statements.

Note that it is up to your use cases as to where you handle exceptions. For some context managers that's the __exit__ method, for others it is only useful to just clean up the context and not suppress exceptions. For example, files are context managers but they would not clear an exception. At that point you'd add a try: .. except SpecificException: to handle such cases:

try:
    with open(filename) as fobj:
        # ...
except IOError:
    # oops, file failed to open

or

try:
    open_file = open(filename)
except IOError:
    # oops, file failed to open
else:
    with open_file as fobj:
        # ...

The point of the context manager there is to ensure that the file object is closed if it had been opened, and nothing else.

Related