Skip to main content

mz_sql/
pure.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! SQL purification.
11//!
12//! See the [crate-level documentation](crate) for details.
13
14use 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
134/// Generates a subsource name by prepending source schema name if present
135///
136/// For eg. if source is `a.b`, then `a` will be prepended to the subsource name
137/// so that it's generated in the same schema as source
138fn 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
147// TODO(database-issues#8620): Remove once subsources are removed
148/// Validates the requested subsources do not have name conflicts with each other
149/// and that the same upstream table is not referenced multiple times.
150fn validate_source_export_names<T>(
151    requested_source_exports: &[RequestedSourceExport<T>],
152) -> Result<(), PlanError> {
153    // This condition would get caught during the catalog transaction, but produces a
154    // vague, non-contextual error. Instead, error here so we can suggest to the user
155    // how to fix the problem.
156    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    // We disallow subsource statements from referencing the same upstream table, but we allow
183    // `CREATE TABLE .. FROM SOURCE` statements to do so. Since `CREATE TABLE .. FROM SOURCE`
184    // purification will only provide 1 `requested_source_export`, we can leave this here without
185    // needing to differentiate between the two types of statements.
186    // TODO(roshan): Remove this when auto-generated subsources are deprecated.
187    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        // The progress subsource, if we are offloading progress info to a separate relation
217        create_progress_subsource: bool,
218        create_source_stmt: CreateSourceStatement<Aug>,
219        // Map of subsource names to external details
220        subsources: BTreeMap<UnresolvedItemName, PurifiedSourceExport>,
221        /// All the available upstream references that can be added as tables
222        /// to this primary source.
223        available_source_references: SourceReferences,
224    },
225    PurifiedAlterSource {
226        alter_source_stmt: AlterSourceStatement<Aug>,
227    },
228    PurifiedAlterSourceAddSubsources {
229        // This just saves us an annoying catalog lookup
230        source_name: ResolvedItemName,
231        /// Options that we will need the values of to update the source's
232        /// definition.
233        options: Vec<AlterSourceAddSubsourceOption<Aug>>,
234        // Map of subsource names to external details
235        subsources: BTreeMap<UnresolvedItemName, PurifiedSourceExport>,
236    },
237    PurifiedAlterSourceRefreshReferences {
238        source_name: ResolvedItemName,
239        /// The updated available upstream references for the primary source.
240        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/// The existing source a statement drives, and how the statement uses it.
284///
285/// Resolved from the statement rather than from `ResolvedIds`: `ALTER SOURCE`
286/// carries its target as an `UnresolvedItemName` that name resolution never
287/// records.
288#[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
302/// Resolves the existing source that `stmt` drives, mirroring how the
303/// purification paths below resolve it themselves.
304///
305/// `None` means the statement drives no such source: it names its connection
306/// outright (`CREATE SOURCE`, `CREATE SINK`), or the name does not resolve to a
307/// source. Purification reports the latter itself, and bails on it before
308/// opening any connection. A [`Statement`] variant that instead reaches
309/// upstream through an existing catalog item needs an arm here, or callers see
310/// no source at all.
311pub 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
332/// Purifies a statement, removing any dependencies on external state.
333///
334/// See the section on [purification](crate#purification) in the crate
335/// documentation for details.
336///
337/// Note that this doesn't handle CREATE MATERIALIZED VIEW, which is
338/// handled by [purify_create_materialized_view_options] instead.
339/// This could be made more consistent by a refactoring discussed here:
340/// <https://github.com/MaterializeInc/materialize/pull/23870#discussion_r1435922709>
341pub 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
375/// Injects `DOC ON` comments into all Avro formats that are using a schema
376/// registry by finding all SQL `COMMENT`s that are attached to the sink's
377/// underlying materialized view (as well as any types referenced by that
378/// underlying materialized view).
379pub(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    // Collect all objects referenced by the sink.
385    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    // Collect all Avro formats that use a schema registry, as well as a set of
394    // all identifiers named in user-provided `DOC ON` options.
395    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 each Avro format in the sink, inject the appropriate `DOC ON` options
413    // for each item referenced by the sink.
414    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        // Adding existing comments if not already provided by user
424        for object_id in &object_ids {
425            // Always add comments to the latest version of the item.
426            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                // Attach comment for the item itself, if the user has not
439                // already provided an overriding `DOC ON` option for the item.
440                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                // Attach comment for each column in the item, if the user has
456                // not already provided an overriding `DOC ON` option for the
457                // column.
458                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
494/// Checks that the sink described in the statement can connect to its external
495/// resources.
496async fn purify_create_sink(
497    catalog: impl SessionCatalog,
498    mut create_sink_stmt: CreateSinkStatement<Aug>,
499    storage_configuration: &StorageConfiguration,
500) -> Result<PurifiedStatement, PlanError> {
501    // General purification
502    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    // The list of options that the user is allowed to specify.
515    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            // We must not leave any state behind in the Kafka broker, so just ensure that
538            // we can connect. This means we don't ensure that we can create the topic and
539            // introduces TOCTOU errors, but creating an inoperable sink is infinitely
540            // preferable to leaking state in users' environments.
541            let scx = StatementContext::new(None, &catalog);
542            let connection = {
543                let item = scx.get_item_by_resolved_name(connection)?;
544                // Get Kafka connection
545                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                    // anyhow doesn't support Clone, so not trivial to move into PlanError
572                    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                // Get Iceberg connection
603                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            // For S3 Tables connections in the Materialize Cloud product, verify the
615            // AWS region matches the environment's region. This check only applies when
616            // the enable_s3_tables_region_check dyncfg is set.
617            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                        // Later on we default to "us-east-1" if the region is not set on the S3 Tables
625                        // connection, so we need to do the same check here.
626                        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            // Validate the sink's (optional) AWS connection even though we never use it.
643            // TODO(kynan): If we do start using the sink's creds, check again that this validation
644            //   accurately reflects what we need.
645            //   Consider rolling the storage creds validation into the catalog connection's "connect" fn,
646            //   which already validates the catalog creds (currently also used for the storage layer).
647            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                    // Get AWS connection
652                    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            // Now that we've validated the sink's storage creds (if they exist)
674            // we _could_ use them to build a complete Iceberg client (both catalog and storage).
675            // TODO(kynan): Actually use those sink-specific creds here instead of ignoring them.
676            // Purification only proves the catalog is reachable, so it needs no table-scoped
677            // storage credentials.
678            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            // Get Kafka connection
710            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
734/// Runs a function on each format within the format specifier.
735///
736/// The function is provided with a `DocOnSchema` that indicates whether the
737/// format is for the key, value, or both.
738///
739// TODO(benesch): rename `DocOnSchema` to the more general `FormatRestriction`.
740fn 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/// Defines whether purification should enforce that at least one valid source
755/// reference is provided on the provided statement.
756#[derive(Debug, Copy, Clone, PartialEq, Eq)]
757pub(crate) enum SourceReferencePolicy {
758    /// Don't allow source references to be provided. This is used for
759    /// enforcing that `CREATE SOURCE` statements don't create subsources.
760    NotAllowed,
761    /// Allow empty references, such as when creating a source that
762    /// will have tables added afterwards.
763    Optional,
764    /// Require that at least one reference is resolved to an upstream
765    /// object.
766    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        // Depending on if the user must or can use the `CREATE TABLE .. FROM SOURCE` statement
852        // to add tables after this source is created we might need to enforce that
853        // auto-generated subsources are created or not created by this source statement.
854        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                    // anyhow doesn't support Clone, so not trivial to move into PlanError
1013                    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                    // Validate that the topic at least exists.
1025                    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                    // Validate the start offsets.
1040                    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                    // Translate `START TIMESTAMP` to a start offset.
1053                    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            // Record whether the upstream server is a physical replica, which changes how we can determine
1118            // the latest LSN.
1119            let is_physical_replica = mz_postgres_util::get_is_in_recovery(&client).await?;
1120
1121            // PostgreSQL 17 introduced failover slots, and rejects creating one on a standby.
1122            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            // Read after the references above, so it bounds the LSN their schemas belong to.
1130            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            // Record the active replication timeline_id to allow detection of a future upstream
1161            // point-in-time-recovery that will put the source into an error state.
1162            let timeline_id = mz_postgres_util::get_timeline_id(&client).await?;
1163
1164            // Remove any old detail references
1165            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            // Connect to the upstream SQL Server instance so we can validate
1187            // we're compatible with CDC.
1188            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            // Update our set of requested source exports.
1234            requested_subsource_map.extend(source_exports);
1235
1236            // Record the most recent restore_history_id, or none if the system has never been
1237            // restored.
1238
1239            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            // Update our 'TEXT' and 'EXCLUDE' column options with the purified and normalized set.
1254            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            // Retrieve the current @gtid_executed value of the server to mark as the effective
1288            // initial snapshot point such that we can ensure consistency if the initial source
1289            // snapshot is broken up over multiple points in time.
1290            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            // We don't have any fields in this details struct but keep this around for
1322            // conformity with postgres and in-case we end up needing it again in the future.
1323            let details = MySqlSourceDetails {};
1324            // Update options with the purified details
1325            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            // Filter to the references that need to be created as 'subsources', which
1357            // doesn't include the default output for single-output sources.
1358            // TODO(database-issues#8620): Remove once subsources are removed
1359            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    // Now that we know which subsources to create alongside this
1446    // statement, remove the references so it is not canonicalized as
1447    // part of the `CREATE SOURCE` statement in the catalog.
1448    *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/// Catalog-derived inputs of upstream reads for `ALTER SOURCE ... ADD SUBSOURCES`
1460/// and `REFRESH REFERENCES`.
1461#[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
1574// TODO(database-issues#8620): Remove once subsources are removed
1575async 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    // Validate this is a source that can have subsources added.
1585    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            // Get PostgresConnection for generating subsources.
1610            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            // Read after the references above, so it bounds the LSN their schemas belong to.
1624            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            // Retrieve the current @gtid_executed value of the server to mark as the effective
1678            // initial snapshot point for these subsources.
1679            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            // Update options with the purified details
1713            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            // Open a connection to the upstream SQL Server instance.
1729            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            // Query the upstream SQL Server instance for available tables to replicate.
1740            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            // Add the new exports to our subsource map.
1771            requested_subsource_map.extend(source_exports);
1772
1773            // Update options on the CREATE SOURCE statement with the purified details.
1774            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        // The source kind was already established earlier in this function;
1789        // reaching this catch-all is an internal invariant violation.
1790        _ => 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            // Get PostgresConnection for generating subsources.
1808            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            // Open a connection to the upstream SQL Server instance.
1856            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            // Query the upstream SQL Server instance for available tables to replicate.
1867            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    // Should be overriden below if a source-specific format is required.
1919    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    // Run purification work specific to each source type: resolve the external reference to
1936    // a fully qualified name and obtain the appropriate details for the source-export statement
1937    let purified_export = match connection {
1938        GenericSourceConnection::Postgres(pg_source_connection) => {
1939            // Get PostgresConnection for generating subsources.
1940            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            // Read after the references above, so it bounds the LSN their schemas belong to.
1954            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                // TODO(database-issues#8620): Remove once subsources are removed
1967                // This `normalized_text_columns` is not relevant for us and is only returned for
1968                // `CREATE SOURCE` statements that automatically generate subsources
1969                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            // There should be exactly one source_export returned for this statement
1984            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            // Retrieve the current @gtid_executed value of the server to mark as the effective
2008            // initial snapshot point for this table.
2009            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                // TODO(database-issues#8620): Remove once subsources are removed
2021                // `normalized_text/exclude_columns` is not relevant for us and is only returned for
2022                // `CREATE SOURCE` statements that automatically generate subsources
2023                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            // There should be exactly one source_export returned for this statement
2040            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            // There should be exactly one source_export returned for this statement
2079            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            // There should be exactly one source_export returned
2091            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            // There should be exactly one source_export returned
2116            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    // Columns and constraints cannot be specified by the user but will be populated below.
2150    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    // Get the source item
2160    let item = match scx.get_item_by_resolved_name(source_name) {
2161        Ok(item) => item,
2162        Err(e) => return Err(e),
2163    };
2164
2165    // Ensure it's an ingestion-based and alterable source.
2166    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    // Our text column values are unqualified (just column names), but the purification methods below
2196    // expect to match the fully-qualified names against the full set of tables in upstream, so we
2197    // need to qualify them using the external reference first.
2198    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    // Update the external reference in the statement to the resolved fully-qualified
2253    // external reference
2254    *external_reference = Some(purified_export.external_reference.clone());
2255
2256    // Update options in the statement using the purified export details
2257    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                    // Purification is what populates the `Defined` variant;
2299                    // reaching this is an internal invariant violation.
2300                    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                    // Purification is what populates the `Defined` variant;
2349                    // reaching this is an internal invariant violation.
2350                    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                    // Purification is what populates the `Defined` variant;
2406                    // reaching this is an internal invariant violation.
2407                    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            // We only determine the table description for multi-output load generator sources here,
2426            // whereas single-output load generators will have their relation description
2427            // determined during statment planning as envelope and format options may affect their
2428            // schema.
2429            if let Some(desc) = desc {
2430                let (gen_columns, gen_constraints) = scx.relation_desc_into_table_defs(&desc)?;
2431                match columns {
2432                    // Purification is what populates the `Defined` variant;
2433                    // reaching this is an internal invariant violation.
2434                    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            // NOTE: Kafka tables have their 'schemas' purified into the statement inside the
2456            // format field, so we don't specify any columns or constraints to be stored
2457            // on the statement here. The RelationDesc will be determined during planning.
2458            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    // TODO: We might as well use the retrieved available references to update the source
2469    // available references table in the catalog, so plumb this through.
2470    // available_source_references: retrieved_source_references.available_source_references(),
2471    Ok(PurifiedStatement::PurifiedCreateTableFromSource { stmt })
2472}
2473
2474#[derive(Clone, Debug)]
2475enum SourceFormatOptions {
2476    Default,
2477    Kafka { topic: String },
2478}
2479
2480/// Generates a progress subsource, choosing an available name in the supplied catalog.
2481pub 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        // Progress subsources do not refer to the source to which they belong.
2565        // Instead the primary source depends on it (the opposite is true of
2566        // ingestion exports, which depend on the primary source).
2567        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    // get the first subsource to determine the connection type
2580    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                // Create the subsource statement
2624                let subsource = CreateSubsourceStatement {
2625                    name: subsource_name,
2626                    columns,
2627                    of_source: Some(source_name.clone()),
2628                    // unlike sources that come from an external upstream, we
2629                    // have more leniency to introduce different constraints
2630                    // every time the load generator is run; i.e. we are not as
2631                    // worried about introducing junk data.
2632                    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            // TODO: as part of database-issues#8322, Kafka sources will begin
2656            // producing data––we'll need to understand the schema
2657            // of the output here.
2658            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        // Use storage decoding's authentication and address policy.
2764        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                // We are in a normal tokio context during purification.
2773                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        // The runtime decode path only handles Avro; reject other formats at
2789        // planning time so the failure mode is a clear SQL error, not a
2790        // permanent decode error on every record.
2791        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    // Structural validation below runs unconditionally — even when a seed is
2905    // already present. The seed is not proof the rest of the statement is
2906    // well-formed: the grammar accepts a user-written `SEED VALUE SCHEMA`, so a
2907    // seeded statement may have arrived straight from a user rather than from
2908    // our own re-parsed `create_sql`. Only the Glue *fetch* is safe to skip
2909    // when seeded; the connection-type and `SCHEMA NAME` checks are not.
2910    let Connection::GlueSchemaRegistry(gsr_connection) = item.connection()? else {
2911        return Err(GluePurificationError::NotGlueConnection(full_name).into());
2912    };
2913
2914    // Pull `SCHEMA NAME` out of the option bag. Required. Use the shared
2915    // extractor (rather than a hand-rolled match) so non-string values get a
2916    // proper "invalid value" error instead of being silently dropped.
2917    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    // The per-side and compatibility options are sink-only. A source reads a
2926    // single writer schema, so reject them here rather than silently ignoring.
2927    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        // The schema is already pinned — either we purified this statement
2942        // before and re-parsed it from persisted create_sql, or a user supplied
2943        // the seed directly. Either way the reader schema is fixed, so skip the
2944        // Glue lookup. (The structural checks above still ran.)
2945        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    /// Reference schemas for the key schema, in dependency order.
2960    pub key_reference_schemas: Vec<String>,
2961    /// Reference schemas for the value schema, in dependency order.
2962    pub value_reference_schemas: Vec<String>,
2963}
2964
2965/// Result of fetching a schema, including any referenced schemas.
2966struct SchemaWithReferences {
2967    /// The primary schema.
2968    schema: String,
2969    /// Reference schemas in dependency order.
2970    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            // Use get_subject_and_references to also fetch referenced schemas
2981            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        // It's possible that CSR was provided with an inline strategy, but without a subject
2995        // or schema id to look up, there isn't a clean way to lookup references for the schema.
2996        // This could be done with extra work, but at this point, it isn't clear this is needed.
2997        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
3039/// Collect protobuf message descriptor from CSR and compile the descriptor.
3040async 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    // Compile .proto files into a file descriptor set.
3053    let mut source_tree = VirtualSourceTree::new();
3054
3055    // Add well-known types (e.g., google/protobuf/timestamp.proto) to the source
3056    // tree. These are implicitly available to protoc but are typically not
3057    // registered in the schema registry.
3058    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    // Ensure there is exactly one message in the file.
3073    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    // Encode the file descriptor set into a SQL byte string.
3081    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
3096/// Purifies a CREATE MATERIALIZED VIEW statement. Additionally, it adjusts `resolved_ids` if
3097/// references to ids appear or disappear during the purification.
3098///
3099/// Note that in contrast with [`purify_statement`], this doesn't need to be async, because
3100/// this function is not making any network calls.
3101pub 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    // 0. Preparations:
3108    // Prepare an expression that calls `mz_now()`, which we can insert in various later steps.
3109    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    // Prepare the `mz_timestamp` type.
3138    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        // 1. Purify `REFRESH AT CREATION` to `REFRESH AT mz_now()`.
3157        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        // 2. If `REFRESH EVERY` doesn't have an `ALIGNED TO`, then add `ALIGNED TO mz_now()`.
3169        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        // 3. Substitute `mz_now()` with the timestamp chosen for the CREATE MATERIALIZED VIEW
3179        // statement. (This has to happen after the above steps, which might introduce `mz_now()`.)
3180        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    // 4. If the user didn't give any REFRESH option, then default to ON COMMIT.
3203    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    // 5. Attend to `resolved_ids`: The purification might have
3219    // - added references to `mz_timestamp`;
3220    // - removed references to `mz_now`.
3221    if introduced_mz_timestamp {
3222        resolved_ids.add_item(mz_timestamp_id);
3223    }
3224    // Even though we always remove `mz_now()` from the `with_options`, there might be `mz_now()`
3225    // remaining in the main query expression of the MV, so let's visit the entire statement to look
3226    // for `mz_now()` everywhere.
3227    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
3234/// Returns true if the [MaterializedViewOption] either already involves `mz_now()` or will involve
3235/// after purification.
3236pub 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            // For a `REFRESH EVERY` without an `ALIGNED TO`, purification will default the
3256            // `ALIGNED TO` to `mz_now()`.
3257            true
3258        }
3259        Some(WithOptionValue::Refresh(RefreshOptionValue::AtCreation)) => {
3260            // `REFRESH AT CREATION` will be purified to `REFRESH AT mz_now()`.
3261            true
3262        }
3263        _ => false,
3264    }
3265}
3266
3267/// Determines whether the AST involves `mz_now()`.
3268struct 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                // We substitute `mz_now()` with number + a cast to `mz_timestamp`. The cast is to
3321                // not alter the type of the expression.
3322                *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}