mz_adapter/coord/

metric_sink.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! Coordinator-installed metric sinks, the curated counterpart to `CREATE METRIC SINK`.
//!
//! A curated metric sink is a [`CURATED`] entry rendered on every replica, publishing its series
//! into that replica's process-local Prometheus registry. Unlike a user's `CREATE METRIC SINK` it
//! is not a catalog item: it gets a transient [`GlobalId`], targets one replica rather than a
//! cluster, and is re-created from the static list on every boot. Modelling the curated set this
//! way keeps it out of the catalog, so adding or removing a definition needs no builtin migration.
//!
//! Every replica means every replica of every cluster, user clusters included. Each definition is
//! therefore a dataflow, with its arrangements, on customer compute, charged to that customer's
//! cluster, and the cost scales with `CURATED`. `coord::introspection` already accepts this for its
//! subscribes.
//!
//! `install_metric_sinks` installs every definition on a newly created replica
//! (`bootstrap_metric_sinks` covers the replicas already present at startup), and
//! `drop_metric_sinks` drops them before a replica is dropped. This mirrors
//! [`crate::coord::introspection`], which installs introspection subscribes on the same triggers.
//!
//! The `disabled_metric_sinks` system var denies definitions by name, and
//! `reconcile_metric_sinks` converges the installed set on it, tearing a denied sink down rather
//! than only gating future installs.

use std::collections::{BTreeMap, BTreeSet};

use anyhow::bail;
use mz_catalog::memory::objects::CatalogItem;
use mz_cluster_client::ReplicaId;
use mz_controller_types::ClusterId;
use mz_ore::collections::CollectionExt;
use mz_ore::{instrument, soft_panic_or_log};
use mz_repr::optimize::OverrideFrom;
use mz_repr::{CatalogItemId, GlobalId, RelationDesc};
use mz_sql::catalog::SessionCatalog;
use mz_sql::plan::{
    HirRelationExpr, METRIC_SINK_CURATED_PREFIX_MARKER, Params, Plan, SubscribeFrom, SubscribePlan,
    validate_metric_sink_desc, validate_metric_sink_prefix,
};
use mz_sql::session::user::{MZ_SYSTEM_ROLE_ID, RoleMetadata};
use mz_sql::session::vars::{ENABLE_METRIC_SINK, SystemVars};
use tracing::{Span, info, warn};

use crate::catalog::Catalog;
use crate::coord::{
    Coordinator, Message, MetricSinkFinish, MetricSinkOptimize, MetricSinkStage, PlanValidity,
    StageResult, Staged,
};
use crate::optimize::Optimize;
use crate::optimize::dataflows::dataflow_import_id_bundle;
use crate::{AdapterError, ExecuteResponse, optimize};

/// A curated metric sink: SQL producing the canonical metric-sink columns, plus the name it is
/// known by in logs.
#[derive(Debug)]
pub(super) struct CuratedMetricSink {
    /// Stable identifier for the definition: used in logs, as the [`Coordinator::metric_sinks`] key,
    /// and as the `sink` label on the health gauges (the `GlobalId` is transient, the name is not).
    /// Must be unique within [`CURATED`].
    name: &'static str,
    /// A `SELECT` producing the canonical metric-sink columns (`metric_name`, `metric_type`,
    /// `labels`, `value`, `help`), the contract `mz_sql::plan::validate_metric_sink_desc` checks.
    ///
    /// The query must read only introspection relations. A catalog-backed relation would put
    /// envd's write frontier on the sink's emission path, which is exactly the coupling these
    /// sinks exist to avoid: the sink would stall whenever envd did, taking the freshness signal
    /// with it.
    source_sql: &'static str,
    /// Prepended to every row's `metric_name` to form the published name, exactly as a user's
    /// `CREATE METRIC SINK ... WITH (PREFIX = ...)`. Every definition in [`CURATED`] uses
    /// [`METRIC_SINK_CURATED_PREFIX_MARKER`], which user sinks are barred from, so nothing a user
    /// publishes can collide with a curated family.
    prefix: &'static str,
}

/// The curated metric sinks, installed on every replica.
///
/// Sources read the raw `..._raw` logging relations, not a derived view that re-aggregates them
/// (`mz_dataflow_arrangement_sizes`) or a join view over the logs (`mz_dataflow_operator_dataflows`),
/// for performance: those churn even on static data and their cost scales with the replica's
/// dataflow-creation rate. A single-relation filter like `mz_compute_exports` carries no aggregation
/// and is read freely.
///
/// Every family sums across workers, so a series carries no `worker_id`, and a multi-process replica
/// reports one number per grouping key rather than one per process. The size families emit one series
/// per dataflow, the errors family one per export.
const CURATED: &[CuratedMetricSink] = &[
    CuratedMetricSink {
        name: "mz_metric_arrangement_sizes",
        prefix: METRIC_SINK_CURATED_PREFIX_MARKER,
        // Key on the export id from `mz_compute_exports`, not `mz_dataflow_global_ids`: a
        // materialized view builds under a transient view id and appears only as an export, so a
        // global-id label names a `t<N>` that maps to no catalog object and churns on every
        // re-render.
        //
        // Logs are per operator, so map operator -> dataflow -> export id. Take the operator ->
        // dataflow step off `mz_dataflow_addresses_per_worker` (`address[1]` is the dataflow id),
        // which avoids the multi-way join behind `mz_dataflow_operator_dataflows`.
        //
        // `min(export_id)` collapses a multi-export dataflow to one series (lexicographic, so
        // `min('u10', 'u2')` is `'u10'`: arbitrary but stable). The group-size hint stops that `min`
        // from rendering the 8-level hierarchy, which would otherwise show up as tuning advice for
        // the sink's own dataflow in `mz_expected_group_size_advice`.
        //
        // NOTE: counting an arrangement this sink reads is a feedback loop. Every change to it
        // changes its logged size, the sink reads that on the next logging tick, and the dataflow
        // re-runs every tick for the replica's lifetime (SQL-730). `NOT LIKE 't%'` keeps out the
        // sink's own dataflow and the replica's introspection subscribes. `NOT LIKE 'si%'` keeps
        // out the logging dataflow: `si<N>` names only the introspection source indexes it
        // exports, and a plain system index prints `s<N>`. Only that filter excludes it under
        // `INTROSPECTION DEBUGGING`, which registers the loggers before the logging dataflow is
        // built, so its own operators get log rows. The cost is that transient dataflows'
        // arrangements go unreported, so these families sum below what the replica holds.
        //
        // `f` joins each raw log separately to reuse its `(operator_id, worker_id)` index
        // (`LogVariant::index_by`). One union would arrange all three logs' rows afresh.
        source_sql: "
WITH ex AS (
    SELECT dataflow_id, min(export_id) AS export_id
    FROM mz_introspection.mz_compute_exports
    WHERE export_id NOT LIKE 't%' AND export_id NOT LIKE 'si%'
    GROUP BY dataflow_id OPTIONS (AGGREGATE INPUT GROUP SIZE = 1)
),
oe AS (
    SELECT a.id, a.worker_id, ex.export_id
    FROM mz_introspection.mz_dataflow_addresses_per_worker a
    JOIN ex ON ex.dataflow_id = a.address[1]
),
f AS (
    SELECT 'arrangement_size_bytes'::text AS metric_name, oe.export_id
    FROM mz_introspection.mz_arrangement_heap_size_raw r
    JOIN oe ON r.operator_id = oe.id AND r.worker_id = oe.worker_id
    UNION ALL
    SELECT 'arrangement_records'::text, oe.export_id
    FROM mz_introspection.mz_arrangement_records_raw r
    JOIN oe ON r.operator_id = oe.id AND r.worker_id = oe.worker_id
    UNION ALL
    SELECT 'arrangement_batches'::text, oe.export_id
    FROM mz_introspection.mz_arrangement_batches_raw r
    JOIN oe ON r.operator_id = oe.id AND r.worker_id = oe.worker_id
)
SELECT metric_name, 'gauge'::text AS metric_type,
       map_build(LIST[ROW('id', export_id)])::map[text=>text] AS labels,
       count(*)::double precision AS value,
       CASE metric_name
           WHEN 'arrangement_size_bytes' THEN 'arrangement heap size in bytes'
           WHEN 'arrangement_records' THEN 'number of records in arrangement heaps'
           WHEN 'arrangement_batches' THEN 'number of batches in arrangements'
       END AS help
FROM f
GROUP BY metric_name, export_id",
    },
    CuratedMetricSink {
        name: "mz_metric_dataflow_errors",
        prefix: METRIC_SINK_CURATED_PREFIX_MARKER,
        // Raw log, not the `mz_compute_error_counts` view: the view joins the storage-managed
        // `mz_internal.mz_compute_dependencies`, which `ensure_reads_only_logs` rejects. `count` is
        // per-worker, so sum per export. `HAVING` drops the healthy ones.
        //
        // NOTE: direct errors only. The raw log attributes an error to the export that raised it. The
        // view also forwards counts onto index-reuse exports, so a broken reuse-index reads 0 here and
        // its errors show under the underlying export, undercounting against the view.
        source_sql: "
SELECT 'dataflow_error_count'::text AS metric_name, 'gauge'::text AS metric_type,
       map_build(LIST[ROW('id', export_id::text)])::map[text=>text] AS labels,
       sum(count)::double precision AS value, 'count of errors in the dataflow'::text AS help
FROM mz_introspection.mz_compute_error_counts_raw
GROUP BY export_id
HAVING sum(count) > 0",
    },
];

/// A [`CuratedMetricSink`] installed on one replica.
#[derive(Debug)]
pub(super) struct InstalledMetricSink {
    /// The cluster the replica belongs to, needed to drop the sink's compute collection.
    cluster_id: ClusterId,
    /// The transient id of the sink's compute export.
    sink_id: GlobalId,
}

/// A [`CuratedMetricSink`] planned once and shared across the replicas it installs on. See
/// [`Coordinator::plan_metric_sink`].
#[derive(Clone, Debug)]
pub(super) struct PlannedMetricSink {
    /// The shaped source query.
    expr: HirRelationExpr,
    /// The shape `expr` produces.
    desc: RelationDesc,
    /// The catalog items the source reads.
    dependencies: BTreeSet<CatalogItemId>,
}

impl Coordinator {
    /// Installs the curated metric sinks on all existing replicas.
    pub(super) async fn bootstrap_metric_sinks(&mut self) {
        for (cluster_id, replica_id) in self.all_cluster_replicas() {
            self.install_metric_sinks(cluster_id, replica_id).await;
        }
    }

    /// Installs the curated metric sinks on the given replica.
    ///
    /// Turning `enable_metric_sink` off stops installing on replicas created from then on. It does
    /// not tear down what is already installed: those keep running until their replica is dropped
    /// or envd restarts. A replica that merely reconnects re-renders them from the controller's
    /// command history, so a replica restart does not clear them either.
    pub(super) async fn install_metric_sinks(
        &mut self,
        cluster_id: ClusterId,
        replica_id: ReplicaId,
    ) {
        if !ENABLE_METRIC_SINK.enabled(self.catalog().system_config()) {
            return;
        }

        // TODO: Skip replicas created with introspection disabled. Their logging dataflows never
        // run, so a `source_sql` reading introspection relations there never advances. That is not
        // just wasted work: the sink publishes its input frontier as its write frontier, so a
        // never-advancing input stalls the sink's frontier at its as-of and pins the read holds it
        // takes on those collections for the replica's whole life (replica-local, released on
        // drop). `coord::introspection` installs subscribes on the same triggers and has the same
        // gap.
        for definition in CURATED {
            if metric_sink_denied(self.catalog().system_config(), definition.name) {
                continue;
            }
            self.install_metric_sink(cluster_id, replica_id, definition)
                .await;
        }
    }

    /// Converges the installed curated sinks on `disabled_metric_sinks`.
    ///
    /// Reconciles the whole set rather than the delta.
    pub(super) async fn reconcile_metric_sinks(&mut self) {
        for entry in self.catalog().system_config().disabled_metric_sinks() {
            if !CURATED.iter().any(|d| d.name == entry) {
                warn!(
                    name = %entry,
                    "disabled_metric_sinks entry matches no curated definition"
                );
            }
        }

        let denied: Vec<_> = self
            .metric_sinks
            .keys()
            .copied()
            .filter(|(_, name)| metric_sink_denied(self.catalog().system_config(), name))
            .collect();
        for (replica_id, name) in denied {
            self.drop_metric_sink(replica_id, name);
        }

        // Reinstall the full non-denied set on every replica. `install_metric_sink` is
        // idempotent: it skips a definition already recorded in `metric_sinks`
        // (`contains_key`, see `install_metric_sink`), so re-running the whole set only
        // installs the ones a preceding `disabled_metric_sinks` edit un-denied.
        for (cluster_id, replica_id) in self.all_cluster_replicas() {
            self.install_metric_sinks(cluster_id, replica_id).await;
        }
    }

    async fn install_metric_sink(
        &mut self,
        cluster_id: ClusterId,
        replica_id: ReplicaId,
        definition: &'static CuratedMetricSink,
    ) {
        // Cheap duplicate check before planning: if the definition is already installed on this
        // replica, there is nothing to do. `metric_sink_finish` keeps a backstop for a double
        // install still in flight (not yet recorded here).
        if self
            .metric_sinks
            .contains_key(&(replica_id, definition.name))
        {
            return;
        }

        let Some(planned) = self.plan_metric_sink(definition) else {
            return;
        };

        let (_, sink_id) = self.allocate_transient_id();
        // Logged only once the definition is known good, so an abandoned install leaves no
        // misleading "installing" line.
        info!(%sink_id, %replica_id, name = definition.name, "installing metric sink");

        let validity = PlanValidity::new(
            &self.catalog,
            planned.dependencies.clone(),
            Some(cluster_id),
            Some(replica_id),
            RoleMetadata::new(MZ_SYSTEM_ROLE_ID),
        );
        let stage = MetricSinkStage::Optimize(MetricSinkOptimize {
            validity,
            definition,
            sink_id,
            expr: planned.expr.clone(),
            desc: planned.desc.clone(),
            cluster_id,
            replica_id,
        });
        self.sequence_staged((), Span::current(), stage).await;
    }

    /// Plans a curated definition once, caching the result in [`Coordinator::metric_sink_plans`].
    ///
    /// The plan depends only on the catalog, never on the replica, so it is shared across every
    /// replica the definition installs on rather than re-planned per replica. Curated sources read
    /// only builtins (enforced by [`ensure_reads_only_logs`]), which do not change while envd runs,
    /// so a cached plan stays valid for envd's lifetime. Returns `None` for an invalid definition,
    /// having soft-panicked.
    fn plan_metric_sink(
        &mut self,
        definition: &'static CuratedMetricSink,
    ) -> Option<PlannedMetricSink> {
        if let Some(planned) = self.metric_sink_plans.get(definition.name) {
            return Some(planned.clone());
        }

        // A user sink's prefix is validated at plan time; a curated one has no such gate, so enforce
        // the same contract here. A failure is a bug in our own definition, hence
        // `soft_panic_or_log!`. User-vs-curated collisions need no check: the curated prefix is
        // reserved against user sinks in `validate_user_metric_sink_prefix`.
        //
        // NOTE: curated definitions are not checked against each other; they stay disjoint by
        // publishing distinct `metric_name`s under the shared curated prefix.
        if let Err(err) = validate_metric_sink_prefix(definition.prefix) {
            soft_panic_or_log!(
                "invalid curated metric sink prefix (name={}): {err}",
                definition.name
            );
            return None;
        }

        let catalog = self.catalog().for_system_session();
        let (expr, desc, dependencies) = match definition.plan_source(&catalog) {
            Ok(planned) => planned,
            Err(err) => {
                soft_panic_or_log!(
                    "invalid curated metric sink (name={}): {err}",
                    definition.name
                );
                return None;
            }
        };

        // Enforce the introspection-only contract before any optimization work, against what the
        // definition reads rather than how the optimizer imports it.
        if let Err(err) = ensure_reads_only_logs(&self.catalog, &dependencies) {
            soft_panic_or_log!(
                "invalid curated metric sink (name={}): {err}",
                definition.name
            );
            return None;
        }

        let planned = PlannedMetricSink {
            expr,
            desc,
            dependencies,
        };
        self.metric_sink_plans
            .insert(definition.name, planned.clone());
        Some(planned)
    }

    #[instrument]
    fn metric_sink_optimize(
        &self,
        stage: MetricSinkOptimize,
    ) -> Result<StageResult<Box<MetricSinkStage>>, AdapterError> {
        let MetricSinkOptimize {
            mut validity,
            definition,
            sink_id,
            expr,
            desc,
            cluster_id,
            replica_id,
        } = stage;

        let compute_instance = self
            .instance_snapshot(cluster_id)
            .expect("compute instance exists");
        // A transient id for the view the optimizer builds to shape the source rows, scoped to this
        // dataflow. See `optimize::metric_sink::shape_metric_sink_source`.
        let (_, view_id) = self.allocate_transient_id();

        let optimizer_config = optimize::OptimizerConfig::from(self.catalog().system_config())
            .override_from(&self.catalog.get_cluster(cluster_id).config.features())
            .override_from(&self.cluster_scoped_optimizer_overrides(cluster_id));

        let mut optimizer = optimize::metric_sink::Optimizer::new(
            self.owned_catalog(),
            compute_instance,
            view_id,
            sink_id,
            optimizer_config,
            self.optimizer_metrics(),
        );
        let catalog = self.owned_catalog();

        let span = Span::current();
        Ok(StageResult::Handle(mz_ore::task::spawn_blocking(
            || "optimize metric sink",
            move || {
                span.in_scope(|| {
                    let metric_sink = optimize::metric_sink::MetricSink::new(
                        format!("metric-sink-{}-{replica_id}", definition.name),
                        optimize::metric_sink::MetricSinkFrom::Query { expr, desc },
                        definition.prefix.to_string(),
                        Some(definition.name.to_string()),
                    );

                    // Both steps run inside one closure so either failure hits the same log.
                    // `sequence_staged` has no session to report to for a coordinator-driven
                    // install, so an error would otherwise vanish.
                    let global_lir_plan = (|| {
                        // MIR ⇒ MIR optimization (global)
                        let global_mir_plan = optimizer.catch_unwind_optimize(metric_sink)?;
                        // The optimizer imports indexes the SQL never named. Fold them into
                        // validity so one dropped before the finish stage fails the recheck rather
                        // than shipping a dataflow that imports a gone collection.
                        let id_bundle =
                            dataflow_import_id_bundle(global_mir_plan.df_desc(), cluster_id);
                        let item_ids = id_bundle.iter().map(|id| catalog.resolve_item_id(&id));
                        validity.extend_dependencies(&catalog, item_ids);
                        // MIR ⇒ LIR lowering and LIR ⇒ LIR optimization (global)
                        optimizer.catch_unwind_optimize(global_mir_plan)
                    })()
                    .inspect_err(|err| {
                        soft_panic_or_log!(
                            "curated metric sink failed to optimize (name={}): {err}",
                            definition.name
                        )
                    })?;

                    let stage = MetricSinkStage::Finish(MetricSinkFinish {
                        validity,
                        definition,
                        sink_id,
                        global_lir_plan,
                        cluster_id,
                        replica_id,
                    });
                    Ok(Box::new(stage))
                })
            },
        )))
    }

    #[instrument]
    async fn metric_sink_finish(
        &mut self,
        stage: MetricSinkFinish,
    ) -> Result<StageResult<Box<MetricSinkStage>>, AdapterError> {
        let MetricSinkFinish {
            validity: _,
            definition,
            sink_id,
            global_lir_plan,
            cluster_id,
            replica_id,
        } = stage;

        // `sequence_staged` rechecked validity before this stage ran, so the replica still exists.
        // The coordinator handles one message at a time, so no replica drop runs between that check
        // and the ship below.

        // The metainfo is dropped rather than persisted: a curated sink is not a catalog item, so
        // there is nothing for `mz_optimizer_notices` to hang its notices off.
        let (mut df_desc, _df_meta) = global_lir_plan.unapply();

        let id_bundle = dataflow_import_id_bundle(&df_desc, cluster_id);

        // Backstop for the introspection-only contract; the real gate is `ensure_reads_only_logs`
        // at install time. A log-only source imports only compute collections, so this should never
        // fire, but a storage import would couple the sink's frontier to envd.
        if !id_bundle.storage_ids.is_empty() {
            soft_panic_or_log!(
                "curated metric sink reads non-introspection relations (name={}): {:?}",
                definition.name,
                id_bundle.storage_ids
            );
            return Ok(StageResult::Response(ExecuteResponse::CreatedMetricSink));
        }

        // `reconcile_metric_sinks` only sees sinks already in `metric_sinks`, so a definition
        // denied while its install was in flight would ship anyway without this recheck.
        if metric_sink_denied(self.catalog().system_config(), definition.name) {
            return Ok(StageResult::Response(ExecuteResponse::CreatedMetricSink));
        }

        // Hold a read on the imports across shipping, so their since cannot advance past the as-of
        // just picked. Compute takes its own holds during `create_dataflow`.
        let read_holds = self.acquire_read_holds(&id_bundle);
        df_desc.set_as_of(read_holds.least_valid_read());

        // Record the install just before shipping. A failed plan or optimize returns earlier, so it
        // leaves no entry behind. `drop_metric_sinks` reads this entry to release the sink's
        // instance-global collection state on replica drop. Recording before the ship is safe because
        // the coordinator runs one message at a time with no await between the two, so no replica drop
        // sees an entry whose dataflow has not shipped.
        let install = InstalledMetricSink {
            cluster_id,
            sink_id,
        };
        if let Some(previous) = self
            .metric_sinks
            .insert((replica_id, definition.name), install)
        {
            // The key is already taken. `curated_names_are_unique` rules out a name collision,
            // so this is the same definition installed twice: reconcile can start a second
            // install while an earlier one is still in flight. Restore the first and abandon this
            // one, else we leak the first's collection (unreachable to `drop_metric_sinks`) and
            // register a second collector under the same `sink` label.
            self.metric_sinks
                .insert((replica_id, definition.name), previous);
            info!(
                %replica_id,
                name = definition.name,
                "abandoning metric sink install, already installed"
            );
            return Ok(StageResult::Response(ExecuteResponse::CreatedMetricSink));
        }

        self.ship_dataflow(df_desc, cluster_id, Some(replica_id))
            .await;

        drop(read_holds);
        // Nobody is waiting on this: `StagedContext for ()` drops the result. Reuses the
        // `CREATE METRIC SINK` response rather than adding a variant no client ever sees.
        Ok(StageResult::Response(ExecuteResponse::CreatedMetricSink))
    }

    /// Drops the curated metric sinks installed on the given replica.
    ///
    /// Called before the replica itself is dropped. Dropping the replica would tear the sink
    /// dataflows down anyway, but the controller's collection state for them is instance-global,
    /// so it has to be released explicitly.
    pub(super) fn drop_metric_sinks(&mut self, replica_id: ReplicaId) {
        for name in metric_sinks_on_replica(&self.metric_sinks, replica_id) {
            self.drop_metric_sink(replica_id, name);
        }
    }

    /// Drops one curated metric sink, if it is installed.
    fn drop_metric_sink(&mut self, replica_id: ReplicaId, name: &'static str) {
        let Some(install) = self.metric_sinks.remove(&(replica_id, name)) else {
            return;
        };
        let InstalledMetricSink {
            cluster_id,
            sink_id,
        } = install;
        info!(%sink_id, %replica_id, name, "dropping metric sink");

        // The entry exists only for a shipped dataflow, so its collection is present and this
        // drop succeeds. Result ignored: a failure during replica teardown is not worth a panic.
        let _ = self
            .controller
            .compute
            .drop_collections(cluster_id, vec![sink_id]);
    }
}

/// Whether `disabled_metric_sinks` denies the curated definition called `name`.
///
/// An entry naming no definition is never asked about, so a stale or misspelled one is inert.
fn metric_sink_denied(system_config: &SystemVars, name: &str) -> bool {
    system_config
        .disabled_metric_sinks()
        .iter()
        .any(|denied| denied == name)
}

/// The names of the definitions installed on `replica_id`, in key order.
///
/// The map is keyed replica-first, so a replica's installs are one contiguous range.
fn metric_sinks_on_replica(
    metric_sinks: &BTreeMap<(ReplicaId, &'static str), InstalledMetricSink>,
    replica_id: ReplicaId,
) -> Vec<&'static str> {
    metric_sinks
        .range((replica_id, "")..)
        .take_while(|((id, _), _)| *id == replica_id)
        .map(|((_, name), _)| *name)
        .collect()
}

/// Enforces the introspection-only contract from [`CuratedMetricSink::source_sql`]: every relation
/// the definition reads, walking views transitively, must be a log collection.
///
/// Checked here against what the definition reads rather than by import kind after optimization: the
/// import split (storage vs index) depends on which indexes the target cluster happens to have, so
/// it gives the same definition different verdicts on different clusters.
fn ensure_reads_only_logs(
    catalog: &Catalog,
    dependencies: &BTreeSet<CatalogItemId>,
) -> Result<(), anyhow::Error> {
    let mut to_visit: Vec<_> = dependencies.iter().copied().collect();
    let mut visited = BTreeSet::new();
    while let Some(id) = to_visit.pop() {
        if !visited.insert(id) {
            continue;
        }
        let entry = catalog.get_entry(&id);
        match entry.item() {
            // The only data leaf allowed.
            CatalogItem::Log(_) => {}
            // Allowed only if everything it reads is, so walk its dependencies.
            CatalogItem::View(_) => to_visit.extend(entry.uses()),
            // No data dependency; a view over logs still references these.
            CatalogItem::Type(_) | CatalogItem::Func(_) => {}
            _ => bail!(
                "curated metric sink reads {}, which is not an introspection log relation \
                 (only logs and views over logs are allowed)",
                catalog.resolve_full_name(entry.name(), None)
            ),
        }
    }
    Ok(())
}

impl CuratedMetricSink {
    /// Plans `source_sql` against a session-less catalog, returning the query, its output shape,
    /// and the catalog items it reads.
    fn plan_source(
        &self,
        catalog: &dyn SessionCatalog,
    ) -> Result<(HirRelationExpr, RelationDesc, BTreeSet<CatalogItemId>), anyhow::Error> {
        // A definition is a single statement. Reject the count explicitly for a clear error.
        let statements = mz_sql::parse::parse(self.source_sql)?;
        if statements.len() != 1 {
            bail!(
                "source SQL must be exactly one statement, got {}",
                statements.len()
            );
        }

        // A metric sink's source is a continuously maintained dataflow, like a SUBSCRIBE, so plan it
        // as one. A maintained lifetime folds any finishing into the expression (an ORDER BY over a
        // maintained collection is dropped, a LIMIT becomes a TopK) rather than leaving it beside the
        // query, so `MetricSinkFrom::Query` gets a self-contained expression whose arity matches its
        // `desc`. This mirrors `coord::introspection`, which plans its specs as subscribes too.
        let subscribe_sql = format!("SUBSCRIBE ({})", self.source_sql);
        let parsed = mz_sql::parse::parse(&subscribe_sql)?.into_element();
        let (stmt, resolved_ids) = mz_sql::names::resolve(catalog, parsed.ast)?;
        let (plan, sql_impl_ids) =
            mz_sql::plan::plan(None, catalog, stmt, &Params::empty(), &resolved_ids)?;
        let Plan::Subscribe(SubscribePlan {
            from: SubscribeFrom::Query { expr, desc },
            ..
        }) = plan
        else {
            bail!("source SQL must be a single SELECT");
        };
        validate_metric_sink_desc(&desc)?;

        // Fold in ids from SQL-implemented function bodies. `plan` keeps them out of `resolved_ids`
        // since a one-shot statement doesn't depend on a function's body, but a metric sink inlines
        // that body into its dataflow, so the body's reads are real imports the gate must check.
        let dependencies = resolved_ids
            .items()
            .chain(sql_impl_ids.items())
            .copied()
            .collect();
        Ok((expr, desc, dependencies))
    }
}

impl Staged for MetricSinkStage {
    type Ctx = ();

    fn validity(&mut self) -> &mut PlanValidity {
        match self {
            Self::Optimize(stage) => &mut stage.validity,
            Self::Finish(stage) => &mut stage.validity,
        }
    }

    async fn stage(
        self,
        coord: &mut Coordinator,
        _ctx: &mut (),
    ) -> Result<StageResult<Box<Self>>, AdapterError> {
        match self {
            Self::Optimize(stage) => coord.metric_sink_optimize(stage),
            Self::Finish(stage) => coord.metric_sink_finish(stage).await,
        }
    }

    fn message(self, _ctx: (), span: Span) -> Message {
        Message::MetricSinkStageReady { span, stage: self }
    }

    fn cancel_enabled(&self) -> bool {
        false
    }
}

#[cfg(test)]
mod tests {
    use std::collections::{BTreeMap, BTreeSet};

    use mz_catalog::memory::objects::CatalogItem;
    use mz_cluster_client::ReplicaId;
    use mz_controller_types::ClusterId;
    use mz_repr::GlobalId;
    use mz_sql::plan::{
        METRIC_SINK_CURATED_PREFIX_MARKER, validate_metric_sink_prefix,
        validate_user_metric_sink_prefix,
    };
    use mz_sql::session::vars::{DISABLED_METRIC_SINKS, SystemVars, Var, VarInput};

    use crate::catalog::Catalog;
    use crate::coord::metric_sink::{
        CURATED, CuratedMetricSink, InstalledMetricSink, ensure_reads_only_logs,
        metric_sink_denied, metric_sinks_on_replica,
    };

    /// `drop_metric_sinks` relies on this range scan returning exactly one replica's installs, with
    /// no bleed into a neighbouring replica's contiguous range.
    #[mz_ore::test]
    fn metric_sinks_on_replica_scans_one_replica() {
        let cluster = ClusterId::user(1).expect("valid cluster id");
        let install = |sink_id| InstalledMetricSink {
            cluster_id: cluster,
            sink_id: GlobalId::Transient(sink_id),
        };
        let r = ReplicaId::User;

        let mut sinks = BTreeMap::new();
        sinks.insert((r(1), "a"), install(10));
        sinks.insert((r(2), "a"), install(20));
        sinks.insert((r(2), "b"), install(21));
        sinks.insert((r(2), "c"), install(22));
        sinks.insert((r(4), "a"), install(40));

        // A replica with several installs: all of them, in key order, and nothing from r(1)/r(4).
        assert_eq!(metric_sinks_on_replica(&sinks, r(2)), vec!["a", "b", "c"]);
        // First and last replicas in the map: the scan stops at each boundary.
        assert_eq!(metric_sinks_on_replica(&sinks, r(1)), vec!["a"]);
        assert_eq!(metric_sinks_on_replica(&sinks, r(4)), vec!["a"]);
        // A replica with no installs, whether ordered between present ones (the r(3) gap) or past
        // the end, returns nothing rather than the next replica's range.
        assert!(metric_sinks_on_replica(&sinks, r(3)).is_empty());
        assert!(metric_sinks_on_replica(&sinks, r(5)).is_empty());
    }

    /// The var is a `Vec<Ident>`, so parsing follows the SQL identifier-list rules: surrounding
    /// whitespace is tolerated, an unquoted name folds to lowercase, and a quoted name keeps its
    /// case. Matching is otherwise exact (no prefix match).
    #[mz_ore::test]
    fn denylist_matches_names_leniently() {
        let denied = |list: &str, name: &str| {
            let mut vars = SystemVars::new();
            vars.set(DISABLED_METRIC_SINKS.name(), VarInput::Flat(list))
                .expect("valid denylist");
            metric_sink_denied(&vars, name)
        };

        assert!(!denied("", "a"));
        assert!(denied("a", "a"));
        assert!(denied("a,b", "b"));
        assert!(denied("  a , b  ", "a"));
        // An unknown name denies nothing but is carried without error.
        assert!(!denied("nope", "a"));
        assert!(denied("nope,a", "a"));
        // Exact match only: no prefix match.
        assert!(!denied("a", "ab"));
        assert!(!denied("ab", "a"));
        // Unquoted names fold to lowercase; quoting pins the case.
        assert!(denied("A", "a"));
        assert!(!denied("\"A\"", "a"));

        // An empty entry between commas is rejected by the identifier parser (unlike the old
        // naive split, which silently dropped it).
        let mut vars = SystemVars::new();
        assert!(
            vars.set(DISABLED_METRIC_SINKS.name(), VarInput::Flat("a,,b"))
                .is_err()
        );
    }

    #[mz_ore::test]
    fn curated_prefixes_are_valid() {
        for definition in CURATED {
            validate_metric_sink_prefix(definition.prefix).unwrap_or_else(|err| {
                panic!(
                    "curated metric sink {:?} has an invalid prefix {:?}: {err}",
                    definition.name, definition.prefix
                )
            });
        }
    }

    #[mz_ore::test]
    fn curated_prefixes_are_reserved_against_user_sinks() {
        for definition in CURATED {
            assert!(
                definition
                    .prefix
                    .starts_with(METRIC_SINK_CURATED_PREFIX_MARKER),
                "curated metric sink {:?} does not use the reserved curated prefix: {:?}",
                definition.name,
                definition.prefix
            );
            assert!(
                validate_user_metric_sink_prefix(definition.prefix).is_err(),
                "a user could claim curated metric sink {:?}'s prefix {:?}",
                definition.name,
                definition.prefix
            );
        }
    }

    /// The registry is keyed on the name, so a duplicate would make one definition's install
    /// unreachable to teardown and both collectors collide on the `sink` label. Guarded at runtime
    /// (`metric_sink_finish`) too, but caught here at build time before it can ship.
    #[mz_ore::test]
    fn curated_names_are_unique() {
        let mut seen = BTreeSet::new();
        for definition in CURATED {
            assert!(
                seen.insert(definition.name),
                "duplicate curated metric sink name {:?}",
                definition.name
            );
        }
    }

    /// Runs the two checks `plan_metric_sink` does at boot, where a failure soft-panics (a hard panic
    /// under debug assertions, so a CI boot crash-loop). The gate is the load-bearing one: a source
    /// can plan cleanly yet still read a storage-backed relation transitively (a builtin view joining
    /// in an introspection source), which only `ensure_reads_only_logs` catches.
    #[mz_ore::test(tokio::test)]
    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `TLS_client_method`
    async fn curated_definitions_plan() {
        Catalog::with_debug(|catalog| async move {
            let session_catalog = catalog.for_system_session();
            for definition in CURATED {
                let (_, _, dependencies) =
                    definition
                        .plan_source(&session_catalog)
                        .unwrap_or_else(|err| {
                            panic!(
                                "curated metric sink {:?} does not plan: {err}",
                                definition.name
                            )
                        });
                ensure_reads_only_logs(&catalog, &dependencies).unwrap_or_else(|err| {
                    panic!(
                        "curated metric sink {:?} reads a non-introspection relation: {err}",
                        definition.name
                    )
                });
            }
        })
        .await
    }

    /// The five canonical columns, no finishing: the shape a definition must produce.
    const VALID_SOURCE: &str = "SELECT 'n'::text AS metric_name, 'gauge'::text AS metric_type, \
        NULL::map[text=>text] AS labels, NULL::double AS value, 'h'::text AS help";

    /// `VALID_SOURCE` with an ORDER BY appended. Maintained-lifetime planning folds it away rather
    /// than rejecting it, since ordering has no meaning for a continuously-consumed collection.
    const ORDERED_SOURCE: &str = "SELECT 'n'::text AS metric_name, 'gauge'::text AS metric_type, \
        NULL::map[text=>text] AS labels, NULL::double AS value, 'h'::text AS help ORDER BY 1";

    /// `plan_source` accepts the canonical column contract (including a source with a finishing,
    /// which maintained-lifetime planning folds in) and rejects a source missing the columns.
    #[mz_ore::test(tokio::test)]
    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `TLS_client_method`
    async fn plan_source_enforces_the_metric_sink_contract() {
        Catalog::with_debug(|catalog| async move {
            let session_catalog = catalog.for_system_session();
            let plan = |source_sql: &'static str| {
                CuratedMetricSink {
                    name: "test",
                    source_sql,
                    prefix: "mz_metric_sink_test_",
                }
                .plan_source(&session_catalog)
            };

            assert!(plan(VALID_SOURCE).is_ok());

            // An ORDER BY is folded away by maintained-lifetime planning, not rejected.
            assert!(plan(ORDERED_SOURCE).is_ok());

            // Missing the canonical columns: rejected by `validate_metric_sink_desc`.
            assert!(plan("SELECT 1 AS foo").is_err());

            // Not exactly one statement: rejected by the explicit count guard.
            assert!(plan("").is_err());
            assert!(plan("SELECT 1; SELECT 2").is_err());
        })
        .await
    }

    /// A SQL-implemented builtin hides its reads: `pg_get_viewdef`'s body reads
    /// `mz_catalog.mz_views`, which the dataflow imports but the statement's resolved ids omit.
    /// `plan_source` must surface those reads so the gate rejects them.
    #[mz_ore::test(tokio::test)]
    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `TLS_client_method`
    async fn ensure_reads_only_logs_sees_sql_impl_function_reads() {
        Catalog::with_debug(|catalog| async move {
            let session_catalog = catalog.for_system_session();
            let (_, _, dependencies) = CuratedMetricSink {
                name: "test",
                source_sql: "SELECT pg_get_viewdef('x') AS metric_name, 'gauge'::text AS metric_type, \
                    NULL::map[text=>text] AS labels, NULL::double AS value, 'h'::text AS help",
                prefix: "mz_metric_sink_test_",
            }
            .plan_source(&session_catalog)
            .expect("plans against the system catalog");
            assert!(ensure_reads_only_logs(&catalog, &dependencies).is_err());
        })
        .await
    }

    /// The introspection-only contract: a log dependency is accepted, a storage-backed one is
    /// rejected. Checked against what the definition reads, so the verdict does not depend on the
    /// target cluster's index layout.
    #[mz_ore::test(tokio::test)]
    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `TLS_client_method`
    async fn ensure_reads_only_logs_accepts_logs_rejects_storage() {
        Catalog::with_debug(|catalog| async move {
            let log_id = catalog
                .entries()
                .find(|e| matches!(e.item(), CatalogItem::Log(_)))
                .expect("debug catalog has a builtin log")
                .id();
            assert!(ensure_reads_only_logs(&catalog, &BTreeSet::from([log_id])).is_ok());

            let storage_id = catalog
                .entries()
                .find(|e| matches!(e.item(), CatalogItem::Table(_) | CatalogItem::Source(_)))
                .expect("debug catalog has a builtin table or source")
                .id();
            assert!(ensure_reads_only_logs(&catalog, &BTreeSet::from([storage_id])).is_err());
        })
        .await
    }
}