1use std::collections::{BTreeMap, BTreeSet};
15use std::fmt;
16use std::iter;
17use std::path::Path;
18use std::sync::Arc;
19
20use anyhow::anyhow;
21use itertools::Itertools;
22use mz_adapter_types::dyncfgs::ENABLE_S3_TABLES_REGION_CHECK;
23use mz_ccsr::{Client, GetBySubjectError};
24use mz_cloud_provider::CloudProvider;
25use mz_controller_types::ClusterId;
26use mz_kafka_util::client::MzClientContext;
27use mz_mysql_util::MySqlTableDesc;
28use mz_ore::collections::CollectionExt;
29use mz_ore::error::ErrorExt;
30use mz_ore::future::InTask;
31use mz_ore::iter::IteratorExt;
32use mz_ore::str::StrExt;
33use mz_postgres_util::desc::PostgresTableDesc;
34use mz_proto::RustType;
35use mz_repr::{CatalogItemId, RelationDesc, RelationVersionSelector, Timestamp, strconv};
36use mz_sql_parser::ast::display::AstDisplay;
37use mz_sql_parser::ast::visit::{Visit, visit_function};
38use mz_sql_parser::ast::visit_mut::{VisitMut, visit_expr_mut};
39use mz_sql_parser::ast::{
40 AlterSourceAction, AlterSourceAddSubsourceOptionName, AlterSourceStatement, AvroDocOn,
41 CastFailureMode, ColumnName, CreateMaterializedViewStatement, CreateSinkConnection,
42 CreateSinkOptionName, CreateSinkStatement, CreateSourceOptionName, CreateSubsourceOption,
43 CreateSubsourceOptionName, CreateTableFromSourceStatement, CsrConfigOption,
44 CsrConfigOptionName, CsrConnection, CsrSeedAvro, CsrSeedProtobuf, CsrSeedProtobufSchema,
45 DeferredItemName, DocOnIdentifier, DocOnSchema, Expr, Function, FunctionArgs, GlueAvroOption,
46 GlueAvroSeed, Ident, KafkaSourceConfigOption, KafkaSourceConfigOptionName, LoadGenerator,
47 LoadGeneratorOption, LoadGeneratorOptionName, MaterializedViewOption,
48 MaterializedViewOptionName, MySqlConfigOption, MySqlConfigOptionName, PgConfigOption,
49 PgConfigOptionName, RawItemName, ReaderSchemaSelectionStrategy, RefreshAtOptionValue,
50 RefreshEveryOptionValue, RefreshOptionValue, SourceEnvelope, SqlServerConfigOption,
51 SqlServerConfigOptionName, Statement, TableFromSourceColumns, TableFromSourceOption,
52 TableFromSourceOptionName, UnresolvedItemName,
53};
54use mz_sql_server_util::desc::SqlServerTableDesc;
55use mz_storage_types::configuration::StorageConfiguration;
56use mz_storage_types::connections::Connection;
57use mz_storage_types::connections::inline::IntoInlineConnection;
58use mz_storage_types::errors::ContextCreationError;
59use mz_storage_types::sources::load_generator::LoadGeneratorOutput;
60use mz_storage_types::sources::mysql::MySqlSourceDetails;
61use mz_storage_types::sources::postgres::PostgresSourcePublicationDetails;
62use mz_storage_types::sources::{
63 GenericSourceConnection, MzOffset, SourceConnection, SourceDesc, SourceExportStatementDetails,
64 SqlServerSourceExtras,
65};
66use prost::Message;
67use protobuf_native::MessageLite;
68use protobuf_native::compiler::{SourceTreeDescriptorDatabase, VirtualSourceTree};
69use rdkafka::admin::AdminClient;
70use references::{RetrievedSourceReferences, SourceReferenceClient};
71use uuid::Uuid;
72
73use crate::ast::{
74 AlterSourceAddSubsourceOption, AvroSchema, CreateSourceConnection, CreateSourceStatement,
75 CreateSubsourceStatement, CsrConnectionAvro, CsrConnectionProtobuf, ExternalReferenceExport,
76 ExternalReferences, Format, FormatSpecifier, ProtobufSchema, Value, WithOptionValue,
77};
78use crate::catalog::{CatalogItemType, SessionCatalog};
79use crate::kafka_util::{KafkaSinkConfigOptionExtracted, KafkaSourceConfigOptionExtracted};
80use crate::names::{
81 Aug, FullItemName, PartialItemName, ResolvedColumnReference, ResolvedDataType, ResolvedIds,
82 ResolvedItemName,
83};
84use crate::plan::error::PlanError;
85use crate::plan::statement::ddl::load_generator_ast_to_generator;
86use crate::plan::{SourceReferences, StatementContext};
87use crate::pure::error::{IcebergSinkPurificationError, SqlServerSourcePurificationError};
88use crate::pure::mysql::{ensure_binlog_full_metadata, is_binlog_full_metadata};
89use crate::{kafka_util, normalize};
90
91use self::error::{
92 CsrPurificationError, KafkaSinkPurificationError, KafkaSourcePurificationError,
93 LoadGeneratorSourcePurificationError, MySqlSourcePurificationError, PgSourcePurificationError,
94};
95
96pub(crate) mod error;
97mod references;
98mod source_preparation;
99
100pub use source_preparation::{
101 SourcePreparation, SourcePreparationRequest, SourcePreparationResponse, SourcePreparationStatus,
102};
103
104pub mod mysql;
105pub mod postgres;
106pub mod sql_server;
107
108pub(crate) struct RequestedSourceExport<T> {
109 external_reference: UnresolvedItemName,
110 name: UnresolvedItemName,
111 meta: T,
112}
113
114impl<T: fmt::Debug> fmt::Debug for RequestedSourceExport<T> {
115 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
116 f.debug_struct("RequestedSourceExport")
117 .field("external_reference", &self.external_reference)
118 .field("name", &self.name)
119 .field("meta", &self.meta)
120 .finish()
121 }
122}
123
124impl<T> RequestedSourceExport<T> {
125 fn change_meta<F>(self, new_meta: F) -> RequestedSourceExport<F> {
126 RequestedSourceExport {
127 external_reference: self.external_reference,
128 name: self.name,
129 meta: new_meta,
130 }
131 }
132}
133
134fn source_export_name_gen(
139 source_name: &UnresolvedItemName,
140 subsource_name: &str,
141) -> Result<UnresolvedItemName, PlanError> {
142 let mut partial = normalize::unresolved_item_name(source_name.clone())?;
143 partial.item = subsource_name.to_string();
144 Ok(UnresolvedItemName::from(partial))
145}
146
147fn validate_source_export_names<T>(
151 requested_source_exports: &[RequestedSourceExport<T>],
152) -> Result<(), PlanError> {
153 if let Some(name) = requested_source_exports
157 .iter()
158 .map(|subsource| &subsource.name)
159 .duplicates()
160 .next()
161 .cloned()
162 {
163 let mut upstream_references: Vec<_> = requested_source_exports
164 .into_iter()
165 .filter_map(|subsource| {
166 if &subsource.name == &name {
167 Some(subsource.external_reference.clone())
168 } else {
169 None
170 }
171 })
172 .collect();
173
174 upstream_references.sort();
175
176 Err(PlanError::SubsourceNameConflict {
177 name,
178 upstream_references,
179 })?;
180 }
181
182 if let Some(name) = requested_source_exports
188 .iter()
189 .map(|export| &export.external_reference)
190 .duplicates()
191 .next()
192 .cloned()
193 {
194 let mut target_names: Vec<_> = requested_source_exports
195 .into_iter()
196 .filter_map(|export| {
197 if &export.external_reference == &name {
198 Some(export.name.clone())
199 } else {
200 None
201 }
202 })
203 .collect();
204
205 target_names.sort();
206
207 Err(PlanError::SubsourceDuplicateReference { name, target_names })?;
208 }
209
210 Ok(())
211}
212
213#[derive(Debug, Clone, PartialEq, Eq)]
214pub enum PurifiedStatement {
215 PurifiedCreateSource {
216 create_progress_subsource: bool,
218 create_source_stmt: CreateSourceStatement<Aug>,
219 subsources: BTreeMap<UnresolvedItemName, PurifiedSourceExport>,
221 available_source_references: SourceReferences,
224 },
225 PurifiedAlterSource {
226 alter_source_stmt: AlterSourceStatement<Aug>,
227 },
228 PurifiedAlterSourceAddSubsources {
229 source_name: ResolvedItemName,
231 options: Vec<AlterSourceAddSubsourceOption<Aug>>,
234 subsources: BTreeMap<UnresolvedItemName, PurifiedSourceExport>,
236 },
237 PurifiedAlterSourceRefreshReferences {
238 source_name: ResolvedItemName,
239 available_source_references: SourceReferences,
241 },
242 PurifiedCreateSink(CreateSinkStatement<Aug>),
243 PurifiedCreateTableFromSource {
244 stmt: CreateTableFromSourceStatement<Aug>,
245 },
246}
247
248#[derive(Debug, Clone, PartialEq, Eq)]
249pub struct PurifiedSourceExport {
250 pub external_reference: UnresolvedItemName,
251 pub details: PurifiedExportDetails,
252}
253
254#[derive(Debug, Clone, PartialEq, Eq)]
255pub enum PurifiedExportDetails {
256 MySql {
257 table: MySqlTableDesc,
258 text_columns: Option<Vec<Ident>>,
259 exclude_columns: Option<Vec<Ident>>,
260 initial_gtid_set: String,
261 binlog_full_metadata: bool,
262 },
263 Postgres {
264 table: PostgresTableDesc,
265 text_columns: Option<Vec<Ident>>,
266 exclude_columns: Option<Vec<Ident>>,
267 initial_lsn: MzOffset,
268 },
269 SqlServer {
270 table: SqlServerTableDesc,
271 text_columns: Option<Vec<Ident>>,
272 excl_columns: Option<Vec<Ident>>,
273 capture_instance: Arc<str>,
274 initial_lsn: mz_sql_server_util::cdc::Lsn,
275 },
276 Kafka {},
277 LoadGenerator {
278 table: Option<RelationDesc>,
279 output: LoadGeneratorOutput,
280 },
281}
282
283#[derive(Debug, Clone, Copy)]
289pub enum StatementSource {
290 Altered(CatalogItemId),
291 Read(CatalogItemId),
292}
293
294impl StatementSource {
295 pub fn id(&self) -> CatalogItemId {
296 match self {
297 StatementSource::Altered(id) | StatementSource::Read(id) => *id,
298 }
299 }
300}
301
302pub fn statement_source(
312 catalog: &impl SessionCatalog,
313 stmt: &Statement<Aug>,
314) -> Option<StatementSource> {
315 let scx = StatementContext::new(None, catalog);
316 match stmt {
317 Statement::AlterSource(stmt) => {
318 let item = scx
319 .resolve_item(RawItemName::Name(stmt.source_name.clone()))
320 .ok()?;
321 (item.item_type() == CatalogItemType::Source)
322 .then(|| StatementSource::Altered(item.id()))
323 }
324 Statement::CreateTableFromSource(stmt) => {
325 let item = scx.get_item_by_resolved_name(&stmt.source).ok()?;
326 (item.item_type() == CatalogItemType::Source).then(|| StatementSource::Read(item.id()))
327 }
328 _ => None,
329 }
330}
331
332pub async fn purify_statement(
342 catalog: impl SessionCatalog,
343 now: u64,
344 stmt: Statement<Aug>,
345 storage_configuration: &StorageConfiguration,
346) -> (Result<PurifiedStatement, PlanError>, Option<ClusterId>) {
347 match stmt {
348 Statement::AlterSource(stmt) => (
349 source_preparation::purify(
350 &catalog,
351 now,
352 Statement::AlterSource(stmt),
353 storage_configuration,
354 )
355 .await,
356 None,
357 ),
358 Statement::CreateSink(stmt) => {
359 let cluster_id = stmt.in_cluster.as_ref().map(|cluster| cluster.id.clone());
360 (
361 purify_create_sink(catalog, stmt, storage_configuration).await,
362 cluster_id,
363 )
364 }
365 o => (
366 Err(internal_err!(
367 "unexpected statement type in purification: {:?}",
368 o
369 )),
370 None,
371 ),
372 }
373}
374
375pub(crate) fn purify_create_sink_avro_doc_on_options(
380 catalog: &dyn SessionCatalog,
381 from_id: CatalogItemId,
382 format: &mut Option<FormatSpecifier<Aug>>,
383) -> Result<(), PlanError> {
384 let from = catalog.get_item(&from_id);
386 let object_ids = from
387 .references()
388 .items()
389 .copied()
390 .chain_one(from.id())
391 .collect::<Vec<_>>();
392
393 let mut avro_format_options = vec![];
396 for_each_format(format, |doc_on_schema, fmt| match fmt {
397 Format::Avro(AvroSchema::InlineSchema { .. })
398 | Format::Avro(AvroSchema::Glue { .. })
399 | Format::Bytes
400 | Format::Csv { .. }
401 | Format::Json { .. }
402 | Format::Protobuf(..)
403 | Format::Regex(..)
404 | Format::Text => (),
405 Format::Avro(AvroSchema::Csr {
406 csr_connection: CsrConnectionAvro { connection, .. },
407 }) => {
408 avro_format_options.push((doc_on_schema, &mut connection.options));
409 }
410 });
411
412 for (for_schema, options) in avro_format_options {
415 let user_provided_comments = options
416 .iter()
417 .filter_map(|CsrConfigOption { name, .. }| match name {
418 CsrConfigOptionName::AvroDocOn(doc_on) => Some(doc_on.clone()),
419 _ => None,
420 })
421 .collect::<BTreeSet<_>>();
422
423 for object_id in &object_ids {
425 let item = catalog
427 .get_item(object_id)
428 .at_version(RelationVersionSelector::Latest);
429 let full_resolved_name = ResolvedItemName::Item {
430 id: *object_id,
431 qualifiers: item.name().qualifiers.clone(),
432 full_name: catalog.resolve_full_name(item.name()),
433 print_id: !matches!(item.item_type(), CatalogItemType::Func),
434 version: RelationVersionSelector::Latest,
435 };
436
437 if let Some(comments_map) = catalog.get_item_comments(object_id) {
438 let doc_on_item_key = AvroDocOn {
441 identifier: DocOnIdentifier::Type(full_resolved_name.clone()),
442 for_schema,
443 };
444 if !user_provided_comments.contains(&doc_on_item_key) {
445 if let Some(root_comment) = comments_map.get(&None) {
446 options.push(CsrConfigOption {
447 name: CsrConfigOptionName::AvroDocOn(doc_on_item_key),
448 value: Some(mz_sql_parser::ast::WithOptionValue::Value(Value::String(
449 root_comment.clone(),
450 ))),
451 });
452 }
453 }
454
455 let column_descs = match item.type_details() {
459 Some(details) => details.typ.desc(catalog).unwrap_or_default(),
460 None => item.relation_desc().map(|d| d.into_owned()),
461 };
462
463 if let Some(desc) = column_descs {
464 for (pos, column_name) in desc.iter_names().enumerate() {
465 if let Some(comment_str) = comments_map.get(&Some(pos + 1)) {
466 let doc_on_column_key = AvroDocOn {
467 identifier: DocOnIdentifier::Column(ColumnName {
468 relation: full_resolved_name.clone(),
469 column: ResolvedColumnReference::Column {
470 name: column_name.to_owned(),
471 index: pos,
472 },
473 }),
474 for_schema,
475 };
476 if !user_provided_comments.contains(&doc_on_column_key) {
477 options.push(CsrConfigOption {
478 name: CsrConfigOptionName::AvroDocOn(doc_on_column_key),
479 value: Some(mz_sql_parser::ast::WithOptionValue::Value(
480 Value::String(comment_str.clone()),
481 )),
482 });
483 }
484 }
485 }
486 }
487 }
488 }
489 }
490
491 Ok(())
492}
493
494async fn purify_create_sink(
497 catalog: impl SessionCatalog,
498 mut create_sink_stmt: CreateSinkStatement<Aug>,
499 storage_configuration: &StorageConfiguration,
500) -> Result<PurifiedStatement, PlanError> {
501 let CreateSinkStatement {
503 connection,
504 format,
505 with_options,
506 name: _,
507 in_cluster: _,
508 if_not_exists: _,
509 from,
510 envelope: _,
511 mode: _,
512 } = &mut create_sink_stmt;
513
514 const USER_ALLOWED_WITH_OPTIONS: &[CreateSinkOptionName] = &[
516 CreateSinkOptionName::Snapshot,
517 CreateSinkOptionName::CommitInterval,
518 ];
519
520 if let Some(op) = with_options
521 .iter()
522 .find(|op| !USER_ALLOWED_WITH_OPTIONS.contains(&op.name))
523 {
524 sql_bail!(
525 "CREATE SINK...WITH ({}..) is not allowed",
526 op.name.to_ast_string_simple(),
527 )
528 }
529
530 match &connection {
531 CreateSinkConnection::Kafka {
532 connection,
533 options,
534 key: _,
535 headers: _,
536 } => {
537 let scx = StatementContext::new(None, &catalog);
542 let connection = {
543 let item = scx.get_item_by_resolved_name(connection)?;
544 match item.connection()? {
546 Connection::Kafka(connection) => {
547 connection.clone().into_inline_connection(scx.catalog)
548 }
549 _ => sql_bail!(
550 "{} is not a kafka connection",
551 scx.catalog.resolve_full_name(item.name())
552 ),
553 }
554 };
555
556 let extracted_options: KafkaSinkConfigOptionExtracted = options.clone().try_into()?;
557
558 if extracted_options.legacy_ids == Some(true) {
559 sql_bail!("LEGACY IDs option is not supported");
560 }
561
562 let client: AdminClient<_> = connection
563 .create_with_context(
564 storage_configuration,
565 MzClientContext::default(),
566 &BTreeMap::new(),
567 InTask::No,
568 )
569 .await
570 .map_err(|e| {
571 KafkaSinkPurificationError::AdminClientError(Arc::new(e))
573 })?;
574
575 let metadata = client
576 .inner()
577 .fetch_metadata(
578 None,
579 storage_configuration
580 .parameters
581 .kafka_timeout_config
582 .fetch_metadata_timeout,
583 )
584 .map_err(|e| {
585 KafkaSinkPurificationError::AdminClientError(Arc::new(
586 ContextCreationError::KafkaError(e),
587 ))
588 })?;
589
590 if metadata.brokers().len() == 0 {
591 Err(KafkaSinkPurificationError::ZeroBrokers)?;
592 }
593 }
594 CreateSinkConnection::Iceberg {
595 catalog_connection,
596 aws_connection,
597 ..
598 } => {
599 let scx = StatementContext::new(None, &catalog);
600 let connection = {
601 let item = scx.get_item_by_resolved_name(catalog_connection)?;
602 match item.connection()? {
604 Connection::IcebergCatalog(connection) => {
605 connection.clone().into_inline_connection(scx.catalog)
606 }
607 _ => sql_bail!(
608 "{} is not an iceberg connection",
609 scx.catalog.resolve_full_name(item.name())
610 ),
611 }
612 };
613
614 if let Some(s3tables) = connection.s3tables_catalog() {
618 let enable_region_check =
619 ENABLE_S3_TABLES_REGION_CHECK.get(scx.catalog.system_vars().dyncfgs());
620 if enable_region_check {
621 let env_id = &catalog.config().environment_id;
622 if matches!(env_id.cloud_provider(), CloudProvider::Aws) {
623 let env_region = env_id.cloud_provider_region();
624 let s3_tables_region = s3tables
627 .aws_connection
628 .connection
629 .region
630 .clone()
631 .unwrap_or_else(|| "us-east-1".to_string());
632 if s3_tables_region != env_region {
633 Err(IcebergSinkPurificationError::S3TablesRegionMismatch {
634 s3_tables_region,
635 environment_region: env_region.to_string(),
636 })?;
637 }
638 }
639 }
640 }
641
642 if let Some(aws_connection) = aws_connection {
648 let aws_conn_id = aws_connection.item_id();
649 let aws_connection = {
650 let item = scx.get_item_by_resolved_name(aws_connection)?;
651 match item.connection()? {
653 Connection::Aws(aws_connection) => aws_connection.clone(),
654 _ => sql_bail!(
655 "{} is not an aws connection",
656 scx.catalog.resolve_full_name(item.name())
657 ),
658 }
659 };
660
661 let _sdk_config = aws_connection
662 .load_sdk_config(
663 &storage_configuration.connection_context,
664 aws_conn_id.clone(),
665 InTask::No,
666 mz_storage_types::dyncfgs::ENFORCE_EXTERNAL_ADDRESSES
667 .get(storage_configuration.config_set()),
668 )
669 .await
670 .map_err(|e| IcebergSinkPurificationError::AwsSdkContextError(Arc::new(e)))?;
671 }
672
673 let _catalog = connection
679 .connect(storage_configuration, InTask::No, None)
680 .await
681 .map_err(|e| IcebergSinkPurificationError::CatalogError(Arc::new(e)))?;
682 }
683 }
684
685 let mut csr_connection_ids = BTreeSet::new();
686 for_each_format(format, |_, fmt| match fmt {
687 Format::Avro(AvroSchema::InlineSchema { .. })
688 | Format::Avro(AvroSchema::Glue { .. })
689 | Format::Bytes
690 | Format::Csv { .. }
691 | Format::Json { .. }
692 | Format::Protobuf(ProtobufSchema::InlineSchema { .. })
693 | Format::Regex(..)
694 | Format::Text => (),
695 Format::Avro(AvroSchema::Csr {
696 csr_connection: CsrConnectionAvro { connection, .. },
697 })
698 | Format::Protobuf(ProtobufSchema::Csr {
699 csr_connection: CsrConnectionProtobuf { connection, .. },
700 }) => {
701 csr_connection_ids.insert(*connection.connection.item_id());
702 }
703 });
704
705 let scx = StatementContext::new(None, &catalog);
706 for csr_connection_id in csr_connection_ids {
707 let connection = {
708 let item = scx.get_item(&csr_connection_id);
709 match item.connection()? {
711 Connection::Csr(connection) => connection.clone().into_inline_connection(&catalog),
712 _ => Err(CsrPurificationError::NotCsrConnection(
713 scx.catalog.resolve_full_name(item.name()),
714 ))?,
715 }
716 };
717
718 let client = connection
719 .connect(storage_configuration, InTask::No)
720 .await
721 .map_err(|e| CsrPurificationError::ClientError(Arc::new(e)))?;
722
723 client
724 .list_subjects()
725 .await
726 .map_err(|e| CsrPurificationError::ListSubjectsError(Arc::new(e)))?;
727 }
728
729 purify_create_sink_avro_doc_on_options(&catalog, *from.item_id(), format)?;
730
731 Ok(PurifiedStatement::PurifiedCreateSink(create_sink_stmt))
732}
733
734fn for_each_format<'a, F>(format: &'a mut Option<FormatSpecifier<Aug>>, mut f: F)
741where
742 F: FnMut(DocOnSchema, &'a mut Format<Aug>),
743{
744 match format {
745 None => (),
746 Some(FormatSpecifier::Bare(fmt)) => f(DocOnSchema::All, fmt),
747 Some(FormatSpecifier::KeyValue { key, value }) => {
748 f(DocOnSchema::KeyOnly, key);
749 f(DocOnSchema::ValueOnly, value);
750 }
751 }
752}
753
754#[derive(Debug, Copy, Clone, PartialEq, Eq)]
757pub(crate) enum SourceReferencePolicy {
758 NotAllowed,
761 Optional,
764 Required,
767}
768
769#[derive(Debug, Clone, PartialEq, Eq)]
770pub(crate) enum FilterConstraints {
771 ExcludeAll,
772 Exclude(BTreeSet<String>),
773}
774
775impl FilterConstraints {
776 pub(crate) fn none() -> Self {
777 FilterConstraints::Exclude(BTreeSet::new())
778 }
779
780 fn from_options(
781 exclude_constraints: Vec<String>,
782 exclude_all_constraints: bool,
783 ) -> Result<Self, PlanError> {
784 if !exclude_all_constraints {
785 return Ok(FilterConstraints::Exclude(
786 exclude_constraints.into_iter().collect(),
787 ));
788 }
789 if !exclude_constraints.is_empty() {
790 sql_bail!("EXCLUDE ALL CONSTRAINTS cannot be combined with EXCLUDE CONSTRAINTS");
791 }
792 Ok(FilterConstraints::ExcludeAll)
793 }
794
795 pub(crate) fn excludes_any(&self) -> bool {
796 match self {
797 FilterConstraints::ExcludeAll => true,
798 FilterConstraints::Exclude(names) => !names.is_empty(),
799 }
800 }
801}
802
803#[derive(Clone, Debug, PartialEq, Eq)]
804struct CreateSourcePreparation {
805 statement: CreateSourceStatement<Aug>,
806 connection: PreparedSourceConnection,
807 reference_policy: SourceReferencePolicy,
808 now: u64,
809 clock_as_of: Option<u64>,
810}
811
812#[derive(Clone, Debug, PartialEq, Eq)]
813enum PreparedSourceConnection {
814 Kafka(mz_storage_types::connections::KafkaConnection),
815 Postgres(
816 CatalogItemId,
817 mz_storage_types::connections::PostgresConnection,
818 ),
819 SqlServer(
820 CatalogItemId,
821 mz_storage_types::connections::SqlServerConnectionDetails,
822 ),
823 MySql(
824 CatalogItemId,
825 mz_storage_types::connections::MySqlConnection,
826 ),
827 LoadGenerator(mz_storage_types::sources::load_generator::LoadGenerator),
828}
829
830#[derive(Clone, Debug)]
831struct PreparedCreateSource {
832 create_progress_subsource: bool,
833 create_source_stmt: CreateSourceStatement<Aug>,
834 subsources: BTreeMap<UnresolvedItemName, PurifiedSourceExport>,
835 available_source_references: crate::plan::SourceReferences,
836 format_options: SourceFormatOptions,
837}
838
839impl CreateSourcePreparation {
840 fn new(
841 catalog: &dyn SessionCatalog,
842 now: u64,
843 clock_as_of: impl FnOnce() -> u64,
844 statement: CreateSourceStatement<Aug>,
845 ) -> Result<Self, PlanError> {
846 if let Some(DeferredItemName::Named(_)) = &statement.progress_subsource {
847 sql_bail!("Cannot manually ID qualify progress subsource")
848 }
849 let scx = StatementContext::new(None, catalog);
850
851 let reference_policy = if scx.catalog.system_vars().enable_create_table_from_source()
855 && scx.catalog.system_vars().force_source_table_syntax()
856 {
857 SourceReferencePolicy::NotAllowed
858 } else if scx.catalog.system_vars().enable_create_table_from_source() {
859 SourceReferencePolicy::Optional
860 } else {
861 SourceReferencePolicy::Required
862 };
863 let connection = match &statement.connection {
864 CreateSourceConnection::Kafka { connection, .. } => {
865 if let Some(references) = &statement.external_references {
866 return Err(KafkaSourcePurificationError::ReferencedSubsources(
867 references.clone(),
868 )
869 .into());
870 }
871 let item = scx.get_item_by_resolved_name(connection)?;
872 let Connection::Kafka(connection) = item.connection()? else {
873 return Err(KafkaSourcePurificationError::NotKafkaConnection(
874 scx.catalog.resolve_full_name(item.name()),
875 )
876 .into());
877 };
878 PreparedSourceConnection::Kafka(connection.clone().into_inline_connection(catalog))
879 }
880 CreateSourceConnection::Postgres { connection, .. } => {
881 let item = scx.get_item_by_resolved_name(connection)?;
882 let Connection::Postgres(connection) = item.connection()? else {
883 return Err(PgSourcePurificationError::NotPgConnection(
884 scx.catalog.resolve_full_name(item.name()),
885 )
886 .into());
887 };
888 PreparedSourceConnection::Postgres(
889 item.id(),
890 connection.clone().into_inline_connection(catalog),
891 )
892 }
893 CreateSourceConnection::SqlServer { connection, .. } => {
894 let item = scx.get_item_by_resolved_name(connection)?;
895 let Connection::SqlServer(connection) = item.connection()? else {
896 return Err(SqlServerSourcePurificationError::NotSqlServerConnection(
897 scx.catalog.resolve_full_name(item.name()),
898 )
899 .into());
900 };
901 PreparedSourceConnection::SqlServer(
902 item.id(),
903 connection.clone().into_inline_connection(catalog),
904 )
905 }
906 CreateSourceConnection::MySql { connection, .. } => {
907 let item = scx.get_item_by_resolved_name(connection)?;
908 let Connection::MySql(connection) = item.connection()? else {
909 return Err(MySqlSourcePurificationError::NotMySqlConnection(
910 scx.catalog.resolve_full_name(item.name()),
911 )
912 .into());
913 };
914 PreparedSourceConnection::MySql(
915 item.id(),
916 connection.clone().into_inline_connection(catalog),
917 )
918 }
919 CreateSourceConnection::LoadGenerator { generator, options } => {
920 PreparedSourceConnection::LoadGenerator(load_generator_ast_to_generator(
921 &scx,
922 generator,
923 options,
924 &statement.include_metadata,
925 )?)
926 }
927 };
928 let clock_as_of = match &statement.connection {
929 CreateSourceConnection::LoadGenerator {
930 generator: LoadGenerator::Clock,
931 options,
932 } if !options
933 .iter()
934 .any(|option| option.name == LoadGeneratorOptionName::AsOf) =>
935 {
936 Some(clock_as_of())
937 }
938 _ => None,
939 };
940 Ok(Self {
941 statement,
942 connection,
943 reference_policy,
944 now,
945 clock_as_of,
946 })
947 }
948}
949
950async fn execute_create_source_preparation(
951 request: CreateSourcePreparation,
952 storage_configuration: &StorageConfiguration,
953) -> Result<PreparedCreateSource, PlanError> {
954 let CreateSourcePreparation {
955 statement: mut create_source_stmt,
956 connection,
957 reference_policy,
958 now,
959 clock_as_of,
960 } = request;
961 let CreateSourceStatement {
962 name: source_name,
963 col_names,
964 key_constraint,
965 connection: source_connection,
966 format,
967 envelope,
968 include_metadata,
969 external_references,
970 progress_subsource,
971 ..
972 } = &mut create_source_stmt;
973
974 let uses_old_syntax = !col_names.is_empty()
975 || key_constraint.is_some()
976 || format.is_some()
977 || envelope.is_some()
978 || !include_metadata.is_empty()
979 || external_references.is_some()
980 || progress_subsource.is_some();
981
982 let mut requested_subsource_map = BTreeMap::new();
983
984 let mut format_options = SourceFormatOptions::Default;
985
986 let retrieved_source_references: RetrievedSourceReferences;
987
988 match (source_connection, connection) {
989 (
990 CreateSourceConnection::Kafka {
991 options: base_with_options,
992 ..
993 },
994 PreparedSourceConnection::Kafka(connection),
995 ) => {
996 let extracted_options: KafkaSourceConfigOptionExtracted =
997 base_with_options.clone().try_into()?;
998
999 let topic = extracted_options
1000 .topic
1001 .ok_or(KafkaSourcePurificationError::ConnectionMissingTopic)?;
1002
1003 let consumer = connection
1004 .create_with_context(
1005 storage_configuration,
1006 MzClientContext::default(),
1007 &BTreeMap::new(),
1008 InTask::No,
1009 )
1010 .await
1011 .map_err(|e| {
1012 KafkaSourcePurificationError::KafkaConsumerError(
1014 e.display_with_causes().to_string(),
1015 )
1016 })?;
1017 let consumer = Arc::new(consumer);
1018
1019 match (
1020 extracted_options.start_offset,
1021 extracted_options.start_timestamp,
1022 ) {
1023 (None, None) => {
1024 kafka_util::ensure_topic_exists(
1026 Arc::clone(&consumer),
1027 &topic,
1028 storage_configuration
1029 .parameters
1030 .kafka_timeout_config
1031 .fetch_metadata_timeout,
1032 )
1033 .await?;
1034 }
1035 (Some(_), Some(_)) => {
1036 sql_bail!("cannot specify START TIMESTAMP and START OFFSET at same time")
1037 }
1038 (Some(start_offsets), None) => {
1039 kafka_util::validate_start_offsets(
1041 Arc::clone(&consumer),
1042 &topic,
1043 start_offsets,
1044 storage_configuration
1045 .parameters
1046 .kafka_timeout_config
1047 .fetch_metadata_timeout,
1048 )
1049 .await?;
1050 }
1051 (None, Some(time_offset)) => {
1052 let start_offsets = kafka_util::lookup_start_offsets(
1054 Arc::clone(&consumer),
1055 &topic,
1056 time_offset,
1057 now,
1058 storage_configuration
1059 .parameters
1060 .kafka_timeout_config
1061 .fetch_metadata_timeout,
1062 )
1063 .await?;
1064
1065 base_with_options.retain(|val| {
1066 !matches!(val.name, KafkaSourceConfigOptionName::StartTimestamp)
1067 });
1068 base_with_options.push(KafkaSourceConfigOption {
1069 name: KafkaSourceConfigOptionName::StartOffset,
1070 value: Some(WithOptionValue::Sequence(
1071 start_offsets
1072 .iter()
1073 .map(|offset| {
1074 WithOptionValue::Value(Value::Number(offset.to_string()))
1075 })
1076 .collect(),
1077 )),
1078 });
1079 }
1080 }
1081
1082 let reference_client = SourceReferenceClient::Kafka { topic: &topic };
1083 retrieved_source_references = reference_client.get_source_references().await?;
1084
1085 format_options = SourceFormatOptions::Kafka { topic };
1086 }
1087 (
1088 CreateSourceConnection::Postgres { options, .. },
1089 PreparedSourceConnection::Postgres(connection_id, connection),
1090 ) => {
1091 let crate::plan::statement::PgConfigOptionExtracted {
1092 publication,
1093 failover,
1094 text_columns,
1095 exclude_columns,
1096 details,
1097 ..
1098 } = options.clone().try_into()?;
1099 let publication =
1100 publication.ok_or(PgSourcePurificationError::ConnectionMissingPublication)?;
1101
1102 if details.is_some() {
1103 Err(PgSourcePurificationError::UserSpecifiedDetails)?;
1104 }
1105
1106 let client = connection
1107 .validate(connection_id, storage_configuration)
1108 .await?;
1109
1110 let reference_client = SourceReferenceClient::Postgres {
1111 client: &client,
1112 publication: &publication,
1113 database: &connection.database,
1114 };
1115 retrieved_source_references = reference_client.get_source_references().await?;
1116
1117 let is_physical_replica = mz_postgres_util::get_is_in_recovery(&client).await?;
1120
1121 if failover
1123 && (is_physical_replica
1124 || mz_postgres_util::schemas::get_pg_major_version(&client).await? < 17)
1125 {
1126 Err(PgSourcePurificationError::FailoverUnsupported)?;
1127 }
1128
1129 let initial_lsn = MzOffset::from(
1131 mz_postgres_util::fetch_max_lsn(&client, is_physical_replica).await?,
1132 );
1133
1134 let postgres::PurifiedSourceExports {
1135 source_exports: subsources,
1136 normalized_text_columns,
1137 } = postgres::purify_source_exports(
1138 &client,
1139 &retrieved_source_references,
1140 external_references,
1141 text_columns,
1142 exclude_columns,
1143 &FilterConstraints::none(),
1144 false,
1145 source_name,
1146 &reference_policy,
1147 initial_lsn,
1148 )
1149 .await?;
1150
1151 if let Some(text_cols_option) = options
1152 .iter_mut()
1153 .find(|option| option.name == PgConfigOptionName::TextColumns)
1154 {
1155 text_cols_option.value = Some(WithOptionValue::Sequence(normalized_text_columns));
1156 }
1157
1158 requested_subsource_map.extend(subsources);
1159
1160 let timeline_id = mz_postgres_util::get_timeline_id(&client).await?;
1163
1164 options.retain(|PgConfigOption { name, .. }| name != &PgConfigOptionName::Details);
1166 let details = PostgresSourcePublicationDetails {
1167 slot: format!(
1168 "materialize_{}",
1169 Uuid::new_v4().to_string().replace('-', "")
1170 ),
1171 timeline_id: Some(timeline_id),
1172 database: connection.database,
1173 is_physical_replica: Some(is_physical_replica),
1174 };
1175 options.push(PgConfigOption {
1176 name: PgConfigOptionName::Details,
1177 value: Some(WithOptionValue::Value(Value::String(hex::encode(
1178 details.into_proto().encode_to_vec(),
1179 )))),
1180 })
1181 }
1182 (
1183 CreateSourceConnection::SqlServer { options, .. },
1184 PreparedSourceConnection::SqlServer(connection_id, connection),
1185 ) => {
1186 let crate::plan::statement::ddl::SqlServerConfigOptionExtracted {
1189 details,
1190 text_columns,
1191 exclude_columns,
1192 seen: _,
1193 } = options.clone().try_into()?;
1194
1195 if details.is_some() {
1196 Err(SqlServerSourcePurificationError::UserSpecifiedDetails)?;
1197 }
1198
1199 let mut client = connection
1200 .validate(connection_id, storage_configuration)
1201 .await?;
1202
1203 let database: Arc<str> = connection.database.into();
1204 let reference_client = SourceReferenceClient::SqlServer {
1205 client: &mut client,
1206 database: Arc::clone(&database),
1207 };
1208 retrieved_source_references = reference_client.get_source_references().await?;
1209 tracing::debug!(?retrieved_source_references, "got source references");
1210
1211 let timeout = mz_storage_types::sources::sql_server::MAX_LSN_WAIT
1212 .get(storage_configuration.config_set());
1213
1214 let purified_source_exports = sql_server::purify_source_exports(
1215 &*database,
1216 &mut client,
1217 &retrieved_source_references,
1218 external_references,
1219 &text_columns,
1220 &exclude_columns,
1221 source_name,
1222 timeout,
1223 &reference_policy,
1224 )
1225 .await?;
1226
1227 let sql_server::PurifiedSourceExports {
1228 source_exports,
1229 normalized_text_columns,
1230 normalized_excl_columns,
1231 } = purified_source_exports;
1232
1233 requested_subsource_map.extend(source_exports);
1235
1236 let restore_history_id =
1240 mz_sql_server_util::inspect::get_latest_restore_history_id(&mut client).await?;
1241 let details = SqlServerSourceExtras { restore_history_id };
1242
1243 options.retain(|SqlServerConfigOption { name, .. }| {
1244 name != &SqlServerConfigOptionName::Details
1245 });
1246 options.push(SqlServerConfigOption {
1247 name: SqlServerConfigOptionName::Details,
1248 value: Some(WithOptionValue::Value(Value::String(hex::encode(
1249 details.into_proto().encode_to_vec(),
1250 )))),
1251 });
1252
1253 if let Some(text_cols_option) = options
1255 .iter_mut()
1256 .find(|option| option.name == SqlServerConfigOptionName::TextColumns)
1257 {
1258 text_cols_option.value = Some(WithOptionValue::Sequence(normalized_text_columns));
1259 }
1260 if let Some(excl_cols_option) = options
1261 .iter_mut()
1262 .find(|option| option.name == SqlServerConfigOptionName::ExcludeColumns)
1263 {
1264 excl_cols_option.value = Some(WithOptionValue::Sequence(normalized_excl_columns));
1265 }
1266 }
1267 (
1268 CreateSourceConnection::MySql { options, .. },
1269 PreparedSourceConnection::MySql(connection_id, connection),
1270 ) => {
1271 let crate::plan::statement::ddl::MySqlConfigOptionExtracted {
1272 details,
1273 text_columns,
1274 exclude_columns,
1275 seen: _,
1276 } = options.clone().try_into()?;
1277
1278 if details.is_some() {
1279 Err(MySqlSourcePurificationError::UserSpecifiedDetails)?;
1280 }
1281
1282 let mut conn = connection
1283 .validate(connection_id, storage_configuration)
1284 .await
1285 .map_err(MySqlSourcePurificationError::InvalidConnection)?;
1286
1287 let initial_gtid_set =
1291 mz_mysql_util::query_sys_var(&mut conn, "global.gtid_executed").await?;
1292
1293 let binlog_full_metadata = is_binlog_full_metadata(&mut conn).await?;
1294
1295 let reference_client = SourceReferenceClient::MySql {
1296 conn: &mut conn,
1297 include_system_schemas: mysql::references_system_schemas(external_references),
1298 };
1299 retrieved_source_references = reference_client.get_source_references().await?;
1300
1301 let mysql::PurifiedSourceExports {
1302 source_exports: subsources,
1303 normalized_text_columns,
1304 normalized_exclude_columns,
1305 } = mysql::purify_source_exports(
1306 &mut conn,
1307 &retrieved_source_references,
1308 external_references,
1309 text_columns,
1310 exclude_columns,
1311 &FilterConstraints::none(),
1312 false,
1313 source_name,
1314 initial_gtid_set.clone(),
1315 &reference_policy,
1316 binlog_full_metadata,
1317 )
1318 .await?;
1319 requested_subsource_map.extend(subsources);
1320
1321 let details = MySqlSourceDetails {};
1324 options
1326 .retain(|MySqlConfigOption { name, .. }| name != &MySqlConfigOptionName::Details);
1327 options.push(MySqlConfigOption {
1328 name: MySqlConfigOptionName::Details,
1329 value: Some(WithOptionValue::Value(Value::String(hex::encode(
1330 details.into_proto().encode_to_vec(),
1331 )))),
1332 });
1333
1334 if let Some(text_cols_option) = options
1335 .iter_mut()
1336 .find(|option| option.name == MySqlConfigOptionName::TextColumns)
1337 {
1338 text_cols_option.value = Some(WithOptionValue::Sequence(normalized_text_columns));
1339 }
1340 if let Some(exclude_cols_option) = options
1341 .iter_mut()
1342 .find(|option| option.name == MySqlConfigOptionName::ExcludeColumns)
1343 {
1344 exclude_cols_option.value =
1345 Some(WithOptionValue::Sequence(normalized_exclude_columns));
1346 }
1347 }
1348 (
1349 CreateSourceConnection::LoadGenerator { generator, options },
1350 PreparedSourceConnection::LoadGenerator(load_generator),
1351 ) => {
1352 let reference_client = SourceReferenceClient::LoadGenerator {
1353 generator: &load_generator,
1354 };
1355 retrieved_source_references = reference_client.get_source_references().await?;
1356 let multi_output_sources =
1360 retrieved_source_references
1361 .all_references()
1362 .iter()
1363 .any(|r| {
1364 matches!(
1365 r.load_generator_output(),
1366 Some(output) if output != &LoadGeneratorOutput::Default
1367 )
1368 });
1369
1370 match external_references {
1371 Some(requested)
1372 if matches!(reference_policy, SourceReferencePolicy::NotAllowed) =>
1373 {
1374 Err(PlanError::UseTablesForSources(requested.to_string()))?
1375 }
1376 Some(requested) if !multi_output_sources => match requested {
1377 ExternalReferences::SubsetTables(_) => {
1378 Err(LoadGeneratorSourcePurificationError::ForTables)?
1379 }
1380 ExternalReferences::SubsetSchemas(_) => {
1381 Err(LoadGeneratorSourcePurificationError::ForSchemas)?
1382 }
1383 ExternalReferences::All => {
1384 Err(LoadGeneratorSourcePurificationError::ForAllTables)?
1385 }
1386 },
1387 Some(requested) => {
1388 let requested_exports = retrieved_source_references
1389 .requested_source_exports(Some(requested), source_name)?;
1390 for export in requested_exports {
1391 requested_subsource_map.insert(
1392 export.name,
1393 PurifiedSourceExport {
1394 external_reference: export.external_reference,
1395 details: PurifiedExportDetails::LoadGenerator {
1396 table: export
1397 .meta
1398 .load_generator_desc()
1399 .ok_or_else(|| {
1400 internal_err!(
1401 "expected load generator source reference"
1402 )
1403 })?
1404 .clone(),
1405 output: export
1406 .meta
1407 .load_generator_output()
1408 .ok_or_else(|| {
1409 internal_err!(
1410 "expected load generator source reference"
1411 )
1412 })?
1413 .clone(),
1414 },
1415 },
1416 );
1417 }
1418 }
1419 None => {
1420 if multi_output_sources
1421 && matches!(reference_policy, SourceReferencePolicy::Required)
1422 {
1423 Err(LoadGeneratorSourcePurificationError::MultiOutputRequiresForAllTables)?
1424 }
1425 }
1426 }
1427
1428 if let LoadGenerator::Clock = generator {
1429 if !options
1430 .iter()
1431 .any(|p| p.name == LoadGeneratorOptionName::AsOf)
1432 {
1433 let now =
1434 clock_as_of.expect("implicit CLOCK AS OF was captured during preparation");
1435 options.push(LoadGeneratorOption {
1436 name: LoadGeneratorOptionName::AsOf,
1437 value: Some(WithOptionValue::Value(Value::Number(now.to_string()))),
1438 });
1439 }
1440 }
1441 }
1442 _ => unreachable!("source preparation inputs must match their statement"),
1443 }
1444
1445 *external_references = None;
1449
1450 Ok(PreparedCreateSource {
1451 format_options,
1452 create_progress_subsource: uses_old_syntax,
1453 create_source_stmt,
1454 subsources: requested_subsource_map,
1455 available_source_references: retrieved_source_references.available_source_references(),
1456 })
1457}
1458
1459#[derive(Clone, Debug, PartialEq, Eq)]
1462enum AlterSourcePreparation {
1463 AddSubsources {
1464 external_references: Vec<ExternalReferenceExport>,
1465 options: Vec<AlterSourceAddSubsourceOption<Aug>>,
1466 desc: SourceDesc,
1467 partial_name: PartialItemName,
1468 unresolved_source_name: UnresolvedItemName,
1469 resolved_source_name: ResolvedItemName,
1470 },
1471 RefreshReferences {
1472 desc: SourceDesc,
1473 resolved_source_name: ResolvedItemName,
1474 },
1475}
1476
1477fn prepare_alter_source(
1478 catalog: &dyn SessionCatalog,
1479 stmt: &AlterSourceStatement<Aug>,
1480) -> Result<Option<AlterSourcePreparation>, PlanError> {
1481 let scx = StatementContext::new(None, catalog);
1482 let AlterSourceStatement {
1483 source_name: unresolved_source_name,
1484 action,
1485 if_exists,
1486 } = stmt;
1487 let item = match scx.resolve_item(RawItemName::Name(unresolved_source_name.clone())) {
1488 Ok(item) => item,
1489 Err(_) if *if_exists => return Ok(None),
1490 Err(e) => return Err(e),
1491 };
1492 let desc = match item.source_desc()? {
1493 Some(desc) => desc.clone().into_inline_connection(scx.catalog),
1494 None => sql_bail!("cannot ALTER this type of source"),
1495 };
1496 let source_name = item.name();
1497 let resolved_source_name = ResolvedItemName::Item {
1498 id: item.id(),
1499 qualifiers: item.name().qualifiers.clone(),
1500 full_name: scx.catalog.resolve_full_name(source_name),
1501 print_id: true,
1502 version: RelationVersionSelector::Latest,
1503 };
1504 let partial_name = scx.catalog.minimal_qualification(source_name);
1505 match action {
1506 AlterSourceAction::AddSubsources {
1507 external_references,
1508 options,
1509 } => {
1510 if scx.catalog.system_vars().enable_create_table_from_source()
1511 && scx.catalog.system_vars().force_source_table_syntax()
1512 {
1513 return Err(PlanError::UseTablesForSources(
1514 "ALTER SOURCE .. ADD SUBSOURCES ..".to_string(),
1515 ));
1516 }
1517 Ok(Some(AlterSourcePreparation::AddSubsources {
1518 external_references: external_references.clone(),
1519 options: options.clone(),
1520 desc,
1521 partial_name,
1522 unresolved_source_name: unresolved_source_name.clone(),
1523 resolved_source_name,
1524 }))
1525 }
1526 AlterSourceAction::RefreshReferences => {
1527 Ok(Some(AlterSourcePreparation::RefreshReferences {
1528 desc,
1529 resolved_source_name,
1530 }))
1531 }
1532 _ => Ok(None),
1533 }
1534}
1535
1536async fn execute_alter_source_preparation(
1537 request: AlterSourcePreparation,
1538 storage_configuration: &StorageConfiguration,
1539) -> Result<PurifiedStatement, PlanError> {
1540 match request {
1541 AlterSourcePreparation::AddSubsources {
1542 external_references,
1543 options,
1544 desc,
1545 partial_name,
1546 unresolved_source_name,
1547 resolved_source_name,
1548 } => {
1549 purify_alter_source_add_subsources(
1550 external_references,
1551 options,
1552 desc,
1553 partial_name,
1554 unresolved_source_name,
1555 resolved_source_name,
1556 storage_configuration,
1557 )
1558 .await
1559 }
1560 AlterSourcePreparation::RefreshReferences {
1561 desc,
1562 resolved_source_name,
1563 } => {
1564 purify_alter_source_refresh_references(
1565 desc,
1566 resolved_source_name,
1567 storage_configuration,
1568 )
1569 .await
1570 }
1571 }
1572}
1573
1574async fn purify_alter_source_add_subsources(
1576 external_references: Vec<ExternalReferenceExport>,
1577 mut options: Vec<AlterSourceAddSubsourceOption<Aug>>,
1578 desc: SourceDesc,
1579 partial_source_name: PartialItemName,
1580 unresolved_source_name: UnresolvedItemName,
1581 resolved_source_name: ResolvedItemName,
1582 storage_configuration: &StorageConfiguration,
1583) -> Result<PurifiedStatement, PlanError> {
1584 let connection_id = match &desc.connection {
1586 GenericSourceConnection::Postgres(c) => c.connection_id,
1587 GenericSourceConnection::MySql(c) => c.connection_id,
1588 GenericSourceConnection::SqlServer(c) => c.connection_id,
1589 _ => sql_bail!(
1590 "source {} does not support ALTER SOURCE.",
1591 partial_source_name
1592 ),
1593 };
1594
1595 let crate::plan::statement::ddl::AlterSourceAddSubsourceOptionExtracted {
1596 text_columns,
1597 exclude_columns,
1598 details,
1599 seen: _,
1600 } = options.clone().try_into()?;
1601 if details.is_some() {
1602 sql_bail!("DETAILS option cannot be explicitly set");
1603 }
1604
1605 let mut requested_subsource_map = BTreeMap::new();
1606
1607 match desc.connection {
1608 GenericSourceConnection::Postgres(pg_source_connection) => {
1609 let pg_connection = &pg_source_connection.connection;
1611
1612 let client = pg_connection
1613 .validate(connection_id, storage_configuration)
1614 .await?;
1615
1616 let reference_client = SourceReferenceClient::Postgres {
1617 client: &client,
1618 publication: &pg_source_connection.publication,
1619 database: &pg_connection.database,
1620 };
1621 let retrieved_source_references = reference_client.get_source_references().await?;
1622
1623 let initial_lsn = MzOffset::from(
1625 mz_postgres_util::fetch_max_lsn(
1626 &client,
1627 pg_source_connection
1628 .publication_details
1629 .get_is_physical_replica(),
1630 )
1631 .await?,
1632 );
1633
1634 let postgres::PurifiedSourceExports {
1635 source_exports: subsources,
1636 normalized_text_columns,
1637 } = postgres::purify_source_exports(
1638 &client,
1639 &retrieved_source_references,
1640 &Some(ExternalReferences::SubsetTables(external_references)),
1641 text_columns,
1642 exclude_columns,
1643 &FilterConstraints::none(),
1644 false,
1645 &unresolved_source_name,
1646 &SourceReferencePolicy::Required,
1647 initial_lsn,
1648 )
1649 .await?;
1650
1651 if let Some(text_cols_option) = options
1652 .iter_mut()
1653 .find(|option| option.name == AlterSourceAddSubsourceOptionName::TextColumns)
1654 {
1655 text_cols_option.value = Some(WithOptionValue::Sequence(normalized_text_columns));
1656 }
1657
1658 requested_subsource_map.extend(subsources);
1659 }
1660 GenericSourceConnection::MySql(mysql_source_connection) => {
1661 let mysql_connection = &mysql_source_connection.connection;
1662 let config = mysql_connection
1663 .config(
1664 &storage_configuration.connection_context.secrets_reader,
1665 storage_configuration,
1666 InTask::No,
1667 )
1668 .await?;
1669
1670 let mut conn = config
1671 .connect(
1672 "mysql purification",
1673 &storage_configuration.connection_context.ssh_tunnel_manager,
1674 )
1675 .await?;
1676
1677 let initial_gtid_set =
1680 mz_mysql_util::query_sys_var(&mut conn, "global.gtid_executed").await?;
1681
1682 let binlog_full_metadata = is_binlog_full_metadata(&mut conn).await?;
1683
1684 let requested_references = Some(ExternalReferences::SubsetTables(external_references));
1685
1686 let reference_client = SourceReferenceClient::MySql {
1687 conn: &mut conn,
1688 include_system_schemas: mysql::references_system_schemas(&requested_references),
1689 };
1690 let retrieved_source_references = reference_client.get_source_references().await?;
1691
1692 let mysql::PurifiedSourceExports {
1693 source_exports: subsources,
1694 normalized_text_columns,
1695 normalized_exclude_columns,
1696 } = mysql::purify_source_exports(
1697 &mut conn,
1698 &retrieved_source_references,
1699 &requested_references,
1700 text_columns,
1701 exclude_columns,
1702 &FilterConstraints::none(),
1703 false,
1704 &unresolved_source_name,
1705 initial_gtid_set,
1706 &SourceReferencePolicy::Required,
1707 binlog_full_metadata,
1708 )
1709 .await?;
1710 requested_subsource_map.extend(subsources);
1711
1712 if let Some(text_cols_option) = options
1714 .iter_mut()
1715 .find(|option| option.name == AlterSourceAddSubsourceOptionName::TextColumns)
1716 {
1717 text_cols_option.value = Some(WithOptionValue::Sequence(normalized_text_columns));
1718 }
1719 if let Some(exclude_cols_option) = options
1720 .iter_mut()
1721 .find(|option| option.name == AlterSourceAddSubsourceOptionName::ExcludeColumns)
1722 {
1723 exclude_cols_option.value =
1724 Some(WithOptionValue::Sequence(normalized_exclude_columns));
1725 }
1726 }
1727 GenericSourceConnection::SqlServer(sql_server_source) => {
1728 let sql_server_connection = &sql_server_source.connection;
1730 let config = sql_server_connection
1731 .resolve_config(
1732 &storage_configuration.connection_context.secrets_reader,
1733 storage_configuration,
1734 InTask::No,
1735 )
1736 .await?;
1737 let mut client = mz_sql_server_util::Client::connect(config).await?;
1738
1739 let database = sql_server_connection.database.clone().into();
1741 let source_references = SourceReferenceClient::SqlServer {
1742 client: &mut client,
1743 database: Arc::clone(&database),
1744 }
1745 .get_source_references()
1746 .await?;
1747 let requested_references = Some(ExternalReferences::SubsetTables(external_references));
1748
1749 let timeout = mz_storage_types::sources::sql_server::MAX_LSN_WAIT
1750 .get(storage_configuration.config_set());
1751
1752 let result = sql_server::purify_source_exports(
1753 &*database,
1754 &mut client,
1755 &source_references,
1756 &requested_references,
1757 &text_columns,
1758 &exclude_columns,
1759 &unresolved_source_name,
1760 timeout,
1761 &SourceReferencePolicy::Required,
1762 )
1763 .await;
1764 let sql_server::PurifiedSourceExports {
1765 source_exports,
1766 normalized_text_columns,
1767 normalized_excl_columns,
1768 } = result?;
1769
1770 requested_subsource_map.extend(source_exports);
1772
1773 if let Some(text_cols_option) = options
1775 .iter_mut()
1776 .find(|option| option.name == AlterSourceAddSubsourceOptionName::TextColumns)
1777 {
1778 text_cols_option.value = Some(WithOptionValue::Sequence(normalized_text_columns));
1779 }
1780 if let Some(exclude_cols_option) = options
1781 .iter_mut()
1782 .find(|option| option.name == AlterSourceAddSubsourceOptionName::ExcludeColumns)
1783 {
1784 exclude_cols_option.value =
1785 Some(WithOptionValue::Sequence(normalized_excl_columns));
1786 }
1787 }
1788 _ => bail_internal!("source does not support ALTER SOURCE...ADD SUBSOURCE"),
1791 };
1792
1793 Ok(PurifiedStatement::PurifiedAlterSourceAddSubsources {
1794 source_name: resolved_source_name,
1795 options,
1796 subsources: requested_subsource_map,
1797 })
1798}
1799
1800async fn purify_alter_source_refresh_references(
1801 desc: SourceDesc,
1802 resolved_source_name: ResolvedItemName,
1803 storage_configuration: &StorageConfiguration,
1804) -> Result<PurifiedStatement, PlanError> {
1805 let retrieved_source_references = match desc.connection {
1806 GenericSourceConnection::Postgres(pg_source_connection) => {
1807 let pg_connection = &pg_source_connection.connection;
1809
1810 let config = pg_connection
1811 .config(
1812 &storage_configuration.connection_context.secrets_reader,
1813 storage_configuration,
1814 InTask::No,
1815 )
1816 .await?;
1817
1818 let client = config
1819 .connect(
1820 "postgres_purification",
1821 &storage_configuration.connection_context.ssh_tunnel_manager,
1822 )
1823 .await?;
1824 let reference_client = SourceReferenceClient::Postgres {
1825 client: &client,
1826 publication: &pg_source_connection.publication,
1827 database: &pg_connection.database,
1828 };
1829 reference_client.get_source_references().await?
1830 }
1831 GenericSourceConnection::MySql(mysql_source_connection) => {
1832 let mysql_connection = &mysql_source_connection.connection;
1833 let config = mysql_connection
1834 .config(
1835 &storage_configuration.connection_context.secrets_reader,
1836 storage_configuration,
1837 InTask::No,
1838 )
1839 .await?;
1840
1841 let mut conn = config
1842 .connect(
1843 "mysql purification",
1844 &storage_configuration.connection_context.ssh_tunnel_manager,
1845 )
1846 .await?;
1847
1848 let reference_client = SourceReferenceClient::MySql {
1849 conn: &mut conn,
1850 include_system_schemas: false,
1851 };
1852 reference_client.get_source_references().await?
1853 }
1854 GenericSourceConnection::SqlServer(sql_server_source) => {
1855 let sql_server_connection = &sql_server_source.connection;
1857 let config = sql_server_connection
1858 .resolve_config(
1859 &storage_configuration.connection_context.secrets_reader,
1860 storage_configuration,
1861 InTask::No,
1862 )
1863 .await?;
1864 let mut client = mz_sql_server_util::Client::connect(config).await?;
1865
1866 let source_references = SourceReferenceClient::SqlServer {
1868 client: &mut client,
1869 database: sql_server_connection.database.clone().into(),
1870 }
1871 .get_source_references()
1872 .await?;
1873 source_references
1874 }
1875 GenericSourceConnection::LoadGenerator(load_gen_connection) => {
1876 let reference_client = SourceReferenceClient::LoadGenerator {
1877 generator: &load_gen_connection.load_generator,
1878 };
1879 reference_client.get_source_references().await?
1880 }
1881 GenericSourceConnection::Kafka(kafka_conn) => {
1882 let reference_client = SourceReferenceClient::Kafka {
1883 topic: &kafka_conn.topic,
1884 };
1885 reference_client.get_source_references().await?
1886 }
1887 };
1888 Ok(PurifiedStatement::PurifiedAlterSourceRefreshReferences {
1889 source_name: resolved_source_name,
1890 available_source_references: retrieved_source_references.available_source_references(),
1891 })
1892}
1893
1894#[derive(Clone, Debug, PartialEq, Eq)]
1895struct TableFromSourcePreparation {
1896 connection: GenericSourceConnection,
1897 requested_references: Option<ExternalReferences>,
1898 qualified_text_columns: Vec<UnresolvedItemName>,
1899 qualified_exclude_columns: Vec<UnresolvedItemName>,
1900 filter_constraints: FilterConstraints,
1901 exclude_all_not_null_constraints: bool,
1902 unresolved_source_name: UnresolvedItemName,
1903}
1904
1905async fn execute_table_from_source_preparation(
1906 request: TableFromSourcePreparation,
1907 storage_configuration: &StorageConfiguration,
1908) -> Result<(PurifiedSourceExport, SourceFormatOptions), PlanError> {
1909 let TableFromSourcePreparation {
1910 connection,
1911 requested_references,
1912 qualified_text_columns,
1913 qualified_exclude_columns,
1914 filter_constraints,
1915 exclude_all_not_null_constraints,
1916 unresolved_source_name,
1917 } = request;
1918 let mut format_options = SourceFormatOptions::Default;
1920
1921 let retrieved_source_references: RetrievedSourceReferences;
1922
1923 if (filter_constraints.excludes_any() || exclude_all_not_null_constraints)
1924 && !matches!(
1925 connection,
1926 GenericSourceConnection::Postgres(_) | GenericSourceConnection::MySql(_)
1927 )
1928 {
1929 sql_bail!(
1930 "EXCLUDE CONSTRAINTS is not supported for {} sources",
1931 connection.name()
1932 );
1933 }
1934
1935 let purified_export = match connection {
1938 GenericSourceConnection::Postgres(pg_source_connection) => {
1939 let pg_connection = &pg_source_connection.connection;
1941
1942 let client = pg_connection
1943 .validate(pg_source_connection.connection_id, storage_configuration)
1944 .await?;
1945
1946 let reference_client = SourceReferenceClient::Postgres {
1947 client: &client,
1948 publication: &pg_source_connection.publication,
1949 database: &pg_connection.database,
1950 };
1951 retrieved_source_references = reference_client.get_source_references().await?;
1952
1953 let initial_lsn = MzOffset::from(
1955 mz_postgres_util::fetch_max_lsn(
1956 &client,
1957 pg_source_connection
1958 .publication_details
1959 .get_is_physical_replica(),
1960 )
1961 .await?,
1962 );
1963
1964 let postgres::PurifiedSourceExports {
1965 source_exports,
1966 normalized_text_columns: _,
1970 } = postgres::purify_source_exports(
1971 &client,
1972 &retrieved_source_references,
1973 &requested_references,
1974 qualified_text_columns,
1975 qualified_exclude_columns,
1976 &filter_constraints,
1977 exclude_all_not_null_constraints,
1978 &unresolved_source_name,
1979 &SourceReferencePolicy::Required,
1980 initial_lsn,
1981 )
1982 .await?;
1983 let (_, purified_export) = source_exports.into_element();
1985 purified_export
1986 }
1987 GenericSourceConnection::MySql(mysql_source_connection) => {
1988 let mysql_connection = &mysql_source_connection.connection;
1989 let config = mysql_connection
1990 .config(
1991 &storage_configuration.connection_context.secrets_reader,
1992 storage_configuration,
1993 InTask::No,
1994 )
1995 .await?;
1996
1997 let mut conn = config
1998 .connect(
1999 "mysql purification",
2000 &storage_configuration.connection_context.ssh_tunnel_manager,
2001 )
2002 .await?;
2003
2004 ensure_binlog_full_metadata(&mut conn).await?;
2005 let binlog_full_metadata = true;
2006
2007 let initial_gtid_set =
2010 mz_mysql_util::query_sys_var(&mut conn, "global.gtid_executed").await?;
2011
2012 let reference_client = SourceReferenceClient::MySql {
2013 conn: &mut conn,
2014 include_system_schemas: mysql::references_system_schemas(&requested_references),
2015 };
2016 retrieved_source_references = reference_client.get_source_references().await?;
2017
2018 let mysql::PurifiedSourceExports {
2019 source_exports,
2020 normalized_text_columns: _,
2024 normalized_exclude_columns: _,
2025 } = mysql::purify_source_exports(
2026 &mut conn,
2027 &retrieved_source_references,
2028 &requested_references,
2029 qualified_text_columns,
2030 qualified_exclude_columns,
2031 &filter_constraints,
2032 exclude_all_not_null_constraints,
2033 &unresolved_source_name,
2034 initial_gtid_set,
2035 &SourceReferencePolicy::Required,
2036 binlog_full_metadata,
2037 )
2038 .await?;
2039 let (_, purified_export) = source_exports.into_element();
2041 purified_export
2042 }
2043 GenericSourceConnection::SqlServer(sql_server_source) => {
2044 let connection = sql_server_source.connection;
2045 let config = connection
2046 .resolve_config(
2047 &storage_configuration.connection_context.secrets_reader,
2048 storage_configuration,
2049 InTask::No,
2050 )
2051 .await?;
2052 let mut client = mz_sql_server_util::Client::connect(config).await?;
2053
2054 let database: Arc<str> = connection.database.into();
2055 let reference_client = SourceReferenceClient::SqlServer {
2056 client: &mut client,
2057 database: Arc::clone(&database),
2058 };
2059 retrieved_source_references = reference_client.get_source_references().await?;
2060 tracing::debug!(?retrieved_source_references, "got source references");
2061
2062 let timeout = mz_storage_types::sources::sql_server::MAX_LSN_WAIT
2063 .get(storage_configuration.config_set());
2064
2065 let purified_source_exports = sql_server::purify_source_exports(
2066 &*database,
2067 &mut client,
2068 &retrieved_source_references,
2069 &requested_references,
2070 &qualified_text_columns,
2071 &qualified_exclude_columns,
2072 &unresolved_source_name,
2073 timeout,
2074 &SourceReferencePolicy::Required,
2075 )
2076 .await?;
2077
2078 let (_, purified_export) = purified_source_exports.source_exports.into_element();
2080 purified_export
2081 }
2082 GenericSourceConnection::LoadGenerator(load_gen_connection) => {
2083 let reference_client = SourceReferenceClient::LoadGenerator {
2084 generator: &load_gen_connection.load_generator,
2085 };
2086 retrieved_source_references = reference_client.get_source_references().await?;
2087
2088 let requested_exports = retrieved_source_references
2089 .requested_source_exports(requested_references.as_ref(), &unresolved_source_name)?;
2090 let export = requested_exports.into_element();
2092 PurifiedSourceExport {
2093 external_reference: export.external_reference,
2094 details: PurifiedExportDetails::LoadGenerator {
2095 table: export
2096 .meta
2097 .load_generator_desc()
2098 .ok_or_else(|| internal_err!("expected load generator source reference"))?
2099 .clone(),
2100 output: export
2101 .meta
2102 .load_generator_output()
2103 .ok_or_else(|| internal_err!("expected load generator source reference"))?
2104 .clone(),
2105 },
2106 }
2107 }
2108 GenericSourceConnection::Kafka(kafka_conn) => {
2109 let reference_client = SourceReferenceClient::Kafka {
2110 topic: &kafka_conn.topic,
2111 };
2112 retrieved_source_references = reference_client.get_source_references().await?;
2113 let requested_exports = retrieved_source_references
2114 .requested_source_exports(requested_references.as_ref(), &unresolved_source_name)?;
2115 let export = requested_exports.into_element();
2117
2118 format_options = SourceFormatOptions::Kafka {
2119 topic: kafka_conn.topic.clone(),
2120 };
2121 PurifiedSourceExport {
2122 external_reference: export.external_reference,
2123 details: PurifiedExportDetails::Kafka {},
2124 }
2125 }
2126 };
2127
2128 Ok((purified_export, format_options))
2129}
2130
2131fn prepare_table_from_source(
2132 catalog: &dyn SessionCatalog,
2133 stmt: &CreateTableFromSourceStatement<Aug>,
2134) -> Result<TableFromSourcePreparation, PlanError> {
2135 let scx = StatementContext::new(None, catalog);
2136 let CreateTableFromSourceStatement {
2137 name: _,
2138 columns,
2139 constraints,
2140 source: source_name,
2141 if_not_exists: _,
2142 external_reference,
2143 format: _,
2144 envelope: _,
2145 include_metadata: _,
2146 with_options,
2147 } = stmt;
2148
2149 if matches!(columns, TableFromSourceColumns::Defined(_)) {
2151 sql_bail!("CREATE TABLE .. FROM SOURCE column definitions cannot be specified directly");
2152 }
2153 if !constraints.is_empty() {
2154 sql_bail!(
2155 "CREATE TABLE .. FROM SOURCE constraint definitions cannot be specified directly"
2156 );
2157 }
2158
2159 let item = match scx.get_item_by_resolved_name(source_name) {
2161 Ok(item) => item,
2162 Err(e) => return Err(e),
2163 };
2164
2165 let desc = match item.source_desc()? {
2167 Some(desc) => desc.clone().into_inline_connection(scx.catalog),
2168 None => {
2169 sql_bail!("cannot ALTER this type of source")
2170 }
2171 };
2172 let unresolved_source_name: UnresolvedItemName = source_name.full_item_name().clone().into();
2173
2174 let crate::plan::statement::ddl::TableFromSourceOptionExtracted {
2175 text_columns,
2176 exclude_columns,
2177 exclude_constraints,
2178 exclude_all_constraints,
2179 exclude_all_not_null_constraints,
2180 retain_history: _,
2181 details,
2182 partition_by: _,
2183 seen: _,
2184 } = with_options.clone().try_into()?;
2185 if details.is_some() {
2186 sql_bail!("DETAILS option cannot be explicitly set");
2187 }
2188
2189 let filter_constraints =
2190 FilterConstraints::from_options(exclude_constraints, exclude_all_constraints)?;
2191 if filter_constraints.excludes_any() || exclude_all_not_null_constraints {
2192 scx.require_feature_flag(&crate::session::vars::ENABLE_EXCLUDE_CONSTRAINTS_OPTION)?;
2193 }
2194
2195 let qualified_text_columns = text_columns
2199 .iter()
2200 .map(|col| {
2201 UnresolvedItemName(
2202 external_reference
2203 .as_ref()
2204 .map(|er| er.0.iter().chain_one(col).map(|i| i.clone()).collect())
2205 .unwrap_or_else(|| vec![col.clone()]),
2206 )
2207 })
2208 .collect_vec();
2209 let qualified_exclude_columns = exclude_columns
2210 .iter()
2211 .map(|col| {
2212 UnresolvedItemName(
2213 external_reference
2214 .as_ref()
2215 .map(|er| er.0.iter().chain_one(col).map(|i| i.clone()).collect())
2216 .unwrap_or_else(|| vec![col.clone()]),
2217 )
2218 })
2219 .collect_vec();
2220
2221 let requested_references = external_reference.as_ref().map(|ref_name| {
2222 ExternalReferences::SubsetTables(vec![ExternalReferenceExport {
2223 reference: ref_name.clone(),
2224 alias: None,
2225 }])
2226 });
2227
2228 Ok(TableFromSourcePreparation {
2229 connection: desc.connection,
2230 requested_references,
2231 qualified_text_columns,
2232 qualified_exclude_columns,
2233 filter_constraints,
2234 exclude_all_not_null_constraints,
2235 unresolved_source_name,
2236 })
2237}
2238
2239fn finish_table_from_source(
2240 catalog: &dyn SessionCatalog,
2241 mut stmt: CreateTableFromSourceStatement<Aug>,
2242 purified_export: PurifiedSourceExport,
2243) -> Result<PurifiedStatement, PlanError> {
2244 let scx = StatementContext::new(None, catalog);
2245 let CreateTableFromSourceStatement {
2246 columns,
2247 constraints,
2248 external_reference,
2249 with_options,
2250 ..
2251 } = &mut stmt;
2252 *external_reference = Some(purified_export.external_reference.clone());
2255
2256 match &purified_export.details {
2258 PurifiedExportDetails::Postgres { .. } => {
2259 let mut unsupported_cols = vec![];
2260 let postgres::PostgresExportStatementValues {
2261 columns: gen_columns,
2262 constraints: gen_constraints,
2263 text_columns: gen_text_columns,
2264 exclude_columns: gen_exclude_columns,
2265 details: gen_details,
2266 external_reference: _,
2267 } = postgres::generate_source_export_statement_values(
2268 &scx,
2269 purified_export,
2270 &mut unsupported_cols,
2271 )?;
2272 if !unsupported_cols.is_empty() {
2273 unsupported_cols.sort();
2274 Err(PgSourcePurificationError::UnrecognizedTypes {
2275 cols: unsupported_cols,
2276 })?;
2277 }
2278
2279 if let Some(text_cols_option) = with_options
2280 .iter_mut()
2281 .find(|option| option.name == TableFromSourceOptionName::TextColumns)
2282 {
2283 if let Some(gen_text_columns) = gen_text_columns {
2284 text_cols_option.value = Some(WithOptionValue::Sequence(gen_text_columns));
2285 }
2286 }
2287 if let Some(exclude_cols_option) = with_options
2288 .iter_mut()
2289 .find(|option| option.name == TableFromSourceOptionName::ExcludeColumns)
2290 {
2291 if let Some(gen_exclude_columns) = gen_exclude_columns {
2292 exclude_cols_option.value =
2293 Some(WithOptionValue::Sequence(gen_exclude_columns));
2294 }
2295 }
2296 match columns {
2297 TableFromSourceColumns::Defined(_) => {
2298 bail_internal!(
2301 "column definitions cannot be explicitly set for this source type"
2302 )
2303 }
2304 TableFromSourceColumns::NotSpecified => {
2305 *columns = TableFromSourceColumns::Defined(gen_columns);
2306 *constraints = gen_constraints;
2307 }
2308 TableFromSourceColumns::Named(_) => {
2309 sql_bail!("columns cannot be named for Postgres sources")
2310 }
2311 }
2312 with_options.push(TableFromSourceOption {
2313 name: TableFromSourceOptionName::Details,
2314 value: Some(WithOptionValue::Value(Value::String(hex::encode(
2315 gen_details.into_proto().encode_to_vec(),
2316 )))),
2317 })
2318 }
2319 PurifiedExportDetails::MySql { .. } => {
2320 let mysql::MySqlExportStatementValues {
2321 columns: gen_columns,
2322 constraints: gen_constraints,
2323 text_columns: gen_text_columns,
2324 exclude_columns: gen_exclude_columns,
2325 details: gen_details,
2326 external_reference: _,
2327 } = mysql::generate_source_export_statement_values(&scx, purified_export)?;
2328
2329 if let Some(text_cols_option) = with_options
2330 .iter_mut()
2331 .find(|option| option.name == TableFromSourceOptionName::TextColumns)
2332 {
2333 if let Some(gen_text_columns) = gen_text_columns {
2334 text_cols_option.value = Some(WithOptionValue::Sequence(gen_text_columns));
2335 }
2336 }
2337 if let Some(exclude_cols_option) = with_options
2338 .iter_mut()
2339 .find(|option| option.name == TableFromSourceOptionName::ExcludeColumns)
2340 {
2341 if let Some(gen_exclude_columns) = gen_exclude_columns {
2342 exclude_cols_option.value =
2343 Some(WithOptionValue::Sequence(gen_exclude_columns));
2344 }
2345 }
2346 match columns {
2347 TableFromSourceColumns::Defined(_) => {
2348 bail_internal!(
2351 "column definitions cannot be explicitly set for this source type"
2352 )
2353 }
2354 TableFromSourceColumns::NotSpecified => {
2355 *columns = TableFromSourceColumns::Defined(gen_columns);
2356 *constraints = gen_constraints;
2357 }
2358 TableFromSourceColumns::Named(_) => {
2359 sql_bail!("columns cannot be named for MySQL sources")
2360 }
2361 }
2362 with_options.push(TableFromSourceOption {
2363 name: TableFromSourceOptionName::Details,
2364 value: Some(WithOptionValue::Value(Value::String(hex::encode(
2365 gen_details.into_proto().encode_to_vec(),
2366 )))),
2367 })
2368 }
2369 PurifiedExportDetails::SqlServer { .. } => {
2370 let sql_server::SqlServerExportStatementValues {
2371 columns: gen_columns,
2372 constraints: gen_constraints,
2373 text_columns: gen_text_columns,
2374 excl_columns: gen_excl_columns,
2375 details: gen_details,
2376 external_reference: _,
2377 } = sql_server::generate_source_export_statement_values(&scx, purified_export)?;
2378
2379 if let Some(text_cols_option) = with_options
2380 .iter_mut()
2381 .find(|opt| opt.name == TableFromSourceOptionName::TextColumns)
2382 {
2383 if let Some(gen_text_columns) = gen_text_columns {
2384 text_cols_option.value = Some(WithOptionValue::Sequence(gen_text_columns));
2385 }
2386 }
2387 if let Some(exclude_cols_option) = with_options
2388 .iter_mut()
2389 .find(|opt| opt.name == TableFromSourceOptionName::ExcludeColumns)
2390 {
2391 if let Some(gen_excl_columns) = gen_excl_columns {
2392 exclude_cols_option.value = Some(WithOptionValue::Sequence(gen_excl_columns));
2393 }
2394 }
2395
2396 match columns {
2397 TableFromSourceColumns::NotSpecified => {
2398 *columns = TableFromSourceColumns::Defined(gen_columns);
2399 *constraints = gen_constraints;
2400 }
2401 TableFromSourceColumns::Named(_) => {
2402 sql_bail!("columns cannot be named for SQL Server sources")
2403 }
2404 TableFromSourceColumns::Defined(_) => {
2405 bail_internal!(
2408 "column definitions cannot be explicitly set for this source type"
2409 )
2410 }
2411 }
2412
2413 with_options.push(TableFromSourceOption {
2414 name: TableFromSourceOptionName::Details,
2415 value: Some(WithOptionValue::Value(Value::String(hex::encode(
2416 gen_details.into_proto().encode_to_vec(),
2417 )))),
2418 })
2419 }
2420 PurifiedExportDetails::LoadGenerator { .. } => {
2421 let (desc, output) = match purified_export.details {
2422 PurifiedExportDetails::LoadGenerator { table, output } => (table, output),
2423 _ => bail_internal!("purified export details must be load generator"),
2424 };
2425 if let Some(desc) = desc {
2430 let (gen_columns, gen_constraints) = scx.relation_desc_into_table_defs(&desc)?;
2431 match columns {
2432 TableFromSourceColumns::Defined(_) => bail_internal!(
2435 "column definitions cannot be explicitly set for this source type"
2436 ),
2437 TableFromSourceColumns::NotSpecified => {
2438 *columns = TableFromSourceColumns::Defined(gen_columns);
2439 *constraints = gen_constraints;
2440 }
2441 TableFromSourceColumns::Named(_) => {
2442 sql_bail!("columns cannot be named for multi-output load generator sources")
2443 }
2444 }
2445 }
2446 let details = SourceExportStatementDetails::LoadGenerator { output };
2447 with_options.push(TableFromSourceOption {
2448 name: TableFromSourceOptionName::Details,
2449 value: Some(WithOptionValue::Value(Value::String(hex::encode(
2450 details.into_proto().encode_to_vec(),
2451 )))),
2452 })
2453 }
2454 PurifiedExportDetails::Kafka {} => {
2455 let details = SourceExportStatementDetails::Kafka {};
2459 with_options.push(TableFromSourceOption {
2460 name: TableFromSourceOptionName::Details,
2461 value: Some(WithOptionValue::Value(Value::String(hex::encode(
2462 details.into_proto().encode_to_vec(),
2463 )))),
2464 })
2465 }
2466 };
2467
2468 Ok(PurifiedStatement::PurifiedCreateTableFromSource { stmt })
2472}
2473
2474#[derive(Clone, Debug)]
2475enum SourceFormatOptions {
2476 Default,
2477 Kafka { topic: String },
2478}
2479
2480pub fn generate_progress_subsource_statement(
2482 scx: &StatementContext,
2483 source: &CreateSourceStatement<Aug>,
2484) -> Result<CreateSubsourceStatement<Aug>, PlanError> {
2485 let CreateSourceStatement {
2486 name: source_name,
2487 connection: source_connection,
2488 progress_subsource,
2489 with_options,
2490 ..
2491 } = source;
2492
2493 let progress_desc = match source_connection {
2494 CreateSourceConnection::Kafka { .. } => {
2495 &mz_storage_types::sources::kafka::KAFKA_PROGRESS_DESC
2496 }
2497 CreateSourceConnection::Postgres { .. } => {
2498 &mz_storage_types::sources::postgres::PG_PROGRESS_DESC
2499 }
2500 CreateSourceConnection::SqlServer { .. } => {
2501 &mz_storage_types::sources::sql_server::SQL_SERVER_PROGRESS_DESC
2502 }
2503 CreateSourceConnection::MySql { .. } => {
2504 &mz_storage_types::sources::mysql::MYSQL_PROGRESS_DESC
2505 }
2506 CreateSourceConnection::LoadGenerator { .. } => {
2507 &mz_storage_types::sources::load_generator::LOAD_GEN_PROGRESS_DESC
2508 }
2509 };
2510
2511 let name = match progress_subsource {
2512 Some(name) => match name {
2513 DeferredItemName::Deferred(name) => name.clone(),
2514 DeferredItemName::Named(_) => {
2515 sql_bail!("progress subsource name cannot be a resolved name")
2516 }
2517 },
2518 None => {
2519 let (item, prefix) = source_name
2520 .0
2521 .split_last()
2522 .ok_or_else(|| sql_err!("source name must have at least one component"))?;
2523 let item_name = Ident::try_generate_name(item.to_string(), "_progress", |candidate| {
2524 let mut suggested_name = prefix.to_vec();
2525 suggested_name.push(candidate.clone());
2526
2527 let partial = normalize::unresolved_item_name(UnresolvedItemName(suggested_name))?;
2528 let qualified = scx.allocate_qualified_name(partial)?;
2529 let item_exists = scx.catalog.get_item_by_name(&qualified).is_some();
2530 let type_exists = scx.catalog.get_type_by_name(&qualified).is_some();
2531 Ok::<_, PlanError>(!item_exists && !type_exists)
2532 })?;
2533
2534 let mut full_name = prefix.to_vec();
2535 full_name.push(item_name);
2536 let full_name = normalize::unresolved_item_name(UnresolvedItemName(full_name))?;
2537 let qualified_name = scx.allocate_qualified_name(full_name)?;
2538 let full_name = scx.catalog.resolve_full_name(&qualified_name);
2539
2540 UnresolvedItemName::from(full_name.clone())
2541 }
2542 };
2543
2544 let (columns, constraints) = scx.relation_desc_into_table_defs(progress_desc)?;
2545
2546 let mut progress_with_options: Vec<_> = with_options
2547 .iter()
2548 .filter_map(|opt| match opt.name {
2549 CreateSourceOptionName::TimestampInterval => None,
2550 CreateSourceOptionName::RetainHistory => Some(CreateSubsourceOption {
2551 name: CreateSubsourceOptionName::RetainHistory,
2552 value: opt.value.clone(),
2553 }),
2554 })
2555 .collect();
2556 progress_with_options.push(CreateSubsourceOption {
2557 name: CreateSubsourceOptionName::Progress,
2558 value: Some(WithOptionValue::Value(Value::Boolean(true))),
2559 });
2560
2561 Ok(CreateSubsourceStatement {
2562 name,
2563 columns,
2564 of_source: None,
2568 constraints,
2569 if_not_exists: false,
2570 with_options: progress_with_options,
2571 })
2572}
2573
2574pub fn generate_subsource_statements(
2575 scx: &StatementContext,
2576 source_name: ResolvedItemName,
2577 subsources: BTreeMap<UnresolvedItemName, PurifiedSourceExport>,
2578) -> Result<Vec<CreateSubsourceStatement<Aug>>, PlanError> {
2579 if subsources.is_empty() {
2581 return Ok(vec![]);
2582 }
2583 let (_, purified_export) = subsources
2584 .iter()
2585 .next()
2586 .ok_or_else(|| internal_err!("expected at least one subsource"))?;
2587
2588 let statements = match &purified_export.details {
2589 PurifiedExportDetails::Postgres { .. } => {
2590 crate::pure::postgres::generate_create_subsource_statements(
2591 scx,
2592 source_name,
2593 subsources,
2594 )?
2595 }
2596 PurifiedExportDetails::MySql { .. } => {
2597 crate::pure::mysql::generate_create_subsource_statements(scx, source_name, subsources)?
2598 }
2599 PurifiedExportDetails::SqlServer { .. } => {
2600 crate::pure::sql_server::generate_create_subsource_statements(
2601 scx,
2602 source_name,
2603 subsources,
2604 )?
2605 }
2606 PurifiedExportDetails::LoadGenerator { .. } => {
2607 let mut subsource_stmts = Vec::with_capacity(subsources.len());
2608 for (subsource_name, purified_export) in subsources {
2609 let (desc, output) = match purified_export.details {
2610 PurifiedExportDetails::LoadGenerator { table, output } => (table, output),
2611 _ => {
2612 bail_internal!("purified export details must be load generator")
2613 }
2614 };
2615 let desc = desc.ok_or_else(|| {
2616 internal_err!(
2617 "subsources cannot be generated for single-output load generators"
2618 )
2619 })?;
2620
2621 let (columns, table_constraints) = scx.relation_desc_into_table_defs(&desc)?;
2622 let details = SourceExportStatementDetails::LoadGenerator { output };
2623 let subsource = CreateSubsourceStatement {
2625 name: subsource_name,
2626 columns,
2627 of_source: Some(source_name.clone()),
2628 constraints: table_constraints,
2633 if_not_exists: false,
2634 with_options: vec![
2635 CreateSubsourceOption {
2636 name: CreateSubsourceOptionName::ExternalReference,
2637 value: Some(WithOptionValue::UnresolvedItemName(
2638 purified_export.external_reference,
2639 )),
2640 },
2641 CreateSubsourceOption {
2642 name: CreateSubsourceOptionName::Details,
2643 value: Some(WithOptionValue::Value(Value::String(hex::encode(
2644 details.into_proto().encode_to_vec(),
2645 )))),
2646 },
2647 ],
2648 };
2649 subsource_stmts.push(subsource);
2650 }
2651
2652 subsource_stmts
2653 }
2654 PurifiedExportDetails::Kafka { .. } => {
2655 if !subsources.is_empty() {
2659 bail_internal!("Kafka sources do not produce data-bearing subsources");
2660 }
2661 vec![]
2662 }
2663 };
2664 Ok(statements)
2665}
2666
2667#[derive(Clone, Debug, PartialEq, Eq)]
2668struct CsrProtobufPreparation {
2669 connection: mz_storage_types::connections::CsrConnection,
2670 topic: String,
2671 require_key: bool,
2672}
2673
2674impl CsrProtobufPreparation {
2675 async fn execute(
2676 self,
2677 storage_configuration: &StorageConfiguration,
2678 ) -> Result<CsrSeedProtobuf, PlanError> {
2679 let Self {
2680 connection: ccsr_connection,
2681 topic,
2682 require_key,
2683 } = self;
2684 let ccsr_client = ccsr_connection
2685 .connect(storage_configuration, InTask::No)
2686 .await
2687 .map_err(|e| CsrPurificationError::ClientError(Arc::new(e)))?;
2688
2689 let value = compile_proto(&format!("{}-value", topic), &ccsr_client).await?;
2690 let key = compile_proto(&format!("{}-key", topic), &ccsr_client)
2691 .await
2692 .ok();
2693
2694 if require_key && key.is_none() {
2695 sql_bail!("Key schema is required for ENVELOPE DEBEZIUM");
2696 }
2697
2698 Ok(CsrSeedProtobuf { value, key })
2699 }
2700}
2701
2702#[derive(Clone, Debug, PartialEq, Eq)]
2703struct CsrAvroPreparation {
2704 connection: mz_storage_types::connections::CsrConnection,
2705 topic: String,
2706 key_strategy: ReaderSchemaSelectionStrategy,
2707 value_strategy: ReaderSchemaSelectionStrategy,
2708 require_key: bool,
2709}
2710
2711impl CsrAvroPreparation {
2712 async fn execute(
2713 self,
2714 storage_configuration: &StorageConfiguration,
2715 ) -> Result<CsrSeedAvro, PlanError> {
2716 let Self {
2717 connection: csr_connection,
2718 topic,
2719 key_strategy,
2720 value_strategy,
2721 require_key,
2722 } = self;
2723 let ccsr_client = csr_connection
2724 .connect(storage_configuration, InTask::No)
2725 .await
2726 .map_err(|e| CsrPurificationError::ClientError(Arc::new(e)))?;
2727
2728 let Schema {
2729 key_schema,
2730 value_schema,
2731 key_reference_schemas,
2732 value_reference_schemas,
2733 } = get_remote_csr_schema(&ccsr_client, key_strategy, value_strategy, &topic).await?;
2734 if require_key && key_schema.is_none() {
2735 sql_bail!("Key schema is required for ENVELOPE DEBEZIUM");
2736 }
2737
2738 Ok(CsrSeedAvro {
2739 key_schema,
2740 value_schema,
2741 key_reference_schemas,
2742 value_reference_schemas,
2743 })
2744 }
2745}
2746
2747#[derive(Clone, Debug, PartialEq, Eq)]
2748struct GlueAvroPreparation {
2749 connection: mz_storage_types::connections::GlueSchemaRegistryConnection,
2750 schema_name: String,
2751}
2752
2753impl GlueAvroPreparation {
2754 async fn execute(
2755 self,
2756 storage_configuration: &StorageConfiguration,
2757 ) -> Result<GlueAvroSeed, PlanError> {
2758 use crate::pure::error::GluePurificationError;
2759 let Self {
2760 connection: gsr_connection,
2761 schema_name,
2762 } = self;
2763 let enforce_external_addresses = mz_storage_types::dyncfgs::ENFORCE_EXTERNAL_ADDRESSES
2765 .get(storage_configuration.config_set());
2766 let sdk_config = gsr_connection
2767 .aws_connection
2768 .connection
2769 .load_sdk_config(
2770 &storage_configuration.connection_context,
2771 gsr_connection.aws_connection.connection_id,
2772 InTask::No,
2774 enforce_external_addresses,
2775 )
2776 .await
2777 .map_err(|e| GluePurificationError::LoadSdkConfigError(Arc::new(e)))?;
2778 let glue_client = mz_aws_glue_schema_registry::ClientConfig::new(sdk_config).build();
2779
2780 let version = glue_client
2781 .get_schema_version_latest_by_name(&gsr_connection.registry_name, &schema_name)
2782 .await
2783 .map_err(|e| GluePurificationError::SchemaLookupError {
2784 registry: gsr_connection.registry_name.clone(),
2785 schema: schema_name.clone(),
2786 cause: Arc::new(e),
2787 })?;
2788 match &version.data_format {
2792 Some(mz_aws_glue_schema_registry::DataFormat::Avro) => {}
2793 other => {
2794 return Err(GluePurificationError::UnsupportedDataFormat {
2795 registry: gsr_connection.registry_name.clone(),
2796 schema: schema_name.clone(),
2797 format: other
2798 .as_ref()
2799 .map(|f| f.as_str().to_string())
2800 .unwrap_or_else(|| "<unspecified>".to_string()),
2801 }
2802 .into());
2803 }
2804 }
2805 let value_schema =
2806 version
2807 .definition
2808 .ok_or_else(|| GluePurificationError::EmptyDefinition {
2809 registry: gsr_connection.registry_name.clone(),
2810 schema: schema_name.clone(),
2811 })?;
2812
2813 Ok(GlueAvroSeed { value_schema })
2814 }
2815}
2816
2817fn prepare_csr_connection_proto(
2818 catalog: &dyn SessionCatalog,
2819 options: &SourceFormatOptions,
2820 csr_connection: &CsrConnectionProtobuf<Aug>,
2821 envelope: &Option<SourceEnvelope>,
2822) -> Result<Option<CsrProtobufPreparation>, PlanError> {
2823 let SourceFormatOptions::Kafka { topic } = options else {
2824 sql_bail!("Confluent Schema Registry is only supported with Kafka sources")
2825 };
2826
2827 let CsrConnectionProtobuf {
2828 seed,
2829 connection: CsrConnection {
2830 connection,
2831 options: _,
2832 },
2833 } = csr_connection;
2834 match seed {
2835 None => {
2836 let scx = StatementContext::new(None, &*catalog);
2837
2838 let ccsr_connection = match scx.get_item_by_resolved_name(connection)?.connection()? {
2839 Connection::Csr(connection) => connection.clone().into_inline_connection(catalog),
2840 _ => sql_bail!("{} is not a schema registry connection", connection),
2841 };
2842
2843 return Ok(Some(CsrProtobufPreparation {
2844 connection: ccsr_connection,
2845 topic: topic.clone(),
2846 require_key: matches!(envelope, Some(SourceEnvelope::Debezium)),
2847 }));
2848 }
2849 Some(_) => (),
2850 }
2851
2852 Ok(None)
2853}
2854
2855fn prepare_csr_connection_avro(
2856 catalog: &dyn SessionCatalog,
2857 options: &SourceFormatOptions,
2858 csr_connection: &CsrConnectionAvro<Aug>,
2859 envelope: &Option<SourceEnvelope>,
2860) -> Result<Option<CsrAvroPreparation>, PlanError> {
2861 let SourceFormatOptions::Kafka { topic } = options else {
2862 sql_bail!("Confluent Schema Registry is only supported with Kafka sources")
2863 };
2864
2865 let CsrConnectionAvro {
2866 connection: CsrConnection { connection, .. },
2867 seed,
2868 key_strategy,
2869 value_strategy,
2870 } = csr_connection;
2871 if seed.is_none() {
2872 let scx = StatementContext::new(None, &*catalog);
2873 let csr_connection = match scx.get_item_by_resolved_name(connection)?.connection()? {
2874 Connection::Csr(connection) => connection.clone().into_inline_connection(catalog),
2875 _ => sql_bail!("{} is not a schema registry connection", connection),
2876 };
2877 return Ok(Some(CsrAvroPreparation {
2878 connection: csr_connection,
2879 topic: topic.clone(),
2880 key_strategy: key_strategy.clone().unwrap_or_default(),
2881 value_strategy: value_strategy.clone().unwrap_or_default(),
2882 require_key: matches!(envelope, Some(SourceEnvelope::Debezium)),
2883 }));
2884 }
2885
2886 Ok(None)
2887}
2888
2889fn prepare_glue_connection_avro(
2890 catalog: &dyn SessionCatalog,
2891 options: &SourceFormatOptions,
2892 connection: &ResolvedItemName,
2893 with_options: &[GlueAvroOption<Aug>],
2894 seed: &Option<GlueAvroSeed>,
2895) -> Result<Option<GlueAvroPreparation>, PlanError> {
2896 use crate::pure::error::GluePurificationError;
2897 let SourceFormatOptions::Kafka { .. } = options else {
2898 sql_bail!("AWS Glue Schema Registry is only supported with Kafka sources")
2899 };
2900
2901 let scx = StatementContext::new(None, &*catalog);
2902 let item = scx.get_item_by_resolved_name(connection)?;
2903 let full_name = scx.catalog.resolve_full_name(item.name());
2904 let Connection::GlueSchemaRegistry(gsr_connection) = item.connection()? else {
2911 return Err(GluePurificationError::NotGlueConnection(full_name).into());
2912 };
2913
2914 let crate::plan::statement::ddl::GlueAvroOptionExtracted {
2918 schema_name,
2919 key_schema_name,
2920 value_schema_name,
2921 key_compatibility_level,
2922 value_compatibility_level,
2923 seen: _,
2924 } = with_options.to_vec().try_into()?;
2925 if key_schema_name.is_some()
2928 || value_schema_name.is_some()
2929 || key_compatibility_level.is_some()
2930 || value_compatibility_level.is_some()
2931 {
2932 sql_bail!(
2933 "KEY SCHEMA NAME, VALUE SCHEMA NAME, KEY COMPATIBILITY LEVEL, and VALUE \
2934 COMPATIBILITY LEVEL are not supported for AWS Glue Schema Registry sources, \
2935 use SCHEMA NAME instead"
2936 );
2937 }
2938 let schema_name = schema_name.ok_or(GluePurificationError::MissingSchemaName)?;
2939
2940 if seed.is_some() {
2941 return Ok(None);
2946 }
2947 let gsr_connection = gsr_connection.into_inline_connection(catalog);
2948
2949 Ok(Some(GlueAvroPreparation {
2950 connection: gsr_connection,
2951 schema_name,
2952 }))
2953}
2954
2955#[derive(Debug)]
2956pub struct Schema {
2957 pub key_schema: Option<String>,
2958 pub value_schema: String,
2959 pub key_reference_schemas: Vec<String>,
2961 pub value_reference_schemas: Vec<String>,
2963}
2964
2965struct SchemaWithReferences {
2967 schema: String,
2969 references: Vec<String>,
2971}
2972
2973async fn get_schema_with_strategy(
2974 client: &Client,
2975 strategy: ReaderSchemaSelectionStrategy,
2976 subject: &str,
2977) -> Result<Option<SchemaWithReferences>, PlanError> {
2978 match strategy {
2979 ReaderSchemaSelectionStrategy::Latest => {
2980 match client.get_subject_and_references(subject).await {
2982 Ok((primary, dependencies)) => Ok(Some(SchemaWithReferences {
2983 schema: primary.schema.raw,
2984 references: dependencies.into_iter().map(|s| s.schema.raw).collect(),
2985 })),
2986 Err(GetBySubjectError::SubjectNotFound)
2987 | Err(GetBySubjectError::VersionNotFound(_)) => Ok(None),
2988 Err(e) => Err(PlanError::FetchingCsrSchemaFailed {
2989 schema_lookup: format!("subject {}", subject.quoted()),
2990 cause: Arc::new(e),
2991 }),
2992 }
2993 }
2994 ReaderSchemaSelectionStrategy::Inline(raw) => Ok(Some(SchemaWithReferences {
2998 schema: raw,
2999 references: vec![],
3000 })),
3001 ReaderSchemaSelectionStrategy::ById(id) => {
3002 match client.get_subject_and_references_by_id(id).await {
3003 Ok((primary, dependencies)) => Ok(Some(SchemaWithReferences {
3004 schema: primary.schema.raw,
3005 references: dependencies.into_iter().map(|s| s.schema.raw).collect(),
3006 })),
3007 Err(GetBySubjectError::SubjectNotFound)
3008 | Err(GetBySubjectError::VersionNotFound(_)) => Ok(None),
3009 Err(e) => Err(PlanError::FetchingCsrSchemaFailed {
3010 schema_lookup: format!("subject {}", subject.quoted()),
3011 cause: Arc::new(e),
3012 }),
3013 }
3014 }
3015 }
3016}
3017
3018async fn get_remote_csr_schema(
3019 ccsr_client: &mz_ccsr::Client,
3020 key_strategy: ReaderSchemaSelectionStrategy,
3021 value_strategy: ReaderSchemaSelectionStrategy,
3022 topic: &str,
3023) -> Result<Schema, PlanError> {
3024 let value_schema_name = format!("{}-value", topic);
3025 let value_result =
3026 get_schema_with_strategy(ccsr_client, value_strategy, &value_schema_name).await?;
3027 let value_result = value_result.ok_or_else(|| anyhow!("No value schema found"))?;
3028
3029 let key_subject = format!("{}-key", topic);
3030 let key_result = get_schema_with_strategy(ccsr_client, key_strategy, &key_subject).await?;
3031 Ok(Schema {
3032 key_schema: key_result.as_ref().map(|r| r.schema.clone()),
3033 value_schema: value_result.schema,
3034 key_reference_schemas: key_result.map(|r| r.references).unwrap_or_default(),
3035 value_reference_schemas: value_result.references,
3036 })
3037}
3038
3039async fn compile_proto(
3041 subject_name: &String,
3042 ccsr_client: &Client,
3043) -> Result<CsrSeedProtobufSchema, PlanError> {
3044 let (primary_subject, dependency_subjects) = ccsr_client
3045 .get_subject_and_references(subject_name)
3046 .await
3047 .map_err(|e| PlanError::FetchingCsrSchemaFailed {
3048 schema_lookup: format!("subject {}", subject_name.quoted()),
3049 cause: Arc::new(e),
3050 })?;
3051
3052 let mut source_tree = VirtualSourceTree::new();
3054
3055 source_tree.as_mut().map_well_known_types();
3059
3060 for subject in iter::once(&primary_subject).chain(dependency_subjects.iter()) {
3061 source_tree.as_mut().add_file(
3062 Path::new(&subject.name),
3063 subject.schema.raw.as_bytes().to_vec(),
3064 );
3065 }
3066 let mut db = SourceTreeDescriptorDatabase::new(source_tree.as_mut());
3067 let fds = db
3068 .as_mut()
3069 .build_file_descriptor_set(&[Path::new(&primary_subject.name)])
3070 .map_err(|cause| PlanError::InvalidProtobufSchema { cause })?;
3071
3072 let primary_fd = fds.file(0);
3074 let message_name = match primary_fd.message_type_size() {
3075 1 => String::from_utf8_lossy(primary_fd.message_type(0).name()).into_owned(),
3076 0 => bail_unsupported!(29603, "Protobuf schemas with no messages"),
3077 _ => bail_unsupported!(29603, "Protobuf schemas with multiple messages"),
3078 };
3079
3080 let bytes = &fds
3082 .serialize()
3083 .map_err(|cause| PlanError::InvalidProtobufSchema { cause })?;
3084 let mut schema = String::new();
3085 strconv::format_bytes(&mut schema, bytes);
3086
3087 Ok(CsrSeedProtobufSchema {
3088 schema,
3089 message_name,
3090 })
3091}
3092
3093const MZ_NOW_NAME: &str = "mz_now";
3094const MZ_NOW_SCHEMA: &str = "mz_catalog";
3095
3096pub fn purify_create_materialized_view_options(
3102 catalog: impl SessionCatalog,
3103 mz_now: Option<Timestamp>,
3104 cmvs: &mut CreateMaterializedViewStatement<Aug>,
3105 resolved_ids: &mut ResolvedIds,
3106) {
3107 let (mz_now_id, mz_now_expr) = {
3110 let item = catalog
3111 .resolve_function(&PartialItemName {
3112 database: None,
3113 schema: Some(MZ_NOW_SCHEMA.to_string()),
3114 item: MZ_NOW_NAME.to_string(),
3115 })
3116 .expect("we should be able to resolve mz_now");
3117 (
3118 item.id(),
3119 Expr::Function(Function {
3120 name: ResolvedItemName::Item {
3121 id: item.id(),
3122 qualifiers: item.name().qualifiers.clone(),
3123 full_name: catalog.resolve_full_name(item.name()),
3124 print_id: false,
3125 version: RelationVersionSelector::Latest,
3126 },
3127 args: FunctionArgs::Args {
3128 args: Vec::new(),
3129 order_by: Vec::new(),
3130 },
3131 filter: None,
3132 over: None,
3133 distinct: false,
3134 }),
3135 )
3136 };
3137 let (mz_timestamp_id, mz_timestamp_type) = {
3139 let item = catalog.get_system_type("mz_timestamp");
3140 let full_name = catalog.resolve_full_name(item.name());
3141 (
3142 item.id(),
3143 ResolvedDataType::Named {
3144 id: item.id(),
3145 qualifiers: item.name().qualifiers.clone(),
3146 full_name,
3147 modifiers: vec![],
3148 print_id: true,
3149 },
3150 )
3151 };
3152
3153 let mut introduced_mz_timestamp = false;
3154
3155 for option in cmvs.with_options.iter_mut() {
3156 if matches!(
3158 option.value,
3159 Some(WithOptionValue::Refresh(RefreshOptionValue::AtCreation))
3160 ) {
3161 option.value = Some(WithOptionValue::Refresh(RefreshOptionValue::At(
3162 RefreshAtOptionValue {
3163 time: mz_now_expr.clone(),
3164 },
3165 )));
3166 }
3167
3168 if let Some(WithOptionValue::Refresh(RefreshOptionValue::Every(
3170 RefreshEveryOptionValue { aligned_to, .. },
3171 ))) = &mut option.value
3172 {
3173 if aligned_to.is_none() {
3174 *aligned_to = Some(mz_now_expr.clone());
3175 }
3176 }
3177
3178 match &mut option.value {
3181 Some(WithOptionValue::Refresh(RefreshOptionValue::At(RefreshAtOptionValue {
3182 time,
3183 }))) => {
3184 let mut visitor = MzNowPurifierVisitor::new(mz_now, mz_timestamp_type.clone());
3185 visitor.visit_expr_mut(time);
3186 introduced_mz_timestamp |= visitor.introduced_mz_timestamp;
3187 }
3188 Some(WithOptionValue::Refresh(RefreshOptionValue::Every(
3189 RefreshEveryOptionValue {
3190 interval: _,
3191 aligned_to: Some(aligned_to),
3192 },
3193 ))) => {
3194 let mut visitor = MzNowPurifierVisitor::new(mz_now, mz_timestamp_type.clone());
3195 visitor.visit_expr_mut(aligned_to);
3196 introduced_mz_timestamp |= visitor.introduced_mz_timestamp;
3197 }
3198 _ => {}
3199 }
3200 }
3201
3202 if !cmvs.with_options.iter().any(|o| {
3204 matches!(
3205 o,
3206 MaterializedViewOption {
3207 value: Some(WithOptionValue::Refresh(..)),
3208 ..
3209 }
3210 )
3211 }) {
3212 cmvs.with_options.push(MaterializedViewOption {
3213 name: MaterializedViewOptionName::Refresh,
3214 value: Some(WithOptionValue::Refresh(RefreshOptionValue::OnCommit)),
3215 })
3216 }
3217
3218 if introduced_mz_timestamp {
3222 resolved_ids.add_item(mz_timestamp_id);
3223 }
3224 let mut visitor = ExprContainsTemporalVisitor::new();
3228 visitor.visit_create_materialized_view_statement(cmvs);
3229 if !visitor.contains_temporal {
3230 resolved_ids.remove_item(&mz_now_id);
3231 }
3232}
3233
3234pub fn materialized_view_option_contains_temporal(mvo: &MaterializedViewOption<Aug>) -> bool {
3237 match &mvo.value {
3238 Some(WithOptionValue::Refresh(RefreshOptionValue::At(RefreshAtOptionValue { time }))) => {
3239 let mut visitor = ExprContainsTemporalVisitor::new();
3240 visitor.visit_expr(time);
3241 visitor.contains_temporal
3242 }
3243 Some(WithOptionValue::Refresh(RefreshOptionValue::Every(RefreshEveryOptionValue {
3244 interval: _,
3245 aligned_to: Some(aligned_to),
3246 }))) => {
3247 let mut visitor = ExprContainsTemporalVisitor::new();
3248 visitor.visit_expr(aligned_to);
3249 visitor.contains_temporal
3250 }
3251 Some(WithOptionValue::Refresh(RefreshOptionValue::Every(RefreshEveryOptionValue {
3252 interval: _,
3253 aligned_to: None,
3254 }))) => {
3255 true
3258 }
3259 Some(WithOptionValue::Refresh(RefreshOptionValue::AtCreation)) => {
3260 true
3262 }
3263 _ => false,
3264 }
3265}
3266
3267struct ExprContainsTemporalVisitor {
3269 pub contains_temporal: bool,
3270}
3271
3272impl ExprContainsTemporalVisitor {
3273 pub fn new() -> ExprContainsTemporalVisitor {
3274 ExprContainsTemporalVisitor {
3275 contains_temporal: false,
3276 }
3277 }
3278}
3279
3280impl Visit<'_, Aug> for ExprContainsTemporalVisitor {
3281 fn visit_function(&mut self, func: &Function<Aug>) {
3282 self.contains_temporal |= func.name.full_item_name().item == MZ_NOW_NAME;
3283 visit_function(self, func);
3284 }
3285}
3286
3287struct MzNowPurifierVisitor {
3288 pub mz_now: Option<Timestamp>,
3289 pub mz_timestamp_type: ResolvedDataType,
3290 pub introduced_mz_timestamp: bool,
3291}
3292
3293impl MzNowPurifierVisitor {
3294 pub fn new(
3295 mz_now: Option<Timestamp>,
3296 mz_timestamp_type: ResolvedDataType,
3297 ) -> MzNowPurifierVisitor {
3298 MzNowPurifierVisitor {
3299 mz_now,
3300 mz_timestamp_type,
3301 introduced_mz_timestamp: false,
3302 }
3303 }
3304}
3305
3306impl VisitMut<'_, Aug> for MzNowPurifierVisitor {
3307 fn visit_expr_mut(&mut self, expr: &'_ mut Expr<Aug>) {
3308 match expr {
3309 Expr::Function(Function {
3310 name:
3311 ResolvedItemName::Item {
3312 full_name: FullItemName { item, .. },
3313 ..
3314 },
3315 ..
3316 }) if item == &MZ_NOW_NAME.to_string() => {
3317 let mz_now = self.mz_now.expect(
3318 "we should have chosen a timestamp if the expression contains mz_now()",
3319 );
3320 *expr = Expr::Cast {
3323 expr: Box::new(Expr::Value(Value::Number(mz_now.to_string()))),
3324 data_type: self.mz_timestamp_type.clone(),
3325 failure_mode: CastFailureMode::Error,
3326 };
3327 self.introduced_mz_timestamp = true;
3328 }
3329 _ => visit_expr_mut(self, expr),
3330 }
3331 }
3332}