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)