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