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 __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 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 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 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 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 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
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