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