Skip to main content

mz_adapter/optimize/
metric_sink.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Row-wise shaping for `MetricSink` sources.
11//!
12//! A metric sink exports the rows of an existing collection into the in-process Prometheus metrics
13//! registry. The compute-side operator (`mz_compute::sink::metric_sink`) reads a canonical row
14//! shape and two planner-computed classification columns rather than parsing `metric_type` strings
15//! or validating metric names on its hot path. `shape_metric_sink_source` produces that shape.
16//!
17//! Only the pure per-row shaping lives here. The cross-row logic (dedup, collision detection,
18//! family-conflict counting) stays in the operator, because it needs the frontier-gated fold that
19//! a per-row `Map` in MIR can't express.
20
21use std::sync::Arc;
22use std::time::{Duration, Instant};
23
24use mz_compute_types::plan::LirRelationExpr;
25use mz_compute_types::sinks::{ComputeSinkConnection, ComputeSinkDesc, MetricSinkConnection};
26use mz_expr::func::variadic::Coalesce;
27use mz_expr::{MirRelationExpr, MirScalarExpr, func};
28use mz_repr::explain::trace_plan;
29use mz_repr::{
30    ColumnName, Datum, GlobalId, RelationDesc, ReprRelationType, ReprScalarType, Row, SqlScalarType,
31};
32use mz_sql::optimizer_metrics::OptimizerMetrics;
33use mz_sql::plan::{HirRelationExpr, HirToMirConfig};
34use mz_transform::TransformCtx;
35use mz_transform::dataflow::DataflowMetainfo;
36use mz_transform::normalize_lets::normalize_lets;
37use mz_transform::typecheck::{SharedTypecheckingContext, empty_typechecking_context};
38use timely::progress::Antichain;
39
40use crate::optimize::dataflows::{
41    ComputeInstanceSnapshot, DataflowBuilder, ExprPrep, ExprPrepMaintained,
42};
43use crate::optimize::{
44    LirDataflowDescription, MirDataflowDescription, Optimize, OptimizerCatalog, OptimizerConfig,
45    OptimizerError, optimize_mir_local,
46};
47
48/// Matches Prometheus's metric name grammar: `[a-zA-Z_:][a-zA-Z0-9_:]*`.
49///
50/// Expressed in MIR (see `shape_metric_sink_source`) rather than parsed from a `&str` on the
51/// operator's hot path.
52const METRIC_NAME_PATTERN: &str = "^[a-zA-Z_:][a-zA-Z0-9_:]*$";
53
54/// Optimizer for metric sinks, both `CREATE METRIC SINK` and the coordinator-installed curated
55/// sinks.
56///
57/// The source is either an existing collection (like `CREATE INDEX`, no HIR to lower) or a planned
58/// query (like a materialized view), see [`MetricSinkFrom`]. Either way the row-wise shaping is
59/// appended in MIR and the dataflow exports a single `MetricSink`. Unlike a materialized view sink
60/// there is no persist shard, so there is no storage-metadata stage.
61pub struct Optimizer {
62    /// A representation typechecking context to use throughout the optimizer pipeline.
63    typecheck_ctx: SharedTypecheckingContext,
64    /// A snapshot of the catalog state.
65    catalog: Arc<dyn OptimizerCatalog>,
66    /// A snapshot of the cluster that will run the dataflow.
67    compute_instance: ComputeInstanceSnapshot,
68    /// A transient GlobalId for the shaped view built over the sink's source relation (see
69    /// `shape_metric_sink_source`).
70    view_id: GlobalId,
71    /// A durable GlobalId to be used with the exported metric sink.
72    sink_id: GlobalId,
73    /// Optimizer config.
74    config: OptimizerConfig,
75    /// Optimizer metrics.
76    metrics: OptimizerMetrics,
77    /// The time spent performing optimization so far.
78    duration: Duration,
79}
80
81impl Optimizer {
82    pub fn new(
83        catalog: Arc<dyn OptimizerCatalog>,
84        compute_instance: ComputeInstanceSnapshot,
85        view_id: GlobalId,
86        sink_id: GlobalId,
87        config: OptimizerConfig,
88        metrics: OptimizerMetrics,
89    ) -> Self {
90        Self {
91            typecheck_ctx: empty_typechecking_context(),
92            catalog,
93            compute_instance,
94            view_id,
95            sink_id,
96            config,
97            metrics,
98            duration: Default::default(),
99        }
100    }
101}
102
103/// A wrapper of metric sink parts needed to start the optimization process.
104pub struct MetricSink {
105    /// Names the assembled dataflow, for debugging.
106    debug_name: String,
107    /// The collection whose rows the sink exports.
108    from: MetricSinkFrom,
109    /// Prepended to every row's `metric_name` to form the published name. Validated as a valid start
110    /// of a Prometheus metric name before it reaches here (plan time for a user sink, install time
111    /// for a curated one).
112    prefix: String,
113    /// Value for the `sink` label on the sink's health gauges. `None` defaults to the sink's
114    /// `GlobalId`, which is what a user sink wants. A curated sink passes its stable name.
115    label: Option<String>,
116}
117
118impl MetricSink {
119    /// Construct a new [`MetricSink`]. Arguments are recorded as-is.
120    pub fn new(
121        debug_name: String,
122        from: MetricSinkFrom,
123        prefix: String,
124        label: Option<String>,
125    ) -> Self {
126        Self {
127            debug_name,
128            from,
129            prefix,
130            label,
131        }
132    }
133}
134
135/// Where a metric sink's rows come from.
136///
137/// Either way the source must expose the canonical metric-sink columns (see
138/// [`shape_metric_sink_source`]).
139pub enum MetricSinkFrom {
140    /// An existing catalog collection, as `CREATE METRIC SINK ... FROM <relation>` resolves to.
141    Id(GlobalId),
142    /// A planned query, as a coordinator-installed sink built from curated SQL uses. The query is
143    /// not a catalog item, so it is lowered and locally optimized here rather than imported.
144    Query {
145        expr: HirRelationExpr,
146        desc: RelationDesc,
147    },
148}
149
150/// The (sealed intermediate) result after embedding a [`MetricSink`] into a
151/// [`MirDataflowDescription`], inlining referenced views, and jointly optimizing the `MIR` plans.
152#[derive(Clone, Debug)]
153pub struct GlobalMirPlan {
154    df_desc: MirDataflowDescription,
155    df_meta: DataflowMetainfo,
156}
157
158impl GlobalMirPlan {
159    pub fn df_desc(&self) -> &MirDataflowDescription {
160        &self.df_desc
161    }
162}
163
164/// The (final) result after MIR ⇒ LIR lowering and optimizing the resulting
165/// `DataflowDescription` with `LIR` plans.
166#[derive(Clone, Debug)]
167pub struct GlobalLirPlan {
168    df_desc: LirDataflowDescription,
169    df_meta: DataflowMetainfo,
170}
171
172impl GlobalLirPlan {
173    pub fn df_desc(&self) -> &LirDataflowDescription {
174        &self.df_desc
175    }
176}
177
178impl Optimize<MetricSink> for Optimizer {
179    type To = GlobalMirPlan;
180
181    fn optimize(&mut self, metric_sink: MetricSink) -> Result<Self::To, OptimizerError> {
182        let time = Instant::now();
183
184        let mut df_builder = {
185            let compute = self.compute_instance.clone();
186            DataflowBuilder::new(&*self.catalog, compute).with_config(&self.config)
187        };
188        let mut df_desc = MirDataflowDescription::new(metric_sink.debug_name);
189        let mut df_meta = DataflowMetainfo::default();
190
191        let (source_expr, source_desc) = match metric_sink.from {
192            MetricSinkFrom::Id(from) => {
193                let from_desc = self
194                    .catalog
195                    .get_entry(&from)
196                    .relation_desc()
197                    .expect("can only create a metric sink on items with a valid description")
198                    .into_owned();
199                let repr_typ = ReprRelationType::from(from_desc.typ());
200                (MirRelationExpr::global_get(from, repr_typ), from_desc)
201            }
202            MetricSinkFrom::Query { expr, desc } => {
203                // HIR ⇒ MIR lowering and decorrelation. The result is inlined under the shaping
204                // below rather than becoming its own build, so the whole source is one view.
205                let expr = expr.lower(HirToMirConfig::from(&self.config), Some(&self.metrics))?;
206                (expr, desc)
207            }
208        };
209
210        // Push the pure row-wise shaping (coalesce identity elements, classify the metric kind,
211        // validate the metric name) into MIR, so the operator only does the cross-row logic
212        // (dedup/collision/family-conflict) that needs the fold. See `shape_metric_sink_source`.
213        let (shaped_expr, shaped_desc) =
214            shape_metric_sink_source(source_expr, &source_desc, &metric_sink.prefix);
215        let mut local_ctx = TransformCtx::local(
216            &self.config.features,
217            &self.typecheck_ctx,
218            &mut df_meta,
219            Some(&mut self.metrics),
220            Some(self.view_id),
221        );
222        let shaped_expr = optimize_mir_local(shaped_expr, &mut local_ctx)?;
223
224        // Imports the source's dependencies (the `Id` variant's collection, or the query's leaf
225        // collections) before inserting the shaped view that reads them.
226        df_builder.import_view_into_dataflow(
227            &self.view_id,
228            &shaped_expr,
229            &mut df_desc,
230            &self.config.features,
231        )?;
232        df_builder.maybe_reoptimize_imported_views(&mut df_desc, &self.config)?;
233
234        let sink_description = ComputeSinkDesc {
235            from: self.view_id,
236            from_desc: shaped_desc,
237            connection: ComputeSinkConnection::MetricSink(MetricSinkConnection {
238                label: metric_sink
239                    .label
240                    .unwrap_or_else(|| self.sink_id.to_string()),
241            }),
242            with_snapshot: true,
243            up_to: Antichain::new(),
244            non_null_assertions: Vec::new(),
245            refresh_schedule: None,
246        };
247        df_desc.export_sink(self.sink_id, sink_description);
248
249        // Prepare expressions in the assembled dataflow.
250        let style = ExprPrepMaintained;
251        df_desc.visit_children(
252            |r| style.prep_relation_expr(r),
253            |s| style.prep_scalar_expr(s),
254        )?;
255
256        // Construct TransformCtx for global optimization.
257        let mut transform_ctx = TransformCtx::global(
258            &df_builder,
259            &mz_transform::EmptyStatisticsOracle,
260            &self.config.features,
261            &self.typecheck_ctx,
262            &mut df_meta,
263            Some(&mut self.metrics),
264        );
265        // Run global optimization.
266        mz_transform::optimize_dataflow(&mut df_desc, &mut transform_ctx, false)?;
267
268        self.duration += time.elapsed();
269
270        Ok(GlobalMirPlan { df_desc, df_meta })
271    }
272}
273
274impl Optimize<GlobalMirPlan> for Optimizer {
275    type To = GlobalLirPlan;
276
277    fn optimize(&mut self, plan: GlobalMirPlan) -> Result<Self::To, OptimizerError> {
278        let time = Instant::now();
279
280        let GlobalMirPlan {
281            mut df_desc,
282            df_meta,
283        } = plan;
284
285        // Ensure all expressions are normalized before finalizing.
286        for build in df_desc.objects_to_build.iter_mut() {
287            normalize_lets(&mut build.plan.0, &self.config.features)?
288        }
289
290        // Finalize the dataflow: MIR ⇒ LIR lowering and LIR ⇒ LIR transforms.
291        let df_desc = LirRelationExpr::finalize_dataflow(
292            df_desc,
293            &self.config.features,
294            Some(self.metrics.lowering()),
295        )?;
296
297        // Trace the pipeline output under `optimize`.
298        trace_plan(&df_desc);
299
300        self.duration += time.elapsed();
301        self.metrics
302            .observe_e2e_optimization_time("metric_sink", self.duration);
303
304        Ok(GlobalLirPlan { df_desc, df_meta })
305    }
306}
307
308impl GlobalLirPlan {
309    /// Unwraps the parts of the final result of the optimization pipeline.
310    pub fn unapply(self) -> (LirDataflowDescription, DataflowMetainfo) {
311        (self.df_desc, self.df_meta)
312    }
313}
314
315/// Extends the metric sink's source expression with the row-wise shaping the operator otherwise
316/// has to do in Rust: prepends the configured `prefix` to `metric_name` to form the published name,
317/// coalesces `labels`/`help` to their identity element, and adds two columns the operator reads
318/// instead of parsing strings on its hot path:
319///
320/// * `metric_kind` (`Int32`, nullable): `0` for `gauge`, `1` for `counter`, `NULL` for any other
321///   `metric_type`.
322/// * `name_valid` (`Bool`, nullable): whether the published name (`prefix + metric_name`) matches
323///   the Prometheus metric-name grammar (see `METRIC_NAME_PATTERN`). The operator treats a `NULL`
324///   the same as `false`. Validating the published name, not the bare `metric_name`, is what lets a
325///   row name start with a digit: the prefix supplies the valid leading character.
326///
327/// No row is dropped or filtered here: the operator still needs every row, including the ones
328/// this marks invalid, to count `skipped`/`null_values`. Only the pure per-row shaping moves to
329/// MIR. Dedup, collision detection, and family-conflict counting stay in the operator, because
330/// they need cross-row state (the frontier-gated fold) that a `Map` can't express.
331///
332/// TODO: A full move would also express the dedup/collision/family-conflict logic in MIR (e.g.
333/// via `Reduce` + `FirstValue`), collapsing the operator to a plain fold over the live set. That
334/// full move is deferred: the tiebreak fidelity that logic needs is easier to keep correct
335/// hand-written and unit-tested for now.
336fn shape_metric_sink_source(
337    source: MirRelationExpr,
338    source_desc: &RelationDesc,
339    prefix: &str,
340) -> (MirRelationExpr, RelationDesc) {
341    // Precondition: `source_desc` describes `source` and exposes the canonical metric-sink columns
342    // (`metric_name`, `metric_type`, `labels`, `value`, `help`).
343    // `mz_sql::plan::validate_metric_sink_desc` enforces this for both `CREATE METRIC SINK` and
344    // the coordinator-installed curated sinks, so a missing column here is a caller bug.
345    let get_idx = |name: &str| {
346        source_desc
347            .get_by_name(&ColumnName::from(name))
348            .expect("metric-sink source relation must expose the canonical columns")
349    };
350    let (metric_name_idx, metric_name_ct) = get_idx("metric_name");
351    let (metric_type_idx, metric_type_ct) = get_idx("metric_type");
352    let (labels_idx, labels_ct) = get_idx("labels");
353    let (value_idx, value_ct) = get_idx("value");
354    let (help_idx, help_ct) = get_idx("help");
355
356    let arity = source_desc.typ().columns().len();
357    // The mapped columns are appended at `arity + N` and the `Project` indexes into `source` by
358    // position, so `source` must have exactly the arity `source_desc` describes. Guaranteed by the
359    // callers (a trivial finishing over the planned query, or a direct `Get` of the source), but a
360    // mismatch would silently read the wrong columns, so assert it here.
361    mz_ore::soft_assert_eq_or_log!(source.arity(), arity);
362    let labels_repr_type = ReprScalarType::from(&labels_ct.scalar_type);
363
364    let empty_map_row = {
365        let mut row = Row::default();
366        row.packer().push_dict_with(|_| {});
367        row
368    };
369    let labels_coalesced = MirScalarExpr::call_variadic(
370        Coalesce,
371        vec![
372            MirScalarExpr::column(labels_idx),
373            MirScalarExpr::literal_from_single_element_row(empty_map_row, labels_repr_type),
374        ],
375    );
376    let help_coalesced = MirScalarExpr::call_variadic(
377        Coalesce,
378        vec![
379            MirScalarExpr::column(help_idx),
380            MirScalarExpr::literal_ok(Datum::String(""), ReprScalarType::String),
381        ],
382    );
383
384    let metric_type_literal = |s: &'static str| {
385        MirScalarExpr::column(metric_type_idx).call_binary(
386            MirScalarExpr::literal_ok(Datum::String(s), ReprScalarType::String),
387            func::Eq,
388        )
389    };
390    let metric_kind = metric_type_literal("gauge").if_then_else(
391        MirScalarExpr::literal_ok(Datum::Int32(0), ReprScalarType::Int32),
392        metric_type_literal("counter").if_then_else(
393            MirScalarExpr::literal_ok(Datum::Int32(1), ReprScalarType::Int32),
394            MirScalarExpr::literal_null(ReprScalarType::Int32),
395        ),
396    );
397
398    // The published name is `prefix + metric_name`. The prefix is validated to start with the
399    // reserved marker (see `validate_metric_sink_prefix`, run at plan time for a user sink and at
400    // install time for a curated one), so every published family lands in the
401    // `mz_metric_sink_` lane nothing else in the process registry writes. `TextConcat` (the `||`
402    // operator) propagates nulls, so a null `metric_name` stays null and is skipped, never
403    // published as the bare prefix.
404    let prefixed_name = MirScalarExpr::literal_ok(Datum::String(prefix), ReprScalarType::String)
405        .call_binary(
406            MirScalarExpr::column(metric_name_idx),
407            func::TextConcatBinary,
408        );
409
410    // The regexp requires at least one leading character, so an empty published name fails it
411    // without a separate `!= ""` check. `is_null().not()` keeps a null name concretely `false`,
412    // not `NULL`. The shaping map appends `prefixed_name` at `arity + 2`, and a map scalar may
413    // reference earlier-appended columns, so this reads that column instead of rebuilding the
414    // concat.
415    let published_name = MirScalarExpr::column(arity + 2);
416    let name_valid = published_name
417        .clone()
418        .call_is_null()
419        .not()
420        .and(published_name.call_binary(
421            MirScalarExpr::literal_ok(Datum::String(METRIC_NAME_PATTERN), ReprScalarType::String),
422            func::IsRegexpMatchCaseSensitive,
423        ));
424
425    let shaped_expr = source
426        .map(vec![
427            labels_coalesced,
428            help_coalesced,
429            prefixed_name,
430            metric_kind,
431            name_valid,
432        ])
433        .project(vec![
434            arity + 2, // prefixed metric_name
435            metric_type_idx,
436            arity, // coalesced labels
437            value_idx,
438            arity + 1, // coalesced help
439            arity + 3, // metric_kind
440            arity + 4, // name_valid
441        ]);
442
443    let mut labels_shaped_ct = labels_ct.clone();
444    labels_shaped_ct.nullable = false;
445    let mut help_shaped_ct = help_ct.clone();
446    help_shaped_ct.nullable = false;
447    let shaped_desc = RelationDesc::from_names_and_types([
448        ("metric_name", metric_name_ct.clone()),
449        ("metric_type", metric_type_ct.clone()),
450        ("labels", labels_shaped_ct),
451        ("value", value_ct.clone()),
452        ("help", help_shaped_ct),
453        ("metric_kind", SqlScalarType::Int32.nullable(true)),
454        ("name_valid", SqlScalarType::Bool.nullable(true)),
455    ]);
456
457    (shaped_expr, shaped_desc)
458}
459
460#[cfg(test)]
461mod tests {
462    use std::collections::BTreeMap;
463
464    use mz_catalog::memory::objects::{CatalogEntry, CatalogItem, Table, TableDataSource};
465    use mz_controller_types::ClusterId;
466    use mz_expr::Eval;
467    use mz_ore::metrics::MetricsRegistry;
468    use mz_repr::adt::mz_acl_item::PrivilegeMap;
469    use mz_repr::role_id::RoleId;
470    use mz_repr::{
471        CatalogItemId, RelationVersion, RelationVersionSelector, RowArena, SqlColumnType,
472        VersionedRelationDesc,
473    };
474    use mz_sql::names::{
475        FullItemName, ItemQualifiers, QualifiedItemName, RawDatabaseSpecifier,
476        ResolvedDatabaseSpecifier, ResolvedIds, SchemaId, SchemaSpecifier,
477    };
478    use mz_sql::session::vars::SystemVars;
479
480    use super::*;
481
482    /// The canonical metric-sink source shape, with `labels`/`help` nullable so the shaping's
483    /// coalesce is observable and an extra trailing column so column resolution is exercised by
484    /// name, not position.
485    fn source_desc() -> RelationDesc {
486        RelationDesc::builder()
487            .with_column("metric_name", SqlScalarType::String.nullable(true))
488            .with_column("metric_type", SqlScalarType::String.nullable(false))
489            .with_column(
490                "labels",
491                SqlScalarType::Map {
492                    value_type: Box::new(SqlScalarType::String),
493                    custom_id: None,
494                }
495                .nullable(true),
496            )
497            .with_column("value", SqlScalarType::Float64.nullable(true))
498            .with_column("help", SqlScalarType::String.nullable(true))
499            .with_column("extra", SqlScalarType::String.nullable(true))
500            .finish()
501    }
502
503    /// A bare `Get` of `TABLE_GID`, the source expression the `MetricSinkFrom::Id` path shapes.
504    fn source_get(desc: &RelationDesc) -> MirRelationExpr {
505        MirRelationExpr::global_get(TABLE_GID, ReprRelationType::from(desc.typ()))
506    }
507
508    #[mz_ore::test]
509    fn shaped_desc_column_contract() {
510        let (_expr, desc) =
511            shape_metric_sink_source(source_get(&source_desc()), &source_desc(), "app_");
512
513        let cols: Vec<(String, SqlColumnType)> = desc
514            .iter()
515            .map(|(name, ty)| (name.as_str().to_string(), ty.clone()))
516            .collect();
517
518        // Exactly the seven canonical columns, in order. The trailing `extra` source column is
519        // projected away.
520        let names: Vec<&str> = cols.iter().map(|(n, _)| n.as_str()).collect();
521        assert_eq!(
522            names,
523            vec![
524                "metric_name",
525                "metric_type",
526                "labels",
527                "value",
528                "help",
529                "metric_kind",
530                "name_valid",
531            ]
532        );
533
534        let by_name = |name: &str| {
535            cols.iter()
536                .find(|(n, _)| n == name)
537                .expect("column present in shaped desc")
538                .1
539                .clone()
540        };
541
542        // `labels`/`help` are coalesced to their identity element, so they are non-null.
543        assert!(!by_name("labels").nullable);
544        assert!(!by_name("help").nullable);
545
546        // `metric_name`/`value` stay nullable (no identity element).
547        assert!(by_name("metric_name").nullable);
548        assert!(by_name("value").nullable);
549
550        // The two classification columns the operator reads.
551        assert_eq!(by_name("metric_kind"), SqlScalarType::Int32.nullable(true));
552        assert_eq!(by_name("name_valid"), SqlScalarType::Bool.nullable(true));
553    }
554
555    #[mz_ore::test]
556    fn shaped_expr_projects_seven_columns() {
557        let (expr, _desc) =
558            shape_metric_sink_source(source_get(&source_desc()), &source_desc(), "app_");
559
560        // The shaping is a `Map` of five new columns followed by a `Project` down to the seven
561        // canonical columns.
562        match &expr {
563            MirRelationExpr::Project { outputs, .. } => {
564                assert_eq!(outputs.len(), 7);
565            }
566            other => panic!("expected a Project at the root of the shaped expr, got {other:?}"),
567        }
568    }
569
570    /// The five scalars the shaping `Map` appends, in order:
571    /// `[labels_coalesced, help_coalesced, prefixed_name, metric_kind, name_valid]`.
572    fn shaped_map_scalars(desc: &RelationDesc, prefix: &str) -> Vec<MirScalarExpr> {
573        let (expr, _desc) = shape_metric_sink_source(source_get(desc), desc, prefix);
574        match expr {
575            MirRelationExpr::Project { input, .. } => match *input {
576                MirRelationExpr::Map { scalars, .. } => scalars,
577                other => panic!("expected a Map under the Project, got {other:?}"),
578            },
579            other => panic!("expected a Project at the root, got {other:?}"),
580        }
581    }
582
583    /// Evaluate the shaping `Map`'s appended scalars in order, extending the row with each result.
584    /// A later scalar may reference an earlier appended column (`name_valid` reads the published-name
585    /// column), so they must be evaluated cumulatively rather than in isolation.
586    fn eval_shaped_row<'a>(
587        scalars: &'a [MirScalarExpr],
588        input: &[Datum<'a>],
589        arena: &'a RowArena,
590    ) -> Vec<Datum<'a>> {
591        let mut row = input.to_vec();
592        for scalar in scalars {
593            let datum = scalar.eval(&row, arena).expect("scalar eval succeeds");
594            row.push(datum);
595        }
596        row
597    }
598
599    #[mz_ore::test]
600    fn metric_kind_classifies_type() {
601        let scalars = shaped_map_scalars(&source_desc(), "app_");
602        let metric_kind = &scalars[3];
603        let arena = RowArena::new();
604        // Row layout matches `source_desc`: [metric_name, metric_type, labels, value, help, extra].
605        for (metric_type, expected) in [
606            ("gauge", Datum::Int32(0)),
607            ("counter", Datum::Int32(1)),
608            ("histogram", Datum::Null),
609            ("summary", Datum::Null),
610        ] {
611            let row = [
612                Datum::Null,
613                Datum::String(metric_type),
614                Datum::Null,
615                Datum::Null,
616                Datum::Null,
617                Datum::Null,
618            ];
619            assert_eq!(
620                metric_kind
621                    .eval(&row, &arena)
622                    .expect("metric_kind eval succeeds"),
623                expected,
624                "metric_type = {metric_type}",
625            );
626        }
627    }
628
629    #[mz_ore::test]
630    fn name_valid_matches_prometheus_grammar() {
631        // Validation is on the published name `prefix + metric_name`, not the bare `metric_name`.
632        // The plan-time-validated prefix supplies the valid leading character, so a row name may
633        // start with a digit (or be empty, publishing the bare prefix). A dash is invalid anywhere,
634        // and a null name stays null.
635        let scalars = shaped_map_scalars(&source_desc(), "app_");
636        let arena = RowArena::new();
637        for (metric_name, expected) in [
638            (Datum::String("http_requests_total"), Datum::True),
639            (Datum::String("with:colons_and_1_digit"), Datum::True),
640            (Datum::String("1_leading_digit"), Datum::True),
641            (Datum::String("has-a-dash"), Datum::False),
642            (Datum::String(""), Datum::True),
643            (Datum::Null, Datum::False),
644        ] {
645            let input = [
646                metric_name,
647                Datum::Null,
648                Datum::Null,
649                Datum::Null,
650                Datum::Null,
651                Datum::Null,
652            ];
653            // `name_valid` is the last appended scalar and reads the published-name column, so
654            // evaluate the whole appended row and read its final column.
655            let row = eval_shaped_row(&scalars, &input, &arena);
656            assert_eq!(
657                *row.last().expect("row has appended columns"),
658                expected,
659                "metric_name = {metric_name:?}",
660            );
661        }
662    }
663
664    /// The smallest catalog the optimizer needs: one table, at `TABLE_GID`, exposing the canonical
665    /// metric-sink columns.
666    #[derive(Debug)]
667    struct SingleTableCatalog {
668        entry: CatalogEntry,
669    }
670
671    const TABLE_ITEM_ID: CatalogItemId = CatalogItemId::User(1);
672    const TABLE_GID: GlobalId = GlobalId::User(1);
673    const SINK_GID: GlobalId = GlobalId::User(2);
674
675    impl SingleTableCatalog {
676        fn new() -> Self {
677            let table = Table {
678                create_sql: None,
679                desc: VersionedRelationDesc::new(source_desc()),
680                collections: BTreeMap::from([(RelationVersion::root(), TABLE_GID)]),
681                conn_id: None,
682                resolved_ids: ResolvedIds::empty(),
683                custom_logical_compaction_window: None,
684                is_retained_metrics_object: false,
685                data_source: TableDataSource::TableWrites {
686                    defaults: Vec::new(),
687                },
688            };
689            let entry = CatalogEntry {
690                item: CatalogItem::Table(table),
691                referenced_by: Vec::new(),
692                used_by: Vec::new(),
693                id: TABLE_ITEM_ID,
694                oid: 20_000,
695                name: QualifiedItemName {
696                    qualifiers: ItemQualifiers {
697                        database_spec: ResolvedDatabaseSpecifier::Ambient,
698                        schema_spec: SchemaSpecifier::Id(SchemaId::User(1)),
699                    },
700                    item: "t".to_string(),
701                },
702                owner_id: RoleId::User(1),
703                privileges: PrivilegeMap::default(),
704            };
705            Self { entry }
706        }
707    }
708
709    impl OptimizerCatalog for SingleTableCatalog {
710        fn get_entry(&self, _id: &GlobalId) -> mz_catalog::memory::objects::CatalogCollectionEntry {
711            mz_catalog::memory::objects::CatalogCollectionEntry {
712                entry: self.entry.clone(),
713                version: RelationVersionSelector::Latest,
714            }
715        }
716
717        fn get_entry_by_item_id(&self, _id: &CatalogItemId) -> &CatalogEntry {
718            &self.entry
719        }
720
721        fn resolve_full_name(
722            &self,
723            name: &QualifiedItemName,
724            _conn_id: Option<&mz_adapter_types::connection::ConnectionId>,
725        ) -> FullItemName {
726            FullItemName {
727                database: RawDatabaseSpecifier::Ambient,
728                schema: "public".to_string(),
729                item: name.item.clone(),
730            }
731        }
732
733        fn get_indexes_on(
734            &self,
735            _id: GlobalId,
736            _cluster: ClusterId,
737        ) -> Box<dyn Iterator<Item = (GlobalId, &mz_catalog::memory::objects::Index)> + '_>
738        {
739            Box::new(std::iter::empty())
740        }
741    }
742
743    const VIEW_GID: GlobalId = GlobalId::Transient(1);
744
745    /// Runs the whole pipeline over `from` and returns the assembled dataflow.
746    fn optimize_from(from: MetricSinkFrom, metric_label: Option<String>) -> LirDataflowDescription {
747        let catalog = Arc::new(SingleTableCatalog::new());
748        let cluster_id = ClusterId::user(1).expect("valid cluster id");
749        let compute_instance = ComputeInstanceSnapshot::new_without_collections(cluster_id);
750        let config = OptimizerConfig::from(&SystemVars::default());
751        let metrics = OptimizerMetrics::register_into(&MetricsRegistry::new(), Duration::MAX);
752
753        let mut optimizer = Optimizer::new(
754            catalog,
755            compute_instance,
756            VIEW_GID,
757            SINK_GID,
758            config,
759            metrics,
760        );
761
762        let global_mir_plan = optimizer
763            .optimize(MetricSink::new(
764                "metric-sink-test".to_string(),
765                from,
766                "app_".to_string(),
767                metric_label,
768            ))
769            .expect("MIR optimization succeeds");
770        let global_lir_plan = optimizer
771            .optimize(global_mir_plan)
772            .expect("LIR optimization succeeds");
773        let (df_desc, _df_meta) = global_lir_plan.unapply();
774        df_desc
775    }
776
777    /// Asserts the dataflow exports exactly one `MetricSink` over the shaped view, whose desc
778    /// carries the operator's column contract.
779    fn assert_one_shaped_metric_sink_export(df_desc: &LirDataflowDescription) {
780        assert!(df_desc.index_exports.is_empty());
781        let sink_exports: Vec<_> = df_desc.sink_exports.iter().collect();
782        assert_eq!(sink_exports.len(), 1);
783        let (sink_id, sink_desc) = sink_exports[0];
784        assert_eq!(*sink_id, SINK_GID);
785        assert!(matches!(
786            sink_desc.connection,
787            ComputeSinkConnection::MetricSink(_)
788        ));
789        assert_eq!(sink_desc.from, VIEW_GID);
790        let shaped_names: Vec<&str> = sink_desc
791            .from_desc
792            .iter_names()
793            .map(|n| n.as_str())
794            .collect();
795        assert_eq!(
796            shaped_names,
797            vec![
798                "metric_name",
799                "metric_type",
800                "labels",
801                "value",
802                "help",
803                "metric_kind",
804                "name_valid",
805            ]
806        );
807    }
808
809    /// The `sink` label carried by the export's connection.
810    fn sink_label(df_desc: &LirDataflowDescription) -> &str {
811        match &df_desc
812            .sink_exports
813            .values()
814            .next()
815            .expect("one export")
816            .connection
817        {
818            ComputeSinkConnection::MetricSink(conn) => &conn.label,
819            other => panic!("expected a metric sink connection, got {other:?}"),
820        }
821    }
822
823    #[mz_ore::test]
824    fn optimizer_exports_one_metric_sink() {
825        let df_desc = optimize_from(MetricSinkFrom::Id(TABLE_GID), None);
826        assert_one_shaped_metric_sink_export(&df_desc);
827        // The source collection is imported, not rebuilt: the only build is the shaped view.
828        assert!(df_desc.source_imports.contains_key(&TABLE_GID));
829        let build_ids: Vec<_> = df_desc.objects_to_build.iter().map(|b| b.id).collect();
830        assert_eq!(build_ids, vec![VIEW_GID]);
831    }
832
833    /// The `Query` source path (what a coordinator-installed curated sink takes) assembles the
834    /// same shape, with the query lowered under the shaping instead of a `Get` of a catalog item.
835    #[mz_ore::test]
836    fn optimizer_shapes_a_query_source() {
837        let desc = source_desc();
838        // The simplest query over the canonical columns. Building richer HIR by hand buys nothing:
839        // what is under test is that a query source is lowered and shaped, not the lowering itself.
840        let expr = HirRelationExpr::Get {
841            id: mz_expr::Id::Global(TABLE_GID),
842            typ: desc.typ().clone(),
843        };
844
845        let df_desc = optimize_from(
846            MetricSinkFrom::Query {
847                expr,
848                desc: desc.clone(),
849            },
850            None,
851        );
852        assert_one_shaped_metric_sink_export(&df_desc);
853        // The query's leaf collection is imported by the shaped view's dependency walk.
854        assert!(df_desc.source_imports.contains_key(&TABLE_GID));
855        let build_ids: Vec<_> = df_desc.objects_to_build.iter().map(|b| b.id).collect();
856        assert_eq!(build_ids, vec![VIEW_GID]);
857    }
858
859    /// With no explicit label a sink is tagged by its `GlobalId`, what a user's `CREATE METRIC
860    /// SINK` relies on. An explicit label (a curated sink's stable name) is used verbatim.
861    #[mz_ore::test]
862    fn metric_sink_label_defaults_to_sink_id_else_override() {
863        let df_desc = optimize_from(MetricSinkFrom::Id(TABLE_GID), None);
864        assert_eq!(sink_label(&df_desc), SINK_GID.to_string());
865
866        let df_desc = optimize_from(
867            MetricSinkFrom::Id(TABLE_GID),
868            Some("mz_curated".to_string()),
869        );
870        assert_eq!(sink_label(&df_desc), "mz_curated");
871    }
872}