Skip to main content

mz_sql/plan/statement/
dml.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//! Data manipulation language (DML).
11//!
12//! This module houses the handlers for statements that manipulate data, like
13//! `INSERT`, `SELECT`, `SUBSCRIBE`, and `COPY`.
14
15use std::borrow::Cow;
16use std::collections::{BTreeMap, BTreeSet};
17
18use itertools::Itertools;
19use mz_arrow_util::builder::ArrowBuilder;
20use mz_expr::{ColumnOrder, RowSetFinishing};
21use mz_ore::error::ErrorExt;
22use mz_ore::num::NonNeg;
23use mz_ore::soft_panic_or_log;
24use mz_ore::str::separated;
25use mz_pgcopy::{CopyCsvFormatParams, CopyFormatParams, CopyTextFormatParams};
26use mz_repr::adt::numeric::NumericMaxScale;
27use mz_repr::bytes::ByteSize;
28use mz_repr::explain::{ExplainConfig, ExplainFormat};
29use mz_repr::optimize::OptimizerFeatureOverrides;
30use mz_repr::{CatalogItemId, Datum, RelationDesc, Row, SqlRelationType, SqlScalarType};
31use mz_sql_parser::ast::{
32    CteBlock, ExplainAnalyzeClusterStatement, ExplainAnalyzeComputationProperties,
33    ExplainAnalyzeComputationProperty, ExplainAnalyzeObjectStatement, ExplainAnalyzeProperty,
34    ExplainPlanOption, ExplainPlanOptionName, ExplainPushdownStatement, ExplainSinkSchemaFor,
35    ExplainSinkSchemaStatement, ExplainTimestampStatement, Expr, IfExistsBehavior, OrderByExpr,
36    SetExpr, SubscribeOutput, UnresolvedItemName,
37};
38use mz_sql_parser::ident;
39use mz_storage_types::sinks::{
40    KafkaSinkConnection, KafkaSinkFormat, KafkaSinkFormatType, MAX_S3_SINK_FILE_SIZE,
41    MIN_S3_SINK_FILE_SIZE, S3SinkFormat, StorageSinkConnection,
42};
43
44use crate::ast::display::{AstDisplay, escaped_string_literal};
45use crate::ast::{
46    AstInfo, CopyDirection, CopyOption, CopyOptionName, CopyRelation, CopyStatement, CopyTarget,
47    DeleteStatement, ExplainPlanStatement, ExplainStage, Explainee, Ident, InsertStatement, Query,
48    SelectStatement, SubscribeOption, SubscribeOptionName, SubscribeRelation, SubscribeStatement,
49    UpdateStatement,
50};
51use crate::catalog::CatalogItemType;
52use crate::names::{Aug, ResolvedItemName};
53use crate::normalize;
54use crate::plan::query::{
55    ExprContext, QueryLifetime, negative_offset_error, offset_into_value, plan_as_of_or_up_to,
56    plan_expr,
57};
58use crate::plan::scope::Scope;
59use crate::plan::statement::show::ShowSelect;
60use crate::plan::statement::{StatementContext, StatementDesc, ddl};
61use crate::plan::{
62    self, CopyFromFilter, CopyToPlan, CreateSinkPlan, ExplainPushdownPlan, ExplainSinkSchemaPlan,
63    ExplainTimestampPlan, HirRelationExpr, side_effecting_func, transform_ast,
64};
65use crate::plan::{
66    CopyFormat, CopyFromPlan, ExplainPlanPlan, InsertPlan, MutationKind, Params, Plan, PlanError,
67    QueryContext, ReadThenWritePlan, SelectPlan, SubscribeFrom, SubscribePlan, query,
68};
69use crate::plan::{CopyFromSource, with_options};
70use crate::session::vars::{self, DISALLOW_UNMATERIALIZABLE_FUNCTIONS_AS_OF};
71
72// TODO(benesch): currently, describing a `SELECT` or `INSERT` query
73// plans the whole query to determine its shape and parameter types,
74// and then throws away that plan. If we were smarter, we'd stash that
75// plan somewhere so we don't have to recompute it when the query is
76// executed.
77
78pub fn describe_insert(
79    scx: &StatementContext,
80    InsertStatement {
81        table_name,
82        columns,
83        source,
84        returning,
85    }: InsertStatement<Aug>,
86) -> Result<StatementDesc, PlanError> {
87    let (_, _, returning) = query::plan_insert_query(scx, table_name, columns, source, returning)?;
88    let desc = if returning.expr.is_empty() {
89        None
90    } else {
91        Some(returning.desc)
92    };
93    Ok(StatementDesc::new(desc))
94}
95
96pub fn plan_insert(
97    scx: &StatementContext,
98    InsertStatement {
99        table_name,
100        columns,
101        source,
102        returning,
103    }: InsertStatement<Aug>,
104    params: &Params,
105) -> Result<Plan, PlanError> {
106    let (id, mut expr, returning) =
107        query::plan_insert_query(scx, table_name, columns, source, returning)?;
108    expr.bind_parameters_and_simplify_offset(scx, QueryLifetime::OneShot, params)?;
109    let returning = returning
110        .expr
111        .into_iter()
112        .map(|mut expr| {
113            expr.bind_parameters_and_simplify_offset(scx, QueryLifetime::OneShot, params)?;
114            expr.lower_uncorrelated(scx.catalog.system_vars())
115        })
116        .collect::<Result<Vec<_>, _>>()?;
117
118    Ok(Plan::Insert(InsertPlan {
119        id,
120        values: expr,
121        returning,
122    }))
123}
124
125pub fn describe_delete(
126    scx: &StatementContext,
127    stmt: DeleteStatement<Aug>,
128) -> Result<StatementDesc, PlanError> {
129    query::plan_delete_query(scx, stmt)?;
130    Ok(StatementDesc::new(None))
131}
132
133pub fn plan_delete(
134    scx: &StatementContext,
135    stmt: DeleteStatement<Aug>,
136    params: &Params,
137) -> Result<Plan, PlanError> {
138    let rtw_plan = query::plan_delete_query(scx, stmt)?;
139    plan_read_then_write(scx, MutationKind::Delete, params, rtw_plan)
140}
141
142pub fn describe_update(
143    scx: &StatementContext,
144    stmt: UpdateStatement<Aug>,
145) -> Result<StatementDesc, PlanError> {
146    query::plan_update_query(scx, stmt)?;
147    Ok(StatementDesc::new(None))
148}
149
150pub fn plan_update(
151    scx: &StatementContext,
152    stmt: UpdateStatement<Aug>,
153    params: &Params,
154) -> Result<Plan, PlanError> {
155    let rtw_plan = query::plan_update_query(scx, stmt)?;
156    plan_read_then_write(scx, MutationKind::Update, params, rtw_plan)
157}
158
159pub fn plan_read_then_write(
160    scx: &StatementContext,
161    kind: MutationKind,
162    params: &Params,
163    query::ReadThenWritePlan {
164        id,
165        mut selection,
166        finishing,
167        assignments,
168    }: query::ReadThenWritePlan,
169) -> Result<Plan, PlanError> {
170    selection.bind_parameters_and_simplify_offset(scx, QueryLifetime::OneShot, params)?;
171    let mut assignments_outer = BTreeMap::new();
172    for (idx, mut set) in assignments {
173        set.bind_parameters_and_simplify_offset(scx, QueryLifetime::OneShot, params)?;
174        let set = set.lower_uncorrelated(scx.catalog.system_vars())?;
175        assignments_outer.insert(idx, set);
176    }
177
178    Ok(Plan::ReadThenWrite(ReadThenWritePlan {
179        id,
180        selection,
181        finishing,
182        assignments: assignments_outer,
183        kind,
184        returning: Vec::new(),
185    }))
186}
187
188pub fn describe_select(
189    scx: &StatementContext,
190    stmt: SelectStatement<Aug>,
191) -> Result<StatementDesc, PlanError> {
192    if let Some(desc) = side_effecting_func::describe_select_if_side_effecting(scx, &stmt)? {
193        return Ok(StatementDesc::new(Some(desc)));
194    }
195
196    let query::PlannedRootQuery { desc, .. } =
197        query::plan_root_query(scx, stmt.query, QueryLifetime::OneShot)?;
198    Ok(StatementDesc::new(Some(desc)))
199}
200
201pub fn plan_select(
202    scx: &StatementContext,
203    select: SelectStatement<Aug>,
204    params: &Params,
205    copy_to: Option<CopyFormat>,
206) -> Result<Plan, PlanError> {
207    if let Some(f) = side_effecting_func::plan_select_if_side_effecting(scx, &select, params)? {
208        return Ok(Plan::SideEffectingFunc(f));
209    }
210
211    let (plan, _desc) = plan_select_inner(scx, select, params, copy_to)?;
212    Ok(Plan::Select(plan))
213}
214
215fn plan_select_inner(
216    scx: &StatementContext,
217    select: SelectStatement<Aug>,
218    params: &Params,
219    copy_to: Option<CopyFormat>,
220) -> Result<(SelectPlan, RelationDesc), PlanError> {
221    let when = query::plan_as_of(scx, select.as_of.clone())?;
222    let lifetime = QueryLifetime::OneShot;
223    let query::PlannedRootQuery {
224        mut expr,
225        desc,
226        finishing,
227        scope: _,
228    } = query::plan_root_query(scx, select.query.clone(), lifetime)?;
229    expr.bind_parameters_and_simplify_offset(scx, lifetime, params)?;
230
231    // We need to concretize the `limit` and `offset` of the RowSetFinishing, so that we go from
232    // `RowSetFinishing<HirScalarExpr, HirScalarExpr>` to `RowSetFinishing`.
233    // This involves binding parameters and evaluating each expression to a number.
234    // (This should be possible even for `limit` here, because we are at the top level of a SELECT,
235    // so this `limit` has to be a constant.)
236    let limit = match finishing.limit {
237        None => None,
238        Some(mut limit) => {
239            limit.bind_parameters_and_simplify_offset(scx, lifetime, params)?;
240            // Evaluate the expression to a literal instead of matching on a bare literal
241            // node. Parameter binding can leave the bound value wrapped in casts, e.g.
242            // `integer_to_bigint($1)`.
243            match limit
244                .try_into_nullable_literal_int64()
245                .map_err(|err| PlanError::InvalidLimit(err.to_string_with_causes()))?
246            {
247                None => None,
248                Some(v) if v >= 0 => NonNeg::<i64>::try_from(v).ok(),
249                Some(_) => sql_bail!("LIMIT must not be negative"),
250            }
251        }
252    };
253    let offset = {
254        let mut offset = finishing.offset.clone();
255        offset.bind_parameters_and_simplify_offset(scx, lifetime, params)?;
256        let offset = offset_into_value(offset.take())?;
257        offset.try_into().map_err(|_| {
258            // We already checked in bind_parameters_and_simplify_offset / offset_into_value.
259            soft_panic_or_log!("unexpectedly negative OFFSET");
260            negative_offset_error(offset)
261        })?
262    };
263
264    // Unmaterializable functions are evaluated based on various information (e.g., the catalog)
265    // during sequencing, without taking AS OF into account. This would be hard to fix, so for now
266    // we just disallow AS OF when there is an unmaterializable function in a query (except mz_now).
267    if scx.is_feature_flag_enabled(&DISALLOW_UNMATERIALIZABLE_FUNCTIONS_AS_OF)
268        && select.as_of.is_some()
269        && expr.contains_unmaterializable_except_temporal()
270    {
271        bail_unsupported!("unmaterializable function (except `mz_now`) in an AS OF query");
272    }
273
274    let plan = SelectPlan {
275        source: expr,
276        when,
277        finishing: RowSetFinishing {
278            limit,
279            offset,
280            project: finishing.project,
281            order_by: finishing.order_by,
282        },
283        copy_to,
284        select: Some(Box::new(select)),
285    };
286
287    Ok((plan, desc))
288}
289
290pub fn describe_explain_plan(
291    scx: &StatementContext,
292    explain: ExplainPlanStatement<Aug>,
293) -> Result<StatementDesc, PlanError> {
294    let mut relation_desc = RelationDesc::builder();
295
296    match explain.stage() {
297        ExplainStage::RawPlan => {
298            let name = "Raw Plan";
299            relation_desc = relation_desc.with_column(name, SqlScalarType::String.nullable(false));
300        }
301        ExplainStage::DecorrelatedPlan => {
302            let name = "Decorrelated Plan";
303            relation_desc = relation_desc.with_column(name, SqlScalarType::String.nullable(false));
304        }
305        ExplainStage::LocalPlan => {
306            let name = "Locally Optimized Plan";
307            relation_desc = relation_desc.with_column(name, SqlScalarType::String.nullable(false));
308        }
309        ExplainStage::GlobalPlan => {
310            let name = "Optimized Plan";
311            relation_desc = relation_desc.with_column(name, SqlScalarType::String.nullable(false));
312        }
313        ExplainStage::PhysicalPlan => {
314            let name = "Physical Plan";
315            relation_desc = relation_desc.with_column(name, SqlScalarType::String.nullable(false));
316        }
317        ExplainStage::Trace => {
318            relation_desc = relation_desc
319                .with_column("Time", SqlScalarType::UInt64.nullable(false))
320                .with_column("Path", SqlScalarType::String.nullable(false))
321                .with_column("Plan", SqlScalarType::String.nullable(false));
322        }
323        ExplainStage::PlanInsights => {
324            let name = "Plan Insights";
325            relation_desc = relation_desc.with_column(name, SqlScalarType::String.nullable(false));
326        }
327    };
328    let relation_desc = relation_desc.finish();
329
330    Ok(
331        StatementDesc::new(Some(relation_desc)).with_params(match explain.explainee {
332            Explainee::Select(select, _) => describe_select(scx, *select)?.param_types,
333            _ => vec![],
334        }),
335    )
336}
337
338pub fn describe_explain_pushdown(
339    scx: &StatementContext,
340    statement: ExplainPushdownStatement<Aug>,
341) -> Result<StatementDesc, PlanError> {
342    let relation_desc = RelationDesc::builder()
343        .with_column("Source", SqlScalarType::String.nullable(false))
344        .with_column("Total Bytes", SqlScalarType::UInt64.nullable(false))
345        .with_column("Selected Bytes", SqlScalarType::UInt64.nullable(false))
346        .with_column("Total Parts", SqlScalarType::UInt64.nullable(false))
347        .with_column("Selected Parts", SqlScalarType::UInt64.nullable(false))
348        .finish();
349
350    Ok(
351        StatementDesc::new(Some(relation_desc)).with_params(match statement.explainee {
352            Explainee::Select(select, _) => describe_select(scx, *select)?.param_types,
353            _ => vec![],
354        }),
355    )
356}
357
358pub fn describe_explain_analyze_object(
359    _scx: &StatementContext,
360    statement: ExplainAnalyzeObjectStatement<Aug>,
361) -> Result<StatementDesc, PlanError> {
362    if statement.as_sql {
363        let relation_desc = RelationDesc::builder()
364            .with_column("SQL", SqlScalarType::String.nullable(false))
365            .finish();
366        return Ok(StatementDesc::new(Some(relation_desc)));
367    }
368
369    match statement.properties {
370        ExplainAnalyzeProperty::Computation(ExplainAnalyzeComputationProperties {
371            properties,
372            skew,
373        }) => {
374            let mut relation_desc = RelationDesc::builder()
375                .with_column("operator", SqlScalarType::String.nullable(false));
376
377            if skew {
378                relation_desc =
379                    relation_desc.with_column("worker_id", SqlScalarType::UInt64.nullable(true));
380            }
381
382            let mut seen_properties = BTreeSet::new();
383            for property in properties {
384                // handle each property only once (belt and suspenders)
385                if !seen_properties.insert(property) {
386                    continue;
387                }
388
389                match property {
390                    ExplainAnalyzeComputationProperty::Memory if skew => {
391                        let numeric = SqlScalarType::Numeric { max_scale: None }.nullable(true);
392                        relation_desc = relation_desc
393                            .with_column("memory_ratio", numeric.clone())
394                            .with_column("worker_memory", SqlScalarType::String.nullable(true))
395                            .with_column("avg_memory", SqlScalarType::String.nullable(true))
396                            .with_column("total_memory", SqlScalarType::String.nullable(true))
397                            .with_column("records_ratio", numeric.clone())
398                            .with_column("worker_records", numeric.clone())
399                            .with_column("avg_records", numeric.clone())
400                            .with_column("total_records", numeric);
401                    }
402                    ExplainAnalyzeComputationProperty::Memory => {
403                        relation_desc = relation_desc
404                            .with_column("total_memory", SqlScalarType::String.nullable(true))
405                            .with_column(
406                                "total_records",
407                                SqlScalarType::Numeric { max_scale: None }.nullable(true),
408                            );
409                    }
410                    ExplainAnalyzeComputationProperty::Cpu => {
411                        if skew {
412                            relation_desc = relation_desc
413                                .with_column(
414                                    "cpu_ratio",
415                                    SqlScalarType::Numeric { max_scale: None }.nullable(true),
416                                )
417                                .with_column(
418                                    "worker_elapsed",
419                                    SqlScalarType::Interval.nullable(true),
420                                )
421                                .with_column("avg_elapsed", SqlScalarType::Interval.nullable(true));
422                        }
423                        relation_desc = relation_desc
424                            .with_column("total_elapsed", SqlScalarType::Interval.nullable(true));
425                    }
426                }
427            }
428
429            let relation_desc = relation_desc.finish();
430            Ok(StatementDesc::new(Some(relation_desc)))
431        }
432        ExplainAnalyzeProperty::Hints => {
433            let relation_desc = RelationDesc::builder()
434                .with_column("operator", SqlScalarType::String.nullable(true))
435                .with_column("levels", SqlScalarType::Int64.nullable(true))
436                .with_column("to_cut", SqlScalarType::Int64.nullable(true))
437                .with_column("hint", SqlScalarType::Float64.nullable(true))
438                .with_column("savings", SqlScalarType::String.nullable(true))
439                .finish();
440            Ok(StatementDesc::new(Some(relation_desc)))
441        }
442    }
443}
444
445pub fn describe_explain_analyze_cluster(
446    _scx: &StatementContext,
447    statement: ExplainAnalyzeClusterStatement,
448) -> Result<StatementDesc, PlanError> {
449    if statement.as_sql {
450        let relation_desc = RelationDesc::builder()
451            .with_column("SQL", SqlScalarType::String.nullable(false))
452            .finish();
453        return Ok(StatementDesc::new(Some(relation_desc)));
454    }
455
456    let ExplainAnalyzeComputationProperties { properties, skew } = statement.properties;
457
458    let mut relation_desc = RelationDesc::builder()
459        .with_column("object", SqlScalarType::String.nullable(false))
460        .with_column("global_id", SqlScalarType::String.nullable(false));
461
462    if skew {
463        relation_desc =
464            relation_desc.with_column("worker_id", SqlScalarType::UInt64.nullable(true));
465    }
466
467    let mut seen_properties = BTreeSet::new();
468    for property in properties {
469        // handle each property only once (belt and suspenders)
470        if !seen_properties.insert(property) {
471            continue;
472        }
473
474        match property {
475            ExplainAnalyzeComputationProperty::Memory if skew => {
476                let numeric = SqlScalarType::Numeric { max_scale: None }.nullable(true);
477                relation_desc = relation_desc
478                    .with_column("max_operator_memory_ratio", numeric.clone())
479                    .with_column("worker_memory", SqlScalarType::String.nullable(true))
480                    .with_column("avg_memory", SqlScalarType::String.nullable(true))
481                    .with_column("total_memory", SqlScalarType::String.nullable(true))
482                    .with_column("max_operator_records_ratio", numeric.clone())
483                    .with_column("worker_records", numeric.clone())
484                    .with_column("avg_records", numeric.clone())
485                    .with_column("total_records", numeric);
486            }
487            ExplainAnalyzeComputationProperty::Memory => {
488                relation_desc = relation_desc
489                    .with_column("total_memory", SqlScalarType::String.nullable(true))
490                    .with_column(
491                        "total_records",
492                        SqlScalarType::Numeric { max_scale: None }.nullable(true),
493                    );
494            }
495            ExplainAnalyzeComputationProperty::Cpu if skew => {
496                relation_desc = relation_desc
497                    .with_column(
498                        "max_operator_cpu_ratio",
499                        SqlScalarType::Numeric { max_scale: None }.nullable(true),
500                    )
501                    .with_column("worker_elapsed", SqlScalarType::Interval.nullable(true))
502                    .with_column("avg_elapsed", SqlScalarType::Interval.nullable(true))
503                    .with_column("total_elapsed", SqlScalarType::Interval.nullable(true));
504            }
505            ExplainAnalyzeComputationProperty::Cpu => {
506                relation_desc = relation_desc
507                    .with_column("total_elapsed", SqlScalarType::Interval.nullable(true));
508            }
509        }
510    }
511
512    Ok(StatementDesc::new(Some(relation_desc.finish())))
513}
514
515pub fn describe_explain_timestamp(
516    scx: &StatementContext,
517    ExplainTimestampStatement { select, .. }: ExplainTimestampStatement<Aug>,
518) -> Result<StatementDesc, PlanError> {
519    let relation_desc = RelationDesc::builder()
520        .with_column("Timestamp", SqlScalarType::String.nullable(false))
521        .finish();
522
523    Ok(StatementDesc::new(Some(relation_desc))
524        .with_params(describe_select(scx, select)?.param_types))
525}
526
527pub fn describe_explain_schema(
528    _: &StatementContext,
529    ExplainSinkSchemaStatement { .. }: ExplainSinkSchemaStatement<Aug>,
530) -> Result<StatementDesc, PlanError> {
531    let relation_desc = RelationDesc::builder()
532        .with_column("Schema", SqlScalarType::String.nullable(false))
533        .finish();
534    Ok(StatementDesc::new(Some(relation_desc)))
535}
536
537// Currently, there are two reasons for why a flag should be `Option<bool>` instead of simply
538// `bool`:
539// - When it's an override of a global feature flag, for example optimizer feature flags. In this
540//   case, we need not just false and true, but also None to say "take the value of the global
541//   flag".
542// - When it's an override of whether SOFT_ASSERTIONS are enabled. For example, when `Arity` is not
543//   explicitly given in the EXPLAIN command, then we'd like staging and prod to default to true,
544//   but otherwise we'd like to default to false.
545generate_extracted_config!(
546    ExplainPlanOption,
547    (Arity, Option<bool>, Default(None)),
548    (Cardinality, bool, Default(false)),
549    (ColumnNames, bool, Default(false)),
550    (FilterPushdown, Option<bool>, Default(None)),
551    (HumanizedExpressions, Option<bool>, Default(None)),
552    (JoinImplementations, bool, Default(false)),
553    (Keys, bool, Default(false)),
554    (LinearChains, bool, Default(false)),
555    (NoFastPath, bool, Default(false)),
556    (NonNegative, bool, Default(false)),
557    (NoNotices, bool, Default(false)),
558    (NodeIdentifiers, bool, Default(false)),
559    (Raw, bool, Default(false)),
560    (RawPlans, bool, Default(false)),
561    (RawSyntax, bool, Default(false)),
562    (Redacted, bool, Default(false)),
563    (SubtreeSize, bool, Default(false)),
564    (Timing, bool, Default(false)),
565    (Types, bool, Default(false)),
566    (Equivalences, bool, Default(false)),
567    (ReoptimizeImportedViews, Option<bool>, Default(None)),
568    (EnableNewOuterJoinLowering, Option<bool>, Default(None)),
569    (EnableEagerDeltaJoins, Option<bool>, Default(None)),
570    (EnableVariadicLeftJoinLowering, Option<bool>, Default(None)),
571    (EnableLetrecFixpointAnalysis, Option<bool>, Default(None)),
572    (EnableJoinPrioritizeArranged, Option<bool>, Default(None)),
573    (
574        EnableProjectionPushdownAfterRelationCse,
575        Option<bool>,
576        Default(None)
577    ),
578    (
579        EnableFixedCorrelatedCteLowering,
580        Option<bool>,
581        Default(None)
582    )
583);
584
585impl TryFrom<ExplainPlanOptionExtracted> for ExplainConfig {
586    type Error = PlanError;
587
588    fn try_from(mut v: ExplainPlanOptionExtracted) -> Result<Self, Self::Error> {
589        // If `WITH(raw)` is specified, ensure that the config will be as
590        // representative for the original plan as possible.
591        if v.raw {
592            v.raw_plans = true;
593            v.raw_syntax = true;
594        }
595
596        // Certain config should default to be enabled in release builds running on
597        // staging or prod (where SOFT_ASSERTIONS are turned off).
598        let enable_on_prod = !mz_ore::assert::soft_assertions_enabled();
599
600        Ok(ExplainConfig {
601            arity: v.arity.unwrap_or(enable_on_prod),
602            cardinality: v.cardinality,
603            column_names: v.column_names,
604            filter_pushdown: v.filter_pushdown.unwrap_or(enable_on_prod),
605            humanized_exprs: !v.raw_plans && (v.humanized_expressions.unwrap_or(enable_on_prod)),
606            join_impls: v.join_implementations,
607            keys: v.keys,
608            linear_chains: !v.raw_plans && v.linear_chains,
609            no_fast_path: v.no_fast_path,
610            no_notices: v.no_notices,
611            node_ids: v.node_identifiers,
612            non_negative: v.non_negative,
613            raw_plans: v.raw_plans,
614            raw_syntax: v.raw_syntax,
615            verbose_syntax: false,
616            redacted: v.redacted,
617            subtree_size: v.subtree_size,
618            equivalences: v.equivalences,
619            timing: v.timing,
620            types: v.types,
621            // The ones that are initialized with `Default::default()` are not wired up to EXPLAIN.
622            features: OptimizerFeatureOverrides {
623                enable_eager_delta_joins: v.enable_eager_delta_joins,
624                enable_new_outer_join_lowering: v.enable_new_outer_join_lowering,
625                enable_variadic_left_join_lowering: v.enable_variadic_left_join_lowering,
626                enable_letrec_fixpoint_analysis: v.enable_letrec_fixpoint_analysis,
627                enable_reduce_mfp_fusion: Default::default(),
628                enable_cardinality_estimates: Default::default(),
629                persist_fast_path_limit: Default::default(),
630                reoptimize_imported_views: v.reoptimize_imported_views,
631                enable_join_prioritize_arranged: v.enable_join_prioritize_arranged,
632                enable_projection_pushdown_after_relation_cse: v
633                    .enable_projection_pushdown_after_relation_cse,
634                enable_less_reduce_in_eqprop: Default::default(),
635                enable_dequadratic_eqprop_map: Default::default(),
636                enable_eq_classes_withholding_errors: Default::default(),
637                enable_fast_path_plan_insights: Default::default(),
638                enable_cast_elimination: Default::default(),
639                enable_case_literal_transform: Default::default(),
640                enable_simplify_quantified_comparisons: Default::default(),
641                enable_simplify_from_less_existence: Default::default(),
642                enable_coalesce_case_transform: Default::default(),
643                enable_will_distinct_propagation: Default::default(),
644                enable_fixed_correlated_cte_lowering: v.enable_fixed_correlated_cte_lowering,
645            },
646        })
647    }
648}
649
650fn plan_explainee(
651    scx: &StatementContext,
652    explainee: Explainee<Aug>,
653    params: &Params,
654) -> Result<plan::Explainee, PlanError> {
655    use crate::plan::ExplaineeStatement;
656
657    let is_replan = matches!(
658        explainee,
659        Explainee::ReplanView(_) | Explainee::ReplanMaterializedView(_) | Explainee::ReplanIndex(_)
660    );
661
662    let explainee = match explainee {
663        Explainee::View(name) | Explainee::ReplanView(name) => {
664            let item = scx.get_item_by_resolved_name(&name)?;
665            let item_type = item.item_type();
666            if item_type != CatalogItemType::View {
667                sql_bail!("Expected {name} to be a view, not a {item_type}");
668            }
669            match is_replan {
670                true => crate::plan::Explainee::ReplanView(item.id()),
671                false => crate::plan::Explainee::View(item.id()),
672            }
673        }
674        Explainee::MaterializedView(name) | Explainee::ReplanMaterializedView(name) => {
675            let item = scx.get_item_by_resolved_name(&name)?;
676            let item_type = item.item_type();
677            if item_type != CatalogItemType::MaterializedView {
678                sql_bail!("Expected {name} to be a materialized view, not a {item_type}");
679            }
680            match is_replan {
681                true => crate::plan::Explainee::ReplanMaterializedView(item.id()),
682                false => crate::plan::Explainee::MaterializedView(item.id()),
683            }
684        }
685        Explainee::Index(name) | Explainee::ReplanIndex(name) => {
686            let item = scx.get_item_by_resolved_name(&name)?;
687            let item_type = item.item_type();
688            if item_type != CatalogItemType::Index {
689                sql_bail!("Expected {name} to be an index, not a {item_type}");
690            }
691            match is_replan {
692                true => crate::plan::Explainee::ReplanIndex(item.id()),
693                false => crate::plan::Explainee::Index(item.id()),
694            }
695        }
696        Explainee::Select(select, broken) => {
697            let (plan, desc) = plan_select_inner(scx, *select, params, None)?;
698            crate::plan::Explainee::Statement(ExplaineeStatement::Select { broken, plan, desc })
699        }
700        Explainee::CreateView(mut stmt, broken) => {
701            if stmt.if_exists != IfExistsBehavior::Skip {
702                // If we don't force this parameter to Skip planning will
703                // fail for names that already exist in the catalog. This
704                // can happen even in `Replace` mode if the existing item
705                // has dependencies.
706                stmt.if_exists = IfExistsBehavior::Skip;
707            } else {
708                sql_bail!(
709                    "Cannot EXPLAIN a CREATE VIEW that explictly sets IF NOT EXISTS \
710                     (the behavior is implied within the scope of an enclosing EXPLAIN)"
711                );
712            }
713
714            let Plan::CreateView(plan) = ddl::plan_create_view(scx, *stmt)? else {
715                sql_bail!("expected CreateViewPlan plan");
716            };
717
718            crate::plan::Explainee::Statement(ExplaineeStatement::CreateView { broken, plan })
719        }
720        Explainee::CreateMaterializedView(mut stmt, broken) => {
721            if stmt.if_exists != IfExistsBehavior::Skip {
722                // If we don't force this parameter to Skip planning will
723                // fail for names that already exist in the catalog. This
724                // can happen even in `Replace` mode if the existing item
725                // has dependencies.
726                stmt.if_exists = IfExistsBehavior::Skip;
727            } else {
728                sql_bail!(
729                    "Cannot EXPLAIN a CREATE MATERIALIZED VIEW that explictly sets IF NOT EXISTS \
730                     (the behavior is implied within the scope of an enclosing EXPLAIN)"
731                );
732            }
733
734            let Plan::CreateMaterializedView(plan) =
735                ddl::plan_create_materialized_view(scx, *stmt)?
736            else {
737                sql_bail!("expected CreateMaterializedViewPlan plan");
738            };
739
740            crate::plan::Explainee::Statement(ExplaineeStatement::CreateMaterializedView {
741                broken,
742                plan,
743            })
744        }
745        Explainee::CreateIndex(mut stmt, broken) => {
746            if !stmt.if_not_exists {
747                // If we don't force this parameter to true planning will
748                // fail for index items that already exist in the catalog.
749                stmt.if_not_exists = true;
750            } else {
751                sql_bail!(
752                    "Cannot EXPLAIN a CREATE INDEX that explictly sets IF NOT EXISTS \
753                     (the behavior is implied within the scope of an enclosing EXPLAIN)"
754                );
755            }
756
757            let Plan::CreateIndex(plan) = ddl::plan_create_index(scx, *stmt)? else {
758                sql_bail!("expected CreateIndexPlan plan");
759            };
760
761            crate::plan::Explainee::Statement(ExplaineeStatement::CreateIndex { broken, plan })
762        }
763        Explainee::Subscribe(stmt, broken) => {
764            let Plan::Subscribe(plan) = plan_subscribe(scx, *stmt, params, None)? else {
765                sql_bail!("expected SubscribePlan");
766            };
767            crate::plan::Explainee::Statement(ExplaineeStatement::Subscribe { broken, plan })
768        }
769    };
770
771    Ok(explainee)
772}
773
774pub fn plan_explain_plan(
775    scx: &StatementContext,
776    explain: ExplainPlanStatement<Aug>,
777    params: &Params,
778) -> Result<Plan, PlanError> {
779    let (format, verbose_syntax) = match explain.format() {
780        mz_sql_parser::ast::ExplainFormat::Text => (ExplainFormat::Text, false),
781        mz_sql_parser::ast::ExplainFormat::VerboseText => (ExplainFormat::Text, true),
782        mz_sql_parser::ast::ExplainFormat::Json => (ExplainFormat::Json, false),
783        mz_sql_parser::ast::ExplainFormat::Dot => (ExplainFormat::Dot, false),
784    };
785    let stage = explain.stage();
786
787    // Plan ExplainConfig.
788    let mut config = {
789        let mut with_options = ExplainPlanOptionExtracted::try_from(explain.with_options)?;
790
791        if !scx.catalog.system_vars().persist_stats_filter_enabled() {
792            // If filtering is disabled, explain plans should not include pushdown info.
793            with_options.filter_pushdown = Some(false);
794        }
795
796        ExplainConfig::try_from(with_options)?
797    };
798    config.verbose_syntax = verbose_syntax;
799
800    let explainee = plan_explainee(scx, explain.explainee, params)?;
801
802    Ok(Plan::ExplainPlan(ExplainPlanPlan {
803        stage,
804        format,
805        config,
806        explainee,
807    }))
808}
809
810pub fn plan_explain_schema(
811    scx: &StatementContext,
812    explain_schema: ExplainSinkSchemaStatement<Aug>,
813) -> Result<Plan, PlanError> {
814    let ExplainSinkSchemaStatement {
815        schema_for,
816        // Parser limits to JSON.
817        format: _,
818        mut statement,
819    } = explain_schema;
820
821    // Force the sink's name to one that's guaranteed not to exist, by virtue of
822    // being a non-existent item in a schema under the system's control, so that
823    // `plan_create_sink` doesn't complain about the name already existing.
824    statement.name = Some(UnresolvedItemName::qualified(&[
825        ident!("mz_catalog"),
826        ident!("mz_explain_schema"),
827    ]));
828
829    crate::pure::purify_create_sink_avro_doc_on_options(
830        scx.catalog,
831        *statement.from.item_id(),
832        &mut statement.format,
833    )?;
834
835    match ddl::plan_create_sink(scx, statement)? {
836        Plan::CreateSink(CreateSinkPlan { sink, .. }) => match sink.connection {
837            StorageSinkConnection::Kafka(KafkaSinkConnection {
838                format:
839                    KafkaSinkFormat {
840                        key_format,
841                        value_format:
842                            KafkaSinkFormatType::Avro {
843                                schema: value_schema,
844                                ..
845                            },
846                        ..
847                    },
848                ..
849            }) => {
850                let schema = match schema_for {
851                    ExplainSinkSchemaFor::Key => key_format
852                        .and_then(|f| match f {
853                            KafkaSinkFormatType::Avro { schema, .. } => Some(schema),
854                            _ => None,
855                        })
856                        .ok_or_else(|| sql_err!("CREATE SINK does not have a key"))?,
857                    ExplainSinkSchemaFor::Value => value_schema,
858                };
859
860                Ok(Plan::ExplainSinkSchema(ExplainSinkSchemaPlan {
861                    sink_from: sink.from,
862                    json_schema: schema,
863                }))
864            }
865            _ => bail_unsupported!(
866                "EXPLAIN SCHEMA is only available for Kafka sinks with Avro schemas"
867            ),
868        },
869        _ => bail_internal!("plan_sink did not produce a CreateSink plan"),
870    }
871}
872
873pub fn plan_explain_pushdown(
874    scx: &StatementContext,
875    statement: ExplainPushdownStatement<Aug>,
876    params: &Params,
877) -> Result<Plan, PlanError> {
878    scx.require_feature_flag(&vars::ENABLE_EXPLAIN_PUSHDOWN)?;
879    let explainee = plan_explainee(scx, statement.explainee, params)?;
880    Ok(Plan::ExplainPushdown(ExplainPushdownPlan { explainee }))
881}
882
883pub fn plan_explain_analyze_object(
884    scx: &StatementContext,
885    statement: ExplainAnalyzeObjectStatement<Aug>,
886    params: &Params,
887) -> Result<Plan, PlanError> {
888    let explainee_name = statement
889        .explainee
890        .name()
891        .ok_or_else(|| sql_err!("EXPLAIN ANALYZE on anonymous dataflows",))?
892        .full_name_str();
893    let explainee = plan_explainee(scx, statement.explainee, params)?;
894
895    let check_ownership = |item_id: &CatalogItemId, item_type: &str| -> Result<(), PlanError> {
896        if scx.catalog.restrict_to_user_objects() {
897            let item = scx.catalog.get_item(item_id);
898            if item.owner_id() != *scx.catalog.active_role_id() {
899                let full_name = scx.catalog.resolve_full_name(item.name());
900                return Err(sql_err!("must be owner of {item_type} {full_name}"));
901            }
902        }
903        Ok(())
904    };
905    match &explainee {
906        plan::Explainee::Index(item_id) => check_ownership(item_id, "INDEX")?,
907        plan::Explainee::MaterializedView(item_id) => {
908            check_ownership(item_id, "MATERIALIZED VIEW")?
909        }
910        _ => return Err(sql_err!("EXPLAIN ANALYZE queries for this explainee type",)),
911    };
912
913    // generate SQL query
914
915    /* WITH {CTEs}
916       SELECT REPEAT(' ', nesting * 2) || operator AS operator
917             {columns}
918        FROM      mz_introspection.mz_lir_mapping mlm
919             JOIN {from} USING (lir_id)
920             JOIN mz_introspection.mz_mappable_objects mo
921               ON (mlm.global_id = mo.global_id)
922       WHERE     mo.name = {escaped explainee_name}
923             AND {predicates}
924       ORDER BY lir_id DESC
925    */
926    let mut ctes = Vec::with_capacity(4); // max 2 per ExplainAnalyzeComputationProperty
927    let mut columns = vec!["REPEAT(' ', nesting * 2) || operator AS operator"];
928    let mut from = vec!["mz_introspection.mz_lir_mapping mlm"];
929    let mut predicates = vec![format!(
930        "mo.name = {}",
931        escaped_string_literal(&explainee_name)
932    )];
933    let mut order_by = vec!["mlm.lir_id DESC"];
934
935    match statement.properties {
936        ExplainAnalyzeProperty::Computation(ExplainAnalyzeComputationProperties {
937            properties,
938            skew,
939        }) => {
940            let mut worker_id = None;
941            let mut seen_properties = BTreeSet::new();
942            for property in properties {
943                // handle each property only once (belt and suspenders)
944                if !seen_properties.insert(property) {
945                    continue;
946                }
947
948                match property {
949                    ExplainAnalyzeComputationProperty::Memory => {
950                        ctes.push((
951                            "summary_memory",
952                            r#"
953  SELECT mlm.global_id AS global_id,
954         mlm.lir_id AS lir_id,
955         SUM(mas.size) AS total_memory,
956         SUM(mas.records) AS total_records,
957         CASE WHEN COUNT(DISTINCT mas.worker_id) <> 0 THEN SUM(mas.size) / COUNT(DISTINCT mas.worker_id) ELSE NULL END AS avg_memory,
958         CASE WHEN COUNT(DISTINCT mas.worker_id) <> 0 THEN SUM(mas.records) / COUNT(DISTINCT mas.worker_id) ELSE NULL END AS avg_records
959    FROM            mz_introspection.mz_lir_mapping mlm
960         CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
961               JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
962                 ON (mas.operator_id = valid_id)
963GROUP BY mlm.global_id, mlm.lir_id"#,
964                        ));
965                        from.push("LEFT JOIN summary_memory sm USING (global_id, lir_id)");
966
967                        if skew {
968                            ctes.push((
969                                "per_worker_memory",
970                                r#"
971  SELECT mlm.global_id AS global_id,
972         mlm.lir_id AS lir_id,
973         mas.worker_id AS worker_id,
974         SUM(mas.size) AS worker_memory,
975         SUM(mas.records) AS worker_records
976    FROM            mz_introspection.mz_lir_mapping mlm
977         CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
978               JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
979                 ON (mas.operator_id = valid_id)
980GROUP BY mlm.global_id, mlm.lir_id, mas.worker_id"#,
981                            ));
982                            from.push("LEFT JOIN per_worker_memory pwm USING (global_id, lir_id)");
983
984                            if let Some(worker_id) = worker_id {
985                                predicates.push(format!(
986                                    "(pwm.worker_id = {worker_id} OR pwm.worker_id IS NULL OR {worker_id} IS NULL)"
987                                ));
988                            } else {
989                                worker_id = Some("pwm.worker_id");
990                                columns.push("pwm.worker_id AS worker_id");
991                                order_by.push("worker_id");
992                            }
993
994                            columns.extend([
995                                "CASE WHEN pwm.worker_id IS NOT NULL AND sm.avg_memory <> 0 THEN ROUND(pwm.worker_memory / sm.avg_memory, 2) ELSE NULL END AS memory_ratio",
996                                "pg_size_pretty(pwm.worker_memory) AS worker_memory",
997                                "pg_size_pretty(sm.avg_memory) AS avg_memory",
998                                "pg_size_pretty(sm.total_memory) AS total_memory",
999                                "CASE WHEN pwm.worker_id IS NOT NULL AND sm.avg_records <> 0 THEN ROUND(pwm.worker_records / sm.avg_records, 2) ELSE NULL END AS records_ratio",
1000                                "pwm.worker_records AS worker_records",
1001                                "sm.avg_records AS avg_records",
1002                                "sm.total_records AS total_records",
1003                            ]);
1004                        } else {
1005                            columns.extend([
1006                                "pg_size_pretty(sm.total_memory) AS total_memory",
1007                                "sm.total_records AS total_records",
1008                            ]);
1009                        }
1010                    }
1011                    ExplainAnalyzeComputationProperty::Cpu => {
1012                        ctes.push((
1013                            "summary_cpu",
1014                            r#"
1015  SELECT mlm.global_id AS global_id,
1016         mlm.lir_id AS lir_id,
1017         SUM(mse.elapsed_ns) AS total_ns,
1018         CASE WHEN COUNT(DISTINCT mse.worker_id) <> 0 THEN SUM(mse.elapsed_ns) / COUNT(DISTINCT mse.worker_id) ELSE NULL END AS avg_ns
1019    FROM            mz_introspection.mz_lir_mapping mlm
1020         CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1021               JOIN mz_introspection.mz_scheduling_elapsed_per_worker mse
1022                 ON (mse.id = valid_id)
1023GROUP BY mlm.global_id, mlm.lir_id"#,
1024                        ));
1025                        from.push("LEFT JOIN summary_cpu sc USING (global_id, lir_id)");
1026
1027                        if skew {
1028                            ctes.push((
1029                                "per_worker_cpu",
1030                                r#"
1031  SELECT mlm.global_id AS global_id,
1032         mlm.lir_id AS lir_id,
1033         mse.worker_id AS worker_id,
1034         SUM(mse.elapsed_ns) AS worker_ns
1035    FROM            mz_introspection.mz_lir_mapping mlm
1036         CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1037               JOIN mz_introspection.mz_scheduling_elapsed_per_worker mse
1038                 ON (mse.id = valid_id)
1039GROUP BY mlm.global_id, mlm.lir_id, mse.worker_id"#,
1040                            ));
1041                            from.push("LEFT JOIN per_worker_cpu pwc USING (global_id, lir_id)");
1042
1043                            if let Some(worker_id) = worker_id {
1044                                predicates.push(format!(
1045                                    "(pwc.worker_id = {worker_id} OR pwc.worker_id IS NULL OR {worker_id} IS NULL)"
1046                                ));
1047                            } else {
1048                                worker_id = Some("pwc.worker_id");
1049                                columns.push("pwc.worker_id AS worker_id");
1050                                order_by.push("worker_id");
1051                            }
1052
1053                            columns.extend([
1054                                "CASE WHEN pwc.worker_id IS NOT NULL AND sc.avg_ns <> 0 THEN ROUND(pwc.worker_ns / sc.avg_ns, 2) ELSE NULL END AS cpu_ratio",
1055                                "pwc.worker_ns / 1000 * '1 microsecond'::INTERVAL AS worker_elapsed",
1056                                "sc.avg_ns / 1000 * '1 microsecond'::INTERVAL AS avg_elapsed",
1057                            ]);
1058                        }
1059                        columns.push(
1060                            "sc.total_ns / 1000 * '1 microsecond'::INTERVAL AS total_elapsed",
1061                        );
1062                    }
1063                }
1064            }
1065        }
1066        ExplainAnalyzeProperty::Hints => {
1067            columns.extend([
1068                "megsa.levels AS levels",
1069                "megsa.to_cut AS to_cut",
1070                "megsa.hint AS hint",
1071                "pg_size_pretty(megsa.savings) AS savings",
1072            ]);
1073            from.extend(["JOIN mz_introspection.mz_dataflow_global_ids mdgi ON (mlm.global_id = mdgi.global_id)",
1074            "LEFT JOIN (generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id JOIN \
1075             mz_introspection.mz_expected_group_size_advice megsa ON (megsa.region_id = valid_id)) ON (megsa.dataflow_id = mdgi.id)"]);
1076        }
1077    }
1078
1079    from.push("JOIN mz_introspection.mz_mappable_objects mo ON (mlm.global_id = mo.global_id)");
1080
1081    let ctes = if !ctes.is_empty() {
1082        format!(
1083            "WITH {}",
1084            separated(
1085                ",\n",
1086                ctes.iter()
1087                    .map(|(name, defn)| format!("{name} AS ({defn})"))
1088            )
1089        )
1090    } else {
1091        String::new()
1092    };
1093    let columns = separated(", ", columns);
1094    let from = separated(" ", from);
1095    let predicates = separated(" AND ", predicates);
1096    let order_by = separated(", ", order_by);
1097    let query = format!(
1098        r#"{ctes}
1099SELECT {columns}
1100FROM {from}
1101WHERE {predicates}
1102ORDER BY {order_by}"#
1103    );
1104
1105    if statement.as_sql {
1106        let rows = vec![Row::pack_slice(&[Datum::String(
1107            &mz_sql_pretty::pretty_str_simple(&query, 80).map_err(|e| {
1108                PlanError::Unstructured(format!("internal error parsing our own SQL: {e}"))
1109            })?,
1110        )])];
1111        let typ = SqlRelationType::new(vec![SqlScalarType::String.nullable(false)]);
1112
1113        Ok(Plan::Select(SelectPlan::immediate(rows, typ)))
1114    } else {
1115        let (show_select, resolved_ids) = ShowSelect::new_from_bare_query(scx, query)?;
1116        scx.record_sql_impl_ids(&resolved_ids);
1117        show_select.plan()
1118    }
1119}
1120
1121pub fn plan_explain_analyze_cluster(
1122    scx: &StatementContext,
1123    statement: ExplainAnalyzeClusterStatement,
1124    _params: &Params,
1125) -> Result<Plan, PlanError> {
1126    // object string
1127    // worker_id uint64        (if           skew)
1128    // memory_ratio numeric    (if memory && skew)
1129    // worker_memory string    (if memory && skew)
1130    // avg_memory string       (if memory && skew)
1131    // total_memory string     (if memory)
1132    // records_ratio numeric   (if memory && skew)
1133    // worker_records          (if memory && skew)
1134    // avg_records numeric     (if memory && skew)
1135    // total_records numeric   (if memory)
1136    // cpu_ratio numeric       (if cpu    && skew)
1137    // worker_elapsed interval (if cpu    && skew)
1138    // avg_elapsed interval    (if cpu    && skew)
1139    // total_elapsed interval  (if cpu)
1140
1141    /* WITH {CTEs}
1142       SELECT mo.name AS object
1143             {columns}
1144        FROM mz_introspection.mz_mappable_objects mo
1145             {from}
1146       WHERE {predicates}
1147       ORDER BY {order_by}, mo.name DESC
1148    */
1149    let mut ctes = Vec::with_capacity(4); // max 2 per ExplainAnalyzeComputationProperty
1150    let mut columns = vec!["mo.name AS object", "mo.global_id AS global_id"];
1151    let mut from = vec!["mz_introspection.mz_mappable_objects mo"];
1152    let mut predicates = vec![];
1153    let mut order_by = vec![];
1154
1155    let ExplainAnalyzeComputationProperties { properties, skew } = statement.properties;
1156    let mut worker_id = None;
1157    let mut seen_properties = BTreeSet::new();
1158    for property in properties {
1159        // handle each property only once (belt and suspenders)
1160        if !seen_properties.insert(property) {
1161            continue;
1162        }
1163
1164        match property {
1165            ExplainAnalyzeComputationProperty::Memory => {
1166                if skew {
1167                    let mut set_worker_id = false;
1168                    if let Some(worker_id) = worker_id {
1169                        // join condition if we're showing skew for more than one property
1170                        predicates.push(format!(
1171                            "(om.worker_id = {worker_id} OR om.worker_id IS NULL OR {worker_id} IS NULL)"
1172                        ));
1173                    } else {
1174                        worker_id = Some("om.worker_id");
1175                        columns.push("om.worker_id AS worker_id");
1176                        set_worker_id = true; // we'll add ourselves to `order_by` later
1177                    };
1178
1179                    // computes the average memory per LIR operator (for per operator ratios)
1180                    ctes.push((
1181                    "per_operator_memory_summary",
1182                    r#"
1183SELECT mlm.global_id AS global_id,
1184       mlm.lir_id AS lir_id,
1185       SUM(mas.size) AS total_memory,
1186       SUM(mas.records) AS total_records,
1187       CASE WHEN COUNT(DISTINCT mas.worker_id) <> 0 THEN SUM(mas.size) / COUNT(DISTINCT mas.worker_id) ELSE NULL END AS avg_memory,
1188       CASE WHEN COUNT(DISTINCT mas.worker_id) <> 0 THEN SUM(mas.records) / COUNT(DISTINCT mas.worker_id) ELSE NULL END AS avg_records
1189FROM        mz_introspection.mz_lir_mapping mlm
1190 CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1191       JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
1192         ON (mas.operator_id = valid_id)
1193GROUP BY mlm.global_id, mlm.lir_id"#,
1194                ));
1195
1196                    // computes the memory per worker in a per operator way
1197                    ctes.push((
1198                    "per_operator_memory_per_worker",
1199                    r#"
1200SELECT mlm.global_id AS global_id,
1201       mlm.lir_id AS lir_id,
1202       mas.worker_id AS worker_id,
1203       SUM(mas.size) AS worker_memory,
1204       SUM(mas.records) AS worker_records
1205FROM        mz_introspection.mz_lir_mapping mlm
1206 CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1207       JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
1208         ON (mas.operator_id = valid_id)
1209GROUP BY mlm.global_id, mlm.lir_id, mas.worker_id"#,
1210                    ));
1211
1212                    // computes memory ratios per worker per operator
1213                    ctes.push((
1214                    "per_operator_memory_ratios",
1215                    r#"
1216SELECT pompw.global_id AS global_id,
1217       pompw.lir_id AS lir_id,
1218       pompw.worker_id AS worker_id,
1219       CASE WHEN pompw.worker_id IS NOT NULL AND poms.avg_memory <> 0 THEN ROUND(pompw.worker_memory / poms.avg_memory, 2) ELSE NULL END AS memory_ratio,
1220       CASE WHEN pompw.worker_id IS NOT NULL AND poms.avg_records <> 0 THEN ROUND(pompw.worker_records / poms.avg_records, 2) ELSE NULL END AS records_ratio
1221  FROM      per_operator_memory_per_worker pompw
1222       JOIN per_operator_memory_summary poms
1223         USING (global_id, lir_id)
1224"#,
1225                    ));
1226
1227                    // summarizes each object, per worker
1228                    ctes.push((
1229                        "object_memory",
1230                        r#"
1231SELECT pompw.global_id AS global_id,
1232       pompw.worker_id AS worker_id,
1233       MAX(pomr.memory_ratio) AS max_operator_memory_ratio,
1234       MAX(pomr.records_ratio) AS max_operator_records_ratio,
1235       SUM(pompw.worker_memory) AS worker_memory,
1236       SUM(pompw.worker_records) AS worker_records
1237FROM        per_operator_memory_per_worker pompw
1238     JOIN   per_operator_memory_ratios pomr
1239     USING (global_id, worker_id, lir_id)
1240GROUP BY pompw.global_id, pompw.worker_id
1241"#,
1242                    ));
1243
1244                    // summarizes each worker
1245                    ctes.push(("object_average_memory", r#"
1246SELECT om.global_id AS global_id,
1247       SUM(om.worker_memory) AS total_memory,
1248       CASE WHEN COUNT(DISTINCT om.worker_id) <> 0 THEN SUM(om.worker_memory) / COUNT(DISTINCT om.worker_id) ELSE NULL END AS avg_memory,
1249       SUM(om.worker_records) AS total_records,
1250       CASE WHEN COUNT(DISTINCT om.worker_id) <> 0 THEN SUM(om.worker_records) / COUNT(DISTINCT om.worker_id) ELSE NULL END AS avg_records
1251  FROM object_memory om
1252GROUP BY om.global_id"#));
1253
1254                    from.push("LEFT JOIN object_memory om USING (global_id)");
1255                    from.push("LEFT JOIN object_average_memory oam USING (global_id)");
1256
1257                    columns.extend([
1258                        "om.max_operator_memory_ratio AS max_operator_memory_ratio",
1259                        "pg_size_pretty(om.worker_memory) AS worker_memory",
1260                        "pg_size_pretty(oam.avg_memory) AS avg_memory",
1261                        "pg_size_pretty(oam.total_memory) AS total_memory",
1262                        "om.max_operator_records_ratio AS max_operator_records_ratio",
1263                        "om.worker_records AS worker_records",
1264                        "oam.avg_records AS avg_records",
1265                        "oam.total_records AS total_records",
1266                    ]);
1267
1268                    order_by.extend([
1269                        "max_operator_memory_ratio DESC NULLS LAST",
1270                        "max_operator_records_ratio DESC NULLS LAST",
1271                        "om.worker_memory DESC NULLS LAST",
1272                        "worker_records DESC NULLS LAST",
1273                    ]);
1274
1275                    if set_worker_id {
1276                        order_by.push("worker_id");
1277                    }
1278                } else {
1279                    // no skew, so just compute totals
1280                    ctes.push((
1281                        "per_operator_memory_totals",
1282                        r#"
1283    SELECT mlm.global_id AS global_id,
1284           mlm.lir_id AS lir_id,
1285           SUM(mas.size) AS total_memory,
1286           SUM(mas.records) AS total_records
1287    FROM        mz_introspection.mz_lir_mapping mlm
1288     CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1289           JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
1290             ON (mas.operator_id = valid_id)
1291    GROUP BY mlm.global_id, mlm.lir_id"#,
1292                    ));
1293
1294                    ctes.push((
1295                        "object_memory_totals",
1296                        r#"
1297SELECT pomt.global_id AS global_id,
1298       SUM(pomt.total_memory) AS total_memory,
1299       SUM(pomt.total_records) AS total_records
1300FROM per_operator_memory_totals pomt
1301GROUP BY pomt.global_id
1302"#,
1303                    ));
1304
1305                    from.push("LEFT JOIN object_memory_totals omt USING (global_id)");
1306                    columns.extend([
1307                        "pg_size_pretty(omt.total_memory) AS total_memory",
1308                        "omt.total_records AS total_records",
1309                    ]);
1310                    order_by.extend([
1311                        "omt.total_memory DESC NULLS LAST",
1312                        "total_records DESC NULLS LAST",
1313                    ]);
1314                }
1315            }
1316            ExplainAnalyzeComputationProperty::Cpu => {
1317                if skew {
1318                    let mut set_worker_id = false;
1319                    if let Some(worker_id) = worker_id {
1320                        // join condition if we're showing skew for more than one property
1321                        predicates.push(format!(
1322                            "(oc.worker_id = {worker_id} OR oc.worker_id IS NULL OR {worker_id} IS NULL)"
1323                        ));
1324                    } else {
1325                        worker_id = Some("oc.worker_id");
1326                        columns.push("oc.worker_id AS worker_id");
1327                        set_worker_id = true; // we'll add ourselves to `order_by` later
1328                    };
1329
1330                    // computes the average memory per LIR operator (for per operator ratios)
1331                    ctes.push((
1332    "per_operator_cpu_summary",
1333    r#"
1334SELECT mlm.global_id AS global_id,
1335       mlm.lir_id AS lir_id,
1336       SUM(mse.elapsed_ns) AS total_ns,
1337       CASE WHEN COUNT(DISTINCT mse.worker_id) <> 0 THEN SUM(mse.elapsed_ns) / COUNT(DISTINCT mse.worker_id) ELSE NULL END AS avg_ns
1338FROM       mz_introspection.mz_lir_mapping mlm
1339CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1340      JOIN mz_introspection.mz_scheduling_elapsed_per_worker mse
1341        ON (mse.id = valid_id)
1342GROUP BY mlm.global_id, mlm.lir_id"#,
1343));
1344
1345                    // computes the CPU per worker in a per operator way
1346                    ctes.push((
1347                        "per_operator_cpu_per_worker",
1348                        r#"
1349SELECT mlm.global_id AS global_id,
1350       mlm.lir_id AS lir_id,
1351       mse.worker_id AS worker_id,
1352       SUM(mse.elapsed_ns) AS worker_ns
1353FROM       mz_introspection.mz_lir_mapping mlm
1354CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1355      JOIN mz_introspection.mz_scheduling_elapsed_per_worker mse
1356        ON (mse.id = valid_id)
1357GROUP BY mlm.global_id, mlm.lir_id, mse.worker_id"#,
1358                    ));
1359
1360                    // computes CPU ratios per worker per operator
1361                    ctes.push((
1362                        "per_operator_cpu_ratios",
1363                        r#"
1364SELECT pocpw.global_id AS global_id,
1365       pocpw.lir_id AS lir_id,
1366       pocpw.worker_id AS worker_id,
1367       CASE WHEN pocpw.worker_id IS NOT NULL AND pocs.avg_ns <> 0 THEN ROUND(pocpw.worker_ns / pocs.avg_ns, 2) ELSE NULL END AS cpu_ratio
1368FROM      per_operator_cpu_per_worker pocpw
1369     JOIN per_operator_cpu_summary pocs
1370     USING (global_id, lir_id)
1371"#,
1372                    ));
1373
1374                    // summarizes each object, per worker
1375                    ctes.push((
1376                        "object_cpu",
1377                        r#"
1378SELECT pocpw.global_id AS global_id,
1379       pocpw.worker_id AS worker_id,
1380       MAX(pomr.cpu_ratio) AS max_operator_cpu_ratio,
1381       SUM(pocpw.worker_ns) AS worker_ns
1382FROM      per_operator_cpu_per_worker pocpw
1383     JOIN per_operator_cpu_ratios pomr
1384     USING (global_id, worker_id, lir_id)
1385GROUP BY pocpw.global_id, pocpw.worker_id
1386"#,
1387                    ));
1388
1389                    // summarizes each worker
1390                    ctes.push((
1391                        "object_average_cpu",
1392                        r#"
1393SELECT oc.global_id AS global_id,
1394       SUM(oc.worker_ns) AS total_ns,
1395       CASE WHEN COUNT(DISTINCT oc.worker_id) <> 0 THEN SUM(oc.worker_ns) / COUNT(DISTINCT oc.worker_id) ELSE NULL END AS avg_ns
1396  FROM object_cpu oc
1397GROUP BY oc.global_id"#,));
1398
1399                    from.push("LEFT JOIN object_cpu oc USING (global_id)");
1400                    from.push("LEFT JOIN object_average_cpu oac USING (global_id)");
1401
1402                    columns.extend([
1403                        "oc.max_operator_cpu_ratio AS max_operator_cpu_ratio",
1404                        "oc.worker_ns / 1000 * '1 microsecond'::interval AS worker_elapsed",
1405                        "oac.avg_ns / 1000 * '1 microsecond'::interval AS avg_elapsed",
1406                        "oac.total_ns / 1000 * '1 microsecond'::interval AS total_elapsed",
1407                    ]);
1408
1409                    order_by.extend([
1410                        "max_operator_cpu_ratio DESC NULLS LAST",
1411                        "worker_elapsed DESC NULLS LAST",
1412                    ]);
1413
1414                    if set_worker_id {
1415                        order_by.push("worker_id");
1416                    }
1417                } else {
1418                    // no skew, so just compute totals
1419                    ctes.push((
1420                        "per_operator_cpu_totals",
1421                        r#"
1422    SELECT mlm.global_id AS global_id,
1423           mlm.lir_id AS lir_id,
1424           SUM(mse.elapsed_ns) AS total_ns
1425    FROM        mz_introspection.mz_lir_mapping mlm
1426     CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1427           JOIN mz_introspection.mz_scheduling_elapsed_per_worker mse
1428             ON (mse.id = valid_id)
1429    GROUP BY mlm.global_id, mlm.lir_id"#,
1430                    ));
1431
1432                    ctes.push((
1433                        "object_cpu_totals",
1434                        r#"
1435SELECT poct.global_id AS global_id,
1436       SUM(poct.total_ns) AS total_ns
1437FROM per_operator_cpu_totals poct
1438GROUP BY poct.global_id
1439"#,
1440                    ));
1441
1442                    from.push("LEFT JOIN object_cpu_totals oct USING (global_id)");
1443                    columns
1444                        .push("oct.total_ns / 1000 * '1 microsecond'::interval AS total_elapsed");
1445                    order_by.extend(["total_elapsed DESC NULLS LAST"]);
1446                }
1447            }
1448        }
1449    }
1450
1451    // generate SQL query text
1452    let ctes = if !ctes.is_empty() {
1453        format!(
1454            "WITH {}",
1455            separated(
1456                ",\n",
1457                ctes.iter()
1458                    .map(|(name, defn)| format!("{name} AS ({defn})"))
1459            )
1460        )
1461    } else {
1462        String::new()
1463    };
1464    let columns = separated(", ", columns);
1465    let from = separated(" ", from);
1466    let predicates = if !predicates.is_empty() {
1467        format!("WHERE {}", separated(" AND ", predicates))
1468    } else {
1469        String::new()
1470    };
1471    // add mo.name last, to break ties only
1472    order_by.push("mo.name DESC");
1473    let order_by = separated(", ", order_by);
1474    let query = format!(
1475        r#"{ctes}
1476SELECT {columns}
1477FROM {from}
1478{predicates}
1479ORDER BY {order_by}"#
1480    );
1481
1482    if statement.as_sql {
1483        let rows = vec![Row::pack_slice(&[Datum::String(
1484            &mz_sql_pretty::pretty_str_simple(&query, 80).map_err(|e| {
1485                PlanError::Unstructured(format!("internal error parsing our own SQL: {e}"))
1486            })?,
1487        )])];
1488        let typ = SqlRelationType::new(vec![SqlScalarType::String.nullable(false)]);
1489
1490        Ok(Plan::Select(SelectPlan::immediate(rows, typ)))
1491    } else {
1492        let (show_select, resolved_ids) = ShowSelect::new_from_bare_query(scx, query)?;
1493        scx.record_sql_impl_ids(&resolved_ids);
1494        show_select.plan()
1495    }
1496}
1497
1498pub fn plan_explain_timestamp(
1499    scx: &StatementContext,
1500    explain: ExplainTimestampStatement<Aug>,
1501) -> Result<Plan, PlanError> {
1502    let (format, _verbose_syntax) = match explain.format() {
1503        mz_sql_parser::ast::ExplainFormat::Text => (ExplainFormat::Text, false),
1504        mz_sql_parser::ast::ExplainFormat::VerboseText => (ExplainFormat::Text, true),
1505        mz_sql_parser::ast::ExplainFormat::Json => (ExplainFormat::Json, false),
1506        mz_sql_parser::ast::ExplainFormat::Dot => (ExplainFormat::Dot, false),
1507    };
1508
1509    let raw_plan = {
1510        let query::PlannedRootQuery {
1511            expr: raw_plan,
1512            desc: _,
1513            finishing: _,
1514            scope: _,
1515        } = query::plan_root_query(scx, explain.select.query, QueryLifetime::OneShot)?;
1516        if raw_plan.contains_parameters()? {
1517            return Err(PlanError::ParameterNotAllowed(
1518                "EXPLAIN TIMESTAMP".to_string(),
1519            ));
1520        }
1521
1522        raw_plan
1523    };
1524    let when = query::plan_as_of(scx, explain.select.as_of)?;
1525
1526    Ok(Plan::ExplainTimestamp(ExplainTimestampPlan {
1527        format,
1528        raw_plan,
1529        when,
1530    }))
1531}
1532
1533generate_extracted_config!(SubscribeOption, (Snapshot, bool), (Progress, bool));
1534
1535pub fn describe_subscribe(
1536    scx: &StatementContext,
1537    stmt: SubscribeStatement<Aug>,
1538) -> Result<StatementDesc, PlanError> {
1539    let relation_desc = match stmt.relation {
1540        SubscribeRelation::Name(name) => {
1541            let item = scx.get_item_by_resolved_name(&name)?;
1542            match item.relation_desc() {
1543                Some(desc) => desc.into_owned(),
1544                None => sql_bail!(
1545                    "'{}' cannot be subscribed to because it is a {}",
1546                    name.full_name_str(),
1547                    item.item_type(),
1548                ),
1549            }
1550        }
1551        SubscribeRelation::Query(query) => {
1552            let query::PlannedRootQuery { desc, .. } =
1553                query::plan_root_query(scx, query, QueryLifetime::Subscribe)?;
1554            desc
1555        }
1556    };
1557    let SubscribeOptionExtracted { progress, .. } = stmt.options.try_into()?;
1558    let progress = progress.unwrap_or(false);
1559    let mut desc = RelationDesc::builder().with_column(
1560        "mz_timestamp",
1561        SqlScalarType::Numeric {
1562            max_scale: Some(NumericMaxScale::ZERO),
1563        }
1564        .nullable(false),
1565    );
1566    if progress {
1567        desc = desc.with_column("mz_progressed", SqlScalarType::Bool.nullable(false));
1568    }
1569
1570    let debezium = matches!(stmt.output, SubscribeOutput::EnvelopeDebezium { .. });
1571    match stmt.output {
1572        SubscribeOutput::Diffs | SubscribeOutput::WithinTimestampOrderBy { .. } => {
1573            desc = desc.with_column("mz_diff", SqlScalarType::Int64.nullable(true));
1574            for (name, mut ty) in relation_desc.into_iter() {
1575                if progress {
1576                    ty.nullable = true;
1577                }
1578                desc = desc.with_column(name, ty);
1579            }
1580        }
1581        SubscribeOutput::EnvelopeUpsert { key_columns }
1582        | SubscribeOutput::EnvelopeDebezium { key_columns } => {
1583            desc = desc.with_column("mz_state", SqlScalarType::String.nullable(true));
1584            let key_columns = key_columns
1585                .into_iter()
1586                .map(normalize::column_name)
1587                .collect_vec();
1588            let mut before_values_desc = RelationDesc::builder();
1589            let mut after_values_desc = RelationDesc::builder();
1590
1591            // Add the key columns in the order that they're specified.
1592            for column_name in &key_columns {
1593                let mut column_ty = relation_desc
1594                    .get_by_name(column_name)
1595                    .map(|(_pos, ty)| ty.clone())
1596                    .ok_or_else(|| PlanError::UnknownColumn {
1597                        table: None,
1598                        column: column_name.clone(),
1599                        similar: Box::new([]),
1600                    })?;
1601                if progress {
1602                    column_ty.nullable = true;
1603                }
1604                desc = desc.with_column(column_name, column_ty);
1605            }
1606
1607            // Then add the remaining columns in the order from the original
1608            // table, filtering out the key columns since we added those above.
1609            for (mut name, mut ty) in relation_desc
1610                .into_iter()
1611                .filter(|(name, _ty)| !key_columns.contains(name))
1612            {
1613                ty.nullable = true;
1614                before_values_desc =
1615                    before_values_desc.with_column(format!("before_{}", name), ty.clone());
1616                if debezium {
1617                    name = format!("after_{}", name).into();
1618                }
1619                after_values_desc = after_values_desc.with_column(name, ty);
1620            }
1621
1622            if debezium {
1623                desc = desc.concat(before_values_desc);
1624            }
1625            desc = desc.concat(after_values_desc);
1626        }
1627    }
1628    Ok(StatementDesc::new(Some(desc.finish())))
1629}
1630
1631pub fn plan_subscribe(
1632    scx: &StatementContext,
1633    SubscribeStatement {
1634        relation,
1635        options,
1636        as_of,
1637        up_to,
1638        output,
1639    }: SubscribeStatement<Aug>,
1640    params: &Params,
1641    copy_to: Option<CopyFormat>,
1642) -> Result<Plan, PlanError> {
1643    let (from, desc, scope) = match relation {
1644        SubscribeRelation::Name(name) => {
1645            let item = scx.get_item_by_resolved_name(&name)?;
1646            let Some(desc) = item.relation_desc() else {
1647                sql_bail!(
1648                    "'{}' cannot be subscribed to because it is a {}",
1649                    name.full_name_str(),
1650                    item.item_type(),
1651                );
1652            };
1653            let item_name = match name {
1654                ResolvedItemName::Item { full_name, .. } => Some(full_name.into()),
1655                _ => None,
1656            };
1657            let scope = Scope::from_source(item_name, desc.iter().map(|(name, _type)| name));
1658            (
1659                SubscribeFrom::Id(item.global_id()),
1660                desc.into_owned(),
1661                scope,
1662            )
1663        }
1664        SubscribeRelation::Query(query) => {
1665            #[allow(deprecated)] // TODO(aalexandrov): Use HirRelationExpr in Subscribe
1666            let query::PlannedRootQuery {
1667                mut expr,
1668                desc,
1669                finishing,
1670                scope,
1671            } = query::plan_root_query(scx, query, QueryLifetime::Subscribe)?;
1672            expr.bind_parameters_and_simplify_offset(scx, QueryLifetime::Subscribe, params)?;
1673            let query = query::PlannedRootQuery {
1674                expr,
1675                desc,
1676                finishing,
1677                scope,
1678            };
1679            // There's no way to apply finishing operations to a `SUBSCRIBE` directly, so the
1680            // finishing should have already been turned into a `TopK` by
1681            // `plan_query` / `plan_root_query`, upon seeing the `QueryLifetime::Subscribe`.
1682            assert!(HirRelationExpr::is_trivial_row_set_finishing_hir(
1683                &query.finishing,
1684                query.desc.arity()
1685            ));
1686            let desc = query.desc.clone();
1687            (
1688                SubscribeFrom::Query {
1689                    expr: query.expr,
1690                    desc: query.desc,
1691                },
1692                desc,
1693                query.scope,
1694            )
1695        }
1696    };
1697
1698    let when = query::plan_as_of(scx, as_of)?;
1699    let up_to = up_to
1700        .map(|up_to| plan_as_of_or_up_to(scx, up_to))
1701        .transpose()?;
1702
1703    let qcx = QueryContext::root(scx, QueryLifetime::Subscribe);
1704    let ecx = ExprContext {
1705        qcx: &qcx,
1706        name: "",
1707        scope: &scope,
1708        relation_type: desc.typ(),
1709        allow_aggregates: false,
1710        allow_subqueries: true,
1711        allow_parameters: true,
1712        allow_windows: false,
1713    };
1714
1715    let output_columns: Vec<_> = scope.column_names().enumerate().collect();
1716    let output = match output {
1717        SubscribeOutput::Diffs => plan::SubscribeOutput::Diffs,
1718        SubscribeOutput::EnvelopeUpsert { key_columns } => {
1719            let order_by = key_columns
1720                .iter()
1721                .map(|ident| OrderByExpr {
1722                    expr: Expr::Identifier(vec![ident.clone()]),
1723                    asc: None,
1724                    nulls_last: None,
1725                })
1726                .collect_vec();
1727            let (order_by, map_exprs) = query::plan_order_by_exprs(
1728                &ExprContext {
1729                    name: "ENVELOPE UPSERT KEY clause",
1730                    ..ecx
1731                },
1732                &order_by[..],
1733                &output_columns[..],
1734            )?;
1735            if !map_exprs.is_empty() {
1736                return Err(PlanError::InvalidKeysInSubscribeEnvelopeUpsert);
1737            }
1738            check_distinct_key_columns(&order_by, &output_columns)?;
1739            plan::SubscribeOutput::EnvelopeUpsert {
1740                order_by_keys: order_by,
1741            }
1742        }
1743        SubscribeOutput::EnvelopeDebezium { key_columns } => {
1744            scx.require_feature_flag(&vars::ENABLE_ENVELOPE_DEBEZIUM_IN_SUBSCRIBE)?;
1745            let order_by = key_columns
1746                .iter()
1747                .map(|ident| OrderByExpr {
1748                    expr: Expr::Identifier(vec![ident.clone()]),
1749                    asc: None,
1750                    nulls_last: None,
1751                })
1752                .collect_vec();
1753            let (order_by, map_exprs) = query::plan_order_by_exprs(
1754                &ExprContext {
1755                    name: "ENVELOPE DEBEZIUM KEY clause",
1756                    ..ecx
1757                },
1758                &order_by[..],
1759                &output_columns[..],
1760            )?;
1761            if !map_exprs.is_empty() {
1762                return Err(PlanError::InvalidKeysInSubscribeEnvelopeDebezium);
1763            }
1764            check_distinct_key_columns(&order_by, &output_columns)?;
1765            plan::SubscribeOutput::EnvelopeDebezium {
1766                order_by_keys: order_by,
1767            }
1768        }
1769        SubscribeOutput::WithinTimestampOrderBy { order_by } => {
1770            scx.require_feature_flag(&vars::ENABLE_WITHIN_TIMESTAMP_ORDER_BY_IN_SUBSCRIBE)?;
1771            let mz_diff = "mz_diff".into();
1772            let output_columns = std::iter::once((0, &mz_diff))
1773                .chain(output_columns.into_iter().map(|(i, c)| (i + 1, c)))
1774                .collect_vec();
1775            match query::plan_order_by_exprs(
1776                &ExprContext {
1777                    name: "WITHIN TIMESTAMP ORDER BY clause",
1778                    ..ecx
1779                },
1780                &order_by[..],
1781                &output_columns[..],
1782            ) {
1783                Err(PlanError::UnknownColumn {
1784                    table: None,
1785                    column,
1786                    similar: _,
1787                }) if &column == &mz_diff => {
1788                    // mz_diff is being used in an expression. Since mz_diff isn't part of the table
1789                    // it looks like an unknown column. Instead, return a better error
1790                    return Err(PlanError::InvalidOrderByInSubscribeWithinTimestampOrderBy);
1791                }
1792                Err(e) => return Err(e),
1793                Ok((order_by, map_exprs)) => {
1794                    if !map_exprs.is_empty() {
1795                        return Err(PlanError::InvalidOrderByInSubscribeWithinTimestampOrderBy);
1796                    }
1797
1798                    plan::SubscribeOutput::WithinTimestampOrderBy { order_by }
1799                }
1800            }
1801        }
1802    };
1803
1804    let SubscribeOptionExtracted {
1805        progress, snapshot, ..
1806    } = options.try_into()?;
1807    Ok(Plan::Subscribe(SubscribePlan {
1808        from,
1809        when,
1810        up_to,
1811        with_snapshot: snapshot.unwrap_or(true),
1812        copy_to,
1813        emit_progress: progress.unwrap_or(false),
1814        output,
1815    }))
1816}
1817
1818/// Ensures each `ColumnOrder` in `order_by` references a distinct column,
1819/// returning `DuplicateKeyColumnInSubscribeEnvelope` on the first repeat.
1820fn check_distinct_key_columns(
1821    order_by: &[ColumnOrder],
1822    output_columns: &[(usize, &mz_repr::ColumnName)],
1823) -> Result<(), PlanError> {
1824    let mut seen = BTreeSet::new();
1825    for co in order_by {
1826        if !seen.insert(co.column) {
1827            return Err(PlanError::DuplicateKeyColumnInSubscribeEnvelope {
1828                column_name: output_columns[co.column].1.to_string(),
1829            });
1830        }
1831    }
1832    Ok(())
1833}
1834
1835pub fn describe_copy_from_table(
1836    scx: &StatementContext,
1837    table_name: <Aug as AstInfo>::ItemName,
1838    columns: Vec<Ident>,
1839) -> Result<StatementDesc, PlanError> {
1840    let (_, desc, _, _) = query::plan_copy_from(scx, table_name, columns)?;
1841    Ok(StatementDesc::new(Some(desc)))
1842}
1843
1844pub fn describe_copy_item(
1845    scx: &StatementContext,
1846    object_name: <Aug as AstInfo>::ItemName,
1847    columns: Vec<Ident>,
1848) -> Result<StatementDesc, PlanError> {
1849    let (_, desc, _, _) = query::plan_copy_item(scx, object_name, columns)?;
1850    Ok(StatementDesc::new(Some(desc)))
1851}
1852
1853pub fn describe_copy(
1854    scx: &StatementContext,
1855    CopyStatement {
1856        relation,
1857        direction,
1858        ..
1859    }: CopyStatement<Aug>,
1860) -> Result<StatementDesc, PlanError> {
1861    Ok(match (relation, direction) {
1862        (CopyRelation::Named { name, columns }, CopyDirection::To) => {
1863            describe_copy_item(scx, name, columns)?
1864        }
1865        (CopyRelation::Named { name, columns }, CopyDirection::From) => {
1866            describe_copy_from_table(scx, name, columns)?
1867        }
1868        (CopyRelation::Select(stmt), _) => describe_select(scx, stmt)?,
1869        (CopyRelation::Subscribe(stmt), _) => describe_subscribe(scx, stmt)?,
1870    }
1871    .with_is_copy())
1872}
1873
1874fn plan_copy_to_expr(
1875    scx: &StatementContext,
1876    select_plan: SelectPlan,
1877    desc: RelationDesc,
1878    to: &Expr<Aug>,
1879    format: CopyFormat,
1880    options: CopyOptionExtracted,
1881) -> Result<Plan, PlanError> {
1882    let conn_id = match options.aws_connection {
1883        Some(conn_id) => CatalogItemId::from(conn_id),
1884        None => sql_bail!("AWS CONNECTION is required for COPY ... TO <expr>"),
1885    };
1886    let connection = scx.get_item(&conn_id).connection()?;
1887
1888    match connection {
1889        mz_storage_types::connections::Connection::Aws(_) => {}
1890        _ => sql_bail!("only AWS CONNECTION is supported for COPY ... TO <expr>"),
1891    }
1892
1893    let format = match format {
1894        CopyFormat::Csv => {
1895            let quote = extract_byte_param_value(options.quote, "quote")?;
1896            let escape = extract_byte_param_value(options.escape, "escape")?;
1897            let delimiter = extract_byte_param_value(options.delimiter, "delimiter")?;
1898            S3SinkFormat::PgCopy(CopyFormatParams::Csv(
1899                CopyCsvFormatParams::try_new(
1900                    delimiter,
1901                    quote,
1902                    escape,
1903                    options.header,
1904                    options.null,
1905                )
1906                .map_err(|e| sql_err!("{}", e))?,
1907            ))
1908        }
1909        CopyFormat::Parquet => {
1910            // Validate that the output desc can be formatted as parquet.
1911            // COPY TO does not apply any type overrides, so pass `|_| None`.
1912            ArrowBuilder::validate_desc_for_parquet(&desc, |_| None)
1913                .map_err(|e| sql_err!("{}", e))?;
1914            S3SinkFormat::Parquet
1915        }
1916        CopyFormat::Binary => bail_unsupported!("FORMAT BINARY"),
1917        CopyFormat::Text => bail_unsupported!("FORMAT TEXT"),
1918    };
1919
1920    // Converting the to expr to a HirScalarExpr
1921    let mut to_expr = to.clone();
1922    transform_ast::transform(scx, &mut to_expr)?;
1923    let relation_type = RelationDesc::empty();
1924    let ecx = &ExprContext {
1925        qcx: &QueryContext::root(scx, QueryLifetime::OneShot),
1926        name: "COPY TO target",
1927        scope: &Scope::empty(),
1928        relation_type: relation_type.typ(),
1929        allow_aggregates: false,
1930        allow_subqueries: false,
1931        allow_parameters: false,
1932        allow_windows: false,
1933    };
1934
1935    let to = plan_expr(ecx, &to_expr)?.type_as(ecx, &SqlScalarType::String)?;
1936
1937    if options.max_file_size.as_bytes() < MIN_S3_SINK_FILE_SIZE.as_bytes() {
1938        sql_bail!(
1939            "MAX FILE SIZE cannot be less than {}",
1940            MIN_S3_SINK_FILE_SIZE
1941        );
1942    }
1943    if options.max_file_size.as_bytes() > MAX_S3_SINK_FILE_SIZE.as_bytes() {
1944        sql_bail!(
1945            "MAX FILE SIZE cannot be greater than {}",
1946            MAX_S3_SINK_FILE_SIZE
1947        );
1948    }
1949
1950    Ok(Plan::CopyTo(CopyToPlan {
1951        select_plan,
1952        desc,
1953        to,
1954        connection: connection.to_owned(),
1955        connection_id: conn_id,
1956        format,
1957        max_file_size: options.max_file_size.as_bytes(),
1958    }))
1959}
1960
1961fn plan_copy_from(
1962    scx: &StatementContext,
1963    target: &CopyTarget<Aug>,
1964    table_name: ResolvedItemName,
1965    columns: Vec<Ident>,
1966    format: Option<CopyFormat>,
1967    options: CopyOptionExtracted,
1968) -> Result<Plan, PlanError> {
1969    fn only_available_with_csv<T>(option: Option<T>, param: &str) -> Result<(), PlanError> {
1970        match option {
1971            Some(_) => sql_bail!("COPY {} available only in CSV mode", param),
1972            None => Ok(()),
1973        }
1974    }
1975
1976    let source = match target {
1977        CopyTarget::Stdin => CopyFromSource::Stdin,
1978        CopyTarget::Expr(from) => {
1979            // Converting the expr to an HirScalarExpr
1980            let mut from_expr = from.clone();
1981            transform_ast::transform(scx, &mut from_expr)?;
1982            let relation_type = RelationDesc::empty();
1983            let ecx = &ExprContext {
1984                qcx: &QueryContext::root(scx, QueryLifetime::OneShot),
1985                name: "COPY FROM target",
1986                scope: &Scope::empty(),
1987                relation_type: relation_type.typ(),
1988                allow_aggregates: false,
1989                allow_subqueries: false,
1990                allow_parameters: false,
1991                allow_windows: false,
1992            };
1993            let from = plan_expr(ecx, &from_expr)?.type_as(ecx, &SqlScalarType::String)?;
1994
1995            match options.aws_connection {
1996                Some(conn_id) => {
1997                    let conn_id = CatalogItemId::from(conn_id);
1998
1999                    // Validate the connection type is one we expect.
2000                    let connection = match scx.get_item(&conn_id).connection()? {
2001                        mz_storage_types::connections::Connection::Aws(conn) => conn,
2002                        _ => sql_bail!("only AWS CONNECTION is supported in COPY ... FROM"),
2003                    };
2004
2005                    CopyFromSource::AwsS3 {
2006                        uri: from,
2007                        connection,
2008                        connection_id: conn_id,
2009                    }
2010                }
2011                None => CopyFromSource::Url(from),
2012            }
2013        }
2014        CopyTarget::Stdout => bail_never_supported!("COPY FROM {} not supported", target),
2015    };
2016
2017    // COPY FROM a URL or S3 bucket only supports CSV and Parquet. Unlike COPY
2018    // FROM STDIN there's no sensible default format, so one must be specified
2019    // explicitly. Reject unsupported formats here in planning; the coordinator
2020    // relies on this and would otherwise soft-panic.
2021    let format = match &source {
2022        CopyFromSource::Stdin => format.unwrap_or(CopyFormat::Text),
2023        CopyFromSource::Url(_) | CopyFromSource::AwsS3 { .. } => match format {
2024            None => sql_bail!("COPY FROM <expr> requires a FORMAT option"),
2025            Some(CopyFormat::Text) => bail_unsupported!("FORMAT TEXT"),
2026            Some(CopyFormat::Binary) => bail_unsupported!("FORMAT BINARY"),
2027            Some(format @ (CopyFormat::Csv | CopyFormat::Parquet)) => format,
2028        },
2029    };
2030
2031    let params = match format {
2032        CopyFormat::Text => {
2033            only_available_with_csv(options.quote, "quote")?;
2034            only_available_with_csv(options.escape, "escape")?;
2035            only_available_with_csv(options.header, "HEADER")?;
2036            let delimiter =
2037                extract_byte_param_value(options.delimiter, "delimiter")?.unwrap_or(b'\t');
2038            let null = match options.null {
2039                Some(null) => Cow::from(null),
2040                None => Cow::from("\\N"),
2041            };
2042            CopyFormatParams::Text(CopyTextFormatParams { null, delimiter })
2043        }
2044        CopyFormat::Csv => {
2045            let quote = extract_byte_param_value(options.quote, "quote")?;
2046            let escape = extract_byte_param_value(options.escape, "escape")?;
2047            let delimiter = extract_byte_param_value(options.delimiter, "delimiter")?;
2048            CopyFormatParams::Csv(
2049                CopyCsvFormatParams::try_new(
2050                    delimiter,
2051                    quote,
2052                    escape,
2053                    options.header,
2054                    options.null,
2055                )
2056                .map_err(|e| sql_err!("{}", e))?,
2057            )
2058        }
2059        CopyFormat::Binary => bail_unsupported!("FORMAT BINARY"),
2060        CopyFormat::Parquet => CopyFormatParams::Parquet,
2061    };
2062
2063    let filter = match (options.files, options.pattern) {
2064        (Some(_), Some(_)) => bail_unsupported!("must specify one of FILES or PATTERN"),
2065        (Some(files), None) => Some(CopyFromFilter::Files(files)),
2066        (None, Some(pattern)) => Some(CopyFromFilter::Pattern(pattern)),
2067        (None, None) => None,
2068    };
2069
2070    if filter.is_some() && matches!(source, CopyFromSource::Stdin) {
2071        bail_unsupported!("COPY FROM ... WITH (FILES ...) only supported from a URL")
2072    }
2073
2074    let table_name_string = table_name.full_name_str();
2075
2076    let (id, source_desc, columns, maybe_mfp) = query::plan_copy_from(scx, table_name, columns)?;
2077
2078    let Some(mfp) = maybe_mfp else {
2079        sql_bail!("[internal error] COPY FROM ... expects an MFP to be produced");
2080    };
2081
2082    Ok(Plan::CopyFrom(CopyFromPlan {
2083        target_id: id,
2084        target_name: table_name_string,
2085        source,
2086        columns,
2087        source_desc,
2088        mfp,
2089        params,
2090        filter,
2091    }))
2092}
2093
2094fn extract_byte_param_value(v: Option<String>, param_name: &str) -> Result<Option<u8>, PlanError> {
2095    match v {
2096        Some(v) if v.len() == 1 => Ok(Some(v.as_bytes()[0])),
2097        Some(..) => sql_bail!("COPY {} must be a single one-byte character", param_name),
2098        None => Ok(None),
2099    }
2100}
2101
2102generate_extracted_config!(
2103    CopyOption,
2104    (Format, String),
2105    (Delimiter, String),
2106    (Null, String),
2107    (Escape, String),
2108    (Quote, String),
2109    (Header, bool),
2110    (AwsConnection, with_options::Object),
2111    (MaxFileSize, ByteSize, Default(ByteSize::mb(256))),
2112    (Files, Vec<String>),
2113    (Pattern, String)
2114);
2115
2116pub fn plan_copy(
2117    scx: &StatementContext,
2118    CopyStatement {
2119        relation,
2120        direction,
2121        target,
2122        options,
2123    }: CopyStatement<Aug>,
2124) -> Result<Plan, PlanError> {
2125    let options = CopyOptionExtracted::try_from(options)?;
2126    // Parse any user-provided FORMAT option. If not provided, will default to
2127    // Text for COPY TO STDOUT and COPY FROM STDIN, but will error for COPY TO <expr>.
2128    let format = options
2129        .format
2130        .as_ref()
2131        .map(|format| match format.to_lowercase().as_str() {
2132            "text" => Ok(CopyFormat::Text),
2133            "csv" => Ok(CopyFormat::Csv),
2134            "binary" => Ok(CopyFormat::Binary),
2135            "parquet" => Ok(CopyFormat::Parquet),
2136            _ => sql_bail!("unknown FORMAT: {}", format),
2137        })
2138        .transpose()?;
2139
2140    match (&direction, &target) {
2141        (CopyDirection::To, CopyTarget::Stdout) => {
2142            if options.delimiter.is_some() {
2143                sql_bail!("COPY TO does not support DELIMITER option yet");
2144            }
2145            if options.quote.is_some() {
2146                sql_bail!("COPY TO does not support QUOTE option yet");
2147            }
2148            if options.escape.is_some() {
2149                sql_bail!("COPY TO does not support ESCAPE option yet");
2150            }
2151            if options.null.is_some() {
2152                sql_bail!("COPY TO does not support NULL option yet");
2153            }
2154            // `HEADER false` is the default and already honored; only an
2155            // enabled header is unimplemented. Silently accepting it would
2156            // make clients strip the first data row as a presumed header.
2157            if options.header == Some(true) {
2158                sql_bail!("COPY TO does not support HEADER option yet");
2159            }
2160            match relation {
2161                CopyRelation::Named { .. } => sql_bail!("named with COPY TO STDOUT unsupported"),
2162                CopyRelation::Select(stmt) => Ok(plan_select(
2163                    scx,
2164                    stmt,
2165                    &Params::empty(),
2166                    Some(format.unwrap_or(CopyFormat::Text)),
2167                )?),
2168                CopyRelation::Subscribe(stmt) => Ok(plan_subscribe(
2169                    scx,
2170                    stmt,
2171                    &Params::empty(),
2172                    Some(format.unwrap_or(CopyFormat::Text)),
2173                )?),
2174            }
2175        }
2176        (CopyDirection::From, target) => match relation {
2177            CopyRelation::Named { name, columns } => {
2178                plan_copy_from(scx, target, name, columns, format, options)
2179            }
2180            _ => sql_bail!("COPY FROM {} not supported", target),
2181        },
2182        (CopyDirection::To, CopyTarget::Expr(to_expr)) => {
2183            let format = match format {
2184                Some(inner) => inner,
2185                _ => sql_bail!("COPY TO <expr> requires a FORMAT option"),
2186            };
2187
2188            let stmt = match relation {
2189                CopyRelation::Named { name, columns } => {
2190                    if !columns.is_empty() {
2191                        // TODO(mouli): Add support for this
2192                        sql_bail!(
2193                            "specifying columns for COPY <table_name> TO commands not yet supported; use COPY (SELECT...) TO ... instead"
2194                        );
2195                    }
2196                    // Generate a synthetic SELECT query that just gets the table
2197                    let query = Query {
2198                        ctes: CteBlock::empty(),
2199                        body: SetExpr::Table(name),
2200                        order_by: vec![],
2201                        limit: None,
2202                        offset: None,
2203                    };
2204                    SelectStatement { query, as_of: None }
2205                }
2206                CopyRelation::Select(stmt) => {
2207                    if !stmt.query.order_by.is_empty() {
2208                        sql_bail!("ORDER BY is not supported in SELECT query for COPY statements")
2209                    }
2210                    stmt
2211                }
2212                CopyRelation::Subscribe(_) => {
2213                    sql_bail!("COPY {} {} not supported", direction, target)
2214                }
2215            };
2216
2217            let (plan, desc) = plan_select_inner(scx, stmt, &Params::empty(), None)?;
2218            plan_copy_to_expr(scx, plan, desc, to_expr, format, options)
2219        }
2220        _ => sql_bail!("COPY {} {} not supported", direction, target),
2221    }
2222}