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 ast_rewrite_strip_unsupported_column_collations(stmt)?;
162 Ok(())
163 })?;
164
165 let commit_ts = tx.upper();
168 let mut item_updates = into_consolidatable_updates_startup(item_updates, commit_ts);
169 let op_item_updates = tx.get_and_commit_op_updates();
170 let op_item_updates = into_consolidatable_updates_startup(op_item_updates, commit_ts);
171 item_updates.extend(op_item_updates);
172 differential_dataflow::consolidation::consolidate_updates(&mut item_updates);
173
174 let (post_item_updates, item_updates): (Vec<_>, Vec<_>) = item_updates
178 .into_iter()
179 .partition(|(kind, _, _)| matches!(kind, StateUpdateKind::StorageCollectionMetadata(_)));
182
183 let item_updates = item_updates
184 .into_iter()
185 .map(|(kind, ts, diff)| StateUpdate {
186 kind,
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
421 | PgConfigOptionName::Publication
422 | PgConfigOptionName::Failover => true,
423 PgConfigOptionName::TextColumns | PgConfigOptionName::ExcludeColumns => false,
424 });
425 }
426 CreateSourceConnection::SqlServer { options, .. } => {
427 options.retain(|o| match o.name {
428 SqlServerConfigOptionName::Details => true,
429 SqlServerConfigOptionName::TextColumns
430 | SqlServerConfigOptionName::ExcludeColumns => false,
431 });
432 }
433 CreateSourceConnection::MySql { options, .. } => {
434 options.retain(|o| match o.name {
435 MySqlConfigOptionName::Details => true,
436 MySqlConfigOptionName::TextColumns | MySqlConfigOptionName::ExcludeColumns => {
437 false
438 }
439 });
440 }
441 CreateSourceConnection::Kafka { .. } | CreateSourceConnection::LoadGenerator { .. } => {
442 }
443 }
444
445 let (new_progress_name, new_progress_stmt, new_source_name, new_source_stmt) =
447 match connection {
448 connection @ (CreateSourceConnection::Postgres { .. }
449 | CreateSourceConnection::MySql { .. }
450 | CreateSourceConnection::SqlServer { .. }
451 | CreateSourceConnection::LoadGenerator {
452 generator:
453 LoadGenerator::Tpch | LoadGenerator::Auction | LoadGenerator::Marketing,
454 ..
455 }) => {
456 assert_eq!(col_names, &[]);
458 assert_eq!(key_constraint, None);
459 assert_eq!(format, None);
460 assert_eq!(envelope, None);
461 assert_eq!(include_metadata, &[]);
462 assert_eq!(external_references, None);
463
464 let dummy_source_stmt = Statement::CreateView(CreateViewStatement {
472 if_exists: IfExistsBehavior::Error,
473 temporary: false,
474 definition: ViewDefinition {
475 name: progress_name,
476 columns: vec![],
477 query: Query {
478 ctes: CteBlock::Simple(vec![]),
479 body: SetExpr::Table(RawItemName::Id(
480 progress_item.id().to_string(),
481 source_name.clone(),
482 None,
483 )),
484 order_by: vec![],
485 limit: None,
486 offset: None,
487 },
488 },
489 });
490
491 let new_progress_stmt = CreateSourceStatement {
492 name: source_name.clone(),
493 in_cluster,
494 col_names: vec![],
495 connection,
496 include_metadata: vec![],
497 format: None,
498 envelope: None,
499 if_not_exists,
500 key_constraint: None,
501 with_options,
502 external_references: None,
503 progress_subsource: None,
504 };
505
506 migrated_source_ids.insert(source_item.id, progress_item.id());
507
508 (
509 full_source_name.item,
510 new_progress_stmt,
511 progress_item.name().item.clone(),
512 dummy_source_stmt,
513 )
514 }
515 CreateSourceConnection::Kafka {
516 options,
517 connection,
518 } => {
519 let constraints = if let Some(_key_constraint) = key_constraint {
520 vec![]
523 } else {
524 vec![]
525 };
526
527 let columns = if col_names.is_empty() {
528 TableFromSourceColumns::NotSpecified
529 } else {
530 TableFromSourceColumns::Named(col_names)
531 };
532
533 let details = SourceExportStatementDetails::Kafka {};
536 let table_with_options = vec![TableFromSourceOption {
537 name: TableFromSourceOptionName::Details,
538 value: Some(WithOptionValue::Value(Value::String(hex::encode(
539 details.into_proto().encode_to_vec(),
540 )))),
541 }];
542 let topic_option = options
544 .iter()
545 .find(|o| matches!(o.name, KafkaSourceConfigOptionName::Topic))
546 .expect("kafka sources must have a topic");
547 let topic = match &topic_option.value {
548 Some(WithOptionValue::Value(Value::String(topic))) => topic,
549 _ => unreachable!("topic must be a string"),
550 };
551 let external_reference = UnresolvedItemName::qualified(&[Ident::new(topic)?]);
552
553 let new_source_stmt =
554 Statement::CreateTableFromSource(CreateTableFromSourceStatement {
555 name: source_name,
556 constraints,
557 columns,
558 if_not_exists,
559 source: raw_progress_name,
560 include_metadata,
561 format,
562 envelope,
563 external_reference: Some(external_reference),
564 with_options: table_with_options,
565 });
566
567 let new_progress_stmt = CreateSourceStatement {
568 name: progress_name,
569 in_cluster,
570 col_names: vec![],
571 connection: CreateSourceConnection::Kafka {
572 options,
573 connection,
574 },
575 include_metadata: vec![],
576 format: None,
577 envelope: None,
578 if_not_exists,
579 key_constraint: None,
580 with_options,
581 external_references: None,
582 progress_subsource: None,
583 };
584 (
585 progress_item.name().item.clone(),
586 new_progress_stmt,
587 full_source_name.item,
588 new_source_stmt,
589 )
590 }
591 CreateSourceConnection::LoadGenerator {
592 generator:
593 generator @ (LoadGenerator::Clock
594 | LoadGenerator::Counter
595 | LoadGenerator::Datums
596 | LoadGenerator::KeyValue),
597 options,
598 } => {
599 let constraints = if let Some(_key_constraint) = key_constraint {
600 vec![]
602 } else {
603 vec![]
604 };
605
606 let columns = if col_names.is_empty() {
607 TableFromSourceColumns::NotSpecified
608 } else {
609 TableFromSourceColumns::Named(col_names)
610 };
611
612 let details = SourceExportStatementDetails::LoadGenerator {
615 output: LoadGeneratorOutput::Default,
616 };
617 let table_with_options = vec![TableFromSourceOption {
618 name: TableFromSourceOptionName::Details,
619 value: Some(WithOptionValue::Value(Value::String(hex::encode(
620 details.into_proto().encode_to_vec(),
621 )))),
622 }];
623 let external_reference = FullItemName {
626 database: mz_sql::names::RawDatabaseSpecifier::Name(
627 mz_storage_types::sources::load_generator::LOAD_GENERATOR_DATABASE_NAME
628 .to_owned(),
629 ),
630 schema: generator.schema_name().to_string(),
631 item: generator.schema_name().to_string(),
632 };
633
634 let new_source_stmt =
635 Statement::CreateTableFromSource(CreateTableFromSourceStatement {
636 name: source_name,
637 constraints,
638 columns,
639 if_not_exists,
640 source: raw_progress_name,
641 include_metadata,
642 format,
643 envelope,
644 external_reference: Some(external_reference.into()),
645 with_options: table_with_options,
646 });
647
648 let new_progress_stmt = CreateSourceStatement {
649 name: progress_name,
650 in_cluster,
651 col_names: vec![],
652 connection: CreateSourceConnection::LoadGenerator { generator, options },
653 include_metadata: vec![],
654 format: None,
655 envelope: None,
656 if_not_exists,
657 key_constraint: None,
658 with_options,
659 external_references: None,
660 progress_subsource: None,
661 };
662 (
663 progress_item.name().item.clone(),
664 new_progress_stmt,
665 full_source_name.item,
666 new_source_stmt,
667 )
668 }
669 };
670
671 info!(
675 "migrate: converted source {} to {}",
676 source_item.create_sql, new_source_stmt
677 );
678 source_item.name = new_source_name.clone();
679 source_item.create_sql = new_source_stmt.to_ast_string_stable();
680 updated_items.insert(source_item.id, source_item);
681 pending_progress_items.insert(progress_item.id(), (new_progress_name, new_progress_stmt));
682 }
683
684 for (mut item, stmt) in subsources {
685 match stmt {
686 CreateSubsourceStatement {
689 of_source: None, ..
690 } => {
691 let Some((new_name, new_stmt)) = pending_progress_items.remove(&item.id) else {
692 panic!("encountered orphan progress subsource id: {}", item.id)
693 };
694 item.name = new_name;
695 item.create_sql = new_stmt.to_ast_string_stable();
696 updated_items.insert(item.id, item);
697 }
698 CreateSubsourceStatement {
701 name,
702 columns,
703 constraints,
704 of_source: Some(raw_source_name),
705 if_not_exists,
706 mut with_options,
707 } => {
708 let new_raw_source_name = match raw_source_name {
709 RawItemName::Id(old_id, name, None) => {
710 let old_id: CatalogItemId = old_id.parse().expect("well formed");
711 let new_id = migrated_source_ids[&old_id].clone();
712 RawItemName::Id(new_id.to_string(), name, None)
713 }
714 _ => unreachable!("unexpected source name: {raw_source_name}"),
715 };
716 let external_reference = match with_options
719 .iter()
720 .position(|opt| opt.name == CreateSubsourceOptionName::ExternalReference)
721 {
722 Some(i) => match with_options.remove(i).value {
723 Some(WithOptionValue::UnresolvedItemName(name)) => name,
724 _ => unreachable!("external reference must be an unresolved item name"),
725 },
726 None => panic!("subsource must have an external reference"),
727 };
728
729 let with_options = with_options
730 .into_iter()
731 .map(|option| {
732 match option.name {
733 CreateSubsourceOptionName::Details => TableFromSourceOption {
734 name: TableFromSourceOptionName::Details,
735 value: option.value,
738 },
739 CreateSubsourceOptionName::TextColumns => TableFromSourceOption {
740 name: TableFromSourceOptionName::TextColumns,
741 value: option.value,
742 },
743 CreateSubsourceOptionName::ExcludeColumns => TableFromSourceOption {
744 name: TableFromSourceOptionName::ExcludeColumns,
745 value: option.value,
746 },
747 CreateSubsourceOptionName::RetainHistory => TableFromSourceOption {
748 name: TableFromSourceOptionName::RetainHistory,
749 value: option.value,
750 },
751 CreateSubsourceOptionName::Progress => {
752 panic!("progress option should not exist on this subsource")
753 }
754 CreateSubsourceOptionName::ExternalReference => {
755 unreachable!("This option is handled separately above.")
756 }
757 }
758 })
759 .collect::<Vec<_>>();
760
761 let table = CreateTableFromSourceStatement {
762 name,
763 constraints,
764 columns: TableFromSourceColumns::Defined(columns),
765 if_not_exists,
766 source: new_raw_source_name,
767 external_reference: Some(external_reference),
768 with_options,
769 envelope: None,
771 include_metadata: vec![],
772 format: None,
773 };
774
775 info!(
776 "migrate: converted subsource {} to table {}",
777 item.create_sql, table
778 );
779 item.create_sql = Statement::CreateTableFromSource(table).to_ast_string_stable();
780 updated_items.insert(item.id, item);
781 }
782 }
783 }
784 assert!(
785 pending_progress_items.is_empty(),
786 "unexpected residual progress items: {pending_progress_items:?}"
787 );
788
789 tx.update_items(updated_items)?;
790
791 Ok(())
792}
793
794pub(crate) fn durable_migrate(
798 tx: &mut Transaction,
799 _organization_id: Uuid,
800 _boot_ts: Timestamp,
801) -> Result<(), anyhow::Error> {
802 const EXPR_CACHE_MIGRATION_KEY: &str = "expr_cache_migration";
805 const EXPR_CACHE_MIGRATION_DONE: u64 = 1;
806 if tx.get_config(EXPR_CACHE_MIGRATION_KEY.to_string()) != Some(EXPR_CACHE_MIGRATION_DONE) {
807 if let Some(shard_id) = tx.get_expression_cache_shard() {
808 tx.insert_unfinalized_shards(btreeset! {shard_id})?;
809 tx.set_expression_cache_shard(ShardId::new())?;
810 }
811 tx.set_config(
812 EXPR_CACHE_MIGRATION_KEY.to_string(),
813 Some(EXPR_CACHE_MIGRATION_DONE),
814 )?;
815 }
816
817 const BUILTIN_MIGRATION_SHARD_MIGRATION_KEY: &str = "migration_shard_migration";
820 const BUILTIN_MIGRATION_SHARD_MIGRATION_DONE: u64 = 1;
821 if tx.get_config(BUILTIN_MIGRATION_SHARD_MIGRATION_KEY.to_string())
822 != Some(BUILTIN_MIGRATION_SHARD_MIGRATION_DONE)
823 {
824 if let Some(shard_id) = tx.get_builtin_migration_shard() {
825 tx.insert_unfinalized_shards(btreeset! {shard_id})?;
826 tx.set_builtin_migration_shard(ShardId::new())?;
827 }
828 tx.set_config(
829 BUILTIN_MIGRATION_SHARD_MIGRATION_KEY.to_string(),
830 Some(BUILTIN_MIGRATION_SHARD_MIGRATION_DONE),
831 )?;
832 }
833
834 if tx
835 .get_setting(MOCK_AUTHENTICATION_NONCE_KEY.to_string())
836 .is_none()
837 {
838 let mut nonce = [0u8; 24];
839 openssl::rand::rand_bytes(&mut nonce).expect("failed to generate nonce");
840 let nonce = BASE64_STANDARD.encode(nonce);
841 tx.set_setting(MOCK_AUTHENTICATION_NONCE_KEY.to_string(), Some(nonce))?;
842 }
843
844 migrate_builtin_tables_to_mvs(tx)?;
845
846 Ok(())
847}
848
849fn migrate_builtin_tables_to_mvs(tx: &mut Transaction) -> Result<(), anyhow::Error> {
853 let expected_mvs: BTreeSet<_> = BUILTINS::materialized_views()
855 .map(|mv| (mv.schema, mv.name))
856 .collect();
857
858 let mut to_remove = BTreeSet::new();
860 let mut to_add = Vec::new();
861 for mapping in tx.get_system_object_mappings() {
862 let desc = &mapping.description;
863 if desc.object_type != CatalogItemType::Table {
864 continue;
865 }
866
867 let key = (&*desc.schema_name, &*desc.object_name);
868 if expected_mvs.contains(&key) {
869 info!(
870 "migrate: builtin {}.{} changed type from table to MV",
871 desc.schema_name, desc.object_name,
872 );
873 to_remove.insert(desc.clone());
874 to_add.push(SystemObjectMapping {
875 description: SystemObjectDescription {
876 schema_name: desc.schema_name.clone(),
877 object_type: CatalogItemType::MaterializedView,
878 object_name: desc.object_name.clone(),
879 },
880 unique_identifier: mapping.unique_identifier,
881 });
882 }
883 }
884
885 if !to_remove.is_empty() {
886 tx.remove_system_object_mappings(to_remove)?;
887 tx.set_system_object_mappings(to_add)?;
888 }
889
890 Ok(())
891}
892
893fn ast_rewrite_create_sink_partition_strategy(
905 stmt: &mut Statement<Raw>,
906) -> Result<(), anyhow::Error> {
907 let Statement::CreateSink(stmt) = stmt else {
908 return Ok(());
909 };
910 stmt.with_options
911 .retain(|op| op.name != CreateSinkOptionName::PartitionStrategy);
912 Ok(())
913}
914
915fn ast_rewrite_sql_server_constraints(stmt: &mut Statement<Raw>) -> Result<(), anyhow::Error> {
917 use mz_sql::ast::{
918 CreateSubsourceOptionName, TableFromSourceOptionName, Value, WithOptionValue,
919 };
920 use mz_sql_server_util::desc::{SqlServerTableConstraint, SqlServerTableConstraintType};
921 use mz_storage_types::sources::ProtoSourceExportStatementDetails;
922 use mz_storage_types::sources::proto_source_export_statement_details::Kind;
923
924 let deets: Option<&mut String> = match stmt {
925 Statement::CreateSubsource(stmt) => stmt.with_options.iter_mut().find_map(|option| {
926 if matches!(option.name, CreateSubsourceOptionName::Details)
927 && let Some(WithOptionValue::Value(Value::String(ref mut details))) = option.value
928 {
929 Some(details)
930 } else {
931 None
932 }
933 }),
934 Statement::CreateTableFromSource(stmt) => stmt.with_options.iter_mut().find_map(|option| {
935 if matches!(option.name, TableFromSourceOptionName::Details)
936 && let Some(WithOptionValue::Value(Value::String(ref mut details))) = option.value
937 {
938 Some(details)
939 } else {
940 None
941 }
942 }),
943 _ => None,
944 };
945 let Some(deets) = deets else {
946 return Ok(());
947 };
948
949 let current_value = hex::decode(&mut *deets)?;
950 let current_value = ProtoSourceExportStatementDetails::decode(&*current_value)?;
951
952 if !matches!(current_value.kind, Some(Kind::SqlServer(_))) {
954 return Ok(());
955 };
956
957 let SourceExportStatementDetails::SqlServer {
958 mut table,
959 capture_instance,
960 initial_lsn,
961 } = SourceExportStatementDetails::from_proto(current_value)?
962 else {
963 unreachable!("statement details must exist for SQL Server");
964 };
965
966 if !table.constraints.is_empty() {
968 return Ok(());
969 }
970
971 let mut migrated_constraints: BTreeMap<_, Vec<_>> = BTreeMap::new();
974 for col in table.columns.iter_mut() {
975 if let Some(constraint_name) = col.primary_key_constraint.take() {
976 migrated_constraints
977 .entry(constraint_name)
978 .or_default()
979 .push(col.name.to_string());
980 }
981 }
982
983 table.constraints = migrated_constraints
984 .into_iter()
985 .map(|(constraint_name, column_names)| SqlServerTableConstraint {
986 constraint_name: constraint_name.to_string(),
987 constraint_type: SqlServerTableConstraintType::PrimaryKey,
988 column_names,
989 })
990 .collect();
991
992 let new_value = SourceExportStatementDetails::SqlServer {
993 table,
994 capture_instance,
995 initial_lsn,
996 };
997 *deets = hex::encode(new_value.into_proto().encode_to_vec());
998
999 Ok(())
1000}
1001
1002fn ast_rewrite_add_missing_index_ids(
1004 tx: &Transaction<'_>,
1005 stmt: &mut Statement<Raw>,
1006) -> Result<(), anyhow::Error> {
1007 let Statement::CreateIndex(stmt) = stmt else {
1008 return Ok(());
1009 };
1010
1011 let unresolved_name = match stmt.on_name.clone() {
1012 mz_sql::ast::RawItemName::Name(name) => name,
1013 mz_sql::ast::RawItemName::Id(..) => return Ok(()),
1015 };
1016
1017 let parts = &unresolved_name.0;
1018 let (db_name, schema_name, item_name) = match parts.len() {
1019 3 => (Some(&parts[0]), &parts[1], &parts[2]),
1020 2 => (None, &parts[0], &parts[1]),
1021 _ => panic!("invalid unresolved name: {unresolved_name:?}"),
1022 };
1023
1024 let db_id = db_name.map(|x| {
1025 let db = tx.get_databases().find(|db| db.name == x.as_str());
1026 let db = db.unwrap_or_else(|| panic!("missing database: {x}"));
1027 db.id
1028 });
1029 let schema_id = {
1030 let schema = tx
1031 .get_schemas()
1032 .find(|s| s.name == schema_name.as_str() && s.database_id == db_id);
1033 let schema = schema.unwrap_or_else(|| panic!("missing schema: {schema_name}, {db_id:?}"));
1034 schema.id
1035 };
1036 let item_id = {
1037 let item = tx
1038 .get_items()
1039 .find(|i| i.name == item_name.as_str() && i.schema_id == schema_id);
1040 let item = item.unwrap_or_else(|| panic!("missing item: {item_name}, {schema_id:?}"));
1041 item.id
1042 };
1043
1044 stmt.on_name = mz_sql::ast::RawItemName::Id(item_id.to_string(), unresolved_name, None);
1045
1046 Ok(())
1047}
1048
1049fn ast_rewrite_add_missing_doc_on_ids(
1061 tx: &Transaction<'_>,
1062 stmt: &mut Statement<Raw>,
1063) -> Result<(), anyhow::Error> {
1064 if !matches!(stmt, Statement::CreateSink(_)) {
1065 return Ok(());
1066 }
1067 rewrite_doc_on_ids(stmt, |name| {
1068 let parts = &name.0;
1069 let (db_name, schema_name, item_name) = match parts.len() {
1070 3 => (Some(&parts[0]), &parts[1], &parts[2]),
1071 2 => (None, &parts[0], &parts[1]),
1074 _ => panic!("invalid doc on reference: {name:?}"),
1075 };
1076 let db_id = db_name.map(|db_name| {
1077 let db = tx.get_databases().find(|db| db.name == db_name.as_str());
1078 let db = db.unwrap_or_else(|| panic!("missing database in doc on reference: {name:?}"));
1079 db.id
1080 });
1081 let schema = tx
1082 .get_schemas()
1083 .find(|s| s.name == schema_name.as_str() && s.database_id == db_id);
1084 let schema =
1085 schema.unwrap_or_else(|| panic!("missing schema in doc on reference: {name:?}"));
1086
1087 let user_type = tx.get_items().find(|i| {
1088 i.name == item_name.as_str()
1089 && i.schema_id == schema.id
1090 && i.item_type() == CatalogItemType::Type
1093 });
1094 if let Some(item) = user_type {
1095 return Some(item.id.to_string());
1096 }
1097 let builtin_type = tx.get_system_object_mappings().find(|m| {
1102 m.description.schema_name == schema.name
1103 && m.description.object_type == CatalogItemType::Type
1104 && m.description.object_name == item_name.as_str()
1105 });
1106 let builtin_type =
1107 builtin_type.unwrap_or_else(|| panic!("missing type in doc on reference: {name:?}"));
1108 Some(builtin_type.unique_identifier.catalog_id.to_string())
1109 });
1110 Ok(())
1111}
1112
1113fn rewrite_doc_on_ids(
1117 stmt: &mut Statement<Raw>,
1118 lookup: impl FnMut(&mz_sql::ast::UnresolvedItemName) -> Option<String>,
1119) {
1120 use mz_sql::ast::visit_mut::{VisitMut, VisitMutNode};
1121 use mz_sql::ast::{ColumnName, DocOnIdentifier, RawItemName, UnresolvedItemName};
1122
1123 struct Rewriter<F> {
1124 lookup: F,
1125 }
1126
1127 impl<'ast, F> VisitMut<'ast, Raw> for Rewriter<F>
1128 where
1129 F: FnMut(&UnresolvedItemName) -> Option<String>,
1130 {
1131 fn visit_doc_on_identifier_mut(&mut self, node: &mut DocOnIdentifier<Raw>) {
1132 let name = match node {
1133 DocOnIdentifier::Type(name) => name,
1134 DocOnIdentifier::Column(ColumnName {
1135 relation,
1136 column: _,
1137 }) => relation,
1138 };
1139 if let RawItemName::Name(unresolved) = name {
1140 if let Some(id) = (self.lookup)(unresolved) {
1141 let unresolved = unresolved.clone();
1142 *name = RawItemName::Id(id, unresolved, None);
1143 }
1144 }
1145 }
1146 }
1147
1148 stmt.visit_mut(&mut Rewriter { lookup });
1149}
1150
1151fn ast_rewrite_strip_builtin_version_pins(stmt: &mut Statement<Raw>) -> Result<(), anyhow::Error> {
1165 use mz_sql::ast::RawItemName;
1166 use mz_sql::ast::visit_mut::{VisitMut, VisitMutNode};
1167
1168 struct StripBuiltinVersionPins;
1169
1170 impl<'ast> VisitMut<'ast, Raw> for StripBuiltinVersionPins {
1171 fn visit_item_name_mut(&mut self, item_name: &mut RawItemName) {
1172 if let RawItemName::Id(id, _, version) = item_name {
1173 if version.is_some() {
1174 if let Ok(parsed) = id.parse::<CatalogItemId>() {
1175 if !parsed.is_user() {
1176 *version = None;
1177 }
1178 }
1179 }
1180 }
1181 }
1182 }
1183
1184 let mut visitor = StripBuiltinVersionPins;
1185 stmt.visit_mut(&mut visitor);
1186 Ok(())
1187}
1188
1189fn ast_rewrite_kafka_metadata_refresh_intervals(
1190 stmt: &mut Statement<Raw>,
1191) -> Result<(), anyhow::Error> {
1192 use mz_sql::ast::{
1193 CreateSinkConnection, CreateSourceConnection, KafkaSinkConfigOptionName,
1194 KafkaSourceConfigOptionName, WithOptionValue,
1195 };
1196 let interval: Option<&mut WithOptionValue<Raw>> = match stmt {
1201 Statement::CreateSource(stmt) => {
1202 if let CreateSourceConnection::Kafka { options, .. } = &mut stmt.connection {
1203 options.iter_mut().find_map(|option| {
1204 if matches!(
1205 option.name,
1206 KafkaSourceConfigOptionName::TopicMetadataRefreshInterval
1207 ) {
1208 option.value.as_mut()
1209 } else {
1210 None
1211 }
1212 })
1213 } else {
1214 None
1215 }
1216 }
1217 Statement::CreateSink(stmt) => {
1218 if let CreateSinkConnection::Kafka { options, .. } = &mut stmt.connection {
1219 options.iter_mut().find_map(|option| {
1220 if matches!(
1221 option.name,
1222 KafkaSinkConfigOptionName::TopicMetadataRefreshInterval
1223 ) {
1224 option.value.as_mut()
1225 } else {
1226 None
1227 }
1228 })
1229 } else {
1230 None
1231 }
1232 }
1233 _ => None,
1234 };
1235
1236 let Some(interval) = interval else {
1237 return Ok(());
1238 };
1239
1240 rewrite_interval_option_floor_1s(interval, "kafka metadata refresh interval")
1241}
1242
1243fn ast_rewrite_small_commit_intervals(stmt: &mut Statement<Raw>) -> Result<(), anyhow::Error> {
1251 use mz_sql::ast::{CreateSinkConnection, CreateSinkOptionName};
1252
1253 let Statement::CreateSink(stmt) = stmt else {
1254 return Ok(());
1255 };
1256 if !matches!(stmt.connection, CreateSinkConnection::Iceberg { .. }) {
1257 return Ok(());
1258 }
1259 let interval = stmt.with_options.iter_mut().find_map(|o| {
1260 if matches!(o.name, CreateSinkOptionName::CommitInterval) {
1261 o.value.as_mut()
1262 } else {
1263 None
1264 }
1265 });
1266 let Some(interval) = interval else {
1267 return Ok(());
1268 };
1269
1270 rewrite_interval_option_floor_1s(interval, "commit interval")
1271}
1272
1273fn ast_rewrite_strip_unsupported_column_collations(
1278 stmt: &mut Statement<Raw>,
1279) -> Result<(), anyhow::Error> {
1280 let Statement::CreateTable(stmt) = stmt else {
1281 return Ok(());
1282 };
1283 for column in &mut stmt.columns {
1284 column
1285 .collation
1286 .take_if(|collation| !mz_sql::plan::is_supported_collation(collation));
1287 }
1288 Ok(())
1289}
1290
1291fn rewrite_interval_option_floor_1s(
1293 value: &mut mz_sql::ast::WithOptionValue<Raw>,
1294 label: &str,
1295) -> Result<(), anyhow::Error> {
1296 use mz_sql::ast::{Value, WithOptionValue};
1297 use mz_sql::plan::TryFromValue;
1298
1299 let dur = Duration::try_from_value(value.clone())
1300 .map_err(|e| anyhow::anyhow!("invalid value for {label}: {value:?}: {e}"))?;
1301
1302 if dur < Duration::from_secs(1) {
1303 *value = WithOptionValue::Value(Value::String("1s".to_string()));
1304 }
1305
1306 Ok(())
1307}
1308
1309#[cfg(test)]
1310mod tests {
1311 use super::*;
1312
1313 fn strip(sql: &str) -> String {
1314 let mut stmt = mz_sql::parse::parse(sql)
1315 .expect("test sql parses")
1316 .into_element()
1317 .ast;
1318 ast_rewrite_strip_builtin_version_pins(&mut stmt).expect("rewrite succeeds");
1319 stmt.to_ast_string_stable()
1320 }
1321
1322 #[mz_ore::test]
1323 #[cfg_attr(miri, ignore)] fn strips_version_from_builtin_reference() {
1325 let out = strip(
1329 r#"CREATE VIEW "materialize"."public"."v" AS SELECT 1 FROM [s518 AS "mz_catalog"."mz_audit_events" VERSION 0]"#,
1330 );
1331 assert!(!out.contains("VERSION"), "version not stripped: {out}");
1332 assert!(out.contains("s518"), "reference dropped: {out}");
1334 }
1335
1336 #[mz_ore::test]
1337 #[cfg_attr(miri, ignore)] fn preserves_version_on_user_reference() {
1339 let out = strip(
1342 r#"CREATE VIEW "materialize"."public"."v" AS SELECT 1 FROM [u5 AS "materialize"."public"."t" VERSION 1]"#,
1343 );
1344 assert!(out.contains("VERSION"), "user version stripped: {out}");
1345 assert!(out.contains("u5"), "reference dropped: {out}");
1346 }
1347
1348 #[mz_ore::test]
1349 #[cfg_attr(miri, ignore)] fn leaves_unpinned_builtin_reference_untouched() {
1351 let out = strip(
1352 r#"CREATE VIEW "materialize"."public"."v" AS SELECT 1 FROM [s518 AS "mz_catalog"."mz_audit_events"]"#,
1353 );
1354 assert!(!out.contains("VERSION"), "unexpected version: {out}");
1355 assert!(out.contains("s518"), "reference dropped: {out}");
1356 }
1357
1358 #[mz_ore::test]
1359 #[cfg_attr(miri, ignore)] fn strips_unsupported_column_collations() {
1361 let mut stmt = mz_sql::parse::parse(
1362 r#"CREATE TABLE "materialize"."public"."t" ("a" [s20 AS "pg_catalog"."text"] COLLATE "unicode", "b" [s21 AS "pg_catalog"."int4"] COLLATE "nonexistent", "c" [s20 AS "pg_catalog"."text"] COLLATE "pg_catalog"."default" NOT NULL)"#,
1363 )
1364 .expect("test sql parses")
1365 .into_element()
1366 .ast;
1367 ast_rewrite_strip_unsupported_column_collations(&mut stmt).expect("rewrite succeeds");
1368 assert_eq!(
1369 stmt.to_ast_string_stable(),
1370 r#"CREATE TABLE "materialize"."public"."t" ("a" [s20 AS "pg_catalog"."text"], "b" [s21 AS "pg_catalog"."int4"], "c" [s20 AS "pg_catalog"."text"] COLLATE "pg_catalog"."default" NOT NULL)"#,
1371 );
1372 }
1373
1374 fn add_doc_on_ids(
1375 sql: &str,
1376 lookup: impl FnMut(&mz_sql::ast::UnresolvedItemName) -> Option<String>,
1377 ) -> String {
1378 let mut stmt = mz_sql::parse::parse(sql)
1379 .expect("test sql parses")
1380 .into_element()
1381 .ast;
1382 rewrite_doc_on_ids(&mut stmt, lookup);
1383 stmt.to_ast_string_stable()
1384 }
1385
1386 #[mz_ore::test]
1387 #[cfg_attr(miri, ignore)] fn adds_ids_to_bare_doc_on_references() {
1389 let out = add_doc_on_ids(
1393 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"#,
1394 |name| match name.0.len() {
1395 3 => {
1396 assert_eq!(name.0[2].as_str(), "point", "unexpected lookup: {name:?}");
1397 Some("u9".into())
1398 }
1399 2 => {
1400 assert_eq!(name.0[1].as_str(), "int4", "unexpected lookup: {name:?}");
1401 Some("s23".into())
1402 }
1403 _ => panic!("unexpected lookup: {name:?}"),
1404 },
1405 );
1406 assert!(
1407 out.contains(r#"DOC ON TYPE [u9 AS "materialize"."public"."point"]"#),
1408 "type reference not rewritten: {out}"
1409 );
1410 assert!(
1411 out.contains(r#"DOC ON COLUMN [u9 AS "materialize"."public"."point"]."x""#),
1412 "column reference not rewritten: {out}"
1413 );
1414 assert!(
1415 out.contains(r#"DOC ON TYPE [s23 AS "pg_catalog"."int4"]"#),
1416 "builtin type reference not rewritten: {out}"
1417 );
1418 }
1419
1420 #[mz_ore::test]
1421 #[cfg_attr(miri, ignore)] fn leaves_id_doc_on_references_untouched() {
1423 let out = add_doc_on_ids(
1424 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"#,
1425 |name| panic!("id reference must not be resolved: {name:?}"),
1426 );
1427 assert!(
1428 out.contains(r#"DOC ON TYPE [u4 AS "materialize"."public"."point"]"#),
1429 "id reference altered: {out}"
1430 );
1431 }
1432}