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, payload=self.CONTENTS)
with self.assertRaises(GCSPseudoLockAlreadyExists):
lock_manager.lock(self.LOCK_NAME, self.CONTENTS2)
self.assertTrue(lock_manager.is_locked(self.LOCK_NAME))
self.assertEqual(self.CONTENTS,
lock_manager.get_lock_payload(self.LOCK_NAME))
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_body = lock_manager.get_lock_payload(self.LOCK_NAME)
self.assertEqual(self.CONTENTS, actual_body)
# lock should be unlocked outside of with
self.assertFalse(lock_manager.is_locked(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_body = lock_manager.get_lock_payload(self.LOCK_NAME)
self.assertEqual(self.CONTENTS, actual_body)
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
