def __init__(
self,
project_id: str,
region: str,
lower_bound_update_datetime: Optional[datetime.datetime],
gcs_destination_path: Optional[GcsfsDirectoryPath] = None,
):
self.project_id = project_id
self.region = region.lower()
self.auth = SftpAuth.for_region(region)
self.delegate = SftpDownloadDelegateFactory.build(region_code=region)
self.gcsfs = DirectIngestGCSFileSystem(GcsfsFactory.build())
self.unable_to_download_items: List[str] = []
self.downloaded_items: List[Tuple[str, datetime.datetime]] = []
self.skipped_files: List[str] = []
self.lower_bound_update_datetime = lower_bound_update_datetime
self.bucket = (gcsfs_sftp_download_bucket_path_for_region(
region, SystemLevel.STATE, project_id=self.project_id) if
gcs_destination_path is None else gcs_destination_path)
self.download_dir = GcsfsDirectoryPath.from_dir_and_subdir(
dir_path=self.bucket, subdir=RAW_INGEST_DIRECTORY)
self.postgres_direct_ingest_file_metadata_manager = (
PostgresDirectIngestRawFileMetadataManager(
region,
DirectIngestInstance.PRIMARY.database_version(
SystemLevel.STATE,
state_code=StateCode(self.region.upper())).name,
))
def __init__(
self,
project_id: str,
region: str,
lower_bound_update_datetime: Optional[datetime.datetime],
gcs_destination_path: Optional[str] = None,
):
self.project_id = project_id
self.region = region.lower()
self.auth = SftpAuth.for_region(region)
self.delegate = SftpDownloadDelegateFactory.build(region_code=region)
self.gcsfs = DirectIngestGCSFileSystem(GcsfsFactory.build())
self.unable_to_download_items: List[str] = []
self.downloaded_items: List[Tuple[str, datetime.datetime]] = []
self.lower_bound_update_datetime = lower_bound_update_datetime
self.bucket = (
GcsfsDirectoryPath.from_absolute_path(
gcsfs_direct_ingest_directory_path_for_region(
region, SystemLevel.STATE, project_id=self.project_id
)
)
if gcs_destination_path is None
else GcsfsDirectoryPath.from_absolute_path(gcs_destination_path)
)
self.download_dir = GcsfsDirectoryPath.from_dir_and_subdir(
dir_path=self.bucket, subdir=RAW_INGEST_DIRECTORY
)
def normalize_raw_file_path() -> Tuple[str, HTTPStatus]:
"""Called from a Cloud Function when a new file is added to a bucket that is configured to rename files but not
ingest them. For example, a bucket that is being used for automatic data transfer testing.
"""
# The bucket name for the file to normalize
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)
if not bucket or not relative_file_path:
return f"Bad parameters [{request.args}]", HTTPStatus.BAD_REQUEST
path = GcsfsPath.from_bucket_and_blob_name(bucket_name=bucket,
blob_name=relative_file_path)
if not isinstance(path, GcsfsFilePath):
raise ValueError(
f"Incorrect type [{type(path)}] for path: {path.uri()}")
fs = DirectIngestGCSFileSystem(GcsfsFactory.build())
fs.mv_path_to_normalized_path(path,
file_type=GcsfsDirectIngestFileType.RAW_DATA)
return "", HTTPStatus.OK
def setUp(self) -> None:
self.fs = DirectIngestGCSFileSystem(FakeGCSFileSystem())
self.prioritizer = GcsfsDirectIngestJobPrioritizer(
self.fs,
self._INGEST_BUCKET_PATH,
["tagA", "tagB"],
file_type_filter=GcsfsDirectIngestFileType.INGEST_VIEW,
)
def __init__(self,
region_name: str,
system_level: SystemLevel,
ingest_directory_path: Optional[str] = None,
storage_directory_path: Optional[str] = None,
max_delay_sec_between_files: Optional[int] = None):
super().__init__(region_name, system_level)
self.fs = DirectIngestGCSFileSystem(GcsfsFactory.build())
self.max_delay_sec_between_files = max_delay_sec_between_files
if not ingest_directory_path:
ingest_directory_path = \
gcsfs_direct_ingest_directory_path_for_region(region_name,
system_level)
self.ingest_directory_path = \
GcsfsDirectoryPath.from_absolute_path(ingest_directory_path)
if not storage_directory_path:
storage_directory_path = \
gcsfs_direct_ingest_storage_directory_path_for_region(
region_name, system_level)
self.storage_directory_path = \
GcsfsDirectoryPath.from_absolute_path(storage_directory_path)
self.temp_output_directory_path = \
GcsfsDirectoryPath.from_absolute_path(gcsfs_direct_ingest_temporary_output_directory_path())
ingest_job_file_type_filter = \
GcsfsDirectIngestFileType.INGEST_VIEW \
if self.region.is_raw_vs_ingest_file_name_detection_enabled() else None
self.file_prioritizer = \
GcsfsDirectIngestJobPrioritizer(
self.fs,
self.ingest_directory_path,
self.get_file_tag_rank_list(),
ingest_job_file_type_filter)
self.ingest_file_split_line_limit = self._INGEST_FILE_SPLIT_LINE_LIMIT
self.file_metadata_manager = PostgresDirectIngestFileMetadataManager(
region_code=self.region.region_code)
self.raw_file_import_manager = DirectIngestRawFileImportManager(
region=self.region,
fs=self.fs,
ingest_directory_path=self.ingest_directory_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,
ingest_directory_path=self.ingest_directory_path,
file_metadata_manager=self.file_metadata_manager,
big_query_client=BigQueryClientImpl(),
view_collector=DirectIngestPreProcessedIngestViewCollector(
self.region, self.get_file_tag_rank_list()))
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 __init__(
self,
paths_with_timestamps: List[Tuple[str, datetime.datetime]],
project_id: str,
region: str,
delegate: UploadStateFilesToIngestBucketDelegate,
destination_bucket_override: Optional[GcsfsBucketPath] = None,
):
self.paths_with_timestamps = paths_with_timestamps
self.project_id = project_id
self.region = region.lower()
self.delegate = delegate
self.gcsfs = DirectIngestGCSFileSystem(GcsfsFactory.build())
# Raw data uploads always default to primary ingest bucket
self.destination_ingest_bucket = (
destination_bucket_override
or gcsfs_direct_ingest_bucket_for_region(
region_code=region,
system_level=SystemLevel.STATE,
ingest_instance=DirectIngestInstance.PRIMARY,
project_id=self.project_id,
))
self.uploaded_files: List[str] = []
self.skipped_files: List[str] = []
self.unable_to_upload_files: List[str] = []
def _move_files(self, from_uri: str):
curr_gcsfs_file_path = GcsfsFilePath.from_absolute_path(from_uri)
previous_date_format = filename_parts_from_path(
curr_gcsfs_file_path).date_str
new_date_format = date.fromisoformat(previous_date_format).strftime(
"%Y/%m/%d/")
path_with_new_file_name = GcsfsFilePath.from_absolute_path(
to_normalized_unprocessed_file_path_from_normalized_path(
from_uri, GcsfsDirectIngestFileType.RAW_DATA))
if DirectIngestGCSFileSystem.is_processed_file(curr_gcsfs_file_path):
path_with_new_file_name = GcsfsFilePath.from_absolute_path(
to_normalized_processed_file_path_from_normalized_path(
from_uri, GcsfsDirectIngestFileType.RAW_DATA))
raw_dir_with_date = GcsfsDirectoryPath.from_dir_and_subdir(
self.region_storage_raw_dir_path, new_date_format)
to_uri = GcsfsFilePath.from_directory_and_file_name(
raw_dir_with_date, path_with_new_file_name.file_name).uri()
if not self.dry_run:
gsutil_mv(from_path=from_uri, to_path=to_uri)
with self.mutex:
self.move_list.append((from_uri, to_uri))
if self.move_progress:
self.move_progress.next()
def test_normalize_file_path(self, mock_fs_factory: mock.MagicMock,
mock_environment: mock.MagicMock) -> None:
mock_environment.return_value = "production"
mock_fs = FakeGCSFileSystem()
mock_fs_factory.return_value = mock_fs
path = GcsfsFilePath.from_absolute_path("bucket-us-xx/file-tag.csv")
mock_fs.test_add_path(path, local_path=None)
request_args = {
"bucket": path.bucket_name,
"relative_file_path": path.blob_name,
}
headers = {"X-Appengine-Cron": "test-cron"}
response = self.client.get("/normalize_raw_file_path",
query_string=request_args,
headers=headers)
self.assertEqual(200, response.status_code)
self.assertEqual(1, len(mock_fs.all_paths))
registered_path = mock_fs.all_paths[0]
if not isinstance(registered_path, GcsfsFilePath):
self.fail(f"Unexpected type for path [{type(registered_path)}]")
self.assertTrue(
DirectIngestGCSFileSystem.is_normalized_file_path(registered_path))
def __init__(
self,
paths_with_timestamps: List[Tuple[str, datetime.datetime]],
project_id: str,
region: str,
gcs_destination_path: Optional[str] = None,
):
self.paths_with_timestamps = paths_with_timestamps
self.project_id = project_id
self.region = region.lower()
self.gcsfs = DirectIngestGCSFileSystem(GcsfsFactory.build())
self.gcs_destination_path = (
GcsfsDirectoryPath.from_absolute_path(
gcsfs_direct_ingest_directory_path_for_region(
region, SystemLevel.STATE, project_id=self.project_id))
if gcs_destination_path is None else
GcsfsDirectoryPath.from_absolute_path(gcs_destination_path))
self.uploaded_files: List[str] = []
self.unable_to_upload_files: List[str] = []
def setUp(self) -> None:
self.project_id = "recidiviz-456"
self.project_id_patcher = patch("recidiviz.utils.metadata.project_id")
self.project_id_patcher.start().return_value = self.project_id
self.test_region = fake_region(
region_code="us_xx", are_raw_data_bq_imports_enabled_in_env=True)
self.region_module_patcher = patch.object(
direct_ingest_raw_table_migration_collector,
"regions",
new=controller_fixtures,
)
self.region_module_patcher.start()
self.fs = DirectIngestGCSFileSystem(FakeGCSFileSystem())
self.ingest_directory_path = GcsfsDirectoryPath(
bucket_name="direct/controllers/fixtures")
self.temp_output_path = GcsfsDirectoryPath(bucket_name="temp_bucket")
self.region_raw_file_config = DirectIngestRegionRawFileConfig(
region_code="us_xx",
yaml_config_file_dir=fixtures.as_filepath("us_xx"),
)
self.mock_big_query_client = create_autospec(BigQueryClient)
self.num_lines_uploaded = 0
self.mock_big_query_client.insert_into_table_from_cloud_storage_async.side_effect = (
self.mock_import_raw_file_to_big_query)
self.import_manager = DirectIngestRawFileImportManager(
region=self.test_region,
fs=self.fs,
ingest_directory_path=self.ingest_directory_path,
temp_output_directory_path=self.temp_output_path,
region_raw_file_config=self.region_raw_file_config,
big_query_client=self.mock_big_query_client,
)
self.import_manager.csv_reader = _TestSafeGcsCsvReader(
self.fs.gcs_file_system)
self.time_patcher = patch(
"recidiviz.ingest.direct.controllers.direct_ingest_raw_file_import_manager.time"
)
self.mock_time = self.time_patcher.start()
def fake_get_dataset_ref(dataset_id: str) -> bigquery.DatasetReference:
return bigquery.DatasetReference(project=self.project_id,
dataset_id=dataset_id)
self.mock_big_query_client.dataset_ref_for_id = fake_get_dataset_ref
def _split_file(
self, path: GcsfsFilePath,
file_contents_handle: GcsfsFileContentsHandle
) -> List[GcsfsFileContentsHandle]:
split_contents_handles = []
for df in pd.read_csv(file_contents_handle.local_file_path,
dtype=str,
chunksize=self.file_split_line_limit,
keep_default_na=False):
local_file_path = DirectIngestGCSFileSystem.generate_random_temp_path(
)
df.to_csv(local_file_path, index=False)
split_contents_handles.append(
GcsfsFileContentsHandle(local_file_path))
return split_contents_handles
def __init__(
self,
*,
state_code: StateCode,
sandbox_dataset_prefix: str,
test_ingest_bucket: GcsfsBucketPath,
):
check_is_valid_sandbox_bucket(test_ingest_bucket)
super().__init__(
region=get_region(state_code.value.lower(), is_direct_ingest=True),
fs=DirectIngestGCSFileSystem(GcsfsFactory.build()),
ingest_bucket_path=test_ingest_bucket,
temp_output_directory_path=GcsfsDirectoryPath.from_dir_and_subdir(
test_ingest_bucket, "temp_raw_data"
),
big_query_client=BigQueryClientImpl(),
)
self.sandbox_dataset = (
f"{sandbox_dataset_prefix}_{super()._raw_tables_dataset()}"
)
class TestGcsfsDirectIngestJobPrioritizerIngestViewFilter(unittest.TestCase):
"""Tests for the GcsfsDirectIngestJobPrioritizer."""
_DAY_1_TIME_1 = datetime.datetime(year=2019,
month=1,
day=2,
hour=3,
minute=4,
second=5,
microsecond=6789,
tzinfo=datetime.timezone.utc)
_DAY_1_TIME_2 = datetime.datetime(year=2019,
month=1,
day=2,
hour=3,
minute=4,
second=5,
microsecond=7789,
tzinfo=datetime.timezone.utc)
_DAY_1_TIME_3 = datetime.datetime(year=2019,
month=1,
day=2,
hour=10,
minute=4,
second=5,
microsecond=678,
tzinfo=datetime.timezone.utc)
_DAY_2_TIME_1 = datetime.datetime(year=2019,
month=1,
day=3,
hour=3,
minute=4,
second=5,
microsecond=6789,
tzinfo=datetime.timezone.utc)
_DAY_1 = _DAY_1_TIME_1.date()
_DAY_2 = _DAY_2_TIME_1.date()
_INGEST_BUCKET_PATH = \
GcsfsDirectoryPath.from_absolute_path('direct/regions/us_nd/fixtures')
def setUp(self) -> None:
self.fs = DirectIngestGCSFileSystem(FakeGCSFileSystem())
self.prioritizer = GcsfsDirectIngestJobPrioritizer(
self.fs,
self._INGEST_BUCKET_PATH, ['tagA', 'tagB'],
file_type_filter=GcsfsDirectIngestFileType.INGEST_VIEW)
FIXTURE_PATH_PREFIX = 'direct/regions/us_nd/fixtures'
def _normalized_path_for_filename(self, filename: str,
file_type: GcsfsDirectIngestFileType,
dt: datetime.datetime) -> GcsfsFilePath:
normalized_path = \
to_normalized_unprocessed_file_path(
original_file_path=os.path.join(self._INGEST_BUCKET_PATH.abs_path(), filename),
file_type=file_type,
dt=dt)
return GcsfsFilePath.from_absolute_path(normalized_path)
def _process_jobs_for_paths_with_no_gaps_in_expected_order(
self, paths: List[GcsfsFilePath]):
for path in paths:
date_str = filename_parts_from_path(path).date_str
next_job_args = self.prioritizer.get_next_job_args()
self.assertIsNotNone(next_job_args)
if next_job_args is None:
# Make mypy happy
self.fail()
self.assertEqual(next_job_args.file_path, path)
self.assertTrue(
self.prioritizer.are_next_args_expected(next_job_args))
self.assertTrue(
self.prioritizer.are_more_jobs_expected_for_day(date_str))
# ... job runs ...
self.fs.mv_path_to_processed_path(path)
def test_empty_fs(self):
self.assertTrue(
self.prioritizer.are_more_jobs_expected_for_day(
self._DAY_1_TIME_1.date().isoformat()))
self.assertIsNone(self.prioritizer.get_next_job_args())
def test_single_expected_file(self):
path = self._normalized_path_for_filename(
'tagA.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_1)
fixture_util.add_direct_ingest_path(self.fs.gcs_file_system,
path,
has_fixture=False)
self._process_jobs_for_paths_with_no_gaps_in_expected_order([path])
self.assertIsNone(self.prioritizer.get_next_job_args())
# We still expect a file for tagB
self.assertTrue(
self.prioritizer.are_more_jobs_expected_for_day(
self._DAY_1.isoformat()))
def test_multiple_files(self):
paths = [
self._normalized_path_for_filename(
'tagA.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_1),
self._normalized_path_for_filename(
'tagB.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_2),
# This file shouldn't ge tpicked up
self._normalized_path_for_filename(
'tagC.csv', GcsfsDirectIngestFileType.RAW_DATA,
self._DAY_1_TIME_3)
]
for path in paths:
fixture_util.add_direct_ingest_path(self.fs.gcs_file_system,
path,
has_fixture=False)
# Exclude last raw file
expected_processed_paths = paths[0:-1]
self._process_jobs_for_paths_with_no_gaps_in_expected_order(
expected_processed_paths)
self.assertIsNone(self.prioritizer.get_next_job_args())
self.assertFalse(
self.prioritizer.are_more_jobs_expected_for_day(
self._DAY_1.isoformat()))
def test_unexpected_file(self):
# Only file is out of order
path = self._normalized_path_for_filename(
'tagB.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_1)
fixture_util.add_direct_ingest_path(self.fs.gcs_file_system,
path,
has_fixture=False)
self.assertTrue(
self.prioritizer.are_more_jobs_expected_for_day(
self._DAY_1.isoformat()))
next_job_args = self.prioritizer.get_next_job_args()
self.assertIsNotNone(next_job_args)
self.assertEqual(next_job_args.file_path, path)
self.assertFalse(
self.prioritizer.are_next_args_expected(next_job_args))
# ... job runs eventually even though unexpected...
self.fs.mv_path_to_processed_path(path)
self.assertIsNone(self.prioritizer.get_next_job_args())
# We still expect a file for tagA
self.assertTrue(
self.prioritizer.are_more_jobs_expected_for_day(
self._DAY_1.isoformat()))
def test_files_on_multiple_days(self):
paths = [
self._normalized_path_for_filename(
'tagA.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_1),
self._normalized_path_for_filename(
'tagB.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_2),
self._normalized_path_for_filename(
'tagA.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_2_TIME_1),
# This file shouldn't ge tpicked up
self._normalized_path_for_filename(
'tagC.csv', GcsfsDirectIngestFileType.RAW_DATA,
self._DAY_1_TIME_3)
]
for path in paths:
fixture_util.add_direct_ingest_path(self.fs.gcs_file_system,
path,
has_fixture=False)
# Exclude last raw file
expected_processed_paths = paths[0:-1]
self._process_jobs_for_paths_with_no_gaps_in_expected_order(
expected_processed_paths)
self.assertIsNone(self.prioritizer.get_next_job_args())
self.assertFalse(
self.prioritizer.are_more_jobs_expected_for_day(
self._DAY_1.isoformat()))
self.assertTrue(
self.prioritizer.are_more_jobs_expected_for_day(
self._DAY_2.isoformat()))
def test_files_on_multiple_days_with_gap(self):
"""Runs a test where there are files on multiple days and there is a gap
in the expected files for the first day.
"""
paths = [
self._normalized_path_for_filename(
'tagB.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_2),
self._normalized_path_for_filename(
'tagA.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_2_TIME_1),
]
for path in paths:
fixture_util.add_direct_ingest_path(self.fs.gcs_file_system,
path,
has_fixture=False)
for i, path in enumerate(paths):
date_str = filename_parts_from_path(path).date_str
next_job_args = self.prioritizer.get_next_job_args()
self.assertIsNotNone(next_job_args)
self.assertEqual(next_job_args.file_path, path)
are_args_expected = \
self.prioritizer.are_next_args_expected(next_job_args)
if i == 0:
self.assertFalse(are_args_expected)
else:
self.assertTrue(are_args_expected)
self.assertTrue(
self.prioritizer.are_more_jobs_expected_for_day(date_str))
# ... job runs ...
self.fs.mv_path_to_processed_path(path)
self.assertIsNone(self.prioritizer.get_next_job_args())
self.assertTrue(
self.prioritizer.are_more_jobs_expected_for_day(
self._DAY_1.isoformat()))
self.assertTrue(
self.prioritizer.are_more_jobs_expected_for_day(
self._DAY_2.isoformat()))
def test_multiple_files_same_tag(self):
paths = [
self._normalized_path_for_filename(
'tagA.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_1),
self._normalized_path_for_filename(
'tagA.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_2),
self._normalized_path_for_filename(
'tagB.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_3),
]
for path in paths:
fixture_util.add_direct_ingest_path(self.fs.gcs_file_system,
path,
has_fixture=False)
self._process_jobs_for_paths_with_no_gaps_in_expected_order(paths)
self.assertIsNone(self.prioritizer.get_next_job_args())
self.assertFalse(
self.prioritizer.are_more_jobs_expected_for_day(
self._DAY_1.isoformat()))
def test_multiple_files_times_out_of_order(self):
"""Runs a test where there are no gaps but the files have been added
(i.e. have creation times) out of order.
"""
paths = [
self._normalized_path_for_filename(
'tagA.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_2),
self._normalized_path_for_filename(
'tagB.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_1),
self._normalized_path_for_filename(
'tagB.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_3),
]
for path in paths:
fixture_util.add_direct_ingest_path(self.fs.gcs_file_system,
path,
has_fixture=False)
for i, path in enumerate(paths):
date_str = filename_parts_from_path(path).date_str
next_job_args = self.prioritizer.get_next_job_args()
self.assertIsNotNone(next_job_args)
self.assertEqual(next_job_args.file_path, path)
self.assertTrue(
self.prioritizer.are_next_args_expected(next_job_args))
are_more_jobs_expected = \
self.prioritizer.are_more_jobs_expected_for_day(date_str)
if i == 2:
self.assertFalse(are_more_jobs_expected)
else:
self.assertTrue(are_more_jobs_expected)
# ... job runs ...
self.fs.mv_path_to_processed_path(path)
self.assertIsNone(self.prioritizer.get_next_job_args())
self.assertFalse(
self.prioritizer.are_more_jobs_expected_for_day(
self._DAY_1.isoformat()))
def test_run_multiple_copies_of_same_tag(self):
paths = [
self._normalized_path_for_filename(
'tagA.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_2),
self._normalized_path_for_filename(
'tagA_2.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_1),
self._normalized_path_for_filename(
'tagB.csv', GcsfsDirectIngestFileType.INGEST_VIEW,
self._DAY_1_TIME_3),
]
for path in paths:
fixture_util.add_direct_ingest_path(self.fs.gcs_file_system,
path,
has_fixture=False)
self._process_jobs_for_paths_with_no_gaps_in_expected_order(paths)
self.assertIsNone(self.prioritizer.get_next_job_args())
self.assertFalse(
self.prioritizer.are_more_jobs_expected_for_day(
self._DAY_1.isoformat()))
class DownloadFilesFromSftpController:
"""Class for interacting with and downloading files from SFTP servers."""
def __init__(
self,
project_id: str,
region: str,
lower_bound_update_datetime: Optional[datetime.datetime],
gcs_destination_path: Optional[GcsfsDirectoryPath] = None,
):
self.project_id = project_id
self.region = region.lower()
self.auth = SftpAuth.for_region(region)
self.delegate = SftpDownloadDelegateFactory.build(region_code=region)
self.gcsfs = DirectIngestGCSFileSystem(GcsfsFactory.build())
self.unable_to_download_items: List[str] = []
self.downloaded_items: List[Tuple[str, datetime.datetime]] = []
self.skipped_files: List[str] = []
self.lower_bound_update_datetime = lower_bound_update_datetime
self.bucket = (gcsfs_sftp_download_bucket_path_for_region(
region, SystemLevel.STATE, project_id=self.project_id) if
gcs_destination_path is None else gcs_destination_path)
self.download_dir = GcsfsDirectoryPath.from_dir_and_subdir(
dir_path=self.bucket, subdir=RAW_INGEST_DIRECTORY)
self.postgres_direct_ingest_file_metadata_manager = (
PostgresDirectIngestRawFileMetadataManager(
region,
DirectIngestInstance.PRIMARY.database_version(
SystemLevel.STATE,
state_code=StateCode(self.region.upper())).name,
))
def _is_after_update_bound(self, sftp_attr: SFTPAttributes) -> bool:
if self.lower_bound_update_datetime is None:
return True
update_time = datetime.datetime.fromtimestamp(sftp_attr.st_mtime)
if self.lower_bound_update_datetime.tzinfo:
update_time = update_time.astimezone(pytz.UTC)
return update_time >= self.lower_bound_update_datetime
def _fetch(
self,
connection: pysftp.Connection,
file_path: str,
file_timestamp: datetime.datetime,
) -> None:
"""Fetches data files from the SFTP, tracking which items downloaded and failed to download."""
normalized_sftp_path = os.path.normpath(file_path)
normalized_upload_path = GcsfsFilePath.from_directory_and_file_name(
dir_path=self.download_dir,
file_name=os.path.basename(
to_normalized_unprocessed_file_path(
normalized_sftp_path,
file_type=GcsfsDirectIngestFileType.RAW_DATA,
dt=file_timestamp,
)),
)
if not self.postgres_direct_ingest_file_metadata_manager.has_raw_file_been_discovered(
normalized_upload_path
) and not self.postgres_direct_ingest_file_metadata_manager.has_raw_file_been_processed(
normalized_upload_path):
logging.info("Downloading %s into %s", normalized_sftp_path,
self.download_dir)
try:
path = GcsfsFilePath.from_directory_and_file_name(
dir_path=self.download_dir, file_name=normalized_sftp_path)
self.gcsfs.upload_from_contents_handle_stream(
path=path,
contents_handle=GcsfsSftpFileContentsHandle(
sftp_connection=connection, local_file_path=file_path),
content_type=BYTES_CONTENT_TYPE,
)
logging.info("Post processing %s", path.uri())
self.downloaded_items.append((
self.delegate.post_process_downloads(path, self.gcsfs),
file_timestamp,
))
except IOError as e:
logging.info(
"Could not download %s into %s: %s",
normalized_sftp_path,
self.download_dir,
e.args,
)
self.unable_to_download_items.append(file_path)
else:
logging.info(
"Skipping downloading %s because it has already been previously downloaded for ingest.",
normalized_sftp_path,
)
self.skipped_files.append(file_path)
def get_paths_to_download(self) -> List[Tuple[str, datetime.datetime]]:
"""Opens a connection to SFTP and based on the delegate, find and recursively list items
that are after the update bound and match the delegate's criteria, returning items and
corresponding timestamps that are to be downloaded."""
with pysftp.Connection(
host=self.auth.hostname,
username=self.auth.username,
password=self.auth.password,
cnopts=self.auth.connection_options,
) as connection:
remote_dirs = connection.listdir()
root = self.delegate.root_directory(remote_dirs)
dirs_with_attributes = connection.listdir_attr(root)
paths_post_timestamp = {
sftp_attr.filename: datetime.datetime.fromtimestamp(
sftp_attr.st_mtime).astimezone(pytz.UTC)
for sftp_attr in dirs_with_attributes
if self._is_after_update_bound(sftp_attr)
}
paths_to_download = self.delegate.filter_paths(
list(paths_post_timestamp.keys()))
files_to_download_with_timestamps: List[Tuple[
str, datetime.datetime]] = []
for path in paths_to_download:
file_timestamp = paths_post_timestamp[path]
if connection.isdir(path):
def set_file(file_to_fetch: str,
file_timestamp: datetime.datetime) -> None:
files_to_download_with_timestamps.append(
(file_to_fetch, file_timestamp))
connection.walktree(
remotepath=path,
fcallback=partial(set_file,
file_timestamp=file_timestamp),
dcallback=lambda _: None,
ucallback=self.unable_to_download_items.append,
recurse=True,
)
else:
files_to_download_with_timestamps.append(
(path, file_timestamp))
return files_to_download_with_timestamps
def clean_up(self) -> None:
"""Attempts to recursively remove any downloaded folders created as part of do_fetch."""
try:
logging.info("Cleaning up items in %s.", self.download_dir.uri())
files_to_delete = self.gcsfs.ls_with_blob_prefix(
bucket_name=self.bucket.abs_path(),
blob_prefix=RAW_INGEST_DIRECTORY)
for file in files_to_delete:
self.gcsfs.delete(
GcsfsFilePath.from_absolute_path(file.abs_path()))
except Exception as e:
logging.info(
"%s could not be cleaned up due to an error %s.",
self.download_dir.uri(),
e.args,
)
def fetch_files(
self, files_to_download_with_timestamps: List[Tuple[str,
datetime.datetime]]
) -> None:
"""Opens up one connection and loops through all of the files with timestamps to upload
to the GCS bucket."""
with pysftp.Connection(
host=self.auth.hostname,
username=self.auth.username,
password=self.auth.password,
cnopts=self.auth.connection_options,
) as connection:
for file_path, file_timestamp in files_to_download_with_timestamps:
self._fetch(connection, file_path, file_timestamp)
def do_fetch(
self,
) -> MultiRequestResultWithSkipped[Tuple[str, datetime.datetime], str,
str]:
"""Attempts to open an SFTP connection and download items, returning the corresponding paths
and the timestamp associated, and also any unable to be downloaded."""
logging.info(
"Downloading raw files from SFTP server [%s] to ingest bucket [%s] for project [%s]",
self.auth.hostname,
self.bucket.uri(),
self.project_id,
)
files_to_download_with_timestamps = self.get_paths_to_download()
logging.info(
"Found %s items to download from SFTP server [%s] to upload to ingest bucket [%s]",
len(files_to_download_with_timestamps),
self.auth.hostname,
self.bucket,
)
self.fetch_files(files_to_download_with_timestamps)
logging.info(
"Download complete, successfully downloaded %s files to ingest bucket [%s] "
"could not download %s files and skipped %s files",
len(self.downloaded_items),
self.download_dir.uri(),
len(self.unable_to_download_items),
len(self.skipped_files),
)
return MultiRequestResultWithSkipped(
successes=self.downloaded_items,
failures=self.unable_to_download_items,
skipped=self.skipped_files,
)
class IngestOperationsStore:
"""
A store for tracking the current state of direct ingest.
"""
def __init__(self, override_project_id: Optional[str] = None) -> None:
self.project_id = (metadata.project_id() if override_project_id is None
else override_project_id)
self.fs = DirectIngestGCSFileSystem(GcsfsFactory.build())
self.cloud_task_manager = DirectIngestCloudTaskManagerImpl()
self.cloud_tasks_client = tasks_v2.CloudTasksClient()
@property
def state_codes_launched_in_env(self) -> List[StateCode]:
return get_direct_ingest_states_launched_in_env()
@staticmethod
def get_queues_for_region(state_code: StateCode) -> List[str]:
"""Returns the list of formatted direct ingest queues for given state"""
queues = set()
for ingest_instance in DirectIngestInstance:
queues.update(
get_direct_ingest_queues_for_state(state_code,
ingest_instance))
return list(queues)
def start_ingest_run(self, state_code: StateCode,
instance_str: str) -> None:
"""This function is called through the Ingest Operations UI in the admin panel.
It calls to start a direct ingest run for the given region_code in the given instance
Requires:
- state_code: (required) State code to start ingest for (i.e. "US_ID")
- instance: (required) Which instance to start ingest for (either PRIMARY or SECONDARY)
"""
try:
instance = DirectIngestInstance[instance_str]
except KeyError as e:
logging.error("Received an invalid instance: %s.", instance_str)
raise ValueError(
f"Invalid instance [{instance_str}] received", ) from e
can_start_ingest = state_code in self.state_codes_launched_in_env
formatted_state_code = state_code.value.lower()
region = get_region(formatted_state_code, is_direct_ingest=True)
# Get the ingest bucket for this region and instance
ingest_bucket_path = gcsfs_direct_ingest_bucket_for_region(
region_code=formatted_state_code,
system_level=SystemLevel.for_region(region),
ingest_instance=instance,
project_id=self.project_id,
)
logging.info(
"Creating cloud task to schedule next job and kick ingest for %s instance in %s.",
instance,
formatted_state_code,
)
self.cloud_task_manager.create_direct_ingest_handle_new_files_task(
region=region,
ingest_instance=instance,
ingest_bucket=ingest_bucket_path,
can_start_ingest=can_start_ingest,
)
def update_ingest_queues_state(self, state_code: StateCode,
new_queue_state: str) -> None:
"""This function is called through the Ingest Operations UI in the admin panel.
It updates the state of the following queues by either pausing or resuming the queues:
- direct-ingest-state-<region_code>-bq-import-export
- direct-ingest-state-<region_code>-process-job-queue
- direct-ingest-state-<region_code>-scheduler
- direct-ingest-state-<region_code>-sftp-queue (for select regions)
Requires:
- state_code: (required) State code to pause queues for
- new_state: (required) Either 'PAUSED' or 'RUNNING'
"""
queues_to_update = self.get_queues_for_region(state_code)
if new_queue_state not in [
QUEUE_STATE_ENUM.RUNNING.name,
QUEUE_STATE_ENUM.PAUSED.name,
]:
logging.error(
"Received an invalid queue state: %s. This method should only be used "
"to update queue states to PAUSED or RUNNING",
new_queue_state,
)
raise ValueError(
f"Invalid queue state [{new_queue_state}] received", )
for queue in queues_to_update:
queue_path = self.cloud_tasks_client.queue_path(
self.project_id, _TASK_LOCATION, queue)
if new_queue_state == QUEUE_STATE_ENUM.PAUSED.name:
logging.info("Pausing queue: %s", new_queue_state)
self.cloud_tasks_client.pause_queue(name=queue_path)
else:
logging.info("Resuming queue: %s", new_queue_state)
self.cloud_tasks_client.resume_queue(name=queue_path)
def get_ingest_queue_states(self,
state_code: StateCode) -> List[Dict[str, str]]:
"""Returns a list of dictionaries that contain the name and states of direct ingest queues for a given region"""
ingest_queue_states: List[Dict[str, str]] = []
queues_to_update = self.get_queues_for_region(state_code)
for queue_name in queues_to_update:
queue_path = self.cloud_tasks_client.queue_path(
self.project_id, _TASK_LOCATION, queue_name)
queue = self.cloud_tasks_client.get_queue(name=queue_path)
queue_state = {
"name": queue_name,
"state": QUEUE_STATE_ENUM(queue.state).name,
}
ingest_queue_states.append(queue_state)
return ingest_queue_states
def _get_bucket_metadata(self, path: GcsfsBucketPath) -> BucketSummaryType:
"""Returns a dictionary containing the following info for a given bucket:
i.e. {
name: bucket_name,
unprocessedFilesRaw: how many unprocessed raw data files in the bucket,
processedFilesRaw: how many processed raw data files are in the bucket (should be zero),
unprocessedFilesIngestView: how many unprocessed ingest view files in the bucket,
processedFilesIngestView: how many processed ingest view files are in the bucket
}
"""
bucket_metadata: BucketSummaryType = {
"name": path.abs_path(),
}
for file_type in GcsfsDirectIngestFileType:
file_type_str = self.get_file_type_api_string(file_type)
unprocessed_files = self.fs.get_unprocessed_file_paths(
path, file_type)
bucket_metadata[f"unprocessedFiles{file_type_str}"] = len(
unprocessed_files)
processed_files = self.fs.get_processed_file_paths(path, file_type)
bucket_metadata[f"processedFiles{file_type_str}"] = len(
processed_files)
return bucket_metadata
@staticmethod
def get_file_type_api_string(file_type: GcsfsDirectIngestFileType) -> str:
"""Get the string representation of the file type to use in the response."""
if file_type == GcsfsDirectIngestFileType.INGEST_VIEW:
return "IngestView"
if file_type == GcsfsDirectIngestFileType.RAW_DATA:
return "Raw"
raise ValueError(f"Unexpected file_type [{file_type}]")
def get_ingest_instance_summaries(
self, state_code: StateCode) -> List[Dict[str, Any]]:
"""Returns a list of dictionaries containing the following info for a given instance:
i.e. {
instance: the direct ingest instance,
dbName: database name for this instance,
storage: storage bucket absolute path,
ingest: {
name: bucket_name,
unprocessedFilesRaw: how many unprocessed raw data files in the bucket,
processedFilesRaw: how many processed raw data files are in the bucket (should be zero),
unprocessedFilesIngestView: how many unprocessed ingest view files in the bucket,
processedFilesIngestView: how many processed ingest view files are in the bucket (should be zero),
},
operations: {
unprocessedFilesRaw: number of unprocessed raw files in the operations database
unprocessedFilesIngestView: number of unprocessed ingest view files in the operations database
dateOfEarliestUnprocessedIngestView: date of earliest unprocessed ingest file, if it exists
}
}
"""
formatted_state_code = state_code.value.lower()
ingest_instance_summaries: List[Dict[str, Any]] = []
for instance in DirectIngestInstance:
# Get the ingest bucket path
ingest_bucket_path = gcsfs_direct_ingest_bucket_for_region(
region_code=formatted_state_code,
system_level=SystemLevel.STATE,
ingest_instance=instance,
project_id=self.project_id,
)
# Get an object containing information about the ingest bucket
ingest_bucket_metadata = self._get_bucket_metadata(
ingest_bucket_path)
# Get the storage bucket for this instance
storage_bucket_path = gcsfs_direct_ingest_storage_directory_path_for_region(
region_code=formatted_state_code,
system_level=SystemLevel.STATE,
ingest_instance=instance,
project_id=self.project_id,
)
# Get the database name corresponding to this instance
ingest_db_name = self._get_database_name_for_state(
state_code, instance)
# Get the operations metadata for this ingest instance
operations_db_metadata = self._get_operations_db_metadata(
state_code, ingest_db_name)
ingest_instance_summary: Dict[str, Any] = {
"instance": instance.value,
"storage": storage_bucket_path.abs_path(),
"ingest": ingest_bucket_metadata,
"dbName": ingest_db_name,
"operations": operations_db_metadata,
}
ingest_instance_summaries.append(ingest_instance_summary)
return ingest_instance_summaries
@staticmethod
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
@staticmethod
def _get_operations_db_metadata(
state_code: StateCode,
ingest_db_name: str) -> Dict[str, Union[int, Optional[datetime]]]:
"""Returns the following dictionary with information about the operations database for the state:
{
unprocessedFilesRaw: <int>
unprocessedFilesIngestView: <int>
dateOfEarliestUnprocessedIngestView: <datetime>
}
If running locally, this does not hit the live DB instance and only returns fake data.
"""
if in_development():
return {
"unprocessedFilesRaw": -1,
"unprocessedFilesIngestView": -2,
"dateOfEarliestUnprocessedIngestView": datetime(2021, 4, 28),
}
file_metadata_manager = PostgresDirectIngestFileMetadataManager(
region_code=state_code.value,
ingest_database_name=ingest_db_name,
)
try:
# Raw files are processed in the primary instance, not secondary
num_unprocessed_raw_files = (
file_metadata_manager.get_num_unprocessed_raw_files())
except DirectIngestInstanceError as _:
num_unprocessed_raw_files = 0
return {
"unprocessedFilesRaw":
num_unprocessed_raw_files,
"unprocessedFilesIngestView":
file_metadata_manager.get_num_unprocessed_ingest_files(),
"dateOfEarliestUnprocessedIngestView":
file_metadata_manager.get_date_of_earliest_unprocessed_ingest_file(
),
}
def __init__(self, override_project_id: Optional[str] = None) -> None:
self.project_id = (metadata.project_id() if override_project_id is None
else override_project_id)
self.fs = DirectIngestGCSFileSystem(GcsfsFactory.build())
self.cloud_task_manager = DirectIngestCloudTaskManagerImpl()
self.cloud_tasks_client = tasks_v2.CloudTasksClient()
class UploadStateFilesToIngestBucketController:
"""Class for uploading files from a local filesystem to a region's ingest bucket."""
SUPPORTED_EXTENSIONS: List[str] = [".csv", ".txt"]
def __init__(
self,
paths_with_timestamps: List[Tuple[str, datetime.datetime]],
project_id: str,
region: str,
gcs_destination_path: Optional[str] = None,
):
self.paths_with_timestamps = paths_with_timestamps
self.project_id = project_id
self.region = region.lower()
self.gcsfs = DirectIngestGCSFileSystem(GcsfsFactory.build())
self.gcs_destination_path = (
GcsfsDirectoryPath.from_absolute_path(
gcsfs_direct_ingest_directory_path_for_region(
region, SystemLevel.STATE, project_id=self.project_id))
if gcs_destination_path is None else
GcsfsDirectoryPath.from_absolute_path(gcs_destination_path))
self.uploaded_files: List[str] = []
self.unable_to_upload_files: List[str] = []
def _copy_to_ingest_bucket(self, path: str,
full_file_upload_path: GcsfsFilePath) -> None:
try:
mimetype, _ = guess_type(os.path.basename(path))
self.gcsfs.mv(
src_path=GcsfsFilePath.from_absolute_path(path),
dst_path=full_file_upload_path,
)
self.gcsfs.set_content_type(full_file_upload_path,
mimetype if mimetype else "text/plain")
logging.info("Copied %s -> %s", path, full_file_upload_path.uri())
self.uploaded_files.append(path)
except BaseException as e:
logging.warning(
"Could not copy %s -> %s due to error %s",
path,
full_file_upload_path.uri(),
e.args,
)
self.unable_to_upload_files.append(path)
def _upload_file(
self, path_with_timestamp: Tuple[str, datetime.datetime]) -> None:
path, timestamp = path_with_timestamp
normalized_file_name = os.path.basename(
to_normalized_unprocessed_file_path(
path,
file_type=GcsfsDirectIngestFileType.RAW_DATA,
dt=timestamp))
full_file_upload_path = GcsfsFilePath.from_directory_and_file_name(
self.gcs_destination_path, normalized_file_name)
self._copy_to_ingest_bucket(path, full_file_upload_path)
def get_paths_to_upload(self) -> List[Tuple[str, datetime.datetime]]:
"""Returns the appropriate paths to upload and the proper associated timestamp that
it is to be normalized with. Skips any files that are not properly supported."""
path_candidates = []
for path, timestamp in self.paths_with_timestamps:
if self.gcsfs.is_dir(path):
directory = GcsfsDirectoryPath.from_absolute_path(path)
files_in_directory = self.gcsfs.ls_with_blob_prefix(
bucket_name=directory.bucket_name,
blob_prefix=directory.relative_path,
)
for file in files_in_directory:
path_candidates.append((file.abs_path(), timestamp))
elif self.gcsfs.is_file(path):
file = GcsfsFilePath.from_absolute_path(path)
path_candidates.append((file.abs_path(), timestamp))
else:
logging.warning(
"Could not indicate %s as a directory or a file in %s. Skipping",
path,
self.gcs_destination_path.uri(),
)
self.unable_to_upload_files.append(path)
continue
result = []
for path, timestamp in path_candidates:
_, ext = os.path.splitext(path)
if not ext or ext not in self.SUPPORTED_EXTENSIONS:
logging.info("Skipping file [%s] - invalid extension %s", path,
ext)
continue
result.append((path, timestamp))
return result
def upload_files(
self,
paths_with_timestamps_to_upload: List[Tuple[str, datetime.datetime]],
thread_pool: ThreadPool,
) -> None:
thread_pool.map(self._upload_file, paths_with_timestamps_to_upload)
def do_upload(self) -> Tuple[List[str], List[str]]:
"""Perform upload to ingest bucket."""
logging.info(
"Uploading raw files to the %s ingest bucket [%s] for project [%s].",
self.region,
self.gcs_destination_path.uri(),
self.project_id,
)
paths_with_timestamps_to_upload = self.get_paths_to_upload()
logging.info(
"Found %s items to upload to ingest bucket [%s]",
len(paths_with_timestamps_to_upload),
self.gcs_destination_path.uri(),
)
thread_pool = ThreadPool(processes=12)
self.upload_files(paths_with_timestamps_to_upload, thread_pool)
thread_pool.close()
thread_pool.join()
logging.info(
"Upload complete, successfully uploaded %s files to ingest bucket [%s], "
"could not upload %s files",
len(self.uploaded_files),
self.gcs_destination_path.uri(),
len(self.unable_to_upload_files),
)
return self.uploaded_files, self.unable_to_upload_files
def _make_processed_path(
self, path_str: str,
file_type: GcsfsDirectIngestFileType) -> GcsfsFilePath:
path = self._make_unprocessed_path(path_str, file_type)
# pylint:disable=protected-access
return DirectIngestGCSFileSystem._to_processed_file_path(path)
def setUp(self) -> None:
self.fs = DirectIngestGCSFileSystem(FakeGCSFileSystem())
class TestDirectIngestGcsFileSystem(TestCase):
"""Tests for the FakeGCSFileSystem."""
STORAGE_DIR_PATH = GcsfsDirectoryPath(bucket_name='storage_bucket',
relative_path='region_subdir')
INGEST_DIR_PATH = GcsfsDirectoryPath(bucket_name='my_bucket')
def setUp(self) -> None:
self.fs = DirectIngestGCSFileSystem(FakeGCSFileSystem())
def fully_process_file(self,
dt: datetime.datetime,
path: GcsfsFilePath,
file_type_differentiation_on: bool = False) -> None:
"""Mimics all the file system calls for a single file in the direct
ingest system, from getting added to the ingest bucket, turning to a
processed file, then getting moved to storage."""
fixture_util.add_direct_ingest_path(self.fs.gcs_file_system,
path,
has_fixture=False)
start_num_total_files = len(self.fs.gcs_file_system.all_paths)
# pylint: disable=protected-access
start_ingest_paths = self.fs._ls_with_file_prefix(
self.INGEST_DIR_PATH, '', None)
start_storage_paths = self.fs._ls_with_file_prefix(
self.STORAGE_DIR_PATH, '', None)
if file_type_differentiation_on:
start_raw_storage_paths = self.fs._ls_with_file_prefix(
self.STORAGE_DIR_PATH,
'',
file_type_filter=GcsfsDirectIngestFileType.RAW_DATA)
start_ingest_view_storage_paths = self.fs._ls_with_file_prefix(
self.STORAGE_DIR_PATH,
'',
file_type_filter=GcsfsDirectIngestFileType.INGEST_VIEW)
else:
start_raw_storage_paths = []
start_ingest_view_storage_paths = []
# File is renamed to normalized path
file_type = GcsfsDirectIngestFileType.RAW_DATA \
if file_type_differentiation_on else GcsfsDirectIngestFileType.UNSPECIFIED
self.fs.mv_path_to_normalized_path(path, file_type, dt)
if file_type_differentiation_on:
raw_unprocessed = self.fs.get_unprocessed_file_paths(
self.INGEST_DIR_PATH,
file_type_filter=GcsfsDirectIngestFileType.RAW_DATA)
self.assertEqual(len(raw_unprocessed), 1)
self.assertTrue(
self.fs.is_seen_unprocessed_file(raw_unprocessed[0]))
# ... raw file imported to BQ
processed_path = self.fs.mv_path_to_processed_path(
raw_unprocessed[0])
processed = self.fs.get_processed_file_paths(
self.INGEST_DIR_PATH, None)
self.assertEqual(len(processed), 1)
self.fs.copy(
processed_path,
GcsfsFilePath.from_absolute_path(
to_normalized_unprocessed_file_path_from_normalized_path(
processed_path.abs_path(),
file_type_override=GcsfsDirectIngestFileType.
INGEST_VIEW)))
self.fs.mv_path_to_storage(processed_path, self.STORAGE_DIR_PATH)
ingest_unprocessed_filter = GcsfsDirectIngestFileType.INGEST_VIEW if file_type_differentiation_on else None
ingest_unprocessed = self.fs.get_unprocessed_file_paths(
self.INGEST_DIR_PATH, file_type_filter=ingest_unprocessed_filter)
self.assertEqual(len(ingest_unprocessed), 1)
self.assertTrue(self.fs.is_seen_unprocessed_file(
ingest_unprocessed[0]))
# ... file is ingested
# File is moved to processed path
self.fs.mv_path_to_processed_path(ingest_unprocessed[0])
processed = self.fs.get_processed_file_paths(self.INGEST_DIR_PATH,
None)
self.assertEqual(len(processed), 1)
self.assertTrue(self.fs.is_processed_file(processed[0]))
unprocessed = self.fs.get_unprocessed_file_paths(
self.INGEST_DIR_PATH, None)
self.assertEqual(len(unprocessed), 0)
# File is moved to storage
ingest_move_type_filter = GcsfsDirectIngestFileType.INGEST_VIEW \
if file_type_differentiation_on else None
self.fs.mv_processed_paths_before_date_to_storage(
self.INGEST_DIR_PATH,
self.STORAGE_DIR_PATH,
date_str_bound=dt.date().isoformat(),
include_bound=True,
file_type_filter=ingest_move_type_filter)
end_ingest_paths = self.fs._ls_with_file_prefix(self.INGEST_DIR_PATH,
'',
file_type_filter=None)
end_storage_paths = self.fs._ls_with_file_prefix(self.STORAGE_DIR_PATH,
'',
file_type_filter=None)
if file_type_differentiation_on:
end_raw_storage_paths = self.fs._ls_with_file_prefix(
self.STORAGE_DIR_PATH,
'',
file_type_filter=GcsfsDirectIngestFileType.RAW_DATA)
end_ingest_view_storage_paths = self.fs._ls_with_file_prefix(
self.STORAGE_DIR_PATH,
'',
file_type_filter=GcsfsDirectIngestFileType.INGEST_VIEW)
else:
end_raw_storage_paths = []
end_ingest_view_storage_paths = []
# Each file gets re-exported as ingest view
splitting_factor = 2 if file_type_differentiation_on else 1
expected_final_total_files = start_num_total_files + splitting_factor - 1
self.assertEqual(len(self.fs.gcs_file_system.all_paths),
expected_final_total_files)
self.assertEqual(len(end_ingest_paths), len(start_ingest_paths) - 1)
self.assertEqual(len(end_storage_paths),
len(start_storage_paths) + 1 * splitting_factor)
if file_type_differentiation_on:
self.assertEqual(
len(end_raw_storage_paths) +
len(end_ingest_view_storage_paths), len(end_storage_paths))
self.assertEqual(len(end_raw_storage_paths),
len(start_raw_storage_paths) + 1)
self.assertEqual(len(end_ingest_view_storage_paths),
len(start_ingest_view_storage_paths) + 1)
for sp in end_storage_paths:
parts = filename_parts_from_path(sp)
if sp.abs_path() not in {
p.abs_path()
for p in start_storage_paths
}:
self.assertTrue(sp.abs_path().startswith(
self.STORAGE_DIR_PATH.abs_path()))
dir_path, storage_file_name = os.path.split(sp.abs_path())
if parts.file_type != GcsfsDirectIngestFileType.UNSPECIFIED:
self.assertTrue(parts.file_type.value in dir_path)
name, _ = path.file_name.split('.')
self.assertTrue(name in storage_file_name)
def test_direct_ingest_file_moves(self) -> None:
self.fully_process_file(
datetime.datetime.now(),
GcsfsFilePath(bucket_name='my_bucket', blob_name='test_file.csv'))
def test_direct_ingest_multiple_file_moves(self) -> None:
self.fully_process_file(
datetime.datetime.now(),
GcsfsFilePath(bucket_name='my_bucket', blob_name='test_file.csv'))
self.fully_process_file(
datetime.datetime.now(),
GcsfsFilePath(bucket_name='my_bucket',
blob_name='test_file_2.csv'))
def test_move_to_storage_with_conflict(self) -> None:
dt = datetime.datetime.now()
self.fully_process_file(
dt,
GcsfsFilePath(bucket_name='my_bucket', blob_name='test_file.csv'))
# Try uploading a file with a duplicate name that has already been
# moved to storage
self.fully_process_file(
dt,
GcsfsFilePath(bucket_name='my_bucket', blob_name='test_file.csv'))
# pylint: disable=protected-access
storage_paths = self.fs._ls_with_file_prefix(self.STORAGE_DIR_PATH,
'',
file_type_filter=None)
self.assertEqual(len(storage_paths), 2)
found_first_file = False
found_second_file = False
for path in storage_paths:
self.assertTrue(filename_parts_from_path(path))
if path.abs_path().endswith('test_file.csv'):
found_first_file = True
if path.abs_path().endswith('test_file-(1).csv'):
found_second_file = True
self.assertTrue(found_first_file)
self.assertTrue(found_second_file)
def test_direct_ingest_file_moves_with_file_types(self) -> None:
self.fully_process_file(datetime.datetime.now(),
GcsfsFilePath(bucket_name='my_bucket',
blob_name='test_file.csv'),
file_type_differentiation_on=True)
def test_direct_ingest_multiple_file_moves_with_file_types(self) -> None:
self.fully_process_file(datetime.datetime.now(),
GcsfsFilePath(bucket_name='my_bucket',
blob_name='test_file.csv'),
file_type_differentiation_on=True)
self.fully_process_file(datetime.datetime.now(),
GcsfsFilePath(bucket_name='my_bucket',
blob_name='test_file_2.csv'),
file_type_differentiation_on=True)
def test_move_to_storage_with_conflict_with_file_types(self) -> None:
dt = datetime.datetime.now()
self.fully_process_file(dt,
GcsfsFilePath(bucket_name='my_bucket',
blob_name='test_file.csv'),
file_type_differentiation_on=True)
# Try uploading a file with a duplicate name that has already been
# moved to storage
self.fully_process_file(dt,
GcsfsFilePath(bucket_name='my_bucket',
blob_name='test_file.csv'),
file_type_differentiation_on=True)
# pylint: disable=protected-access
storage_paths = self.fs._ls_with_file_prefix(self.STORAGE_DIR_PATH,
'',
file_type_filter=None)
self.assertEqual(len(storage_paths), 4)
found_first_file = False
found_second_file = False
for path in storage_paths:
if path.abs_path().endswith('test_file.csv'):
found_first_file = True
if path.abs_path().endswith('test_file-(1).csv'):
found_second_file = True
self.assertTrue(found_first_file)
self.assertTrue(found_second_file)
class GcsfsDirectIngestController(
BaseDirectIngestController[GcsfsIngestArgs, GcsfsFileContentsHandle]):
"""Controller for parsing and persisting a file in the GCS filesystem."""
_MAX_STORAGE_FILE_RENAME_TRIES = 10
_DEFAULT_MAX_PROCESS_JOB_WAIT_TIME_SEC = 300
_INGEST_FILE_SPLIT_LINE_LIMIT = 2500
def __init__(
self,
region_name: str,
system_level: SystemLevel,
ingest_directory_path: Optional[str] = None,
storage_directory_path: Optional[str] = None,
max_delay_sec_between_files: Optional[int] = None,
):
super().__init__(region_name, system_level)
self.fs = DirectIngestGCSFileSystem(GcsfsFactory.build())
self.max_delay_sec_between_files = max_delay_sec_between_files
if not ingest_directory_path:
ingest_directory_path = gcsfs_direct_ingest_directory_path_for_region(
region_name, system_level)
self.ingest_directory_path = GcsfsDirectoryPath.from_absolute_path(
ingest_directory_path)
if not storage_directory_path:
storage_directory_path = (
gcsfs_direct_ingest_storage_directory_path_for_region(
region_name, system_level))
self.storage_directory_path = GcsfsDirectoryPath.from_absolute_path(
storage_directory_path)
self.temp_output_directory_path = GcsfsDirectoryPath.from_absolute_path(
gcsfs_direct_ingest_temporary_output_directory_path())
ingest_job_file_type_filter = (
GcsfsDirectIngestFileType.INGEST_VIEW
if self.region.is_raw_vs_ingest_file_name_detection_enabled() else
None)
self.file_prioritizer = GcsfsDirectIngestJobPrioritizer(
self.fs,
self.ingest_directory_path,
self.get_file_tag_rank_list(),
ingest_job_file_type_filter,
)
self.ingest_file_split_line_limit = self._INGEST_FILE_SPLIT_LINE_LIMIT
self.file_metadata_manager = PostgresDirectIngestFileMetadataManager(
region_code=self.region.region_code)
self.raw_file_import_manager = DirectIngestRawFileImportManager(
region=self.region,
fs=self.fs,
ingest_directory_path=self.ingest_directory_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,
ingest_directory_path=self.ingest_directory_path,
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(),
)
# ================= #
# NEW FILE HANDLING #
# ================= #
def handle_file(self, path: GcsfsFilePath, start_ingest: bool) -> None:
"""Called when a single new file is added to an ingest bucket (may also
be called as a result of a rename).
May be called from any worker/queue.
"""
if self.fs.is_processed_file(path):
logging.info("File [%s] is already processed, returning.",
path.abs_path())
return
if self.fs.is_normalized_file_path(path):
parts = filename_parts_from_path(path)
if (parts.is_file_split and parts.file_split_size
and parts.file_split_size <=
self.ingest_file_split_line_limit):
self.kick_scheduler(just_finished_job=False)
logging.info(
"File [%s] is already normalized and split split "
"with correct size, kicking scheduler.",
path.abs_path(),
)
return
logging.info("Creating cloud task to schedule next job.")
self.cloud_task_manager.create_direct_ingest_handle_new_files_task(
region=self.region, can_start_ingest=start_ingest)
def _register_all_new_paths_in_metadata(
self, paths: List[GcsfsFilePath]) -> None:
for path in paths:
if not self.file_metadata_manager.has_file_been_discovered(path):
self.file_metadata_manager.mark_file_as_discovered(path)
@trace.span
def handle_new_files(self, can_start_ingest: bool) -> None:
"""Searches the ingest directory for new/unprocessed files. Normalizes
file names and splits files as necessary, schedules the next ingest job
if allowed.
Should only be called from the scheduler queue.
"""
if not can_start_ingest and self.region.is_ingest_launched_in_env():
raise ValueError(
"The can_start_ingest flag should only be used for regions where ingest is not yet launched in a "
"particular environment. If we want to be able to selectively pause ingest processing for a state, we "
"will first have to build a config that is respected by both the /ensure_all_file_paths_normalized "
"endpoint and any cloud functions that trigger ingest.")
unnormalized_paths = self.fs.get_unnormalized_file_paths(
self.ingest_directory_path)
unnormalized_path_file_type = (
GcsfsDirectIngestFileType.RAW_DATA
if self.region.is_raw_vs_ingest_file_name_detection_enabled() else
GcsfsDirectIngestFileType.UNSPECIFIED)
for path in unnormalized_paths:
logging.info("File [%s] is not yet seen, normalizing.",
path.abs_path())
self.fs.mv_path_to_normalized_path(
path, file_type=unnormalized_path_file_type)
if unnormalized_paths:
logging.info(
"Normalized at least one path - returning, will handle "
"normalized files separately.")
# Normalizing file paths will cause the cloud function that calls
# this function to be re-triggered.
return
if not can_start_ingest:
logging.warning(
"Ingest not configured to start post-file normalization - returning."
)
return
check_is_region_launched_in_env(self.region)
unprocessed_raw_paths = []
ingest_file_type_filter = (
GcsfsDirectIngestFileType.INGEST_VIEW
if self.region.is_raw_vs_ingest_file_name_detection_enabled() else
None)
unprocessed_ingest_view_paths = self.fs.get_unprocessed_file_paths(
self.ingest_directory_path,
file_type_filter=ingest_file_type_filter)
if self.region.is_raw_vs_ingest_file_name_detection_enabled():
unprocessed_raw_paths = self.fs.get_unprocessed_file_paths(
self.ingest_directory_path,
file_type_filter=GcsfsDirectIngestFileType.RAW_DATA,
)
self._register_all_new_paths_in_metadata(unprocessed_raw_paths)
if self.region.are_ingest_view_exports_enabled_in_env():
self._register_all_new_paths_in_metadata(
unprocessed_ingest_view_paths)
unprocessed_paths = unprocessed_raw_paths + unprocessed_ingest_view_paths
did_split = False
for path in unprocessed_ingest_view_paths:
if self._split_file_if_necessary(path):
did_split = True
if did_split:
if self.region.are_ingest_view_exports_enabled_in_env():
post_split_unprocessed_ingest_view_paths = (
self.fs.get_unprocessed_file_paths(
self.ingest_directory_path,
file_type_filter=GcsfsDirectIngestFileType.INGEST_VIEW,
))
self._register_all_new_paths_in_metadata(
post_split_unprocessed_ingest_view_paths)
logging.info(
"Split at least one path - returning, will handle split "
"files separately.")
# Writing new split files to storage will cause the cloud function
# that calls this function to be re-triggered.
return
if unprocessed_paths:
self.schedule_next_ingest_job_or_wait_if_necessary(
just_finished_job=False)
def do_raw_data_import(self,
data_import_args: GcsfsRawDataBQImportArgs) -> None:
"""Process a raw incoming file by importing it to BQ, tracking it in our metadata tables, and moving it to
storage on completion.
"""
check_is_region_launched_in_env(self.region)
if not self.region.are_raw_data_bq_imports_enabled_in_env():
raise ValueError(
f"Raw data imports not enabled for region [{self.region.region_code}]"
)
if not self.fs.exists(data_import_args.raw_data_file_path):
logging.warning(
"File path [%s] no longer exists - might have already been "
"processed or deleted",
data_import_args.raw_data_file_path,
)
self.kick_scheduler(just_finished_job=True)
return
file_metadata = self.file_metadata_manager.get_file_metadata(
data_import_args.raw_data_file_path)
if file_metadata.processed_time:
logging.warning(
"File [%s] is already marked as processed. Skipping file processing.",
data_import_args.raw_data_file_path.file_name,
)
self.kick_scheduler(just_finished_job=True)
return
self.raw_file_import_manager.import_raw_file_to_big_query(
data_import_args.raw_data_file_path, file_metadata)
if not self.region.are_ingest_view_exports_enabled_in_env():
# TODO(#3162) This is a stopgap measure for regions that have only partially launched. Delete once SQL
# pre-processing is enabled for all direct ingest regions.
parts = filename_parts_from_path(
data_import_args.raw_data_file_path)
ingest_file_tags = self.get_file_tag_rank_list()
if parts.file_tag in ingest_file_tags:
self.fs.copy(
data_import_args.raw_data_file_path,
GcsfsFilePath.from_absolute_path(
to_normalized_unprocessed_file_path_from_normalized_path(
data_import_args.raw_data_file_path.abs_path(),
file_type_override=GcsfsDirectIngestFileType.
INGEST_VIEW,
)),
)
processed_path = self.fs.mv_path_to_processed_path(
data_import_args.raw_data_file_path)
self.file_metadata_manager.mark_file_as_processed(
path=data_import_args.raw_data_file_path)
self.fs.mv_path_to_storage(processed_path, self.storage_directory_path)
self.kick_scheduler(just_finished_job=True)
def do_ingest_view_export(
self, ingest_view_export_args: GcsfsIngestViewExportArgs) -> None:
check_is_region_launched_in_env(self.region)
if not self.region.are_ingest_view_exports_enabled_in_env():
raise ValueError(
f"Ingest view exports not enabled for region [{self.region.region_code}]. Passed args: "
f"{ingest_view_export_args}")
did_export = self.ingest_view_export_manager.export_view_for_args(
ingest_view_export_args)
if (not did_export or not self.file_metadata_manager.
get_ingest_view_metadata_pending_export()):
logging.info("Creating cloud task to schedule next job.")
self.cloud_task_manager.create_direct_ingest_handle_new_files_task(
region=self.region, can_start_ingest=True)
# ============== #
# JOB SCHEDULING #
# ============== #
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:
if not self.region.are_raw_data_bq_imports_enabled_in_env():
return False
queue_info = self.cloud_task_manager.get_bq_import_export_queue_info(
self.region)
did_schedule = False
tasks_to_schedule = [
GcsfsRawDataBQImportArgs(path) for path in
self.raw_file_import_manager.get_unprocessed_raw_files_to_import()
]
for task_args in tasks_to_schedule:
# If the file path has not actually been discovered by the metadata manager yet, it likely was just added
# and a subsequent call to handle_files will register it and trigger another call to this function so we can
# schedule the appropriate job.
discovered = self.file_metadata_manager.has_file_been_discovered(
task_args.raw_data_file_path)
if discovered and not queue_info.has_task_already_scheduled(
task_args):
self.cloud_task_manager.create_direct_ingest_raw_data_import_task(
self.region, task_args)
did_schedule = True
return queue_info.has_raw_data_import_jobs_queued() or did_schedule
def _schedule_ingest_view_export_tasks(self) -> bool:
"""Schedules all pending ingest view export tasks for launched ingest view tags, if they have not been
scheduled. If tasks are scheduled or are still running, returns True. Otherwise, if it's safe to proceed with
next steps of ingest, returns False."""
if not self.region.are_ingest_view_exports_enabled_in_env():
return False
queue_info = self.cloud_task_manager.get_bq_import_export_queue_info(
self.region)
if queue_info.has_ingest_view_export_jobs_queued():
# Since we schedule all export jobs at once, after all raw files have been processed, we wait for all of the
# export jobs to be done before checking if we need to schedule more.
return True
did_schedule = False
tasks_to_schedule = (
self.ingest_view_export_manager.get_ingest_view_export_task_args())
rank_list = self.get_file_tag_rank_list()
ingest_view_name_rank = {
ingest_view_name: i
for i, ingest_view_name in enumerate(rank_list)
}
# Filter out views that aren't in ingest view tags.
filtered_tasks_to_schedule = []
for args in tasks_to_schedule:
if args.ingest_view_name not in ingest_view_name_rank:
logging.warning(
"Skipping ingest view task export for [%s] - not in controller ingest tags.",
args.ingest_view_name,
)
continue
filtered_tasks_to_schedule.append(args)
tasks_to_schedule = filtered_tasks_to_schedule
# Sort by tag order and export datetime
tasks_to_schedule.sort(key=lambda args: (
ingest_view_name_rank[args.ingest_view_name],
args.upper_bound_datetime_to_export,
))
for task_args in tasks_to_schedule:
if not queue_info.has_task_already_scheduled(task_args):
self.cloud_task_manager.create_direct_ingest_ingest_view_export_task(
self.region, task_args)
did_schedule = True
return did_schedule
@classmethod
@abc.abstractmethod
def get_file_tag_rank_list(cls) -> List[str]:
pass
def _get_next_job_args(self) -> Optional[GcsfsIngestArgs]:
args = self.file_prioritizer.get_next_job_args()
if not self.region.are_ingest_view_exports_enabled_in_env():
return args
if not args:
return None
discovered = self.file_metadata_manager.has_file_been_discovered(
args.file_path)
if not discovered:
# If the file path has not actually been discovered by the controller yet, it likely was just added and a
# subsequent call to handle_files will register it and trigger another call to this function so we can
# schedule the appropriate job.
logging.info(
"Found args [%s] for a file that has not been discovered by the metadata manager yet - not scheduling.",
args,
)
return None
return args
def _wait_time_sec_for_next_args(self, args: GcsfsIngestArgs) -> int:
if self.file_prioritizer.are_next_args_expected(args):
# Run job immediately
return 0
now = datetime.datetime.utcnow()
file_upload_time: datetime.datetime = filename_parts_from_path(
args.file_path).utc_upload_datetime
max_delay_sec = (self.max_delay_sec_between_files
if self.max_delay_sec_between_files is not None else
self._DEFAULT_MAX_PROCESS_JOB_WAIT_TIME_SEC)
max_wait_from_file_upload_time = file_upload_time + datetime.timedelta(
seconds=max_delay_sec)
if max_wait_from_file_upload_time <= now:
wait_time = 0
else:
wait_time = (max_wait_from_file_upload_time - now).seconds
logging.info("Waiting [%s] sec for [%s]", wait_time,
self._job_tag(args))
return wait_time
def _on_job_scheduled(self, ingest_args: GcsfsIngestArgs) -> None:
pass
# =================== #
# SINGLE JOB RUN CODE #
# =================== #
def _job_tag(self, args: GcsfsIngestArgs) -> str:
return (f"{self.region.region_code}/{args.file_path.file_name}:"
f"{args.ingest_time}")
def _get_contents_handle(
self, args: GcsfsIngestArgs) -> Optional[GcsfsFileContentsHandle]:
return self._get_contents_handle_from_path(args.file_path)
def _get_contents_handle_from_path(
self, path: GcsfsFilePath) -> Optional[GcsfsFileContentsHandle]:
return self.fs.download_to_temp_file(path)
@abc.abstractmethod
def _are_contents_empty(self, args: GcsfsIngestArgs,
contents_handle: GcsfsFileContentsHandle) -> bool:
pass
def _can_proceed_with_ingest_for_contents(
self, args: GcsfsIngestArgs,
contents_handle: GcsfsFileContentsHandle) -> bool:
parts = filename_parts_from_path(args.file_path)
return self._are_contents_empty(
args, contents_handle) or not self._must_split_contents(
parts.file_type, args.file_path)
def _must_split_contents(self, file_type: GcsfsDirectIngestFileType,
path: GcsfsFilePath) -> bool:
if (self.region.is_raw_vs_ingest_file_name_detection_enabled()
and file_type == GcsfsDirectIngestFileType.RAW_DATA):
return False
return not self._file_meets_file_line_limit(
self.ingest_file_split_line_limit, path)
@abc.abstractmethod
def _file_meets_file_line_limit(self, line_limit: int,
path: GcsfsFilePath) -> bool:
"""Subclasses should implement to determine whether the file meets the
expected line limit"""
@abc.abstractmethod
def _parse(self, args: GcsfsIngestArgs,
contents_handle: GcsfsFileContentsHandle) -> IngestInfo:
pass
def _should_split_file(self, path: GcsfsFilePath) -> bool:
"""Returns a handle to the contents of this path if this file should be split, None otherwise."""
parts = filename_parts_from_path(path)
if (self.region.is_raw_vs_ingest_file_name_detection_enabled()
and parts.file_type != GcsfsDirectIngestFileType.INGEST_VIEW):
raise ValueError(
f"Should not be attempting to split files other than ingest view files, found path with "
f"file type: {parts.file_type}")
if parts.file_tag not in self.get_file_tag_rank_list():
logging.info(
"File tag [%s] for path [%s] not in rank list - not splitting.",
parts.file_tag,
path.abs_path(),
)
return False
if (parts.is_file_split and parts.file_split_size and
parts.file_split_size <= self.ingest_file_split_line_limit):
logging.info(
"File [%s] already split with size [%s].",
path.abs_path(),
parts.file_split_size,
)
return False
return self._must_split_contents(parts.file_type, path)
@trace.span
def _split_file_if_necessary(self, path: GcsfsFilePath) -> bool:
"""Checks if the given file needs to be split according to this controller's |file_split_line_limit|.
Returns True if the file was split, False if splitting was not necessary.
"""
should_split = self._should_split_file(path)
if not should_split:
logging.info("No need to split file path [%s].", path.abs_path())
return False
logging.info("Proceeding to file splitting for path [%s].",
path.abs_path())
original_metadata = None
if self.region.are_ingest_view_exports_enabled_in_env():
original_metadata = self.file_metadata_manager.get_file_metadata(
path)
output_dir = GcsfsDirectoryPath.from_file_path(path)
split_contents_paths = self._split_file(path)
upload_paths = []
for i, split_contents_path in enumerate(split_contents_paths):
upload_path = self._create_split_file_path(path,
output_dir,
split_num=i)
logging.info(
"Copying split [%s] to direct ingest directory at path [%s].",
i,
upload_path.abs_path(),
)
upload_paths.append(upload_path)
try:
self.fs.mv(split_contents_path, upload_path)
except Exception as e:
logging.error(
"Threw error while copying split files from temp bucket - attempting to clean up before rethrowing."
" [%s]",
e,
)
for p in upload_paths:
self.fs.delete(p)
raise e
# We wait to register files with metadata manager until all files have been successfully copied to avoid leaving
# the metadata manager in an inconsistent state.
if self.region.are_ingest_view_exports_enabled_in_env():
if not isinstance(original_metadata,
DirectIngestIngestFileMetadata):
raise ValueError(
"Attempting to split a non-ingest view type file")
logging.info(
"Registering [%s] split files with the metadata manager.",
len(upload_paths),
)
for upload_path in upload_paths:
ingest_file_metadata = (
self.file_metadata_manager.register_ingest_file_split(
original_metadata, upload_path))
self.file_metadata_manager.mark_ingest_view_exported(
ingest_file_metadata)
self.file_metadata_manager.mark_file_as_processed(path)
logging.info(
"Done splitting file [%s] into [%s] paths, moving it to storage.",
path.abs_path(),
len(split_contents_paths),
)
self.fs.mv_path_to_storage(path, self.storage_directory_path)
return True
def _create_split_file_path(
self,
original_file_path: GcsfsFilePath,
output_dir: GcsfsDirectoryPath,
split_num: int,
) -> GcsfsFilePath:
parts = filename_parts_from_path(original_file_path)
rank_str = str(split_num + 1).zfill(5)
updated_file_name = (
f"{parts.stripped_file_name}_{rank_str}"
f"_{SPLIT_FILE_SUFFIX}_size{self.ingest_file_split_line_limit}"
f".{parts.extension}")
file_type = (
GcsfsDirectIngestFileType.INGEST_VIEW
if self.region.is_raw_vs_ingest_file_name_detection_enabled() else
GcsfsDirectIngestFileType.UNSPECIFIED)
return GcsfsFilePath.from_directory_and_file_name(
output_dir,
to_normalized_unprocessed_file_path(updated_file_name,
file_type=file_type,
dt=parts.utc_upload_datetime),
)
@abc.abstractmethod
def _split_file(self, path: GcsfsFilePath) -> List[GcsfsFilePath]:
"""Should be implemented by subclasses to split a file accessible via the provided path into multiple
files and upload those files to GCS. Returns the list of upload paths."""
def _do_cleanup(self, args: GcsfsIngestArgs) -> None:
self.fs.mv_path_to_processed_path(args.file_path)
if self.region.are_ingest_view_exports_enabled_in_env():
self.file_metadata_manager.mark_file_as_processed(args.file_path)
parts = filename_parts_from_path(args.file_path)
self._move_processed_files_to_storage_as_necessary(
last_processed_date_str=parts.date_str)
def _is_last_job_for_day(self, args: GcsfsIngestArgs) -> bool:
"""Returns True if the file handled in |args| is the last file for that
upload date."""
parts = filename_parts_from_path(args.file_path)
upload_date, date_str = parts.utc_upload_datetime, parts.date_str
more_jobs_expected = self.file_prioritizer.are_more_jobs_expected_for_day(
date_str)
if more_jobs_expected:
return False
next_job_args = self.file_prioritizer.get_next_job_args(date_str)
if next_job_args:
next_job_date = filename_parts_from_path(
next_job_args.file_path).utc_upload_datetime
return next_job_date > upload_date
return True
def _move_processed_files_to_storage_as_necessary(
self, last_processed_date_str: str) -> None:
"""Moves files that have already been ingested/processed, up to and including the given date, into storage,
if there is nothing more left to ingest/process, i.e. we are not expecting more files."""
next_args = self.file_prioritizer.get_next_job_args()
should_move_last_processed_date = False
if not next_args:
are_more_jobs_expected = (
self.file_prioritizer.are_more_jobs_expected_for_day(
last_processed_date_str))
if not are_more_jobs_expected:
should_move_last_processed_date = True
else:
next_date_str = filename_parts_from_path(
next_args.file_path).date_str
if next_date_str < last_processed_date_str:
logging.info("Found a file [%s] from a date previous to our "
"last processed date - not moving anything to "
"storage.")
return
# If there are still more to process on this day, do not move files
# from this day.
should_move_last_processed_date = next_date_str != last_processed_date_str
# Note: at this point, we expect RAW file type files to already have been moved once they were imported to BQ.
file_type_to_move = (
GcsfsDirectIngestFileType.INGEST_VIEW
if self.region.is_raw_vs_ingest_file_name_detection_enabled() else
None)
self.fs.mv_processed_paths_before_date_to_storage(
self.ingest_directory_path,
self.storage_directory_path,
file_type_filter=file_type_to_move,
date_str_bound=last_processed_date_str,
include_bound=should_move_last_processed_date,
)
@staticmethod
def file_tag(file_path: GcsfsFilePath) -> str:
return filename_parts_from_path(file_path).file_tag
def setUp(self) -> None:
self.fs = DirectIngestGCSFileSystem(FakeGCSFileSystem())
self.prioritizer = GcsfsDirectIngestJobPrioritizer(
self.fs,
self._INGEST_BUCKET_PATH, ['tagA', 'tagB'],
file_type_filter=None)
class BaseDirectIngestController(Ingestor):
"""Parses and persists individual-level info from direct ingest partners."""
_INGEST_FILE_SPLIT_LINE_LIMIT = 2500
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)
@property
def region(self) -> Region:
return regions.get_region(self.region_code().lower(),
is_direct_ingest=True)
@classmethod
@abc.abstractmethod
def region_code(cls) -> str:
pass
@abc.abstractmethod
def get_file_tag_rank_list(self) -> List[str]:
pass
@property
def system_level(self) -> SystemLevel:
return SystemLevel.for_region(self.region)
@property
def ingest_database_key(self) -> SQLAlchemyDatabaseKey:
schema_type = self.system_level.schema_type()
if schema_type == SchemaType.STATE:
state_code = StateCode(self.region_code().upper())
return SQLAlchemyDatabaseKey.for_state_code(
state_code,
self.ingest_instance.database_version(self.system_level,
state_code=state_code),
)
return SQLAlchemyDatabaseKey.for_schema(schema_type)
# ============== #
# 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,
ingest_instance=self.ingest_instance,
ingest_bucket=self.ingest_bucket_path,
just_finished_job=just_finished_job,
)
def schedule_next_ingest_job(self, just_finished_job: bool) -> None:
"""Creates a cloud task to run a /process_job request for the file, which will
process and commit the contents to Postgres."""
check_is_region_launched_in_env(self.region)
if self.ingest_instance_status_manager.is_instance_paused():
logging.info("Ingest out of [%s] is currently paused.",
self.ingest_bucket_path.uri())
return
if self._schedule_any_pre_ingest_tasks():
logging.info("Found pre-ingest tasks to schedule - returning.")
return
if self.region_lock_manager.is_locked():
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,
self.ingest_instance,
)
if (process_job_queue_info.tasks_for_instance(
region_code=self.region_code(),
ingest_instance=self.ingest_instance)
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 not self.region_lock_manager.can_proceed():
logging.info(
"Postgres to BigQuery export is running, cannot run ingest - returning"
)
return
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_instance=self.ingest_instance,
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
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:
"""Schedules all pending ingest view export tasks for launched ingest view tags,
if they have not been scheduled. If tasks are scheduled or are still running,
returns True. Otherwise, if it's safe to proceed with next steps of ingest,
returns False."""
queue_info = self.cloud_task_manager.get_bq_import_export_queue_info(
self.region, self.ingest_instance)
did_schedule = False
tasks_to_schedule = [
GcsfsRawDataBQImportArgs(path) for path in
self.raw_file_import_manager.get_unprocessed_raw_files_to_import()
]
for task_args in tasks_to_schedule:
# If the file path has not actually been discovered by the metadata manager yet, it likely was just added
# and a subsequent call to handle_files will register it and trigger another call to this function so we can
# schedule the appropriate job.
discovered = self.file_metadata_manager.has_raw_file_been_discovered(
task_args.raw_data_file_path)
# If the file path has been processed, but still in the GCS bucket, it's likely due
# to either a manual move or an accidental duplicate uploading. In either case, we
# trust the database to have the source of truth.
processed = self.file_metadata_manager.has_raw_file_been_processed(
task_args.raw_data_file_path)
if processed:
logging.warning(
"File [%s] is already marked as processed. Skipping file processing.",
task_args.raw_data_file_path,
)
if (discovered and not processed
and not queue_info.has_task_already_scheduled(task_args)):
self.cloud_task_manager.create_direct_ingest_raw_data_import_task(
self.region, self.ingest_instance, task_args)
did_schedule = True
return queue_info.has_raw_data_import_jobs_queued() or did_schedule
def _schedule_ingest_view_export_tasks(self) -> bool:
"""Schedules all pending ingest view export tasks for launched ingest view tags,
if they have not been scheduled. If tasks are scheduled or are still running,
returns True. Otherwise, if it's safe to proceed with next steps of ingest,
returns False.
"""
queue_info = self.cloud_task_manager.get_bq_import_export_queue_info(
self.region, self.ingest_instance)
if queue_info.has_ingest_view_export_jobs_queued():
# Since we schedule all export jobs at once, after all raw files have been processed, we wait for all of the
# export jobs to be done before checking if we need to schedule more.
return True
did_schedule = False
tasks_to_schedule = (
self.ingest_view_export_manager.get_ingest_view_export_task_args())
rank_list = self.get_file_tag_rank_list()
ingest_view_name_rank = {
ingest_view_name: i
for i, ingest_view_name in enumerate(rank_list)
}
# Filter out views that aren't in ingest view tags.
filtered_tasks_to_schedule = []
for args in tasks_to_schedule:
if args.ingest_view_name not in ingest_view_name_rank:
logging.warning(
"Skipping ingest view task export for [%s] - not in controller ingest tags.",
args.ingest_view_name,
)
continue
filtered_tasks_to_schedule.append(args)
tasks_to_schedule = filtered_tasks_to_schedule
# Sort by tag order and export datetime
tasks_to_schedule.sort(key=lambda args_: (
ingest_view_name_rank[args_.ingest_view_name],
args_.upper_bound_datetime_to_export,
))
for task_args in tasks_to_schedule:
if not queue_info.has_task_already_scheduled(task_args):
self.cloud_task_manager.create_direct_ingest_ingest_view_export_task(
self.region, self.ingest_instance, task_args)
did_schedule = True
return did_schedule
def _get_next_job_args(self) -> Optional[GcsfsIngestArgs]:
"""Returns args for the next ingest job, or None if there is nothing to process."""
args = self.file_prioritizer.get_next_job_args()
if not args:
return None
discovered = self.file_metadata_manager.has_ingest_view_file_been_discovered(
args.file_path)
if not discovered:
# If the file path has not actually been discovered by the controller yet, it likely was just added and a
# subsequent call to handle_files will register it and trigger another call to this function so we can
# schedule the appropriate job.
logging.info(
"Found args [%s] for a file that has not been discovered by the metadata manager yet - not scheduling.",
args,
)
return None
return args
def _on_job_scheduled(self, ingest_args: GcsfsIngestArgs) -> None:
"""Called from the scheduler queue when an individual direct ingest job
is scheduled.
"""
# =================== #
# SINGLE JOB RUN CODE #
# =================== #
def default_job_lock_timeout_in_seconds(self) -> int:
"""This method can be overridden by subclasses that need more (or less)
time to process jobs to completion, but by default enforces a
one hour timeout on locks.
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
def run_ingest_job_and_kick_scheduler_on_completion(
self, args: GcsfsIngestArgs) -> None:
check_is_region_launched_in_env(self.region)
if self.ingest_instance_status_manager.is_instance_paused():
logging.info("Ingest out of [%s] is currently paused.",
self.ingest_bucket_path.uri())
return
if not self.region_lock_manager.can_proceed():
logging.warning(
"Postgres to BigQuery export is running, can not run ingest")
raise GCSPseudoLockAlreadyExists(
"Postgres to BigQuery export is running, can not run ingest")
with self.region_lock_manager.using_region_lock(
expiration_in_seconds=self.default_job_lock_timeout_in_seconds(
), ):
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: GcsfsIngestArgs) -> 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: GcsfsIngestArgs,
contents_handle: GcsfsFileContentsHandle) -> None:
"""
Runs the full ingest process for this controller for files with
non-empty contents.
"""
ii = self._parse(args, contents_handle)
if not ii:
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 = serialization.convert_ingest_info_to_proto(ii)
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: GcsfsIngestArgs) -> IngestMetadata:
return IngestMetadata(
region=self.region.region_code,
jurisdiction_id=self.region.jurisdiction_id,
ingest_time=args.ingest_time,
enum_overrides=self.get_enum_overrides(),
system_level=self.system_level,
database_key=self.ingest_database_key,
)
def _job_tag(self, args: GcsfsIngestArgs) -> str:
"""Returns a (short) string tag to identify an ingest run in logs."""
return (f"{self.region.region_code}/{args.file_path.file_name}:"
f"{args.ingest_time}")
def _get_contents_handle(
self, args: GcsfsIngestArgs) -> Optional[GcsfsFileContentsHandle]:
"""Returns a handle to the contents allows us to iterate 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).
"""
return self._get_contents_handle_from_path(args.file_path)
def _get_contents_handle_from_path(
self, path: GcsfsFilePath) -> Optional[GcsfsFileContentsHandle]:
return self.fs.download_to_temp_file(path)
@abc.abstractmethod
def _are_contents_empty(self, args: GcsfsIngestArgs,
contents_handle: GcsfsFileContentsHandle) -> 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: GcsfsIngestArgs,
contents_handle: GcsfsFileContentsHandle
) -> ingest_info.IngestInfo:
"""Parses ingest view file contents into an IngestInfo object."""
def _do_cleanup(self, args: GcsfsIngestArgs) -> None:
"""Does necessary cleanup once file contents have been successfully persisted to
Postgres.
"""
self.fs.mv_path_to_processed_path(args.file_path)
self.file_metadata_manager.mark_ingest_view_file_as_processed(
args.file_path)
parts = filename_parts_from_path(args.file_path)
self._move_processed_files_to_storage_as_necessary(
last_processed_date_str=parts.date_str)
def _can_proceed_with_ingest_for_contents(
self, args: GcsfsIngestArgs,
contents_handle: GcsfsFileContentsHandle) -> bool:
"""Given a pointer to the contents, returns whether the controller can continue
ingest.
"""
parts = filename_parts_from_path(args.file_path)
return self._are_contents_empty(
args, contents_handle) or not self._must_split_contents(
parts.file_type, args.file_path)
def _must_split_contents(self, file_type: GcsfsDirectIngestFileType,
path: GcsfsFilePath) -> bool:
if file_type == GcsfsDirectIngestFileType.RAW_DATA:
return False
return not self._file_meets_file_line_limit(
self.ingest_file_split_line_limit, path)
@abc.abstractmethod
def _file_meets_file_line_limit(self, line_limit: int,
path: GcsfsFilePath) -> bool:
"""Subclasses should implement to determine whether the file meets the
expected line limit"""
def _move_processed_files_to_storage_as_necessary(
self, last_processed_date_str: str) -> None:
"""Moves files that have already been ingested/processed, up to and including the given date, into storage,
if there is nothing more left to ingest/process, i.e. we are not expecting more files."""
next_args = self.file_prioritizer.get_next_job_args()
should_move_last_processed_date = False
if not next_args:
are_more_jobs_expected = (
self.file_prioritizer.are_more_jobs_expected_for_day(
last_processed_date_str))
if not are_more_jobs_expected:
should_move_last_processed_date = True
else:
next_date_str = filename_parts_from_path(
next_args.file_path).date_str
if next_date_str < last_processed_date_str:
logging.info("Found a file [%s] from a date previous to our "
"last processed date - not moving anything to "
"storage.")
return
# If there are still more to process on this day, do not move files
# from this day.
should_move_last_processed_date = next_date_str != last_processed_date_str
# Note: at this point, we expect RAW file type files to already have been moved once they were imported to BQ.
self.fs.mv_processed_paths_before_date_to_storage(
self.ingest_bucket_path,
self.storage_directory_path,
file_type_filter=GcsfsDirectIngestFileType.INGEST_VIEW,
date_str_bound=last_processed_date_str,
include_bound=should_move_last_processed_date,
)
# ================= #
# NEW FILE HANDLING #
# ================= #
def handle_file(self, path: GcsfsFilePath, start_ingest: bool) -> None:
"""Called when a single new file is added to an ingest bucket (may also
be called as a result of a rename).
May be called from any worker/queue.
"""
if self.fs.is_processed_file(path):
logging.info("File [%s] is already processed, returning.",
path.abs_path())
return
if self.fs.is_normalized_file_path(path):
parts = filename_parts_from_path(path)
if (parts.is_file_split and parts.file_split_size
and parts.file_split_size <=
self.ingest_file_split_line_limit):
self.kick_scheduler(just_finished_job=False)
logging.info(
"File [%s] is already normalized and split split "
"with correct size, kicking scheduler.",
path.abs_path(),
)
return
logging.info("Creating cloud task to schedule next job.")
self.cloud_task_manager.create_direct_ingest_handle_new_files_task(
region=self.region,
ingest_instance=self.ingest_instance,
ingest_bucket=self.ingest_bucket_path,
can_start_ingest=start_ingest,
)
def _register_all_new_paths_in_metadata(
self, paths: List[GcsfsFilePath]) -> None:
for path in paths:
parts = filename_parts_from_path(path)
if parts.file_type == GcsfsDirectIngestFileType.RAW_DATA:
if not self.file_metadata_manager.has_raw_file_been_discovered(
path):
self.file_metadata_manager.mark_raw_file_as_discovered(
path)
elif parts.file_type == GcsfsDirectIngestFileType.INGEST_VIEW:
if not self.file_metadata_manager.has_ingest_view_file_been_discovered(
path):
self.file_metadata_manager.mark_ingest_view_file_as_discovered(
path)
else:
raise ValueError(f"Unexpected file type [{parts.file_type}]")
@trace.span
def handle_new_files(self, can_start_ingest: bool) -> None:
"""Searches the ingest directory for new/unprocessed files. Normalizes
file names and splits files as necessary, schedules the next ingest job
if allowed.
Should only be called from the scheduler queue.
"""
if not can_start_ingest and self.region.is_ingest_launched_in_env():
raise ValueError(
"The can_start_ingest flag should only be used for regions where ingest is not yet launched in a "
"particular environment. If we want to be able to selectively pause ingest processing for a state, we "
"will first have to build a config that is respected by both the /ensure_all_raw_file_paths_normalized "
"endpoint and any cloud functions that trigger ingest.")
if self.ingest_instance_status_manager.is_instance_paused():
logging.info("Ingest out of [%s] is currently paused.",
self.ingest_bucket_path.uri())
return
unnormalized_paths = self.fs.get_unnormalized_file_paths(
self.ingest_bucket_path)
for path in unnormalized_paths:
logging.info("File [%s] is not yet seen, normalizing.",
path.abs_path())
self.fs.mv_path_to_normalized_path(
path, file_type=GcsfsDirectIngestFileType.RAW_DATA)
if unnormalized_paths:
logging.info(
"Normalized at least one path - returning, will handle "
"normalized files separately.")
# Normalizing file paths will cause the cloud function that calls
# this function to be re-triggered.
return
if not can_start_ingest:
logging.warning(
"Ingest not configured to start post-file normalization - returning."
)
return
check_is_region_launched_in_env(self.region)
unprocessed_ingest_view_paths = self.fs.get_unprocessed_file_paths(
self.ingest_bucket_path,
file_type_filter=GcsfsDirectIngestFileType.INGEST_VIEW,
)
unprocessed_raw_paths = self.fs.get_unprocessed_file_paths(
self.ingest_bucket_path,
file_type_filter=GcsfsDirectIngestFileType.RAW_DATA,
)
if (unprocessed_raw_paths
and self.ingest_instance == DirectIngestInstance.SECONDARY):
raise ValueError(
f"Raw data import not supported from SECONDARY ingest bucket "
f"[{self.ingest_bucket_path}], but found {len(unprocessed_raw_paths)} "
f"raw files. All raw files should be removed from this bucket and "
f"uploaded to the primary ingest bucket, if appropriate.")
self._register_all_new_paths_in_metadata(unprocessed_raw_paths)
self._register_all_new_paths_in_metadata(unprocessed_ingest_view_paths)
unprocessed_paths = unprocessed_raw_paths + unprocessed_ingest_view_paths
did_split = False
for path in unprocessed_ingest_view_paths:
if self._split_file_if_necessary(path):
did_split = True
if did_split:
post_split_unprocessed_ingest_view_paths = (
self.fs.get_unprocessed_file_paths(
self.ingest_bucket_path,
file_type_filter=GcsfsDirectIngestFileType.INGEST_VIEW,
))
self._register_all_new_paths_in_metadata(
post_split_unprocessed_ingest_view_paths)
logging.info(
"Split at least one path - returning, will handle split "
"files separately.")
# Writing new split files to storage will cause the cloud function
# that calls this function to be re-triggered.
return
if unprocessed_paths:
self.schedule_next_ingest_job(just_finished_job=False)
def do_raw_data_import(self,
data_import_args: GcsfsRawDataBQImportArgs) -> None:
"""Process a raw incoming file by importing it to BQ, tracking it in our metadata tables, and moving it to
storage on completion.
"""
check_is_region_launched_in_env(self.region)
if self.ingest_instance_status_manager.is_instance_paused():
logging.info("Ingest out of [%s] is currently paused.",
self.ingest_bucket_path.uri())
return
if self.ingest_instance == DirectIngestInstance.SECONDARY:
raise ValueError(
f"Raw data import not supported from SECONDARY ingest bucket "
f"[{self.ingest_bucket_path}]. Raw data task for "
f"[{data_import_args.raw_data_file_path}] should never have been "
f"scheduled.")
if not self.fs.exists(data_import_args.raw_data_file_path):
logging.warning(
"File path [%s] no longer exists - might have already been "
"processed or deleted",
data_import_args.raw_data_file_path,
)
self.kick_scheduler(just_finished_job=True)
return
file_metadata = self.file_metadata_manager.get_raw_file_metadata(
data_import_args.raw_data_file_path)
if file_metadata.processed_time:
logging.warning(
"File [%s] is already marked as processed. Skipping file processing.",
data_import_args.raw_data_file_path.file_name,
)
self.kick_scheduler(just_finished_job=True)
return
self.raw_file_import_manager.import_raw_file_to_big_query(
data_import_args.raw_data_file_path, file_metadata)
processed_path = self.fs.mv_path_to_processed_path(
data_import_args.raw_data_file_path)
self.file_metadata_manager.mark_raw_file_as_processed(
path=data_import_args.raw_data_file_path)
self.fs.mv_path_to_storage(processed_path, self.storage_directory_path)
self.kick_scheduler(just_finished_job=True)
def do_ingest_view_export(
self, ingest_view_export_args: GcsfsIngestViewExportArgs) -> None:
check_is_region_launched_in_env(self.region)
if self.ingest_instance_status_manager.is_instance_paused():
logging.info("Ingest out of [%s] is currently paused.",
self.ingest_bucket_path.uri())
return
did_export = self.ingest_view_export_manager.export_view_for_args(
ingest_view_export_args)
if (not did_export or not self.file_metadata_manager.
get_ingest_view_metadata_pending_export()):
logging.info("Creating cloud task to schedule next job.")
self.cloud_task_manager.create_direct_ingest_handle_new_files_task(
region=self.region,
ingest_instance=self.ingest_instance,
ingest_bucket=self.ingest_bucket_path,
can_start_ingest=True,
)
def _should_split_file(self, path: GcsfsFilePath) -> bool:
"""Returns a handle to the contents of this path if this file should be split, None otherwise."""
parts = filename_parts_from_path(path)
if parts.file_type != GcsfsDirectIngestFileType.INGEST_VIEW:
raise ValueError(
f"Should not be attempting to split files other than ingest view files, found path with "
f"file type: {parts.file_type}")
if parts.file_tag not in self.get_file_tag_rank_list():
logging.info(
"File tag [%s] for path [%s] not in rank list - not splitting.",
parts.file_tag,
path.abs_path(),
)
return False
if (parts.is_file_split and parts.file_split_size and
parts.file_split_size <= self.ingest_file_split_line_limit):
logging.info(
"File [%s] already split with size [%s].",
path.abs_path(),
parts.file_split_size,
)
return False
return self._must_split_contents(parts.file_type, path)
@trace.span
def _split_file_if_necessary(self, path: GcsfsFilePath) -> bool:
"""Checks if the given file needs to be split according to this controller's |file_split_line_limit|.
Returns True if the file was split, False if splitting was not necessary.
"""
should_split = self._should_split_file(path)
if not should_split:
logging.info("No need to split file path [%s].", path.abs_path())
return False
logging.info("Proceeding to file splitting for path [%s].",
path.abs_path())
original_metadata = self.file_metadata_manager.get_ingest_view_file_metadata(
path)
output_dir = GcsfsDirectoryPath.from_file_path(path)
split_contents_paths = self._split_file(path)
upload_paths = []
for i, split_contents_path in enumerate(split_contents_paths):
upload_path = self._create_split_file_path(path,
output_dir,
split_num=i)
logging.info(
"Copying split [%s] to direct ingest directory at path [%s].",
i,
upload_path.abs_path(),
)
upload_paths.append(upload_path)
try:
self.fs.mv(split_contents_path, upload_path)
except Exception as e:
logging.error(
"Threw error while copying split files from temp bucket - attempting to clean up before rethrowing."
" [%s]",
e,
)
for p in upload_paths:
self.fs.delete(p)
raise e
# We wait to register files with metadata manager until all files have been successfully copied to avoid leaving
# the metadata manager in an inconsistent state.
if not isinstance(original_metadata, DirectIngestIngestFileMetadata):
raise ValueError("Attempting to split a non-ingest view type file")
logging.info(
"Registering [%s] split files with the metadata manager.",
len(upload_paths),
)
for upload_path in upload_paths:
ingest_file_metadata = (
self.file_metadata_manager.register_ingest_file_split(
original_metadata, upload_path))
self.file_metadata_manager.mark_ingest_view_exported(
ingest_file_metadata)
self.file_metadata_manager.mark_ingest_view_file_as_processed(path)
logging.info(
"Done splitting file [%s] into [%s] paths, moving it to storage.",
path.abs_path(),
len(split_contents_paths),
)
self.fs.mv_path_to_storage(path, self.storage_directory_path)
return True
def _create_split_file_path(
self,
original_file_path: GcsfsFilePath,
output_dir: GcsfsDirectoryPath,
split_num: int,
) -> GcsfsFilePath:
parts = filename_parts_from_path(original_file_path)
rank_str = str(split_num + 1).zfill(5)
updated_file_name = (
f"{parts.stripped_file_name}_{rank_str}"
f"_{SPLIT_FILE_SUFFIX}_size{self.ingest_file_split_line_limit}"
f".{parts.extension}")
return GcsfsFilePath.from_directory_and_file_name(
output_dir,
to_normalized_unprocessed_file_path(
updated_file_name,
file_type=parts.file_type,
dt=parts.utc_upload_datetime,
),
)
@abc.abstractmethod
def _split_file(self, path: GcsfsFilePath) -> List[GcsfsFilePath]:
"""Should be implemented by subclasses to split a file accessible via the provided path into multiple
files and upload those files to GCS. Returns the list of upload paths."""
@staticmethod
def file_tag(file_path: GcsfsFilePath) -> str:
return filename_parts_from_path(file_path).file_tag
