def test_lock_two_diff(self) -> None:
"""Locks two different locks, asserts both locked"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(self.LOCK_NAME)
lock_manager.lock(self.LOCK_NAME2)
self.assertTrue(lock_manager.is_locked(self.LOCK_NAME))
self.assertTrue(lock_manager.is_locked(self.LOCK_NAME2))
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_unlock_locks_with_prefix(self) -> None:
"""Tests that all locks with prefix are unlocked"""
lock_manager = GCSPseudoLockManager()
lock_manager.lock(self.PREFIX + self.LOCK_NAME)
lock_manager.lock(self.PREFIX + self.LOCK_NAME2)
lock_manager.unlock_locks_with_prefix(self.PREFIX)
self.assertFalse(lock_manager.is_locked(self.PREFIX + self.LOCK_NAME))
self.assertFalse(lock_manager.is_locked(self.PREFIX + self.LOCK_NAME2))
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_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(self) -> None:
"""Locks and then locks again, asserts its still locked and an error is raised"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(self.LOCK_NAME)
with self.assertRaises(GCSPseudoLockAlreadyExists):
lock_manager.lock(self.LOCK_NAME)
self.assertTrue(lock_manager.is_locked(self.LOCK_NAME))
def test_raise_from_using_lock(self) -> None:
lock_manager = GCSPseudoLockManager()
with self.assertRaises(ValueError):
with lock_manager.using_lock(self.LOCK_NAME, self.CONTENTS):
raise ValueError
# lock should be unlocked outside of with
self.assertFalse(lock_manager.is_locked(self.LOCK_NAME))
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 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_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_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 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_lock_expiration_not_met(self) -> None:
now = datetime.now()
lock_manager = GCSPseudoLockManager()
path = GcsfsFilePath(bucket_name=lock_manager.bucket_name,
blob_name=self.LOCK_NAME)
self.fs.upload_from_string(
path,
json.dumps(
GCSPseudoLockBody(lock_time=now,
expiration_in_seconds=60).to_json(),
default=str,
),
content_type="text/text",
)
self.assertTrue(lock_manager.is_locked(self.LOCK_NAME))
def test_monitor_refresh_bq_tasks_requeue_with_topic_and_message(
self,
mock_task_manager: mock.MagicMock,
mock_pubsub_helper: mock.MagicMock,
mock_supported_region_codes: mock.MagicMock,
) -> None:
"""Test that a new bq monitor task is added to the queue when there are
still unfinished tasks on the bq queue, with topic/message to publish."""
queue_path = "test-queue-path"
lock_manager = GCSPseudoLockManager()
mock_supported_region_codes.return_value = []
schema = "schema"
topic = "fake_topic"
message = "fake_message"
route = "/monitor_refresh_bq_tasks"
data = {"schema": schema, "topic": topic, "message": message}
lock_manager.lock(POSTGRES_TO_BQ_EXPORT_RUNNING_LOCK_NAME +
schema.upper())
mock_task_manager.return_value.get_bq_queue_info.return_value = (
CloudTaskQueueInfo(
queue_name="queue_name",
task_names=[
f"{queue_path}/tasks/table_name-123-{schema}",
f"{queue_path}/tasks/table_name-456-{schema}",
f"{queue_path}/tasks/table_name-789-{schema}",
],
))
response = self.mock_flask_client.post(
route,
data=json.dumps(data),
content_type="application/json",
headers={"X-Appengine-Inbound-Appid": "recidiviz-123"},
)
self.assertEqual(response.status_code, HTTPStatus.OK)
mock_task_manager.return_value.create_bq_refresh_monitor_task.assert_called_with(
schema, topic, message)
mock_pubsub_helper.publish_message_to_topic.assert_not_called()
self.assertTrue(
lock_manager.is_locked(POSTGRES_TO_BQ_EXPORT_RUNNING_LOCK_NAME +
schema.upper()))
def test_monitor_refresh_bq_tasks_requeue_unlock_no_publish(
self,
mock_task_manager: mock.MagicMock,
mock_pubsub_helper: mock.MagicMock,
mock_supported_region_codes: mock.MagicMock,
) -> None:
"""Test that a bq monitor task does not publish topic/message
with empty topic/message and that it unlocks export lock"""
lock_manager = GCSPseudoLockManager()
mock_supported_region_codes.return_value = []
schema = "schema"
topic = ""
message = ""
route = "/monitor_refresh_bq_tasks"
data = {"schema": schema, "topic": topic, "message": message}
lock_manager.lock(POSTGRES_TO_BQ_EXPORT_RUNNING_LOCK_NAME +
schema.upper())
mock_task_manager.return_value.get_bq_queue_info.return_value = (
CloudTaskQueueInfo(queue_name="queue_name", task_names=[]))
response = self.mock_flask_client.post(
route,
data=json.dumps(data),
content_type="application/json",
headers={"X-Appengine-Inbound-Appid": "recidiviz-123"},
)
self.assertEqual(response.status_code, HTTPStatus.OK)
mock_task_manager.return_value.create_bq_refresh_monitor_task.assert_not_called(
)
mock_pubsub_helper.publish_message_to_topic.assert_not_called()
self.assertFalse(
lock_manager.is_locked(POSTGRES_TO_BQ_EXPORT_RUNNING_LOCK_NAME +
schema.upper()))
class BaseDirectIngestController(Ingestor, Generic[IngestArgsType, ContentsHandleType]):
"""Parses and persists individual-level info from direct ingest partners."""
def __init__(self, region_name: str, system_level: SystemLevel):
"""Initialize the controller.
Args:
region_name: (str) the name of the region to be collected.
"""
self.region = regions.get_region(region_name, is_direct_ingest=True)
self.system_level = system_level
self.cloud_task_manager = DirectIngestCloudTaskManagerImpl()
self.lock_manager = GCSPseudoLockManager()
# ============== #
# JOB SCHEDULING #
# ============== #
def kick_scheduler(self, just_finished_job: bool) -> None:
logging.info("Creating cloud task to schedule next job.")
self.cloud_task_manager.create_direct_ingest_scheduler_queue_task(
region=self.region, just_finished_job=just_finished_job, delay_sec=0
)
def schedule_next_ingest_job_or_wait_if_necessary(
self, just_finished_job: bool
) -> None:
"""Creates a cloud task to run the next ingest job. Depending on the
next job's IngestArgs, we either post a task to direct/scheduler/ if
a wait_time is specified or direct/process_job/ if we can run the next
job immediately."""
check_is_region_launched_in_env(self.region)
if self._schedule_any_pre_ingest_tasks():
logging.info("Found pre-ingest tasks to schedule - returning.")
return
if self.lock_manager.is_locked(self.ingest_process_lock_for_region()):
logging.info("Direct ingest is already locked on region [%s]", self.region)
return
process_job_queue_info = self.cloud_task_manager.get_process_job_queue_info(
self.region
)
if process_job_queue_info.size() and not just_finished_job:
logging.info(
"Already running job [%s] - will not schedule another job for "
"region [%s]",
process_job_queue_info.task_names[0],
self.region.region_code,
)
return
next_job_args = self._get_next_job_args()
if not next_job_args:
logging.info(
"No more jobs to run for region [%s] - returning",
self.region.region_code,
)
return
if process_job_queue_info.is_task_queued(self.region, next_job_args):
logging.info(
"Already have task queued for next job [%s] - returning.",
self._job_tag(next_job_args),
)
return
if self.lock_manager.is_locked(
postgres_to_bq_lock_name_for_schema(
schema_type_for_system_level(self.system_level)
)
) or self.lock_manager.is_locked(
postgres_to_bq_lock_name_for_schema(SchemaType.OPERATIONS)
):
logging.info(
"Postgres to BigQuery export is running, cannot run ingest - returning"
)
return
# TODO(#3020): Add similar logic between the raw data BQ import and ingest view export tasks
# TODO(#3162): Delete this wait logic from here once all regions have been transitioned to a SQL
# preprocessing model.
wait_time_sec = self._wait_time_sec_for_next_args(next_job_args)
logging.info(
"Found next ingest job to run [%s] with wait time [%s].",
self._job_tag(next_job_args),
wait_time_sec,
)
if wait_time_sec:
scheduler_queue_info = self.cloud_task_manager.get_scheduler_queue_info(
self.region
)
if scheduler_queue_info.size() <= 1:
logging.info(
"Creating cloud task to fire timer in [%s] seconds", wait_time_sec
)
self.cloud_task_manager.create_direct_ingest_scheduler_queue_task(
region=self.region, just_finished_job=False, delay_sec=wait_time_sec
)
else:
logging.info(
"[%s] tasks already in the scheduler queue for region "
"[%s] - not queueing another task.",
str(scheduler_queue_info.size),
self.region.region_code,
)
else:
logging.info(
"Creating cloud task to run job [%s]", self._job_tag(next_job_args)
)
self.cloud_task_manager.create_direct_ingest_process_job_task(
region=self.region, ingest_args=next_job_args
)
self._on_job_scheduled(next_job_args)
def _schedule_any_pre_ingest_tasks(self) -> bool:
"""Schedules any tasks related to SQL preprocessing of new files in preparation for ingest of those files into
our Postgres database.
Returns True if any jobs were scheduled or if there were already any pre-ingest jobs scheduled. Returns False if
there are no remaining ingest jobs to schedule and it is safe to proceed with ingest.
"""
if self._schedule_raw_data_import_tasks():
logging.info("Found pre-ingest raw data import tasks to schedule.")
return True
# TODO(#3020): We have logic to ensure that we wait 10 min for all files to upload properly before moving on to
# ingest. We probably actually need this to happen between raw data import and ingest view export steps - if we
# haven't seen all files yet and most recent raw data file came in sometime in the last 10 min, we should wait
# to do view exports.
if self._schedule_ingest_view_export_tasks():
logging.info("Found pre-ingest view export tasks to schedule.")
return True
return False
def _schedule_raw_data_import_tasks(self) -> bool:
return False
def _schedule_ingest_view_export_tasks(self) -> bool:
return False
@abc.abstractmethod
def _get_next_job_args(self) -> Optional[IngestArgsType]:
"""Should be overridden to return args for the next ingest job, or
None if there is nothing to process."""
@abc.abstractmethod
def _on_job_scheduled(self, ingest_args: IngestArgsType) -> None:
"""Called from the scheduler queue when an individual direct ingest job
is scheduled.
"""
def _wait_time_sec_for_next_args(self, _: IngestArgsType) -> int:
"""Should be overwritten to return the number of seconds we should
wait to try to run another job, given the args of the next job we would
run if we had to run a job right now. This gives controllers the ability
to back off if we want to attempt to enforce an ingest order for files
that might all be uploaded around the same time, but in inconsistent
order.
"""
return 0
# =================== #
# SINGLE JOB RUN CODE #
# =================== #
def run_ingest_job_and_kick_scheduler_on_completion(
self, args: IngestArgsType
) -> None:
check_is_region_launched_in_env(self.region)
if self.lock_manager.is_locked(
postgres_to_bq_lock_name_for_schema(
schema_type_for_system_level(self.system_level)
)
) or self.lock_manager.is_locked(
postgres_to_bq_lock_name_for_schema(SchemaType.OPERATIONS)
):
raise GCSPseudoLockAlreadyExists(
"Postgres to BigQuery export is running, can not run ingest"
)
with self.lock_manager.using_lock(self.ingest_process_lock_for_region()):
should_schedule = self._run_ingest_job(args)
if should_schedule:
self.kick_scheduler(just_finished_job=True)
logging.info("Done running task. Returning.")
def _run_ingest_job(self, args: IngestArgsType) -> bool:
"""
Runs the full ingest process for this controller - reading and parsing
raw input data, transforming it to our schema, then writing to the
database.
Returns:
True if we should try to schedule the next job on completion. False,
otherwise.
"""
check_is_region_launched_in_env(self.region)
start_time = datetime.datetime.now()
logging.info("Starting ingest for ingest run [%s]", self._job_tag(args))
contents_handle = self._get_contents_handle(args)
if contents_handle is None:
logging.warning(
"Failed to get contents handle for ingest run [%s] - " "returning.",
self._job_tag(args),
)
# If the file no-longer exists, we do want to kick the scheduler
# again to pick up the next file to run. We expect this to happen
# occasionally as a race when the scheduler picks up a file before
# it has been properly moved.
return True
if not self._can_proceed_with_ingest_for_contents(args, contents_handle):
logging.warning(
"Cannot proceed with contents for ingest run [%s] - returning.",
self._job_tag(args),
)
# If we get here, we've failed to properly split a file picked up
# by the scheduler. We don't want to schedule a new job after
# returning here, otherwise we'll get ourselves in a loop where we
# continually try to schedule this file.
return False
logging.info(
"Successfully read contents for ingest run [%s]", self._job_tag(args)
)
if not self._are_contents_empty(args, contents_handle):
self._parse_and_persist_contents(args, contents_handle)
else:
logging.warning(
"Contents are empty for ingest run [%s] - skipping parse and "
"persist steps.",
self._job_tag(args),
)
self._do_cleanup(args)
duration_sec = (datetime.datetime.now() - start_time).total_seconds()
logging.info(
"Finished ingest in [%s] sec for ingest run [%s].",
str(duration_sec),
self._job_tag(args),
)
return True
@trace.span
def _parse_and_persist_contents(
self, args: IngestArgsType, contents_handle: ContentsHandleType
) -> None:
"""
Runs the full ingest process for this controller for files with
non-empty contents.
"""
ingest_info = self._parse(args, contents_handle)
if not ingest_info:
raise DirectIngestError(
error_type=DirectIngestErrorType.PARSE_ERROR,
msg="No IngestInfo after parse.",
)
logging.info(
"Successfully parsed data for ingest run [%s]", self._job_tag(args)
)
ingest_info_proto = ingest_utils.convert_ingest_info_to_proto(ingest_info)
logging.info(
"Successfully converted ingest_info to proto for ingest " "run [%s]",
self._job_tag(args),
)
ingest_metadata = self._get_ingest_metadata(args)
persist_success = persistence.write(ingest_info_proto, ingest_metadata)
if not persist_success:
raise DirectIngestError(
error_type=DirectIngestErrorType.PERSISTENCE_ERROR,
msg="Persist step failed",
)
logging.info("Successfully persisted for ingest run [%s]", self._job_tag(args))
def _get_ingest_metadata(self, args: IngestArgsType) -> IngestMetadata:
return IngestMetadata(
self.region.region_code,
self.region.jurisdiction_id,
args.ingest_time,
self.get_enum_overrides(),
self.system_level,
)
def ingest_process_lock_for_region(self) -> str:
return (
GCS_TO_POSTGRES_INGEST_RUNNING_LOCK_NAME + self.region.region_code.upper()
)
@abc.abstractmethod
def _job_tag(self, args: IngestArgsType) -> str:
"""Should be overwritten to return a (short) string tag to identify an
ingest run in logs.
"""
@abc.abstractmethod
def _get_contents_handle(
self, args: IngestArgsType
) -> Optional[ContentsHandleType]:
"""Should be overridden by subclasses to return a handle to the contents
that can return an iterator over the contents and also manages cleanup
of resources once we are done with the contents.
Will return None if the contents could not be read (i.e. if they no
longer exist).
"""
@abc.abstractmethod
def _are_contents_empty(
self, args: IngestArgsType, contents_handle: ContentsHandleType
) -> bool:
"""Should be overridden by subclasses to return True if the contents
for the given args should be considered "empty" and not parsed. For
example, a CSV might have a single header line but no actual data.
"""
@abc.abstractmethod
def _parse(
self, args: IngestArgsType, contents_handle: ContentsHandleType
) -> IngestInfo:
"""Should be overridden by subclasses to parse raw ingested contents
into an IngestInfo object.
"""
@abc.abstractmethod
def _do_cleanup(self, args: IngestArgsType) -> None:
"""Should be overridden by subclasses to do any necessary cleanup in the
ingested source once contents have been successfully persisted.
"""
@abc.abstractmethod
def _can_proceed_with_ingest_for_contents(
self, args: IngestArgsType, contents_handle: ContentsHandleType
) -> bool:
"""Given a pointer to the contents, can the controller continue with
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()
)
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_lock(self) -> None:
"""Locks temp and then checks if locked"""
lock_manager = GCSPseudoLockManager(self.PROJECT_ID)
lock_manager.lock(self.LOCK_NAME)
self.assertTrue(lock_manager.is_locked(self.LOCK_NAME))
def test_lock_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))
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
