from flask import Blueprint, Response, abort, current_app, request from prometheus_client import CONTENT_TYPE_LATEST, CollectorRegistry, generate_latest from .enterprise_billing_collector import AzureEABillingCollector from .allocated_vm_collector import AzureAllocatedVMCollector from .reserved_vm_collector import AzureReservedVMCollector bp = Blueprint('views', __name__) DEFAULT_SCRAPE_TIMEOUT = 10 def _get_timeout(): try: return float(request.headers.get('X-Prometheus-Scrape-Timeout-Seconds')) except Exception: return current_app.config.get('BILLING_SCRAPE_TIMEOUT', DEFAULT_SCRAPE_TIMEOUT) def _register_billing_collector(registry): timeout = _get_timeout() collector = AzureEABillingCollector( current_app.config['BILLING_METRIC_NAME'], current_app.config['ENROLLMENT_NUMBER'], current_app.config['BILLING_API_ACCESS_KEY'], timeout ) registry.register(collector) def _register_allocated_vm_collector(registry): collector = AzureAllocatedVMCollector( current_app.config['APPLICATION_ID'], current_app.config['APPLICATION_SECRET'], current_app.config['AD_TENANT_ID'], current_app.config['SUBSCRIPTION_IDS'], current_app.config['ALLOCATED_VM_METRIC_NAME'], ) registry.register(collector) def _register_reserved_vm_collector(registry): collector = AzureReservedVMCollector( current_app.config['APPLICATION_ID'], current_app.config['APPLICATION_SECRET'], current_app.config['AD_TENANT_ID'], current_app.config['RESERVED_VM_METRIC_NAME'], ) registry.register(collector) def _register_collectors(registry): if 'BILLING_METRIC_NAME' in current_app.config: _register_billing_collector(registry) if 'ALLOCATED_VM_METRIC_NAME' in current_app.config: _register_allocated_vm_collector(registry) if 'RESERVED_VM_METRIC_NAME' in current_app.config: _register_reserved_vm_collector(registry) @bp.route("/health") def health(): return 'ok' @bp.route("/metrics") def metrics(): registry = CollectorRegistry() _register_collectors(registry) try: content = generate_latest(registry) return content, 200, {'Content-Type': CONTENT_TYPE_LATEST} except Exception as e: abort(Response("Scrape failed: {}".format(e), status=502))