pdk-metrics-lib 1.7.0

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

use crate::MetricsInstance;
use classy::extract::context::ConfigureContext;
use classy::extract::{Extract, FromContext};
use pdk_core::classy::{MetricType, MetricsHost};
use pdk_core::policy_context::api::Metadata;
use std::collections::BTreeMap;
use std::convert::Infallible;
use std::rc::Rc;

const SOURCE_TAG: &str = "source";
const SOURCE_TAG_VALUE: &str = "custom-metrics";
const CATEGORY_TAG: &str = "category";
const NAME_FORBIDDEN: char = '.';
const TAG_SEPARATOR: char = ',';
const TAG_VALUE_SEPARATOR: char = '=';
const REPLACE_CHAR: char = '_';

/// Struct that can be injected into your context to handle a [`Metric`](super::Metric).
pub struct MetricsBuilder {
    metrics: Rc<dyn MetricsHost>,
    metadata: Rc<Metadata>,
}

/// It can injected by the framework during the configuration phase of the policy lifecycle:
///
/// ``` rust
/// #[entrypoint]
/// async fn configure(
///     launcher: Launcher,
///     Configuration(configuration): Configuration,
///     metrics_builder: MetricsBuilder,
/// ) -> Result<()> {
///     // your code here.
/// }
impl FromContext<ConfigureContext> for MetricsBuilder {
    type Error = Infallible;

    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
        let metrics: Rc<dyn MetricsHost> = context.extract()?;
        let metadata: Rc<Metadata> = Rc::new(context.extract()?);

        Ok(MetricsBuilder { metrics, metadata })
    }
}

impl MetricsBuilder {
    fn add_default_tags(&self, metric: MetricsInstanceBuilder) -> MetricsInstanceBuilder {
        metric.tag(SOURCE_TAG, SOURCE_TAG_VALUE).tag(
            CATEGORY_TAG,
            self.metadata.policy_metadata.filter_name.as_str(),
        )
    }

    pub fn counter(&self, name: &str) -> MetricsInstanceBuilder {
        self.add_default_tags(MetricsInstanceBuilder::new(
            name,
            MetricType::Counter,
            Rc::clone(&self.metrics),
        ))
    }

    pub fn gauge(&self, name: &str) -> MetricsInstanceBuilder {
        self.add_default_tags(MetricsInstanceBuilder::new(
            name,
            MetricType::Gauge,
            Rc::clone(&self.metrics),
        ))
    }
}

/// A particular instance of the [`MetricsBuilder`].
pub struct MetricsInstanceBuilder {
    metrics: Rc<dyn MetricsHost>,
    metric_type: MetricType,
    name: String,
    tags: BTreeMap<String, String>,
}

impl MetricsInstanceBuilder {
    pub fn new(name: &str, metric_type: MetricType, metrics: Rc<dyn MetricsHost>) -> Self {
        Self {
            metrics,
            metric_type,
            name: sanitize_name(name),
            tags: BTreeMap::new(),
        }
    }

    pub fn tag(mut self, key: &str, value: &str) -> Self {
        self.tags.insert(sanitize_tag(key), sanitize_tag(value));
        self
    }

    pub fn skip_default_tags(mut self) -> Self {
        self.tags
            .retain(|k, _| *k != CATEGORY_TAG && *k != SOURCE_TAG);
        self
    }

    pub fn build(self) -> MetricsInstance {
        let mut string = String::with_capacity(
            self.name.len()
                + 2 * self.tags.len() // TAG_SEPARATOR && TAG_VALUE_SEPARATOR
                + self
                    .tags
                    .iter()
                    .map(|(k, v)| k.len() + v.len())
                    .sum::<usize>(),
        );
        string.push_str(&self.name);
        for (k, v) in self.tags.into_iter() {
            string.push(TAG_SEPARATOR);
            string.push_str(&k);
            string.push(TAG_VALUE_SEPARATOR);
            string.push_str(&v);
        }

        let id = self.metrics.define_metric(self.metric_type, &string);

        pdk_core::logger::debug!("Defined metric: {string}, with id {id}");

        MetricsInstance {
            id,
            metrics: self.metrics,
        }
    }
}

fn sanitize_name(value: &str) -> String {
    sanitize(
        value,
        &[NAME_FORBIDDEN, TAG_VALUE_SEPARATOR, TAG_SEPARATOR],
        REPLACE_CHAR,
    )
}

fn sanitize_tag(value: &str) -> String {
    sanitize(value, &[TAG_VALUE_SEPARATOR, TAG_SEPARATOR], REPLACE_CHAR)
}

fn sanitize(value: &str, matching_chars: &[char], replace: char) -> String {
    let mut result = String::with_capacity(value.len());
    for char in value.chars() {
        if matching_chars.contains(&char) {
            result.push(replace)
        } else {
            result.push(char)
        }
    }
    result
}