Merge pull request #418 from benjaminysmith/ingestion_logging

Retrieve the modified DB row count using cursor.rowcount

Author: krivard

src/acquisition/covidcast/database.py

@@ -212,17 +212,18 @@ def insert_or_update_batch(self, cc_rows, batch_size=2**20, commit_partial=False
         ) for row in cc_rows[start:end]]
 
 
-        result = self._cursor.executemany(insert_into_tmp_sql, args)
+        self._cursor.executemany(insert_into_tmp_sql, args)
         self._cursor.execute(insert_or_update_sql)
+        modified_row_count = self._cursor.rowcount
         self._cursor.execute(zero_is_latest_issue_sql)
         self._cursor.execute(set_is_latest_issue_sql)
         self._cursor.execute(truncate_tmp_table_sql)
 
-        if result is None:
-          # the SQL connector does not support returning number of rows affected
+        if modified_row_count is None or modified_row_count == -1:
+          # the SQL connector does not support returning number of rows affected (see PEP 249)
           total = None
         else:
-          total += result
+          total += modified_row_count
         if commit_partial:
           self._connection.commit()
     except Exception as e:

tests/acquisition/covidcast/test_database.py

@@ -111,3 +111,29 @@ def test_insert_or_update_batch_exception_reraised(self):
 
     cc_rows = {MagicMock(geo_id='CA', val=1, se=0, sample_size=0)}
     self.assertRaises(Exception, database.insert_or_update_batch, cc_rows)
+  
+  def test_insert_or_update_batch_row_count_returned(self):
+    """Test that the row count is returned"""
+    mock_connector = MagicMock()
+    database = Database()
+    database.connect(connector_impl=mock_connector)
+    connection = mock_connector.connect()
+    cursor = connection.cursor() 
+    cursor.rowcount = 3
+
+    cc_rows = [MagicMock(geo_id='CA', val=1, se=0, sample_size=0)]
+    result = database.insert_or_update_batch(cc_rows)
+    self.assertEqual(result, 3)
+
+  def test_insert_or_update_batch_none_returned(self):
+    """Test that None is returned when row count cannot be returned"""
+    mock_connector = MagicMock()
+    database = Database()
+    database.connect(connector_impl=mock_connector)
+    connection = mock_connector.connect()
+    cursor = connection.cursor() 
+    cursor.rowcount = -1
+
+    cc_rows = [MagicMock(geo_id='CA', val=1, se=0, sample_size=0)]
+    result = database.insert_or_update_batch(cc_rows)
+    self.assertIsNone(result)