Replace "g.db" object with a proper database connection

* wqflask/tests/unit/wqflask/api/test_correlation.py: Use proper
database connection instead of the db connection attached at "g.db".
* wqflask/tests/unit/wqflask/snp_browser/test_snp_browser.py: Ditto.
* wqflask/wqflask/api/correlation.py: Ditto.
* wqflask/wqflask/snp_browser/snp_browser.py: Ditto.

4 files changed, 202 insertions, 196 deletions

diff --git a/wqflask/tests/unit/wqflask/api/test_correlation.py b/wqflask/tests/unit/wqflask/api/test_correlation.py
index 1089a36f..57986e80 100644
--- a/wqflask/tests/unit/wqflask/api/test_correlation.py
+++ b/wqflask/tests/unit/wqflask/api/test_correlation.py
@@ -52,28 +52,28 @@ class TestCorrelations(unittest.TestCase):
 
         self.assertEqual(corr_params_results, expected_results)
 
-    @mock.patch("wqflask.api.correlation.g")
+    @mock.patch("wqflask.api.correlation.database_connection")
     def test_convert_to_mouse_gene_id(self, mock_db):
-
-        results = convert_to_mouse_gene_id(species="Other", gene_id="")
-        self.assertEqual(results, None)
-
-        rat_species_results = convert_to_mouse_gene_id(
-            species="rat", gene_id="GH1")
-
-        mock_db.db.execute.return_value.fetchone.side_effect = [
-            AttributeSetter({"mouse": "MG-1"}), AttributeSetter({"mouse": "MG-2"})]
-
-        self.assertEqual(convert_to_mouse_gene_id(
-            species="mouse", gene_id="MG-4"), "MG-4")
-        self.assertEqual(convert_to_mouse_gene_id(
-            species="rat", gene_id="R1"), "MG-1")
-        self.assertEqual(convert_to_mouse_gene_id(
-            species="human", gene_id="H1"), "MG-2")
-
-    @mock.patch("wqflask.api.correlation.g")
+        conn = mock.MagicMock()
+        mock_db.return_value.__enter__.return_value = conn
+        with conn.cursor() as cursor:
+            cursor.fetchone.side_effect = [
+                ("MG-1",), ("MG-2",)]
+
+            self.assertEqual(convert_to_mouse_gene_id(
+                species="Other", gene_id=""), None)
+            self.assertEqual(convert_to_mouse_gene_id(
+                species="mouse", gene_id="MG-4"), "MG-4")
+            self.assertEqual(convert_to_mouse_gene_id(
+                species="rat", gene_id="R1"), "MG-1")
+            self.assertEqual(convert_to_mouse_gene_id(
+                species="human", gene_id="H1"), "MG-2")
+
+    @mock.patch("wqflask.api.correlation.database_connection")
     @mock.patch("wqflask.api.correlation.convert_to_mouse_gene_id")
-    def test_do_literature_correlation_for_all_traits(self, mock_convert_to_mouse_geneid, mock_db):
+    def test_do_literature_correlation_for_all_traits(
+            self,
+            mock_convert_to_mouse_geneid, mock_db):
         mock_convert_to_mouse_geneid.side_effect = [
             "MG-1", "MG-2;", "MG-3", "MG-4"]
 
@@ -83,19 +83,22 @@ class TestCorrelations(unittest.TestCase):
             "TT-3": "GH-3"
 
         }
-        mock_db.db.execute.return_value.fetchone.side_effect = [AttributeSetter(
-            {"value": "V1"}), AttributeSetter({"value": "V2"}), AttributeSetter({"value": "V3"})]
-
-        this_trait = AttributeSetter({"geneid": "GH-1"})
-
-        target_dataset = AttributeSetter(
-            {"group": AttributeSetter({"species": "rat"})})
-        results = do_literature_correlation_for_all_traits(
-            this_trait=this_trait, target_dataset=target_dataset, trait_geneid_dict=trait_geneid_dict, corr_params={})
-
-        expected_results = {'TT-1': ['GH-1', 0],
-                            'TT-2': ['GH-2', 'V1'], 'TT-3': ['GH-3', 'V2']}
-        self.assertEqual(results, expected_results)
+        conn = mock.MagicMock()
+        mock_db.return_value.__enter__.return_value = conn
+        with conn.cursor() as cursor:
+            cursor.fetchone.side_effect = [
+                AttributeSetter({"value": "V1"}),
+                AttributeSetter({"value": "V2"}),
+                AttributeSetter({"value": "V3"})]
+            this_trait = AttributeSetter({"geneid": "GH-1"})
+            target_dataset = AttributeSetter(
+                {"group": AttributeSetter({"species": "rat"})})
+            results = do_literature_correlation_for_all_traits(
+                this_trait=this_trait, target_dataset=target_dataset,
+                trait_geneid_dict=trait_geneid_dict, corr_params={})
+            expected_results = {'TT-1': ['GH-1', 0],
+                                'TT-2': ['GH-2', 'V1'], 'TT-3': ['GH-3', 'V2']}
+            self.assertEqual(results, expected_results)
 
     @mock.patch("wqflask.api.correlation.corr_result_helpers.normalize_values")
     def test_get_sample_r_and_p_values(self, mock_normalize):
diff --git a/wqflask/tests/unit/wqflask/snp_browser/test_snp_browser.py b/wqflask/tests/unit/wqflask/snp_browser/test_snp_browser.py
index 89442c47..46a74231 100644
--- a/wqflask/tests/unit/wqflask/snp_browser/test_snp_browser.py
+++ b/wqflask/tests/unit/wqflask/snp_browser/test_snp_browser.py
@@ -40,66 +40,88 @@ class TestSnpBrowser(unittest.TestCase):
         self.assertEqual(expected_results, results_with_snp)
         self.assertEqual(expected_results_with_indel, results_with_indel)
 
-    @mock.patch("wqflask.snp_browser.snp_browser.g")
+    @mock.patch("wqflask.snp_browser.snp_browser.database_connection")
     def test_get_gene_id(self, mock_db):
-        mock_db.db.execute.return_value.fetchone.return_value = "517d729f-aa13-4413-a885-40a3f7ff768a"
-        db_query_value = """
-                SELECT
-                        geneId
-                FROM
-                        GeneList
-                WHERE
-                        SpeciesId = c9c0f59e-1259-4cba-91e6-831ef1a99c83 AND geneSymbol = 'INSR'
-            """
-        results = get_gene_id(
-            species_id="c9c0f59e-1259-4cba-91e6-831ef1a99c83", gene_name="INSR")
-        mock_db.db.execute.assert_called_once_with(db_query_value)
-        self.assertEqual(results, "517d729f-aa13-4413-a885-40a3f7ff768a")
+        db_query_value = ("SELECT geneId FROM GeneList WHERE "
+                          "SpeciesId = %s AND geneSymbol = %s")
+        conn = mock.MagicMock()
+        mock_db.return_value.__enter__.return_value = conn
+        with conn.cursor() as cursor:
+            cursor.fetchone.return_value = (("517d729f-aa13-4413"
+                                             "-a885-40a3f7ff768a"),)
 
-    @mock.patch("wqflask.snp_browser.snp_browser.g")
+            results = get_gene_id(
+                species_id="c9c0f59e-1259-4cba-91e6-831ef1a99c83",
+                gene_name="INSR")
+            cursor.execute.assert_called_once_with(
+                db_query_value,
+                ("c9c0f59e-1259-4cba-91e6-831ef1a99c83",
+                 "INSR"))
+            self.assertEqual(results,
+                             "517d729f-aa13-4413-a885-40a3f7ff768a")
+
+    @mock.patch("wqflask.snp_browser.snp_browser.database_connection")
     def test_gene_id_name_dict(self, mock_db):
         no_gene_names = []
-        self.assertEqual("", get_gene_id_name_dict(
-            species_id="fregb343bui43g4", gene_name_list=no_gene_names))
-        gene_name_list = ["GH1", "GH2", "GH3"]
-        mock_db.db.execute.return_value.fetchall.side_effect = [[], [("fsdf43-fseferger-f22", "GH1"), ("1sdf43-fsewferger-f22", "GH2"),
-                                                                     ("fwdj43-fstferger-f22", "GH3")]]
-        no_results = get_gene_id_name_dict(
-            species_id="ret3-32rf32", gene_name_list=gene_name_list)
-        results_found = get_gene_id_name_dict(
-            species_id="ret3-32rf32", gene_name_list=gene_name_list)
-        expected_found = {'GH1': 'fsdf43-fseferger-f22',
-                          'GH2': '1sdf43-fsewferger-f22', 'GH3': 'fwdj43-fstferger-f22'}
-        db_query_value = """
-                SELECT
-                        geneId, geneSymbol
-                FROM
-                        GeneList
-                WHERE
-                        SpeciesId = ret3-32rf32 AND geneSymbol in ('GH1','GH2','GH3')
-            """
-        mock_db.db.execute.assert_called_with(db_query_value)
-        self.assertEqual(results_found, expected_found)
-        self.assertEqual(no_results, {})
+        conn = mock.MagicMock()
+        mock_db.return_value.__enter__.return_value = conn
+        with conn.cursor() as cursor:
+            cursor.fetchall.side_effect = [
+                [],
+                [("fsdf43-fseferger-f22", "GH1"),
+                 ("1sdf43-fsewferger-f22", "GH2"),
+                 ("fwdj43-fstferger-f22", "GH3")]]
+            self.assertEqual("", get_gene_id_name_dict(
+                species_id="fregb343bui43g4",
+                gene_name_list=no_gene_names))
+            gene_name_list = ["GH1", "GH2", "GH3"]
+            no_results = get_gene_id_name_dict(
+                species_id="ret3-32rf32", gene_name_list=gene_name_list)
+            results_found = get_gene_id_name_dict(
+                species_id="ret3-32rf32", gene_name_list=gene_name_list)
+            expected_found = {'GH1': 'fsdf43-fseferger-f22',
+                              'GH2': '1sdf43-fsewferger-f22',
+                              'GH3': 'fwdj43-fstferger-f22'}
+            db_query_value = (
+                "SELECT geneId, geneSymbol FROM GeneList WHERE "
+                "SpeciesId = %s AND geneSymbol in (%s, %s, %s)")
+            cursor.execute.assert_called_with(
+                db_query_value, ("ret3-32rf32", "GH1", "GH2", "GH3"))
+            self.assertEqual(results_found, expected_found)
+            self.assertEqual(no_results, {})
 
-    @mock.patch("wqflask.snp_browser.snp_browser.g")
+    @mock.patch("wqflask.snp_browser.snp_browser.database_connection")
     def test_check_if_in_gene(self, mock_db):
-        mock_db.db.execute.return_value.fetchone.side_effect = [
-            ("fsdf-232sdf-sdf", "GHA"), ""]
-        results_found = check_if_in_gene(
-            species_id="517d729f-aa13-4413-a885-40a3f7ff768a", chr="CH1", mb=12.09)
-        db_query_value = """SELECT geneId, geneSymbol
-                   FROM GeneList
-                   WHERE SpeciesId = 517d729f-aa13-4413-a885-40a3f7ff768a AND chromosome = 'CH1' AND
-                        (txStart < 12.09 AND txEnd > 12.09); """
-        gene_not_found = check_if_in_gene(
-            species_id="517d729f-aa13-4413-a885-40a3f7ff768a", chr="CH1", mb=12.09)
-        mock_db.db.execute.assert_called_with(db_query_value)
-        self.assertEqual(gene_not_found, "")
+        conn = mock.MagicMock()
+        mock_db.return_value.__enter__.return_value = conn
+        with conn.cursor() as cursor:
+            cursor.fetchone.side_effect = [
+                ("fsdf-232sdf-sdf", "GHA"), ""]
+            results_found = check_if_in_gene(
+                species_id="517d729f-aa13-4413-a885-40a3f7ff768a",
+                chr_="CH1", mb=12.09)
+            self.assertEqual(results_found, ["fsdf-232sdf-sdf", "GHA"])
+            db_query_value = (
+                "SELECT geneId, geneSymbol FROM GeneList "
+                "WHERE SpeciesId = %s AND chromosome = %s "
+                "AND (txStart < %s AND txEnd > %s)")
+            gene_not_found = check_if_in_gene(
+                species_id="517d729f-aa13-4413-a885-40a3f7ff768a",
+                chr_="CH1", mb=12.09)
+            cursor.execute.assert_has_calls(
+                [mock.call(db_query_value,
+                           ("517d729f-aa13-4413-a885-40a3f7ff768a",
+                            "CH1", 12.09, 12.09)),
+                 mock.call(db_query_value,
+                           ("517d729f-aa13-4413-a885-40a3f7ff768a",
+                            "CH1", 12.09, 12.09))])
+            self.assertEqual(gene_not_found, "")
 
-    @mock.patch("wqflask.snp_browser.snp_browser.g")
+    @mock.patch("wqflask.snp_browser.snp_browser.database_connection")
     def test_get_browser_sample_lists(self, mock_db):
-        mock_db.db.execute.return_value.fetchall.return_value = []
-
-        results = get_browser_sample_lists(species_id="12")
-        self.assertEqual(results, {'mouse': [], 'rat': []})
+        conn = mock.MagicMock()
+        mock_db.return_value.__enter__.return_value = conn
+        with conn.cursor() as cursor:
+            cursor.execute.return_value.fetchall.return_value = []
+            results = get_browser_sample_lists(species_id="12")
+            self.assertEqual(results, {'mouse': [], 'rat': []})
diff --git a/wqflask/wqflask/api/correlation.py b/wqflask/wqflask/api/correlation.py
index 9b875c99..ab1e772a 100644
--- a/wqflask/wqflask/api/correlation.py
+++ b/wqflask/wqflask/api/correlation.py
@@ -3,11 +3,9 @@ import scipy
 
 from base import data_set
 from base.trait import create_trait, retrieve_sample_data
-from flask import g
 from utility import corr_result_helpers
-from utility.db_tools import escape
 from wqflask.correlation import correlation_functions
-
+from wqflask.database import database_connection
 
 def do_correlation(start_vars):
     assert('db' in start_vars)
@@ -125,22 +123,24 @@ def do_literature_correlation_for_all_traits(this_trait, target_dataset, trait_g
             target_dataset.group.species.lower(), gene_id)
 
         if mouse_gene_id and str(mouse_gene_id).find(";") == -1:
-            result = g.db.execute(
-                """SELECT value
-                    FROM LCorrRamin3
-                    WHERE GeneId1='%s' and
-                            GeneId2='%s'
-                """ % (escape(mouse_gene_id), escape(input_trait_mouse_gene_id))
-            ).fetchone()
-            if not result:
-                result = g.db.execute("""SELECT value
-                    FROM LCorrRamin3
-                    WHERE GeneId2='%s' and
-                            GeneId1='%s'
-                """ % (escape(mouse_gene_id), escape(input_trait_mouse_gene_id))
-                ).fetchone()
+            result = ""
+            with database_connection() as conn:
+                with conn.cursor() as cursor:
+                    cursor.execute(
+                        ("SELECT value FROM LCorrRamin3 "
+                         "WHERE GeneId1=%s AND GeneId2=%s"),
+                        (mouse_gene_id,
+                         input_trait_mouse_gene_id))
+                    result = cursor.fetchone()
+                    if not result:
+                        cursor.execute(
+                            ("SELECT value FROM LCorrRamin3 "
+                             "WHERE GeneId2=%s AND GeneId1=%s"),
+                            (mouse_gene_id,
+                             input_trait_mouse_gene_id))
+                        result = cursor.fetchone()
             if result:
-                lit_corr = result.value
+                lit_corr = result[0]
                 lit_corr_data[trait] = [gene_id, lit_corr]
             else:
                 lit_corr_data[trait] = [gene_id, 0]
@@ -195,30 +195,24 @@ def convert_to_mouse_gene_id(species=None, gene_id=None):
         return None
 
     mouse_gene_id = None
-
-    if species == 'mouse':
-        mouse_gene_id = gene_id
-
-    elif species == 'rat':
-
-        query = """SELECT mouse
-                FROM GeneIDXRef
-                WHERE rat='%s'""" % escape(gene_id)
-
-        result = g.db.execute(query).fetchone()
-        if result != None:
-            mouse_gene_id = result.mouse
-
-    elif species == 'human':
-
-        query = """SELECT mouse
-                FROM GeneIDXRef
-                WHERE human='%s'""" % escape(gene_id)
-
-        result = g.db.execute(query).fetchone()
-        if result != None:
-            mouse_gene_id = result.mouse
-
+    with database_connection() as conn:
+        with conn.cursor() as cursor:
+            if species == 'mouse':
+                mouse_gene_id = gene_id
+            elif species == 'rat':
+                cursor.execute(
+                    ("SELECT mouse FROM GeneIDXRef "
+                     "WHERE rat=%s"), gene_id)
+                result = cursor.fetchone()
+                if result:
+                    mouse_gene_id = result[0]
+            elif species == 'human':
+                cursor.execute(
+                    "SELECT mouse FROM GeneIDXRef "
+                    "WHERE human=%s", gene_id)
+                result = cursor.fetchone()
+                if result:
+                    mouse_gene_id = result[0]
     return mouse_gene_id
 
 
diff --git a/wqflask/wqflask/snp_browser/snp_browser.py b/wqflask/wqflask/snp_browser/snp_browser.py
index cbdaf901..43bb55b5 100644
--- a/wqflask/wqflask/snp_browser/snp_browser.py
+++ b/wqflask/wqflask/snp_browser/snp_browser.py
@@ -6,6 +6,8 @@ from PIL import (Image)
 from base import species
 from base import webqtlConfig
 
+from wqflask.database import database_connection
+
 
 class SnpBrowser:
 
@@ -662,24 +664,22 @@ class SnpBrowser:
         click_mb_step = (self.end_mb - self.start_mb) / n_click
 
 
-
 def get_browser_sample_lists(species_id=1):
     strain_lists = {}
     mouse_strain_list = []
-    query = "SHOW COLUMNS FROM SnpPattern;"
-    results = g.db.execute(query).fetchall()
-    for result in results[1:]:
-        mouse_strain_list.append(result[0])
-
     rat_strain_list = []
-    query = "SHOW COLUMNS FROM RatSnpPattern;"
-    results = g.db.execute(query).fetchall()
-    for result in results[2:]:
-        rat_strain_list.append(result[0])
-
-    strain_lists['mouse'] = mouse_strain_list
-    strain_lists['rat'] = rat_strain_list
-
+    with database_connection() as conn:
+        with conn.cursor() as cursor:
+            cursor.execute("SHOW COLUMNS FROM SnpPattern")
+            _mouse_snp_pattern = cursor.fetchall()
+            cursor.execute("SHOW COLUMNS FROM RatSnpPattern")
+            _rats_snp_pattern = cursor.fetchall()
+            for result in _mouse_snp_pattern[1:]:
+                mouse_strain_list.append(result[0])
+            for result in _rats_snp_pattern[2:]:
+                rat_strain_list.append(result[0])
+            strain_lists['mouse'] = mouse_strain_list
+            strain_lists['rat'] = rat_strain_list
     return strain_lists
 
 
@@ -885,64 +885,51 @@ def get_effect_info(effect_list):
 
 
 def get_gene_id(species_id, gene_name):
-    query = """
-                SELECT
-                        geneId
-                FROM
-                        GeneList
-                WHERE
-                        SpeciesId = %s AND geneSymbol = '%s'
-            """ % (species_id, gene_name)
-
-    result = g.db.execute(query).fetchone()
-
-    if len(result) > 0:
-        return result
-    else:
-        return ""
+    query = ("SELECT geneId FROM GeneList WHERE "
+             "SpeciesId = %s AND geneSymbol = %s")
+
+    with database_connection() as conn:
+        with conn.cursor() as cursor:
+            cursor.execute(query, (species_id, gene_name))
+            if (result := cursor.fetchone()):
+                return result[0]
+    return ""
 
 
 def get_gene_id_name_dict(species_id, gene_name_list):
     gene_id_name_dict = {}
     if len(gene_name_list) == 0:
         return ""
-    gene_name_str_list = ["'" + gene_name + \
-                          "'" for gene_name in gene_name_list]
-    gene_name_str = ",".join(gene_name_str_list)
-
-    query = """
-                SELECT
-                        geneId, geneSymbol
-                FROM
-                        GeneList
-                WHERE
-                        SpeciesId = %s AND geneSymbol in (%s)
-            """ % (species_id, gene_name_str)
-
-    results = g.db.execute(query).fetchall()
-
-    if len(results) > 0:
-        for item in results:
-            gene_id_name_dict[item[1]] = item[0]
-
+    query = ("SELECT geneId, geneSymbol FROM "
+             "GeneList WHERE SpeciesId = %s AND "
+             f"geneSymbol in ({', '.join(['%s'] * len(gene_name_list))})")
+    with database_connection() as conn:
+        with conn.cursor() as cursor:
+            cursor.execute(query, (species_id, *gene_name_list))
+            results = cursor.fetchall()
+            if results:
+                for item in results:
+                    gene_id_name_dict[item[1]] = item[0]
     return gene_id_name_dict
 
 
-def check_if_in_gene(species_id, chr, mb):
-    if species_id != 0:  # ZS: Check if this is necessary
-        query = """SELECT geneId, geneSymbol
-                   FROM GeneList
-                   WHERE SpeciesId = {0} AND chromosome = '{1}' AND
-                        (txStart < {2} AND txEnd > {2}); """.format(species_id, chr, mb)
-    else:
-        query = """SELECT geneId,geneSymbol
-                   FROM GeneList
-                   WHERE chromosome = '{0}' AND
-                        (txStart < {1} AND txEnd > {1}); """.format(chr, mb)
-
-    result = g.db.execute(query).fetchone()
-
-    if result:
-        return [result[0], result[1]]
-    else:
-        return ""
+def check_if_in_gene(species_id, chr_, mb):
+    with database_connection() as conn:
+        with conn.cursor() as cursor:
+            if species_id != 0:  # ZS: Check if this is necessary
+                cursor.execute(
+                    "SELECT geneId, geneSymbol "
+                    "FROM GeneList WHERE "
+                    "SpeciesId = %s AND chromosome = %s "
+                    "AND (txStart < %s AND txEnd > %s)",
+                    (species_id, chr_, mb, mb))
+            else:
+                cursor.execute(
+                    "SELECT geneId,geneSymbol "
+                    "FROM GeneList WHERE "
+                    "chromosome = %s AND "
+                    "(txStart < %s AND txEnd > %s)",
+                    (chr_, mb, mb))
+            if (result := cursor.fetchone()):
+                return [result[0], result[1]]
+    return ""