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
890// The `EXPLAIN ANALYZE ... CPU` queries attribute scheduling time to LIR nodes through each
891// node's operator span in `mz_lir_mapping`. A span includes the regions the node renders into,
892// and a region's `elapsed_ns` already contains the time of the operators inside it, so only the
893// span's outermost operators count: its operators at the smallest address depth. Channels have
894// their scope's address and no elapsed time, so the join with `mz_dataflow_operators` keeps them
895// out of the depth comparison. This relies on each span being recorded at a single scope level,
896// see `render_letfree_plan` in mz-compute.
897const LIR_SPAN_OPERATORS_CTE: (&str, &str) = (
898    "lir_span_operators",
899    r#"
900  SELECT mlm.global_id AS global_id,
901         mlm.lir_id AS lir_id,
902         valid_id AS id,
903         list_length(mda.address) AS depth
904    FROM            mz_introspection.mz_lir_mapping mlm
905         CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
906               JOIN mz_introspection.mz_dataflow_operators mdo
907                 ON (mdo.id = valid_id)
908               JOIN mz_introspection.mz_dataflow_addresses mda
909                 ON (mda.id = valid_id)"#,
910);
911
912/// Per LIR node and worker, the scheduling time of the node's outermost operators. Reads
913/// `LIR_SPAN_OPERATORS_CTE`, which must precede it.
914const LIR_CPU_PER_WORKER_CTE: (&str, &str) = (
915    "lir_cpu_per_worker",
916    r#"
917  SELECT lso.global_id AS global_id,
918         lso.lir_id AS lir_id,
919         mse.worker_id AS worker_id,
920         SUM(mse.elapsed_ns) AS worker_ns
921    FROM            lir_span_operators lso
922               JOIN mz_introspection.mz_scheduling_elapsed_per_worker mse
923                 ON (mse.id = lso.id)
924   WHERE (lso.global_id, lso.lir_id, lso.depth) IN
925         (SELECT global_id, lir_id, MIN(depth) FROM lir_span_operators GROUP BY global_id, lir_id)
926GROUP BY lso.global_id, lso.lir_id, mse.worker_id"#,
927);
928
929pub fn plan_explain_analyze_object(
930    scx: &StatementContext,
931    statement: ExplainAnalyzeObjectStatement<Aug>,
932    params: &Params,
933) -> Result<Plan, PlanError> {
934    let explainee_name = statement
935        .explainee
936        .name()
937        .ok_or_else(|| sql_err!("EXPLAIN ANALYZE on anonymous dataflows",))?
938        .full_name_str();
939    let explainee = plan_explainee(scx, statement.explainee, params)?;
940
941    let check_ownership = |item_id: &CatalogItemId, item_type: &str| -> Result<(), PlanError> {
942        if scx.catalog.restrict_to_user_objects() {
943            let item = scx.catalog.get_item(item_id);
944            if item.owner_id() != *scx.catalog.active_role_id() {
945                let full_name = scx.catalog.resolve_full_name(item.name());
946                return Err(sql_err!("must be owner of {item_type} {full_name}"));
947            }
948        }
949        Ok(())
950    };
951    match &explainee {
952        plan::Explainee::Index(item_id) => check_ownership(item_id, "INDEX")?,
953        plan::Explainee::MaterializedView(item_id) => {
954            check_ownership(item_id, "MATERIALIZED VIEW")?
955        }
956        _ => return Err(sql_err!("EXPLAIN ANALYZE queries for this explainee type",)),
957    };
958
959    // generate SQL query
960
961    /* WITH {CTEs}
962       SELECT REPEAT(' ', nesting * 2) || operator AS operator
963             {columns}
964        FROM      mz_introspection.mz_lir_mapping mlm
965             JOIN {from} USING (lir_id)
966             JOIN mz_introspection.mz_mappable_objects mo
967               ON (mlm.global_id = mo.global_id)
968       WHERE     mo.name = {escaped explainee_name}
969             AND {predicates}
970       ORDER BY lir_id DESC
971    */
972    let mut ctes = Vec::with_capacity(5); // max 3 per ExplainAnalyzeComputationProperty
973    let mut columns = vec!["REPEAT(' ', nesting * 2) || operator AS operator"];
974    let mut from = vec!["mz_introspection.mz_lir_mapping mlm"];
975    let mut predicates = vec![format!(
976        "mo.name = {}",
977        escaped_string_literal(&explainee_name)
978    )];
979    let mut order_by = vec!["mlm.lir_id DESC"];
980
981    match statement.properties {
982        ExplainAnalyzeProperty::Computation(ExplainAnalyzeComputationProperties {
983            properties,
984            skew,
985        }) => {
986            let mut worker_id = None;
987            let mut seen_properties = BTreeSet::new();
988            for property in properties {
989                // handle each property only once (belt and suspenders)
990                if !seen_properties.insert(property) {
991                    continue;
992                }
993
994                match property {
995                    ExplainAnalyzeComputationProperty::Memory => {
996                        ctes.push((
997                            "summary_memory",
998                            r#"
999  SELECT mlm.global_id AS global_id,
1000         mlm.lir_id AS lir_id,
1001         SUM(mas.size) AS total_memory,
1002         SUM(mas.records) AS total_records,
1003         CASE WHEN COUNT(DISTINCT mas.worker_id) <> 0 THEN SUM(mas.size) / COUNT(DISTINCT mas.worker_id) ELSE NULL END AS avg_memory,
1004         CASE WHEN COUNT(DISTINCT mas.worker_id) <> 0 THEN SUM(mas.records) / COUNT(DISTINCT mas.worker_id) ELSE NULL END AS avg_records
1005    FROM            mz_introspection.mz_lir_mapping mlm
1006         CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1007               JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
1008                 ON (mas.operator_id = valid_id)
1009GROUP BY mlm.global_id, mlm.lir_id"#,
1010                        ));
1011                        from.push("LEFT JOIN summary_memory sm USING (global_id, lir_id)");
1012
1013                        if skew {
1014                            ctes.push((
1015                                "per_worker_memory",
1016                                r#"
1017  SELECT mlm.global_id AS global_id,
1018         mlm.lir_id AS lir_id,
1019         mas.worker_id AS worker_id,
1020         SUM(mas.size) AS worker_memory,
1021         SUM(mas.records) AS worker_records
1022    FROM            mz_introspection.mz_lir_mapping mlm
1023         CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1024               JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
1025                 ON (mas.operator_id = valid_id)
1026GROUP BY mlm.global_id, mlm.lir_id, mas.worker_id"#,
1027                            ));
1028                            from.push("LEFT JOIN per_worker_memory pwm USING (global_id, lir_id)");
1029
1030                            if let Some(worker_id) = worker_id {
1031                                predicates.push(format!(
1032                                    "(pwm.worker_id = {worker_id} OR pwm.worker_id IS NULL OR {worker_id} IS NULL)"
1033                                ));
1034                            } else {
1035                                worker_id = Some("pwm.worker_id");
1036                                columns.push("pwm.worker_id AS worker_id");
1037                                order_by.push("worker_id");
1038                            }
1039
1040                            columns.extend([
1041                                "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",
1042                                "pg_size_pretty(pwm.worker_memory) AS worker_memory",
1043                                "pg_size_pretty(sm.avg_memory) AS avg_memory",
1044                                "pg_size_pretty(sm.total_memory) AS total_memory",
1045                                "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",
1046                                "pwm.worker_records AS worker_records",
1047                                "sm.avg_records AS avg_records",
1048                                "sm.total_records AS total_records",
1049                            ]);
1050                        } else {
1051                            columns.extend([
1052                                "pg_size_pretty(sm.total_memory) AS total_memory",
1053                                "sm.total_records AS total_records",
1054                            ]);
1055                        }
1056                    }
1057                    ExplainAnalyzeComputationProperty::Cpu => {
1058                        ctes.extend([LIR_SPAN_OPERATORS_CTE, LIR_CPU_PER_WORKER_CTE]);
1059                        ctes.push((
1060                            "summary_cpu",
1061                            r#"
1062  SELECT global_id,
1063         lir_id,
1064         SUM(worker_ns) AS total_ns,
1065         SUM(worker_ns) / COUNT(*) AS avg_ns
1066    FROM lir_cpu_per_worker
1067GROUP BY global_id, lir_id"#,
1068                        ));
1069                        from.push("LEFT JOIN summary_cpu sc USING (global_id, lir_id)");
1070
1071                        if skew {
1072                            from.push("LEFT JOIN lir_cpu_per_worker pwc USING (global_id, lir_id)");
1073
1074                            if let Some(worker_id) = worker_id {
1075                                predicates.push(format!(
1076                                    "(pwc.worker_id = {worker_id} OR pwc.worker_id IS NULL OR {worker_id} IS NULL)"
1077                                ));
1078                            } else {
1079                                worker_id = Some("pwc.worker_id");
1080                                columns.push("pwc.worker_id AS worker_id");
1081                                order_by.push("worker_id");
1082                            }
1083
1084                            columns.extend([
1085                                "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",
1086                                "pwc.worker_ns / 1000 * '1 microsecond'::INTERVAL AS worker_elapsed",
1087                                "sc.avg_ns / 1000 * '1 microsecond'::INTERVAL AS avg_elapsed",
1088                            ]);
1089                        }
1090                        columns.push(
1091                            "sc.total_ns / 1000 * '1 microsecond'::INTERVAL AS total_elapsed",
1092                        );
1093                    }
1094                }
1095            }
1096        }
1097        ExplainAnalyzeProperty::Hints => {
1098            columns.extend([
1099                "megsa.levels AS levels",
1100                "megsa.to_cut AS to_cut",
1101                "megsa.hint AS hint",
1102                "pg_size_pretty(megsa.savings) AS savings",
1103            ]);
1104            from.extend(["JOIN mz_introspection.mz_dataflow_global_ids mdgi ON (mlm.global_id = mdgi.global_id)",
1105            "LEFT JOIN (generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id JOIN \
1106             mz_introspection.mz_expected_group_size_advice megsa ON (megsa.region_id = valid_id)) ON (megsa.dataflow_id = mdgi.id)"]);
1107        }
1108    }
1109
1110    from.push("JOIN mz_introspection.mz_mappable_objects mo ON (mlm.global_id = mo.global_id)");
1111
1112    let ctes = if !ctes.is_empty() {
1113        format!(
1114            "WITH {}",
1115            separated(
1116                ",\n",
1117                ctes.iter()
1118                    .map(|(name, defn)| format!("{name} AS ({defn})"))
1119            )
1120        )
1121    } else {
1122        String::new()
1123    };
1124    let columns = separated(", ", columns);
1125    let from = separated(" ", from);
1126    let predicates = separated(" AND ", predicates);
1127    let order_by = separated(", ", order_by);
1128    let query = format!(
1129        r#"{ctes}
1130SELECT {columns}
1131FROM {from}
1132WHERE {predicates}
1133ORDER BY {order_by}"#
1134    );
1135
1136    if statement.as_sql {
1137        let rows = vec![Row::pack_slice(&[Datum::String(
1138            &mz_sql_pretty::pretty_str_simple(&query, 80).map_err(|e| {
1139                PlanError::Unstructured(format!("internal error parsing our own SQL: {e}"))
1140            })?,
1141        )])];
1142        let typ = SqlRelationType::new(vec![SqlScalarType::String.nullable(false)]);
1143
1144        Ok(Plan::Select(SelectPlan::immediate(rows, typ)))
1145    } else {
1146        let (show_select, resolved_ids) = ShowSelect::new_from_bare_query(scx, query)?;
1147        scx.record_sql_impl_ids(&resolved_ids);
1148        show_select.plan()
1149    }
1150}
1151
1152pub fn plan_explain_analyze_cluster(
1153    scx: &StatementContext,
1154    statement: ExplainAnalyzeClusterStatement,
1155    _params: &Params,
1156) -> Result<Plan, PlanError> {
1157    // object string
1158    // worker_id uint64        (if           skew)
1159    // memory_ratio numeric    (if memory && skew)
1160    // worker_memory string    (if memory && skew)
1161    // avg_memory string       (if memory && skew)
1162    // total_memory string     (if memory)
1163    // records_ratio numeric   (if memory && skew)
1164    // worker_records          (if memory && skew)
1165    // avg_records numeric     (if memory && skew)
1166    // total_records numeric   (if memory)
1167    // cpu_ratio numeric       (if cpu    && skew)
1168    // worker_elapsed interval (if cpu    && skew)
1169    // avg_elapsed interval    (if cpu    && skew)
1170    // total_elapsed interval  (if cpu)
1171
1172    /* WITH {CTEs}
1173       SELECT mo.name AS object
1174             {columns}
1175        FROM mz_introspection.mz_mappable_objects mo
1176             {from}
1177       WHERE {predicates}
1178       ORDER BY {order_by}, mo.name DESC
1179    */
1180    let mut ctes = Vec::with_capacity(4); // max 2 per ExplainAnalyzeComputationProperty
1181    let mut columns = vec!["mo.name AS object", "mo.global_id AS global_id"];
1182    let mut from = vec!["mz_introspection.mz_mappable_objects mo"];
1183    let mut predicates = vec![];
1184    let mut order_by = vec![];
1185
1186    let ExplainAnalyzeComputationProperties { properties, skew } = statement.properties;
1187    let mut worker_id = None;
1188    let mut seen_properties = BTreeSet::new();
1189    for property in properties {
1190        // handle each property only once (belt and suspenders)
1191        if !seen_properties.insert(property) {
1192            continue;
1193        }
1194
1195        match property {
1196            ExplainAnalyzeComputationProperty::Memory => {
1197                if skew {
1198                    let mut set_worker_id = false;
1199                    if let Some(worker_id) = worker_id {
1200                        // join condition if we're showing skew for more than one property
1201                        predicates.push(format!(
1202                            "(om.worker_id = {worker_id} OR om.worker_id IS NULL OR {worker_id} IS NULL)"
1203                        ));
1204                    } else {
1205                        worker_id = Some("om.worker_id");
1206                        columns.push("om.worker_id AS worker_id");
1207                        set_worker_id = true; // we'll add ourselves to `order_by` later
1208                    };
1209
1210                    // computes the average memory per LIR operator (for per operator ratios)
1211                    ctes.push((
1212                    "per_operator_memory_summary",
1213                    r#"
1214SELECT mlm.global_id AS global_id,
1215       mlm.lir_id AS lir_id,
1216       SUM(mas.size) AS total_memory,
1217       SUM(mas.records) AS total_records,
1218       CASE WHEN COUNT(DISTINCT mas.worker_id) <> 0 THEN SUM(mas.size) / COUNT(DISTINCT mas.worker_id) ELSE NULL END AS avg_memory,
1219       CASE WHEN COUNT(DISTINCT mas.worker_id) <> 0 THEN SUM(mas.records) / COUNT(DISTINCT mas.worker_id) ELSE NULL END AS avg_records
1220FROM        mz_introspection.mz_lir_mapping mlm
1221 CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1222       JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
1223         ON (mas.operator_id = valid_id)
1224GROUP BY mlm.global_id, mlm.lir_id"#,
1225                ));
1226
1227                    // computes the memory per worker in a per operator way
1228                    ctes.push((
1229                    "per_operator_memory_per_worker",
1230                    r#"
1231SELECT mlm.global_id AS global_id,
1232       mlm.lir_id AS lir_id,
1233       mas.worker_id AS worker_id,
1234       SUM(mas.size) AS worker_memory,
1235       SUM(mas.records) AS worker_records
1236FROM        mz_introspection.mz_lir_mapping mlm
1237 CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1238       JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
1239         ON (mas.operator_id = valid_id)
1240GROUP BY mlm.global_id, mlm.lir_id, mas.worker_id"#,
1241                    ));
1242
1243                    // computes memory ratios per worker per operator
1244                    ctes.push((
1245                    "per_operator_memory_ratios",
1246                    r#"
1247SELECT pompw.global_id AS global_id,
1248       pompw.lir_id AS lir_id,
1249       pompw.worker_id AS worker_id,
1250       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,
1251       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
1252  FROM      per_operator_memory_per_worker pompw
1253       JOIN per_operator_memory_summary poms
1254         USING (global_id, lir_id)
1255"#,
1256                    ));
1257
1258                    // summarizes each object, per worker
1259                    ctes.push((
1260                        "object_memory",
1261                        r#"
1262SELECT pompw.global_id AS global_id,
1263       pompw.worker_id AS worker_id,
1264       MAX(pomr.memory_ratio) AS max_operator_memory_ratio,
1265       MAX(pomr.records_ratio) AS max_operator_records_ratio,
1266       SUM(pompw.worker_memory) AS worker_memory,
1267       SUM(pompw.worker_records) AS worker_records
1268FROM        per_operator_memory_per_worker pompw
1269     JOIN   per_operator_memory_ratios pomr
1270     USING (global_id, worker_id, lir_id)
1271GROUP BY pompw.global_id, pompw.worker_id
1272"#,
1273                    ));
1274
1275                    // summarizes each worker
1276                    ctes.push(("object_average_memory", r#"
1277SELECT om.global_id AS global_id,
1278       SUM(om.worker_memory) AS total_memory,
1279       CASE WHEN COUNT(DISTINCT om.worker_id) <> 0 THEN SUM(om.worker_memory) / COUNT(DISTINCT om.worker_id) ELSE NULL END AS avg_memory,
1280       SUM(om.worker_records) AS total_records,
1281       CASE WHEN COUNT(DISTINCT om.worker_id) <> 0 THEN SUM(om.worker_records) / COUNT(DISTINCT om.worker_id) ELSE NULL END AS avg_records
1282  FROM object_memory om
1283GROUP BY om.global_id"#));
1284
1285                    from.push("LEFT JOIN object_memory om USING (global_id)");
1286                    from.push("LEFT JOIN object_average_memory oam USING (global_id)");
1287
1288                    columns.extend([
1289                        "om.max_operator_memory_ratio AS max_operator_memory_ratio",
1290                        "pg_size_pretty(om.worker_memory) AS worker_memory",
1291                        "pg_size_pretty(oam.avg_memory) AS avg_memory",
1292                        "pg_size_pretty(oam.total_memory) AS total_memory",
1293                        "om.max_operator_records_ratio AS max_operator_records_ratio",
1294                        "om.worker_records AS worker_records",
1295                        "oam.avg_records AS avg_records",
1296                        "oam.total_records AS total_records",
1297                    ]);
1298
1299                    order_by.extend([
1300                        "max_operator_memory_ratio DESC NULLS LAST",
1301                        "max_operator_records_ratio DESC NULLS LAST",
1302                        "om.worker_memory DESC NULLS LAST",
1303                        "worker_records DESC NULLS LAST",
1304                    ]);
1305
1306                    if set_worker_id {
1307                        order_by.push("worker_id");
1308                    }
1309                } else {
1310                    // no skew, so just compute totals
1311                    ctes.push((
1312                        "per_operator_memory_totals",
1313                        r#"
1314    SELECT mlm.global_id AS global_id,
1315           mlm.lir_id AS lir_id,
1316           SUM(mas.size) AS total_memory,
1317           SUM(mas.records) AS total_records
1318    FROM        mz_introspection.mz_lir_mapping mlm
1319     CROSS JOIN generate_series((mlm.operator_id_start) :: int8, (mlm.operator_id_end - 1) :: int8) AS valid_id
1320           JOIN mz_introspection.mz_arrangement_sizes_per_worker mas
1321             ON (mas.operator_id = valid_id)
1322    GROUP BY mlm.global_id, mlm.lir_id"#,
1323                    ));
1324
1325                    ctes.push((
1326                        "object_memory_totals",
1327                        r#"
1328SELECT pomt.global_id AS global_id,
1329       SUM(pomt.total_memory) AS total_memory,
1330       SUM(pomt.total_records) AS total_records
1331FROM per_operator_memory_totals pomt
1332GROUP BY pomt.global_id
1333"#,
1334                    ));
1335
1336                    from.push("LEFT JOIN object_memory_totals omt USING (global_id)");
1337                    columns.extend([
1338                        "pg_size_pretty(omt.total_memory) AS total_memory",
1339                        "omt.total_records AS total_records",
1340                    ]);
1341                    order_by.extend([
1342                        "omt.total_memory DESC NULLS LAST",
1343                        "total_records DESC NULLS LAST",
1344                    ]);
1345                }
1346            }
1347            ExplainAnalyzeComputationProperty::Cpu => {
1348                ctes.extend([LIR_SPAN_OPERATORS_CTE, LIR_CPU_PER_WORKER_CTE]);
1349                if skew {
1350                    let mut set_worker_id = false;
1351                    if let Some(worker_id) = worker_id {
1352                        // join condition if we're showing skew for more than one property
1353                        predicates.push(format!(
1354                            "(oc.worker_id = {worker_id} OR oc.worker_id IS NULL OR {worker_id} IS NULL)"
1355                        ));
1356                    } else {
1357                        worker_id = Some("oc.worker_id");
1358                        columns.push("oc.worker_id AS worker_id");
1359                        set_worker_id = true; // we'll add ourselves to `order_by` later
1360                    };
1361
1362                    // computes the average CPU per LIR operator (for per operator ratios)
1363                    ctes.push((
1364                        "per_operator_cpu_summary",
1365                        r#"
1366SELECT global_id,
1367       lir_id,
1368       SUM(worker_ns) AS total_ns,
1369       SUM(worker_ns) / COUNT(*) AS avg_ns
1370FROM lir_cpu_per_worker
1371GROUP BY global_id, lir_id"#,
1372                    ));
1373
1374                    // computes CPU ratios per worker per operator
1375                    ctes.push((
1376                        "per_operator_cpu_ratios",
1377                        r#"
1378SELECT pocpw.global_id AS global_id,
1379       pocpw.lir_id AS lir_id,
1380       pocpw.worker_id AS worker_id,
1381       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
1382FROM      lir_cpu_per_worker pocpw
1383     JOIN per_operator_cpu_summary pocs
1384     USING (global_id, lir_id)
1385"#,
1386                    ));
1387
1388                    // summarizes each object, per worker
1389                    ctes.push((
1390                        "object_cpu",
1391                        r#"
1392SELECT pocpw.global_id AS global_id,
1393       pocpw.worker_id AS worker_id,
1394       MAX(pomr.cpu_ratio) AS max_operator_cpu_ratio,
1395       SUM(pocpw.worker_ns) AS worker_ns
1396FROM      lir_cpu_per_worker pocpw
1397     JOIN per_operator_cpu_ratios pomr
1398     USING (global_id, worker_id, lir_id)
1399GROUP BY pocpw.global_id, pocpw.worker_id
1400"#,
1401                    ));
1402
1403                    // summarizes each worker
1404                    ctes.push((
1405                        "object_average_cpu",
1406                        r#"
1407SELECT oc.global_id AS global_id,
1408       SUM(oc.worker_ns) AS total_ns,
1409       CASE WHEN COUNT(DISTINCT oc.worker_id) <> 0 THEN SUM(oc.worker_ns) / COUNT(DISTINCT oc.worker_id) ELSE NULL END AS avg_ns
1410  FROM object_cpu oc
1411GROUP BY oc.global_id"#,));
1412
1413                    from.push("LEFT JOIN object_cpu oc USING (global_id)");
1414                    from.push("LEFT JOIN object_average_cpu oac USING (global_id)");
1415
1416                    columns.extend([
1417                        "oc.max_operator_cpu_ratio AS max_operator_cpu_ratio",
1418                        "oc.worker_ns / 1000 * '1 microsecond'::interval AS worker_elapsed",
1419                        "oac.avg_ns / 1000 * '1 microsecond'::interval AS avg_elapsed",
1420                        "oac.total_ns / 1000 * '1 microsecond'::interval AS total_elapsed",
1421                    ]);
1422
1423                    order_by.extend([
1424                        "max_operator_cpu_ratio DESC NULLS LAST",
1425                        "worker_elapsed DESC NULLS LAST",
1426                    ]);
1427
1428                    if set_worker_id {
1429                        order_by.push("worker_id");
1430                    }
1431                } else {
1432                    // no skew, so just compute totals
1433                    ctes.push((
1434                        "per_operator_cpu_totals",
1435                        r#"
1436SELECT global_id, lir_id, SUM(worker_ns) AS total_ns
1437FROM lir_cpu_per_worker
1438GROUP BY global_id, lir_id"#,
1439                    ));
1440
1441                    ctes.push((
1442                        "object_cpu_totals",
1443                        r#"
1444SELECT poct.global_id AS global_id,
1445       SUM(poct.total_ns) AS total_ns
1446FROM per_operator_cpu_totals poct
1447GROUP BY poct.global_id
1448"#,
1449                    ));
1450
1451                    from.push("LEFT JOIN object_cpu_totals oct USING (global_id)");
1452                    columns
1453                        .push("oct.total_ns / 1000 * '1 microsecond'::interval AS total_elapsed");
1454                    order_by.extend(["total_elapsed DESC NULLS LAST"]);
1455                }
1456            }
1457        }
1458    }
1459
1460    // generate SQL query text
1461    let ctes = if !ctes.is_empty() {
1462        format!(
1463            "WITH {}",
1464            separated(
1465                ",\n",
1466                ctes.iter()
1467                    .map(|(name, defn)| format!("{name} AS ({defn})"))
1468            )
1469        )
1470    } else {
1471        String::new()
1472    };
1473    let columns = separated(", ", columns);
1474    let from = separated(" ", from);
1475    let predicates = if !predicates.is_empty() {
1476        format!("WHERE {}", separated(" AND ", predicates))
1477    } else {
1478        String::new()
1479    };
1480    // add mo.name last, to break ties only
1481    order_by.push("mo.name DESC");
1482    let order_by = separated(", ", order_by);
1483    let query = format!(
1484        r#"{ctes}
1485SELECT {columns}
1486FROM {from}
1487{predicates}
1488ORDER BY {order_by}"#
1489    );
1490
1491    if statement.as_sql {
1492        let rows = vec![Row::pack_slice(&[Datum::String(
1493            &mz_sql_pretty::pretty_str_simple(&query, 80).map_err(|e| {
1494                PlanError::Unstructured(format!("internal error parsing our own SQL: {e}"))
1495            })?,
1496        )])];
1497        let typ = SqlRelationType::new(vec![SqlScalarType::String.nullable(false)]);
1498
1499        Ok(Plan::Select(SelectPlan::immediate(rows, typ)))
1500    } else {
1501        let (show_select, resolved_ids) = ShowSelect::new_from_bare_query(scx, query)?;
1502        scx.record_sql_impl_ids(&resolved_ids);
1503        show_select.plan()
1504    }
1505}
1506
1507pub fn plan_explain_timestamp(
1508    scx: &StatementContext,
1509    explain: ExplainTimestampStatement<Aug>,
1510) -> Result<Plan, PlanError> {
1511    let (format, _verbose_syntax) = match explain.format() {
1512        mz_sql_parser::ast::ExplainFormat::Text => (ExplainFormat::Text, false),
1513        mz_sql_parser::ast::ExplainFormat::VerboseText => (ExplainFormat::Text, true),
1514        mz_sql_parser::ast::ExplainFormat::Json => (ExplainFormat::Json, false),
1515        mz_sql_parser::ast::ExplainFormat::Dot => (ExplainFormat::Dot, false),
1516    };
1517
1518    let raw_plan = {
1519        let query::PlannedRootQuery {
1520            expr: raw_plan,
1521            desc: _,
1522            finishing: _,
1523            scope: _,
1524        } = query::plan_root_query(scx, explain.select.query, QueryLifetime::OneShot)?;
1525        if raw_plan.contains_parameters()? {
1526            return Err(PlanError::ParameterNotAllowed(
1527                "EXPLAIN TIMESTAMP".to_string(),
1528            ));
1529        }
1530
1531        raw_plan
1532    };
1533    let when = query::plan_as_of(scx, explain.select.as_of)?;
1534
1535    Ok(Plan::ExplainTimestamp(ExplainTimestampPlan {
1536        format,
1537        raw_plan,
1538        when,
1539    }))
1540}
1541
1542generate_extracted_config!(SubscribeOption, (Snapshot, bool), (Progress, bool));
1543
1544pub fn describe_subscribe(
1545    scx: &StatementContext,
1546    stmt: SubscribeStatement<Aug>,
1547) -> Result<StatementDesc, PlanError> {
1548    let relation_desc = match stmt.relation {
1549        SubscribeRelation::Name(name) => {
1550            let item = scx.get_item_by_resolved_name(&name)?;
1551            match item.relation_desc() {
1552                Some(desc) => desc.into_owned(),
1553                None => sql_bail!(
1554                    "'{}' cannot be subscribed to because it is a {}",
1555                    name.full_name_str(),
1556                    item.item_type(),
1557                ),
1558            }
1559        }
1560        SubscribeRelation::Query(query) => {
1561            let query::PlannedRootQuery { desc, .. } =
1562                query::plan_root_query(scx, query, QueryLifetime::Subscribe)?;
1563            desc
1564        }
1565    };
1566    let SubscribeOptionExtracted { progress, .. } = stmt.options.try_into()?;
1567    let progress = progress.unwrap_or(false);
1568    let mut desc = RelationDesc::builder().with_column(
1569        "mz_timestamp",
1570        SqlScalarType::Numeric {
1571            max_scale: Some(NumericMaxScale::ZERO),
1572        }
1573        .nullable(false),
1574    );
1575    if progress {
1576        desc = desc.with_column("mz_progressed", SqlScalarType::Bool.nullable(false));
1577    }
1578
1579    let debezium = matches!(stmt.output, SubscribeOutput::EnvelopeDebezium { .. });
1580    match stmt.output {
1581        SubscribeOutput::Diffs | SubscribeOutput::WithinTimestampOrderBy { .. } => {
1582            desc = desc.with_column("mz_diff", SqlScalarType::Int64.nullable(true));
1583            for (name, mut ty) in relation_desc.into_iter() {
1584                if progress {
1585                    ty.nullable = true;
1586                }
1587                desc = desc.with_column(name, ty);
1588            }
1589        }
1590        SubscribeOutput::EnvelopeUpsert { key_columns }
1591        | SubscribeOutput::EnvelopeDebezium { key_columns } => {
1592            desc = desc.with_column("mz_state", SqlScalarType::String.nullable(true));
1593            let key_columns = key_columns
1594                .into_iter()
1595                .map(normalize::column_name)
1596                .collect_vec();
1597            let mut before_values_desc = RelationDesc::builder();
1598            let mut after_values_desc = RelationDesc::builder();
1599
1600            // Add the key columns in the order that they're specified.
1601            for column_name in &key_columns {
1602                let mut column_ty = relation_desc
1603                    .get_by_name(column_name)
1604                    .map(|(_pos, ty)| ty.clone())
1605                    .ok_or_else(|| PlanError::UnknownColumn {
1606                        table: None,
1607                        column: column_name.clone(),
1608                        similar: Box::new([]),
1609                    })?;
1610                if progress {
1611                    column_ty.nullable = true;
1612                }
1613                desc = desc.with_column(column_name, column_ty);
1614            }
1615
1616            // Then add the remaining columns in the order from the original
1617            // table, filtering out the key columns since we added those above.
1618            for (mut name, mut ty) in relation_desc
1619                .into_iter()
1620                .filter(|(name, _ty)| !key_columns.contains(name))
1621            {
1622                ty.nullable = true;
1623                before_values_desc =
1624                    before_values_desc.with_column(format!("before_{}", name), ty.clone());
1625                if debezium {
1626                    name = format!("after_{}", name).into();
1627                }
1628                after_values_desc = after_values_desc.with_column(name, ty);
1629            }
1630
1631            if debezium {
1632                desc = desc.concat(before_values_desc);
1633            }
1634            desc = desc.concat(after_values_desc);
1635        }
1636    }
1637    Ok(StatementDesc::new(Some(desc.finish())))
1638}
1639
1640pub fn plan_subscribe(
1641    scx: &StatementContext,
1642    SubscribeStatement {
1643        relation,
1644        options,
1645        as_of,
1646        up_to,
1647        output,
1648    }: SubscribeStatement<Aug>,
1649    params: &Params,
1650    copy_to: Option<CopyFormat>,
1651) -> Result<Plan, PlanError> {
1652    let (from, desc, scope) = match relation {
1653        SubscribeRelation::Name(name) => {
1654            let item = scx.get_item_by_resolved_name(&name)?;
1655            let Some(desc) = item.relation_desc() else {
1656                sql_bail!(
1657                    "'{}' cannot be subscribed to because it is a {}",
1658                    name.full_name_str(),
1659                    item.item_type(),
1660                );
1661            };
1662            let item_name = match name {
1663                ResolvedItemName::Item { full_name, .. } => Some(full_name.into()),
1664                _ => None,
1665            };
1666            let scope = Scope::from_source(item_name, desc.iter().map(|(name, _type)| name));
1667            (
1668                SubscribeFrom::Id(item.global_id()),
1669                desc.into_owned(),
1670                scope,
1671            )
1672        }
1673        SubscribeRelation::Query(query) => {
1674            #[allow(deprecated)] // TODO(aalexandrov): Use HirRelationExpr in Subscribe
1675            let query::PlannedRootQuery {
1676                mut expr,
1677                desc,
1678                finishing,
1679                scope,
1680            } = query::plan_root_query(scx, query, QueryLifetime::Subscribe)?;
1681            expr.bind_parameters_and_simplify_offset(scx, QueryLifetime::Subscribe, params)?;
1682            let query = query::PlannedRootQuery {
1683                expr,
1684                desc,
1685                finishing,
1686                scope,
1687            };
1688            // There's no way to apply finishing operations to a `SUBSCRIBE` directly, so the
1689            // finishing should have already been turned into a `TopK` by
1690            // `plan_query` / `plan_root_query`, upon seeing the `QueryLifetime::Subscribe`.
1691            assert!(HirRelationExpr::is_trivial_row_set_finishing_hir(
1692                &query.finishing,
1693                query.desc.arity()
1694            ));
1695            let desc = query.desc.clone();
1696            (
1697                SubscribeFrom::Query {
1698                    expr: query.expr,
1699                    desc: query.desc,
1700                },
1701                desc,
1702                query.scope,
1703            )
1704        }
1705    };
1706
1707    let when = query::plan_as_of(scx, as_of)?;
1708    let up_to = up_to
1709        .map(|up_to| plan_as_of_or_up_to(scx, up_to))
1710        .transpose()?;
1711
1712    let qcx = QueryContext::root(scx, QueryLifetime::Subscribe);
1713    let ecx = ExprContext {
1714        qcx: &qcx,
1715        name: "",
1716        scope: &scope,
1717        relation_type: desc.typ(),
1718        allow_aggregates: false,
1719        allow_subqueries: true,
1720        allow_parameters: true,
1721        allow_windows: false,
1722    };
1723
1724    let output_columns: Vec<_> = scope.column_names().enumerate().collect();
1725    let output = match output {
1726        SubscribeOutput::Diffs => plan::SubscribeOutput::Diffs,
1727        SubscribeOutput::EnvelopeUpsert { key_columns } => {
1728            let order_by = key_columns
1729                .iter()
1730                .map(|ident| OrderByExpr {
1731                    expr: Expr::Identifier(vec![ident.clone()]),
1732                    asc: None,
1733                    nulls_last: None,
1734                })
1735                .collect_vec();
1736            let (order_by, map_exprs) = query::plan_order_by_exprs(
1737                &ExprContext {
1738                    name: "ENVELOPE UPSERT KEY clause",
1739                    ..ecx
1740                },
1741                &order_by[..],
1742                &output_columns[..],
1743            )?;
1744            if !map_exprs.is_empty() {
1745                return Err(PlanError::InvalidKeysInSubscribeEnvelopeUpsert);
1746            }
1747            check_distinct_key_columns(&order_by, &output_columns)?;
1748            plan::SubscribeOutput::EnvelopeUpsert {
1749                order_by_keys: order_by,
1750            }
1751        }
1752        SubscribeOutput::EnvelopeDebezium { key_columns } => {
1753            scx.require_feature_flag(&vars::ENABLE_ENVELOPE_DEBEZIUM_IN_SUBSCRIBE)?;
1754            let order_by = key_columns
1755                .iter()
1756                .map(|ident| OrderByExpr {
1757                    expr: Expr::Identifier(vec![ident.clone()]),
1758                    asc: None,
1759                    nulls_last: None,
1760                })
1761                .collect_vec();
1762            let (order_by, map_exprs) = query::plan_order_by_exprs(
1763                &ExprContext {
1764                    name: "ENVELOPE DEBEZIUM KEY clause",
1765                    ..ecx
1766                },
1767                &order_by[..],
1768                &output_columns[..],
1769            )?;
1770            if !map_exprs.is_empty() {
1771                return Err(PlanError::InvalidKeysInSubscribeEnvelopeDebezium);
1772            }
1773            check_distinct_key_columns(&order_by, &output_columns)?;
1774            plan::SubscribeOutput::EnvelopeDebezium {
1775                order_by_keys: order_by,
1776            }
1777        }
1778        SubscribeOutput::WithinTimestampOrderBy { order_by } => {
1779            scx.require_feature_flag(&vars::ENABLE_WITHIN_TIMESTAMP_ORDER_BY_IN_SUBSCRIBE)?;
1780            let mz_diff = "mz_diff".into();
1781            let output_columns = std::iter::once((0, &mz_diff))
1782                .chain(output_columns.into_iter().map(|(i, c)| (i + 1, c)))
1783                .collect_vec();
1784            match query::plan_order_by_exprs(
1785                &ExprContext {
1786                    name: "WITHIN TIMESTAMP ORDER BY clause",
1787                    ..ecx
1788                },
1789                &order_by[..],
1790                &output_columns[..],
1791            ) {
1792                Err(PlanError::UnknownColumn {
1793                    table: None,
1794                    column,
1795                    similar: _,
1796                }) if &column == &mz_diff => {
1797                    // mz_diff is being used in an expression. Since mz_diff isn't part of the table
1798                    // it looks like an unknown column. Instead, return a better error
1799                    return Err(PlanError::InvalidOrderByInSubscribeWithinTimestampOrderBy);
1800                }
1801                Err(e) => return Err(e),
1802                Ok((order_by, map_exprs)) => {
1803                    if !map_exprs.is_empty() {
1804                        return Err(PlanError::InvalidOrderByInSubscribeWithinTimestampOrderBy);
1805                    }
1806
1807                    plan::SubscribeOutput::WithinTimestampOrderBy { order_by }
1808                }
1809            }
1810        }
1811    };
1812
1813    let SubscribeOptionExtracted {
1814        progress, snapshot, ..
1815    } = options.try_into()?;
1816    Ok(Plan::Subscribe(SubscribePlan {
1817        from,
1818        when,
1819        up_to,
1820        with_snapshot: snapshot.unwrap_or(true),
1821        copy_to,
1822        emit_progress: progress.unwrap_or(false),
1823        output,
1824    }))
1825}
1826
1827/// Ensures each `ColumnOrder` in `order_by` references a distinct column,
1828/// returning `DuplicateKeyColumnInSubscribeEnvelope` on the first repeat.
1829fn check_distinct_key_columns(
1830    order_by: &[ColumnOrder],
1831    output_columns: &[(usize, &mz_repr::ColumnName)],
1832) -> Result<(), PlanError> {
1833    let mut seen = BTreeSet::new();
1834    for co in order_by {
1835        if !seen.insert(co.column) {
1836            return Err(PlanError::DuplicateKeyColumnInSubscribeEnvelope {
1837                column_name: output_columns[co.column].1.to_string(),
1838            });
1839        }
1840    }
1841    Ok(())
1842}
1843
1844pub fn describe_copy_from_table(
1845    scx: &StatementContext,
1846    table_name: <Aug as AstInfo>::ItemName,
1847    columns: Vec<Ident>,
1848) -> Result<StatementDesc, PlanError> {
1849    let (_, desc, _, _) = query::plan_copy_from(scx, table_name, columns)?;
1850    Ok(StatementDesc::new(Some(desc)))
1851}
1852
1853pub fn describe_copy_item(
1854    scx: &StatementContext,
1855    object_name: <Aug as AstInfo>::ItemName,
1856    columns: Vec<Ident>,
1857) -> Result<StatementDesc, PlanError> {
1858    let (_, desc, _, _) = query::plan_copy_item(scx, object_name, columns)?;
1859    Ok(StatementDesc::new(Some(desc)))
1860}
1861
1862pub fn describe_copy(
1863    scx: &StatementContext,
1864    CopyStatement {
1865        relation,
1866        direction,
1867        ..
1868    }: CopyStatement<Aug>,
1869) -> Result<StatementDesc, PlanError> {
1870    Ok(match (relation, direction) {
1871        (CopyRelation::Named { name, columns }, CopyDirection::To) => {
1872            describe_copy_item(scx, name, columns)?
1873        }
1874        (CopyRelation::Named { name, columns }, CopyDirection::From) => {
1875            describe_copy_from_table(scx, name, columns)?
1876        }
1877        (CopyRelation::Select(stmt), _) => describe_select(scx, stmt)?,
1878        (CopyRelation::Subscribe(stmt), _) => describe_subscribe(scx, stmt)?,
1879    }
1880    .with_is_copy())
1881}
1882
1883fn plan_copy_to_expr(
1884    scx: &StatementContext,
1885    select_plan: SelectPlan,
1886    desc: RelationDesc,
1887    to: &Expr<Aug>,
1888    format: CopyFormat,
1889    options: CopyOptionExtracted,
1890) -> Result<Plan, PlanError> {
1891    let conn_id = match options.aws_connection {
1892        Some(conn_id) => CatalogItemId::from(conn_id),
1893        None => sql_bail!("AWS CONNECTION is required for COPY ... TO <expr>"),
1894    };
1895    let connection = scx.get_item(&conn_id).connection()?;
1896
1897    match connection {
1898        mz_storage_types::connections::Connection::Aws(_) => {}
1899        _ => sql_bail!("only AWS CONNECTION is supported for COPY ... TO <expr>"),
1900    }
1901
1902    let format = match format {
1903        CopyFormat::Csv => {
1904            let quote = extract_byte_param_value(options.quote, "quote")?;
1905            let escape = extract_byte_param_value(options.escape, "escape")?;
1906            let delimiter = extract_byte_param_value(options.delimiter, "delimiter")?;
1907            S3SinkFormat::PgCopy(CopyFormatParams::Csv(
1908                CopyCsvFormatParams::try_new(
1909                    delimiter,
1910                    quote,
1911                    escape,
1912                    options.header,
1913                    options.null,
1914                )
1915                .map_err(|e| sql_err!("{}", e))?,
1916            ))
1917        }
1918        CopyFormat::Parquet => {
1919            // Validate that the output desc can be formatted as parquet.
1920            // COPY TO does not apply any type overrides, so pass `|_| None`.
1921            ArrowBuilder::validate_desc_for_parquet(&desc, |_| None)
1922                .map_err(|e| sql_err!("{}", e))?;
1923            S3SinkFormat::Parquet
1924        }
1925        CopyFormat::Binary => bail_unsupported!("FORMAT BINARY"),
1926        CopyFormat::Text => bail_unsupported!("FORMAT TEXT"),
1927    };
1928
1929    // Converting the to expr to a HirScalarExpr
1930    let mut to_expr = to.clone();
1931    transform_ast::transform(scx, &mut to_expr)?;
1932    let relation_type = RelationDesc::empty();
1933    let ecx = &ExprContext {
1934        qcx: &QueryContext::root(scx, QueryLifetime::OneShot),
1935        name: "COPY TO target",
1936        scope: &Scope::empty(),
1937        relation_type: relation_type.typ(),
1938        allow_aggregates: false,
1939        allow_subqueries: false,
1940        allow_parameters: false,
1941        allow_windows: false,
1942    };
1943
1944    let to = plan_expr(ecx, &to_expr)?.type_as(ecx, &SqlScalarType::String)?;
1945
1946    if options.max_file_size.as_bytes() < MIN_S3_SINK_FILE_SIZE.as_bytes() {
1947        sql_bail!(
1948            "MAX FILE SIZE cannot be less than {}",
1949            MIN_S3_SINK_FILE_SIZE
1950        );
1951    }
1952    if options.max_file_size.as_bytes() > MAX_S3_SINK_FILE_SIZE.as_bytes() {
1953        sql_bail!(
1954            "MAX FILE SIZE cannot be greater than {}",
1955            MAX_S3_SINK_FILE_SIZE
1956        );
1957    }
1958
1959    Ok(Plan::CopyTo(CopyToPlan {
1960        select_plan,
1961        desc,
1962        to,
1963        connection: connection.to_owned(),
1964        connection_id: conn_id,
1965        format,
1966        max_file_size: options.max_file_size.as_bytes(),
1967    }))
1968}
1969
1970fn plan_copy_from(
1971    scx: &StatementContext,
1972    target: &CopyTarget<Aug>,
1973    table_name: ResolvedItemName,
1974    columns: Vec<Ident>,
1975    format: Option<CopyFormat>,
1976    options: CopyOptionExtracted,
1977) -> Result<Plan, PlanError> {
1978    fn only_available_with_csv<T>(option: Option<T>, param: &str) -> Result<(), PlanError> {
1979        match option {
1980            Some(_) => sql_bail!("COPY {} available only in CSV mode", param),
1981            None => Ok(()),
1982        }
1983    }
1984
1985    let source = match target {
1986        CopyTarget::Stdin => CopyFromSource::Stdin,
1987        CopyTarget::Expr(from) => {
1988            // Converting the expr to an HirScalarExpr
1989            let mut from_expr = from.clone();
1990            transform_ast::transform(scx, &mut from_expr)?;
1991            let relation_type = RelationDesc::empty();
1992            let ecx = &ExprContext {
1993                qcx: &QueryContext::root(scx, QueryLifetime::OneShot),
1994                name: "COPY FROM target",
1995                scope: &Scope::empty(),
1996                relation_type: relation_type.typ(),
1997                allow_aggregates: false,
1998                allow_subqueries: false,
1999                allow_parameters: false,
2000                allow_windows: false,
2001            };
2002            let from = plan_expr(ecx, &from_expr)?.type_as(ecx, &SqlScalarType::String)?;
2003
2004            match options.aws_connection {
2005                Some(conn_id) => {
2006                    let conn_id = CatalogItemId::from(conn_id);
2007
2008                    // Validate the connection type is one we expect.
2009                    let connection = match scx.get_item(&conn_id).connection()? {
2010                        mz_storage_types::connections::Connection::Aws(conn) => conn,
2011                        _ => sql_bail!("only AWS CONNECTION is supported in COPY ... FROM"),
2012                    };
2013
2014                    CopyFromSource::AwsS3 {
2015                        uri: from,
2016                        connection,
2017                        connection_id: conn_id,
2018                    }
2019                }
2020                None => CopyFromSource::Url(from),
2021            }
2022        }
2023        CopyTarget::Stdout => bail_never_supported!("COPY FROM {} not supported", target),
2024    };
2025
2026    // COPY FROM a URL or S3 bucket only supports CSV and Parquet. Unlike COPY
2027    // FROM STDIN there's no sensible default format, so one must be specified
2028    // explicitly. Reject unsupported formats here in planning; the coordinator
2029    // relies on this and would otherwise soft-panic.
2030    let format = match &source {
2031        CopyFromSource::Stdin => format.unwrap_or(CopyFormat::Text),
2032        CopyFromSource::Url(_) | CopyFromSource::AwsS3 { .. } => match format {
2033            None => sql_bail!("COPY FROM <expr> requires a FORMAT option"),
2034            Some(CopyFormat::Text) => bail_unsupported!("FORMAT TEXT"),
2035            Some(CopyFormat::Binary) => bail_unsupported!("FORMAT BINARY"),
2036            Some(format @ (CopyFormat::Csv | CopyFormat::Parquet)) => format,
2037        },
2038    };
2039
2040    let params = match format {
2041        CopyFormat::Text => {
2042            only_available_with_csv(options.quote, "quote")?;
2043            only_available_with_csv(options.escape, "escape")?;
2044            only_available_with_csv(options.header, "HEADER")?;
2045            let delimiter =
2046                extract_byte_param_value(options.delimiter, "delimiter")?.unwrap_or(b'\t');
2047            let null = match options.null {
2048                Some(null) => Cow::from(null),
2049                None => Cow::from("\\N"),
2050            };
2051            CopyFormatParams::Text(CopyTextFormatParams { null, delimiter })
2052        }
2053        CopyFormat::Csv => {
2054            let quote = extract_byte_param_value(options.quote, "quote")?;
2055            let escape = extract_byte_param_value(options.escape, "escape")?;
2056            let delimiter = extract_byte_param_value(options.delimiter, "delimiter")?;
2057            CopyFormatParams::Csv(
2058                CopyCsvFormatParams::try_new(
2059                    delimiter,
2060                    quote,
2061                    escape,
2062                    options.header,
2063                    options.null,
2064                )
2065                .map_err(|e| sql_err!("{}", e))?,
2066            )
2067        }
2068        CopyFormat::Binary => bail_unsupported!("FORMAT BINARY"),
2069        CopyFormat::Parquet => CopyFormatParams::Parquet,
2070    };
2071
2072    let filter = match (options.files, options.pattern) {
2073        (Some(_), Some(_)) => bail_unsupported!("must specify one of FILES or PATTERN"),
2074        (Some(files), None) => Some(CopyFromFilter::Files(files)),
2075        (None, Some(pattern)) => Some(CopyFromFilter::Pattern(pattern)),
2076        (None, None) => None,
2077    };
2078
2079    if filter.is_some() && matches!(source, CopyFromSource::Stdin) {
2080        bail_unsupported!("COPY FROM ... WITH (FILES ...) only supported from a URL")
2081    }
2082
2083    let table_name_string = table_name.full_name_str();
2084
2085    let (id, source_desc, columns, maybe_mfp) = query::plan_copy_from(scx, table_name, columns)?;
2086
2087    let Some(mfp) = maybe_mfp else {
2088        sql_bail!("[internal error] COPY FROM ... expects an MFP to be produced");
2089    };
2090
2091    Ok(Plan::CopyFrom(CopyFromPlan {
2092        target_id: id,
2093        target_name: table_name_string,
2094        source,
2095        columns,
2096        source_desc,
2097        mfp,
2098        params,
2099        filter,
2100    }))
2101}
2102
2103fn extract_byte_param_value(v: Option<String>, param_name: &str) -> Result<Option<u8>, PlanError> {
2104    match v {
2105        Some(v) if v.len() == 1 => Ok(Some(v.as_bytes()[0])),
2106        Some(..) => sql_bail!("COPY {} must be a single one-byte character", param_name),
2107        None => Ok(None),
2108    }
2109}
2110
2111generate_extracted_config!(
2112    CopyOption,
2113    (Format, String),
2114    (Delimiter, String),
2115    (Null, String),
2116    (Escape, String),
2117    (Quote, String),
2118    (Header, bool),
2119    (AwsConnection, with_options::Object),
2120    (MaxFileSize, ByteSize, Default(ByteSize::mb(256))),
2121    (Files, Vec<String>),
2122    (Pattern, String)
2123);
2124
2125pub fn plan_copy(
2126    scx: &StatementContext,
2127    CopyStatement {
2128        relation,
2129        direction,
2130        target,
2131        options,
2132    }: CopyStatement<Aug>,
2133) -> Result<Plan, PlanError> {
2134    let options = CopyOptionExtracted::try_from(options)?;
2135    // Parse any user-provided FORMAT option. If not provided, will default to
2136    // Text for COPY TO STDOUT and COPY FROM STDIN, but will error for COPY TO <expr>.
2137    let format = options
2138        .format
2139        .as_ref()
2140        .map(|format| match format.to_lowercase().as_str() {
2141            "text" => Ok(CopyFormat::Text),
2142            "csv" => Ok(CopyFormat::Csv),
2143            "binary" => Ok(CopyFormat::Binary),
2144            "parquet" => Ok(CopyFormat::Parquet),
2145            _ => sql_bail!("unknown FORMAT: {}", format),
2146        })
2147        .transpose()?;
2148
2149    match (&direction, &target) {
2150        (CopyDirection::To, CopyTarget::Stdout) => {
2151            if options.delimiter.is_some() {
2152                sql_bail!("COPY TO does not support DELIMITER option yet");
2153            }
2154            if options.quote.is_some() {
2155                sql_bail!("COPY TO does not support QUOTE option yet");
2156            }
2157            if options.escape.is_some() {
2158                sql_bail!("COPY TO does not support ESCAPE option yet");
2159            }
2160            if options.null.is_some() {
2161                sql_bail!("COPY TO does not support NULL option yet");
2162            }
2163            // `HEADER false` is the default and already honored; only an
2164            // enabled header is unimplemented. Silently accepting it would
2165            // make clients strip the first data row as a presumed header.
2166            if options.header == Some(true) {
2167                sql_bail!("COPY TO does not support HEADER option yet");
2168            }
2169            match relation {
2170                CopyRelation::Named { .. } => sql_bail!("named with COPY TO STDOUT unsupported"),
2171                CopyRelation::Select(stmt) => Ok(plan_select(
2172                    scx,
2173                    stmt,
2174                    &Params::empty(),
2175                    Some(format.unwrap_or(CopyFormat::Text)),
2176                )?),
2177                CopyRelation::Subscribe(stmt) => Ok(plan_subscribe(
2178                    scx,
2179                    stmt,
2180                    &Params::empty(),
2181                    Some(format.unwrap_or(CopyFormat::Text)),
2182                )?),
2183            }
2184        }
2185        (CopyDirection::From, target) => match relation {
2186            CopyRelation::Named { name, columns } => {
2187                plan_copy_from(scx, target, name, columns, format, options)
2188            }
2189            _ => sql_bail!("COPY FROM {} not supported", target),
2190        },
2191        (CopyDirection::To, CopyTarget::Expr(to_expr)) => {
2192            let format = match format {
2193                Some(inner) => inner,
2194                _ => sql_bail!("COPY TO <expr> requires a FORMAT option"),
2195            };
2196
2197            let stmt = match relation {
2198                CopyRelation::Named { name, columns } => {
2199                    if !columns.is_empty() {
2200                        // TODO(mouli): Add support for this
2201                        sql_bail!(
2202                            "specifying columns for COPY <table_name> TO commands not yet supported; use COPY (SELECT...) TO ... instead"
2203                        );
2204                    }
2205                    // Generate a synthetic SELECT query that just gets the table
2206                    let query = Query {
2207                        ctes: CteBlock::empty(),
2208                        body: SetExpr::Table(name),
2209                        order_by: vec![],
2210                        limit: None,
2211                        offset: None,
2212                    };
2213                    SelectStatement { query, as_of: None }
2214                }
2215                CopyRelation::Select(stmt) => {
2216                    if !stmt.query.order_by.is_empty() {
2217                        sql_bail!("ORDER BY is not supported in SELECT query for COPY statements")
2218                    }
2219                    stmt
2220                }
2221                CopyRelation::Subscribe(_) => {
2222                    sql_bail!("COPY {} {} not supported", direction, target)
2223                }
2224            };
2225
2226            let (plan, desc) = plan_select_inner(scx, stmt, &Params::empty(), None)?;
2227            plan_copy_to_expr(scx, plan, desc, to_expr, format, options)
2228        }
2229        _ => sql_bail!("COPY {} {} not supported", direction, target),
2230    }
2231}