def load_cloudasset_data(session, config):
"""Export asset data from Cloud Asset API and load into storage.
Args:
session (object): Database session.
config (object): Inventory configuration on server.
Returns:
int: The count of assets imported into the database, or None if there
is an error.
"""
# Start by ensuring that there is no existing CAI data in storage.
_clear_cai_data(session)
cloudasset_client = cloudasset.CloudAssetClient(
config.get_api_quota_configs())
imported_assets = 0
root_resources = []
if config.use_composite_root():
root_resources.extend(config.get_composite_root_resources())
else:
root_resources.append(config.get_root_resource_id())
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
futures = []
for root_id in root_resources:
for content_type in CONTENT_TYPES:
futures.append(executor.submit(_export_assets,
cloudasset_client,
config,
root_id,
content_type))
for future in concurrent.futures.as_completed(futures):
temporary_file = ''
try:
temporary_file = future.result()
if not temporary_file:
return _clear_cai_data(session)
LOGGER.debug('Importing Cloud Asset data from %s to database.',
temporary_file)
with open(temporary_file, 'r') as cai_data:
rows = CaiDataAccess.populate_cai_data(cai_data, session)
imported_assets += rows
LOGGER.info('%s assets imported to database.', rows)
finally:
if temporary_file:
os.unlink(temporary_file)
return imported_assets
def _add_iam_policies(self):
"""Add CAI IAM Policies to temporary table."""
iam_policy_data = StringIO(CAI_IAM_POLICY_DATA)
rows = CaiDataAccess.populate_cai_data(iam_policy_data, self.session)
expected_rows = len(CAI_IAM_POLICY_DATA.split('\n'))
self.assertEqual(expected_rows, rows)
def _add_resources(self):
"""Add CAI resources to temporary table."""
resource_data = StringIO(CAI_RESOURCE_DATA)
rows = CaiDataAccess.populate_cai_data(resource_data, self.session)
expected_rows = len(CAI_RESOURCE_DATA.split('\n'))
self.assertEqual(expected_rows, rows)