Create database connections within context managers

Use the `with` context manager to open database connections, so as to ensure
that those connections are closed once the call is completed. This hopefully
avoids the 'too many connections' error

2 files changed, 17 insertions, 19 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