def test_create_wait_to_refresh_bq_tasks_state_ingest_locked(
self, mock_task_manager):
# Arrange
mock_table = Mock()
mock_table.name = "test_table"
self.mock_bq_refresh_config.for_schema_type.return_value.get_tables_to_export.return_value = [
mock_table
]
lock_manager = GCSPseudoLockManager()
lock_manager.lock(GCS_TO_POSTGRES_INGEST_RUNNING_LOCK_NAME)
# Act
response = self.mock_flask_client.get(
"/create_refresh_bq_tasks/state",
headers={"X-Appengine-Inbound-Appid": "recidiviz-123"},
)
# Assert
self.assertEqual(response.status_code, HTTPStatus.OK)
self.assertFalse(
lock_manager.no_active_locks_with_prefix(
POSTGRES_TO_BQ_EXPORT_RUNNING_LOCK_NAME))
self.assertTrue(
mock_task_manager.return_value.
job_monitor_cloud_task_queue_manager.create_task.called)
mock_task_manager.return_value.create_refresh_bq_table_task.assert_not_called(
)
mock_task_manager.return_value.create_bq_refresh_monitor_task.assert_not_called(
)
def test_create_refresh_bq_tasks_state(self, mock_task_manager):
# Arrange
mock_table = Mock()
mock_table.name = "test_table"
self.mock_bq_refresh_config.for_schema_type.return_value.get_tables_to_export.return_value = [
mock_table
]
lock_manager = GCSPseudoLockManager()
# Act
response = self.mock_flask_client.get(
"/create_refresh_bq_tasks/state",
headers={"X-Appengine-Inbound-Appid": "recidiviz-123"},
)
# Assert
self.assertFalse(
lock_manager.no_active_locks_with_prefix(
POSTGRES_TO_BQ_EXPORT_RUNNING_LOCK_NAME))
self.assertEqual(response.status_code, HTTPStatus.OK)
self.mock_bq_refresh_config.for_schema_type.assert_called_with(
SchemaType.STATE)
mock_task_manager.return_value.create_refresh_bq_table_task.assert_called_with(
"test_table", SchemaType.STATE)
mock_task_manager.return_value.create_bq_refresh_monitor_task.assert_called_with(
SchemaType.STATE.value, "v1.calculator.trigger_daily_pipelines",
ANY)
def test_region_are_running(self) -> None:
"""Ensures lock manager can see regions are running"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(GCS_TO_POSTGRES_INGEST_RUNNING_LOCK_NAME +
self.REGION.upper())
self.assertFalse(
lock_manager.no_active_locks_with_prefix(
GCS_TO_POSTGRES_INGEST_RUNNING_LOCK_NAME))
def test_locks_with_prefix_ignore_expired(self) -> None:
"""Ensures lock manager can see if locks with a prefix exist."""
prefix = "SOME_LOCK_PREFIX"
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
self._upload_fake_expired_lock(lock_manager, prefix + "some_suffix")
self.assertTrue(lock_manager.no_active_locks_with_prefix(prefix))
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 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_unlock_empty_locks_with_prefix(self) -> None:
"""Tests that nonexistent locks with prefix, asserts error raised"""
lock_manager = GCSPseudoLockManager()
with self.assertRaises(GCSPseudoLockDoesNotExist):
lock_manager.unlock_locks_with_prefix(self.PREFIX)
self.assertTrue(lock_manager.no_active_locks_with_prefix(self.PREFIX))
def test_locks_with_prefix_exist(self) -> None:
"""Ensures lock manager can see if locks with a prefix exist."""
prefix = "SOME_LOCK_PREFIX"
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(prefix + "some_suffix")
self.assertFalse(lock_manager.no_active_locks_with_prefix(prefix))
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
