def create_all_operations_bq_export_tasks():
"""Creates an export task for each table to be exported.
A task is created for each table defined in export_config.OPERATIONS_TABLES_TO_EXPORT.
Re-creates all tasks if any task fails to be created.
"""
schema_type_str = SchemaType.OPERATIONS.value
logging.info("Beginning BQ export for operations schema tables.")
task_manager = BQExportCloudTaskManager()
for table in export_config.OPERATIONS_TABLES_TO_EXPORT:
task_manager.create_bq_task(table.name, schema_type_str)
return ('', HTTPStatus.OK)
def handle_bq_monitor_task():
"""Worker function to publish a message to a Pub/Sub topic once all tasks in
the BIGQUERY_QUEUE queue have completed.
"""
json_data = request.get_data(as_text=True)
data = json.loads(json_data)
topic = data['topic']
message = data['message']
task_manager = BQExportCloudTaskManager()
bq_tasks_in_queue = task_manager.get_bq_queue_info().size() > 0
# If there are BQ tasks in the queue, then re-queue this task in a minute
if bq_tasks_in_queue:
logging.info("Tasks still in bigquery queue. Re-queuing bq monitor"
" task.")
task_manager.create_bq_monitor_task(topic, message)
return ('', HTTPStatus.OK)
# Publish a message to the Pub/Sub topic once all BQ exports are complete
pubsub_helper.publish_message_to_topic(message=message, topic=topic)
return ('', HTTPStatus.OK)
def test_create_bq_monitor_test(self, mock_client, mock_uuid):
# Arrange
delay_sec = 60
now_utc_timestamp = int(datetime.datetime.now().timestamp())
uuid = 'random-uuid'
mock_uuid.uuid4.return_value = uuid
project_id = 'recidiviz-456'
topic = 'fake.topic'
message = 'A fake message'
queue_path = f'queue_path/{project_id}/{QUEUES_REGION}'
task_id = 'fake-topic-2019-04-13-random-uuid'
task_path = f'{queue_path}/{task_id}'
body = {
'topic': topic,
'message': message,
}
task = tasks_v2.types.task_pb2.Task(
name=task_path,
schedule_time=timestamp_pb2.Timestamp(seconds=(now_utc_timestamp +
delay_sec)),
app_engine_http_request={
'http_method': 'POST',
'relative_uri': '/export_manager/bq_monitor',
'body': json.dumps(body).encode()
})
mock_client.return_value.task_path.return_value = task_path
mock_client.return_value.queue_path.return_value = queue_path
# Act
BQExportCloudTaskManager(project_id=project_id). \
create_bq_monitor_task(topic=topic, message=message)
# Assert
mock_client.return_value.queue_path.assert_called_with(
project_id, QUEUES_REGION, JOB_MONITOR_QUEUE_V2)
mock_client.return_value.task_path.assert_called_with(
project_id, QUEUES_REGION, JOB_MONITOR_QUEUE_V2, task_id)
mock_client.return_value.create_task.assert_called_with(
queue_path, task)
def test_create_bq_test(self, mock_client, mock_uuid):
# Arrange
uuid = 'random-uuid'
mock_uuid.uuid4.return_value = uuid
project_id = 'recidiviz-456'
table_name = 'test_table'
schema_type = 'test_schema_type'
queue_path = f'queue_path/{project_id}/{QUEUES_REGION}'
task_id = 'test_table-test_schema_type-2019-04-12-random-uuid'
task_path = f'{queue_path}/{task_id}'
body = {'table_name': table_name, 'schema_type': schema_type}
task = tasks_v2.types.task_pb2.Task(name=task_path,
app_engine_http_request={
'http_method': 'POST',
'relative_uri':
'/export_manager/export',
'body':
json.dumps(body).encode()
})
mock_client.return_value.task_path.return_value = task_path
mock_client.return_value.queue_path.return_value = queue_path
# Act
BQExportCloudTaskManager(project_id=project_id). \
create_bq_task(table_name=table_name, schema_type=schema_type)
# Assert
mock_client.return_value.queue_path.assert_called_with(
project_id, QUEUES_REGION, BIGQUERY_QUEUE_V2)
mock_client.return_value.task_path.assert_called_with(
project_id, QUEUES_REGION, BIGQUERY_QUEUE_V2, task_id)
mock_client.return_value.create_task.assert_called_with(
queue_path, task)
def create_all_state_bq_export_tasks():
"""Creates an export task for each table to be exported.
A task is created for each table defined in export_config.STATE_TABLES_TO_EXPORT.
Re-creates all tasks if any task fails to be created.
"""
schema_type_str = SchemaType.STATE.value
logging.info("Beginning BQ export for state schema tables.")
task_manager = BQExportCloudTaskManager()
for table in export_config.STATE_TABLES_TO_EXPORT:
task_manager.create_bq_task(table.name, schema_type_str)
pub_sub_topic = 'v1.calculator.recidivism'
pub_sub_message = 'State export to BQ complete'
task_manager.create_bq_monitor_task(pub_sub_topic, pub_sub_message)
return ('', HTTPStatus.OK)