Python

Control-flow structure for database connections

Control-flow structure for database connections

cover photo: Emma Döbken 2018.

With Python, creating a database connection is straightforward. Yet, I often see the following case go wrong, while a simple solution is easily at hand by using the context manager pattern.

For database connections, you’ll need at least one secret. Let’s say you get this secret from a secret manager by running the get_secret() method. You also use an utility like JayDeBeApi to setup the connection and you are smart enough to close the connection after querying and deleting the password:

my_secret = get_secret(username)
conn = jaydebeapi.connect(
    jclassname='org.netezza.Driver',
    url=url,
    driver_args=[username, my_secret])

# query all tables owned by JOOST
sql = "SELECT * FROM _v_table WHERE OWNER like '%JOOST%'"
df = pd.read_sql(sql, con=conn)

# close the connection
conn.close()
# forget the secret
del my_secret

This only works if everything succeeds; and your code is only as good as how it behaves when things go wrong. So you are going to catch the error with a try-except pattern, something like this:

try:
    my_secret = get_secret(username)
    conn = jaydebeapi.connect(
        jclassname='org.netezza.Driver',
        url=url,
        driver_args=[username, my_secret])

    # query all tables owned by JOOST
    sql = "SELECT * FROM _v_table WHERE OWNER like '%JOOST%'"
    df = pd.read_sql(sql, con=conn)
except Exception as e:
    # close the connection
    conn.close()
    # forget the secret
    del my_secret
    raise e
# close the connection
conn.close()
# forget the secret
del my_secret

Well… the except statement should be replaced with a finally statement to be run no-matter-what, making the last lines redundant; but the above is what I have seen more than once.

Query within a Context

The with statement provides a control-flow structure for objects where the __enter__() method is called before the with-block end the __exit__() method after the with-block.

class NetezzaConnection:
    """Provide a Netezza connection context."""
    
    def __init__(self, url, username):
        self.url = url
        self.username
        self.jclassname = 'org.netezza.Driver'
    
    def __enter__(self):
        _secret = get_secret(self.username)
        self.connection = jaydebeapi.connect(
            jclassname=self.jclassname,
            url=self.url,
            driver_args=[self.username, _secret])
        return self.connection
    
    def __exit__(self, type, value, traceback):
        self.connection.close()

This provides a context in which the connection and secrets exist. Now, keep the query in this context and the connection will be closed, whatever happens:

with NetezzaConnection(url, username) as conn:
    # query all tables owned by JOOST
    sql = "SELECT * FROM _v_table WHERE OWNER like '%JOOST%'"
    df = pd.read_sql(sql, con=conn)

The complexity has moved to the object and the query itself is easier to read and maintain.