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_add_missing_doc_on_ids(tx, stmt)?;
158 ast_rewrite_kafka_metadata_refresh_intervals(stmt)?;
159 ast_rewrite_small_commit_intervals(stmt)?;
160 ast_rewrite_strip_builtin_version_pins(stmt)?;
161 Ok(())
162 })?;
163
164 let commit_ts = tx.upper();
167 let mut item_updates = into_consolidatable_updates_startup(item_updates, commit_ts);
168 let op_item_updates = tx.get_and_commit_op_updates();
169 let op_item_updates = into_consolidatable_updates_startup(op_item_updates, commit_ts);
170 item_updates.extend(op_item_updates);
171 differential_dataflow::consolidation::consolidate_updates(&mut item_updates);
172
173 let (post_item_updates, item_updates): (Vec<_>, Vec<_>) = item_updates
177 .into_iter()
178 .partition(|(kind, _, _)| matches!(kind, StateUpdateKind::StorageCollectionMetadata(_)));
181
182 let item_updates = item_updates
183 .into_iter()
184 .map(|(kind, ts, diff)| StateUpdate {
185 kind,
186 ts,
187 diff: diff.try_into().expect("valid diff"),
188 })
189 .collect();
190
191 let force_source_table_syntax = state.system_config().force_source_table_syntax();
192 if force_source_table_syntax {
196 state
197 .system_config_mut()
198 .set(FORCE_SOURCE_TABLE_SYNTAX.name(), VarInput::Flat("off"))
199 .expect("known parameter");
200 }
201
202 let (mut ast_builtin_table_updates, mut ast_catalog_updates) =
203 state.apply_updates(item_updates, local_expr_cache).await;
204
205 info!("migrating from catalog version {:?}", catalog_version);
206
207 let conn_cat = state.for_system_session();
208
209 if force_source_table_syntax {
212 rewrite_sources_to_tables(tx, &conn_cat)?;
213 }
214
215 rewrite_items(tx, &conn_cat, |_tx, _conn_cat, _id, _stmt| {
216 let _catalog_version = catalog_version.clone();
217 Ok(())
232 })?;
233
234 if force_source_table_syntax {
235 state
236 .system_config_mut()
237 .set(FORCE_SOURCE_TABLE_SYNTAX.name(), VarInput::Flat("on"))
238 .expect("known parameter");
239 }
240
241 let op_item_updates = tx.get_and_commit_op_updates();
247 let (item_builtin_table_updates, item_catalog_updates) =
248 state.apply_updates(op_item_updates, local_expr_cache).await;
249
250 ast_builtin_table_updates.extend(item_builtin_table_updates);
251 ast_catalog_updates.extend(item_catalog_updates);
252
253 info!(
254 "migration from catalog version {:?} complete",
255 catalog_version
256 );
257
258 Ok(MigrateResult {
259 builtin_table_updates: ast_builtin_table_updates,
260 catalog_updates: ast_catalog_updates,
261 post_item_updates,
262 })
263}
264
265fn rewrite_sources_to_tables(
344 tx: &mut Transaction<'_>,
345 catalog: &ConnCatalog<'_>,
346) -> Result<(), anyhow::Error> {
347 use mz_sql::ast::{
348 CreateSourceConnection, CreateSourceStatement, CreateSubsourceOptionName,
349 CreateSubsourceStatement, CreateTableFromSourceStatement, Ident,
350 KafkaSourceConfigOptionName, LoadGenerator, MySqlConfigOptionName, PgConfigOptionName,
351 RawItemName, TableFromSourceColumns, TableFromSourceOption, TableFromSourceOptionName,
352 UnresolvedItemName, Value, WithOptionValue,
353 };
354
355 let mut updated_items = BTreeMap::new();
356
357 let mut sources = vec![];
358 let mut subsources = vec![];
359
360 for item in tx.get_items() {
361 let stmt = mz_sql::parse::parse(&item.create_sql)?.into_element().ast;
362 match stmt {
363 Statement::CreateSubsource(stmt) => subsources.push((item, stmt)),
364 Statement::CreateSource(stmt) => sources.push((item, stmt)),
365 _ => {}
366 }
367 }
368
369 let mut pending_progress_items = BTreeMap::new();
370 let mut migrated_source_ids = BTreeMap::new();
371 for (mut source_item, source_stmt) in sources {
374 let CreateSourceStatement {
375 name,
376 in_cluster,
377 col_names,
378 mut connection,
379 include_metadata,
380 format,
381 envelope,
382 if_not_exists,
383 key_constraint,
384 with_options,
385 external_references,
386 progress_subsource,
387 } = source_stmt;
388
389 let (progress_name, progress_item) = match progress_subsource {
390 Some(DeferredItemName::Named(RawItemName::Name(name))) => {
391 let partial_name = normalize::unresolved_item_name(name.clone())?;
392 (name, catalog.resolve_item(&partial_name)?)
393 }
394 Some(DeferredItemName::Named(RawItemName::Id(id, name, _))) => {
395 let gid = id.parse()?;
396 (name, catalog.get_item(&gid))
397 }
398 Some(DeferredItemName::Deferred(_)) => {
399 unreachable!("invalid progress subsource")
400 }
401 None => {
402 info!("migrate: skipping already migrated source: {name}");
403 continue;
404 }
405 };
406 let raw_progress_name =
407 RawItemName::Id(progress_item.id().to_string(), progress_name.clone(), None);
408
409 let catalog_item = catalog.get_item(&source_item.id);
411 let source_name: &QualifiedItemName = catalog_item.name();
412 let full_source_name: FullItemName = catalog.resolve_full_name(source_name);
413 let source_name: UnresolvedItemName = normalize::unresolve(full_source_name.clone());
414
415 match &mut connection {
417 CreateSourceConnection::Postgres { options, .. } => {
418 options.retain(|o| match o.name {
419 PgConfigOptionName::Details | PgConfigOptionName::Publication => true,
420 PgConfigOptionName::TextColumns | PgConfigOptionName::ExcludeColumns => false,
421 });
422 }
423 CreateSourceConnection::SqlServer { options, .. } => {
424 options.retain(|o| match o.name {
425 SqlServerConfigOptionName::Details => true,
426 SqlServerConfigOptionName::TextColumns
427 | SqlServerConfigOptionName::ExcludeColumns => false,
428 });
429 }
430 CreateSourceConnection::MySql { options, .. } => {
431 options.retain(|o| match o.name {
432 MySqlConfigOptionName::Details => true,
433 MySqlConfigOptionName::TextColumns | MySqlConfigOptionName::ExcludeColumns => {
434 false
435 }
436 });
437 }
438 CreateSourceConnection::Kafka { .. } | CreateSourceConnection::LoadGenerator { .. } => {
439 }
440 }
441
442 let (new_progress_name, new_progress_stmt, new_source_name, new_source_stmt) =
444 match connection {
445 connection @ (CreateSourceConnection::Postgres { .. }
446 | CreateSourceConnection::MySql { .. }
447 | CreateSourceConnection::SqlServer { .. }
448 | CreateSourceConnection::LoadGenerator {
449 generator:
450 LoadGenerator::Tpch | LoadGenerator::Auction | LoadGenerator::Marketing,
451 ..
452 }) => {
453 assert_eq!(col_names, &[]);
455 assert_eq!(key_constraint, None);
456 assert_eq!(format, None);
457 assert_eq!(envelope, None);
458 assert_eq!(include_metadata, &[]);
459 assert_eq!(external_references, None);
460
461 let dummy_source_stmt = Statement::CreateView(CreateViewStatement {
469 if_exists: IfExistsBehavior::Error,
470 temporary: false,
471 definition: ViewDefinition {
472 name: progress_name,
473 columns: vec![],
474 query: Query {
475 ctes: CteBlock::Simple(vec![]),
476 body: SetExpr::Table(RawItemName::Id(
477 progress_item.id().to_string(),
478 source_name.clone(),
479 None,
480 )),
481 order_by: vec![],
482 limit: None,
483 offset: None,
484 },
485 },
486 });
487
488 let new_progress_stmt = CreateSourceStatement {
489 name: source_name.clone(),
490 in_cluster,
491 col_names: vec![],
492 connection,
493 include_metadata: vec![],
494 format: None,
495 envelope: None,
496 if_not_exists,
497 key_constraint: None,
498 with_options,
499 external_references: None,
500 progress_subsource: None,
501 };
502
503 migrated_source_ids.insert(source_item.id, progress_item.id());
504
505 (
506 full_source_name.item,
507 new_progress_stmt,
508 progress_item.name().item.clone(),
509 dummy_source_stmt,
510 )
511 }
512 CreateSourceConnection::Kafka {
513 options,
514 connection,
515 } => {
516 let constraints = if let Some(_key_constraint) = key_constraint {
517 vec![]
520 } else {
521 vec![]
522 };
523
524 let columns = if col_names.is_empty() {
525 TableFromSourceColumns::NotSpecified
526 } else {
527 TableFromSourceColumns::Named(col_names)
528 };
529
530 let details = SourceExportStatementDetails::Kafka {};
533 let table_with_options = vec![TableFromSourceOption {
534 name: TableFromSourceOptionName::Details,
535 value: Some(WithOptionValue::Value(Value::String(hex::encode(
536 details.into_proto().encode_to_vec(),
537 )))),
538 }];
539 let topic_option = options
541 .iter()
542 .find(|o| matches!(o.name, KafkaSourceConfigOptionName::Topic))
543 .expect("kafka sources must have a topic");
544 let topic = match &topic_option.value {
545 Some(WithOptionValue::Value(Value::String(topic))) => topic,
546 _ => unreachable!("topic must be a string"),
547 };
548 let external_reference = UnresolvedItemName::qualified(&[Ident::new(topic)?]);
549
550 let new_source_stmt =
551 Statement::CreateTableFromSource(CreateTableFromSourceStatement {
552 name: source_name,
553 constraints,
554 columns,
555 if_not_exists,
556 source: raw_progress_name,
557 include_metadata,
558 format,
559 envelope,
560 external_reference: Some(external_reference),
561 with_options: table_with_options,
562 });
563
564 let new_progress_stmt = CreateSourceStatement {
565 name: progress_name,
566 in_cluster,
567 col_names: vec![],
568 connection: CreateSourceConnection::Kafka {
569 options,
570 connection,
571 },
572 include_metadata: vec![],
573 format: None,
574 envelope: None,
575 if_not_exists,
576 key_constraint: None,
577 with_options,
578 external_references: None,
579 progress_subsource: None,
580 };
581 (
582 progress_item.name().item.clone(),
583 new_progress_stmt,
584 full_source_name.item,
585 new_source_stmt,
586 )
587 }
588 CreateSourceConnection::LoadGenerator {
589 generator:
590 generator @ (LoadGenerator::Clock
591 | LoadGenerator::Counter
592 | LoadGenerator::Datums
593 | LoadGenerator::KeyValue),
594 options,
595 } => {
596 let constraints = if let Some(_key_constraint) = key_constraint {
597 vec![]
599 } else {
600 vec![]
601 };
602
603 let columns = if col_names.is_empty() {
604 TableFromSourceColumns::NotSpecified
605 } else {
606 TableFromSourceColumns::Named(col_names)
607 };
608
609 let details = SourceExportStatementDetails::LoadGenerator {
612 output: LoadGeneratorOutput::Default,
613 };
614 let table_with_options = vec![TableFromSourceOption {
615 name: TableFromSourceOptionName::Details,
616 value: Some(WithOptionValue::Value(Value::String(hex::encode(
617 details.into_proto().encode_to_vec(),
618 )))),
619 }];
620 let external_reference = FullItemName {
623 database: mz_sql::names::RawDatabaseSpecifier::Name(
624 mz_storage_types::sources::load_generator::LOAD_GENERATOR_DATABASE_NAME
625 .to_owned(),
626 ),
627 schema: generator.schema_name().to_string(),
628 item: generator.schema_name().to_string(),
629 };
630
631 let new_source_stmt =
632 Statement::CreateTableFromSource(CreateTableFromSourceStatement {
633 name: source_name,
634 constraints,
635 columns,
636 if_not_exists,
637 source: raw_progress_name,
638 include_metadata,
639 format,
640 envelope,
641 external_reference: Some(external_reference.into()),
642 with_options: table_with_options,
643 });
644
645 let new_progress_stmt = CreateSourceStatement {
646 name: progress_name,
647 in_cluster,
648 col_names: vec![],
649 connection: CreateSourceConnection::LoadGenerator { generator, options },
650 include_metadata: vec![],
651 format: None,
652 envelope: None,
653 if_not_exists,
654 key_constraint: None,
655 with_options,
656 external_references: None,
657 progress_subsource: None,
658 };
659 (
660 progress_item.name().item.clone(),
661 new_progress_stmt,
662 full_source_name.item,
663 new_source_stmt,
664 )
665 }
666 };
667
668 info!(
672 "migrate: converted source {} to {}",
673 source_item.create_sql, new_source_stmt
674 );
675 source_item.name = new_source_name.clone();
676 source_item.create_sql = new_source_stmt.to_ast_string_stable();
677 updated_items.insert(source_item.id, source_item);
678 pending_progress_items.insert(progress_item.id(), (new_progress_name, new_progress_stmt));
679 }
680
681 for (mut item, stmt) in subsources {
682 match stmt {
683 CreateSubsourceStatement {
686 of_source: None, ..
687 } => {
688 let Some((new_name, new_stmt)) = pending_progress_items.remove(&item.id) else {
689 panic!("encountered orphan progress subsource id: {}", item.id)
690 };
691 item.name = new_name;
692 item.create_sql = new_stmt.to_ast_string_stable();
693 updated_items.insert(item.id, item);
694 }
695 CreateSubsourceStatement {
698 name,
699 columns,
700 constraints,
701 of_source: Some(raw_source_name),
702 if_not_exists,
703 mut with_options,
704 } => {
705 let new_raw_source_name = match raw_source_name {
706 RawItemName::Id(old_id, name, None) => {
707 let old_id: CatalogItemId = old_id.parse().expect("well formed");
708 let new_id = migrated_source_ids[&old_id].clone();
709 RawItemName::Id(new_id.to_string(), name, None)
710 }
711 _ => unreachable!("unexpected source name: {raw_source_name}"),
712 };
713 let external_reference = match with_options
716 .iter()
717 .position(|opt| opt.name == CreateSubsourceOptionName::ExternalReference)
718 {
719 Some(i) => match with_options.remove(i).value {
720 Some(WithOptionValue::UnresolvedItemName(name)) => name,
721 _ => unreachable!("external reference must be an unresolved item name"),
722 },
723 None => panic!("subsource must have an external reference"),
724 };
725
726 let with_options = with_options
727 .into_iter()
728 .map(|option| {
729 match option.name {
730 CreateSubsourceOptionName::Details => TableFromSourceOption {
731 name: TableFromSourceOptionName::Details,
732 value: option.value,
735 },
736 CreateSubsourceOptionName::TextColumns => TableFromSourceOption {
737 name: TableFromSourceOptionName::TextColumns,
738 value: option.value,
739 },
740 CreateSubsourceOptionName::ExcludeColumns => TableFromSourceOption {
741 name: TableFromSourceOptionName::ExcludeColumns,
742 value: option.value,
743 },
744 CreateSubsourceOptionName::RetainHistory => TableFromSourceOption {
745 name: TableFromSourceOptionName::RetainHistory,
746 value: option.value,
747 },
748 CreateSubsourceOptionName::Progress => {
749 panic!("progress option should not exist on this subsource")
750 }
751 CreateSubsourceOptionName::ExternalReference => {
752 unreachable!("This option is handled separately above.")
753 }
754 }
755 })
756 .collect::<Vec<_>>();
757
758 let table = CreateTableFromSourceStatement {
759 name,
760 constraints,
761 columns: TableFromSourceColumns::Defined(columns),
762 if_not_exists,
763 source: new_raw_source_name,
764 external_reference: Some(external_reference),
765 with_options,
766 envelope: None,
768 include_metadata: vec![],
769 format: None,
770 };
771
772 info!(
773 "migrate: converted subsource {} to table {}",
774 item.create_sql, table
775 );
776 item.create_sql = Statement::CreateTableFromSource(table).to_ast_string_stable();
777 updated_items.insert(item.id, item);
778 }
779 }
780 }
781 assert!(
782 pending_progress_items.is_empty(),
783 "unexpected residual progress items: {pending_progress_items:?}"
784 );
785
786 tx.update_items(updated_items)?;
787
788 Ok(())
789}
790
791pub(crate) fn durable_migrate(
795 tx: &mut Transaction,
796 _organization_id: Uuid,
797 _boot_ts: Timestamp,
798) -> Result<(), anyhow::Error> {
799 const EXPR_CACHE_MIGRATION_KEY: &str = "expr_cache_migration";
802 const EXPR_CACHE_MIGRATION_DONE: u64 = 1;
803 if tx.get_config(EXPR_CACHE_MIGRATION_KEY.to_string()) != Some(EXPR_CACHE_MIGRATION_DONE) {
804 if let Some(shard_id) = tx.get_expression_cache_shard() {
805 tx.insert_unfinalized_shards(btreeset! {shard_id})?;
806 tx.set_expression_cache_shard(ShardId::new())?;
807 }
808 tx.set_config(
809 EXPR_CACHE_MIGRATION_KEY.to_string(),
810 Some(EXPR_CACHE_MIGRATION_DONE),
811 )?;
812 }
813
814 const BUILTIN_MIGRATION_SHARD_MIGRATION_KEY: &str = "migration_shard_migration";
817 const BUILTIN_MIGRATION_SHARD_MIGRATION_DONE: u64 = 1;
818 if tx.get_config(BUILTIN_MIGRATION_SHARD_MIGRATION_KEY.to_string())
819 != Some(BUILTIN_MIGRATION_SHARD_MIGRATION_DONE)
820 {
821 if let Some(shard_id) = tx.get_builtin_migration_shard() {
822 tx.insert_unfinalized_shards(btreeset! {shard_id})?;
823 tx.set_builtin_migration_shard(ShardId::new())?;
824 }
825 tx.set_config(
826 BUILTIN_MIGRATION_SHARD_MIGRATION_KEY.to_string(),
827 Some(BUILTIN_MIGRATION_SHARD_MIGRATION_DONE),
828 )?;
829 }
830
831 if tx
832 .get_setting(MOCK_AUTHENTICATION_NONCE_KEY.to_string())
833 .is_none()
834 {
835 let mut nonce = [0u8; 24];
836 openssl::rand::rand_bytes(&mut nonce).expect("failed to generate nonce");
837 let nonce = BASE64_STANDARD.encode(nonce);
838 tx.set_setting(MOCK_AUTHENTICATION_NONCE_KEY.to_string(), Some(nonce))?;
839 }
840
841 migrate_builtin_tables_to_mvs(tx)?;
842
843 Ok(())
844}
845
846fn migrate_builtin_tables_to_mvs(tx: &mut Transaction) -> Result<(), anyhow::Error> {
850 let expected_mvs: BTreeSet<_> = BUILTINS::materialized_views()
852 .map(|mv| (mv.schema, mv.name))
853 .collect();
854
855 let mut to_remove = BTreeSet::new();
857 let mut to_add = Vec::new();
858 for mapping in tx.get_system_object_mappings() {
859 let desc = &mapping.description;
860 if desc.object_type != CatalogItemType::Table {
861 continue;
862 }
863
864 let key = (&*desc.schema_name, &*desc.object_name);
865 if expected_mvs.contains(&key) {
866 info!(
867 "migrate: builtin {}.{} changed type from table to MV",
868 desc.schema_name, desc.object_name,
869 );
870 to_remove.insert(desc.clone());
871 to_add.push(SystemObjectMapping {
872 description: SystemObjectDescription {
873 schema_name: desc.schema_name.clone(),
874 object_type: CatalogItemType::MaterializedView,
875 object_name: desc.object_name.clone(),
876 },
877 unique_identifier: mapping.unique_identifier,
878 });
879 }
880 }
881
882 if !to_remove.is_empty() {
883 tx.remove_system_object_mappings(to_remove)?;
884 tx.set_system_object_mappings(to_add)?;
885 }
886
887 Ok(())
888}
889
890fn ast_rewrite_create_sink_partition_strategy(
902 stmt: &mut Statement<Raw>,
903) -> Result<(), anyhow::Error> {
904 let Statement::CreateSink(stmt) = stmt else {
905 return Ok(());
906 };
907 stmt.with_options
908 .retain(|op| op.name != CreateSinkOptionName::PartitionStrategy);
909 Ok(())
910}
911
912fn ast_rewrite_sql_server_constraints(stmt: &mut Statement<Raw>) -> Result<(), anyhow::Error> {
914 use mz_sql::ast::{
915 CreateSubsourceOptionName, TableFromSourceOptionName, Value, WithOptionValue,
916 };
917 use mz_sql_server_util::desc::{SqlServerTableConstraint, SqlServerTableConstraintType};
918 use mz_storage_types::sources::ProtoSourceExportStatementDetails;
919 use mz_storage_types::sources::proto_source_export_statement_details::Kind;
920
921 let deets: Option<&mut String> = match stmt {
922 Statement::CreateSubsource(stmt) => stmt.with_options.iter_mut().find_map(|option| {
923 if matches!(option.name, CreateSubsourceOptionName::Details)
924 && let Some(WithOptionValue::Value(Value::String(ref mut details))) = option.value
925 {
926 Some(details)
927 } else {
928 None
929 }
930 }),
931 Statement::CreateTableFromSource(stmt) => stmt.with_options.iter_mut().find_map(|option| {
932 if matches!(option.name, TableFromSourceOptionName::Details)
933 && let Some(WithOptionValue::Value(Value::String(ref mut details))) = option.value
934 {
935 Some(details)
936 } else {
937 None
938 }
939 }),
940 _ => None,
941 };
942 let Some(deets) = deets else {
943 return Ok(());
944 };
945
946 let current_value = hex::decode(&mut *deets)?;
947 let current_value = ProtoSourceExportStatementDetails::decode(&*current_value)?;
948
949 if !matches!(current_value.kind, Some(Kind::SqlServer(_))) {
951 return Ok(());
952 };
953
954 let SourceExportStatementDetails::SqlServer {
955 mut table,
956 capture_instance,
957 initial_lsn,
958 } = SourceExportStatementDetails::from_proto(current_value)?
959 else {
960 unreachable!("statement details must exist for SQL Server");
961 };
962
963 if !table.constraints.is_empty() {
965 return Ok(());
966 }
967
968 let mut migrated_constraints: BTreeMap<_, Vec<_>> = BTreeMap::new();
971 for col in table.columns.iter_mut() {
972 if let Some(constraint_name) = col.primary_key_constraint.take() {
973 migrated_constraints
974 .entry(constraint_name)
975 .or_default()
976 .push(col.name.to_string());
977 }
978 }
979
980 table.constraints = migrated_constraints
981 .into_iter()
982 .map(|(constraint_name, column_names)| SqlServerTableConstraint {
983 constraint_name: constraint_name.to_string(),
984 constraint_type: SqlServerTableConstraintType::PrimaryKey,
985 column_names,
986 })
987 .collect();
988
989 let new_value = SourceExportStatementDetails::SqlServer {
990 table,
991 capture_instance,
992 initial_lsn,
993 };
994 *deets = hex::encode(new_value.into_proto().encode_to_vec());
995
996 Ok(())
997}
998
999fn ast_rewrite_add_missing_index_ids(
1001 tx: &Transaction<'_>,
1002 stmt: &mut Statement<Raw>,
1003) -> Result<(), anyhow::Error> {
1004 let Statement::CreateIndex(stmt) = stmt else {
1005 return Ok(());
1006 };
1007
1008 let unresolved_name = match stmt.on_name.clone() {
1009 mz_sql::ast::RawItemName::Name(name) => name,
1010 mz_sql::ast::RawItemName::Id(..) => return Ok(()),
1012 };
1013
1014 let parts = &unresolved_name.0;
1015 let (db_name, schema_name, item_name) = match parts.len() {
1016 3 => (Some(&parts[0]), &parts[1], &parts[2]),
1017 2 => (None, &parts[0], &parts[1]),
1018 _ => panic!("invalid unresolved name: {unresolved_name:?}"),
1019 };
1020
1021 let db_id = db_name.map(|x| {
1022 let db = tx.get_databases().find(|db| db.name == x.as_str());
1023 let db = db.unwrap_or_else(|| panic!("missing database: {x}"));
1024 db.id
1025 });
1026 let schema_id = {
1027 let schema = tx
1028 .get_schemas()
1029 .find(|s| s.name == schema_name.as_str() && s.database_id == db_id);
1030 let schema = schema.unwrap_or_else(|| panic!("missing schema: {schema_name}, {db_id:?}"));
1031 schema.id
1032 };
1033 let item_id = {
1034 let item = tx
1035 .get_items()
1036 .find(|i| i.name == item_name.as_str() && i.schema_id == schema_id);
1037 let item = item.unwrap_or_else(|| panic!("missing item: {item_name}, {schema_id:?}"));
1038 item.id
1039 };
1040
1041 stmt.on_name = mz_sql::ast::RawItemName::Id(item_id.to_string(), unresolved_name, None);
1042
1043 Ok(())
1044}
1045
1046fn ast_rewrite_add_missing_doc_on_ids(
1058 tx: &Transaction<'_>,
1059 stmt: &mut Statement<Raw>,
1060) -> Result<(), anyhow::Error> {
1061 if !matches!(stmt, Statement::CreateSink(_)) {
1062 return Ok(());
1063 }
1064 rewrite_doc_on_ids(stmt, |name| {
1065 let parts = &name.0;
1066 let (db_name, schema_name, item_name) = match parts.len() {
1067 3 => (Some(&parts[0]), &parts[1], &parts[2]),
1068 2 => (None, &parts[0], &parts[1]),
1071 _ => panic!("invalid doc on reference: {name:?}"),
1072 };
1073 let db_id = db_name.map(|db_name| {
1074 let db = tx.get_databases().find(|db| db.name == db_name.as_str());
1075 let db = db.unwrap_or_else(|| panic!("missing database in doc on reference: {name:?}"));
1076 db.id
1077 });
1078 let schema = tx
1079 .get_schemas()
1080 .find(|s| s.name == schema_name.as_str() && s.database_id == db_id);
1081 let schema =
1082 schema.unwrap_or_else(|| panic!("missing schema in doc on reference: {name:?}"));
1083
1084 let user_type = tx.get_items().find(|i| {
1085 i.name == item_name.as_str()
1086 && i.schema_id == schema.id
1087 && i.item_type() == CatalogItemType::Type
1090 });
1091 if let Some(item) = user_type {
1092 return Some(item.id.to_string());
1093 }
1094 let builtin_type = tx.get_system_object_mappings().find(|m| {
1099 m.description.schema_name == schema.name
1100 && m.description.object_type == CatalogItemType::Type
1101 && m.description.object_name == item_name.as_str()
1102 });
1103 let builtin_type =
1104 builtin_type.unwrap_or_else(|| panic!("missing type in doc on reference: {name:?}"));
1105 Some(builtin_type.unique_identifier.catalog_id.to_string())
1106 });
1107 Ok(())
1108}
1109
1110fn rewrite_doc_on_ids(
1114 stmt: &mut Statement<Raw>,
1115 lookup: impl FnMut(&mz_sql::ast::UnresolvedItemName) -> Option<String>,
1116) {
1117 use mz_sql::ast::visit_mut::{VisitMut, VisitMutNode};
1118 use mz_sql::ast::{ColumnName, DocOnIdentifier, RawItemName, UnresolvedItemName};
1119
1120 struct Rewriter<F> {
1121 lookup: F,
1122 }
1123
1124 impl<'ast, F> VisitMut<'ast, Raw> for Rewriter<F>
1125 where
1126 F: FnMut(&UnresolvedItemName) -> Option<String>,
1127 {
1128 fn visit_doc_on_identifier_mut(&mut self, node: &mut DocOnIdentifier<Raw>) {
1129 let name = match node {
1130 DocOnIdentifier::Type(name) => name,
1131 DocOnIdentifier::Column(ColumnName {
1132 relation,
1133 column: _,
1134 }) => relation,
1135 };
1136 if let RawItemName::Name(unresolved) = name {
1137 if let Some(id) = (self.lookup)(unresolved) {
1138 let unresolved = unresolved.clone();
1139 *name = RawItemName::Id(id, unresolved, None);
1140 }
1141 }
1142 }
1143 }
1144
1145 stmt.visit_mut(&mut Rewriter { lookup });
1146}
1147
1148fn ast_rewrite_strip_builtin_version_pins(stmt: &mut Statement<Raw>) -> Result<(), anyhow::Error> {
1162 use mz_sql::ast::RawItemName;
1163 use mz_sql::ast::visit_mut::{VisitMut, VisitMutNode};
1164
1165 struct StripBuiltinVersionPins;
1166
1167 impl<'ast> VisitMut<'ast, Raw> for StripBuiltinVersionPins {
1168 fn visit_item_name_mut(&mut self, item_name: &mut RawItemName) {
1169 if let RawItemName::Id(id, _, version) = item_name {
1170 if version.is_some() {
1171 if let Ok(parsed) = id.parse::<CatalogItemId>() {
1172 if !parsed.is_user() {
1173 *version = None;
1174 }
1175 }
1176 }
1177 }
1178 }
1179 }
1180
1181 let mut visitor = StripBuiltinVersionPins;
1182 stmt.visit_mut(&mut visitor);
1183 Ok(())
1184}
1185
1186fn ast_rewrite_kafka_metadata_refresh_intervals(
1187 stmt: &mut Statement<Raw>,
1188) -> Result<(), anyhow::Error> {
1189 use mz_sql::ast::{
1190 CreateSinkConnection, CreateSourceConnection, KafkaSinkConfigOptionName,
1191 KafkaSourceConfigOptionName, WithOptionValue,
1192 };
1193 let interval: Option<&mut WithOptionValue<Raw>> = match stmt {
1198 Statement::CreateSource(stmt) => {
1199 if let CreateSourceConnection::Kafka { options, .. } = &mut stmt.connection {
1200 options.iter_mut().find_map(|option| {
1201 if matches!(
1202 option.name,
1203 KafkaSourceConfigOptionName::TopicMetadataRefreshInterval
1204 ) {
1205 option.value.as_mut()
1206 } else {
1207 None
1208 }
1209 })
1210 } else {
1211 None
1212 }
1213 }
1214 Statement::CreateSink(stmt) => {
1215 if let CreateSinkConnection::Kafka { options, .. } = &mut stmt.connection {
1216 options.iter_mut().find_map(|option| {
1217 if matches!(
1218 option.name,
1219 KafkaSinkConfigOptionName::TopicMetadataRefreshInterval
1220 ) {
1221 option.value.as_mut()
1222 } else {
1223 None
1224 }
1225 })
1226 } else {
1227 None
1228 }
1229 }
1230 _ => None,
1231 };
1232
1233 let Some(interval) = interval else {
1234 return Ok(());
1235 };
1236
1237 rewrite_interval_option_floor_1s(interval, "kafka metadata refresh interval")
1238}
1239
1240fn ast_rewrite_small_commit_intervals(stmt: &mut Statement<Raw>) -> Result<(), anyhow::Error> {
1248 use mz_sql::ast::{CreateSinkConnection, CreateSinkOptionName};
1249
1250 let Statement::CreateSink(stmt) = stmt else {
1251 return Ok(());
1252 };
1253 if !matches!(stmt.connection, CreateSinkConnection::Iceberg { .. }) {
1254 return Ok(());
1255 }
1256 let interval = stmt.with_options.iter_mut().find_map(|o| {
1257 if matches!(o.name, CreateSinkOptionName::CommitInterval) {
1258 o.value.as_mut()
1259 } else {
1260 None
1261 }
1262 });
1263 let Some(interval) = interval else {
1264 return Ok(());
1265 };
1266
1267 rewrite_interval_option_floor_1s(interval, "commit interval")
1268}
1269
1270fn rewrite_interval_option_floor_1s(
1272 value: &mut mz_sql::ast::WithOptionValue<Raw>,
1273 label: &str,
1274) -> Result<(), anyhow::Error> {
1275 use mz_sql::ast::{Value, WithOptionValue};
1276 use mz_sql::plan::TryFromValue;
1277
1278 let dur = Duration::try_from_value(value.clone())
1279 .map_err(|e| anyhow::anyhow!("invalid value for {label}: {value:?}: {e}"))?;
1280
1281 if dur < Duration::from_secs(1) {
1282 *value = WithOptionValue::Value(Value::String("1s".to_string()));
1283 }
1284
1285 Ok(())
1286}
1287
1288#[cfg(test)]
1289mod tests {
1290 use super::*;
1291
1292 fn strip(sql: &str) -> String {
1293 let mut stmt = mz_sql::parse::parse(sql)
1294 .expect("test sql parses")
1295 .into_element()
1296 .ast;
1297 ast_rewrite_strip_builtin_version_pins(&mut stmt).expect("rewrite succeeds");
1298 stmt.to_ast_string_stable()
1299 }
1300
1301 #[mz_ore::test]
1302 #[cfg_attr(miri, ignore)] fn strips_version_from_builtin_reference() {
1304 let out = strip(
1308 r#"CREATE VIEW "materialize"."public"."v" AS SELECT 1 FROM [s518 AS "mz_catalog"."mz_audit_events" VERSION 0]"#,
1309 );
1310 assert!(!out.contains("VERSION"), "version not stripped: {out}");
1311 assert!(out.contains("s518"), "reference dropped: {out}");
1313 }
1314
1315 #[mz_ore::test]
1316 #[cfg_attr(miri, ignore)] fn preserves_version_on_user_reference() {
1318 let out = strip(
1321 r#"CREATE VIEW "materialize"."public"."v" AS SELECT 1 FROM [u5 AS "materialize"."public"."t" VERSION 1]"#,
1322 );
1323 assert!(out.contains("VERSION"), "user version stripped: {out}");
1324 assert!(out.contains("u5"), "reference dropped: {out}");
1325 }
1326
1327 #[mz_ore::test]
1328 #[cfg_attr(miri, ignore)] fn leaves_unpinned_builtin_reference_untouched() {
1330 let out = strip(
1331 r#"CREATE VIEW "materialize"."public"."v" AS SELECT 1 FROM [s518 AS "mz_catalog"."mz_audit_events"]"#,
1332 );
1333 assert!(!out.contains("VERSION"), "unexpected version: {out}");
1334 assert!(out.contains("s518"), "reference dropped: {out}");
1335 }
1336
1337 fn add_doc_on_ids(
1338 sql: &str,
1339 lookup: impl FnMut(&mz_sql::ast::UnresolvedItemName) -> Option<String>,
1340 ) -> String {
1341 let mut stmt = mz_sql::parse::parse(sql)
1342 .expect("test sql parses")
1343 .into_element()
1344 .ast;
1345 rewrite_doc_on_ids(&mut stmt, lookup);
1346 stmt.to_ast_string_stable()
1347 }
1348
1349 #[mz_ore::test]
1350 #[cfg_attr(miri, ignore)] fn adds_ids_to_bare_doc_on_references() {
1352 let out = add_doc_on_ids(
1356 r#"CREATE SINK "materialize"."public"."s" FROM [u1 AS "materialize"."public"."t"] INTO KAFKA CONNECTION [u2 AS "materialize"."public"."kc"] (TOPIC = 'top') FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION [u3 AS "materialize"."public"."csr"] (DOC ON TYPE "materialize"."public"."point" = 'p', DOC ON COLUMN "materialize"."public"."point"."x" = 'x', DOC ON TYPE "pg_catalog"."int4" = 'i') ENVELOPE UPSERT"#,
1357 |name| match name.0.len() {
1358 3 => {
1359 assert_eq!(name.0[2].as_str(), "point", "unexpected lookup: {name:?}");
1360 Some("u9".into())
1361 }
1362 2 => {
1363 assert_eq!(name.0[1].as_str(), "int4", "unexpected lookup: {name:?}");
1364 Some("s23".into())
1365 }
1366 _ => panic!("unexpected lookup: {name:?}"),
1367 },
1368 );
1369 assert!(
1370 out.contains(r#"DOC ON TYPE [u9 AS "materialize"."public"."point"]"#),
1371 "type reference not rewritten: {out}"
1372 );
1373 assert!(
1374 out.contains(r#"DOC ON COLUMN [u9 AS "materialize"."public"."point"]."x""#),
1375 "column reference not rewritten: {out}"
1376 );
1377 assert!(
1378 out.contains(r#"DOC ON TYPE [s23 AS "pg_catalog"."int4"]"#),
1379 "builtin type reference not rewritten: {out}"
1380 );
1381 }
1382
1383 #[mz_ore::test]
1384 #[cfg_attr(miri, ignore)] fn leaves_id_doc_on_references_untouched() {
1386 let out = add_doc_on_ids(
1387 r#"CREATE SINK "materialize"."public"."s" FROM [u1 AS "materialize"."public"."t"] INTO KAFKA CONNECTION [u2 AS "materialize"."public"."kc"] (TOPIC = 'top') FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION [u3 AS "materialize"."public"."csr"] (DOC ON TYPE [u4 AS "materialize"."public"."point"] = 'p', DOC ON COLUMN [u1 AS "materialize"."public"."t"]."c1" = 'c') ENVELOPE UPSERT"#,
1388 |name| panic!("id reference must not be resolved: {name:?}"),
1389 );
1390 assert!(
1391 out.contains(r#"DOC ON TYPE [u4 AS "materialize"."public"."point"]"#),
1392 "id reference altered: {out}"
1393 );
1394 }
1395}