[covid_hosp:f] Fix geocode acquisition and running time.

* Switch to executemany
* Add new limited_geocode datatype
* Fix test prep missing from #1030

Author: krivard

integrations/acquisition/covid_hosp/facility/test_scenarios.py

@@ -38,6 +38,7 @@ def setUp(self):
     with Database.connect() as db:
       with db.new_cursor() as cur:
         cur.execute('truncate table covid_hosp_facility')
+        cur.execute('truncate table covid_hosp_facility_key')
         cur.execute('truncate table covid_hosp_meta')
 
   @freeze_time("2021-03-16")

src/acquisition/covid_hosp/common/database.py

@@ -186,21 +186,37 @@ def nan_safe_dtype(dtype, value):
           f'VALUES ({value_placeholders})'
     id_and_publication_date = (0, publication_date)
     if logger:
-      logger.info("updating values")
+      logger.info('updating values', count=len(dataframe.index))
+    n = 0
+    many_values = []
     with self.new_cursor() as cursor:
-      for _, row in dataframe.iterrows():
+      for index, row in dataframe.iterrows():
         values = []
         for c in dataframe_columns_and_types:
           values.append(nan_safe_dtype(c.dtype, row[c.csv_name]))
-        cursor.execute(sql,
-                       id_and_publication_date +
-                       tuple(values) +
-                       tuple(i.csv_name for i in self.additional_fields))
+        many_values.append(id_and_publication_date +
+          tuple(values) +
+          tuple(i.csv_name for i in self.additional_fields))
+        n += 1
+        # insert in batches because one at a time is slow and all at once makes
+        # the connection drop :(
+        if n % 5_000 == 0:
+          try:
+            cursor.executemany(sql, many_values)
+            many_values = []
+          except Exception as e:
+            if logger:
+              logger.info('error on insert', index=index, values=values)
+              logger.error(e)
+            raise e
+      # insert final batch
+      if many_values:
+        cursor.executemany(sql, many_values)
 
     # deal with non/seldomly updated columns used like a fk table (if this database needs it)
     if hasattr(self, 'AGGREGATE_KEY_COLS'):
       if logger:
-        logger.info("updating keys")
+        logger.info('updating keys')
       ak_cols = self.AGGREGATE_KEY_COLS
 
       # restrict data to just the key columns and remove duplicate rows
@@ -227,6 +243,8 @@ def nan_safe_dtype(dtype, value):
       ak_table = self.table_name + '_key'
       # assemble full SQL statement
       ak_insert_sql = f'INSERT INTO `{ak_table}` ({ak_cols_str}) VALUES ({values_str}) AS v ON DUPLICATE KEY UPDATE {ak_updates_str}'
+      if logger:
+        logger.info("database query", sql=ak_insert_sql)
 
       # commit the data
       with self.new_cursor() as cur:

src/acquisition/covid_hosp/common/network.py

@@ -13,6 +13,8 @@ def fetch_metadata_for_dataset(dataset_id, logger=False):
     ----------
     dataset_id : str
       healthdata.gov dataset identifier of the dataset.
+    logger : structlog.Logger [optional; default False]
+      Logger to receive messages.
 
     Returns
     -------
@@ -39,6 +41,8 @@ def fetch_dataset(url, pandas_impl=pandas, logger=False):
     ----------
     url : str
       URL to the dataset in CSV format.
+    logger : structlog.Logger [optional; default False]
+      Logger to receive messages.
 
     Returns
     -------

src/acquisition/covid_hosp/common/utils.py

@@ -79,6 +79,17 @@ def limited_string(value):
       return value
     return limited_string
 
+  GEOCODE_LENGTH = 32
+  GEOCODE_PATTERN = re.compile(r'POINT \(([0-9.-]*) ([0-9.-]*)\)')
+  def limited_geocode(value):
+    if len(value) < Utils.GEOCODE_LENGTH:
+      return value
+    # otherwise parse and reduce precision to 5
+    m = Utils.GEOCODE_PATTERN.match(value)
+    if not m:
+      raise CovidHospException(f"Couldn't parse geocode '{value}'")
+    return f'POINT ({" ".join(map(lambda x: f"{float(x):.6f}", m.groups()))})'
+
   def issues_to_fetch(metadata, newer_than, older_than, logger=False):
     """
     Construct all issue dates and URLs to be ingested based on metadata.
@@ -100,6 +111,7 @@ def issues_to_fetch(metadata, newer_than, older_than, logger=False):
     """
     daily_issues = {}
     n_beyond = 0
+    n_selected = 0
     for index in sorted(set(metadata.index)):
       day = index.date()
       if day > newer_than and day < older_than:
@@ -109,11 +121,13 @@ def issues_to_fetch(metadata, newer_than, older_than, logger=False):
           daily_issues[day] = urls_list
         else:
           daily_issues[day] += urls_list
+        n_selected += len(urls_list)
       elif day >= older_than:
         n_beyond += 1
-    if n_beyond > 0:
-      if logger:
-        logger.info("issues available", on_or_newer=older_than, count=n_beyond)
+    if logger:
+      if n_beyond > 0:
+        logger.info("issues available beyond selection", on_or_newer=older_than, count=n_beyond)
+      logger.info("issues selected", newer_than=str(newer_than), older_than=str(older_than), count=n_selected)
     return daily_issues
 
   @staticmethod
@@ -173,7 +187,7 @@ def update_dataset(database, network, newer_than=None, older_than=None):
     bool
       Whether a new dataset was acquired.
     """
-    logger = get_structured_logger(f"{database.__class__.__module__}.{database.__class__.__name__}.update_dataset")
+    logger = get_structured_logger(f"{database.__module__}.{database.__name__}.update_dataset")
 
     metadata = network.fetch_metadata(logger=logger)
     datasets = []
@@ -189,7 +203,7 @@ def update_dataset(database, network, newer_than=None, older_than=None):
     for issue, revisions in daily_issues.items():
       issue_int = int(issue.strftime("%Y%m%d"))
       # download the dataset and add it to the database
-      dataset = Utils.merge_by_key_cols([network.fetch_dataset(url) for url, _ in revisions],
+      dataset = Utils.merge_by_key_cols([network.fetch_dataset(url, logger=logger) for url, _ in revisions],
                                         db.KEY_COLS,
                                         logger=logger)
       # add metadata to the database

src/acquisition/covid_hosp/facility/database.py

@@ -40,7 +40,7 @@ class Database(BaseDatabase):
       Columndef('ccn', 'ccn', str),
       Columndef('city', 'city', str),
       Columndef('fips_code', 'fips_code', str),
-      Columndef('geocoded_hospital_address', 'geocoded_hospital_address', Utils.limited_string_fn(32)),
+      Columndef('geocoded_hospital_address', 'geocoded_hospital_address', Utils.limited_geocode),
       Columndef('hhs_ids', 'hhs_ids', str),
       Columndef('hospital_name', 'hospital_name', str),
       Columndef('hospital_subtype', 'hospital_subtype', str),

tests/acquisition/covid_hosp/common/test_database.py

@@ -144,9 +144,9 @@ def test_insert_dataset(self):
     result = database.insert_dataset(sentinel.publication_date, dataset)
 
     self.assertIsNone(result)
-    self.assertEqual(mock_cursor.execute.call_count, 6)
+    self.assertEqual(mock_cursor.executemany.call_count, 1)
 
-    actual_sql = mock_cursor.execute.call_args[0][0]
+    actual_sql = mock_cursor.executemany.call_args[0][0]
     self.assertIn(
       'INSERT INTO `test_table` (`id`, `publication_date`, `sql_str_col`, `sql_int_col`, `sql_float_col`)',
       actual_sql)
@@ -162,5 +162,9 @@ def test_insert_dataset(self):
 
     for i, expected in enumerate(expected_values):
       with self.subTest(name=f'row {i + 1}'):
-        actual = mock_cursor.execute.call_args_list[i][0][1]
+        # [0]: the first call() object
+        # [0]: get positional args out of the call() object
+        # [-1]: the last arg of the executemany call
+        # [i]: the ith row inserted in the executemany
+        actual = mock_cursor.executemany.call_args_list[0][0][-1][i]
         self.assertEqual(actual, (0, sentinel.publication_date) + expected)

tests/acquisition/covid_hosp/common/test_utils.py

@@ -97,7 +97,7 @@ def test_run_skip_old_dataset(self):
     mock_network = MagicMock()
     mock_network.fetch_metadata.return_value = \
         self.test_utils.load_sample_metadata()
-    mock_database = MagicMock()
+    mock_database = MagicMock(**{"__module__":"test_module", "__name__":"test_name"})
     with mock_database.connect() as mock_connection:
       pass
     mock_connection.get_max_issue.return_value = pd.Timestamp("2200/1/1")
@@ -117,7 +117,7 @@ def test_run_acquire_new_dataset(self):
         self.test_utils.load_sample_metadata()
     fake_dataset = pd.DataFrame({"date": [pd.Timestamp("2020/1/1")], "state": ["ca"]})
     mock_network.fetch_dataset.return_value = fake_dataset
-    mock_database = MagicMock()
+    mock_database = MagicMock(**{"__module__":"test_module", "__name__":"test_name"})
     with mock_database.connect() as mock_connection:
       pass
     type(mock_connection).KEY_COLS = PropertyMock(return_value=["state", "date"])

tests/acquisition/covid_hosp/facility/test_database.py

@@ -35,9 +35,14 @@ def test_insert_dataset(self):
     result = database.insert_dataset(sentinel.publication_date, dataset)
 
     self.assertIsNone(result)
-    self.assertEqual(mock_cursor.execute.call_count, 22)
-
-    last_query_values = mock_cursor.execute.call_args[0][-1]
+    # once for the values, once for the keys
+    self.assertEqual(mock_cursor.executemany.call_count, 2)
+
+    # [0]: the first call() object
+    # [0]: get the positional args out of the call() object
+    # [-1]: the last arg of the executemany call
+    # [-1]: the last row inserted in the executemany
+    last_query_values = mock_cursor.executemany.call_args_list[0][0][-1][-1]
     expected_query_values = (
         0, sentinel.publication_date, '450822', 20201130,
         '6800 N MACARTHUR BLVD', 61.1, 7, 428, 60.9, 7, 426, 61.1, 7, 428,

tests/acquisition/covid_hosp/state_daily/test_database.py

@@ -38,9 +38,9 @@ def test_insert_dataset(self):
     result = database.insert_dataset(sentinel.issue, dataset)
 
     self.assertIsNone(result)
-    self.assertEqual(mock_cursor.execute.call_count, 53)
+    self.assertEqual(mock_cursor.executemany.call_count, 1)
 
-    last_query_values = mock_cursor.execute.call_args[0][-1]
+    last_query_values = mock_cursor.executemany.call_args[0][-1][-1]
     expected_query_values = (
         0, sentinel.issue,  'WY', 20201209,
         0.2519685039370078, 29, 127, 32, 0.4233576642335766, 31, 137, 58, 22, 2,

tests/acquisition/covid_hosp/state_timeseries/test_database.py

@@ -36,9 +36,9 @@ def test_insert_dataset(self):
     result = database.insert_dataset(sentinel.issue, dataset)
 
     self.assertIsNone(result)
-    self.assertEqual(mock_cursor.execute.call_count, 22)
+    self.assertEqual(mock_cursor.executemany.call_count, 1)
 
-    last_query_values = mock_cursor.execute.call_args[0][-1]
+    last_query_values = mock_cursor.executemany.call_args[0][-1][-1]
     expected_query_values = (
         0, sentinel.issue, 'WY', 20200826, 0.0934579439252336, 26, 107, 10,
         0.4298245614035088, 28, 114, 49, 19, 7, 2, None, 4, 2, 0, 1, '2', 0, 26,