about summary refs log tree commit diff
path: root/tests/performance
diff options
context:
space:
mode:
Diffstat (limited to 'tests/performance')
-rw-r--r--tests/performance/perf_query.py14
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):