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

1 files changed, 10 insertions, 10 deletions

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