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 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(
),
}