def _download_cloudasset_data(config, inventory_index_id):
"""Download cloud asset data.
Args:
config (InventoryConfig): Inventory config.
inventory_index_id (int): The inventory index ID for this export.
Yields:
str: GCS path of the cloud asset file.
"""
root_resources = []
if config.use_composite_root():
root_resources.extend(config.get_composite_root_resources())
else:
root_resources.append(config.get_root_resource_id())
cloudasset_client = cloudasset.CloudAssetClient(
config.get_api_quota_configs())
with concurrent.futures.ThreadPoolExecutor(max_workers=4) 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,
inventory_index_id))
for future in concurrent.futures.as_completed(futures):
yield future.result()
def setUpClass(cls, mock_google_credential):
"""Set up."""
fake_global_configs = {'cloudasset': {'max_calls': 1, 'period': 1.0}}
cls.asset_api_client = cloudasset.CloudAssetClient(
global_configs=fake_global_configs, use_rate_limiter=False)
cls.default_output_config = (
cls.asset_api_client.build_gcs_object_output(
fake_cloudasset.DESTINATION))
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 test_no_quota(self, mock_google_credential):
"""Verify no rate limiter is used if the configuration is missing."""
asset_api_client = cloudasset.CloudAssetClient(global_configs={})
self.assertEqual(None, asset_api_client.repository._rate_limiter)
def setUpClass(cls, mock_google_credential):
"""Set up."""
fake_global_configs = {'cloudasset': {'max_calls': 1, 'period': 1.0}}
cls.asset_api_client = cloudasset.CloudAssetClient(
global_configs=fake_global_configs, use_rate_limiter=False)
