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
890const 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
912const 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 let mut ctes = Vec::with_capacity(5); 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 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 let mut ctes = Vec::with_capacity(4); 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 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 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; };
1209
1210 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 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 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 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 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 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 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; };
1361
1362 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 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 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 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 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 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 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 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 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)] 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 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 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
1827fn 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 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 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 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 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 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 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 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 sql_bail!(
2202 "specifying columns for COPY <table_name> TO commands not yet supported; use COPY (SELECT...) TO ... instead"
2203 );
2204 }
2205 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}