def test_count_all_rows_query(self):
"""Query to count all rows looks sensible.
NOTE: Actual behavior is tested by integration test.
"""
mock_connector = MagicMock()
database = Database()
database.connect(connector_impl=mock_connector)
connection = mock_connector.connect()
cursor = connection.cursor()
cursor.__iter__.return_value = [(123, )]
num = database.count_all_rows()
self.assertEqual(num, 123)
self.assertTrue(cursor.execute.called)
sql = cursor.execute.call_args[0][0].lower()
self.assertIn('select count(1)', sql)
self.assertIn('from `covidcast`', sql)
def main(args):
"""Delete rows from covidcast."""
logger = get_structured_logger("csv_deletion", filename=args.log_file)
start_time = time.time()
database = Database()
database.connect()
all_n = 0
try:
for deletion_file in sorted(
glob.glob(os.path.join(args.deletion_dir, '*.csv'))):
n = handle_file(deletion_file, database, logger)
if n is not None:
all_n += n
else:
all_n = "rowcount unsupported"
finally:
database.disconnect(True)
logger.info("Deleted CSVs from database",
total_runtime_in_seconds=round(time.time() - start_time, 2),
row_count=all_n)