python - Python - Manage cursor connection outside context manager

Question

I'm new at Python and I'm trying to build a start project to get into this language.

I created a SQLite3 DB and managed to make transactions with it. Everything works fine.

I wanted to get deeper in Python so I've been searching and discovered Decorators and Context Manager and I was trying to implement these concepts on my Query Execution's functions. However, I'm stuck with a problem.

I've created a class that handles the open and close connection tasks.

DB_ContextManager.py class:

class DB_ContextManager():
    def __init__(self, db_connection):
        self.db_connection = db_connection

    def __enter__(self):
        self.conn = sqlite3.connect(self.db_connection)        
        return self.conn

    def __exit__(self, exc_type, exc_val, exc_tb): # obligatory params
        self.conn.close()

And also created ConnectionDB.py which is responsible for executing queries.

from Database.DB_ContextManager import DB_ContextManager as DB_CM

# Handles SELECT queries
def ExecuteSelectQuery(self, pQuery):
    try:
        with DB_CM(db_connection_string) as conn:
            cur = conn.cursor()
            cur.execute(pQuery)
            result = cur.fetchall()        
            return result

    except Exception as e:
        LH.Handler(log_folder, 'ConnectionDB', 'Queries', 'ExecuteSelectQuery', e)            
        raise DE.ConnectionDB_Exception()

# Handles INSERTs, UPDATEs, DELETEs queries
def ExecuteNonQuery(self, pQuery):
    try:            
        with DB_CM(db_connection_string) as conn:
            cur = conn.cursor()
            cur.execute(pQuery)

    except Exception as e:
        LH.Handler(log_folder, 'ConnectionDB', 'Queries', 'ExecuteSelectNonQuery', e)            
        raise DE.ConnectionDB_Exception()

As you can see

with DB_CM(db_connection_string) as conn:
            cur = conn.cursor()
            cur.execute(pQuery)

is repeated in each function

To avoid this situation, I'd like to create a Decorator function that encapsulates this piece of code. My problem is that the cursor 'dies' inside the ContextManager and, for example, ExecuteSelectQuery needs the cursor to fetch the return data after the query was executed.

I know it's a small project and thinking so long term in future may not be necessary. But, remember, it's a start project and I'm learning to apply new concepts.

---

SOLUTION

As @blhsing suggested, I return the connection object instead of the cursor in the ContextManager.

Also I handles commit() and rollback() in it.

So, summarizing:

ConnectionDB.py

def ExecuteSelectQuery(self, pQuery):
    with DB_CM(db_connection_string, pQuery) as cur:
        result = cur.fetchall()
        return result

def ExecuteSelectNonQuery(self, pQuery):
    with DB_CM(db_connection_string, pQuery) as cur: 
        pass

and ConnectionDB.py

class DB_ContextManager():
    def __init__(self, db_connection, pQuery):
        self.db_connection = db_connection
        self.query = pQuery

    def __enter__(self):
        try:
            self.conn = sqlite3.connect(self.db_connection)
            cur = self.conn.cursor()
            cur.execute(self.query)
            self.conn.commit()
            return cur            

        except Exception as e:
            LH.Handler(log_folder, 'DB_ContextManager', 'DB_ContextManager', '__enter__', e)
            self.conn.rollback()
            raise DE.ConnectionDB_Exception()

    def __exit__(self, exc_type, exc_val, exc_tb): # obligatory params
        self.conn.close()

标签： pythonsqliteconnectiondecorator

Answer 1

You can make the context manager return the cursor instead of the connection object:

class DB_CM():
    def __init__(self, db_connection):
        self.db_connection = db_connection

    def __enter__(self):
        self.conn = sqlite3.connect(self.db_connection)
        cur = self.conn.cursor()
        cur.execute(pQuery)
        return cur

    def __exit__(self, exc_type, exc_val, exc_tb): # obligatory params
        self.conn.close()

so that the try block of ExecuteSelectQuery can be revised as:

with DB_CM(db_connection_string) as cur:
    result = cur.fetchall()        
    return result

and the try block of ExecuteNonQuery can be simply:

with DB_CM(db_connection_string):
    pass