#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You 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.

#

# pytype: skip-file

import collections

import time

from functools import reduce

from typing import FrozenSet

from typing import Hashable

from typing import List

from typing import Union

from apache_beam.coders import coder_impl

from apache_beam.coders import coders

from apache_beam.metrics.cells import BoundedTrieData

from apache_beam.metrics.cells import DistributionData

from apache_beam.metrics.cells import DistributionResult

from apache_beam.metrics.cells import GaugeData

from apache_beam.metrics.cells import GaugeResult

from apache_beam.metrics.cells import HistogramData

from apache_beam.metrics.cells import StringSetData

from apache_beam.portability import common_urns

from apache_beam.portability.api import metrics_pb2

SAMPLED_BYTE_SIZE_URN = (

common_urns.monitoring_info_specs.SAMPLED_BYTE_SIZE.spec.urn)

ELEMENT_COUNT_URN = common_urns.monitoring_info_specs.ELEMENT_COUNT.spec.urn

START_BUNDLE_MSECS_URN = (

common_urns.monitoring_info_specs.START_BUNDLE_MSECS.spec.urn)

PROCESS_BUNDLE_MSECS_URN = (

common_urns.monitoring_info_specs.PROCESS_BUNDLE_MSECS.spec.urn)

FINISH_BUNDLE_MSECS_URN = (

common_urns.monitoring_info_specs.FINISH_BUNDLE_MSECS.spec.urn)

TOTAL_MSECS_URN = common_urns.monitoring_info_specs.TOTAL_MSECS.spec.urn

USER_COUNTER_URN = common_urns.monitoring_info_specs.USER_SUM_INT64.spec.urn

USER_HISTOGRAM_URN = common_urns.monitoring_info_specs.USER_HISTOGRAM.spec.urn

USER_DISTRIBUTION_URN = (

common_urns.monitoring_info_specs.USER_DISTRIBUTION_INT64.spec.urn)

USER_GAUGE_URN = common_urns.monitoring_info_specs.USER_LATEST_INT64.spec.urn

USER_STRING_SET_URN = common_urns.monitoring_info_specs.USER_SET_STRING.spec.urn

USER_BOUNDED_TRIE_URN = (

common_urns.monitoring_info_specs.USER_BOUNDED_TRIE.spec.urn)

USER_METRIC_URNS = set([

USER_COUNTER_URN,

USER_DISTRIBUTION_URN,

USER_GAUGE_URN,

USER_STRING_SET_URN,

USER_BOUNDED_TRIE_URN,

USER_HISTOGRAM_URN

])

WORK_REMAINING_URN = common_urns.monitoring_info_specs.WORK_REMAINING.spec.urn

WORK_COMPLETED_URN = common_urns.monitoring_info_specs.WORK_COMPLETED.spec.urn

DATA_CHANNEL_READ_INDEX = (

common_urns.monitoring_info_specs.DATA_CHANNEL_READ_INDEX.spec.urn)

API_REQUEST_COUNT_URN = (

common_urns.monitoring_info_specs.API_REQUEST_COUNT.spec.urn)

# TODO(ajamato): Implement the remaining types, i.e. Double types

# Extrema types, etc. See:

# https://s.apache.org/beam-fn-api-metrics

SUM_INT64_TYPE = common_urns.monitoring_info_types.SUM_INT64_TYPE.urn

DISTRIBUTION_INT64_TYPE = (

common_urns.monitoring_info_types.DISTRIBUTION_INT64_TYPE.urn)

LATEST_INT64_TYPE = common_urns.monitoring_info_types.LATEST_INT64_TYPE.urn

PROGRESS_TYPE = common_urns.monitoring_info_types.PROGRESS_TYPE.urn

STRING_SET_TYPE = common_urns.monitoring_info_types.SET_STRING_TYPE.urn

BOUNDED_TRIE_TYPE = common_urns.monitoring_info_types.BOUNDED_TRIE_TYPE.urn

HISTOGRAM_TYPE = common_urns.monitoring_info_types.HISTOGRAM.urn

COUNTER_TYPES = set([SUM_INT64_TYPE])

DISTRIBUTION_TYPES = set([DISTRIBUTION_INT64_TYPE])

GAUGE_TYPES = set([LATEST_INT64_TYPE])

STRING_SET_TYPES = set([STRING_SET_TYPE])

BOUNDED_TRIE_TYPES = set([BOUNDED_TRIE_TYPE])

HISTOGRAM_TYPES = set([HISTOGRAM_TYPE])

# TODO(migryz) extract values from beam_fn_api.proto::MonitoringInfoLabels

PCOLLECTION_LABEL = (

common_urns.monitoring_info_labels.PCOLLECTION.label_props.name)

PTRANSFORM_LABEL = (

common_urns.monitoring_info_labels.TRANSFORM.label_props.name)

NAMESPACE_LABEL = (

common_urns.monitoring_info_labels.NAMESPACE.label_props.name)

NAME_LABEL = (common_urns.monitoring_info_labels.NAME.label_props.name)

SERVICE_LABEL = (common_urns.monitoring_info_labels.SERVICE.label_props.name)

METHOD_LABEL = (common_urns.monitoring_info_labels.METHOD.label_props.name)

RESOURCE_LABEL = (common_urns.monitoring_info_labels.RESOURCE.label_props.name)

STATUS_LABEL = (common_urns.monitoring_info_labels.STATUS.label_props.name)

BIGQUERY_PROJECT_ID_LABEL = (

common_urns.monitoring_info_labels.BIGQUERY_PROJECT_ID.label_props.name)

BIGQUERY_DATASET_LABEL = (

common_urns.monitoring_info_labels.BIGQUERY_DATASET.label_props.name)

BIGQUERY_TABLE_LABEL = (

common_urns.monitoring_info_labels.BIGQUERY_TABLE.label_props.name)

BIGQUERY_VIEW_LABEL = (

common_urns.monitoring_info_labels.BIGQUERY_VIEW.label_props.name)

BIGQUERY_QUERY_NAME_LABEL = (

common_urns.monitoring_info_labels.BIGQUERY_QUERY_NAME.label_props.name)

GCS_PROJECT_ID_LABEL = (

common_urns.monitoring_info_labels.GCS_PROJECT_ID.label_props.name)

GCS_BUCKET_LABEL = (

common_urns.monitoring_info_labels.GCS_BUCKET.label_props.name)

DATASTORE_PROJECT_ID_LABEL = (

common_urns.monitoring_info_labels.DATASTORE_PROJECT.label_props.name)

DATASTORE_NAMESPACE_LABEL = (

common_urns.monitoring_info_labels.DATASTORE_NAMESPACE.label_props.name)

SPANNER_PROJECT_ID = (

common_urns.monitoring_info_labels.SPANNER_PROJECT_ID.label_props.name)

SPANNER_DATABASE_ID = (

common_urns.monitoring_info_labels.SPANNER_DATABASE_ID.label_props.name)

SPANNER_TABLE_ID = (

common_urns.monitoring_info_labels.SPANNER_TABLE_ID.label_props.name)

SPANNER_QUERY_NAME = (

common_urns.monitoring_info_labels.SPANNER_QUERY_NAME.label_props.name)

BIGTABLE_PROJECT_ID_LABEL = (

common_urns.monitoring_info_labels.BIGTABLE_PROJECT_ID.label_props.name)

INSTANCE_ID_LABEL = (

common_urns.monitoring_info_labels.INSTANCE_ID.label_props.name)

TABLE_ID_LABEL = common_urns.monitoring_info_labels.TABLE_ID.label_props.name

def extract_counter_value(monitoring_info_proto):

"""Returns the counter value of the monitoring info."""

if not is_counter(monitoring_info_proto):

raise ValueError('Unsupported type %s' % monitoring_info_proto.type)

# Only SUM_INT64_TYPE is currently supported.

return coders.VarIntCoder().decode(monitoring_info_proto.payload)

def extract_gauge_value(monitoring_info_proto):

"""Returns a tuple containing (timestamp, value)"""

if not is_gauge(monitoring_info_proto):

raise ValueError('Unsupported type %s' % monitoring_info_proto.type)

# Only LATEST_INT64_TYPE is currently supported.

return _decode_gauge(coders.VarIntCoder(), monitoring_info_proto.payload)

def extract_distribution(monitoring_info_proto):

"""Returns a tuple of (count, sum, min, max).

Args:

proto: The monitoring info for the distribution.

"""

if not is_distribution(monitoring_info_proto):

raise ValueError('Unsupported type %s' % monitoring_info_proto.type)

# Only DISTRIBUTION_INT64_TYPE is currently supported.

return _decode_distribution(

coders.VarIntCoder(), monitoring_info_proto.payload)

def extract_string_set_value(monitoring_info_proto):

if not is_string_set(monitoring_info_proto):

raise ValueError('Unsupported type %s' % monitoring_info_proto.type)

coder = coders.IterableCoder(coders.StrUtf8Coder())

return set(coder.decode(monitoring_info_proto.payload))

def extract_bounded_trie_value(monitoring_info_proto):

if not is_bounded_trie(monitoring_info_proto):

raise ValueError('Unsupported type %s' % monitoring_info_proto.type)

return BoundedTrieData.from_proto(

metrics_pb2.BoundedTrie.FromString(monitoring_info_proto.payload))

def extract_histogram_value(monitoring_info_proto):

if not is_histogram(monitoring_info_proto):

raise ValueError('Unsupported type %s' % monitoring_info_proto.type)

return HistogramData.from_proto(

metrics_pb2.HistogramValue.FromString(monitoring_info_proto.payload))

def create_labels(ptransform=None, namespace=None, name=None, pcollection=None):

"""Create the label dictionary based on the provided values.

Args:

ptransform: The ptransform id used as a label.

pcollection: The pcollection id used as a label.

"""

labels = {}

if ptransform:

labels[PTRANSFORM_LABEL] = ptransform

if namespace:

labels[NAMESPACE_LABEL] = namespace

if name:

labels[NAME_LABEL] = name

if pcollection:

labels[PCOLLECTION_LABEL] = pcollection

return labels

def int64_user_counter(

namespace, name, metric, ptransform=None) -> metrics_pb2.MonitoringInfo:

"""Return the counter monitoring info for the specifed URN, metric and labels.

Args:

urn: The URN of the monitoring info/metric.

metric: The payload field to use in the monitoring info or an int value.

ptransform: The ptransform id used as a label.

"""

labels = create_labels(ptransform=ptransform, namespace=namespace, name=name)

if isinstance(metric, int):

metric = coders.VarIntCoder().encode(metric)

return create_monitoring_info(

USER_COUNTER_URN, SUM_INT64_TYPE, metric, labels)

def int64_counter(

urn,

metric,

ptransform=None,

pcollection=None,

labels=None) -> metrics_pb2.MonitoringInfo:

"""Return the counter monitoring info for the specifed URN, metric and labels.

Args:

urn: The URN of the monitoring info/metric.

metric: The payload field to use in the monitoring info or an int value.

ptransform: The ptransform id used as a label.

pcollection: The pcollection id used as a label.

"""

labels = labels or {}

labels.update(create_labels(ptransform=ptransform, pcollection=pcollection))

if isinstance(metric, int):

metric = coders.VarIntCoder().encode(metric)

return create_monitoring_info(urn, SUM_INT64_TYPE, metric, labels)

def int64_user_distribution(

namespace, name, metric, ptransform=None) -> metrics_pb2.MonitoringInfo:

"""Return the distribution monitoring info for the URN, metric and labels.

Args:

urn: The URN of the monitoring info/metric.

metric: The DistributionData for the metric.

ptransform: The ptransform id used as a label.

"""

labels = create_labels(ptransform=ptransform, namespace=namespace, name=name)

payload = _encode_distribution(

coders.VarIntCoder(), metric.count, metric.sum, metric.min, metric.max)

return create_monitoring_info(

USER_DISTRIBUTION_URN, DISTRIBUTION_INT64_TYPE, payload, labels)

def int64_distribution(

urn,

metric,

ptransform=None,

pcollection=None) -> metrics_pb2.MonitoringInfo:

"""Return a distribution monitoring info for the URN, metric and labels.

Args:

urn: The URN of the monitoring info/metric.

metric: The DistributionData for the metric.

ptransform: The ptransform id used as a label.

pcollection: The pcollection id used as a label.

"""

labels = create_labels(ptransform=ptransform, pcollection=pcollection)

payload = _encode_distribution(

coders.VarIntCoder(), metric.count, metric.sum, metric.min, metric.max)

return create_monitoring_info(urn, DISTRIBUTION_INT64_TYPE, payload, labels)

def int64_user_gauge(

namespace, name, metric, ptransform=None) -> metrics_pb2.MonitoringInfo:

"""Return the gauge monitoring info for the URN, metric and labels.

Args:

namespace: User-defined namespace of gauge metric.

name: Name of gauge metric.

metric: The GaugeData containing the metrics.

ptransform: The ptransform id used as a label.

"""

labels = create_labels(ptransform=ptransform, namespace=namespace, name=name)

if isinstance(metric, GaugeData):

coder = coders.VarIntCoder()

value = metric.value

timestamp = metric.timestamp

else:

raise TypeError(

'Expected GaugeData metric type but received %s with value %s' %

(type(metric), metric))

payload = _encode_gauge(coder, timestamp, value)

return create_monitoring_info(

USER_GAUGE_URN, LATEST_INT64_TYPE, payload, labels)

def int64_gauge(urn, metric, ptransform=None) -> metrics_pb2.MonitoringInfo:

"""Return the gauge monitoring info for the URN, metric and labels.

Args:

urn: The URN of the monitoring info/metric.

metric: An int representing the value. The current time will be used for

the timestamp.

ptransform: The ptransform id used as a label.

"""

labels = create_labels(ptransform=ptransform)

if isinstance(metric, int):

value = metric

time_ms = int(time.time()) * 1000

else:

raise TypeError(

'Expected int metric type but received %s with value %s' %

(type(metric), metric))

coder = coders.VarIntCoder()

payload = coder.encode(time_ms) + coder.encode(value)

return create_monitoring_info(urn, LATEST_INT64_TYPE, payload, labels)

def user_set_string(namespace, name, metric, ptransform=None):

"""Return the string set monitoring info for the URN, metric and labels.

Args:

namespace: User-defined namespace of StringSet.

name: Name of StringSet.

metric: The StringSetData representing the metrics.

ptransform: The ptransform id used as a label.

"""

labels = create_labels(ptransform=ptransform, namespace=namespace, name=name)

if isinstance(metric, StringSetData):

metric = metric.string_set

if isinstance(metric, set):

metric = list(metric)

if isinstance(metric, list):

metric = coders.IterableCoder(coders.StrUtf8Coder()).encode(metric)

return create_monitoring_info(

USER_STRING_SET_URN, STRING_SET_TYPE, metric, labels)

def user_histogram(namespace, name, metric: HistogramData, ptransform=None):

"""Return the histogram monitoring info for the URN, metric and labels.

Args:

namespace: User-defined namespace of Histogram.

name: Name of Histogram.

metric: The Histogram representing the metrics.

ptransform: The ptransform id used as a label.

"""

labels = create_labels(ptransform=ptransform, namespace=namespace, name=name)

metric_proto = metric.to_proto()

return create_monitoring_info(

USER_HISTOGRAM_URN,

HISTOGRAM_TYPE,

metric_proto.SerializeToString(),

labels)

def user_bounded_trie(namespace, name, metric, ptransform=None):

"""Return the string set monitoring info for the URN, metric and labels.

Args:

namespace: User-defined namespace of BoundedTrie.

name: Name of BoundedTrie.

metric: The BoundedTrieData representing the metrics.

ptransform: The ptransform id used as a label.

"""

labels = create_labels(ptransform=ptransform, namespace=namespace, name=name)

return create_monitoring_info(

USER_BOUNDED_TRIE_URN,

BOUNDED_TRIE_TYPE,

metric.to_proto().SerializeToString(),

labels)

def create_monitoring_info(

urn, type_urn, payload, labels=None) -> metrics_pb2.MonitoringInfo:

"""Return the monitoring info for the URN, type, metric and labels.

Args:

urn: The URN of the monitoring info/metric.

type_urn: The URN of the type of the monitoring info/metric.

i.e. beam:metrics:sum_int_64, beam:metrics:latest_int_64.

payload: The payload field to use in the monitoring info.

labels: The label dictionary to use in the MonitoringInfo.

"""

try:

return metrics_pb2.MonitoringInfo(

urn=urn, type=type_urn, labels=labels or {}, payload=payload)

except TypeError as e:

raise RuntimeError(

f'Failed to create MonitoringInfo for urn {urn} type {type_urn} '

f'labels {labels} and payload {payload}') from e

def is_counter(monitoring_info_proto):

"""Returns true if the monitoring info is a coutner metric."""

return monitoring_info_proto.type in COUNTER_TYPES

def is_gauge(monitoring_info_proto):

"""Returns true if the monitoring info is a gauge metric."""

return monitoring_info_proto.type in GAUGE_TYPES

def is_distribution(monitoring_info_proto):

"""Returns true if the monitoring info is a distrbution metric."""

return monitoring_info_proto.type in DISTRIBUTION_TYPES

def is_histogram(monitoring_info_proto):

"""Returns true if the monitoring info is a distrbution metric."""

return monitoring_info_proto.type in HISTOGRAM_TYPES

def is_string_set(monitoring_info_proto):

"""Returns true if the monitoring info is a StringSet metric."""

return monitoring_info_proto.type in STRING_SET_TYPES

def is_bounded_trie(monitoring_info_proto):

"""Returns true if the monitoring info is a BoundedTrie metric."""

return monitoring_info_proto.type in BOUNDED_TRIE_TYPES

def is_user_monitoring_info(monitoring_info_proto):

"""Returns true if the monitoring info is a user metric."""

return monitoring_info_proto.urn in USER_METRIC_URNS

def extract_metric_result_map_value(

monitoring_info_proto

) -> Union[None, int, DistributionResult, GaugeResult, set, BoundedTrieData]:

"""Returns the relevant GaugeResult, DistributionResult or int value for

counter metric, set for StringSet metric.

These are the proper format for use in the MetricResult.query() result.

"""

# Returns a metric result (AKA the legacy format).

# from the MonitoringInfo

if is_counter(monitoring_info_proto):

return extract_counter_value(monitoring_info_proto)

if is_distribution(monitoring_info_proto):

(count, sum, min, max) = extract_distribution(monitoring_info_proto)

return DistributionResult(DistributionData(sum, count, min, max))

if is_gauge(monitoring_info_proto):

(timestamp, value) = extract_gauge_value(monitoring_info_proto)

return GaugeResult(GaugeData(value, timestamp))

if is_string_set(monitoring_info_proto):

return extract_string_set_value(monitoring_info_proto)

if is_bounded_trie(monitoring_info_proto):

return extract_bounded_trie_value(monitoring_info_proto)

return None

def parse_namespace_and_name(monitoring_info_proto):

"""Returns the (namespace, name) tuple of the URN in the monitoring info."""

# Remove the URN prefix which indicates that it is a user counter.

if is_user_monitoring_info(monitoring_info_proto):

labels = monitoring_info_proto.labels

return labels[NAMESPACE_LABEL], labels[NAME_LABEL]

# If it is not a user counter, just use the first part of the URN, i.e. 'beam'

split = monitoring_info_proto.urn.split(':', 1)

return split[0], split[1]

def get_step_name(monitoring_info_proto):

"""Returns a step name for the given monitoring info or None if step name

cannot be specified."""

# Right now only metrics that have a PTRANSFORM are taken into account

return monitoring_info_proto.labels.get(PTRANSFORM_LABEL)

def to_key(

monitoring_info_proto: metrics_pb2.MonitoringInfo) -> FrozenSet[Hashable]:

"""Returns a key based on the URN and labels.

This is useful in maps to prevent reporting the same MonitoringInfo twice.

"""

key_items: List[Hashable] = list(monitoring_info_proto.labels.items())

key_items.append(monitoring_info_proto.urn)

return frozenset(key_items)

def sum_payload_combiner(payload_a, payload_b):

coder = coders.VarIntCoder()

return coder.encode(coder.decode(payload_a) + coder.decode(payload_b))

def distribution_payload_combiner(payload_a, payload_b):

coder = coders.VarIntCoder()

(count_a, sum_a, min_a, max_a) = _decode_distribution(coder, payload_a)

(count_b, sum_b, min_b, max_b) = _decode_distribution(coder, payload_b)

return _encode_distribution(

coder,

count_a + count_b,

sum_a + sum_b,

min(min_a, min_b),

max(max_a, max_b))

_KNOWN_COMBINERS = {

SUM_INT64_TYPE: sum_payload_combiner,

DISTRIBUTION_INT64_TYPE: distribution_payload_combiner,

}

def consolidate(metrics, key=to_key):

grouped = collections.defaultdict(list)

for metric in metrics:

grouped[key(metric)].append(metric)

for values in grouped.values():

if len(values) == 1:

yield values[0]

else:

combiner = _KNOWN_COMBINERS.get(values[0].type)

if combiner:

def merge(a, b):

# pylint: disable=cell-var-from-loop

return metrics_pb2.MonitoringInfo(

urn=a.urn,

type=a.type,

labels=dict((label, value) for label, value in a.labels.items()

if b.labels.get(label) == value),

payload=combiner(a.payload, b.payload))

yield reduce(merge, values)

else:

for value in values:

yield value

def _decode_gauge(coder, payload):

"""Returns a tuple of (timestamp, value)."""

timestamp_coder = coders.VarIntCoder().get_impl()

stream = coder_impl.create_InputStream(payload)

time_ms = timestamp_coder.decode_from_stream(stream, True)

return (time_ms / 1000.0, coder.get_impl().decode_from_stream(stream, True))

def _encode_gauge(coder, timestamp, value):

timestamp_coder = coders.VarIntCoder().get_impl()

stream = coder_impl.create_OutputStream()

timestamp_coder.encode_to_stream(int(timestamp * 1000), stream, True)

coder.get_impl().encode_to_stream(value, stream, True)

return stream.get()

def _decode_distribution(value_coder, payload):

"""Returns a tuple of (count, sum, min, max)."""

count_coder = coders.VarIntCoder().get_impl()

value_coder = value_coder.get_impl()

stream = coder_impl.create_InputStream(payload)

return (

count_coder.decode_from_stream(stream, True),

value_coder.decode_from_stream(stream, True),

value_coder.decode_from_stream(stream, True),

value_coder.decode_from_stream(stream, True))

def _encode_distribution(value_coder, count, sum, min, max):

count_coder = coders.VarIntCoder().get_impl()

value_coder = value_coder.get_impl()

stream = coder_impl.create_OutputStream()

count_coder.encode_to_stream(count, stream, True)

value_coder.encode_to_stream(sum, stream, True)

value_coder.encode_to_stream(min, stream, True)

value_coder.encode_to_stream(max, stream, True)

return stream.get()