Use context manager with database connection

Use the `with` context manager with database connections and cursors
to ensure that they are closed once they are no longer needed.

Where it was not feasible to use the `with` context manager without a
huge refactor/rewrite, the cursors and connections are closed manually.

4 files changed, 53 insertions, 71 deletions

diff --git a/wqflask/maintenance/gen_select_dataset.py b/wqflask/maintenance/gen_select_dataset.py
index db65a11f..9f4b670d 100644
--- a/wqflask/maintenance/gen_select_dataset.py
+++ b/wqflask/maintenance/gen_select_dataset.py
@@ -39,21 +39,13 @@ from wqflask import app
 
 from utility.tools import locate, locate_ignore_error, TEMPDIR, SQL_URI
 
-import MySQLdb
-
 import simplejson as json
 import urllib.parse
 
 
-#import sqlalchemy as sa
-
 from pprint import pformat as pf
 
-#Engine = sa.create_engine(zach_settings.SQL_URI)
-
-# build MySql database connection
-
-#conn = Engine.connect()
+from wqflask.database import database_connection
 
 
 def parse_db_uri():
@@ -71,19 +63,19 @@ def parse_db_uri():
     return db_conn_info
 
 
-def get_species():
+def get_species(cursor):
     """Build species list"""
-    #Cursor.execute("select Name, MenuName from Species where Species.Name != 'macaque monkey' order by OrderId")
-    Cursor.execute("select Name, MenuName from Species order by OrderId")
-    species = list(Cursor.fetchall())
+    #cursor.execute("select Name, MenuName from Species where Species.Name != 'macaque monkey' order by OrderId")
+    cursor.execute("select Name, MenuName from Species order by OrderId")
+    species = list(cursor.fetchall())
     return species
 
 
-def get_groups(species):
+def get_groups(cursor, species):
     """Build groups list"""
     groups = {}
     for species_name, _species_full_name in species:
-        Cursor.execute("""select InbredSet.Name, InbredSet.FullName from InbredSet,
+        cursor.execute("""select InbredSet.Name, InbredSet.FullName from InbredSet,
                        Species,
                        ProbeFreeze, GenoFreeze, PublishFreeze where Species.Name = '%s'
                        and InbredSet.SpeciesId = Species.Id and
@@ -92,7 +84,7 @@ def get_groups(species):
                         or ProbeFreeze.InbredSetId = InbredSet.Id)
                         group by InbredSet.Name
                         order by InbredSet.FullName""" % species_name)
-        results = Cursor.fetchall()
+        results = cursor.fetchall()
         groups[species_name] = list(results)
     return groups
 
@@ -273,13 +265,13 @@ def build_datasets(species, group, type_name):
     return datasets
 
 
-def main():
+def main(cursor):
     """Generates and outputs (as json file) the data for the main dropdown menus on the home page"""
 
     parse_db_uri()
 
-    species = get_species()
-    groups = get_groups(species)
+    species = get_species(cursor)
+    groups = get_groups(cursor, species)
     types = get_types(groups)
     datasets = get_datasets(types)
 
@@ -316,6 +308,6 @@ def _test_it():
 
 
 if __name__ == '__main__':
-    Conn = MySQLdb.Connect(**parse_db_uri())
-    Cursor = Conn.cursor()
-    main()
+    with database_connection() as conn:
+        with conn.cursor() as cursor:
+            main(cursor)
diff --git a/wqflask/maintenance/generate_probesetfreeze_file.py b/wqflask/maintenance/generate_probesetfreeze_file.py
index e964c8ed..f43f952b 100644
--- a/wqflask/maintenance/generate_probesetfreeze_file.py
+++ b/wqflask/maintenance/generate_probesetfreeze_file.py
@@ -8,20 +8,11 @@ import os
 import collections
 import csv
 
-import MySQLdb
-
 from base import webqtlConfig
 
 from pprint import pformat as pf
 
-
-def get_cursor():
-    con = MySQLdb.Connect(db=webqtlConfig.DB_UPDNAME,
-                          host=webqtlConfig.MYSQL_UPDSERVER,
-                          user=webqtlConfig.DB_UPDUSER,
-                          passwd=webqtlConfig.DB_UPDPASSWD)
-    cursor = con.cursor()
-    return cursor
+from wqflask.database import database_connection
 
 
 def show_progress(process, counter):
@@ -116,13 +107,14 @@ def main():
         "(Oct08)_RankInv_Beta.txt")
     dataset_name = "Eye_AXBXA_1008_RankInv"
 
-    cursor = get_cursor()
-    strains = get_strains(cursor)
-    print("Getting probset_vals")
-    probeset_vals = get_probeset_vals(cursor, dataset_name)
-    print("Finished getting probeset_vals")
-    trimmed_strains = trim_strains(strains, probeset_vals)
-    write_data_matrix_file(trimmed_strains, probeset_vals, filename)
+    with database_connection as conn:
+        with conn.cursor() as cursor:
+            strains = get_strains(cursor)
+            print("Getting probset_vals")
+            probeset_vals = get_probeset_vals(cursor, dataset_name)
+            print("Finished getting probeset_vals")
+            trimmed_strains = trim_strains(strains, probeset_vals)
+            write_data_matrix_file(trimmed_strains, probeset_vals, filename)
 
 
 if __name__ == '__main__':
diff --git a/wqflask/maintenance/quantile_normalize.py b/wqflask/maintenance/quantile_normalize.py
index 32780ca6..2e2b0ec3 100644
--- a/wqflask/maintenance/quantile_normalize.py
+++ b/wqflask/maintenance/quantile_normalize.py
@@ -1,6 +1,5 @@
 import sys
 sys.path.insert(0, './')
-import MySQLdb
 import urllib.parse
 
 import numpy as np
@@ -9,6 +8,7 @@ import pandas as pd
 from flask import Flask, g, request
 
 from wqflask import app
+from wqflask.database import database_connection
 
 
 def parse_db_uri():
@@ -52,7 +52,7 @@ def quantileNormalize(df_input):
     return df
 
 
-def set_data(dataset_name):
+def set_data(cursor, dataset_name):
     orig_file = "/home/zas1024/cfw_data/" + dataset_name + ".txt"
 
     sample_list = []
@@ -80,8 +80,8 @@ def set_data(dataset_name):
                                  ProbeSetFreeze.Id = ProbeSetXRef.ProbeSetFreezeId and
                                  ProbeSetXRef.ProbeSetId = ProbeSet.Id and
                                  ProbeSet.Name = '%s'""" % (dataset_name, line1.split('\t')[0])
-                Cursor.execute(query)
-                result_info = Cursor.fetchone()
+                cursor.execute(query)
+                result_info = cursor.fetchone()
 
                 yield {
                     "_index": "traits",
@@ -99,15 +99,14 @@ def set_data(dataset_name):
 
 
 if __name__ == '__main__':
-    Conn = MySQLdb.Connect(**parse_db_uri())
-    Cursor = Conn.cursor()
+    with database_connection as conn:
+        with conn.cursor as cursor:
+            success, _ = bulk(es, set_data(cursor, sys.argv[1]))
 
-    success, _ = bulk(es, set_data(sys.argv[1]))
-
-    response = es.search(
-        index="traits", doc_type="trait", body={
-            "query": {"match": {"name": "ENSMUSG00000028982"}}
-        }
-    )
+            response = es.search(
+                index="traits", doc_type="trait", body={
+                    "query": {"match": {"name": "ENSMUSG00000028982"}}
+                }
+            )
 
-    print(response)
+            print(response)
diff --git a/wqflask/maintenance/set_resource_defaults.py b/wqflask/maintenance/set_resource_defaults.py
index 0f472494..22d73ba3 100644
--- a/wqflask/maintenance/set_resource_defaults.py
+++ b/wqflask/maintenance/set_resource_defaults.py
@@ -30,10 +30,9 @@ from utility.tools import SQL_URI
 from utility.redis_tools import get_redis_conn, get_user_id, add_resource, get_resources, get_resource_info
 Redis = get_redis_conn()
 
-import MySQLdb
-
 import urllib.parse
 
+from wqflask.database import database_connection
 from utility.logger import getLogger
 logger = getLogger(__name__)
 
@@ -53,14 +52,14 @@ def parse_db_uri():
     return db_conn_info
 
 
-def insert_probeset_resources(default_owner_id):
+def insert_probeset_resources(cursor, default_owner_id):
     current_resources = Redis.hgetall("resources")
-    Cursor.execute("""  SELECT
+    cursor.execute("""  SELECT
                             ProbeSetFreeze.Id, ProbeSetFreeze.Name, ProbeSetFreeze.confidentiality, ProbeSetFreeze.public
                         FROM
                             ProbeSetFreeze""")
 
-    resource_results = Cursor.fetchall()
+    resource_results = cursor.fetchall()
     for i, resource in enumerate(resource_results):
         resource_ob = {}
         resource_ob['name'] = resource[1]
@@ -80,9 +79,9 @@ def insert_probeset_resources(default_owner_id):
         add_resource(resource_ob, update=False)
 
 
-def insert_publish_resources(default_owner_id):
+def insert_publish_resources(cursor, default_owner_id):
     current_resources = Redis.hgetall("resources")
-    Cursor.execute("""  SELECT 
+    cursor.execute("""  SELECT 
                             PublishXRef.Id, PublishFreeze.Id, InbredSet.InbredSetCode
                         FROM
                             PublishXRef, PublishFreeze, InbredSet, Publication
@@ -91,7 +90,7 @@ def insert_publish_resources(default_owner_id):
                             InbredSet.Id = PublishXRef.InbredSetId AND
                             Publication.Id = PublishXRef.PublicationId""")
 
-    resource_results = Cursor.fetchall()
+    resource_results = cursor.fetchall()
     for resource in resource_results:
         if resource[2]:
             resource_ob = {}
@@ -114,14 +113,14 @@ def insert_publish_resources(default_owner_id):
             continue
 
 
-def insert_geno_resources(default_owner_id):
+def insert_geno_resources(cursor, default_owner_id):
     current_resources = Redis.hgetall("resources")
-    Cursor.execute("""  SELECT
+    cursor.execute("""  SELECT
                             GenoFreeze.Id, GenoFreeze.ShortName, GenoFreeze.confidentiality
                         FROM
                             GenoFreeze""")
 
-    resource_results = Cursor.fetchall()
+    resource_results = cursor.fetchall()
     for i, resource in enumerate(resource_results):
         resource_ob = {}
         resource_ob['name'] = resource[1]
@@ -147,15 +146,15 @@ def insert_geno_resources(default_owner_id):
 def insert_resources(default_owner_id):
     current_resources = get_resources()
     print("START")
-    insert_publish_resources(default_owner_id)
+    insert_publish_resources(cursor, default_owner_id)
     print("AFTER PUBLISH")
-    insert_geno_resources(default_owner_id)
+    insert_geno_resources(cursor, default_owner_id)
     print("AFTER GENO")
-    insert_probeset_resources(default_owner_id)
+    insert_probeset_resources(cursor, default_owner_id)
     print("AFTER PROBESET")
 
 
-def main():
+def main(cursor):
     """Generates and outputs (as json file) the data for the main dropdown menus on the home page"""
 
     Redis.delete("resources")
@@ -166,6 +165,6 @@ def main():
 
 
 if __name__ == '__main__':
-    Conn = MySQLdb.Connect(**parse_db_uri())
-    Cursor = Conn.cursor()
-    main()
+    with database_connection() as conn:
+        with conn.cursor() as cursor:
+            main(cursor)