1use std::collections::{BTreeMap, BTreeSet};
11use std::time::Duration;
12
13use base64::prelude::*;
14use maplit::btreeset;
15use mz_catalog::builtin::{BUILTINS, BuiltinTable};
16use mz_catalog::durable::objects::{SystemObjectDescription, SystemObjectMapping};
17use mz_catalog::durable::{MOCK_AUTHENTICATION_NONCE_KEY, Transaction};
18use mz_catalog::memory::objects::{StateUpdate, StateUpdateKind};
19use mz_ore::collections::CollectionExt;
20use mz_ore::now::NowFn;
21use mz_persist_types::ShardId;
22use mz_proto::RustType;
23use mz_repr::{CatalogItemId, Diff, Timestamp};
24use mz_sql::ast::display::AstDisplay;
25use mz_sql::ast::{
26 CreateSinkOptionName, CreateViewStatement, CteBlock, DeferredItemName, IfExistsBehavior, Query,
27 SetExpr, SqlServerConfigOptionName, ViewDefinition,
28};
29use mz_sql::catalog::{CatalogItemType, SessionCatalog};
30use mz_sql::names::{FullItemName, QualifiedItemName};
31use mz_sql::normalize;
32use mz_sql::session::vars::{FORCE_SOURCE_TABLE_SYNTAX, Var, VarInput};
33use mz_sql_parser::ast::{Raw, Statement};
34use mz_storage_client::controller::StorageTxn;
35use mz_storage_types::sources::SourceExportStatementDetails;
36use mz_storage_types::sources::load_generator::LoadGeneratorOutput;
37use prost::Message;
38use semver::Version;
39use tracing::info;
40use uuid::Uuid;
41
42use crate::catalog::open::into_consolidatable_updates_startup;
44use crate::catalog::state::LocalExpressionCache;
45use crate::catalog::{BuiltinTableUpdate, CatalogState, ConnCatalog};
46use crate::coord::catalog_implications::parsed_state_updates::ParsedStateUpdate;
47
48const MIGRATION_VERSION_KEY: &str = "migration_version";
58
59pub(crate) fn get_migration_version(txn: &Transaction<'_>) -> Option<Version> {
60 txn.get_setting(MIGRATION_VERSION_KEY.into())
61 .map(|s| s.parse().expect("valid migration version"))
62}
63
64pub(crate) fn set_migration_version(
65 txn: &mut Transaction<'_>,
66 version: Version,
67) -> Result<(), mz_catalog::durable::CatalogError> {
68 txn.set_setting(MIGRATION_VERSION_KEY.into(), Some(version.to_string()))
69}
70
71fn rewrite_ast_items<F>(tx: &mut Transaction<'_>, mut f: F) -> Result<(), anyhow::Error>
72where
73 F: for<'a> FnMut(
74 &'a mut Transaction<'_>,
75 CatalogItemId,
76 &'a mut Statement<Raw>,
77 ) -> Result<(), anyhow::Error>,
78{
79 let mut updated_items = BTreeMap::new();
80
81 for mut item in tx.get_items() {
82 let mut stmt = mz_sql::parse::parse(&item.create_sql)?.into_element().ast;
83 f(tx, item.id, &mut stmt)?;
84
85 item.create_sql = stmt.to_ast_string_stable();
86
87 updated_items.insert(item.id, item);
88 }
89 tx.update_items(updated_items)?;
90 Ok(())
91}
92
93fn rewrite_items<F>(
94 tx: &mut Transaction<'_>,
95 cat: &ConnCatalog<'_>,
96 mut f: F,
97) -> Result<(), anyhow::Error>
98where
99 F: for<'a> FnMut(
100 &'a mut Transaction<'_>,
101 &'a &ConnCatalog<'_>,
102 CatalogItemId,
103 &'a mut Statement<Raw>,
104 ) -> Result<(), anyhow::Error>,
105{
106 let mut updated_items = BTreeMap::new();
107 let items = tx.get_items();
108 for mut item in items {
109 let mut stmt = mz_sql::parse::parse(&item.create_sql)?.into_element().ast;
110
111 f(tx, &cat, item.id, &mut stmt)?;
112
113 item.create_sql = stmt.to_ast_string_stable();
114
115 updated_items.insert(item.id, item);
116 }
117 tx.update_items(updated_items)?;
118 Ok(())
119}
120
121pub(crate) struct MigrateResult {
122 pub(crate) builtin_table_updates: Vec<BuiltinTableUpdate<&'static BuiltinTable>>,
123 pub(crate) catalog_updates: Vec<ParsedStateUpdate>,
124 pub(crate) post_item_updates: Vec<(StateUpdateKind, Timestamp, Diff)>,
125}
126
127pub(crate) async fn migrate(
131 state: &mut CatalogState,
132 tx: &mut Transaction<'_>,
133 local_expr_cache: &mut LocalExpressionCache,
134 item_updates: Vec<StateUpdate>,
135 _now: NowFn,
136 _boot_ts: Timestamp,
137) -> Result<MigrateResult, anyhow::Error> {
138 let catalog_version = get_migration_version(tx).unwrap_or(Version::new(0, 0, 0));
139
140 info!(
141 "migrating statements from catalog version {:?}",
142 catalog_version
143 );
144
145 rewrite_ast_items(tx, |tx, _id, stmt| {
146 ast_rewrite_create_sink_partition_strategy(stmt)?;
155 ast_rewrite_sql_server_constraints(stmt)?;
156 ast_rewrite_add_missing_index_ids(tx, stmt)?;
157 ast_rewrite_kafka_metadata_refresh_intervals(stmt)?;
158 ast_rewrite_small_commit_intervals(stmt)?;
159 ast_rewrite_strip_builtin_version_pins(stmt)?;
160 Ok(())
161 })?;
162
163 let commit_ts = tx.upper();
166 let mut item_updates = into_consolidatable_updates_startup(item_updates, commit_ts);
167 let op_item_updates = tx.get_and_commit_op_updates();
168 let op_item_updates = into_consolidatable_updates_startup(op_item_updates, commit_ts);
169 item_updates.extend(op_item_updates);
170 differential_dataflow::consolidation::consolidate_updates(&mut item_updates);
171
172 let (post_item_updates, item_updates): (Vec<_>, Vec<_>) = item_updates
176 .into_iter()
177 .partition(|(kind, _, _)| matches!(kind, StateUpdateKind::StorageCollectionMetadata(_)));
180
181 let item_updates = item_updates
182 .into_iter()
183 .map(|(kind, ts, diff)| StateUpdate {
184 kind,
185 ts,
186 diff: diff.try_into().expect("valid diff"),
187 })
188 .collect();
189
190 let force_source_table_syntax = state.system_config().force_source_table_syntax();
191 if force_source_table_syntax {
195 state
196 .system_config_mut()
197 .set(FORCE_SOURCE_TABLE_SYNTAX.name(), VarInput::Flat("off"))
198 .expect("known parameter");
199 }
200
201 let (mut ast_builtin_table_updates, mut ast_catalog_updates) =
202 state.apply_updates(item_updates, local_expr_cache).await;
203
204 info!("migrating from catalog version {:?}", catalog_version);
205
206 let conn_cat = state.for_system_session();
207
208 if force_source_table_syntax {
211 rewrite_sources_to_tables(tx, &conn_cat)?;
212 }
213
214 rewrite_items(tx, &conn_cat, |_tx, _conn_cat, _id, _stmt| {
215 let _catalog_version = catalog_version.clone();
216 Ok(())
231 })?;
232
233 if force_source_table_syntax {
234 state
235 .system_config_mut()
236 .set(FORCE_SOURCE_TABLE_SYNTAX.name(), VarInput::Flat("on"))
237 .expect("known parameter");
238 }
239
240 let op_item_updates = tx.get_and_commit_op_updates();
246 let (item_builtin_table_updates, item_catalog_updates) =
247 state.apply_updates(op_item_updates, local_expr_cache).await;
248
249 ast_builtin_table_updates.extend(item_builtin_table_updates);
250 ast_catalog_updates.extend(item_catalog_updates);
251
252 info!(
253 "migration from catalog version {:?} complete",
254 catalog_version
255 );
256
257 Ok(MigrateResult {
258 builtin_table_updates: ast_builtin_table_updates,
259 catalog_updates: ast_catalog_updates,
260 post_item_updates,
261 })
262}
263
264fn rewrite_sources_to_tables(
343 tx: &mut Transaction<'_>,
344 catalog: &ConnCatalog<'_>,
345) -> Result<(), anyhow::Error> {
346 use mz_sql::ast::{
347 CreateSourceConnection, CreateSourceStatement, CreateSubsourceOptionName,
348 CreateSubsourceStatement, CreateTableFromSourceStatement, Ident,
349 KafkaSourceConfigOptionName, LoadGenerator, MySqlConfigOptionName, PgConfigOptionName,
350 RawItemName, TableFromSourceColumns, TableFromSourceOption, TableFromSourceOptionName,
351 UnresolvedItemName, Value, WithOptionValue,
352 };
353
354 let mut updated_items = BTreeMap::new();
355
356 let mut sources = vec![];
357 let mut subsources = vec![];
358
359 for item in tx.get_items() {
360 let stmt = mz_sql::parse::parse(&item.create_sql)?.into_element().ast;
361 match stmt {
362 Statement::CreateSubsource(stmt) => subsources.push((item, stmt)),
363 Statement::CreateSource(stmt) => sources.push((item, stmt)),
364 _ => {}
365 }
366 }
367
368 let mut pending_progress_items = BTreeMap::new();
369 let mut migrated_source_ids = BTreeMap::new();
370 for (mut source_item, source_stmt) in sources {
373 let CreateSourceStatement {
374 name,
375 in_cluster,
376 col_names,
377 mut connection,
378 include_metadata,
379 format,
380 envelope,
381 if_not_exists,
382 key_constraint,
383 with_options,
384 external_references,
385 progress_subsource,
386 } = source_stmt;
387
388 let (progress_name, progress_item) = match progress_subsource {
389 Some(DeferredItemName::Named(RawItemName::Name(name))) => {
390 let partial_name = normalize::unresolved_item_name(name.clone())?;
391 (name, catalog.resolve_item(&partial_name)?)
392 }
393 Some(DeferredItemName::Named(RawItemName::Id(id, name, _))) => {
394 let gid = id.parse()?;
395 (name, catalog.get_item(&gid))
396 }
397 Some(DeferredItemName::Deferred(_)) => {
398 unreachable!("invalid progress subsource")
399 }
400 None => {
401 info!("migrate: skipping already migrated source: {name}");
402 continue;
403 }
404 };
405 let raw_progress_name =
406 RawItemName::Id(progress_item.id().to_string(), progress_name.clone(), None);
407
408 let catalog_item = catalog.get_item(&source_item.id);
410 let source_name: &QualifiedItemName = catalog_item.name();
411 let full_source_name: FullItemName = catalog.resolve_full_name(source_name);
412 let source_name: UnresolvedItemName = normalize::unresolve(full_source_name.clone());
413
414 match &mut connection {
416 CreateSourceConnection::Postgres { options, .. } => {
417 options.retain(|o| match o.name {
418 PgConfigOptionName::Details | PgConfigOptionName::Publication => true,
419 PgConfigOptionName::TextColumns | PgConfigOptionName::ExcludeColumns => false,
420 });
421 }
422 CreateSourceConnection::SqlServer { options, .. } => {
423 options.retain(|o| match o.name {
424 SqlServerConfigOptionName::Details => true,
425 SqlServerConfigOptionName::TextColumns
426 | SqlServerConfigOptionName::ExcludeColumns => false,
427 });
428 }
429 CreateSourceConnection::MySql { options, .. } => {
430 options.retain(|o| match o.name {
431 MySqlConfigOptionName::Details => true,
432 MySqlConfigOptionName::TextColumns | MySqlConfigOptionName::ExcludeColumns => {
433 false
434 }
435 });
436 }
437 CreateSourceConnection::Kafka { .. } | CreateSourceConnection::LoadGenerator { .. } => {
438 }
439 }
440
441 let (new_progress_name, new_progress_stmt, new_source_name, new_source_stmt) =
443 match connection {
444 connection @ (CreateSourceConnection::Postgres { .. }
445 | CreateSourceConnection::MySql { .. }
446 | CreateSourceConnection::SqlServer { .. }
447 | CreateSourceConnection::LoadGenerator {
448 generator:
449 LoadGenerator::Tpch | LoadGenerator::Auction | LoadGenerator::Marketing,
450 ..
451 }) => {
452 assert_eq!(col_names, &[]);
454 assert_eq!(key_constraint, None);
455 assert_eq!(format, None);
456 assert_eq!(envelope, None);
457 assert_eq!(include_metadata, &[]);
458 assert_eq!(external_references, None);
459
460 let dummy_source_stmt = Statement::CreateView(CreateViewStatement {
468 if_exists: IfExistsBehavior::Error,
469 temporary: false,
470 definition: ViewDefinition {
471 name: progress_name,
472 columns: vec![],
473 query: Query {
474 ctes: CteBlock::Simple(vec![]),
475 body: SetExpr::Table(RawItemName::Id(
476 progress_item.id().to_string(),
477 source_name.clone(),
478 None,
479 )),
480 order_by: vec![],
481 limit: None,
482 offset: None,
483 },
484 },
485 });
486
487 let new_progress_stmt = CreateSourceStatement {
488 name: source_name.clone(),
489 in_cluster,
490 col_names: vec![],
491 connection,
492 include_metadata: vec![],
493 format: None,
494 envelope: None,
495 if_not_exists,
496 key_constraint: None,
497 with_options,
498 external_references: None,
499 progress_subsource: None,
500 };
501
502 migrated_source_ids.insert(source_item.id, progress_item.id());
503
504 (
505 full_source_name.item,
506 new_progress_stmt,
507 progress_item.name().item.clone(),
508 dummy_source_stmt,
509 )
510 }
511 CreateSourceConnection::Kafka {
512 options,
513 connection,
514 } => {
515 let constraints = if let Some(_key_constraint) = key_constraint {
516 vec![]
519 } else {
520 vec![]
521 };
522
523 let columns = if col_names.is_empty() {
524 TableFromSourceColumns::NotSpecified
525 } else {
526 TableFromSourceColumns::Named(col_names)
527 };
528
529 let details = SourceExportStatementDetails::Kafka {};
532 let table_with_options = vec![TableFromSourceOption {
533 name: TableFromSourceOptionName::Details,
534 value: Some(WithOptionValue::Value(Value::String(hex::encode(
535 details.into_proto().encode_to_vec(),
536 )))),
537 }];
538 let topic_option = options
540 .iter()
541 .find(|o| matches!(o.name, KafkaSourceConfigOptionName::Topic))
542 .expect("kafka sources must have a topic");
543 let topic = match &topic_option.value {
544 Some(WithOptionValue::Value(Value::String(topic))) => topic,
545 _ => unreachable!("topic must be a string"),
546 };
547 let external_reference = UnresolvedItemName::qualified(&[Ident::new(topic)?]);
548
549 let new_source_stmt =
550 Statement::CreateTableFromSource(CreateTableFromSourceStatement {
551 name: source_name,
552 constraints,
553 columns,
554 if_not_exists,
555 source: raw_progress_name,
556 include_metadata,
557 format,
558 envelope,
559 external_reference: Some(external_reference),
560 with_options: table_with_options,
561 });
562
563 let new_progress_stmt = CreateSourceStatement {
564 name: progress_name,
565 in_cluster,
566 col_names: vec![],
567 connection: CreateSourceConnection::Kafka {
568 options,
569 connection,
570 },
571 include_metadata: vec![],
572 format: None,
573 envelope: None,
574 if_not_exists,
575 key_constraint: None,
576 with_options,
577 external_references: None,
578 progress_subsource: None,
579 };
580 (
581 progress_item.name().item.clone(),
582 new_progress_stmt,
583 full_source_name.item,
584 new_source_stmt,
585 )
586 }
587 CreateSourceConnection::LoadGenerator {
588 generator:
589 generator @ (LoadGenerator::Clock
590 | LoadGenerator::Counter
591 | LoadGenerator::Datums
592 | LoadGenerator::KeyValue),
593 options,
594 } => {
595 let constraints = if let Some(_key_constraint) = key_constraint {
596 vec![]
598 } else {
599 vec![]
600 };
601
602 let columns = if col_names.is_empty() {
603 TableFromSourceColumns::NotSpecified
604 } else {
605 TableFromSourceColumns::Named(col_names)
606 };
607
608 let details = SourceExportStatementDetails::LoadGenerator {
611 output: LoadGeneratorOutput::Default,
612 };
613 let table_with_options = vec![TableFromSourceOption {
614 name: TableFromSourceOptionName::Details,
615 value: Some(WithOptionValue::Value(Value::String(hex::encode(
616 details.into_proto().encode_to_vec(),
617 )))),
618 }];
619 let external_reference = FullItemName {
622 database: mz_sql::names::RawDatabaseSpecifier::Name(
623 mz_storage_types::sources::load_generator::LOAD_GENERATOR_DATABASE_NAME
624 .to_owned(),
625 ),
626 schema: generator.schema_name().to_string(),
627 item: generator.schema_name().to_string(),
628 };
629
630 let new_source_stmt =
631 Statement::CreateTableFromSource(CreateTableFromSourceStatement {
632 name: source_name,
633 constraints,
634 columns,
635 if_not_exists,
636 source: raw_progress_name,
637 include_metadata,
638 format,
639 envelope,
640 external_reference: Some(external_reference.into()),
641 with_options: table_with_options,
642 });
643
644 let new_progress_stmt = CreateSourceStatement {
645 name: progress_name,
646 in_cluster,
647 col_names: vec![],
648 connection: CreateSourceConnection::LoadGenerator { generator, options },
649 include_metadata: vec![],
650 format: None,
651 envelope: None,
652 if_not_exists,
653 key_constraint: None,
654 with_options,
655 external_references: None,
656 progress_subsource: None,
657 };
658 (
659 progress_item.name().item.clone(),
660 new_progress_stmt,
661 full_source_name.item,
662 new_source_stmt,
663 )
664 }
665 };
666
667 info!(
671 "migrate: converted source {} to {}",
672 source_item.create_sql, new_source_stmt
673 );
674 source_item.name = new_source_name.clone();
675 source_item.create_sql = new_source_stmt.to_ast_string_stable();
676 updated_items.insert(source_item.id, source_item);
677 pending_progress_items.insert(progress_item.id(), (new_progress_name, new_progress_stmt));
678 }
679
680 for (mut item, stmt) in subsources {
681 match stmt {
682 CreateSubsourceStatement {
685 of_source: None, ..
686 } => {
687 let Some((new_name, new_stmt)) = pending_progress_items.remove(&item.id) else {
688 panic!("encountered orphan progress subsource id: {}", item.id)
689 };
690 item.name = new_name;
691 item.create_sql = new_stmt.to_ast_string_stable();
692 updated_items.insert(item.id, item);
693 }
694 CreateSubsourceStatement {
697 name,
698 columns,
699 constraints,
700 of_source: Some(raw_source_name),
701 if_not_exists,
702 mut with_options,
703 } => {
704 let new_raw_source_name = match raw_source_name {
705 RawItemName::Id(old_id, name, None) => {
706 let old_id: CatalogItemId = old_id.parse().expect("well formed");
707 let new_id = migrated_source_ids[&old_id].clone();
708 RawItemName::Id(new_id.to_string(), name, None)
709 }
710 _ => unreachable!("unexpected source name: {raw_source_name}"),
711 };
712 let external_reference = match with_options
715 .iter()
716 .position(|opt| opt.name == CreateSubsourceOptionName::ExternalReference)
717 {
718 Some(i) => match with_options.remove(i).value {
719 Some(WithOptionValue::UnresolvedItemName(name)) => name,
720 _ => unreachable!("external reference must be an unresolved item name"),
721 },
722 None => panic!("subsource must have an external reference"),
723 };
724
725 let with_options = with_options
726 .into_iter()
727 .map(|option| {
728 match option.name {
729 CreateSubsourceOptionName::Details => TableFromSourceOption {
730 name: TableFromSourceOptionName::Details,
731 value: option.value,
734 },
735 CreateSubsourceOptionName::TextColumns => TableFromSourceOption {
736 name: TableFromSourceOptionName::TextColumns,
737 value: option.value,
738 },
739 CreateSubsourceOptionName::ExcludeColumns => TableFromSourceOption {
740 name: TableFromSourceOptionName::ExcludeColumns,
741 value: option.value,
742 },
743 CreateSubsourceOptionName::RetainHistory => TableFromSourceOption {
744 name: TableFromSourceOptionName::RetainHistory,
745 value: option.value,
746 },
747 CreateSubsourceOptionName::Progress => {
748 panic!("progress option should not exist on this subsource")
749 }
750 CreateSubsourceOptionName::ExternalReference => {
751 unreachable!("This option is handled separately above.")
752 }
753 }
754 })
755 .collect::<Vec<_>>();
756
757 let table = CreateTableFromSourceStatement {
758 name,
759 constraints,
760 columns: TableFromSourceColumns::Defined(columns),
761 if_not_exists,
762 source: new_raw_source_name,
763 external_reference: Some(external_reference),
764 with_options,
765 envelope: None,
767 include_metadata: vec![],
768 format: None,
769 };
770
771 info!(
772 "migrate: converted subsource {} to table {}",
773 item.create_sql, table
774 );
775 item.create_sql = Statement::CreateTableFromSource(table).to_ast_string_stable();
776 updated_items.insert(item.id, item);
777 }
778 }
779 }
780 assert!(
781 pending_progress_items.is_empty(),
782 "unexpected residual progress items: {pending_progress_items:?}"
783 );
784
785 tx.update_items(updated_items)?;
786
787 Ok(())
788}
789
790pub(crate) fn durable_migrate(
794 tx: &mut Transaction,
795 _organization_id: Uuid,
796 _boot_ts: Timestamp,
797) -> Result<(), anyhow::Error> {
798 const EXPR_CACHE_MIGRATION_KEY: &str = "expr_cache_migration";
801 const EXPR_CACHE_MIGRATION_DONE: u64 = 1;
802 if tx.get_config(EXPR_CACHE_MIGRATION_KEY.to_string()) != Some(EXPR_CACHE_MIGRATION_DONE) {
803 if let Some(shard_id) = tx.get_expression_cache_shard() {
804 tx.insert_unfinalized_shards(btreeset! {shard_id})?;
805 tx.set_expression_cache_shard(ShardId::new())?;
806 }
807 tx.set_config(
808 EXPR_CACHE_MIGRATION_KEY.to_string(),
809 Some(EXPR_CACHE_MIGRATION_DONE),
810 )?;
811 }
812
813 const BUILTIN_MIGRATION_SHARD_MIGRATION_KEY: &str = "migration_shard_migration";
816 const BUILTIN_MIGRATION_SHARD_MIGRATION_DONE: u64 = 1;
817 if tx.get_config(BUILTIN_MIGRATION_SHARD_MIGRATION_KEY.to_string())
818 != Some(BUILTIN_MIGRATION_SHARD_MIGRATION_DONE)
819 {
820 if let Some(shard_id) = tx.get_builtin_migration_shard() {
821 tx.insert_unfinalized_shards(btreeset! {shard_id})?;
822 tx.set_builtin_migration_shard(ShardId::new())?;
823 }
824 tx.set_config(
825 BUILTIN_MIGRATION_SHARD_MIGRATION_KEY.to_string(),
826 Some(BUILTIN_MIGRATION_SHARD_MIGRATION_DONE),
827 )?;
828 }
829
830 if tx
831 .get_setting(MOCK_AUTHENTICATION_NONCE_KEY.to_string())
832 .is_none()
833 {
834 let mut nonce = [0u8; 24];
835 openssl::rand::rand_bytes(&mut nonce).expect("failed to generate nonce");
836 let nonce = BASE64_STANDARD.encode(nonce);
837 tx.set_setting(MOCK_AUTHENTICATION_NONCE_KEY.to_string(), Some(nonce))?;
838 }
839
840 migrate_builtin_tables_to_mvs(tx)?;
841
842 Ok(())
843}
844
845fn migrate_builtin_tables_to_mvs(tx: &mut Transaction) -> Result<(), anyhow::Error> {
849 let expected_mvs: BTreeSet<_> = BUILTINS::materialized_views()
851 .map(|mv| (mv.schema, mv.name))
852 .collect();
853
854 let mut to_remove = BTreeSet::new();
856 let mut to_add = Vec::new();
857 for mapping in tx.get_system_object_mappings() {
858 let desc = &mapping.description;
859 if desc.object_type != CatalogItemType::Table {
860 continue;
861 }
862
863 let key = (&*desc.schema_name, &*desc.object_name);
864 if expected_mvs.contains(&key) {
865 info!(
866 "migrate: builtin {}.{} changed type from table to MV",
867 desc.schema_name, desc.object_name,
868 );
869 to_remove.insert(desc.clone());
870 to_add.push(SystemObjectMapping {
871 description: SystemObjectDescription {
872 schema_name: desc.schema_name.clone(),
873 object_type: CatalogItemType::MaterializedView,
874 object_name: desc.object_name.clone(),
875 },
876 unique_identifier: mapping.unique_identifier,
877 });
878 }
879 }
880
881 if !to_remove.is_empty() {
882 tx.remove_system_object_mappings(to_remove)?;
883 tx.set_system_object_mappings(to_add)?;
884 }
885
886 Ok(())
887}
888
889fn ast_rewrite_create_sink_partition_strategy(
901 stmt: &mut Statement<Raw>,
902) -> Result<(), anyhow::Error> {
903 let Statement::CreateSink(stmt) = stmt else {
904 return Ok(());
905 };
906 stmt.with_options
907 .retain(|op| op.name != CreateSinkOptionName::PartitionStrategy);
908 Ok(())
909}
910
911fn ast_rewrite_sql_server_constraints(stmt: &mut Statement<Raw>) -> Result<(), anyhow::Error> {
913 use mz_sql::ast::{
914 CreateSubsourceOptionName, TableFromSourceOptionName, Value, WithOptionValue,
915 };
916 use mz_sql_server_util::desc::{SqlServerTableConstraint, SqlServerTableConstraintType};
917 use mz_storage_types::sources::ProtoSourceExportStatementDetails;
918 use mz_storage_types::sources::proto_source_export_statement_details::Kind;
919
920 let deets: Option<&mut String> = match stmt {
921 Statement::CreateSubsource(stmt) => stmt.with_options.iter_mut().find_map(|option| {
922 if matches!(option.name, CreateSubsourceOptionName::Details)
923 && let Some(WithOptionValue::Value(Value::String(ref mut details))) = option.value
924 {
925 Some(details)
926 } else {
927 None
928 }
929 }),
930 Statement::CreateTableFromSource(stmt) => stmt.with_options.iter_mut().find_map(|option| {
931 if matches!(option.name, TableFromSourceOptionName::Details)
932 && let Some(WithOptionValue::Value(Value::String(ref mut details))) = option.value
933 {
934 Some(details)
935 } else {
936 None
937 }
938 }),
939 _ => None,
940 };
941 let Some(deets) = deets else {
942 return Ok(());
943 };
944
945 let current_value = hex::decode(&mut *deets)?;
946 let current_value = ProtoSourceExportStatementDetails::decode(&*current_value)?;
947
948 if !matches!(current_value.kind, Some(Kind::SqlServer(_))) {
950 return Ok(());
951 };
952
953 let SourceExportStatementDetails::SqlServer {
954 mut table,
955 capture_instance,
956 initial_lsn,
957 } = SourceExportStatementDetails::from_proto(current_value)?
958 else {
959 unreachable!("statement details must exist for SQL Server");
960 };
961
962 if !table.constraints.is_empty() {
964 return Ok(());
965 }
966
967 let mut migrated_constraints: BTreeMap<_, Vec<_>> = BTreeMap::new();
970 for col in table.columns.iter_mut() {
971 if let Some(constraint_name) = col.primary_key_constraint.take() {
972 migrated_constraints
973 .entry(constraint_name)
974 .or_default()
975 .push(col.name.to_string());
976 }
977 }
978
979 table.constraints = migrated_constraints
980 .into_iter()
981 .map(|(constraint_name, column_names)| SqlServerTableConstraint {
982 constraint_name: constraint_name.to_string(),
983 constraint_type: SqlServerTableConstraintType::PrimaryKey,
984 column_names,
985 })
986 .collect();
987
988 let new_value = SourceExportStatementDetails::SqlServer {
989 table,
990 capture_instance,
991 initial_lsn,
992 };
993 *deets = hex::encode(new_value.into_proto().encode_to_vec());
994
995 Ok(())
996}
997
998fn ast_rewrite_add_missing_index_ids(
1000 tx: &Transaction<'_>,
1001 stmt: &mut Statement<Raw>,
1002) -> Result<(), anyhow::Error> {
1003 let Statement::CreateIndex(stmt) = stmt else {
1004 return Ok(());
1005 };
1006
1007 let unresolved_name = match stmt.on_name.clone() {
1008 mz_sql::ast::RawItemName::Name(name) => name,
1009 mz_sql::ast::RawItemName::Id(..) => return Ok(()),
1011 };
1012
1013 let parts = &unresolved_name.0;
1014 let (db_name, schema_name, item_name) = match parts.len() {
1015 3 => (Some(&parts[0]), &parts[1], &parts[2]),
1016 2 => (None, &parts[0], &parts[1]),
1017 _ => panic!("invalid unresolved name: {unresolved_name:?}"),
1018 };
1019
1020 let db_id = db_name.map(|x| {
1021 let db = tx.get_databases().find(|db| db.name == x.as_str());
1022 let db = db.unwrap_or_else(|| panic!("missing database: {x}"));
1023 db.id
1024 });
1025 let schema_id = {
1026 let schema = tx
1027 .get_schemas()
1028 .find(|s| s.name == schema_name.as_str() && s.database_id == db_id);
1029 let schema = schema.unwrap_or_else(|| panic!("missing schema: {schema_name}, {db_id:?}"));
1030 schema.id
1031 };
1032 let item_id = {
1033 let item = tx
1034 .get_items()
1035 .find(|i| i.name == item_name.as_str() && i.schema_id == schema_id);
1036 let item = item.unwrap_or_else(|| panic!("missing item: {item_name}, {schema_id:?}"));
1037 item.id
1038 };
1039
1040 stmt.on_name = mz_sql::ast::RawItemName::Id(item_id.to_string(), unresolved_name, None);
1041
1042 Ok(())
1043}
1044
1045fn ast_rewrite_strip_builtin_version_pins(stmt: &mut Statement<Raw>) -> Result<(), anyhow::Error> {
1059 use mz_sql::ast::RawItemName;
1060 use mz_sql::ast::visit_mut::{VisitMut, VisitMutNode};
1061
1062 struct StripBuiltinVersionPins;
1063
1064 impl<'ast> VisitMut<'ast, Raw> for StripBuiltinVersionPins {
1065 fn visit_item_name_mut(&mut self, item_name: &mut RawItemName) {
1066 if let RawItemName::Id(id, _, version) = item_name {
1067 if version.is_some() {
1068 if let Ok(parsed) = id.parse::<CatalogItemId>() {
1069 if !parsed.is_user() {
1070 *version = None;
1071 }
1072 }
1073 }
1074 }
1075 }
1076 }
1077
1078 let mut visitor = StripBuiltinVersionPins;
1079 stmt.visit_mut(&mut visitor);
1080 Ok(())
1081}
1082
1083fn ast_rewrite_kafka_metadata_refresh_intervals(
1084 stmt: &mut Statement<Raw>,
1085) -> Result<(), anyhow::Error> {
1086 use mz_sql::ast::{
1087 CreateSinkConnection, CreateSourceConnection, KafkaSinkConfigOptionName,
1088 KafkaSourceConfigOptionName, WithOptionValue,
1089 };
1090 let interval: Option<&mut WithOptionValue<Raw>> = match stmt {
1095 Statement::CreateSource(stmt) => {
1096 if let CreateSourceConnection::Kafka { options, .. } = &mut stmt.connection {
1097 options.iter_mut().find_map(|option| {
1098 if matches!(
1099 option.name,
1100 KafkaSourceConfigOptionName::TopicMetadataRefreshInterval
1101 ) {
1102 option.value.as_mut()
1103 } else {
1104 None
1105 }
1106 })
1107 } else {
1108 None
1109 }
1110 }
1111 Statement::CreateSink(stmt) => {
1112 if let CreateSinkConnection::Kafka { options, .. } = &mut stmt.connection {
1113 options.iter_mut().find_map(|option| {
1114 if matches!(
1115 option.name,
1116 KafkaSinkConfigOptionName::TopicMetadataRefreshInterval
1117 ) {
1118 option.value.as_mut()
1119 } else {
1120 None
1121 }
1122 })
1123 } else {
1124 None
1125 }
1126 }
1127 _ => None,
1128 };
1129
1130 let Some(interval) = interval else {
1131 return Ok(());
1132 };
1133
1134 rewrite_interval_option_floor_1s(interval, "kafka metadata refresh interval")
1135}
1136
1137fn ast_rewrite_small_commit_intervals(stmt: &mut Statement<Raw>) -> Result<(), anyhow::Error> {
1145 use mz_sql::ast::{CreateSinkConnection, CreateSinkOptionName};
1146
1147 let Statement::CreateSink(stmt) = stmt else {
1148 return Ok(());
1149 };
1150 if !matches!(stmt.connection, CreateSinkConnection::Iceberg { .. }) {
1151 return Ok(());
1152 }
1153 let interval = stmt.with_options.iter_mut().find_map(|o| {
1154 if matches!(o.name, CreateSinkOptionName::CommitInterval) {
1155 o.value.as_mut()
1156 } else {
1157 None
1158 }
1159 });
1160 let Some(interval) = interval else {
1161 return Ok(());
1162 };
1163
1164 rewrite_interval_option_floor_1s(interval, "commit interval")
1165}
1166
1167fn rewrite_interval_option_floor_1s(
1169 value: &mut mz_sql::ast::WithOptionValue<Raw>,
1170 label: &str,
1171) -> Result<(), anyhow::Error> {
1172 use mz_sql::ast::{Value, WithOptionValue};
1173 use mz_sql::plan::TryFromValue;
1174
1175 let dur = Duration::try_from_value(value.clone())
1176 .map_err(|e| anyhow::anyhow!("invalid value for {label}: {value:?}: {e}"))?;
1177
1178 if dur < Duration::from_secs(1) {
1179 *value = WithOptionValue::Value(Value::String("1s".to_string()));
1180 }
1181
1182 Ok(())
1183}
1184
1185#[cfg(test)]
1186mod tests {
1187 use super::*;
1188
1189 fn strip(sql: &str) -> String {
1190 let mut stmt = mz_sql::parse::parse(sql)
1191 .expect("test sql parses")
1192 .into_element()
1193 .ast;
1194 ast_rewrite_strip_builtin_version_pins(&mut stmt).expect("rewrite succeeds");
1195 stmt.to_ast_string_stable()
1196 }
1197
1198 #[mz_ore::test]
1199 #[cfg_attr(miri, ignore)] fn strips_version_from_builtin_reference() {
1201 let out = strip(
1205 r#"CREATE VIEW "materialize"."public"."v" AS SELECT 1 FROM [s518 AS "mz_catalog"."mz_audit_events" VERSION 0]"#,
1206 );
1207 assert!(!out.contains("VERSION"), "version not stripped: {out}");
1208 assert!(out.contains("s518"), "reference dropped: {out}");
1210 }
1211
1212 #[mz_ore::test]
1213 #[cfg_attr(miri, ignore)] fn preserves_version_on_user_reference() {
1215 let out = strip(
1218 r#"CREATE VIEW "materialize"."public"."v" AS SELECT 1 FROM [u5 AS "materialize"."public"."t" VERSION 1]"#,
1219 );
1220 assert!(out.contains("VERSION"), "user version stripped: {out}");
1221 assert!(out.contains("u5"), "reference dropped: {out}");
1222 }
1223
1224 #[mz_ore::test]
1225 #[cfg_attr(miri, ignore)] fn leaves_unpinned_builtin_reference_untouched() {
1227 let out = strip(
1228 r#"CREATE VIEW "materialize"."public"."v" AS SELECT 1 FROM [s518 AS "mz_catalog"."mz_audit_events"]"#,
1229 );
1230 assert!(!out.contains("VERSION"), "unexpected version: {out}");
1231 assert!(out.contains("s518"), "reference dropped: {out}");
1232 }
1233}