90 lines
3.2 KiB
Python
90 lines
3.2 KiB
Python
# Copyright 2022 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
'''Discovers project quota via Compute API and overlay user overrides.
|
|
|
|
This plugin discovers project quota via batch Compute API requests. Project and
|
|
network quotas are then optionally overlaid with custom quota modifiers passed
|
|
in as options. Region quota discovery is partially implemented but not active.
|
|
'''
|
|
|
|
import logging
|
|
|
|
from . import Level, Resource, register_init, register_discovery
|
|
from .utils import batched, poor_man_mp_request, poor_man_mp_response
|
|
|
|
LOGGER = logging.getLogger('net-dash.discovery.compute-quota')
|
|
NAME = 'quota'
|
|
|
|
API_GLOBAL_URL = '/compute/v1/projects/{}'
|
|
API_REGION_URL = '/compute/v1/projects/{}/regions/{}'
|
|
|
|
|
|
def _handle_discovery(resources, response):
|
|
'Processes asset batch response and overlays custom quota.'
|
|
LOGGER.info('discovery handle request')
|
|
content_type = response.headers['content-type']
|
|
per_project_quota = resources['config:custom_quota'].get('projects', {})
|
|
# process batch response
|
|
for part in poor_man_mp_response(content_type, response.content):
|
|
kind = part.get('kind')
|
|
quota = {
|
|
q['metric']: int(q['limit'])
|
|
for q in sorted(part.get('quotas', []), key=lambda v: v['metric'])
|
|
}
|
|
self_link = part.get('selfLink')
|
|
if not self_link:
|
|
logging.warn('invalid quota response')
|
|
continue
|
|
self_link = self_link.split('/')
|
|
if kind == 'compute#project':
|
|
project_id = self_link[-1]
|
|
region = 'global'
|
|
elif kind == 'compute#region':
|
|
project_id = self_link[-3]
|
|
region = self_link[-1]
|
|
# custom quota overrides
|
|
for k, v in per_project_quota.get(project_id, {}).get(region, {}).items():
|
|
quota[k] = int(v)
|
|
if project_id not in resources[NAME]:
|
|
resources[NAME][project_id] = {}
|
|
yield Resource(NAME, project_id, quota, region)
|
|
|
|
|
|
@register_init
|
|
def init(resources):
|
|
'Prepares quota datastructures in the shared resource map.'
|
|
LOGGER.info('init')
|
|
resources.setdefault(NAME, {})
|
|
|
|
|
|
@register_discovery(Level.DERIVED, 0)
|
|
def start_discovery(resources, response=None):
|
|
'Plugin entry point, triggers discovery and handles requests and responses.'
|
|
LOGGER.info(f'discovery (has response: {response is not None})')
|
|
if response is None:
|
|
# TODO: regions
|
|
urls = [API_GLOBAL_URL.format(p) for p in resources['projects']]
|
|
if not urls:
|
|
return
|
|
for batch in batched(urls, 10):
|
|
yield poor_man_mp_request(batch)
|
|
else:
|
|
for result in _handle_discovery(resources, response):
|
|
yield result
|
|
# store custom network-level quota
|
|
per_network_quota = resources['config:custom_quota'].get('networks', {})
|
|
for network_id, overrides in per_network_quota.items():
|
|
quota = {k: int(v) for k, v in overrides.items()}
|
|
yield Resource(NAME, network_id, quota)
|