def test_lock_one_unlock_other(self) -> None:
"""Locks one lock and unlocks another, asserts both have correct status"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(self.LOCK_NAME)
with self.assertRaises(GCSPseudoLockDoesNotExist):
lock_manager.unlock(self.LOCK_NAME2)
self.assertTrue(lock_manager.is_locked(self.LOCK_NAME))
self.assertFalse(lock_manager.is_locked(self.LOCK_NAME2))
def test_locks_with_prefix_do_not_exist(self) -> None:
"""Ensures lock manager can see regions are not running"""
prefix = "SOME_LOCK_PREFIX"
lock_name = prefix + "some_suffix"
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(lock_name)
lock_manager.unlock(lock_name)
self.assertTrue(lock_manager.no_active_locks_with_prefix(prefix))
def test_lock_two_diff_unlock_one(self) -> None:
"""Locks two different locks, unlocks one, asserts both in correct place"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(self.LOCK_NAME)
lock_manager.lock(self.LOCK_NAME2)
lock_manager.unlock(self.LOCK_NAME)
self.assertFalse(lock_manager.is_locked(self.LOCK_NAME))
self.assertTrue(lock_manager.is_locked(self.LOCK_NAME2))
def test_double_unlock(self) -> None:
"""Unlocks and then unlocks gain, asserts its still unlocked and raises an error"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
self.assertFalse(lock_manager.is_locked(self.LOCK_NAME))
with self.assertRaises(GCSPseudoLockDoesNotExist):
lock_manager.unlock(self.LOCK_NAME)
with self.assertRaises(GCSPseudoLockDoesNotExist):
lock_manager.unlock(self.LOCK_NAME)
self.assertFalse(lock_manager.is_locked(self.LOCK_NAME))
def test_lock_unlock_two_diff(self) -> None:
"""Locks two different locks, unlocks both, asserts both unlocked"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(self.LOCK_NAME)
lock_manager.lock(self.LOCK_NAME2)
lock_manager.unlock(self.LOCK_NAME)
lock_manager.unlock(self.LOCK_NAME2)
self.assertFalse(lock_manager.is_locked(self.LOCK_NAME))
self.assertFalse(lock_manager.is_locked(self.LOCK_NAME2))
def test_unlock_expired(self) -> None:
lock_manager = GCSPseudoLockManager()
self._upload_fake_expired_lock(lock_manager, self.LOCK_NAME)
self.assertFalse(lock_manager.is_locked(self.LOCK_NAME))
# Should not raise an error
lock_manager.unlock(self.LOCK_NAME)
def test_unlock_delete_fails(self) -> None:
self.gcs_factory_patcher.stop()
self.gcs_factory_patcher.start().return_value = _FailingDeleteFs()
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(self.LOCK_NAME)
with self.assertRaises(GCSPseudoLockFailedUnlock):
lock_manager.unlock(self.LOCK_NAME)
def test_region_are_not_running(self) -> None:
"""Ensures lock manager can see regions are not running"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(GCS_TO_POSTGRES_INGEST_RUNNING_LOCK_NAME +
self.REGION.upper())
lock_manager.unlock(GCS_TO_POSTGRES_INGEST_RUNNING_LOCK_NAME +
self.REGION.upper())
self.assertTrue(
lock_manager.no_active_locks_with_prefix(
GCS_TO_POSTGRES_INGEST_RUNNING_LOCK_NAME))
def test_contents_of_unlocked_and_relocked(self) -> None:
"""Locks with pre-specified contents and asserts the lockfile contains those contents"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(self.LOCK_NAME, self.CONTENTS)
lock_manager.unlock(self.LOCK_NAME)
lock_manager.lock(self.LOCK_NAME, self.CONTENTS2)
path = GcsfsFilePath(bucket_name=lock_manager.bucket_name,
blob_name=self.LOCK_NAME)
actual_contents = self.fs.download_as_string(path)
self.assertEqual(self.CONTENTS2, actual_contents)
def test_lock_unlock_with_retry(self) -> None:
"""Locks then unlocks temp, checks if still locked"""
self.gcs_factory_patcher.stop()
self.gcs_factory_patcher.start(
).return_value = _MultipleAttemptDeleteFs()
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(self.LOCK_NAME)
lock_manager.unlock(self.LOCK_NAME)
self.assertFalse(lock_manager.is_locked(self.LOCK_NAME))
def monitor_refresh_bq_tasks() -> Tuple[str, int]:
"""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)
schema = data["schema"]
topic = data["topic"]
message = data["message"]
task_manager = BQRefreshCloudTaskManager()
# If any of the tasks in the queue have task_name containing schema, consider BQ tasks in queue
bq_tasks_in_queue = False
bq_task_list = task_manager.get_bq_queue_info().task_names
for task_name in bq_task_list:
task_id = task_name[task_name.find("/tasks/"):]
if schema in task_id:
bq_tasks_in_queue = True
# 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_refresh_monitor_task(schema, topic, message)
return "", HTTPStatus.OK
# Publish a message to the Pub/Sub topic once state BQ export is complete
if topic:
pubsub_helper.publish_message_to_topic(message=message, topic=topic)
# Unlock export lock when all BQ exports complete
lock_manager = GCSPseudoLockManager()
lock_manager.unlock(postgres_to_bq_lock_name_with_suffix(schema))
logging.info(
"Done running export for %s, unlocking Postgres to BigQuery export",
schema)
# Kick scheduler to restart ingest
kick_all_schedulers()
return ("", HTTPStatus.OK)
def test_lock_unlock(self) -> None:
"""Locks then unlocks temp, checks if still locked"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(self.LOCK_NAME)
lock_manager.unlock(self.LOCK_NAME)
self.assertFalse(lock_manager.is_locked(self.LOCK_NAME))
def test_unlock_never_locked(self) -> None:
"""Unlocks a lock that has never been locked before, check it raises an error and remains locked"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
with self.assertRaises(GCSPseudoLockDoesNotExist):
lock_manager.unlock(self.LOCK_NAME2)
self.assertFalse(lock_manager.is_locked(self.LOCK_NAME))
class DirectIngestRegionLockManager:
"""Manages acquiring and releasing the lock for the ingest process that writes
data to Postgres for a given region's ingest instance.
"""
def __init__(
self,
region_code: str,
ingest_instance: DirectIngestInstance,
blocking_locks: List[str],
) -> None:
"""
Args:
region_code: The region code for the region to lock / unlock ingest for.
blocking_locks: Any locks that, if present, mean ingest into Postgres
cannot proceed for this region.
"""
self.region_code = region_code
self.ingest_instance = ingest_instance
self.blocking_locks = blocking_locks
self.lock_manager = GCSPseudoLockManager()
def is_locked(self) -> bool:
"""Returns True if the ingest lock is held for the region associated with this
lock manager.
"""
return self.lock_manager.is_locked(self._ingest_lock_name_for_instance())
def can_proceed(self) -> bool:
"""Returns True if ingest can proceed for the region associated with this
lock manager.
"""
for lock in self.blocking_locks:
if self.lock_manager.is_locked(lock):
return False
return True
def acquire_lock(self) -> None:
self.lock_manager.lock(self._ingest_lock_name_for_instance())
def release_lock(self) -> None:
self.lock_manager.unlock(self._ingest_lock_name_for_instance())
@contextmanager
def using_region_lock(
self,
*,
expiration_in_seconds: int,
) -> Iterator[None]:
"""A context manager for acquiring the lock for a given region. Usage:
with lock_manager.using_region_lock(expiration_in_seconds=60):
... do work requiring the lock
"""
with self.lock_manager.using_lock(
self._ingest_lock_name_for_instance(),
expiration_in_seconds=expiration_in_seconds,
):
yield
@staticmethod
def for_state_ingest(
state_code: StateCode, ingest_instance: DirectIngestInstance
) -> "DirectIngestRegionLockManager":
return DirectIngestRegionLockManager.for_direct_ingest(
region_code=state_code.value,
ingest_instance=ingest_instance,
schema_type=SchemaType.STATE,
)
@staticmethod
def for_direct_ingest(
region_code: str,
ingest_instance: DirectIngestInstance,
schema_type: DirectIngestSchemaType,
) -> "DirectIngestRegionLockManager":
return DirectIngestRegionLockManager(
region_code=region_code,
ingest_instance=ingest_instance,
blocking_locks=[
postgres_to_bq_lock_name_for_schema(schema_type),
postgres_to_bq_lock_name_for_schema(SchemaType.OPERATIONS),
],
)
def _ingest_lock_name_for_instance(self) -> str:
if StateCode.is_state_code(self.region_code):
return (
STATE_GCS_TO_POSTGRES_INGEST_RUNNING_LOCK_PREFIX
+ self.region_code.upper()
+ f"_{self.ingest_instance.name}"
)
return (
JAILS_GCS_TO_POSTGRES_INGEST_RUNNING_LOCK_PREFIX + self.region_code.upper()
)
class CloudSqlToBQLockManager:
"""Manages acquiring and releasing the lock for the Cloud SQL -> BQ refresh, as well
as determining if the refresh can proceed given other ongoing processes.
"""
def __init__(self) -> None:
self.lock_manager = GCSPseudoLockManager()
def acquire_lock(self, lock_id: str, schema_type: SchemaType) -> None:
"""Acquires the CloudSQL -> BQ refresh lock for a given schema, or refreshes the
timeout of the lock if a lock with the given |lock_id| already exists. The
presence of the lock tells other ongoing processes to yield until the lock has
been released.
Acquiring the lock does NOT tell us if we can proceed with the refresh. You
must call can_proceed() to determine if all blocking processes have
successfully yielded.
Throws if a lock with a different lock_id exists for this schema.
"""
lock_name = postgres_to_bq_lock_name_for_schema(schema_type)
try:
self.lock_manager.lock(
lock_name,
payload=lock_id,
expiration_in_seconds=self._export_lock_timeout_for_schema(schema_type),
)
except GCSPseudoLockAlreadyExists as e:
previous_lock_id = self.lock_manager.get_lock_payload(lock_name)
logging.info("Lock contents: %s", previous_lock_id)
if lock_id != previous_lock_id:
raise GCSPseudoLockAlreadyExists(
f"UUID {lock_id} does not match existing lock's UUID {previous_lock_id}"
) from e
def can_proceed(self, schema_type: SchemaType) -> bool:
"""Returns True if all blocking processes have stopped and we can proceed with
the export, False otherwise.
"""
if not self.is_locked(schema_type):
raise GCSPseudoLockDoesNotExist(
f"Must acquire the lock for [{schema_type}] before checking if can proceed"
)
if schema_type not in (
SchemaType.STATE,
SchemaType.JAILS,
SchemaType.OPERATIONS,
):
return True
if schema_type == SchemaType.STATE:
blocking_lock_prefix = STATE_GCS_TO_POSTGRES_INGEST_RUNNING_LOCK_PREFIX
elif schema_type == SchemaType.JAILS:
blocking_lock_prefix = JAILS_GCS_TO_POSTGRES_INGEST_RUNNING_LOCK_PREFIX
elif schema_type == SchemaType.OPERATIONS:
# The operations export yields for all types of ingest
blocking_lock_prefix = GCS_TO_POSTGRES_INGEST_RUNNING_LOCK_PREFIX
else:
raise ValueError(f"Unexpected schema type [{schema_type}]")
no_blocking_locks = self.lock_manager.no_active_locks_with_prefix(
blocking_lock_prefix
)
return no_blocking_locks
def release_lock(self, schema_type: SchemaType) -> None:
"""Releases the CloudSQL -> BQ refresh lock for a given schema."""
self.lock_manager.unlock(postgres_to_bq_lock_name_for_schema(schema_type))
def is_locked(self, schema_type: SchemaType) -> bool:
return self.lock_manager.is_locked(
postgres_to_bq_lock_name_for_schema(schema_type)
)
@staticmethod
def _export_lock_timeout_for_schema(_schema_type: SchemaType) -> int:
"""Defines the exported lock timeouts permitted based on the schema arg.
For the moment all lock timeouts are set to one hour in length.
Export jobs may take longer than the alotted time, but if they do so, they
will de facto relinquish their hold on the acquired lock."""
return 3600
