diff options
Diffstat (limited to 'tests/performance/perf_query.py')
-rw-r--r-- | tests/performance/perf_query.py | 14 |
1 files changed, 6 insertions, 8 deletions
diff --git a/tests/performance/perf_query.py b/tests/performance/perf_query.py index c22dcf5..e534e9b 100644 --- a/tests/performance/perf_query.py +++ b/tests/performance/perf_query.py @@ -28,15 +28,13 @@ def timer(func): def query_executor(query: str, fetch_all: bool = True): """function to execute a query""" - conn, _ = database_connector() + with database_connector() as conn: + with conn.cursor() as cursor: + cursor.execute(query) - with conn: - cursor = conn.cursor() - cursor.execute(query) - - if fetch_all: - return cursor.fetchall() - return cursor.fetchone() + if fetch_all: + return cursor.fetchall() + return cursor.fetchone() def fetch_probeset_query(dataset_name: str): |