Search code examples
pythondatabasecontextmanager

database context manager and query execution exception handling


I would like to connect to a database, execute a query inside a context manager and handle the exceptions

import pyodbc

class db_connection(object):

    def __init__(self, connection_string):
        self.connection_string = connection_string
        self.connector = None

    def __enter__(self):
        print("connecting to DB ...")
        try:
            self.connector = pyodbc.connect(self.connection_string)
        except Exception:
            print("connection failed")
        else:
            print("connected")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb): 
        pass 

    def runQuery(self, query):
        result_list = None
        if self.connector == None:
            pass
        else:
            try:
                cursor = self.connector.cursor()
                result_list = cursor.execute(query).fetchall() 
            except Exception:
                print("query failed")
            finally:
                return result_list

To call this class like so:

connect_string = "DRIVER={ODBC SQL};SERVER=LOCALSERVER;DATABASE=LOCALDB; Trusted_Connection=yes"
query = "SELECT ALL FROM LOCALDB"

with db_connection(connect_string) as connection:
    connection.runQuery(query)

With two questions (besides the fact I am not to sure this is a decent pythonic solution):

1 - I don't seem to need the _ exit _ statement at all

2 - The runQuery() method needs to be called separately, is it possible to also include this in the _ enter _ statement


Solution

  • The pythonic way would be to close in the _ exit _ statement the resource you opened in the _ enter _ statement (to prevent memory leaks etc). But connections are meant to be reused so we don't need to close them everytime. If you want to reuse connections try researching how connection pooling works in pyodbc and if you don't then close the connection in the _ exit _ method. Anyway you still need the _ exit _ method to fulfill the context manager protocol. In your case the _ exit _ statement manages exceptions.

    I don't think that including the runQuery() method in in the _ enter _ statement is a good idea.

    # !BAD example
    with db_connection_run_query(connect_string, query) as query_result:
        ...
    

    Query_result is just a list of Rows - you don't need a context manager to work with it.

    BTW your runQuery method doesn't close the cursor. This is also a potential memory leak. You need to close it in the finally clause before returning the result. (context manager might also help here)

    try:
        cursor = self.connector.cursor()
        result_list = cursor.execute(query).fetchall() 
    except Exception:
        print("query failed")
    finally:
        cursor.close()
        return result_list
    

    One more thing. Your runQuery method will only work with select queries and not with inserts, updates and deletes (To perform those you need to commit a transaction by calling commit() on cursor). So you may consider renaming runQuery to runSelectQuery.