3 files changed, 19 insertions, 21 deletions
diff --git a/gn3/api/correlation.py b/gn3/api/correlation.py
index 14c029c..f2ac4d7 100644
--- a/gn3/api/correlation.py
+++ b/gn3/api/correlation.py
@@ -68,17 +68,15 @@ def compute_lit_corr(species=None, gene_id=None):
     might be needed for actual computing of the correlation results
     """
 
-    conn, _cursor_object = database_connector()
-    target_traits_gene_ids = request.get_json()
-    target_trait_gene_list = list(target_traits_gene_ids.items())
+    with database_connector() as conn:
+        target_traits_gene_ids = request.get_json()
+        target_trait_gene_list = list(target_traits_gene_ids.items())
 
-    lit_corr_results = compute_all_lit_correlation(
-        conn=conn, trait_lists=target_trait_gene_list,
-        species=species, gene_id=gene_id)
+        lit_corr_results = compute_all_lit_correlation(
+            conn=conn, trait_lists=target_trait_gene_list,
+            species=species, gene_id=gene_id)
 
-    conn.close()
-
-    return jsonify(lit_corr_results)
+        return jsonify(lit_corr_results)
 
 
 @correlation.route("/tissue_corr/<string:corr_method>", methods=["POST"])
diff --git a/gn3/api/heatmaps.py b/gn3/api/heatmaps.py
index b2511c3..80c8ca8 100644
--- a/gn3/api/heatmaps.py
+++ b/gn3/api/heatmaps.py
@@ -24,14 +24,14 @@ def clustered_heatmaps():
         return jsonify({
             "message": "You need to provide at least two trait names."
         }), 400
-    conn, _cursor = database_connector()
-    def parse_trait_fullname(trait):
-        name_parts = trait.split(":")
-        return f"{name_parts[1]}::{name_parts[0]}"
-    traits_fullnames = [parse_trait_fullname(trait) for trait in traits_names]
+    with database_connector() as conn:
+        def parse_trait_fullname(trait):
+            name_parts = trait.split(":")
+            return f"{name_parts[1]}::{name_parts[0]}"
+        traits_fullnames = [parse_trait_fullname(trait) for trait in traits_names]
 
-    with io.StringIO() as io_str:
-        figure = build_heatmap(traits_fullnames, conn, vertical=vertical)
-        figure.write_json(io_str)
-        fig_json = io_str.getvalue()
-    return fig_json, 200
+        with io.StringIO() as io_str:
+            figure = build_heatmap(traits_fullnames, conn, vertical=vertical)
+            figure.write_json(io_str)
+            fig_json = io_str.getvalue()
+        return fig_json, 200
diff --git a/gn3/db_utils.py b/gn3/db_utils.py
index 7263705..3703cbb 100644
--- a/gn3/db_utils.py
+++ b/gn3/db_utils.py
@@ -14,10 +14,10 @@ def parse_db_url() -> Tuple:
             parsed_db.password, parsed_db.path[1:])
 
 
-def database_connector() -> Tuple:
+def database_connector() -> mdb.Connection:
     """function to create db connector"""
     host, user, passwd, db_name = parse_db_url()
     conn = mdb.connect(host, user, passwd, db_name)
     cursor = conn.cursor()
 
-    return (conn, cursor)
+    return conn