text: # SPDX-FileCopyrightText: 2020 The Magma Authors.
  # SPDX-FileCopyrightText: 2022 Open Networking Foundation <support@opennetworking.org>
  #
  # SPDX-License-Identifier: BSD-3-Clause
  
  import logging
  import time
  
  import metrics_pb2
  from orc8r.protos import metricsd_pb2
  from prometheus_client import REGISTRY
  
  
  def get_metrics(registry=REGISTRY, verbose=False):
      """
      Collects timeseries samples from prometheus metric collector registry
      adds a common timestamp, and encodes them to protobuf
  
      Arguments:
          regsitry: a prometheus CollectorRegistry instance
          verbose: whether to optimize for bandwidth and ignore metric name/help
  
      Returns:
          a prometheus MetricFamily protobuf stream
      """
      timestamp_ms = int(time.time() * 1000)
      for metric_family in registry.collect():
          if metric_family.type in ('counter', 'gauge'):
              family_proto = encode_counter_gauge(metric_family, timestamp_ms)
          elif metric_family.type == 'summary':
              family_proto = encode_summary(metric_family, timestamp_ms)
          elif metric_family.type == 'histogram':
              family_proto = encode_histogram(metric_family, timestamp_ms)
  
          if verbose:
              family_proto.help = metric_family.documentation
              family_proto.name = metric_family.name
          else:
              try:
                  family_proto.name = \
                      str(metricsd_pb2.MetricName.Value(metric_family.name))
              except ValueError as e:
                  logging.debug(e)  # If enum is not defined
                  family_proto.name = metric_family.name
          yield family_proto
  
  
  def encode_counter_gauge(family, timestamp_ms):
      """
      Takes a Counter/Gauge family which is a collection of timeseries
      samples that share a name (uniquely identified by labels) and yields
      equivalent protobufs.
  
      Each timeseries corresponds to a single sample tuple of the format:
          (NAME, LABELS, VALUE)
  
      Arguments:
          family: a prometheus gauge metric family
          timestamp_ms: the timestamp to attach to the samples
      Raises:
          ValueError if metric name is not defined in MetricNames protobuf
      Returns:
          A Counter or Gauge prometheus MetricFamily protobuf
      """
      family_proto = metrics_pb2.MetricFamily()
      family_proto.type = \
          metrics_pb2.MetricType.Value(family.type.upper())
      for sample in family.samples:
          metric_proto = metrics_pb2.Metric()
          if family_proto.type == metrics_pb2.COUNTER:
              metric_proto.counter.value = sample[2]
          elif family_proto.type == metrics_pb2.GAUGE:
              metric_proto.gauge.value = sample[2]
          # Add meta-data to the timeseries
          metric_proto.timestamp_ms = timestamp_ms
          metric_proto.label.extend(_convert_labels_to_enums(sample[1].items()))
          # Append metric sample to family
          family_proto.metric.extend([metric_proto])
      return family_proto
  
  
  def encode_summary(family, timestamp_ms):
      """
      Takes a Summary Metric family which is a collection of timeseries
      samples that share a name (uniquely identified by labels) and yields
      equivalent protobufs.
  
      Each summary timeseries consists of sample tuples for the count, sum,
      and quantiles in the format (NAME,LABELS,VALUE). The NAME is suffixed
      with either _count, _sum to indicate count and sum respectively.
      Quantile samples will be of the same NAME with quantile label.
  
      Arguments:
          family: a prometheus summary metric family
          timestamp_ms: the timestamp to attach to the samples
      Raises:
          ValueError if metric name is not defined in MetricNames protobuf
      Returns:
          a Summary prometheus MetricFamily protobuf
      """
      family_proto = metrics_pb2.MetricFamily()
      family_proto.type = metrics_pb2.SUMMARY
      metric_protos = {}
      # Build a map of each of the summary timeseries from the samples
      for sample in family.samples:
          quantile = sample[1].pop('quantile', None)  # Remove from label set
          # Each time series identified by label set excluding the quantile
          metric_proto = \
              metric_protos.setdefault(
                  frozenset(sample[1].items()),
                  metrics_pb2.Metric(),
              )
          if sample[0].endswith('_count'):
              metric_proto.summary.sample_count = int(sample[2])
          elif sample[0].endswith('_sum'):
              metric_proto.summary.sample_sum = sample[2]
          elif quantile:
              quantile = metric_proto.summary.quantile.add()
              quantile.value = sample[2]
              quantile.quantile = _goStringToFloat(quantile)
      # Go back and add meta-data to the timeseries
      for labels, metric_proto in metric_protos.items():
          metric_proto.timestamp_ms = timestamp_ms
          metric_proto.label.extend(_convert_labels_to_enums(labels))
          # Add it to the family
          family_proto.metric.extend([metric_proto])
      return family_proto
  
  
  def encode_histogram(family, timestamp_ms):
      """
      Takes a Histogram Metric family which is a collection of timeseries
      samples that share a name (uniquely identified by labels) and yields
      equivalent protobufs.
  
      Each summary timeseries consists of sample tuples for the count, sum,
      and quantiles in the format (NAME,LABELS,VALUE). The NAME is suffixed
      with either _count, _sum, _buckets to indicate count, sum and buckets
      respectively. Bucket samples will also contain a le to indicate its
      upper bound.
  
      Arguments:
          family: a prometheus histogram metric family
          timestamp_ms: the timestamp to attach to the samples
      Raises:
          ValueError if metric name is not defined in MetricNames protobuf
      Returns:
          a Histogram prometheus MetricFamily protobuf
      """
      family_proto = metrics_pb2.MetricFamily()
      family_proto.type = metrics_pb2.HISTOGRAM
      metric_protos = {}
      for sample in family.samples:
          upper_bound = sample[1].pop('le', None)  # Remove from label set
          metric_proto = \
              metric_protos.setdefault(
                  frozenset(sample[1].items()),
                  metrics_pb2.Metric(),
              )
          if sample[0].endswith('_count'):
              metric_proto.histogram.sample_count = int(sample[2])
          elif sample[0].endswith('_sum'):
              metric_proto.histogram.sample_sum = sample[2]
          elif sample[0].endswith('_bucket'):
              quantile = metric_proto.histogram.bucket.add()
              quantile.cumulative_count = int(sample[2])
              quantile.upper_bound = _goStringToFloat(upper_bound)
      # Go back and add meta-data to the timeseries
      for labels, metric_proto in metric_protos.items():
          metric_proto.timestamp_ms = timestamp_ms
          metric_proto.label.extend(_convert_labels_to_enums(labels))
          # Add it to the family
          family_proto.metric.extend([metric_proto])
      return family_proto
  
  
  def _goStringToFloat(s):
      if s == '+Inf':
          return float("inf")
      elif s == '-Inf':
          return float("-inf")
      elif s == 'NaN':
          return float('nan')
      else:
          return float(s)
  
  
  def _convert_labels_to_enums(labels):
      """
      Try to convert both the label names and label values to enum values.
      Defaults to the given name and value if it fails to convert.
      Arguments:
          labels: an array of label pairs that may contain enum names
      Returns:
          an array of label pairs with enum names converted to enum values
      """
      new_labels = []
      for name, value in labels:
          try:
              name = str(metricsd_pb2.MetricLabelName.Value(name))
          except ValueError as e:
              logging.debug(e)
          new_labels.append(metrics_pb2.LabelPair(name=name, value=value))
      return new_labels
  