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::{BootstrapStateUpdateKind, StateUpdate};
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<(BootstrapStateUpdateKind, 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, _, _)| {
180 matches!(kind, BootstrapStateUpdateKind::StorageCollectionMetadata(_))
181 });
182
183 let item_updates = item_updates
184 .into_iter()
185 .map(|(kind, ts, diff)| StateUpdate {
186 kind: kind.into(),
187 ts,
188 diff: diff.try_into().expect("valid diff"),
189 })
190 .collect();
191
192 let force_source_table_syntax = state.system_config().force_source_table_syntax();
193 if force_source_table_syntax {
197 state
198 .system_config_mut()
199 .set(FORCE_SOURCE_TABLE_SYNTAX.name(), VarInput::Flat("off"))
200 .expect("known parameter");
201 }
202
203 let (mut ast_builtin_table_updates, mut ast_catalog_updates) =
204 state.apply_updates(item_updates, local_expr_cache).await;
205
206 info!("migrating from catalog version {:?}", catalog_version);
207
208 let conn_cat = state.for_system_session();
209
210 if force_source_table_syntax {
213 rewrite_sources_to_tables(tx, &conn_cat)?;
214 }
215
216 rewrite_items(tx, &conn_cat, |_tx, _conn_cat, _id, _stmt| {
217 let _catalog_version = catalog_version.clone();
218 Ok(())
233 })?;
234
235 if force_source_table_syntax {
236 state
237 .system_config_mut()
238 .set(FORCE_SOURCE_TABLE_SYNTAX.name(), VarInput::Flat("on"))
239 .expect("known parameter");
240 }
241
242 let op_item_updates = tx.get_and_commit_op_updates();
248 let (item_builtin_table_updates, item_catalog_updates) =
249 state.apply_updates(op_item_updates, local_expr_cache).await;
250
251 ast_builtin_table_updates.extend(item_builtin_table_updates);
252 ast_catalog_updates.extend(item_catalog_updates);
253
254 info!(
255 "migration from catalog version {:?} complete",
256 catalog_version
257 );
258
259 Ok(MigrateResult {
260 builtin_table_updates: ast_builtin_table_updates,
261 catalog_updates: ast_catalog_updates,
262 post_item_updates,
263 })
264}
265
266fn rewrite_sources_to_tables(
345 tx: &mut Transaction<'_>,
346 catalog: &ConnCatalog<'_>,
347) -> Result<(), anyhow::Error> {
348 use mz_sql::ast::{
349 CreateSourceConnection, CreateSourceStatement, CreateSubsourceOptionName,
350 CreateSubsourceStatement, CreateTableFromSourceStatement, Ident,
351 KafkaSourceConfigOptionName, LoadGenerator, MySqlConfigOptionName, PgConfigOptionName,
352 RawItemName, TableFromSourceColumns, TableFromSourceOption, TableFromSourceOptionName,
353 UnresolvedItemName, Value, WithOptionValue,
354 };
355
356 let mut updated_items = BTreeMap::new();
357
358 let mut sources = vec![];
359 let mut subsources = vec![];
360
361 for item in tx.get_items() {
362 let stmt = mz_sql::parse::parse(&item.create_sql)?.into_element().ast;
363 match stmt {
364 Statement::CreateSubsource(stmt) => subsources.push((item, stmt)),
365 Statement::CreateSource(stmt) => sources.push((item, stmt)),
366 _ => {}
367 }
368 }
369
370 let mut pending_progress_items = BTreeMap::new();
371 let mut migrated_source_ids = BTreeMap::new();
372 for (mut source_item, source_stmt) in sources {
375 let CreateSourceStatement {
376 name,
377 in_cluster,
378 col_names,
379 mut connection,
380 include_metadata,
381 format,
382 envelope,
383 if_not_exists,
384 key_constraint,
385 with_options,
386 external_references,
387 progress_subsource,
388 } = source_stmt;
389
390 let (progress_name, progress_item) = match progress_subsource {
391 Some(DeferredItemName::Named(RawItemName::Name(name))) => {
392 let partial_name = normalize::unresolved_item_name(name.clone())?;
393 (name, catalog.resolve_item(&partial_name)?)
394 }
395 Some(DeferredItemName::Named(RawItemName::Id(id, name, _))) => {
396 let gid = id.parse()?;
397 (name, catalog.get_item(&gid))
398 }
399 Some(DeferredItemName::Deferred(_)) => {
400 unreachable!("invalid progress subsource")
401 }
402 None => {
403 info!("migrate: skipping already migrated source: {name}");
404 continue;
405 }
406 };
407 let raw_progress_name =
408 RawItemName::Id(progress_item.id().to_string(), progress_name.clone(), None);
409
410 let catalog_item = catalog.get_item(&source_item.id);
412 let source_name: &QualifiedItemName = catalog_item.name();
413 let full_source_name: FullItemName = catalog.resolve_full_name(source_name);
414 let source_name: UnresolvedItemName = normalize::unresolve(full_source_name.clone());
415
416 match &mut connection {
418 CreateSourceConnection::Postgres { options, .. } => {
419 options.retain(|o| match o.name {
420 PgConfigOptionName::Details | PgConfigOptionName::Publication => true,
421 PgConfigOptionName::TextColumns | PgConfigOptionName::ExcludeColumns => false,
422 });
423 }
424 CreateSourceConnection::SqlServer { options, .. } => {
425 options.retain(|o| match o.name {
426 SqlServerConfigOptionName::Details => true,
427 SqlServerConfigOptionName::TextColumns
428 | SqlServerConfigOptionName::ExcludeColumns => false,
429 });
430 }
431 CreateSourceConnection::MySql { options, .. } => {
432 options.retain(|o| match o.name {
433 MySqlConfigOptionName::Details => true,
434 MySqlConfigOptionName::TextColumns | MySqlConfigOptionName::ExcludeColumns => {
435 false
436 }
437 });
438 }
439 CreateSourceConnection::Kafka { .. } | CreateSourceConnection::LoadGenerator { .. } => {
440 }
441 }
442
443 let (new_progress_name, new_progress_stmt, new_source_name, new_source_stmt) =
445 match connection {
446 connection @ (CreateSourceConnection::Postgres { .. }
447 | CreateSourceConnection::MySql { .. }
448 | CreateSourceConnection::SqlServer { .. }
449 | CreateSourceConnection::LoadGenerator {
450 generator:
451 LoadGenerator::Tpch | LoadGenerator::Auction | LoadGenerator::Marketing,
452 ..
453 }) => {
454 assert_eq!(col_names, &[]);
456 assert_eq!(key_constraint, None);
457 assert_eq!(format, None);
458 assert_eq!(envelope, None);
459 assert_eq!(include_metadata, &[]);
460 assert_eq!(external_references, None);
461
462 let dummy_source_stmt = Statement::CreateView(CreateViewStatement {
470 if_exists: IfExistsBehavior::Error,
471 temporary: false,
472 definition: ViewDefinition {
473 name: progress_name,
474 columns: vec![],
475 query: Query {
476 ctes: CteBlock::Simple(vec![]),
477 body: SetExpr::Table(RawItemName::Id(
478 progress_item.id().to_string(),
479 source_name.clone(),
480 None,
481 )),
482 order_by: vec![],
483 limit: None,
484 offset: None,
485 },
486 },
487 });
488
489 let new_progress_stmt = CreateSourceStatement {
490 name: source_name.clone(),
491 in_cluster,
492 col_names: vec![],
493 connection,
494 include_metadata: vec![],
495 format: None,
496 envelope: None,
497 if_not_exists,
498 key_constraint: None,
499 with_options,
500 external_references: None,
501 progress_subsource: None,
502 };
503
504 migrated_source_ids.insert(source_item.id, progress_item.id());
505
506 (
507 full_source_name.item,
508 new_progress_stmt,
509 progress_item.name().item.clone(),
510 dummy_source_stmt,
511 )
512 }
513 CreateSourceConnection::Kafka {
514 options,
515 connection,
516 } => {
517 let constraints = if let Some(_key_constraint) = key_constraint {
518 vec![]
521 } else {
522 vec![]
523 };
524
525 let columns = if col_names.is_empty() {
526 TableFromSourceColumns::NotSpecified
527 } else {
528 TableFromSourceColumns::Named(col_names)
529 };
530
531 let details = SourceExportStatementDetails::Kafka {};
534 let table_with_options = vec![TableFromSourceOption {
535 name: TableFromSourceOptionName::Details,
536 value: Some(WithOptionValue::Value(Value::String(hex::encode(
537 details.into_proto().encode_to_vec(),
538 )))),
539 }];
540 let topic_option = options
542 .iter()
543 .find(|o| matches!(o.name, KafkaSourceConfigOptionName::Topic))
544 .expect("kafka sources must have a topic");
545 let topic = match &topic_option.value {
546 Some(WithOptionValue::Value(Value::String(topic))) => topic,
547 _ => unreachable!("topic must be a string"),
548 };
549 let external_reference = UnresolvedItemName::qualified(&[Ident::new(topic)?]);
550
551 let new_source_stmt =
552 Statement::CreateTableFromSource(CreateTableFromSourceStatement {
553 name: source_name,
554 constraints,
555 columns,
556 if_not_exists,
557 source: raw_progress_name,
558 include_metadata,
559 format,
560 envelope,
561 external_reference: Some(external_reference),
562 with_options: table_with_options,
563 });
564
565 let new_progress_stmt = CreateSourceStatement {
566 name: progress_name,
567 in_cluster,
568 col_names: vec![],
569 connection: CreateSourceConnection::Kafka {
570 options,
571 connection,
572 },
573 include_metadata: vec![],
574 format: None,
575 envelope: None,
576 if_not_exists,
577 key_constraint: None,
578 with_options,
579 external_references: None,
580 progress_subsource: None,
581 };
582 (
583 progress_item.name().item.clone(),
584 new_progress_stmt,
585 full_source_name.item,
586 new_source_stmt,
587 )
588 }
589 CreateSourceConnection::LoadGenerator {
590 generator:
591 generator @ (LoadGenerator::Clock
592 | LoadGenerator::Counter
593 | LoadGenerator::Datums
594 | LoadGenerator::KeyValue),
595 options,
596 } => {
597 let constraints = if let Some(_key_constraint) = key_constraint {
598 vec![]
600 } else {
601 vec![]
602 };
603
604 let columns = if col_names.is_empty() {
605 TableFromSourceColumns::NotSpecified
606 } else {
607 TableFromSourceColumns::Named(col_names)
608 };
609
610 let details = SourceExportStatementDetails::LoadGenerator {
613 output: LoadGeneratorOutput::Default,
614 };
615 let table_with_options = vec![TableFromSourceOption {
616 name: TableFromSourceOptionName::Details,
617 value: Some(WithOptionValue::Value(Value::String(hex::encode(
618 details.into_proto().encode_to_vec(),
619 )))),
620 }];
621 let external_reference = FullItemName {
624 database: mz_sql::names::RawDatabaseSpecifier::Name(
625 mz_storage_types::sources::load_generator::LOAD_GENERATOR_DATABASE_NAME
626 .to_owned(),
627 ),
628 schema: generator.schema_name().to_string(),
629 item: generator.schema_name().to_string(),
630 };
631
632 let new_source_stmt =
633 Statement::CreateTableFromSource(CreateTableFromSourceStatement {
634 name: source_name,
635 constraints,
636 columns,
637 if_not_exists,
638 source: raw_progress_name,
639 include_metadata,
640 format,
641 envelope,
642 external_reference: Some(external_reference.into()),
643 with_options: table_with_options,
644 });
645
646 let new_progress_stmt = CreateSourceStatement {
647 name: progress_name,
648 in_cluster,
649 col_names: vec![],
650 connection: CreateSourceConnection::LoadGenerator { generator, options },
651 include_metadata: vec![],
652 format: None,
653 envelope: None,
654 if_not_exists,
655 key_constraint: None,
656 with_options,
657 external_references: None,
658 progress_subsource: None,
659 };
660 (
661 progress_item.name().item.clone(),
662 new_progress_stmt,
663 full_source_name.item,
664 new_source_stmt,
665 )
666 }
667 };
668
669 info!(
673 "migrate: converted source {} to {}",
674 source_item.create_sql, new_source_stmt
675 );
676 source_item.name = new_source_name.clone();
677 source_item.create_sql = new_source_stmt.to_ast_string_stable();
678 updated_items.insert(source_item.id, source_item);
679 pending_progress_items.insert(progress_item.id(), (new_progress_name, new_progress_stmt));
680 }
681
682 for (mut item, stmt) in subsources {
683 match stmt {
684 CreateSubsourceStatement {
687 of_source: None, ..
688 } => {
689 let Some((new_name, new_stmt)) = pending_progress_items.remove(&item.id) else {
690 panic!("encountered orphan progress subsource id: {}", item.id)
691 };
692 item.name = new_name;
693 item.create_sql = new_stmt.to_ast_string_stable();
694 updated_items.insert(item.id, item);
695 }
696 CreateSubsourceStatement {
699 name,
700 columns,
701 constraints,
702 of_source: Some(raw_source_name),
703 if_not_exists,
704 mut with_options,
705 } => {
706 let new_raw_source_name = match raw_source_name {
707 RawItemName::Id(old_id, name, None) => {
708 let old_id: CatalogItemId = old_id.parse().expect("well formed");
709 let new_id = migrated_source_ids[&old_id].clone();
710 RawItemName::Id(new_id.to_string(), name, None)
711 }
712 _ => unreachable!("unexpected source name: {raw_source_name}"),
713 };
714 let external_reference = match with_options
717 .iter()
718 .position(|opt| opt.name == CreateSubsourceOptionName::ExternalReference)
719 {
720 Some(i) => match with_options.remove(i).value {
721 Some(WithOptionValue::UnresolvedItemName(name)) => name,
722 _ => unreachable!("external reference must be an unresolved item name"),
723 },
724 None => panic!("subsource must have an external reference"),
725 };
726
727 let with_options = with_options
728 .into_iter()
729 .map(|option| {
730 match option.name {
731 CreateSubsourceOptionName::Details => TableFromSourceOption {
732 name: TableFromSourceOptionName::Details,
733 value: option.value,
736 },
737 CreateSubsourceOptionName::TextColumns => TableFromSourceOption {
738 name: TableFromSourceOptionName::TextColumns,
739 value: option.value,
740 },
741 CreateSubsourceOptionName::ExcludeColumns => TableFromSourceOption {
742 name: TableFromSourceOptionName::ExcludeColumns,
743 value: option.value,
744 },
745 CreateSubsourceOptionName::RetainHistory => TableFromSourceOption {
746 name: TableFromSourceOptionName::RetainHistory,
747 value: option.value,
748 },
749 CreateSubsourceOptionName::Progress => {
750 panic!("progress option should not exist on this subsource")
751 }
752 CreateSubsourceOptionName::ExternalReference => {
753 unreachable!("This option is handled separately above.")
754 }
755 }
756 })
757 .collect::<Vec<_>>();
758
759 let table = CreateTableFromSourceStatement {
760 name,
761 constraints,
762 columns: TableFromSourceColumns::Defined(columns),
763 if_not_exists,
764 source: new_raw_source_name,
765 external_reference: Some(external_reference),
766 with_options,
767 envelope: None,
769 include_metadata: vec![],
770 format: None,
771 };
772
773 info!(
774 "migrate: converted subsource {} to table {}",
775 item.create_sql, table
776 );
777 item.create_sql = Statement::CreateTableFromSource(table).to_ast_string_stable();
778 updated_items.insert(item.id, item);
779 }
780 }
781 }
782 assert!(
783 pending_progress_items.is_empty(),
784 "unexpected residual progress items: {pending_progress_items:?}"
785 );
786
787 tx.update_items(updated_items)?;
788
789 Ok(())
790}
791
792pub(crate) fn durable_migrate(
796 tx: &mut Transaction,
797 _organization_id: Uuid,
798 _boot_ts: Timestamp,
799) -> Result<(), anyhow::Error> {
800 const EXPR_CACHE_MIGRATION_KEY: &str = "expr_cache_migration";
803 const EXPR_CACHE_MIGRATION_DONE: u64 = 1;
804 if tx.get_config(EXPR_CACHE_MIGRATION_KEY.to_string()) != Some(EXPR_CACHE_MIGRATION_DONE) {
805 if let Some(shard_id) = tx.get_expression_cache_shard() {
806 tx.insert_unfinalized_shards(btreeset! {shard_id})?;
807 tx.set_expression_cache_shard(ShardId::new())?;
808 }
809 tx.set_config(
810 EXPR_CACHE_MIGRATION_KEY.to_string(),
811 Some(EXPR_CACHE_MIGRATION_DONE),
812 )?;
813 }
814
815 const BUILTIN_MIGRATION_SHARD_MIGRATION_KEY: &str = "migration_shard_migration";
818 const BUILTIN_MIGRATION_SHARD_MIGRATION_DONE: u64 = 1;
819 if tx.get_config(BUILTIN_MIGRATION_SHARD_MIGRATION_KEY.to_string())
820 != Some(BUILTIN_MIGRATION_SHARD_MIGRATION_DONE)
821 {
822 if let Some(shard_id) = tx.get_builtin_migration_shard() {
823 tx.insert_unfinalized_shards(btreeset! {shard_id})?;
824 tx.set_builtin_migration_shard(ShardId::new())?;
825 }
826 tx.set_config(
827 BUILTIN_MIGRATION_SHARD_MIGRATION_KEY.to_string(),
828 Some(BUILTIN_MIGRATION_SHARD_MIGRATION_DONE),
829 )?;
830 }
831
832 if tx
833 .get_setting(MOCK_AUTHENTICATION_NONCE_KEY.to_string())
834 .is_none()
835 {
836 let mut nonce = [0u8; 24];
837 openssl::rand::rand_bytes(&mut nonce).expect("failed to generate nonce");
838 let nonce = BASE64_STANDARD.encode(nonce);
839 tx.set_setting(MOCK_AUTHENTICATION_NONCE_KEY.to_string(), Some(nonce))?;
840 }
841
842 migrate_builtin_tables_to_mvs(tx)?;
843
844 Ok(())
845}
846
847fn migrate_builtin_tables_to_mvs(tx: &mut Transaction) -> Result<(), anyhow::Error> {
851 let expected_mvs: BTreeSet<_> = BUILTINS::materialized_views()
853 .map(|mv| (mv.schema, mv.name))
854 .collect();
855
856 let mut to_remove = BTreeSet::new();
858 let mut to_add = Vec::new();
859 for mapping in tx.get_system_object_mappings() {
860 let desc = &mapping.description;
861 if desc.object_type != CatalogItemType::Table {
862 continue;
863 }
864
865 let key = (&*desc.schema_name, &*desc.object_name);
866 if expected_mvs.contains(&key) {
867 info!(
868 "migrate: builtin {}.{} changed type from table to MV",
869 desc.schema_name, desc.object_name,
870 );
871 to_remove.insert(desc.clone());
872 to_add.push(SystemObjectMapping {
873 description: SystemObjectDescription {
874 schema_name: desc.schema_name.clone(),
875 object_type: CatalogItemType::MaterializedView,
876 object_name: desc.object_name.clone(),
877 },
878 unique_identifier: mapping.unique_identifier,
879 });
880 }
881 }
882
883 if !to_remove.is_empty() {
884 tx.remove_system_object_mappings(to_remove)?;
885 tx.set_system_object_mappings(to_add)?;
886 }
887
888 Ok(())
889}
890
891fn ast_rewrite_create_sink_partition_strategy(
903 stmt: &mut Statement<Raw>,
904) -> Result<(), anyhow::Error> {
905 let Statement::CreateSink(stmt) = stmt else {
906 return Ok(());
907 };
908 stmt.with_options
909 .retain(|op| op.name != CreateSinkOptionName::PartitionStrategy);
910 Ok(())
911}
912
913fn ast_rewrite_sql_server_constraints(stmt: &mut Statement<Raw>) -> Result<(), anyhow::Error> {
915 use mz_sql::ast::{
916 CreateSubsourceOptionName, TableFromSourceOptionName, Value, WithOptionValue,
917 };
918 use mz_sql_server_util::desc::{SqlServerTableConstraint, SqlServerTableConstraintType};
919 use mz_storage_types::sources::ProtoSourceExportStatementDetails;
920 use mz_storage_types::sources::proto_source_export_statement_details::Kind;
921
922 let deets: Option<&mut String> = match stmt {
923 Statement::CreateSubsource(stmt) => stmt.with_options.iter_mut().find_map(|option| {
924 if matches!(option.name, CreateSubsourceOptionName::Details)
925 && let Some(WithOptionValue::Value(Value::String(ref mut details))) = option.value
926 {
927 Some(details)
928 } else {
929 None
930 }
931 }),
932 Statement::CreateTableFromSource(stmt) => stmt.with_options.iter_mut().find_map(|option| {
933 if matches!(option.name, TableFromSourceOptionName::Details)
934 && let Some(WithOptionValue::Value(Value::String(ref mut details))) = option.value
935 {
936 Some(details)
937 } else {
938 None
939 }
940 }),
941 _ => None,
942 };
943 let Some(deets) = deets else {
944 return Ok(());
945 };
946
947 let current_value = hex::decode(&mut *deets)?;
948 let current_value = ProtoSourceExportStatementDetails::decode(&*current_value)?;
949
950 if !matches!(current_value.kind, Some(Kind::SqlServer(_))) {
952 return Ok(());
953 };
954
955 let SourceExportStatementDetails::SqlServer {
956 mut table,
957 capture_instance,
958 initial_lsn,
959 } = SourceExportStatementDetails::from_proto(current_value)?
960 else {
961 unreachable!("statement details must exist for SQL Server");
962 };
963
964 if !table.constraints.is_empty() {
966 return Ok(());
967 }
968
969 let mut migrated_constraints: BTreeMap<_, Vec<_>> = BTreeMap::new();
972 for col in table.columns.iter_mut() {
973 if let Some(constraint_name) = col.primary_key_constraint.take() {
974 migrated_constraints
975 .entry(constraint_name)
976 .or_default()
977 .push(col.name.to_string());
978 }
979 }
980
981 table.constraints = migrated_constraints
982 .into_iter()
983 .map(|(constraint_name, column_names)| SqlServerTableConstraint {
984 constraint_name: constraint_name.to_string(),
985 constraint_type: SqlServerTableConstraintType::PrimaryKey,
986 column_names,
987 })
988 .collect();
989
990 let new_value = SourceExportStatementDetails::SqlServer {
991 table,
992 capture_instance,
993 initial_lsn,
994 };
995 *deets = hex::encode(new_value.into_proto().encode_to_vec());
996
997 Ok(())
998}
999
1000fn ast_rewrite_add_missing_index_ids(
1002 tx: &Transaction<'_>,
1003 stmt: &mut Statement<Raw>,
1004) -> Result<(), anyhow::Error> {
1005 let Statement::CreateIndex(stmt) = stmt else {
1006 return Ok(());
1007 };
1008
1009 let unresolved_name = match stmt.on_name.clone() {
1010 mz_sql::ast::RawItemName::Name(name) => name,
1011 mz_sql::ast::RawItemName::Id(..) => return Ok(()),
1013 };
1014
1015 let parts = &unresolved_name.0;
1016 let (db_name, schema_name, item_name) = match parts.len() {
1017 3 => (Some(&parts[0]), &parts[1], &parts[2]),
1018 2 => (None, &parts[0], &parts[1]),
1019 _ => panic!("invalid unresolved name: {unresolved_name:?}"),
1020 };
1021
1022 let db_id = db_name.map(|x| {
1023 let db = tx.get_databases().find(|db| db.name == x.as_str());
1024 let db = db.unwrap_or_else(|| panic!("missing database: {x}"));
1025 db.id
1026 });
1027 let schema_id = {
1028 let schema = tx
1029 .get_schemas()
1030 .find(|s| s.name == schema_name.as_str() && s.database_id == db_id);
1031 let schema = schema.unwrap_or_else(|| panic!("missing schema: {schema_name}, {db_id:?}"));
1032 schema.id
1033 };
1034 let item_id = {
1035 let item = tx
1036 .get_items()
1037 .find(|i| i.name == item_name.as_str() && i.schema_id == schema_id);
1038 let item = item.unwrap_or_else(|| panic!("missing item: {item_name}, {schema_id:?}"));
1039 item.id
1040 };
1041
1042 stmt.on_name = mz_sql::ast::RawItemName::Id(item_id.to_string(), unresolved_name, None);
1043
1044 Ok(())
1045}
1046
1047fn ast_rewrite_strip_builtin_version_pins(stmt: &mut Statement<Raw>) -> Result<(), anyhow::Error> {
1061 use mz_sql::ast::RawItemName;
1062 use mz_sql::ast::visit_mut::{VisitMut, VisitMutNode};
1063
1064 struct StripBuiltinVersionPins;
1065
1066 impl<'ast> VisitMut<'ast, Raw> for StripBuiltinVersionPins {
1067 fn visit_item_name_mut(&mut self, item_name: &mut RawItemName) {
1068 if let RawItemName::Id(id, _, version) = item_name {
1069 if version.is_some() {
1070 if let Ok(parsed) = id.parse::<CatalogItemId>() {
1071 if !parsed.is_user() {
1072 *version = None;
1073 }
1074 }
1075 }
1076 }
1077 }
1078 }
1079
1080 let mut visitor = StripBuiltinVersionPins;
1081 stmt.visit_mut(&mut visitor);
1082 Ok(())
1083}
1084
1085fn ast_rewrite_kafka_metadata_refresh_intervals(
1086 stmt: &mut Statement<Raw>,
1087) -> Result<(), anyhow::Error> {
1088 use mz_sql::ast::{
1089 CreateSinkConnection, CreateSourceConnection, KafkaSinkConfigOptionName,
1090 KafkaSourceConfigOptionName, WithOptionValue,
1091 };
1092 let interval: Option<&mut WithOptionValue<Raw>> = match stmt {
1097 Statement::CreateSource(stmt) => {
1098 if let CreateSourceConnection::Kafka { options, .. } = &mut stmt.connection {
1099 options.iter_mut().find_map(|option| {
1100 if matches!(
1101 option.name,
1102 KafkaSourceConfigOptionName::TopicMetadataRefreshInterval
1103 ) {
1104 option.value.as_mut()
1105 } else {
1106 None
1107 }
1108 })
1109 } else {
1110 None
1111 }
1112 }
1113 Statement::CreateSink(stmt) => {
1114 if let CreateSinkConnection::Kafka { options, .. } = &mut stmt.connection {
1115 options.iter_mut().find_map(|option| {
1116 if matches!(
1117 option.name,
1118 KafkaSinkConfigOptionName::TopicMetadataRefreshInterval
1119 ) {
1120 option.value.as_mut()
1121 } else {
1122 None
1123 }
1124 })
1125 } else {
1126 None
1127 }
1128 }
1129 _ => None,
1130 };
1131
1132 let Some(interval) = interval else {
1133 return Ok(());
1134 };
1135
1136 rewrite_interval_option_floor_1s(interval, "kafka metadata refresh interval")
1137}
1138
1139fn ast_rewrite_small_commit_intervals(stmt: &mut Statement<Raw>) -> Result<(), anyhow::Error> {
1147 use mz_sql::ast::{CreateSinkConnection, CreateSinkOptionName};
1148
1149 let Statement::CreateSink(stmt) = stmt else {
1150 return Ok(());
1151 };
1152 if !matches!(stmt.connection, CreateSinkConnection::Iceberg { .. }) {
1153 return Ok(());
1154 }
1155 let interval = stmt.with_options.iter_mut().find_map(|o| {
1156 if matches!(o.name, CreateSinkOptionName::CommitInterval) {
1157 o.value.as_mut()
1158 } else {
1159 None
1160 }
1161 });
1162 let Some(interval) = interval else {
1163 return Ok(());
1164 };
1165
1166 rewrite_interval_option_floor_1s(interval, "commit interval")
1167}
1168
1169fn rewrite_interval_option_floor_1s(
1171 value: &mut mz_sql::ast::WithOptionValue<Raw>,
1172 label: &str,
1173) -> Result<(), anyhow::Error> {
1174 use mz_sql::ast::{Value, WithOptionValue};
1175 use mz_sql::plan::TryFromValue;
1176
1177 let dur = Duration::try_from_value(value.clone())
1178 .map_err(|e| anyhow::anyhow!("invalid value for {label}: {value:?}: {e}"))?;
1179
1180 if dur < Duration::from_secs(1) {
1181 *value = WithOptionValue::Value(Value::String("1s".to_string()));
1182 }
1183
1184 Ok(())
1185}
1186
1187#[cfg(test)]
1188mod tests {
1189 use super::*;
1190
1191 fn strip(sql: &str) -> String {
1192 let mut stmt = mz_sql::parse::parse(sql)
1193 .expect("test sql parses")
1194 .into_element()
1195 .ast;
1196 ast_rewrite_strip_builtin_version_pins(&mut stmt).expect("rewrite succeeds");
1197 stmt.to_ast_string_stable()
1198 }
1199
1200 #[mz_ore::test]
1201 #[cfg_attr(miri, ignore)] fn strips_version_from_builtin_reference() {
1203 let out = strip(
1207 r#"CREATE VIEW "materialize"."public"."v" AS SELECT 1 FROM [s518 AS "mz_catalog"."mz_audit_events" VERSION 0]"#,
1208 );
1209 assert!(!out.contains("VERSION"), "version not stripped: {out}");
1210 assert!(out.contains("s518"), "reference dropped: {out}");
1212 }
1213
1214 #[mz_ore::test]
1215 #[cfg_attr(miri, ignore)] fn preserves_version_on_user_reference() {
1217 let out = strip(
1220 r#"CREATE VIEW "materialize"."public"."v" AS SELECT 1 FROM [u5 AS "materialize"."public"."t" VERSION 1]"#,
1221 );
1222 assert!(out.contains("VERSION"), "user version stripped: {out}");
1223 assert!(out.contains("u5"), "reference dropped: {out}");
1224 }
1225
1226 #[mz_ore::test]
1227 #[cfg_attr(miri, ignore)] fn leaves_unpinned_builtin_reference_untouched() {
1229 let out = strip(
1230 r#"CREATE VIEW "materialize"."public"."v" AS SELECT 1 FROM [s518 AS "mz_catalog"."mz_audit_events"]"#,
1231 );
1232 assert!(!out.contains("VERSION"), "unexpected version: {out}");
1233 assert!(out.contains("s518"), "reference dropped: {out}");
1234 }
1235}