def test_disconnect_with_commit(self):
"""Disconnect from the database and commit."""
mock_connector = MagicMock()
database = Database()
database.connect(connector_impl=mock_connector)
# commit
database.disconnect(True)
connection = mock_connector.connect()
self.assertTrue(connection.commit.called)
self.assertTrue(connection.close.called)
def test_disconnect_with_rollback(self):
"""Disconnect from the database and rollback."""
mock_connector = MagicMock()
database = Database()
database.connect(connector_impl=mock_connector)
# rollback
database.disconnect(False)
connection = mock_connector.connect()
self.assertFalse(connection.commit.called)
self.assertTrue(connection.close.called)
def main():
response = Epidata.covidcast_meta()
print(response['result'])
if response['result'] == 1:
commit = False
database = Database()
database.connect()
try:
database.update_covidcast_meta_cache(
json.dumps(response['epidata']))
commit = True
print('successfully cached epidata')
finally:
database.disconnect(commit)
else:
print('metadata is not available')
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)