Merge pull request #150 from melange396/csv_acquisition

Improvements to acquisition / CSV file loading.

Author: krivard

integrations/acquisition/covidcast/test_csv_uploading.py

@@ -75,6 +75,11 @@ def test_uploading(self):
       f.write('nd,20,0.02,200\n')
       f.write('wa,30,0.03,300\n')
 
+    # invalid
+    with open(source_receiving_dir + '/20200419_state_wip_really_long_name_that_will_get_truncated.csv', 'w') as f:
+      f.write('geo_id,val,se,sample_size\n')
+      f.write('pa,100,5.4,624\n')
+
     # invalid
     with open(source_receiving_dir + '/20200420_state_test.csv', 'w') as f:
       f.write('this,header,is,wrong\n')
@@ -83,11 +88,6 @@ def test_uploading(self):
     with open(source_receiving_dir + '/hello.csv', 'w') as f:
       f.write('file name is wrong\n')
 
-    # invalid
-    with open(source_receiving_dir + '/20200419_state_wip_really_long_name_that_will_get_truncated.csv', 'w') as f:
-      f.write('geo_id,val,se,sample_size\n')
-      f.write('pa,100,5.4,624\n')
-
     # upload CSVs
     args = MagicMock(data_dir=data_dir)
     main(args)

src/acquisition/covidcast/csv_importer.py

@@ -92,7 +92,8 @@ def find_csv_files(scan_dir, issue=(date.today(), epi.Week.fromdate(date.today()
     issue_value=-1
     lag_value=-1
 
-    for path in glob.glob(os.path.join(scan_dir, '*', '*')):
+    for path in sorted(glob.glob(os.path.join(scan_dir, '*', '*'))):
+
       if not path.lower().endswith('.csv'):
         # safe to ignore this file
         continue

src/acquisition/covidcast/csv_to_database.py

@@ -6,7 +6,7 @@
 
 # first party
 from delphi.epidata.acquisition.covidcast.csv_importer import CsvImporter
-from delphi.epidata.acquisition.covidcast.database import Database
+from delphi.epidata.acquisition.covidcast.database import Database, CovidcastRow
 from delphi.epidata.acquisition.covidcast.file_archiver import FileArchiver
 
 
@@ -45,7 +45,7 @@ def scan_upload_archive(
 
   # helper to archive a failed file without compression
   def archive_as_failed(path_src, filename, source):
-    print('archiving as failed')
+    print('archiving as failed - '+source)
     path_dst = os.path.join(archive_failed_dir, source)
     compress = False
     file_archiver_impl.archive_file(path_src, path_dst, filename, compress)
@@ -74,29 +74,22 @@ def archive_as_successful(path_src, filename, source):
 
     (source, signal, time_type, geo_type, time_value, issue, lag) = details
 
-    # iterate over rows and upload to the database
-    all_rows_valid = True
-    for row_values in csv_importer_impl.load_csv(path, geo_type):
-      if not row_values:
-        all_rows_valid = False
-        continue
-
-      try:
-        database.insert_or_update(
-            source,
-            signal,
-            time_type,
-            geo_type,
-            time_value,
-            row_values.geo_value,
-            row_values.value,
-            row_values.stderr,
-            row_values.sample_size,
-            issue,
-            lag)
-      except Exception as e:
-        all_rows_valid = False
-        print('exception while inserting row:', e, row_values)
+    csv_rows = csv_importer_impl.load_csv(path, geo_type)
+
+    all_rows_valid = False
+    try:
+      cc_rows = CovidcastRow.fromCsvRows(csv_rows, source, signal, time_type, geo_type, time_value, issue, lag)
+      rows_list = list(cc_rows)
+      if not rows_list:
+        raise ValueError("No data")
+      result = database.insert_or_update_bulk(rows_list)
+      if result is None or result: # else would indicate zero rows inserted
+        database.commit()
+        all_rows_valid = True
+    except Exception as e:
+      all_rows_valid = False
+      print('exception while inserting rows:', e)
+      database.rollback()
 
     # archive the current file based on validation results
     if all_rows_valid:

src/acquisition/covidcast/database.py

@@ -10,6 +10,43 @@
 # first party
 import delphi.operations.secrets as secrets
 
+from math import ceil
+
+
+class CovidcastRow():
+  """A container for all the values of a single covidcast row."""
+
+  @staticmethod
+  def fromCsvRowValue(row_value, source, signal, time_type, geo_type, time_value, issue, lag):
+    return CovidcastRow(source, signal, time_type, geo_type, time_value,
+                        row_value.geo_value,
+                        row_value.value,
+                        row_value.stderr,
+                        row_value.sample_size,
+                        issue, lag)
+
+  @staticmethod
+  def fromCsvRows(row_values, source, signal, time_type, geo_type, time_value, issue, lag):
+    # NOTE: returns a generator, as row_values is expected to be a generator
+    return (CovidcastRow.fromCsvRowValue(row_value, source, signal, time_type, geo_type, time_value, issue, lag)
+            for row_value in row_values)
+
+  def __init__(self, source, signal, time_type, geo_type, time_value, geo_value, value, stderr, sample_size, issue, lag):
+    self.id = None
+    self.source = source
+    self.signal = signal
+    self.time_type = time_type
+    self.geo_type = geo_type
+    self.time_value = time_value
+    self.geo_value = geo_value      # from CSV row
+    self.value = value              # ...
+    self.stderr = stderr            # ...
+    self.sample_size = sample_size  # from CSV row
+    self.timestamp2 = 0
+    self.direction = None
+    self.issue = issue
+    self.lag = lag
+
 
 class Database:
   """A collection of covidcast database operations."""
@@ -27,6 +64,12 @@ def connect(self, connector_impl=mysql.connector):
         database=Database.DATABASE_NAME)
     self._cursor = self._connection.cursor()
 
+  def commit(self):
+    self._connection.commit()
+
+  def rollback(self):
+    self._connection.rollback()
+
   def disconnect(self, commit):
     """Close the database connection.
 
@@ -46,6 +89,69 @@ def count_all_rows(self):
     for (num,) in self._cursor:
       return num
 
+  def insert_or_update_bulk(self, cc_rows):
+    return self.insert_or_update_batch(cc_rows)
+
+  def insert_or_update_batch(self, cc_rows, batch_size=0, commit_partial=False):
+    """
+    Insert new rows (or update existing) in the `covidcast` table.
+
+    This has the intentional side effect of updating the primary timestamp.
+    """
+    sql = '''
+      INSERT INTO `covidcast`
+        (`id`, `source`, `signal`, `time_type`, `geo_type`, `time_value`, `geo_value`,
+        `timestamp1`, `value`, `stderr`, `sample_size`,
+        `timestamp2`, `direction`,
+        `issue`, `lag`)
+      VALUES
+        (0, %s, %s, %s, %s, %s, %s,
+        UNIX_TIMESTAMP(NOW()), %s, %s, %s,
+        0, NULL,
+        %s, %s)
+      ON DUPLICATE KEY UPDATE
+        `timestamp1` = VALUES(`timestamp1`),
+        `value` = VALUES(`value`),
+        `stderr` = VALUES(`stderr`),
+        `sample_size` = VALUES(`sample_size`)
+    '''
+    # TODO: ^ do we want to reset `timestamp2` and `direction` in the duplicate key case?
+
+    # TODO: consider handling cc_rows as a generator instead of a list
+    num_rows = len(cc_rows)
+    total = 0
+    if not batch_size:
+      batch_size = num_rows
+    num_batches = ceil(num_rows/batch_size)
+    for batch_num in range(num_batches):
+      start = batch_num * batch_size
+      end = min(num_rows, start + batch_size)
+      length = end - start
+
+      args = [(
+        row.source,
+        row.signal,
+        row.time_type,
+        row.geo_type,
+        row.time_value,
+        row.geo_value,
+        row.value,
+        row.stderr,
+        row.sample_size,
+        row.issue,
+        row.lag
+      ) for row in cc_rows[start:end]]
+
+      result = self._cursor.executemany(sql, args)
+      if result is None:
+        # the SQL connector does not support returning number of rows affected
+        total = None
+      else:
+        total += result
+      if commit_partial:
+        self._connection.commit()
+    return total
+
   def insert_or_update(
       self,
       source,
@@ -429,6 +535,7 @@ def update_timeseries_timestamp2(
     args = (source, signal, time_type, geo_type, geo_value)
     self._cursor.execute(sql, args)
 
+
   def get_covidcast_meta(self):
     """Compute and return metadata on all non-WIP COVIDcast signals."""
 
@@ -491,7 +598,7 @@ def get_covidcast_meta(self):
         t.`signal` ASC,
         t.`time_type` ASC,
         t.`geo_type` ASC
-'''
+    '''
     self._cursor.execute(sql)
     return list(dict(zip(self._cursor.column_names,x)) for x in self._cursor)
 

tests/acquisition/covidcast/test_csv_importer.py

@@ -57,19 +57,19 @@ def test_find_csv_files(self):
     mock_glob = MagicMock()
     mock_glob.glob.return_value = glob_paths
 
-    found = list(CsvImporter.find_csv_files(path_prefix, glob=mock_glob))
+    found = set(CsvImporter.find_csv_files(path_prefix, glob=mock_glob))
 
     expected_issue_day=int(date.today().strftime("%Y%m%d"))
     expected_issue_week=int(str(epi.Week.fromdate(date.today())))
     time_value_day = 20200408
-    expected = [
+    expected = set([
       (glob_paths[0], ('fb_survey', 'cli', 'week', 'county', 202015, expected_issue_week, delta_epiweeks(202015, expected_issue_week))),
       (glob_paths[1], ('ght', 'rawsearch', 'day', 'state', time_value_day, expected_issue_day, (date.today() - date(year=time_value_day // 10000, month=(time_value_day // 100) % 100, day=time_value_day % 100)).days)),
       (glob_paths[2], None),
       (glob_paths[3], None),
       (glob_paths[4], None),
       (glob_paths[5], None),
-    ]
+    ])
     self.assertEqual(found, expected)
 
   def test_is_header_valid_allows_extra_columns(self):

tests/acquisition/covidcast/test_csv_to_database.py

@@ -39,6 +39,8 @@ def load_csv_impl(path, *args):
         yield make_row('b1')
         yield None
         yield make_row('b3')
+      elif path == 'path/d.csv':
+        yield make_row('d1')
       else:
         # fail the test for any other path
         raise Exception('unexpected path')
@@ -47,10 +49,14 @@ def load_csv_impl(path, *args):
     mock_database = MagicMock()
     mock_csv_importer = MagicMock()
     mock_csv_importer.find_csv_files.return_value = [
+      # a good file
       ('path/a.csv', ('src_a', 'sig_a', 'day', 'hrr', 20200419, 20200420, 1)),
+      # a file with a data error
       ('path/b.csv', ('src_b', 'sig_b', 'week', 'msa', 202016, 202017, 1)),
       # emulate a file that's named incorrectly
       ('path/c.csv', None),
+      # another good file
+      ('path/d.csv', ('src_d', 'sig_d', 'week', 'msa', 202016, 202017, 1)),
     ]
     mock_csv_importer.load_csv = load_csv_impl
     mock_file_archiver = MagicMock()
@@ -61,27 +67,29 @@ def load_csv_impl(path, *args):
         csv_importer_impl=mock_csv_importer,
         file_archiver_impl=mock_file_archiver)
 
-    # verify that five rows were added to the database
-    self.assertEqual(mock_database.insert_or_update.call_count, 5)
-    call_args_list = mock_database.insert_or_update.call_args_list
-    actual_args = [args for (args, kwargs) in call_args_list]
+    # verify that appropriate rows were added to the database
+    self.assertEqual(mock_database.insert_or_update_bulk.call_count, 2)
+    call_args_list = mock_database.insert_or_update_bulk.call_args_list
+    actual_args = [[(a.source, a.signal, a.time_type, a.geo_type, a.time_value,
+                     a.geo_value, a.value, a.stderr, a.sample_size, a.issue, a.lag)
+                    for a in call.args[0]] for call in call_args_list]
     expected_args = [
-      ('src_a', 'sig_a', 'day', 'hrr', 20200419, 'a1', 'a1', 'a1', 'a1', 20200420, 1),
-      ('src_a', 'sig_a', 'day', 'hrr', 20200419, 'a2', 'a2', 'a2', 'a2', 20200420, 1),
-      ('src_a', 'sig_a', 'day', 'hrr', 20200419, 'a3', 'a3', 'a3', 'a3', 20200420, 1),
-      ('src_b', 'sig_b', 'week', 'msa', 202016, 'b1', 'b1', 'b1', 'b1', 202017, 1),
-      ('src_b', 'sig_b', 'week', 'msa', 202016, 'b3', 'b3', 'b3', 'b3', 202017, 1),
+      [('src_a', 'sig_a', 'day', 'hrr', 20200419, 'a1', 'a1', 'a1', 'a1', 20200420, 1),
+       ('src_a', 'sig_a', 'day', 'hrr', 20200419, 'a2', 'a2', 'a2', 'a2', 20200420, 1),
+       ('src_a', 'sig_a', 'day', 'hrr', 20200419, 'a3', 'a3', 'a3', 'a3', 20200420, 1)],
+      [('src_d', 'sig_d', 'week', 'msa', 202016, 'd1', 'd1', 'd1', 'd1', 202017, 1)]
     ]
     self.assertEqual(actual_args, expected_args)
 
-    # verify that one file was successful (a) and two failed (b, c)
-    self.assertEqual(mock_file_archiver.archive_file.call_count, 3)
+    # verify that two files were successful (a, d) and two failed (b, c)
+    self.assertEqual(mock_file_archiver.archive_file.call_count, 4)
     call_args_list = mock_file_archiver.archive_file.call_args_list
     actual_args = [args for (args, kwargs) in call_args_list]
     expected_args = [
       ('path', 'data_dir/archive/successful/src_a', 'a.csv', True),
       ('path', 'data_dir/archive/failed/src_b', 'b.csv', False),
       ('path', 'data_dir/archive/failed/unknown', 'c.csv', False),
+      ('path', 'data_dir/archive/successful/src_d', 'd.csv', True),
     ]
     self.assertEqual(actual_args, expected_args)
 
@@ -133,7 +141,7 @@ def test_database_exception_is_handled(self):
 
     data_dir = 'data_dir'
     mock_database = MagicMock()
-    mock_database.insert_or_update.side_effect = Exception('testing')
+    mock_database.insert_or_update_bulk.side_effect = Exception('testing')
     mock_csv_importer = MagicMock()
     mock_csv_importer.find_csv_files.return_value = [
       ('path/file.csv', ('src', 'sig', 'day', 'hrr', 20200423, 20200424, 1)),
@@ -149,8 +157,8 @@ def test_database_exception_is_handled(self):
         csv_importer_impl=mock_csv_importer,
         file_archiver_impl=mock_file_archiver)
 
-    # verify that a row insertion was attempted
-    self.assertTrue(mock_database.insert_or_update.called)
+    # verify that insertions were attempted
+    self.assertTrue(mock_database.insert_or_update_bulk.called)
 
     # verify that the file was archived as having failed
     self.assertTrue(mock_file_archiver.archive_file.called)