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