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        EnableUnionCancellationAfterRelationCse,
585        Option<bool>,
586        Default(None)
587    )
588);
589
590impl TryFrom<ExplainPlanOptionExtracted> for ExplainConfig {
591    type Error = PlanError;
592
593    fn try_from(mut v: ExplainPlanOptionExtracted) -> Result<Self, Self::Error> {
594        // If `WITH(raw)` is specified, ensure that the config will be as
595        // representative for the original plan as possible.
596        if v.raw {
597            v.raw_plans = true;
598            v.raw_syntax = true;
599        }
600
601        // Certain config should default to be enabled in release builds running on
602        // staging or prod (where SOFT_ASSERTIONS are turned off).
603        let enable_on_prod = !mz_ore::assert::soft_assertions_enabled();
604
605        Ok(ExplainConfig {
606            arity: v.arity.unwrap_or(enable_on_prod),
607            cardinality: v.cardinality,
608            column_names: v.column_names,
609            filter_pushdown: v.filter_pushdown.unwrap_or(enable_on_prod),
610            humanized_exprs: !v.raw_plans && (v.humanized_expressions.unwrap_or(enable_on_prod)),
611            join_impls: v.join_implementations,
612            keys: v.keys,
613            linear_chains: !v.raw_plans && v.linear_chains,
614            no_fast_path: v.no_fast_path,
615            no_notices: v.no_notices,
616            node_ids: v.node_identifiers,
617            non_negative: v.non_negative,
618            raw_plans: v.raw_plans,
619            raw_syntax: v.raw_syntax,
620            verbose_syntax: false,
621            redacted: v.redacted,
622            subtree_size: v.subtree_size,
623            equivalences: v.equivalences,
624            timing: v.timing,
625            types: v.types,
626            // The ones that are initialized with `Default::default()` are not wired up to EXPLAIN.
627            features: OptimizerFeatureOverrides {
628                enable_eager_delta_joins: v.enable_eager_delta_joins,
629                enable_new_outer_join_lowering: v.enable_new_outer_join_lowering,
630                enable_variadic_left_join_lowering: v.enable_variadic_left_join_lowering,
631                enable_letrec_fixpoint_analysis: v.enable_letrec_fixpoint_analysis,
632                enable_reduce_mfp_fusion: Default::default(),
633                enable_cardinality_estimates: Default::default(),
634                persist_fast_path_limit: Default::default(),
635                reoptimize_imported_views: v.reoptimize_imported_views,
636                enable_join_prioritize_arranged: v.enable_join_prioritize_arranged,
637                enable_projection_pushdown_after_relation_cse: v
638                    .enable_projection_pushdown_after_relation_cse,
639                enable_union_cancellation_after_relation_cse: v
640                    .enable_union_cancellation_after_relation_cse,
641                enable_less_reduce_in_eqprop: Default::default(),
642                enable_dequadratic_eqprop_map: Default::default(),
643                enable_eq_classes_withholding_errors: Default::default(),
644                enable_fast_path_plan_insights: Default::default(),
645                enable_cast_elimination: Default::default(),
646                enable_case_literal_transform: Default::default(),
647                enable_simplify_quantified_comparisons: Default::default(),
648                enable_simplify_from_less_existence: Default::default(),
649                enable_coalesce_case_transform: Default::default(),
650                enable_will_distinct_propagation: Default::default(),
651                enable_fixed_correlated_cte_lowering: v.enable_fixed_correlated_cte_lowering,
652            },
653        })
654    }
655}
656
657fn plan_explainee(
658    scx: &StatementContext,
659    explainee: Explainee<Aug>,
660    params: &Params,
661) -> Result<plan::Explainee, PlanError> {
662    use crate::plan::ExplaineeStatement;
663
664    let is_replan = matches!(
665        explainee,
666        Explainee::ReplanView(_) | Explainee::ReplanMaterializedView(_) | Explainee::ReplanIndex(_)
667    );
668
669    let explainee = match explainee {
670        Explainee::View(name) | Explainee::ReplanView(name) => {
671            let item = scx.get_item_by_resolved_name(&name)?;
672            let item_type = item.item_type();
673            if item_type != CatalogItemType::View {
674                sql_bail!("Expected {name} to be a view, not a {item_type}");
675            }
676            match is_replan {
677                true => crate::plan::Explainee::ReplanView(item.id()),
678                false => crate::plan::Explainee::View(item.id()),
679            }
680        }
681        Explainee::MaterializedView(name) | Explainee::ReplanMaterializedView(name) => {
682            let item = scx.get_item_by_resolved_name(&name)?;
683            let item_type = item.item_type();
684            if item_type != CatalogItemType::MaterializedView {
685                sql_bail!("Expected {name} to be a materialized view, not a {item_type}");
686            }
687            match is_replan {
688                true => crate::plan::Explainee::ReplanMaterializedView(item.id()),
689                false => crate::plan::Explainee::MaterializedView(item.id()),
690            }
691        }
692        Explainee::Index(name) | Explainee::ReplanIndex(name) => {
693            let item = scx.get_item_by_resolved_name(&name)?;
694            let item_type = item.item_type();
695            if item_type != CatalogItemType::Index {
696                sql_bail!("Expected {name} to be an index, not a {item_type}");
697            }
698            match is_replan {
699                true => crate::plan::Explainee::ReplanIndex(item.id()),
700                false => crate::plan::Explainee::Index(item.id()),
701            }
702        }
703        Explainee::Select(select, broken) => {
704            let (plan, desc) = plan_select_inner(scx, *select, params, None)?;
705            crate::plan::Explainee::Statement(ExplaineeStatement::Select { broken, plan, desc })
706        }
707        Explainee::CreateView(mut stmt, broken) => {
708            if stmt.if_exists != IfExistsBehavior::Skip {
709                // If we don't force this parameter to Skip planning will
710                // fail for names that already exist in the catalog. This
711                // can happen even in `Replace` mode if the existing item
712                // has dependencies.
713                stmt.if_exists = IfExistsBehavior::Skip;
714            } else {
715                sql_bail!(
716                    "Cannot EXPLAIN a CREATE VIEW that explictly sets IF NOT EXISTS \
717                     (the behavior is implied within the scope of an enclosing EXPLAIN)"
718                );
719            }
720
721            let Plan::CreateView(plan) = ddl::plan_create_view(scx, *stmt)? else {
722                sql_bail!("expected CreateViewPlan plan");
723            };
724
725            crate::plan::Explainee::Statement(ExplaineeStatement::CreateView { broken, plan })
726        }
727        Explainee::CreateMaterializedView(mut stmt, broken) => {
728            if stmt.if_exists != IfExistsBehavior::Skip {
729                // If we don't force this parameter to Skip planning will
730                // fail for names that already exist in the catalog. This
731                // can happen even in `Replace` mode if the existing item
732                // has dependencies.
733                stmt.if_exists = IfExistsBehavior::Skip;
734            } else {
735                sql_bail!(
736                    "Cannot EXPLAIN a CREATE MATERIALIZED VIEW that explictly sets IF NOT EXISTS \
737                     (the behavior is implied within the scope of an enclosing EXPLAIN)"
738                );
739            }
740
741            let Plan::CreateMaterializedView(plan) =
742                ddl::plan_create_materialized_view(scx, *stmt)?
743            else {
744                sql_bail!("expected CreateMaterializedViewPlan plan");
745            };
746
747            crate::plan::Explainee::Statement(ExplaineeStatement::CreateMaterializedView {
748                broken,
749                plan,
750            })
751        }
752        Explainee::CreateIndex(mut stmt, broken) => {
753            if !stmt.if_not_exists {
754                // If we don't force this parameter to true planning will
755                // fail for index items that already exist in the catalog.
756                stmt.if_not_exists = true;
757            } else {
758                sql_bail!(
759                    "Cannot EXPLAIN a CREATE INDEX that explictly sets IF NOT EXISTS \
760                     (the behavior is implied within the scope of an enclosing EXPLAIN)"
761                );
762            }
763
764            let Plan::CreateIndex(plan) = ddl::plan_create_index(scx, *stmt)? else {
765                sql_bail!("expected CreateIndexPlan plan");
766            };
767
768            crate::plan::Explainee::Statement(ExplaineeStatement::CreateIndex { broken, plan })
769        }
770        Explainee::Subscribe(stmt, broken) => {
771            let Plan::Subscribe(plan) = plan_subscribe(scx, *stmt, params, None)? else {
772                sql_bail!("expected SubscribePlan");
773            };
774            crate::plan::Explainee::Statement(ExplaineeStatement::Subscribe { broken, plan })
775        }
776    };
777
778    Ok(explainee)
779}
780
781pub fn plan_explain_plan(
782    scx: &StatementContext,
783    explain: ExplainPlanStatement<Aug>,
784    params: &Params,
785) -> Result<Plan, PlanError> {
786    let (format, verbose_syntax) = match explain.format() {
787        mz_sql_parser::ast::ExplainFormat::Text => (ExplainFormat::Text, false),
788        mz_sql_parser::ast::ExplainFormat::VerboseText => (ExplainFormat::Text, true),
789        mz_sql_parser::ast::ExplainFormat::Json => (ExplainFormat::Json, false),
790        mz_sql_parser::ast::ExplainFormat::Dot => (ExplainFormat::Dot, false),
791    };
792    let stage = explain.stage();
793
794    // Plan ExplainConfig.
795    let mut config = {
796        let mut with_options = ExplainPlanOptionExtracted::try_from(explain.with_options)?;
797
798        if !scx.catalog.system_vars().persist_stats_filter_enabled() {
799            // If filtering is disabled, explain plans should not include pushdown info.
800            with_options.filter_pushdown = Some(false);
801        }
802
803        ExplainConfig::try_from(with_options)?
804    };
805    config.verbose_syntax = verbose_syntax;
806
807    let explainee = plan_explainee(scx, explain.explainee, params)?;
808
809    Ok(Plan::ExplainPlan(ExplainPlanPlan {
810        stage,
811        format,
812        config,
813        explainee,
814    }))
815}
816
817pub fn plan_explain_schema(
818    scx: &StatementContext,
819    explain_schema: ExplainSinkSchemaStatement<Aug>,
820) -> Result<Plan, PlanError> {
821    let ExplainSinkSchemaStatement {
822        schema_for,
823        // Parser limits to JSON.
824        format: _,
825        mut statement,
826    } = explain_schema;
827
828    // Force the sink's name to one that's guaranteed not to exist, by virtue of
829    // being a non-existent item in a schema under the system's control, so that
830    // `plan_create_sink` doesn't complain about the name already existing.
831    statement.name = Some(UnresolvedItemName::qualified(&[
832        ident!("mz_catalog"),
833        ident!("mz_explain_schema"),
834    ]));
835
836    crate::pure::purify_create_sink_avro_doc_on_options(
837        scx.catalog,
838        *statement.from.item_id(),
839        &mut statement.format,
840    )?;
841
842    match ddl::plan_create_sink(scx, statement)? {
843        Plan::CreateSink(CreateSinkPlan { sink, .. }) => match sink.connection {
844            StorageSinkConnection::Kafka(KafkaSinkConnection {
845                format:
846                    KafkaSinkFormat {
847                        key_format,
848                        value_format:
849                            KafkaSinkFormatType::Avro {
850                                schema: value_schema,
851                                ..
852                            },
853                        ..
854                    },
855                ..
856            }) => {
857                let schema = match schema_for {
858                    ExplainSinkSchemaFor::Key => key_format
859                        .and_then(|f| match f {
860                            KafkaSinkFormatType::Avro { schema, .. } => Some(schema),
861                            _ => None,
862                        })
863                        .ok_or_else(|| sql_err!("CREATE SINK does not have a key"))?,
864                    ExplainSinkSchemaFor::Value => value_schema,
865                };
866
867                Ok(Plan::ExplainSinkSchema(ExplainSinkSchemaPlan {
868                    sink_from: sink.from,
869                    json_schema: schema,
870                }))
871            }
872            _ => bail_unsupported!(
873                "EXPLAIN SCHEMA is only available for Kafka sinks with Avro schemas"
874            ),
875        },
876        _ => bail_internal!("plan_sink did not produce a CreateSink plan"),
877    }
878}
879
880pub fn plan_explain_pushdown(
881    scx: &StatementContext,
882    statement: ExplainPushdownStatement<Aug>,
883    params: &Params,
884) -> Result<Plan, PlanError> {
885    scx.require_feature_flag(&vars::ENABLE_EXPLAIN_PUSHDOWN)?;
886    let explainee = plan_explainee(scx, statement.explainee, params)?;
887    Ok(Plan::ExplainPushdown(ExplainPushdownPlan { explainee }))
888}
889
890pub fn plan_explain_analyze_object(
891    scx: &StatementContext,
892    statement: ExplainAnalyzeObjectStatement<Aug>,
893    params: &Params,
894) -> Result<Plan, PlanError> {
895    let explainee_name = statement
896        .explainee
897        .name()
898        .ok_or_else(|| sql_err!("EXPLAIN ANALYZE on anonymous dataflows",))?
899        .full_name_str();
900    let explainee = plan_explainee(scx, statement.explainee, params)?;
901
902    let check_ownership = |item_id: &CatalogItemId, item_type: &str| -> Result<(), PlanError> {
903        if scx.catalog.restrict_to_user_objects() {
904            let item = scx.catalog.get_item(item_id);
905            if item.owner_id() != *scx.catalog.active_role_id() {
906                let full_name = scx.catalog.resolve_full_name(item.name());
907                return Err(sql_err!("must be owner of {item_type} {full_name}"));
908            }
909        }
910        Ok(())
911    };
912    match &explainee {
913        plan::Explainee::Index(item_id) => check_ownership(item_id, "INDEX")?,
914        plan::Explainee::MaterializedView(item_id) => {
915            check_ownership(item_id, "MATERIALIZED VIEW")?
916        }
917        _ => return Err(sql_err!("EXPLAIN ANALYZE queries for this explainee type",)),
918    };
919
920    // generate SQL query
921
922    /* WITH {CTEs}
923       SELECT REPEAT(' ', nesting * 2) || operator AS operator
924             {columns}
925        FROM      mz_introspection.mz_lir_mapping mlm
926             JOIN {from} USING (lir_id)
927             JOIN mz_introspection.mz_mappable_objects mo
928               ON (mlm.global_id = mo.global_id)
929       WHERE     mo.name = {escaped explainee_name}
930             AND {predicates}
931       ORDER BY lir_id DESC
932    */
933    let mut ctes = Vec::with_capacity(4); // max 2 per ExplainAnalyzeComputationProperty
934    let mut columns = vec!["REPEAT(' ', nesting * 2) || operator AS operator"];
935    let mut from = vec!["mz_introspection.mz_lir_mapping mlm"];
936    let mut predicates = vec![format!(
937        "mo.name = {}",
938        escaped_string_literal(&explainee_name)
939    )];
940    let mut order_by = vec!["mlm.lir_id DESC"];
941
942    match statement.properties {
943        ExplainAnalyzeProperty::Computation(ExplainAnalyzeComputationProperties {
944            properties,
945            skew,
946        }) => {
947            let mut worker_id = None;
948            let mut seen_properties = BTreeSet::new();
949            for property in properties {
950                // handle each property only once (belt and suspenders)
951                if !seen_properties.insert(property) {
952                    continue;
953                }
954
955                match property {
956                    ExplainAnalyzeComputationProperty::Memory => {
957                        ctes.push((
958                            "summary_memory",
959                            r#"
960  SELECT mlm.global_id AS global_id,
961         mlm.lir_id AS lir_id,
962         SUM(mas.size) AS total_memory,
963         SUM(mas.records) AS total_records,
964         CASE WHEN COUNT(DISTINCT mas.worker_id) <> 0 THEN SUM(mas.size) / COUNT(DISTINCT mas.worker_id) ELSE NULL END AS avg_memory,
965         CASE WHEN COUNT(DISTINCT mas.worker_id) <> 0 THEN SUM(mas.records) / COUNT(DISTINCT mas.worker_id) ELSE NULL END AS avg_records
966    FROM            mz_introspection.mz_lir_mapping mlm
967         CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
968               JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
969                 ON (mas.operator_id = valid_id)
970GROUP BY mlm.global_id, mlm.lir_id"#,
971                        ));
972                        from.push("LEFT JOIN summary_memory sm USING (global_id, lir_id)");
973
974                        if skew {
975                            ctes.push((
976                                "per_worker_memory",
977                                r#"
978  SELECT mlm.global_id AS global_id,
979         mlm.lir_id AS lir_id,
980         mas.worker_id AS worker_id,
981         SUM(mas.size) AS worker_memory,
982         SUM(mas.records) AS worker_records
983    FROM            mz_introspection.mz_lir_mapping mlm
984         CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
985               JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
986                 ON (mas.operator_id = valid_id)
987GROUP BY mlm.global_id, mlm.lir_id, mas.worker_id"#,
988                            ));
989                            from.push("LEFT JOIN per_worker_memory pwm USING (global_id, lir_id)");
990
991                            if let Some(worker_id) = worker_id {
992                                predicates.push(format!(
993                                    "(pwm.worker_id = {worker_id} OR pwm.worker_id IS NULL OR {worker_id} IS NULL)"
994                                ));
995                            } else {
996                                worker_id = Some("pwm.worker_id");
997                                columns.push("pwm.worker_id AS worker_id");
998                                order_by.push("worker_id");
999                            }
1000
1001                            columns.extend([
1002                                "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",
1003                                "pg_size_pretty(pwm.worker_memory) AS worker_memory",
1004                                "pg_size_pretty(sm.avg_memory) AS avg_memory",
1005                                "pg_size_pretty(sm.total_memory) AS total_memory",
1006                                "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",
1007                                "pwm.worker_records AS worker_records",
1008                                "sm.avg_records AS avg_records",
1009                                "sm.total_records AS total_records",
1010                            ]);
1011                        } else {
1012                            columns.extend([
1013                                "pg_size_pretty(sm.total_memory) AS total_memory",
1014                                "sm.total_records AS total_records",
1015                            ]);
1016                        }
1017                    }
1018                    ExplainAnalyzeComputationProperty::Cpu => {
1019                        ctes.push((
1020                            "summary_cpu",
1021                            r#"
1022  SELECT mlm.global_id AS global_id,
1023         mlm.lir_id AS lir_id,
1024         SUM(mse.elapsed_ns) AS total_ns,
1025         CASE WHEN COUNT(DISTINCT mse.worker_id) <> 0 THEN SUM(mse.elapsed_ns) / COUNT(DISTINCT mse.worker_id) ELSE NULL END AS avg_ns
1026    FROM            mz_introspection.mz_lir_mapping mlm
1027         CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1028               JOIN mz_introspection.mz_scheduling_elapsed_per_worker mse
1029                 ON (mse.id = valid_id)
1030GROUP BY mlm.global_id, mlm.lir_id"#,
1031                        ));
1032                        from.push("LEFT JOIN summary_cpu sc USING (global_id, lir_id)");
1033
1034                        if skew {
1035                            ctes.push((
1036                                "per_worker_cpu",
1037                                r#"
1038  SELECT mlm.global_id AS global_id,
1039         mlm.lir_id AS lir_id,
1040         mse.worker_id AS worker_id,
1041         SUM(mse.elapsed_ns) AS worker_ns
1042    FROM            mz_introspection.mz_lir_mapping mlm
1043         CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1044               JOIN mz_introspection.mz_scheduling_elapsed_per_worker mse
1045                 ON (mse.id = valid_id)
1046GROUP BY mlm.global_id, mlm.lir_id, mse.worker_id"#,
1047                            ));
1048                            from.push("LEFT JOIN per_worker_cpu pwc USING (global_id, lir_id)");
1049
1050                            if let Some(worker_id) = worker_id {
1051                                predicates.push(format!(
1052                                    "(pwc.worker_id = {worker_id} OR pwc.worker_id IS NULL OR {worker_id} IS NULL)"
1053                                ));
1054                            } else {
1055                                worker_id = Some("pwc.worker_id");
1056                                columns.push("pwc.worker_id AS worker_id");
1057                                order_by.push("worker_id");
1058                            }
1059
1060                            columns.extend([
1061                                "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",
1062                                "pwc.worker_ns / 1000 * '1 microsecond'::INTERVAL AS worker_elapsed",
1063                                "sc.avg_ns / 1000 * '1 microsecond'::INTERVAL AS avg_elapsed",
1064                            ]);
1065                        }
1066                        columns.push(
1067                            "sc.total_ns / 1000 * '1 microsecond'::INTERVAL AS total_elapsed",
1068                        );
1069                    }
1070                }
1071            }
1072        }
1073        ExplainAnalyzeProperty::Hints => {
1074            columns.extend([
1075                "megsa.levels AS levels",
1076                "megsa.to_cut AS to_cut",
1077                "megsa.hint AS hint",
1078                "pg_size_pretty(megsa.savings) AS savings",
1079            ]);
1080            from.extend(["JOIN mz_introspection.mz_dataflow_global_ids mdgi ON (mlm.global_id = mdgi.global_id)",
1081            "LEFT JOIN (generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id JOIN \
1082             mz_introspection.mz_expected_group_size_advice megsa ON (megsa.region_id = valid_id)) ON (megsa.dataflow_id = mdgi.id)"]);
1083        }
1084    }
1085
1086    from.push("JOIN mz_introspection.mz_mappable_objects mo ON (mlm.global_id = mo.global_id)");
1087
1088    let ctes = if !ctes.is_empty() {
1089        format!(
1090            "WITH {}",
1091            separated(
1092                ",\n",
1093                ctes.iter()
1094                    .map(|(name, defn)| format!("{name} AS ({defn})"))
1095            )
1096        )
1097    } else {
1098        String::new()
1099    };
1100    let columns = separated(", ", columns);
1101    let from = separated(" ", from);
1102    let predicates = separated(" AND ", predicates);
1103    let order_by = separated(", ", order_by);
1104    let query = format!(
1105        r#"{ctes}
1106SELECT {columns}
1107FROM {from}
1108WHERE {predicates}
1109ORDER BY {order_by}"#
1110    );
1111
1112    if statement.as_sql {
1113        let rows = vec![Row::pack_slice(&[Datum::String(
1114            &mz_sql_pretty::pretty_str_simple(&query, 80).map_err(|e| {
1115                PlanError::Unstructured(format!("internal error parsing our own SQL: {e}"))
1116            })?,
1117        )])];
1118        let typ = SqlRelationType::new(vec![SqlScalarType::String.nullable(false)]);
1119
1120        Ok(Plan::Select(SelectPlan::immediate(rows, typ)))
1121    } else {
1122        let (show_select, resolved_ids) = ShowSelect::new_from_bare_query(scx, query)?;
1123        scx.record_sql_impl_ids(&resolved_ids);
1124        show_select.plan()
1125    }
1126}
1127
1128pub fn plan_explain_analyze_cluster(
1129    scx: &StatementContext,
1130    statement: ExplainAnalyzeClusterStatement,
1131    _params: &Params,
1132) -> Result<Plan, PlanError> {
1133    // object string
1134    // worker_id uint64        (if           skew)
1135    // memory_ratio numeric    (if memory && skew)
1136    // worker_memory string    (if memory && skew)
1137    // avg_memory string       (if memory && skew)
1138    // total_memory string     (if memory)
1139    // records_ratio numeric   (if memory && skew)
1140    // worker_records          (if memory && skew)
1141    // avg_records numeric     (if memory && skew)
1142    // total_records numeric   (if memory)
1143    // cpu_ratio numeric       (if cpu    && skew)
1144    // worker_elapsed interval (if cpu    && skew)
1145    // avg_elapsed interval    (if cpu    && skew)
1146    // total_elapsed interval  (if cpu)
1147
1148    /* WITH {CTEs}
1149       SELECT mo.name AS object
1150             {columns}
1151        FROM mz_introspection.mz_mappable_objects mo
1152             {from}
1153       WHERE {predicates}
1154       ORDER BY {order_by}, mo.name DESC
1155    */
1156    let mut ctes = Vec::with_capacity(4); // max 2 per ExplainAnalyzeComputationProperty
1157    let mut columns = vec!["mo.name AS object", "mo.global_id AS global_id"];
1158    let mut from = vec!["mz_introspection.mz_mappable_objects mo"];
1159    let mut predicates = vec![];
1160    let mut order_by = vec![];
1161
1162    let ExplainAnalyzeComputationProperties { properties, skew } = statement.properties;
1163    let mut worker_id = None;
1164    let mut seen_properties = BTreeSet::new();
1165    for property in properties {
1166        // handle each property only once (belt and suspenders)
1167        if !seen_properties.insert(property) {
1168            continue;
1169        }
1170
1171        match property {
1172            ExplainAnalyzeComputationProperty::Memory => {
1173                if skew {
1174                    let mut set_worker_id = false;
1175                    if let Some(worker_id) = worker_id {
1176                        // join condition if we're showing skew for more than one property
1177                        predicates.push(format!(
1178                            "(om.worker_id = {worker_id} OR om.worker_id IS NULL OR {worker_id} IS NULL)"
1179                        ));
1180                    } else {
1181                        worker_id = Some("om.worker_id");
1182                        columns.push("om.worker_id AS worker_id");
1183                        set_worker_id = true; // we'll add ourselves to `order_by` later
1184                    };
1185
1186                    // computes the average memory per LIR operator (for per operator ratios)
1187                    ctes.push((
1188                    "per_operator_memory_summary",
1189                    r#"
1190SELECT mlm.global_id AS global_id,
1191       mlm.lir_id AS lir_id,
1192       SUM(mas.size) AS total_memory,
1193       SUM(mas.records) AS total_records,
1194       CASE WHEN COUNT(DISTINCT mas.worker_id) <> 0 THEN SUM(mas.size) / COUNT(DISTINCT mas.worker_id) ELSE NULL END AS avg_memory,
1195       CASE WHEN COUNT(DISTINCT mas.worker_id) <> 0 THEN SUM(mas.records) / COUNT(DISTINCT mas.worker_id) ELSE NULL END AS avg_records
1196FROM        mz_introspection.mz_lir_mapping mlm
1197 CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1198       JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
1199         ON (mas.operator_id = valid_id)
1200GROUP BY mlm.global_id, mlm.lir_id"#,
1201                ));
1202
1203                    // computes the memory per worker in a per operator way
1204                    ctes.push((
1205                    "per_operator_memory_per_worker",
1206                    r#"
1207SELECT mlm.global_id AS global_id,
1208       mlm.lir_id AS lir_id,
1209       mas.worker_id AS worker_id,
1210       SUM(mas.size) AS worker_memory,
1211       SUM(mas.records) AS worker_records
1212FROM        mz_introspection.mz_lir_mapping mlm
1213 CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1214       JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
1215         ON (mas.operator_id = valid_id)
1216GROUP BY mlm.global_id, mlm.lir_id, mas.worker_id"#,
1217                    ));
1218
1219                    // computes memory ratios per worker per operator
1220                    ctes.push((
1221                    "per_operator_memory_ratios",
1222                    r#"
1223SELECT pompw.global_id AS global_id,
1224       pompw.lir_id AS lir_id,
1225       pompw.worker_id AS worker_id,
1226       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,
1227       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
1228  FROM      per_operator_memory_per_worker pompw
1229       JOIN per_operator_memory_summary poms
1230         USING (global_id, lir_id)
1231"#,
1232                    ));
1233
1234                    // summarizes each object, per worker
1235                    ctes.push((
1236                        "object_memory",
1237                        r#"
1238SELECT pompw.global_id AS global_id,
1239       pompw.worker_id AS worker_id,
1240       MAX(pomr.memory_ratio) AS max_operator_memory_ratio,
1241       MAX(pomr.records_ratio) AS max_operator_records_ratio,
1242       SUM(pompw.worker_memory) AS worker_memory,
1243       SUM(pompw.worker_records) AS worker_records
1244FROM        per_operator_memory_per_worker pompw
1245     JOIN   per_operator_memory_ratios pomr
1246     USING (global_id, worker_id, lir_id)
1247GROUP BY pompw.global_id, pompw.worker_id
1248"#,
1249                    ));
1250
1251                    // summarizes each worker
1252                    ctes.push(("object_average_memory", r#"
1253SELECT om.global_id AS global_id,
1254       SUM(om.worker_memory) AS total_memory,
1255       CASE WHEN COUNT(DISTINCT om.worker_id) <> 0 THEN SUM(om.worker_memory) / COUNT(DISTINCT om.worker_id) ELSE NULL END AS avg_memory,
1256       SUM(om.worker_records) AS total_records,
1257       CASE WHEN COUNT(DISTINCT om.worker_id) <> 0 THEN SUM(om.worker_records) / COUNT(DISTINCT om.worker_id) ELSE NULL END AS avg_records
1258  FROM object_memory om
1259GROUP BY om.global_id"#));
1260
1261                    from.push("LEFT JOIN object_memory om USING (global_id)");
1262                    from.push("LEFT JOIN object_average_memory oam USING (global_id)");
1263
1264                    columns.extend([
1265                        "om.max_operator_memory_ratio AS max_operator_memory_ratio",
1266                        "pg_size_pretty(om.worker_memory) AS worker_memory",
1267                        "pg_size_pretty(oam.avg_memory) AS avg_memory",
1268                        "pg_size_pretty(oam.total_memory) AS total_memory",
1269                        "om.max_operator_records_ratio AS max_operator_records_ratio",
1270                        "om.worker_records AS worker_records",
1271                        "oam.avg_records AS avg_records",
1272                        "oam.total_records AS total_records",
1273                    ]);
1274
1275                    order_by.extend([
1276                        "max_operator_memory_ratio DESC NULLS LAST",
1277                        "max_operator_records_ratio DESC NULLS LAST",
1278                        "om.worker_memory DESC NULLS LAST",
1279                        "worker_records DESC NULLS LAST",
1280                    ]);
1281
1282                    if set_worker_id {
1283                        order_by.push("worker_id");
1284                    }
1285                } else {
1286                    // no skew, so just compute totals
1287                    ctes.push((
1288                        "per_operator_memory_totals",
1289                        r#"
1290    SELECT mlm.global_id AS global_id,
1291           mlm.lir_id AS lir_id,
1292           SUM(mas.size) AS total_memory,
1293           SUM(mas.records) AS total_records
1294    FROM        mz_introspection.mz_lir_mapping mlm
1295     CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1296           JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
1297             ON (mas.operator_id = valid_id)
1298    GROUP BY mlm.global_id, mlm.lir_id"#,
1299                    ));
1300
1301                    ctes.push((
1302                        "object_memory_totals",
1303                        r#"
1304SELECT pomt.global_id AS global_id,
1305       SUM(pomt.total_memory) AS total_memory,
1306       SUM(pomt.total_records) AS total_records
1307FROM per_operator_memory_totals pomt
1308GROUP BY pomt.global_id
1309"#,
1310                    ));
1311
1312                    from.push("LEFT JOIN object_memory_totals omt USING (global_id)");
1313                    columns.extend([
1314                        "pg_size_pretty(omt.total_memory) AS total_memory",
1315                        "omt.total_records AS total_records",
1316                    ]);
1317                    order_by.extend([
1318                        "omt.total_memory DESC NULLS LAST",
1319                        "total_records DESC NULLS LAST",
1320                    ]);
1321                }
1322            }
1323            ExplainAnalyzeComputationProperty::Cpu => {
1324                if skew {
1325                    let mut set_worker_id = false;
1326                    if let Some(worker_id) = worker_id {
1327                        // join condition if we're showing skew for more than one property
1328                        predicates.push(format!(
1329                            "(oc.worker_id = {worker_id} OR oc.worker_id IS NULL OR {worker_id} IS NULL)"
1330                        ));
1331                    } else {
1332                        worker_id = Some("oc.worker_id");
1333                        columns.push("oc.worker_id AS worker_id");
1334                        set_worker_id = true; // we'll add ourselves to `order_by` later
1335                    };
1336
1337                    // computes the average memory per LIR operator (for per operator ratios)
1338                    ctes.push((
1339    "per_operator_cpu_summary",
1340    r#"
1341SELECT mlm.global_id AS global_id,
1342       mlm.lir_id AS lir_id,
1343       SUM(mse.elapsed_ns) AS total_ns,
1344       CASE WHEN COUNT(DISTINCT mse.worker_id) <> 0 THEN SUM(mse.elapsed_ns) / COUNT(DISTINCT mse.worker_id) ELSE NULL END AS avg_ns
1345FROM       mz_introspection.mz_lir_mapping mlm
1346CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1347      JOIN mz_introspection.mz_scheduling_elapsed_per_worker mse
1348        ON (mse.id = valid_id)
1349GROUP BY mlm.global_id, mlm.lir_id"#,
1350));
1351
1352                    // computes the CPU per worker in a per operator way
1353                    ctes.push((
1354                        "per_operator_cpu_per_worker",
1355                        r#"
1356SELECT mlm.global_id AS global_id,
1357       mlm.lir_id AS lir_id,
1358       mse.worker_id AS worker_id,
1359       SUM(mse.elapsed_ns) AS worker_ns
1360FROM       mz_introspection.mz_lir_mapping mlm
1361CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1362      JOIN mz_introspection.mz_scheduling_elapsed_per_worker mse
1363        ON (mse.id = valid_id)
1364GROUP BY mlm.global_id, mlm.lir_id, mse.worker_id"#,
1365                    ));
1366
1367                    // computes CPU ratios per worker per operator
1368                    ctes.push((
1369                        "per_operator_cpu_ratios",
1370                        r#"
1371SELECT pocpw.global_id AS global_id,
1372       pocpw.lir_id AS lir_id,
1373       pocpw.worker_id AS worker_id,
1374       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
1375FROM      per_operator_cpu_per_worker pocpw
1376     JOIN per_operator_cpu_summary pocs
1377     USING (global_id, lir_id)
1378"#,
1379                    ));
1380
1381                    // summarizes each object, per worker
1382                    ctes.push((
1383                        "object_cpu",
1384                        r#"
1385SELECT pocpw.global_id AS global_id,
1386       pocpw.worker_id AS worker_id,
1387       MAX(pomr.cpu_ratio) AS max_operator_cpu_ratio,
1388       SUM(pocpw.worker_ns) AS worker_ns
1389FROM      per_operator_cpu_per_worker pocpw
1390     JOIN per_operator_cpu_ratios pomr
1391     USING (global_id, worker_id, lir_id)
1392GROUP BY pocpw.global_id, pocpw.worker_id
1393"#,
1394                    ));
1395
1396                    // summarizes each worker
1397                    ctes.push((
1398                        "object_average_cpu",
1399                        r#"
1400SELECT oc.global_id AS global_id,
1401       SUM(oc.worker_ns) AS total_ns,
1402       CASE WHEN COUNT(DISTINCT oc.worker_id) <> 0 THEN SUM(oc.worker_ns) / COUNT(DISTINCT oc.worker_id) ELSE NULL END AS avg_ns
1403  FROM object_cpu oc
1404GROUP BY oc.global_id"#,));
1405
1406                    from.push("LEFT JOIN object_cpu oc USING (global_id)");
1407                    from.push("LEFT JOIN object_average_cpu oac USING (global_id)");
1408
1409                    columns.extend([
1410                        "oc.max_operator_cpu_ratio AS max_operator_cpu_ratio",
1411                        "oc.worker_ns / 1000 * '1 microsecond'::interval AS worker_elapsed",
1412                        "oac.avg_ns / 1000 * '1 microsecond'::interval AS avg_elapsed",
1413                        "oac.total_ns / 1000 * '1 microsecond'::interval AS total_elapsed",
1414                    ]);
1415
1416                    order_by.extend([
1417                        "max_operator_cpu_ratio DESC NULLS LAST",
1418                        "worker_elapsed DESC NULLS LAST",
1419                    ]);
1420
1421                    if set_worker_id {
1422                        order_by.push("worker_id");
1423                    }
1424                } else {
1425                    // no skew, so just compute totals
1426                    ctes.push((
1427                        "per_operator_cpu_totals",
1428                        r#"
1429    SELECT mlm.global_id AS global_id,
1430           mlm.lir_id AS lir_id,
1431           SUM(mse.elapsed_ns) AS total_ns
1432    FROM        mz_introspection.mz_lir_mapping mlm
1433     CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1434           JOIN mz_introspection.mz_scheduling_elapsed_per_worker mse
1435             ON (mse.id = valid_id)
1436    GROUP BY mlm.global_id, mlm.lir_id"#,
1437                    ));
1438
1439                    ctes.push((
1440                        "object_cpu_totals",
1441                        r#"
1442SELECT poct.global_id AS global_id,
1443       SUM(poct.total_ns) AS total_ns
1444FROM per_operator_cpu_totals poct
1445GROUP BY poct.global_id
1446"#,
1447                    ));
1448
1449                    from.push("LEFT JOIN object_cpu_totals oct USING (global_id)");
1450                    columns
1451                        .push("oct.total_ns / 1000 * '1 microsecond'::interval AS total_elapsed");
1452                    order_by.extend(["total_elapsed DESC NULLS LAST"]);
1453                }
1454            }
1455        }
1456    }
1457
1458    // generate SQL query text
1459    let ctes = if !ctes.is_empty() {
1460        format!(
1461            "WITH {}",
1462            separated(
1463                ",\n",
1464                ctes.iter()
1465                    .map(|(name, defn)| format!("{name} AS ({defn})"))
1466            )
1467        )
1468    } else {
1469        String::new()
1470    };
1471    let columns = separated(", ", columns);
1472    let from = separated(" ", from);
1473    let predicates = if !predicates.is_empty() {
1474        format!("WHERE {}", separated(" AND ", predicates))
1475    } else {
1476        String::new()
1477    };
1478    // add mo.name last, to break ties only
1479    order_by.push("mo.name DESC");
1480    let order_by = separated(", ", order_by);
1481    let query = format!(
1482        r#"{ctes}
1483SELECT {columns}
1484FROM {from}
1485{predicates}
1486ORDER BY {order_by}"#
1487    );
1488
1489    if statement.as_sql {
1490        let rows = vec![Row::pack_slice(&[Datum::String(
1491            &mz_sql_pretty::pretty_str_simple(&query, 80).map_err(|e| {
1492                PlanError::Unstructured(format!("internal error parsing our own SQL: {e}"))
1493            })?,
1494        )])];
1495        let typ = SqlRelationType::new(vec![SqlScalarType::String.nullable(false)]);
1496
1497        Ok(Plan::Select(SelectPlan::immediate(rows, typ)))
1498    } else {
1499        let (show_select, resolved_ids) = ShowSelect::new_from_bare_query(scx, query)?;
1500        scx.record_sql_impl_ids(&resolved_ids);
1501        show_select.plan()
1502    }
1503}
1504
1505pub fn plan_explain_timestamp(
1506    scx: &StatementContext,
1507    explain: ExplainTimestampStatement<Aug>,
1508) -> Result<Plan, PlanError> {
1509    let (format, _verbose_syntax) = match explain.format() {
1510        mz_sql_parser::ast::ExplainFormat::Text => (ExplainFormat::Text, false),
1511        mz_sql_parser::ast::ExplainFormat::VerboseText => (ExplainFormat::Text, true),
1512        mz_sql_parser::ast::ExplainFormat::Json => (ExplainFormat::Json, false),
1513        mz_sql_parser::ast::ExplainFormat::Dot => (ExplainFormat::Dot, false),
1514    };
1515
1516    let raw_plan = {
1517        let query::PlannedRootQuery {
1518            expr: raw_plan,
1519            desc: _,
1520            finishing: _,
1521            scope: _,
1522        } = query::plan_root_query(scx, explain.select.query, QueryLifetime::OneShot)?;
1523        if raw_plan.contains_parameters()? {
1524            return Err(PlanError::ParameterNotAllowed(
1525                "EXPLAIN TIMESTAMP".to_string(),
1526            ));
1527        }
1528
1529        raw_plan
1530    };
1531    let when = query::plan_as_of(scx, explain.select.as_of)?;
1532
1533    Ok(Plan::ExplainTimestamp(ExplainTimestampPlan {
1534        format,
1535        raw_plan,
1536        when,
1537    }))
1538}
1539
1540generate_extracted_config!(SubscribeOption, (Snapshot, bool), (Progress, bool));
1541
1542pub fn describe_subscribe(
1543    scx: &StatementContext,
1544    stmt: SubscribeStatement<Aug>,
1545) -> Result<StatementDesc, PlanError> {
1546    let relation_desc = match stmt.relation {
1547        SubscribeRelation::Name(name) => {
1548            let item = scx.get_item_by_resolved_name(&name)?;
1549            match item.relation_desc() {
1550                Some(desc) => desc.into_owned(),
1551                None => sql_bail!(
1552                    "'{}' cannot be subscribed to because it is a {}",
1553                    name.full_name_str(),
1554                    item.item_type(),
1555                ),
1556            }
1557        }
1558        SubscribeRelation::Query(query) => {
1559            let query::PlannedRootQuery { desc, .. } =
1560                query::plan_root_query(scx, query, QueryLifetime::Subscribe)?;
1561            desc
1562        }
1563    };
1564    let SubscribeOptionExtracted { progress, .. } = stmt.options.try_into()?;
1565    let progress = progress.unwrap_or(false);
1566    let mut desc = RelationDesc::builder().with_column(
1567        "mz_timestamp",
1568        SqlScalarType::Numeric {
1569            max_scale: Some(NumericMaxScale::ZERO),
1570        }
1571        .nullable(false),
1572    );
1573    if progress {
1574        desc = desc.with_column("mz_progressed", SqlScalarType::Bool.nullable(false));
1575    }
1576
1577    let debezium = matches!(stmt.output, SubscribeOutput::EnvelopeDebezium { .. });
1578    match stmt.output {
1579        SubscribeOutput::Diffs | SubscribeOutput::WithinTimestampOrderBy { .. } => {
1580            desc = desc.with_column("mz_diff", SqlScalarType::Int64.nullable(true));
1581            for (name, mut ty) in relation_desc.into_iter() {
1582                if progress {
1583                    ty.nullable = true;
1584                }
1585                desc = desc.with_column(name, ty);
1586            }
1587        }
1588        SubscribeOutput::EnvelopeUpsert { key_columns }
1589        | SubscribeOutput::EnvelopeDebezium { key_columns } => {
1590            desc = desc.with_column("mz_state", SqlScalarType::String.nullable(true));
1591            let key_columns = key_columns
1592                .into_iter()
1593                .map(normalize::column_name)
1594                .collect_vec();
1595            let mut before_values_desc = RelationDesc::builder();
1596            let mut after_values_desc = RelationDesc::builder();
1597
1598            // Add the key columns in the order that they're specified.
1599            for column_name in &key_columns {
1600                let mut column_ty = relation_desc
1601                    .get_by_name(column_name)
1602                    .map(|(_pos, ty)| ty.clone())
1603                    .ok_or_else(|| PlanError::UnknownColumn {
1604                        table: None,
1605                        column: column_name.clone(),
1606                        similar: Box::new([]),
1607                    })?;
1608                if progress {
1609                    column_ty.nullable = true;
1610                }
1611                desc = desc.with_column(column_name, column_ty);
1612            }
1613
1614            // Then add the remaining columns in the order from the original
1615            // table, filtering out the key columns since we added those above.
1616            for (mut name, mut ty) in relation_desc
1617                .into_iter()
1618                .filter(|(name, _ty)| !key_columns.contains(name))
1619            {
1620                ty.nullable = true;
1621                before_values_desc =
1622                    before_values_desc.with_column(format!("before_{}", name), ty.clone());
1623                if debezium {
1624                    name = format!("after_{}", name).into();
1625                }
1626                after_values_desc = after_values_desc.with_column(name, ty);
1627            }
1628
1629            if debezium {
1630                desc = desc.concat(before_values_desc);
1631            }
1632            desc = desc.concat(after_values_desc);
1633        }
1634    }
1635    Ok(StatementDesc::new(Some(desc.finish())))
1636}
1637
1638pub fn plan_subscribe(
1639    scx: &StatementContext,
1640    SubscribeStatement {
1641        relation,
1642        options,
1643        as_of,
1644        up_to,
1645        output,
1646    }: SubscribeStatement<Aug>,
1647    params: &Params,
1648    copy_to: Option<CopyFormat>,
1649) -> Result<Plan, PlanError> {
1650    let (from, desc, scope) = match relation {
1651        SubscribeRelation::Name(name) => {
1652            let item = scx.get_item_by_resolved_name(&name)?;
1653            let Some(desc) = item.relation_desc() else {
1654                sql_bail!(
1655                    "'{}' cannot be subscribed to because it is a {}",
1656                    name.full_name_str(),
1657                    item.item_type(),
1658                );
1659            };
1660            let item_name = match name {
1661                ResolvedItemName::Item { full_name, .. } => Some(full_name.into()),
1662                _ => None,
1663            };
1664            let scope = Scope::from_source(item_name, desc.iter().map(|(name, _type)| name));
1665            (
1666                SubscribeFrom::Id(item.global_id()),
1667                desc.into_owned(),
1668                scope,
1669            )
1670        }
1671        SubscribeRelation::Query(query) => {
1672            #[allow(deprecated)] // TODO(aalexandrov): Use HirRelationExpr in Subscribe
1673            let query::PlannedRootQuery {
1674                mut expr,
1675                desc,
1676                finishing,
1677                scope,
1678            } = query::plan_root_query(scx, query, QueryLifetime::Subscribe)?;
1679            expr.bind_parameters_and_simplify_offset(scx, QueryLifetime::Subscribe, params)?;
1680            let query = query::PlannedRootQuery {
1681                expr,
1682                desc,
1683                finishing,
1684                scope,
1685            };
1686            // There's no way to apply finishing operations to a `SUBSCRIBE` directly, so the
1687            // finishing should have already been turned into a `TopK` by
1688            // `plan_query` / `plan_root_query`, upon seeing the `QueryLifetime::Subscribe`.
1689            assert!(HirRelationExpr::is_trivial_row_set_finishing_hir(
1690                &query.finishing,
1691                query.desc.arity()
1692            ));
1693            let desc = query.desc.clone();
1694            (
1695                SubscribeFrom::Query {
1696                    expr: query.expr,
1697                    desc: query.desc,
1698                },
1699                desc,
1700                query.scope,
1701            )
1702        }
1703    };
1704
1705    let when = query::plan_as_of(scx, as_of)?;
1706    let up_to = up_to
1707        .map(|up_to| plan_as_of_or_up_to(scx, up_to))
1708        .transpose()?;
1709
1710    let qcx = QueryContext::root(scx, QueryLifetime::Subscribe);
1711    let ecx = ExprContext {
1712        qcx: &qcx,
1713        name: "",
1714        scope: &scope,
1715        relation_type: desc.typ(),
1716        allow_aggregates: false,
1717        allow_subqueries: true,
1718        allow_parameters: true,
1719        allow_windows: false,
1720    };
1721
1722    let output_columns: Vec<_> = scope.column_names().enumerate().collect();
1723    let output = match output {
1724        SubscribeOutput::Diffs => plan::SubscribeOutput::Diffs,
1725        SubscribeOutput::EnvelopeUpsert { key_columns } => {
1726            let order_by = key_columns
1727                .iter()
1728                .map(|ident| OrderByExpr {
1729                    expr: Expr::Identifier(vec![ident.clone()]),
1730                    asc: None,
1731                    nulls_last: None,
1732                })
1733                .collect_vec();
1734            let (order_by, map_exprs) = query::plan_order_by_exprs(
1735                &ExprContext {
1736                    name: "ENVELOPE UPSERT KEY clause",
1737                    ..ecx
1738                },
1739                &order_by[..],
1740                &output_columns[..],
1741            )?;
1742            if !map_exprs.is_empty() {
1743                return Err(PlanError::InvalidKeysInSubscribeEnvelopeUpsert);
1744            }
1745            check_distinct_key_columns(&order_by, &output_columns)?;
1746            plan::SubscribeOutput::EnvelopeUpsert {
1747                order_by_keys: order_by,
1748            }
1749        }
1750        SubscribeOutput::EnvelopeDebezium { key_columns } => {
1751            scx.require_feature_flag(&vars::ENABLE_ENVELOPE_DEBEZIUM_IN_SUBSCRIBE)?;
1752            let order_by = key_columns
1753                .iter()
1754                .map(|ident| OrderByExpr {
1755                    expr: Expr::Identifier(vec![ident.clone()]),
1756                    asc: None,
1757                    nulls_last: None,
1758                })
1759                .collect_vec();
1760            let (order_by, map_exprs) = query::plan_order_by_exprs(
1761                &ExprContext {
1762                    name: "ENVELOPE DEBEZIUM KEY clause",
1763                    ..ecx
1764                },
1765                &order_by[..],
1766                &output_columns[..],
1767            )?;
1768            if !map_exprs.is_empty() {
1769                return Err(PlanError::InvalidKeysInSubscribeEnvelopeDebezium);
1770            }
1771            check_distinct_key_columns(&order_by, &output_columns)?;
1772            plan::SubscribeOutput::EnvelopeDebezium {
1773                order_by_keys: order_by,
1774            }
1775        }
1776        SubscribeOutput::WithinTimestampOrderBy { order_by } => {
1777            scx.require_feature_flag(&vars::ENABLE_WITHIN_TIMESTAMP_ORDER_BY_IN_SUBSCRIBE)?;
1778            let mz_diff = "mz_diff".into();
1779            let output_columns = std::iter::once((0, &mz_diff))
1780                .chain(output_columns.into_iter().map(|(i, c)| (i + 1, c)))
1781                .collect_vec();
1782            match query::plan_order_by_exprs(
1783                &ExprContext {
1784                    name: "WITHIN TIMESTAMP ORDER BY clause",
1785                    ..ecx
1786                },
1787                &order_by[..],
1788                &output_columns[..],
1789            ) {
1790                Err(PlanError::UnknownColumn {
1791                    table: None,
1792                    column,
1793                    similar: _,
1794                }) if &column == &mz_diff => {
1795                    // mz_diff is being used in an expression. Since mz_diff isn't part of the table
1796                    // it looks like an unknown column. Instead, return a better error
1797                    return Err(PlanError::InvalidOrderByInSubscribeWithinTimestampOrderBy);
1798                }
1799                Err(e) => return Err(e),
1800                Ok((order_by, map_exprs)) => {
1801                    if !map_exprs.is_empty() {
1802                        return Err(PlanError::InvalidOrderByInSubscribeWithinTimestampOrderBy);
1803                    }
1804
1805                    plan::SubscribeOutput::WithinTimestampOrderBy { order_by }
1806                }
1807            }
1808        }
1809    };
1810
1811    let SubscribeOptionExtracted {
1812        progress, snapshot, ..
1813    } = options.try_into()?;
1814    Ok(Plan::Subscribe(SubscribePlan {
1815        from,
1816        when,
1817        up_to,
1818        with_snapshot: snapshot.unwrap_or(true),
1819        copy_to,
1820        emit_progress: progress.unwrap_or(false),
1821        output,
1822    }))
1823}
1824
1825/// Ensures each `ColumnOrder` in `order_by` references a distinct column,
1826/// returning `DuplicateKeyColumnInSubscribeEnvelope` on the first repeat.
1827fn check_distinct_key_columns(
1828    order_by: &[ColumnOrder],
1829    output_columns: &[(usize, &mz_repr::ColumnName)],
1830) -> Result<(), PlanError> {
1831    let mut seen = BTreeSet::new();
1832    for co in order_by {
1833        if !seen.insert(co.column) {
1834            return Err(PlanError::DuplicateKeyColumnInSubscribeEnvelope {
1835                column_name: output_columns[co.column].1.to_string(),
1836            });
1837        }
1838    }
1839    Ok(())
1840}
1841
1842pub fn describe_copy_from_table(
1843    scx: &StatementContext,
1844    table_name: <Aug as AstInfo>::ItemName,
1845    columns: Vec<Ident>,
1846) -> Result<StatementDesc, PlanError> {
1847    let (_, desc, _, _) = query::plan_copy_from(scx, table_name, columns)?;
1848    Ok(StatementDesc::new(Some(desc)))
1849}
1850
1851pub fn describe_copy_item(
1852    scx: &StatementContext,
1853    object_name: <Aug as AstInfo>::ItemName,
1854    columns: Vec<Ident>,
1855) -> Result<StatementDesc, PlanError> {
1856    let (_, desc, _, _) = query::plan_copy_item(scx, object_name, columns)?;
1857    Ok(StatementDesc::new(Some(desc)))
1858}
1859
1860pub fn describe_copy(
1861    scx: &StatementContext,
1862    CopyStatement {
1863        relation,
1864        direction,
1865        ..
1866    }: CopyStatement<Aug>,
1867) -> Result<StatementDesc, PlanError> {
1868    Ok(match (relation, direction) {
1869        (CopyRelation::Named { name, columns }, CopyDirection::To) => {
1870            describe_copy_item(scx, name, columns)?
1871        }
1872        (CopyRelation::Named { name, columns }, CopyDirection::From) => {
1873            describe_copy_from_table(scx, name, columns)?
1874        }
1875        (CopyRelation::Select(stmt), _) => describe_select(scx, stmt)?,
1876        (CopyRelation::Subscribe(stmt), _) => describe_subscribe(scx, stmt)?,
1877    }
1878    .with_is_copy())
1879}
1880
1881fn plan_copy_to_expr(
1882    scx: &StatementContext,
1883    select_plan: SelectPlan,
1884    desc: RelationDesc,
1885    to: &Expr<Aug>,
1886    format: CopyFormat,
1887    options: CopyOptionExtracted,
1888) -> Result<Plan, PlanError> {
1889    let conn_id = match options.aws_connection {
1890        Some(conn_id) => CatalogItemId::from(conn_id),
1891        None => sql_bail!("AWS CONNECTION is required for COPY ... TO <expr>"),
1892    };
1893    let connection = scx.get_item(&conn_id).connection()?;
1894
1895    match connection {
1896        mz_storage_types::connections::Connection::Aws(_) => {}
1897        _ => sql_bail!("only AWS CONNECTION is supported for COPY ... TO <expr>"),
1898    }
1899
1900    let format = match format {
1901        CopyFormat::Csv => {
1902            let quote = extract_byte_param_value(options.quote, "quote")?;
1903            let escape = extract_byte_param_value(options.escape, "escape")?;
1904            let delimiter = extract_byte_param_value(options.delimiter, "delimiter")?;
1905            S3SinkFormat::PgCopy(CopyFormatParams::Csv(
1906                CopyCsvFormatParams::try_new(
1907                    delimiter,
1908                    quote,
1909                    escape,
1910                    options.header,
1911                    options.null,
1912                )
1913                .map_err(|e| sql_err!("{}", e))?,
1914            ))
1915        }
1916        CopyFormat::Parquet => {
1917            // Validate that the output desc can be formatted as parquet.
1918            // COPY TO does not apply any type overrides, so pass `|_| None`.
1919            ArrowBuilder::validate_desc_for_parquet(&desc, |_| None)
1920                .map_err(|e| sql_err!("{}", e))?;
1921            S3SinkFormat::Parquet
1922        }
1923        CopyFormat::Binary => bail_unsupported!("FORMAT BINARY"),
1924        CopyFormat::Text => bail_unsupported!("FORMAT TEXT"),
1925    };
1926
1927    // Converting the to expr to a HirScalarExpr
1928    let mut to_expr = to.clone();
1929    transform_ast::transform(scx, &mut to_expr)?;
1930    let relation_type = RelationDesc::empty();
1931    let ecx = &ExprContext {
1932        qcx: &QueryContext::root(scx, QueryLifetime::OneShot),
1933        name: "COPY TO target",
1934        scope: &Scope::empty(),
1935        relation_type: relation_type.typ(),
1936        allow_aggregates: false,
1937        allow_subqueries: false,
1938        allow_parameters: false,
1939        allow_windows: false,
1940    };
1941
1942    let to = plan_expr(ecx, &to_expr)?.type_as(ecx, &SqlScalarType::String)?;
1943
1944    if options.max_file_size.as_bytes() < MIN_S3_SINK_FILE_SIZE.as_bytes() {
1945        sql_bail!(
1946            "MAX FILE SIZE cannot be less than {}",
1947            MIN_S3_SINK_FILE_SIZE
1948        );
1949    }
1950    if options.max_file_size.as_bytes() > MAX_S3_SINK_FILE_SIZE.as_bytes() {
1951        sql_bail!(
1952            "MAX FILE SIZE cannot be greater than {}",
1953            MAX_S3_SINK_FILE_SIZE
1954        );
1955    }
1956
1957    Ok(Plan::CopyTo(CopyToPlan {
1958        select_plan,
1959        desc,
1960        to,
1961        connection: connection.to_owned(),
1962        connection_id: conn_id,
1963        format,
1964        max_file_size: options.max_file_size.as_bytes(),
1965    }))
1966}
1967
1968fn plan_copy_from(
1969    scx: &StatementContext,
1970    target: &CopyTarget<Aug>,
1971    table_name: ResolvedItemName,
1972    columns: Vec<Ident>,
1973    format: Option<CopyFormat>,
1974    options: CopyOptionExtracted,
1975) -> Result<Plan, PlanError> {
1976    fn only_available_with_csv<T>(option: Option<T>, param: &str) -> Result<(), PlanError> {
1977        match option {
1978            Some(_) => sql_bail!("COPY {} available only in CSV mode", param),
1979            None => Ok(()),
1980        }
1981    }
1982
1983    let source = match target {
1984        CopyTarget::Stdin => CopyFromSource::Stdin,
1985        CopyTarget::Expr(from) => {
1986            // Converting the expr to an HirScalarExpr
1987            let mut from_expr = from.clone();
1988            transform_ast::transform(scx, &mut from_expr)?;
1989            let relation_type = RelationDesc::empty();
1990            let ecx = &ExprContext {
1991                qcx: &QueryContext::root(scx, QueryLifetime::OneShot),
1992                name: "COPY FROM target",
1993                scope: &Scope::empty(),
1994                relation_type: relation_type.typ(),
1995                allow_aggregates: false,
1996                allow_subqueries: false,
1997                allow_parameters: false,
1998                allow_windows: false,
1999            };
2000            let from = plan_expr(ecx, &from_expr)?.type_as(ecx, &SqlScalarType::String)?;
2001
2002            match options.aws_connection {
2003                Some(conn_id) => {
2004                    let conn_id = CatalogItemId::from(conn_id);
2005
2006                    // Validate the connection type is one we expect.
2007                    let connection = match scx.get_item(&conn_id).connection()? {
2008                        mz_storage_types::connections::Connection::Aws(conn) => conn,
2009                        _ => sql_bail!("only AWS CONNECTION is supported in COPY ... FROM"),
2010                    };
2011
2012                    CopyFromSource::AwsS3 {
2013                        uri: from,
2014                        connection,
2015                        connection_id: conn_id,
2016                    }
2017                }
2018                None => CopyFromSource::Url(from),
2019            }
2020        }
2021        CopyTarget::Stdout => bail_never_supported!("COPY FROM {} not supported", target),
2022    };
2023
2024    // COPY FROM a URL or S3 bucket only supports CSV and Parquet. Unlike COPY
2025    // FROM STDIN there's no sensible default format, so one must be specified
2026    // explicitly. Reject unsupported formats here in planning; the coordinator
2027    // relies on this and would otherwise soft-panic.
2028    let format = match &source {
2029        CopyFromSource::Stdin => format.unwrap_or(CopyFormat::Text),
2030        CopyFromSource::Url(_) | CopyFromSource::AwsS3 { .. } => match format {
2031            None => sql_bail!("COPY FROM <expr> requires a FORMAT option"),
2032            Some(CopyFormat::Text) => bail_unsupported!("FORMAT TEXT"),
2033            Some(CopyFormat::Binary) => bail_unsupported!("FORMAT BINARY"),
2034            Some(format @ (CopyFormat::Csv | CopyFormat::Parquet)) => format,
2035        },
2036    };
2037
2038    let params = match format {
2039        CopyFormat::Text => {
2040            only_available_with_csv(options.quote, "quote")?;
2041            only_available_with_csv(options.escape, "escape")?;
2042            only_available_with_csv(options.header, "HEADER")?;
2043            let delimiter =
2044                extract_byte_param_value(options.delimiter, "delimiter")?.unwrap_or(b'\t');
2045            let null = match options.null {
2046                Some(null) => Cow::from(null),
2047                None => Cow::from("\\N"),
2048            };
2049            CopyFormatParams::Text(CopyTextFormatParams { null, delimiter })
2050        }
2051        CopyFormat::Csv => {
2052            let quote = extract_byte_param_value(options.quote, "quote")?;
2053            let escape = extract_byte_param_value(options.escape, "escape")?;
2054            let delimiter = extract_byte_param_value(options.delimiter, "delimiter")?;
2055            CopyFormatParams::Csv(
2056                CopyCsvFormatParams::try_new(
2057                    delimiter,
2058                    quote,
2059                    escape,
2060                    options.header,
2061                    options.null,
2062                )
2063                .map_err(|e| sql_err!("{}", e))?,
2064            )
2065        }
2066        CopyFormat::Binary => bail_unsupported!("FORMAT BINARY"),
2067        CopyFormat::Parquet => CopyFormatParams::Parquet,
2068    };
2069
2070    let filter = match (options.files, options.pattern) {
2071        (Some(_), Some(_)) => bail_unsupported!("must specify one of FILES or PATTERN"),
2072        (Some(files), None) => Some(CopyFromFilter::Files(files)),
2073        (None, Some(pattern)) => Some(CopyFromFilter::Pattern(pattern)),
2074        (None, None) => None,
2075    };
2076
2077    if filter.is_some() && matches!(source, CopyFromSource::Stdin) {
2078        bail_unsupported!("COPY FROM ... WITH (FILES ...) only supported from a URL")
2079    }
2080
2081    let table_name_string = table_name.full_name_str();
2082
2083    let (id, source_desc, columns, maybe_mfp) = query::plan_copy_from(scx, table_name, columns)?;
2084
2085    let Some(mfp) = maybe_mfp else {
2086        sql_bail!("[internal error] COPY FROM ... expects an MFP to be produced");
2087    };
2088
2089    Ok(Plan::CopyFrom(CopyFromPlan {
2090        target_id: id,
2091        target_name: table_name_string,
2092        source,
2093        columns,
2094        source_desc,
2095        mfp,
2096        params,
2097        filter,
2098    }))
2099}
2100
2101fn extract_byte_param_value(v: Option<String>, param_name: &str) -> Result<Option<u8>, PlanError> {
2102    match v {
2103        Some(v) if v.len() == 1 => Ok(Some(v.as_bytes()[0])),
2104        Some(..) => sql_bail!("COPY {} must be a single one-byte character", param_name),
2105        None => Ok(None),
2106    }
2107}
2108
2109generate_extracted_config!(
2110    CopyOption,
2111    (Format, String),
2112    (Delimiter, String),
2113    (Null, String),
2114    (Escape, String),
2115    (Quote, String),
2116    (Header, bool),
2117    (AwsConnection, with_options::Object),
2118    (MaxFileSize, ByteSize, Default(ByteSize::mb(256))),
2119    (Files, Vec<String>),
2120    (Pattern, String)
2121);
2122
2123pub fn plan_copy(
2124    scx: &StatementContext,
2125    CopyStatement {
2126        relation,
2127        direction,
2128        target,
2129        options,
2130    }: CopyStatement<Aug>,
2131) -> Result<Plan, PlanError> {
2132    let options = CopyOptionExtracted::try_from(options)?;
2133    // Parse any user-provided FORMAT option. If not provided, will default to
2134    // Text for COPY TO STDOUT and COPY FROM STDIN, but will error for COPY TO <expr>.
2135    let format = options
2136        .format
2137        .as_ref()
2138        .map(|format| match format.to_lowercase().as_str() {
2139            "text" => Ok(CopyFormat::Text),
2140            "csv" => Ok(CopyFormat::Csv),
2141            "binary" => Ok(CopyFormat::Binary),
2142            "parquet" => Ok(CopyFormat::Parquet),
2143            _ => sql_bail!("unknown FORMAT: {}", format),
2144        })
2145        .transpose()?;
2146
2147    match (&direction, &target) {
2148        (CopyDirection::To, CopyTarget::Stdout) => {
2149            if options.delimiter.is_some() {
2150                sql_bail!("COPY TO does not support DELIMITER option yet");
2151            }
2152            if options.quote.is_some() {
2153                sql_bail!("COPY TO does not support QUOTE option yet");
2154            }
2155            if options.escape.is_some() {
2156                sql_bail!("COPY TO does not support ESCAPE option yet");
2157            }
2158            if options.null.is_some() {
2159                sql_bail!("COPY TO does not support NULL option yet");
2160            }
2161            // `HEADER false` is the default and already honored; only an
2162            // enabled header is unimplemented. Silently accepting it would
2163            // make clients strip the first data row as a presumed header.
2164            if options.header == Some(true) {
2165                sql_bail!("COPY TO does not support HEADER option yet");
2166            }
2167            match relation {
2168                CopyRelation::Named { .. } => sql_bail!("named with COPY TO STDOUT unsupported"),
2169                CopyRelation::Select(stmt) => Ok(plan_select(
2170                    scx,
2171                    stmt,
2172                    &Params::empty(),
2173                    Some(format.unwrap_or(CopyFormat::Text)),
2174                )?),
2175                CopyRelation::Subscribe(stmt) => Ok(plan_subscribe(
2176                    scx,
2177                    stmt,
2178                    &Params::empty(),
2179                    Some(format.unwrap_or(CopyFormat::Text)),
2180                )?),
2181            }
2182        }
2183        (CopyDirection::From, target) => match relation {
2184            CopyRelation::Named { name, columns } => {
2185                plan_copy_from(scx, target, name, columns, format, options)
2186            }
2187            _ => sql_bail!("COPY FROM {} not supported", target),
2188        },
2189        (CopyDirection::To, CopyTarget::Expr(to_expr)) => {
2190            let format = match format {
2191                Some(inner) => inner,
2192                _ => sql_bail!("COPY TO <expr> requires a FORMAT option"),
2193            };
2194
2195            let stmt = match relation {
2196                CopyRelation::Named { name, columns } => {
2197                    if !columns.is_empty() {
2198                        // TODO(mouli): Add support for this
2199                        sql_bail!(
2200                            "specifying columns for COPY <table_name> TO commands not yet supported; use COPY (SELECT...) TO ... instead"
2201                        );
2202                    }
2203                    // Generate a synthetic SELECT query that just gets the table
2204                    let query = Query {
2205                        ctes: CteBlock::empty(),
2206                        body: SetExpr::Table(name),
2207                        order_by: vec![],
2208                        limit: None,
2209                        offset: None,
2210                    };
2211                    SelectStatement { query, as_of: None }
2212                }
2213                CopyRelation::Select(stmt) => {
2214                    if !stmt.query.order_by.is_empty() {
2215                        sql_bail!("ORDER BY is not supported in SELECT query for COPY statements")
2216                    }
2217                    stmt
2218                }
2219                CopyRelation::Subscribe(_) => {
2220                    sql_bail!("COPY {} {} not supported", direction, target)
2221                }
2222            };
2223
2224            let (plan, desc) = plan_select_inner(scx, stmt, &Params::empty(), None)?;
2225            plan_copy_to_expr(scx, plan, desc, to_expr, format, options)
2226        }
2227        _ => sql_bail!("COPY {} {} not supported", direction, target),
2228    }
2229}