def _acquire_ingest_lock() -> Tuple[str, HTTPStatus]:
try:
state_code = StateCode(request.json["stateCode"])
ingest_instance = DirectIngestInstance(
request.json["ingestInstance"])
except ValueError:
return "invalid parameters provided", HTTPStatus.BAD_REQUEST
lock_manager = DirectIngestRegionLockManager.for_state_ingest(
state_code, ingest_instance=ingest_instance)
try:
lock_manager.acquire_lock()
except GCSPseudoLockAlreadyExists:
return "lock already exists", HTTPStatus.CONFLICT
if not lock_manager.can_proceed():
try:
lock_manager.release_lock()
except Exception as e:
logging.exception(e)
return (
"other locks blocking ingest have been acquired; releasing lock",
HTTPStatus.CONFLICT,
)
return "", HTTPStatus.OK
def scheduler() -> Tuple[str, HTTPStatus]:
"""Checks the state of the ingest instance and schedules any tasks to be run."""
logging.info("Received request for direct ingest scheduler: %s",
request.values)
region_code = get_str_param_value("region", request.values)
just_finished_job = get_bool_param_value("just_finished_job",
request.values,
default=False)
# The bucket name for ingest instance to schedule work out of
bucket = get_str_param_value("bucket", request.args)
if not region_code or just_finished_job is None or not bucket:
response = f"Bad parameters [{request.values}]"
logging.error(response)
return response, HTTPStatus.BAD_REQUEST
bucket_path = GcsfsBucketPath(bucket)
with monitoring.push_region_tag(
region_code,
ingest_instance=DirectIngestInstance.for_ingest_bucket(
bucket_path).value,
):
try:
controller = DirectIngestControllerFactory.build(
ingest_bucket_path=bucket_path, allow_unlaunched=False)
except DirectIngestError as e:
if e.is_bad_request():
logging.error(str(e))
return str(e), HTTPStatus.BAD_REQUEST
raise e
controller.schedule_next_ingest_job(just_finished_job)
return "", HTTPStatus.OK
def _get_database_name_for_state(state_code: StateCode,
instance: DirectIngestInstance) -> str:
"""Returns the database name for the given state and instance"""
return SQLAlchemyDatabaseKey.for_state_code(
state_code,
instance.database_version(SystemLevel.STATE,
state_code=state_code),
).db_name
def __init__(self, ingest_bucket_path: GcsfsBucketPath) -> None:
"""Initialize the controller."""
self.cloud_task_manager = DirectIngestCloudTaskManagerImpl()
self.ingest_instance = DirectIngestInstance.for_ingest_bucket(
ingest_bucket_path)
self.region_lock_manager = DirectIngestRegionLockManager.for_direct_ingest(
region_code=self.region.region_code,
schema_type=self.system_level.schema_type(),
ingest_instance=self.ingest_instance,
)
self.fs = DirectIngestGCSFileSystem(GcsfsFactory.build())
self.ingest_bucket_path = ingest_bucket_path
self.storage_directory_path = (
gcsfs_direct_ingest_storage_directory_path_for_region(
region_code=self.region_code(),
system_level=self.system_level,
ingest_instance=self.ingest_instance,
))
self.temp_output_directory_path = (
gcsfs_direct_ingest_temporary_output_directory_path())
self.file_prioritizer = GcsfsDirectIngestJobPrioritizer(
self.fs,
self.ingest_bucket_path,
self.get_file_tag_rank_list(),
)
self.ingest_file_split_line_limit = self._INGEST_FILE_SPLIT_LINE_LIMIT
self.file_metadata_manager = PostgresDirectIngestFileMetadataManager(
region_code=self.region.region_code,
ingest_database_name=self.ingest_database_key.db_name,
)
self.raw_file_import_manager = DirectIngestRawFileImportManager(
region=self.region,
fs=self.fs,
ingest_bucket_path=self.ingest_bucket_path,
temp_output_directory_path=self.temp_output_directory_path,
big_query_client=BigQueryClientImpl(),
)
self.ingest_view_export_manager = DirectIngestIngestViewExportManager(
region=self.region,
fs=self.fs,
output_bucket_name=self.ingest_bucket_path.bucket_name,
file_metadata_manager=self.file_metadata_manager,
big_query_client=BigQueryClientImpl(),
view_collector=DirectIngestPreProcessedIngestViewCollector(
self.region, self.get_file_tag_rank_list()),
launched_file_tags=self.get_file_tag_rank_list(),
)
self.ingest_instance_status_manager = DirectIngestInstanceStatusManager(
self.region_code(), self.ingest_instance)
def _export_database_to_gcs() -> Tuple[str, HTTPStatus]:
try:
state_code = StateCode(request.json["stateCode"])
ingest_instance = DirectIngestInstance(
request.json["ingestInstance"].upper())
db_version = ingest_instance.database_version(
system_level=SystemLevel.STATE, state_code=state_code)
except ValueError:
return "invalid parameters provided", HTTPStatus.BAD_REQUEST
lock_manager = DirectIngestRegionLockManager.for_state_ingest(
state_code, ingest_instance)
if not lock_manager.can_proceed():
return (
"other locks blocking ingest have been acquired; aborting operation",
HTTPStatus.CONFLICT,
)
db_key = SQLAlchemyDatabaseKey.for_state_code(state_code, db_version)
cloud_sql_client = CloudSQLClientImpl(project_id=project_id)
operation_id = cloud_sql_client.export_to_gcs_sql(
db_key,
GcsfsFilePath.from_absolute_path(
f"{STATE_INGEST_EXPORT_URI}/{db_version.value}/{state_code.value}"
),
)
if operation_id is None:
return (
"Cloud SQL export operation was not started successfully.",
HTTPStatus.INTERNAL_SERVER_ERROR,
)
operation_succeeded = cloud_sql_client.wait_until_operation_completed(
operation_id, seconds_to_wait=GCS_IMPORT_EXPORT_TIMEOUT_SEC)
if not operation_succeeded:
return (
"Cloud SQL import did not complete within 60 seconds",
HTTPStatus.INTERNAL_SERVER_ERROR,
)
return operation_id, HTTPStatus.OK
def test_from_county_ingest_bucket(self) -> None:
ingest_bucket_path = gcsfs_direct_ingest_bucket_for_region(
region_code="us_xx_yyyyy",
system_level=SystemLevel.COUNTY,
ingest_instance=DirectIngestInstance.PRIMARY,
project_id="recidiviz-456",
)
self.assertEqual(
DirectIngestInstance.PRIMARY,
DirectIngestInstance.for_ingest_bucket(ingest_bucket_path),
)
def test_direct_ingest_instance_status_contains_data_for_all_states(
self) -> None:
'''Enforces that after all migrations the set of direct ingest instance statuses
matches the list of known states.
If this test fails, you will likely have to add a new migration because a new state
was recently created. To do so, first run:
```
python -m recidiviz.tools.migrations.autogenerate_migration \
--database OPERATIONS \
--message add_us_xx
```
This will generate a blank migration. You should then modify the migration, changing
the `upgrade` method to look like:
```
def upgrade() -> None:
op.execute("""
INSERT INTO direct_ingest_instance_status (region_code, instance, is_paused) VALUES
('US_XX', 'PRIMARY', TRUE),
('US_XX', 'SECONDARY', TRUE);
""")
```
Afterwards, this test should ideally pass.
'''
with runner(self.default_config(), self.engine) as r:
r.migrate_up_to("head")
engine = create_engine(
local_postgres_helpers.postgres_db_url_from_env_vars())
conn = engine.connect()
rows = conn.execute(
"SELECT region_code, instance FROM direct_ingest_instance_status;"
)
instance_to_state_codes = defaultdict(set)
for row in rows:
instance_to_state_codes[DirectIngestInstance(row[1])].add(
row[0])
required_states = {
name.upper()
for name in get_existing_region_dir_names()
}
for instance in DirectIngestInstance:
self.assertEqual(required_states,
instance_to_state_codes[instance])
def is_task_queued(self, region: Region,
ingest_args: GcsfsIngestArgs) -> bool:
"""Returns true if the ingest_args correspond to a task currently in
the queue.
"""
task_id_prefix = _build_task_id(
region.region_code,
DirectIngestInstance.for_ingest_bucket(
ingest_args.file_path.bucket_path),
ingest_args.task_id_tag(),
prefix_only=True,
)
return bool(next(self._tasks_for_prefix(task_id_prefix), None))
def test_from_state_ingest_bucket(self) -> None:
ingest_bucket_path = gcsfs_direct_ingest_bucket_for_region(
region_code="us_xx",
system_level=SystemLevel.STATE,
ingest_instance=DirectIngestInstance.PRIMARY,
project_id="recidiviz-456",
)
self.assertEqual(
DirectIngestInstance.PRIMARY,
DirectIngestInstance.for_ingest_bucket(ingest_bucket_path),
)
ingest_bucket_path = gcsfs_direct_ingest_bucket_for_region(
region_code="us_xx",
system_level=SystemLevel.STATE,
ingest_instance=DirectIngestInstance.SECONDARY,
project_id="recidiviz-456",
)
self.assertEqual(
DirectIngestInstance.SECONDARY,
DirectIngestInstance.for_ingest_bucket(ingest_bucket_path),
)
def _release_ingest_lock() -> Tuple[str, HTTPStatus]:
try:
state_code = StateCode(request.json["stateCode"])
ingest_instance = DirectIngestInstance(
request.json["ingestInstance"])
except ValueError:
return "invalid parameters provided", HTTPStatus.BAD_REQUEST
lock_manager = DirectIngestRegionLockManager.for_state_ingest(
state_code, ingest_instance=ingest_instance)
try:
lock_manager.release_lock()
except GCSPseudoLockDoesNotExist:
return "lock does not exist", HTTPStatus.NOT_FOUND
return "", HTTPStatus.OK
def handle_direct_ingest_file() -> Tuple[str, HTTPStatus]:
"""Called from a Cloud Function when a new file is added to a direct ingest
bucket. Will trigger a job that deals with normalizing and splitting the
file as is appropriate, then start the scheduler if allowed.
"""
region_code = get_str_param_value("region", request.args)
# The bucket name for the file to ingest
bucket = get_str_param_value("bucket", request.args)
# The relative path to the file, not including the bucket name
relative_file_path = get_str_param_value("relative_file_path",
request.args,
preserve_case=True)
start_ingest = get_bool_param_value("start_ingest",
request.args,
default=False)
if not region_code or not bucket or not relative_file_path or start_ingest is None:
response = f"Bad parameters [{request.args}]"
logging.error(response)
return response, HTTPStatus.BAD_REQUEST
bucket_path = GcsfsBucketPath(bucket_name=bucket)
with monitoring.push_region_tag(
region_code,
ingest_instance=DirectIngestInstance.for_ingest_bucket(
bucket_path).value,
):
try:
controller = DirectIngestControllerFactory.build(
ingest_bucket_path=bucket_path,
allow_unlaunched=True,
)
except DirectIngestError as e:
if e.is_bad_request():
logging.error(str(e))
return str(e), HTTPStatus.BAD_REQUEST
raise e
path = GcsfsPath.from_bucket_and_blob_name(
bucket_name=bucket, blob_name=relative_file_path)
if isinstance(path, GcsfsFilePath):
controller.handle_file(path, start_ingest=start_ingest)
return "", HTTPStatus.OK
def _pause_direct_ingest_instance() -> Tuple[str, HTTPStatus]:
try:
state_code = StateCode(request.json["stateCode"])
ingest_instance = DirectIngestInstance(
request.json["ingestInstance"])
except ValueError:
return "invalid parameters provided", HTTPStatus.BAD_REQUEST
ingest_status_manager = DirectIngestInstanceStatusManager(
region_code=state_code.value, ingest_instance=ingest_instance)
try:
ingest_status_manager.pause_instance()
except Exception:
return (
"something went wrong pausing the intance",
HTTPStatus.INTERNAL_SERVER_ERROR,
)
return "", HTTPStatus.OK
def main(state_code: StateCode, ingest_instance: DirectIngestInstance) -> None:
"""Executes the main flow of the script."""
print(
f"RUN THE FOLLOWING COMMANDS IN ORDER TO DELETE ALL DATA FOR REGION [{state_code.value}]"
)
print(
"********************************************************************************"
)
db_version = ingest_instance.database_version(SystemLevel.STATE,
state_code)
db_key = SQLAlchemyDatabaseKey.for_state_code(state_code=state_code,
db_version=db_version)
# Connect to correct database for instance first
print(f"\\c {db_key.db_name}")
# Then run deletion commands
for cmd in generate_region_deletion_commands(state_code, db_version):
print(cmd)
print(
"********************************************************************************"
)
print("HOW TO PERFORM DELETION:")
print(
"1) Log into prod data client (`gcloud compute ssh prod-data-client --project=recidiviz-123`)"
)
print("\n> For production deletion:")
print(
"2) Go to secret manager to get login credentials stored in `state_db_user` and `state_db_password` secrets:"
"\n\thttps://console.cloud.google.com/security/secret-manager?organizationId=448885369991&"
"project=recidiviz-123")
print("3) Log into postgres database (`prod-state-psql`)")
print("\n> For staging deletion:")
print(
"2) Go to secret manager to get login credentials stored in `state_db_user` and `state_db_password` secrets:"
"\n\thttps://console.cloud.google.com/security/secret-manager?organizationId=448885369991&"
"project=recidiviz-staging")
print("3) Log into postgres database (`dev-state-psql`)")
print("\n> For all:")
print(
"4) Paste full list of commands listed above in postgres command line and run. Some commands may take a "
"while to run.")
def handle_new_files() -> Tuple[str, HTTPStatus]:
"""Normalizes and splits files in the ingest bucket for a given region as
is appropriate. Will schedule the next process_job task if no renaming /
splitting work has been done that will trigger subsequent calls to this
endpoint.
"""
logging.info("Received request for direct ingest handle_new_files: %s",
request.values)
region_code = get_str_param_value("region", request.values)
can_start_ingest = get_bool_param_value("can_start_ingest",
request.values,
default=False)
bucket = get_str_param_value("bucket", request.values)
if not region_code or can_start_ingest is None or not bucket:
response = f"Bad parameters [{request.values}]"
logging.error(response)
return response, HTTPStatus.BAD_REQUEST
bucket_path = GcsfsBucketPath(bucket_name=bucket)
with monitoring.push_region_tag(
region_code,
ingest_instance=DirectIngestInstance.for_ingest_bucket(
bucket_path).value,
):
try:
controller = DirectIngestControllerFactory.build(
ingest_bucket_path=bucket_path,
allow_unlaunched=True,
)
except DirectIngestError as e:
if e.is_bad_request():
logging.error(str(e))
return str(e), HTTPStatus.BAD_REQUEST
raise e
controller.handle_new_files(can_start_ingest=can_start_ingest)
return "", HTTPStatus.OK
def process_job() -> Tuple[str, HTTPStatus]:
"""Processes a single direct ingest file, specified in the provided ingest
arguments.
"""
logging.info("Received request to process direct ingest job: [%s]",
request.values)
region_code = get_str_param_value("region", request.values)
file_path = get_str_param_value("file_path",
request.values,
preserve_case=True)
if not region_code or not file_path:
response = f"Bad parameters [{request.values}]"
logging.error(response)
return response, HTTPStatus.BAD_REQUEST
gcsfs_path = GcsfsFilePath.from_absolute_path(file_path)
with monitoring.push_region_tag(
region_code,
ingest_instance=DirectIngestInstance.for_ingest_bucket(
gcsfs_path.bucket_path).value,
):
json_data = request.get_data(as_text=True)
ingest_args = _parse_cloud_task_args(json_data)
if not ingest_args:
raise DirectIngestError(
msg="process_job was called with no GcsfsIngestArgs.",
error_type=DirectIngestErrorType.INPUT_ERROR,
)
if not isinstance(ingest_args, GcsfsIngestArgs):
raise DirectIngestError(
msg=
f"process_job was called with incorrect args type [{type(ingest_args)}].",
error_type=DirectIngestErrorType.INPUT_ERROR,
)
if gcsfs_path != ingest_args.file_path:
raise DirectIngestError(
msg=f"Different paths were passed in the url and request body\n"
f"url: {gcsfs_path.uri()}\n"
f"body: {ingest_args.file_path.uri()}",
error_type=DirectIngestErrorType.INPUT_ERROR,
)
with monitoring.push_tags(
{TagKey.INGEST_TASK_TAG: ingest_args.task_id_tag()}):
try:
controller = DirectIngestControllerFactory.build(
ingest_bucket_path=ingest_args.file_path.bucket_path,
allow_unlaunched=False,
)
except DirectIngestError as e:
if e.is_bad_request():
logging.error(str(e))
return str(e), HTTPStatus.BAD_REQUEST
raise e
try:
controller.run_ingest_job_and_kick_scheduler_on_completion(
ingest_args)
except GCSPseudoLockAlreadyExists as e:
logging.warning(str(e))
return str(e), HTTPStatus.CONFLICT
return "", HTTPStatus.OK
def main() -> None:
"""Runs the move_state_files_to_deprecated script."""
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser.add_argument(
"--file-type",
required=True,
choices=[file_type.value for file_type in GcsfsDirectIngestFileType],
help="Defines whether we should move raw files or generated ingest_view files",
)
parser.add_argument("--region", required=True, help="E.g. 'us_nd'")
parser.add_argument(
"--ingest-instance",
required=True,
choices=[instance.value for instance in DirectIngestInstance],
help="Defines which ingest instance we should be deprecating files for.",
)
parser.add_argument(
"--dry-run",
default=True,
type=str_to_bool,
help="Runs move in dry-run mode, only prints the file moves it would do.",
)
parser.add_argument(
"--start-date-bound",
help="The lower bound date to start from, inclusive. For partial moving of ingested files. "
"E.g. 2019-09-23.",
)
parser.add_argument(
"--end-date-bound",
help="The upper bound date to end at, inclusive. For partial moving of ingested files. "
"E.g. 2019-09-23.",
)
parser.add_argument(
"--project-id", help="The id for this particular project, E.g. 'recidiviz-123'"
)
parser.add_argument(
"--file-filter",
default=None,
help="Regex name filter - when set, will only move files that match this regex.",
)
args = parser.parse_args()
logging.basicConfig(level=logging.INFO, format="%(message)s")
MoveFilesToDeprecatedController(
file_type=GcsfsDirectIngestFileType(args.file_type),
region_code=args.region,
ingest_instance=DirectIngestInstance(args.ingest_instance),
start_date_bound=args.start_date_bound,
end_date_bound=args.end_date_bound,
project_id=args.project_id,
dry_run=args.dry_run,
file_filter=args.file_filter,
).run()
def raw_data_import() -> Tuple[str, HTTPStatus]:
"""Imports a single raw direct ingest CSV file from a location in GCS File System to its corresponding raw data
table in BQ.
"""
logging.info("Received request to do direct ingest raw data import: [%s]",
request.values)
region_code = get_str_param_value("region", request.values)
file_path = get_str_param_value("file_path",
request.values,
preserve_case=True)
if not region_code or not file_path:
response = f"Bad parameters [{request.values}]"
logging.error(response)
return response, HTTPStatus.BAD_REQUEST
gcsfs_path = GcsfsFilePath.from_absolute_path(file_path)
with monitoring.push_region_tag(
region_code,
ingest_instance=DirectIngestInstance.for_ingest_bucket(
gcsfs_path.bucket_path).value,
):
json_data = request.get_data(as_text=True)
data_import_args = _parse_cloud_task_args(json_data)
if not data_import_args:
raise DirectIngestError(
msg=
"raw_data_import was called with no GcsfsRawDataBQImportArgs.",
error_type=DirectIngestErrorType.INPUT_ERROR,
)
if not isinstance(data_import_args, GcsfsRawDataBQImportArgs):
raise DirectIngestError(
msg=
f"raw_data_import was called with incorrect args type [{type(data_import_args)}].",
error_type=DirectIngestErrorType.INPUT_ERROR,
)
if gcsfs_path != data_import_args.raw_data_file_path:
raise DirectIngestError(
msg=f"Different paths were passed in the url and request body\n"
f"url: {gcsfs_path.uri()}\n"
f"body: {data_import_args.raw_data_file_path.uri()}",
error_type=DirectIngestErrorType.INPUT_ERROR,
)
with monitoring.push_tags(
{TagKey.RAW_DATA_IMPORT_TAG: data_import_args.task_id_tag()}):
try:
controller = DirectIngestControllerFactory.build(
ingest_bucket_path=data_import_args.raw_data_file_path.
bucket_path,
allow_unlaunched=False,
)
except DirectIngestError as e:
if e.is_bad_request():
logging.error(str(e))
return str(e), HTTPStatus.BAD_REQUEST
raise e
controller.do_raw_data_import(data_import_args)
return "", HTTPStatus.OK
def _is_legacy_instance(system_level: SystemLevel, region_code: str,
ingest_instance: DirectIngestInstance) -> bool:
return (ingest_instance.database_version(system_level,
StateCode.get(region_code)) is
SQLAlchemyStateDatabaseVersion.LEGACY)
def ingest_view_export() -> Tuple[str, HTTPStatus]:
"""Exports an ingest view from BQ to a file in the region's GCS File System ingest bucket that is ready to be
processed and ingested into our Recidiviz DB.
"""
logging.info("Received request to do direct ingest view export: [%s]",
request.values)
region_code = get_str_param_value("region", request.values)
output_bucket_name = get_str_param_value("output_bucket",
request.values,
preserve_case=True)
if not region_code or not output_bucket_name:
response = f"Bad parameters [{request.values}]"
logging.error(response)
return response, HTTPStatus.BAD_REQUEST
with monitoring.push_region_tag(
region_code,
ingest_instance=DirectIngestInstance.for_ingest_bucket(
GcsfsBucketPath(output_bucket_name)).value,
):
json_data = request.get_data(as_text=True)
ingest_view_export_args = _parse_cloud_task_args(json_data)
if not ingest_view_export_args:
raise DirectIngestError(
msg=
"raw_data_import was called with no GcsfsIngestViewExportArgs.",
error_type=DirectIngestErrorType.INPUT_ERROR,
)
if not isinstance(ingest_view_export_args, GcsfsIngestViewExportArgs):
raise DirectIngestError(
msg=
f"raw_data_import was called with incorrect args type [{type(ingest_view_export_args)}].",
error_type=DirectIngestErrorType.INPUT_ERROR,
)
if output_bucket_name != ingest_view_export_args.output_bucket_name:
raise DirectIngestError(
msg=
f"Different buckets were passed in the url and request body\n"
f"url: {output_bucket_name}\n"
f"body: {ingest_view_export_args.output_bucket_name}",
error_type=DirectIngestErrorType.INPUT_ERROR,
)
with monitoring.push_tags({
TagKey.INGEST_VIEW_EXPORT_TAG:
ingest_view_export_args.task_id_tag()
}):
try:
controller = DirectIngestControllerFactory.build(
ingest_bucket_path=GcsfsBucketPath(
ingest_view_export_args.output_bucket_name),
allow_unlaunched=False,
)
except DirectIngestError as e:
if e.is_bad_request():
logging.error(str(e))
return str(e), HTTPStatus.BAD_REQUEST
raise e
controller.do_ingest_view_export(ingest_view_export_args)
return "", HTTPStatus.OK