API refactor

venv/lib/python3.12/site-packages/prometheus_client/exposition.py

@@ -1,5 +1,6 @@
 import base64
 from contextlib import closing
+from functools import partial
 import gzip
 from http.server import BaseHTTPRequestHandler
 import os
@@ -19,10 +20,12 @@ from wsgiref.simple_server import make_server, WSGIRequestHandler, WSGIServer
 
 from .openmetrics import exposition as openmetrics
 from .registry import CollectorRegistry, REGISTRY
-from .utils import floatToGoString
+from .utils import floatToGoString, parse_version
 
 __all__ = (
     'CONTENT_TYPE_LATEST',
+    'CONTENT_TYPE_PLAIN_0_0_4',
+    'CONTENT_TYPE_PLAIN_1_0_0',
     'delete_from_gateway',
     'generate_latest',
     'instance_ip_grouping_key',
@@ -36,8 +39,13 @@ __all__ = (
     'write_to_textfile',
 )
 
-CONTENT_TYPE_LATEST = 'text/plain; version=0.0.4; charset=utf-8'
-"""Content type of the latest text format"""
+CONTENT_TYPE_PLAIN_0_0_4 = 'text/plain; version=0.0.4; charset=utf-8'
+"""Content type of the compatibility format"""
+
+CONTENT_TYPE_PLAIN_1_0_0 = 'text/plain; version=1.0.0; charset=utf-8'
+"""Content type of the latest format"""
+
+CONTENT_TYPE_LATEST = CONTENT_TYPE_PLAIN_1_0_0
 
 
 class _PrometheusRedirectHandler(HTTPRedirectHandler):
@@ -118,12 +126,24 @@ def make_wsgi_app(registry: CollectorRegistry = REGISTRY, disable_compression: b
         accept_header = environ.get('HTTP_ACCEPT')
         accept_encoding_header = environ.get('HTTP_ACCEPT_ENCODING')
         params = parse_qs(environ.get('QUERY_STRING', ''))
-        if environ['PATH_INFO'] == '/favicon.ico':
+        method = environ['REQUEST_METHOD']
+
+        if method == 'OPTIONS':
+            status = '200 OK'
+            headers = [('Allow', 'OPTIONS,GET')]
+            output = b''
+        elif method != 'GET':
+            status = '405 Method Not Allowed'
+            headers = [('Allow', 'OPTIONS,GET')]
+            output = '# HTTP {}: {}; use OPTIONS or GET\n'.format(status, method).encode()
+        elif environ['PATH_INFO'] == '/favicon.ico':
             # Serve empty response for browsers
             status = '200 OK'
-            headers = [('', '')]
+            headers = []
             output = b''
         else:
+            # Note: For backwards compatibility, the URI path for GET is not
+            # constrained to the documented /metrics, but any path is allowed.
             # Bake output
             status, headers, output = _bake_output(registry, accept_header, accept_encoding_header, params, disable_compression)
         # Return output
@@ -154,12 +174,63 @@ def _get_best_family(address, port):
     # binding an ipv6 address is requested.
     # This function is based on what upstream python did for http.server
     # in https://github.com/python/cpython/pull/11767
-    infos = socket.getaddrinfo(address, port)
+    infos = socket.getaddrinfo(address, port, type=socket.SOCK_STREAM, flags=socket.AI_PASSIVE)
     family, _, _, _, sockaddr = next(iter(infos))
     return family, sockaddr[0]
 
 
-def start_wsgi_server(port: int, addr: str = '0.0.0.0', registry: CollectorRegistry = REGISTRY) -> None:
+def _get_ssl_ctx(
+        certfile: str,
+        keyfile: str,
+        protocol: int,
+        cafile: Optional[str] = None,
+        capath: Optional[str] = None,
+        client_auth_required: bool = False,
+) -> ssl.SSLContext:
+    """Load context supports SSL."""
+    ssl_cxt = ssl.SSLContext(protocol=protocol)
+
+    if cafile is not None or capath is not None:
+        try:
+            ssl_cxt.load_verify_locations(cafile, capath)
+        except IOError as exc:
+            exc_type = type(exc)
+            msg = str(exc)
+            raise exc_type(f"Cannot load CA certificate chain from file "
+                           f"{cafile!r} or directory {capath!r}: {msg}")
+    else:
+        try:
+            ssl_cxt.load_default_certs(purpose=ssl.Purpose.CLIENT_AUTH)
+        except IOError as exc:
+            exc_type = type(exc)
+            msg = str(exc)
+            raise exc_type(f"Cannot load default CA certificate chain: {msg}")
+
+    if client_auth_required:
+        ssl_cxt.verify_mode = ssl.CERT_REQUIRED
+
+    try:
+        ssl_cxt.load_cert_chain(certfile=certfile, keyfile=keyfile)
+    except IOError as exc:
+        exc_type = type(exc)
+        msg = str(exc)
+        raise exc_type(f"Cannot load server certificate file {certfile!r} or "
+                       f"its private key file {keyfile!r}: {msg}")
+
+    return ssl_cxt
+
+
+def start_wsgi_server(
+        port: int,
+        addr: str = '0.0.0.0',
+        registry: CollectorRegistry = REGISTRY,
+        certfile: Optional[str] = None,
+        keyfile: Optional[str] = None,
+        client_cafile: Optional[str] = None,
+        client_capath: Optional[str] = None,
+        protocol: int = ssl.PROTOCOL_TLS_SERVER,
+        client_auth_required: bool = False,
+) -> Tuple[WSGIServer, threading.Thread]:
     """Starts a WSGI server for prometheus metrics as a daemon thread."""
 
     class TmpServer(ThreadingWSGIServer):
@@ -168,30 +239,51 @@ def start_wsgi_server(port: int, addr: str = '0.0.0.0', registry: CollectorRegis
     TmpServer.address_family, addr = _get_best_family(addr, port)
     app = make_wsgi_app(registry)
     httpd = make_server(addr, port, app, TmpServer, handler_class=_SilentHandler)
+    if certfile and keyfile:
+        context = _get_ssl_ctx(certfile, keyfile, protocol, client_cafile, client_capath, client_auth_required)
+        httpd.socket = context.wrap_socket(httpd.socket, server_side=True)
     t = threading.Thread(target=httpd.serve_forever)
     t.daemon = True
     t.start()
 
+    return httpd, t
+
 
 start_http_server = start_wsgi_server
 
 
-def generate_latest(registry: CollectorRegistry = REGISTRY) -> bytes:
-    """Returns the metrics from the registry in latest text format as a string."""
+def generate_latest(registry: CollectorRegistry = REGISTRY, escaping: str = openmetrics.UNDERSCORES) -> bytes:
+    """
+    Generates the exposition format using the basic Prometheus text format.
 
-    def sample_line(line):
-        if line.labels:
-            labelstr = '{{{0}}}'.format(','.join(
+    Params:
+        registry: CollectorRegistry to export data from.
+        escaping: Escaping scheme used for metric and label names.
+
+    Returns: UTF-8 encoded string containing the metrics in text format.
+    """
+
+    def sample_line(samples):
+        if samples.labels:
+            labelstr = '{0}'.format(','.join(
+                # Label values always support UTF-8
                 ['{}="{}"'.format(
-                    k, v.replace('\\', r'\\').replace('\n', r'\n').replace('"', r'\"'))
-                    for k, v in sorted(line.labels.items())]))
+                    openmetrics.escape_label_name(k, escaping), openmetrics._escape(v, openmetrics.ALLOWUTF8, False))
+                    for k, v in sorted(samples.labels.items())]))
         else:
             labelstr = ''
         timestamp = ''
-        if line.timestamp is not None:
+        if samples.timestamp is not None:
             # Convert to milliseconds.
-            timestamp = f' {int(float(line.timestamp) * 1000):d}'
-        return f'{line.name}{labelstr} {floatToGoString(line.value)}{timestamp}\n'
+            timestamp = f' {int(float(samples.timestamp) * 1000):d}'
+        if escaping != openmetrics.ALLOWUTF8 or openmetrics._is_valid_legacy_metric_name(samples.name):
+            if labelstr:
+                labelstr = '{{{0}}}'.format(labelstr)
+            return f'{openmetrics.escape_metric_name(samples.name, escaping)}{labelstr} {floatToGoString(samples.value)}{timestamp}\n'
+        maybe_comma = ''
+        if labelstr:
+            maybe_comma = ','
+        return f'{{{openmetrics.escape_metric_name(samples.name, escaping)}{maybe_comma}{labelstr}}} {floatToGoString(samples.value)}{timestamp}\n'
 
     output = []
     for metric in registry.collect():
@@ -214,8 +306,8 @@ def generate_latest(registry: CollectorRegistry = REGISTRY) -> bytes:
                 mtype = 'untyped'
 
             output.append('# HELP {} {}\n'.format(
-                mname, metric.documentation.replace('\\', r'\\').replace('\n', r'\n')))
-            output.append(f'# TYPE {mname} {mtype}\n')
+                openmetrics.escape_metric_name(mname, escaping), metric.documentation.replace('\\', r'\\').replace('\n', r'\n')))
+            output.append(f'# TYPE {openmetrics.escape_metric_name(mname, escaping)} {mtype}\n')
 
             om_samples: Dict[str, List[str]] = {}
             for s in metric.samples:
@@ -231,20 +323,79 @@ def generate_latest(registry: CollectorRegistry = REGISTRY) -> bytes:
             raise
 
         for suffix, lines in sorted(om_samples.items()):
-            output.append('# HELP {}{} {}\n'.format(metric.name, suffix,
-                                                    metric.documentation.replace('\\', r'\\').replace('\n', r'\n')))
-            output.append(f'# TYPE {metric.name}{suffix} gauge\n')
+            output.append('# HELP {} {}\n'.format(openmetrics.escape_metric_name(metric.name + suffix, escaping),
+                                                  metric.documentation.replace('\\', r'\\').replace('\n', r'\n')))
+            output.append(f'# TYPE {openmetrics.escape_metric_name(metric.name + suffix, escaping)} gauge\n')
             output.extend(lines)
     return ''.join(output).encode('utf-8')
 
 
 def choose_encoder(accept_header: str) -> Tuple[Callable[[CollectorRegistry], bytes], str]:
+    # Python client library accepts a narrower range of content-types than
+    # Prometheus does.
     accept_header = accept_header or ''
+    escaping = openmetrics.UNDERSCORES
     for accepted in accept_header.split(','):
         if accepted.split(';')[0].strip() == 'application/openmetrics-text':
-            return (openmetrics.generate_latest,
-                    openmetrics.CONTENT_TYPE_LATEST)
-    return generate_latest, CONTENT_TYPE_LATEST
+            toks = accepted.split(';')
+            version = _get_version(toks)
+            escaping = _get_escaping(toks)
+            # Only return an escaping header if we have a good version and
+            # mimetype.
+            if not version:
+                return (partial(openmetrics.generate_latest, escaping=openmetrics.UNDERSCORES, version="1.0.0"), openmetrics.CONTENT_TYPE_LATEST)
+            if version and parse_version(version) >= (1, 0, 0):
+                return (partial(openmetrics.generate_latest, escaping=escaping, version=version),
+                        f'application/openmetrics-text; version={version}; charset=utf-8; escaping=' + str(escaping))
+        elif accepted.split(';')[0].strip() == 'text/plain':
+            toks = accepted.split(';')
+            version = _get_version(toks)
+            escaping = _get_escaping(toks)
+            # Only return an escaping header if we have a good version and
+            # mimetype.
+            if version and parse_version(version) >= (1, 0, 0):
+                return (partial(generate_latest, escaping=escaping),
+                        CONTENT_TYPE_LATEST + '; escaping=' + str(escaping))
+    return generate_latest, CONTENT_TYPE_PLAIN_0_0_4
+
+
+def _get_version(accept_header: List[str]) -> str:
+    """Return the version tag from the Accept header.
+
+    If no version is specified, returns empty string."""
+
+    for tok in accept_header:
+        if '=' not in tok:
+            continue
+        key, value = tok.strip().split('=', 1)
+        if key == 'version':
+            return value
+    return ""
+
+
+def _get_escaping(accept_header: List[str]) -> str:
+    """Return the escaping scheme from the Accept header.
+
+    If no escaping scheme is specified or the scheme is not one of the allowed
+    strings, defaults to UNDERSCORES."""
+
+    for tok in accept_header:
+        if '=' not in tok:
+            continue
+        key, value = tok.strip().split('=', 1)
+        if key != 'escaping':
+            continue
+        if value == openmetrics.ALLOWUTF8:
+            return openmetrics.ALLOWUTF8
+        elif value == openmetrics.UNDERSCORES:
+            return openmetrics.UNDERSCORES
+        elif value == openmetrics.DOTS:
+            return openmetrics.DOTS
+        elif value == openmetrics.VALUES:
+            return openmetrics.VALUES
+        else:
+            return openmetrics.UNDERSCORES
+    return openmetrics.UNDERSCORES
 
 
 def gzip_accepted(accept_encoding_header: str) -> bool:
@@ -293,20 +444,34 @@ class MetricsHandler(BaseHTTPRequestHandler):
         return MyMetricsHandler
 
 
-def write_to_textfile(path: str, registry: CollectorRegistry) -> None:
+def write_to_textfile(path: str, registry: CollectorRegistry, escaping: str = openmetrics.ALLOWUTF8, tmpdir: Optional[str] = None) -> None:
     """Write metrics to the given path.
 
     This is intended for use with the Node exporter textfile collector.
-    The path must end in .prom for the textfile collector to process it."""
-    tmppath = f'{path}.{os.getpid()}.{threading.current_thread().ident}'
-    with open(tmppath, 'wb') as f:
-        f.write(generate_latest(registry))
+    The path must end in .prom for the textfile collector to process it.
 
-    # rename(2) is atomic but fails on Windows if the destination file exists
-    if os.name == 'nt':
-        os.replace(tmppath, path)
+    An optional tmpdir parameter can be set to determine where the
+    metrics will be temporarily written to. If not set, it will be in
+    the same directory as the .prom file. If provided, the path MUST be
+    on the same filesystem."""
+    if tmpdir is not None:
+        filename = os.path.basename(path)
+        tmppath = f'{os.path.join(tmpdir, filename)}.{os.getpid()}.{threading.current_thread().ident}'
     else:
-        os.rename(tmppath, path)
+        tmppath = f'{path}.{os.getpid()}.{threading.current_thread().ident}'
+    try:
+        with open(tmppath, 'wb') as f:
+            f.write(generate_latest(registry, escaping))
+
+        # rename(2) is atomic but fails on Windows if the destination file exists
+        if os.name == 'nt':
+            os.replace(tmppath, path)
+        else:
+            os.rename(tmppath, path)
+    except Exception:
+        if os.path.exists(tmppath):
+            os.remove(tmppath)
+        raise
 
 
 def _make_handler(
@@ -407,7 +572,7 @@ def tls_auth_handler(
     The default protocol (ssl.PROTOCOL_TLS_CLIENT) will also enable
     ssl.CERT_REQUIRED and SSLContext.check_hostname by default. This can be
     disabled by setting insecure_skip_verify to True.
-    
+
     Both this handler and the TLS feature on pushgateay are experimental."""
     context = ssl.SSLContext(protocol=protocol)
     if cafile is not None:
@@ -564,7 +729,7 @@ def _use_gateway(
 
     handler(
         url=url, method=method, timeout=timeout,
-        headers=[('Content-Type', CONTENT_TYPE_LATEST)], data=data,
+        headers=[('Content-Type', CONTENT_TYPE_PLAIN_0_0_4)], data=data,
     )()