def wait_for_ingest_to_create_tasks(schema_arg: str) -> Tuple[str, HTTPStatus]:
"""Worker function to wait until ingest is not running to create_all_bq_refresh_tasks_for_schema.
When ingest is not running/locked, creates task to create_all_bq_refresh_tasks_for_schema.
When ingest is running/locked, re-enqueues this task to run again in 60 seconds.
"""
task_manager = BQRefreshCloudTaskManager()
lock_manager = GCSPseudoLockManager()
json_data_text = request.get_data(as_text=True)
try:
json_data = json.loads(json_data_text)
except (TypeError, json.decoder.JSONDecodeError):
json_data = {}
if "lock_id" not in json_data:
lock_id = str(uuid.uuid4())
else:
lock_id = json_data["lock_id"]
logging.info("Request lock id: %s", lock_id)
if not lock_manager.is_locked(
postgres_to_bq_lock_name_with_suffix(schema_arg)):
time = datetime.now().strftime("%m/%d/%Y, %H:%M:%S")
contents_as_json = {"time": time, "lock_id": lock_id}
contents = json.dumps(contents_as_json)
lock_manager.lock(postgres_to_bq_lock_name_with_suffix(schema_arg),
contents)
else:
contents = lock_manager.get_lock_contents(
postgres_to_bq_lock_name_with_suffix(schema_arg))
try:
contents_json = json.loads(contents)
except (TypeError, json.decoder.JSONDecodeError):
contents_json = {}
logging.info("Lock contents: %s", contents_json)
if lock_id != contents_json.get("lock_id"):
raise GCSPseudoLockAlreadyExists(
f"UUID {lock_id} does not match existing lock's UUID")
no_regions_running = lock_manager.no_active_locks_with_prefix(
GCS_TO_POSTGRES_INGEST_RUNNING_LOCK_NAME)
if not no_regions_running:
logging.info("Regions running, renqueuing this task.")
task_id = "{}-{}-{}".format("renqueue_wait_task",
str(datetime.utcnow().date()),
uuid.uuid4())
body = {"schema_type": schema_arg, "lock_id": lock_id}
task_manager.job_monitor_cloud_task_queue_manager.create_task(
task_id=task_id,
body=body,
relative_uri=
f"/cloud_sql_to_bq/create_refresh_bq_tasks/{schema_arg}",
schedule_delay_seconds=60,
)
return "", HTTPStatus.OK
logging.info("No regions running, calling create_refresh_bq_tasks")
create_all_bq_refresh_tasks_for_schema(schema_arg)
return "", HTTPStatus.OK
def test_using_lock(self) -> None:
lock_manager = GCSPseudoLockManager()
with lock_manager.using_lock(self.LOCK_NAME, self.CONTENTS):
self.assertTrue(lock_manager.is_locked(self.LOCK_NAME))
# Contents appropriately set
actual_contents = lock_manager.get_lock_contents(self.LOCK_NAME)
self.assertEqual(self.CONTENTS, actual_contents)
# lock should be unlocked outside of with
self.assertFalse(lock_manager.is_locked(self.LOCK_NAME))
def test_double_lock_diff_contents(self) -> None:
"""Locks and then locks again with unique contents, asserts its still locked and an error is raised"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(self.LOCK_NAME)
time = datetime.now().strftime("%m/%d/%Y, %H:%M:%S")
lock_id = str(uuid.uuid4())
contents_as_json = {"time": time, "uuid": lock_id}
contents = json.dumps(contents_as_json)
with self.assertRaises(GCSPseudoLockAlreadyExists):
lock_manager.lock(self.LOCK_NAME, contents)
self.assertTrue(lock_manager.is_locked(self.LOCK_NAME))
self.assertEqual(time, lock_manager.get_lock_contents(self.LOCK_NAME))
def test_get_lock_contents(self) -> None:
"""Tests that the get_lock_contents gets the correct contents from the lock"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(self.LOCK_NAME, self.CONTENTS)
actual_contents = lock_manager.get_lock_contents(self.LOCK_NAME)
self.assertEqual(self.CONTENTS, actual_contents)
