Merge pull request #64 from CSCfi/feature/different_loader

Optimise init load in DB for parsing VCF file

Author: teemukataja

beacon_api/utils/db_load.py

@@ -101,24 +101,18 @@ def _unpack(self, variant):
         me_type = ['dup:tandem', 'del:me', 'ins:me']
         # sv_type = ['dup', 'inv', 'ins', 'del', 'cnv']
         # supported_vt = ['snp', 'indel', 'mnp', 'dup', 'inv', 'ins', 'del']
-        for k, v in variant.INFO:
-            if k == 'AC':
-                ac = self._handle_type(v, int)
-            if k == 'AN':
-                an = v
-            else:
-                an = variant.num_called*2
-            if k == 'AF':
-                aaf = self._handle_type(v, float)
-            else:
-                aaf = [float(ac_value) / float(an) for ac_value in ac]
-            if variant.is_sv:
-                alt = [elem.strip("<>") for elem in variant.ALT]
-                if k == 'SVTYPE':
-                    vt = [self._rchop(e, ":"+v) if e.lower().startswith(tuple(me_type)) else v for e in alt]
-            else:
-                if k == 'VT':
-                    vt = [self._transform_vt(var_type.lower(), variant) for var_type in v.split(',')]
+        ac = self._handle_type(variant.INFO.get('AC'), int) if variant.INFO.get('AC') else []
+        an = variant.INFO.get('AN') if variant.INFO.get('AN') else variant.num_called * 2
+        aaf = self._handle_type(variant.INFO.get('AF'), float) if variant.INFO.get('AF') else [float(ac_value) / float(an) for ac_value in ac]
+        if variant.is_sv:
+            alt = [elem.strip("<>") for elem in variant.ALT]
+            if variant.INFO.get('SVTYPE'):
+                v = variant.INFO.get('SVTYPE')
+                vt = [self._rchop(e, ":"+v) if e.lower().startswith(tuple(me_type)) else v for e in alt]
+        else:
+            if variant.INFO.get('VT'):
+                v = variant.INFO.get('VT')
+                vt = [self._transform_vt(var_type.lower(), variant) for var_type in v.split(',')]
 
         return (aaf, ac, vt, alt, an)
 
@@ -235,6 +229,7 @@ async def insert_variants(self, dataset_id, variants):
             async with self._conn.transaction():
                 LOG.info('Insert variants into the database')
                 for variant in variants:
+                    # print(variant)
                     # params = (frequency, count, actual variant Type)
                     if variant.aaf > 0:
                         params = self._unpack(variant)

tests/test_basic.py

@@ -1,6 +1,7 @@
 import asynctest
 from beacon_api.utils.db_load import parse_arguments, init_beacon_db, main
 from beacon_api.conf.config import init_db_pool
+from beacon_api.permissions.rems import get_rems_controlled
 from testfixtures import TempDirectory
 
 
@@ -94,6 +95,21 @@ def test_main_db(self, mock_init):
         main()
         mock_init.assert_called()
 
+    def test_rems_controlled(self):
+        """Test rems permissions claim parsing."""
+        claim = [{"affiliation": "",
+                  "datasets": ["EGAD01", "urn:hg:example-controlled"],
+                  "source_signature": "",
+                  "url_prefix": ""},
+                 {"affiliation": "",
+                  "datasets": ["urn:hg:example-controlled", "EGAD02",
+                               "urn:hg:example-controlled3"],
+                  "source_signature": "",
+                  "url_prefix": ""}]
+        self.assertCountEqual(get_rems_controlled(claim),
+                              ['EGAD01', 'urn:hg:example-controlled',
+                               'urn:hg:example-controlled3', 'EGAD02'])
+
 
 if __name__ == '__main__':
     asynctest.main()

tests/test_db_load.py

@@ -7,12 +7,12 @@
 class Variant:
     """Variant Class.
 
-    Mock this for Variant caculations.
+    Mock this for Variant calculations.
     """
 
-    def __init__(self, ALT, REF, INFO, call_rate, var_type, num_called):
+    def __init__(self, ALT, REF, INF, call_rate, var_type, num_called):
         """Initialize class."""
-        self.INFO = INFO.items()
+        self.INFO = INF
         self.ALT = ALT
         self.REF = REF
         self.call_rate = call_rate
@@ -21,6 +21,24 @@ def __init__(self, ALT, REF, INFO, call_rate, var_type, num_called):
         self.is_sv = False
 
 
+class INFO:
+    """INFO CLass.
+
+    Mock this for storing VCF info.
+    """
+
+    def __init__(self, AC, VT, AN):
+        """Initialize class."""
+        self.AC = AC
+        self.VT = VT
+        self.AN = AN
+        self.AF = 0
+
+    def get(self, key):
+        """Inside `__getitem__` method."""
+        return getattr(self, key)
+
+
 class Transaction:
     """Class Transaction.
 
@@ -234,10 +252,13 @@ async def test_unpack(self, db_mock, mock_log):
         """Test database URL fetching."""
         db_mock.return_value = Connection()
         await self._db.connection()
-        variant = Variant('TC', 'T', {'AC': (1, 2), 'VT': 'M,S,I', 'AN': 3}, 0.7, 'snp', 3)
+        inf1 = INFO((1, 2), 'M,S,I', 3)
+        variant = Variant('TC', 'T', inf1, 0.7, 'snp', 3)
         result = self._db._unpack(variant)
         self.assertEqual(([0.3333333333333333, 0.6666666666666666], [1, 2], ['MNP', 'SNP', 'INS'], 'TC', 3), result)
-        variant = Variant('TC', 'T', {'AC': 1, 'VT': 'S', 'AN': 3}, 0.7, 'snp', 3)
+        inf2 = INFO(1, 'S', 3)
+        variant = Variant('TC', 'T', inf2, 0.7, 'snp', 3)
+        # print(variant.INFO.get('AC'))
         result = self._db._unpack(variant)
         self.assertEqual(([0.3333333333333333], [1], ['SNP'], 'TC', 3), result)