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)