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                text_columns,
1094                exclude_columns,
1095                details,
1096                ..
1097            } = options.clone().try_into()?;
1098            let publication =
1099                publication.ok_or(PgSourcePurificationError::ConnectionMissingPublication)?;
1100
1101            if details.is_some() {
1102                Err(PgSourcePurificationError::UserSpecifiedDetails)?;
1103            }
1104
1105            let client = connection
1106                .validate(connection_id, storage_configuration)
1107                .await?;
1108
1109            let reference_client = SourceReferenceClient::Postgres {
1110                client: &client,
1111                publication: &publication,
1112                database: &connection.database,
1113            };
1114            retrieved_source_references = reference_client.get_source_references().await?;
1115
1116            // Record whether the upstream server is a physical replica, which changes how we can determine
1117            // the latest LSN.
1118            let is_physical_replica = mz_postgres_util::get_is_in_recovery(&client).await?;
1119
1120            // Read after the references above, so it bounds the LSN their schemas belong to.
1121            let initial_lsn = MzOffset::from(
1122                mz_postgres_util::fetch_max_lsn(&client, is_physical_replica).await?,
1123            );
1124
1125            let postgres::PurifiedSourceExports {
1126                source_exports: subsources,
1127                normalized_text_columns,
1128            } = postgres::purify_source_exports(
1129                &client,
1130                &retrieved_source_references,
1131                external_references,
1132                text_columns,
1133                exclude_columns,
1134                &FilterConstraints::none(),
1135                false,
1136                source_name,
1137                &reference_policy,
1138                initial_lsn,
1139            )
1140            .await?;
1141
1142            if let Some(text_cols_option) = options
1143                .iter_mut()
1144                .find(|option| option.name == PgConfigOptionName::TextColumns)
1145            {
1146                text_cols_option.value = Some(WithOptionValue::Sequence(normalized_text_columns));
1147            }
1148
1149            requested_subsource_map.extend(subsources);
1150
1151            // Record the active replication timeline_id to allow detection of a future upstream
1152            // point-in-time-recovery that will put the source into an error state.
1153            let timeline_id = mz_postgres_util::get_timeline_id(&client).await?;
1154
1155            // Remove any old detail references
1156            options.retain(|PgConfigOption { name, .. }| name != &PgConfigOptionName::Details);
1157            let details = PostgresSourcePublicationDetails {
1158                slot: format!(
1159                    "materialize_{}",
1160                    Uuid::new_v4().to_string().replace('-', "")
1161                ),
1162                timeline_id: Some(timeline_id),
1163                database: connection.database,
1164                is_physical_replica: Some(is_physical_replica),
1165            };
1166            options.push(PgConfigOption {
1167                name: PgConfigOptionName::Details,
1168                value: Some(WithOptionValue::Value(Value::String(hex::encode(
1169                    details.into_proto().encode_to_vec(),
1170                )))),
1171            })
1172        }
1173        (
1174            CreateSourceConnection::SqlServer { options, .. },
1175            PreparedSourceConnection::SqlServer(connection_id, connection),
1176        ) => {
1177            // Connect to the upstream SQL Server instance so we can validate
1178            // we're compatible with CDC.
1179            let crate::plan::statement::ddl::SqlServerConfigOptionExtracted {
1180                details,
1181                text_columns,
1182                exclude_columns,
1183                seen: _,
1184            } = options.clone().try_into()?;
1185
1186            if details.is_some() {
1187                Err(SqlServerSourcePurificationError::UserSpecifiedDetails)?;
1188            }
1189
1190            let mut client = connection
1191                .validate(connection_id, storage_configuration)
1192                .await?;
1193
1194            let database: Arc<str> = connection.database.into();
1195            let reference_client = SourceReferenceClient::SqlServer {
1196                client: &mut client,
1197                database: Arc::clone(&database),
1198            };
1199            retrieved_source_references = reference_client.get_source_references().await?;
1200            tracing::debug!(?retrieved_source_references, "got source references");
1201
1202            let timeout = mz_storage_types::sources::sql_server::MAX_LSN_WAIT
1203                .get(storage_configuration.config_set());
1204
1205            let purified_source_exports = sql_server::purify_source_exports(
1206                &*database,
1207                &mut client,
1208                &retrieved_source_references,
1209                external_references,
1210                &text_columns,
1211                &exclude_columns,
1212                source_name,
1213                timeout,
1214                &reference_policy,
1215            )
1216            .await?;
1217
1218            let sql_server::PurifiedSourceExports {
1219                source_exports,
1220                normalized_text_columns,
1221                normalized_excl_columns,
1222            } = purified_source_exports;
1223
1224            // Update our set of requested source exports.
1225            requested_subsource_map.extend(source_exports);
1226
1227            // Record the most recent restore_history_id, or none if the system has never been
1228            // restored.
1229
1230            let restore_history_id =
1231                mz_sql_server_util::inspect::get_latest_restore_history_id(&mut client).await?;
1232            let details = SqlServerSourceExtras { restore_history_id };
1233
1234            options.retain(|SqlServerConfigOption { name, .. }| {
1235                name != &SqlServerConfigOptionName::Details
1236            });
1237            options.push(SqlServerConfigOption {
1238                name: SqlServerConfigOptionName::Details,
1239                value: Some(WithOptionValue::Value(Value::String(hex::encode(
1240                    details.into_proto().encode_to_vec(),
1241                )))),
1242            });
1243
1244            // Update our 'TEXT' and 'EXCLUDE' column options with the purified and normalized set.
1245            if let Some(text_cols_option) = options
1246                .iter_mut()
1247                .find(|option| option.name == SqlServerConfigOptionName::TextColumns)
1248            {
1249                text_cols_option.value = Some(WithOptionValue::Sequence(normalized_text_columns));
1250            }
1251            if let Some(excl_cols_option) = options
1252                .iter_mut()
1253                .find(|option| option.name == SqlServerConfigOptionName::ExcludeColumns)
1254            {
1255                excl_cols_option.value = Some(WithOptionValue::Sequence(normalized_excl_columns));
1256            }
1257        }
1258        (
1259            CreateSourceConnection::MySql { options, .. },
1260            PreparedSourceConnection::MySql(connection_id, connection),
1261        ) => {
1262            let crate::plan::statement::ddl::MySqlConfigOptionExtracted {
1263                details,
1264                text_columns,
1265                exclude_columns,
1266                seen: _,
1267            } = options.clone().try_into()?;
1268
1269            if details.is_some() {
1270                Err(MySqlSourcePurificationError::UserSpecifiedDetails)?;
1271            }
1272
1273            let mut conn = connection
1274                .validate(connection_id, storage_configuration)
1275                .await
1276                .map_err(MySqlSourcePurificationError::InvalidConnection)?;
1277
1278            // Retrieve the current @gtid_executed value of the server to mark as the effective
1279            // initial snapshot point such that we can ensure consistency if the initial source
1280            // snapshot is broken up over multiple points in time.
1281            let initial_gtid_set =
1282                mz_mysql_util::query_sys_var(&mut conn, "global.gtid_executed").await?;
1283
1284            let binlog_full_metadata = is_binlog_full_metadata(&mut conn).await?;
1285
1286            let reference_client = SourceReferenceClient::MySql {
1287                conn: &mut conn,
1288                include_system_schemas: mysql::references_system_schemas(external_references),
1289            };
1290            retrieved_source_references = reference_client.get_source_references().await?;
1291
1292            let mysql::PurifiedSourceExports {
1293                source_exports: subsources,
1294                normalized_text_columns,
1295                normalized_exclude_columns,
1296            } = mysql::purify_source_exports(
1297                &mut conn,
1298                &retrieved_source_references,
1299                external_references,
1300                text_columns,
1301                exclude_columns,
1302                &FilterConstraints::none(),
1303                false,
1304                source_name,
1305                initial_gtid_set.clone(),
1306                &reference_policy,
1307                binlog_full_metadata,
1308            )
1309            .await?;
1310            requested_subsource_map.extend(subsources);
1311
1312            // We don't have any fields in this details struct but keep this around for
1313            // conformity with postgres and in-case we end up needing it again in the future.
1314            let details = MySqlSourceDetails {};
1315            // Update options with the purified details
1316            options
1317                .retain(|MySqlConfigOption { name, .. }| name != &MySqlConfigOptionName::Details);
1318            options.push(MySqlConfigOption {
1319                name: MySqlConfigOptionName::Details,
1320                value: Some(WithOptionValue::Value(Value::String(hex::encode(
1321                    details.into_proto().encode_to_vec(),
1322                )))),
1323            });
1324
1325            if let Some(text_cols_option) = options
1326                .iter_mut()
1327                .find(|option| option.name == MySqlConfigOptionName::TextColumns)
1328            {
1329                text_cols_option.value = Some(WithOptionValue::Sequence(normalized_text_columns));
1330            }
1331            if let Some(exclude_cols_option) = options
1332                .iter_mut()
1333                .find(|option| option.name == MySqlConfigOptionName::ExcludeColumns)
1334            {
1335                exclude_cols_option.value =
1336                    Some(WithOptionValue::Sequence(normalized_exclude_columns));
1337            }
1338        }
1339        (
1340            CreateSourceConnection::LoadGenerator { generator, options },
1341            PreparedSourceConnection::LoadGenerator(load_generator),
1342        ) => {
1343            let reference_client = SourceReferenceClient::LoadGenerator {
1344                generator: &load_generator,
1345            };
1346            retrieved_source_references = reference_client.get_source_references().await?;
1347            // Filter to the references that need to be created as 'subsources', which
1348            // doesn't include the default output for single-output sources.
1349            // TODO(database-issues#8620): Remove once subsources are removed
1350            let multi_output_sources =
1351                retrieved_source_references
1352                    .all_references()
1353                    .iter()
1354                    .any(|r| {
1355                        matches!(
1356                            r.load_generator_output(),
1357                            Some(output) if output != &LoadGeneratorOutput::Default
1358                        )
1359                    });
1360
1361            match external_references {
1362                Some(requested)
1363                    if matches!(reference_policy, SourceReferencePolicy::NotAllowed) =>
1364                {
1365                    Err(PlanError::UseTablesForSources(requested.to_string()))?
1366                }
1367                Some(requested) if !multi_output_sources => match requested {
1368                    ExternalReferences::SubsetTables(_) => {
1369                        Err(LoadGeneratorSourcePurificationError::ForTables)?
1370                    }
1371                    ExternalReferences::SubsetSchemas(_) => {
1372                        Err(LoadGeneratorSourcePurificationError::ForSchemas)?
1373                    }
1374                    ExternalReferences::All => {
1375                        Err(LoadGeneratorSourcePurificationError::ForAllTables)?
1376                    }
1377                },
1378                Some(requested) => {
1379                    let requested_exports = retrieved_source_references
1380                        .requested_source_exports(Some(requested), source_name)?;
1381                    for export in requested_exports {
1382                        requested_subsource_map.insert(
1383                            export.name,
1384                            PurifiedSourceExport {
1385                                external_reference: export.external_reference,
1386                                details: PurifiedExportDetails::LoadGenerator {
1387                                    table: export
1388                                        .meta
1389                                        .load_generator_desc()
1390                                        .ok_or_else(|| {
1391                                            internal_err!(
1392                                                "expected load generator source reference"
1393                                            )
1394                                        })?
1395                                        .clone(),
1396                                    output: export
1397                                        .meta
1398                                        .load_generator_output()
1399                                        .ok_or_else(|| {
1400                                            internal_err!(
1401                                                "expected load generator source reference"
1402                                            )
1403                                        })?
1404                                        .clone(),
1405                                },
1406                            },
1407                        );
1408                    }
1409                }
1410                None => {
1411                    if multi_output_sources
1412                        && matches!(reference_policy, SourceReferencePolicy::Required)
1413                    {
1414                        Err(LoadGeneratorSourcePurificationError::MultiOutputRequiresForAllTables)?
1415                    }
1416                }
1417            }
1418
1419            if let LoadGenerator::Clock = generator {
1420                if !options
1421                    .iter()
1422                    .any(|p| p.name == LoadGeneratorOptionName::AsOf)
1423                {
1424                    let now =
1425                        clock_as_of.expect("implicit CLOCK AS OF was captured during preparation");
1426                    options.push(LoadGeneratorOption {
1427                        name: LoadGeneratorOptionName::AsOf,
1428                        value: Some(WithOptionValue::Value(Value::Number(now.to_string()))),
1429                    });
1430                }
1431            }
1432        }
1433        _ => unreachable!("source preparation inputs must match their statement"),
1434    }
1435
1436    // Now that we know which subsources to create alongside this
1437    // statement, remove the references so it is not canonicalized as
1438    // part of the `CREATE SOURCE` statement in the catalog.
1439    *external_references = None;
1440
1441    Ok(PreparedCreateSource {
1442        format_options,
1443        create_progress_subsource: uses_old_syntax,
1444        create_source_stmt,
1445        subsources: requested_subsource_map,
1446        available_source_references: retrieved_source_references.available_source_references(),
1447    })
1448}
1449
1450/// Catalog-derived inputs of upstream reads for `ALTER SOURCE ... ADD SUBSOURCES`
1451/// and `REFRESH REFERENCES`.
1452#[derive(Clone, Debug, PartialEq, Eq)]
1453enum AlterSourcePreparation {
1454    AddSubsources {
1455        external_references: Vec<ExternalReferenceExport>,
1456        options: Vec<AlterSourceAddSubsourceOption<Aug>>,
1457        desc: SourceDesc,
1458        partial_name: PartialItemName,
1459        unresolved_source_name: UnresolvedItemName,
1460        resolved_source_name: ResolvedItemName,
1461    },
1462    RefreshReferences {
1463        desc: SourceDesc,
1464        resolved_source_name: ResolvedItemName,
1465    },
1466}
1467
1468fn prepare_alter_source(
1469    catalog: &dyn SessionCatalog,
1470    stmt: &AlterSourceStatement<Aug>,
1471) -> Result<Option<AlterSourcePreparation>, PlanError> {
1472    let scx = StatementContext::new(None, catalog);
1473    let AlterSourceStatement {
1474        source_name: unresolved_source_name,
1475        action,
1476        if_exists,
1477    } = stmt;
1478    let item = match scx.resolve_item(RawItemName::Name(unresolved_source_name.clone())) {
1479        Ok(item) => item,
1480        Err(_) if *if_exists => return Ok(None),
1481        Err(e) => return Err(e),
1482    };
1483    let desc = match item.source_desc()? {
1484        Some(desc) => desc.clone().into_inline_connection(scx.catalog),
1485        None => sql_bail!("cannot ALTER this type of source"),
1486    };
1487    let source_name = item.name();
1488    let resolved_source_name = ResolvedItemName::Item {
1489        id: item.id(),
1490        qualifiers: item.name().qualifiers.clone(),
1491        full_name: scx.catalog.resolve_full_name(source_name),
1492        print_id: true,
1493        version: RelationVersionSelector::Latest,
1494    };
1495    let partial_name = scx.catalog.minimal_qualification(source_name);
1496    match action {
1497        AlterSourceAction::AddSubsources {
1498            external_references,
1499            options,
1500        } => {
1501            if scx.catalog.system_vars().enable_create_table_from_source()
1502                && scx.catalog.system_vars().force_source_table_syntax()
1503            {
1504                return Err(PlanError::UseTablesForSources(
1505                    "ALTER SOURCE .. ADD SUBSOURCES ..".to_string(),
1506                ));
1507            }
1508            Ok(Some(AlterSourcePreparation::AddSubsources {
1509                external_references: external_references.clone(),
1510                options: options.clone(),
1511                desc,
1512                partial_name,
1513                unresolved_source_name: unresolved_source_name.clone(),
1514                resolved_source_name,
1515            }))
1516        }
1517        AlterSourceAction::RefreshReferences => {
1518            Ok(Some(AlterSourcePreparation::RefreshReferences {
1519                desc,
1520                resolved_source_name,
1521            }))
1522        }
1523        _ => Ok(None),
1524    }
1525}
1526
1527async fn execute_alter_source_preparation(
1528    request: AlterSourcePreparation,
1529    storage_configuration: &StorageConfiguration,
1530) -> Result<PurifiedStatement, PlanError> {
1531    match request {
1532        AlterSourcePreparation::AddSubsources {
1533            external_references,
1534            options,
1535            desc,
1536            partial_name,
1537            unresolved_source_name,
1538            resolved_source_name,
1539        } => {
1540            purify_alter_source_add_subsources(
1541                external_references,
1542                options,
1543                desc,
1544                partial_name,
1545                unresolved_source_name,
1546                resolved_source_name,
1547                storage_configuration,
1548            )
1549            .await
1550        }
1551        AlterSourcePreparation::RefreshReferences {
1552            desc,
1553            resolved_source_name,
1554        } => {
1555            purify_alter_source_refresh_references(
1556                desc,
1557                resolved_source_name,
1558                storage_configuration,
1559            )
1560            .await
1561        }
1562    }
1563}
1564
1565// TODO(database-issues#8620): Remove once subsources are removed
1566async fn purify_alter_source_add_subsources(
1567    external_references: Vec<ExternalReferenceExport>,
1568    mut options: Vec<AlterSourceAddSubsourceOption<Aug>>,
1569    desc: SourceDesc,
1570    partial_source_name: PartialItemName,
1571    unresolved_source_name: UnresolvedItemName,
1572    resolved_source_name: ResolvedItemName,
1573    storage_configuration: &StorageConfiguration,
1574) -> Result<PurifiedStatement, PlanError> {
1575    // Validate this is a source that can have subsources added.
1576    let connection_id = match &desc.connection {
1577        GenericSourceConnection::Postgres(c) => c.connection_id,
1578        GenericSourceConnection::MySql(c) => c.connection_id,
1579        GenericSourceConnection::SqlServer(c) => c.connection_id,
1580        _ => sql_bail!(
1581            "source {} does not support ALTER SOURCE.",
1582            partial_source_name
1583        ),
1584    };
1585
1586    let crate::plan::statement::ddl::AlterSourceAddSubsourceOptionExtracted {
1587        text_columns,
1588        exclude_columns,
1589        details,
1590        seen: _,
1591    } = options.clone().try_into()?;
1592    if details.is_some() {
1593        sql_bail!("DETAILS option cannot be explicitly set");
1594    }
1595
1596    let mut requested_subsource_map = BTreeMap::new();
1597
1598    match desc.connection {
1599        GenericSourceConnection::Postgres(pg_source_connection) => {
1600            // Get PostgresConnection for generating subsources.
1601            let pg_connection = &pg_source_connection.connection;
1602
1603            let client = pg_connection
1604                .validate(connection_id, storage_configuration)
1605                .await?;
1606
1607            let reference_client = SourceReferenceClient::Postgres {
1608                client: &client,
1609                publication: &pg_source_connection.publication,
1610                database: &pg_connection.database,
1611            };
1612            let retrieved_source_references = reference_client.get_source_references().await?;
1613
1614            // Read after the references above, so it bounds the LSN their schemas belong to.
1615            let initial_lsn = MzOffset::from(
1616                mz_postgres_util::fetch_max_lsn(
1617                    &client,
1618                    pg_source_connection
1619                        .publication_details
1620                        .get_is_physical_replica(),
1621                )
1622                .await?,
1623            );
1624
1625            let postgres::PurifiedSourceExports {
1626                source_exports: subsources,
1627                normalized_text_columns,
1628            } = postgres::purify_source_exports(
1629                &client,
1630                &retrieved_source_references,
1631                &Some(ExternalReferences::SubsetTables(external_references)),
1632                text_columns,
1633                exclude_columns,
1634                &FilterConstraints::none(),
1635                false,
1636                &unresolved_source_name,
1637                &SourceReferencePolicy::Required,
1638                initial_lsn,
1639            )
1640            .await?;
1641
1642            if let Some(text_cols_option) = options
1643                .iter_mut()
1644                .find(|option| option.name == AlterSourceAddSubsourceOptionName::TextColumns)
1645            {
1646                text_cols_option.value = Some(WithOptionValue::Sequence(normalized_text_columns));
1647            }
1648
1649            requested_subsource_map.extend(subsources);
1650        }
1651        GenericSourceConnection::MySql(mysql_source_connection) => {
1652            let mysql_connection = &mysql_source_connection.connection;
1653            let config = mysql_connection
1654                .config(
1655                    &storage_configuration.connection_context.secrets_reader,
1656                    storage_configuration,
1657                    InTask::No,
1658                )
1659                .await?;
1660
1661            let mut conn = config
1662                .connect(
1663                    "mysql purification",
1664                    &storage_configuration.connection_context.ssh_tunnel_manager,
1665                )
1666                .await?;
1667
1668            // Retrieve the current @gtid_executed value of the server to mark as the effective
1669            // initial snapshot point for these subsources.
1670            let initial_gtid_set =
1671                mz_mysql_util::query_sys_var(&mut conn, "global.gtid_executed").await?;
1672
1673            let binlog_full_metadata = is_binlog_full_metadata(&mut conn).await?;
1674
1675            let requested_references = Some(ExternalReferences::SubsetTables(external_references));
1676
1677            let reference_client = SourceReferenceClient::MySql {
1678                conn: &mut conn,
1679                include_system_schemas: mysql::references_system_schemas(&requested_references),
1680            };
1681            let retrieved_source_references = reference_client.get_source_references().await?;
1682
1683            let mysql::PurifiedSourceExports {
1684                source_exports: subsources,
1685                normalized_text_columns,
1686                normalized_exclude_columns,
1687            } = mysql::purify_source_exports(
1688                &mut conn,
1689                &retrieved_source_references,
1690                &requested_references,
1691                text_columns,
1692                exclude_columns,
1693                &FilterConstraints::none(),
1694                false,
1695                &unresolved_source_name,
1696                initial_gtid_set,
1697                &SourceReferencePolicy::Required,
1698                binlog_full_metadata,
1699            )
1700            .await?;
1701            requested_subsource_map.extend(subsources);
1702
1703            // Update options with the purified details
1704            if let Some(text_cols_option) = options
1705                .iter_mut()
1706                .find(|option| option.name == AlterSourceAddSubsourceOptionName::TextColumns)
1707            {
1708                text_cols_option.value = Some(WithOptionValue::Sequence(normalized_text_columns));
1709            }
1710            if let Some(exclude_cols_option) = options
1711                .iter_mut()
1712                .find(|option| option.name == AlterSourceAddSubsourceOptionName::ExcludeColumns)
1713            {
1714                exclude_cols_option.value =
1715                    Some(WithOptionValue::Sequence(normalized_exclude_columns));
1716            }
1717        }
1718        GenericSourceConnection::SqlServer(sql_server_source) => {
1719            // Open a connection to the upstream SQL Server instance.
1720            let sql_server_connection = &sql_server_source.connection;
1721            let config = sql_server_connection
1722                .resolve_config(
1723                    &storage_configuration.connection_context.secrets_reader,
1724                    storage_configuration,
1725                    InTask::No,
1726                )
1727                .await?;
1728            let mut client = mz_sql_server_util::Client::connect(config).await?;
1729
1730            // Query the upstream SQL Server instance for available tables to replicate.
1731            let database = sql_server_connection.database.clone().into();
1732            let source_references = SourceReferenceClient::SqlServer {
1733                client: &mut client,
1734                database: Arc::clone(&database),
1735            }
1736            .get_source_references()
1737            .await?;
1738            let requested_references = Some(ExternalReferences::SubsetTables(external_references));
1739
1740            let timeout = mz_storage_types::sources::sql_server::MAX_LSN_WAIT
1741                .get(storage_configuration.config_set());
1742
1743            let result = sql_server::purify_source_exports(
1744                &*database,
1745                &mut client,
1746                &source_references,
1747                &requested_references,
1748                &text_columns,
1749                &exclude_columns,
1750                &unresolved_source_name,
1751                timeout,
1752                &SourceReferencePolicy::Required,
1753            )
1754            .await;
1755            let sql_server::PurifiedSourceExports {
1756                source_exports,
1757                normalized_text_columns,
1758                normalized_excl_columns,
1759            } = result?;
1760
1761            // Add the new exports to our subsource map.
1762            requested_subsource_map.extend(source_exports);
1763
1764            // Update options on the CREATE SOURCE statement with the purified details.
1765            if let Some(text_cols_option) = options
1766                .iter_mut()
1767                .find(|option| option.name == AlterSourceAddSubsourceOptionName::TextColumns)
1768            {
1769                text_cols_option.value = Some(WithOptionValue::Sequence(normalized_text_columns));
1770            }
1771            if let Some(exclude_cols_option) = options
1772                .iter_mut()
1773                .find(|option| option.name == AlterSourceAddSubsourceOptionName::ExcludeColumns)
1774            {
1775                exclude_cols_option.value =
1776                    Some(WithOptionValue::Sequence(normalized_excl_columns));
1777            }
1778        }
1779        // The source kind was already established earlier in this function;
1780        // reaching this catch-all is an internal invariant violation.
1781        _ => bail_internal!("source does not support ALTER SOURCE...ADD SUBSOURCE"),
1782    };
1783
1784    Ok(PurifiedStatement::PurifiedAlterSourceAddSubsources {
1785        source_name: resolved_source_name,
1786        options,
1787        subsources: requested_subsource_map,
1788    })
1789}
1790
1791async fn purify_alter_source_refresh_references(
1792    desc: SourceDesc,
1793    resolved_source_name: ResolvedItemName,
1794    storage_configuration: &StorageConfiguration,
1795) -> Result<PurifiedStatement, PlanError> {
1796    let retrieved_source_references = match desc.connection {
1797        GenericSourceConnection::Postgres(pg_source_connection) => {
1798            // Get PostgresConnection for generating subsources.
1799            let pg_connection = &pg_source_connection.connection;
1800
1801            let config = pg_connection
1802                .config(
1803                    &storage_configuration.connection_context.secrets_reader,
1804                    storage_configuration,
1805                    InTask::No,
1806                )
1807                .await?;
1808
1809            let client = config
1810                .connect(
1811                    "postgres_purification",
1812                    &storage_configuration.connection_context.ssh_tunnel_manager,
1813                )
1814                .await?;
1815            let reference_client = SourceReferenceClient::Postgres {
1816                client: &client,
1817                publication: &pg_source_connection.publication,
1818                database: &pg_connection.database,
1819            };
1820            reference_client.get_source_references().await?
1821        }
1822        GenericSourceConnection::MySql(mysql_source_connection) => {
1823            let mysql_connection = &mysql_source_connection.connection;
1824            let config = mysql_connection
1825                .config(
1826                    &storage_configuration.connection_context.secrets_reader,
1827                    storage_configuration,
1828                    InTask::No,
1829                )
1830                .await?;
1831
1832            let mut conn = config
1833                .connect(
1834                    "mysql purification",
1835                    &storage_configuration.connection_context.ssh_tunnel_manager,
1836                )
1837                .await?;
1838
1839            let reference_client = SourceReferenceClient::MySql {
1840                conn: &mut conn,
1841                include_system_schemas: false,
1842            };
1843            reference_client.get_source_references().await?
1844        }
1845        GenericSourceConnection::SqlServer(sql_server_source) => {
1846            // Open a connection to the upstream SQL Server instance.
1847            let sql_server_connection = &sql_server_source.connection;
1848            let config = sql_server_connection
1849                .resolve_config(
1850                    &storage_configuration.connection_context.secrets_reader,
1851                    storage_configuration,
1852                    InTask::No,
1853                )
1854                .await?;
1855            let mut client = mz_sql_server_util::Client::connect(config).await?;
1856
1857            // Query the upstream SQL Server instance for available tables to replicate.
1858            let source_references = SourceReferenceClient::SqlServer {
1859                client: &mut client,
1860                database: sql_server_connection.database.clone().into(),
1861            }
1862            .get_source_references()
1863            .await?;
1864            source_references
1865        }
1866        GenericSourceConnection::LoadGenerator(load_gen_connection) => {
1867            let reference_client = SourceReferenceClient::LoadGenerator {
1868                generator: &load_gen_connection.load_generator,
1869            };
1870            reference_client.get_source_references().await?
1871        }
1872        GenericSourceConnection::Kafka(kafka_conn) => {
1873            let reference_client = SourceReferenceClient::Kafka {
1874                topic: &kafka_conn.topic,
1875            };
1876            reference_client.get_source_references().await?
1877        }
1878    };
1879    Ok(PurifiedStatement::PurifiedAlterSourceRefreshReferences {
1880        source_name: resolved_source_name,
1881        available_source_references: retrieved_source_references.available_source_references(),
1882    })
1883}
1884
1885#[derive(Clone, Debug, PartialEq, Eq)]
1886struct TableFromSourcePreparation {
1887    connection: GenericSourceConnection,
1888    requested_references: Option<ExternalReferences>,
1889    qualified_text_columns: Vec<UnresolvedItemName>,
1890    qualified_exclude_columns: Vec<UnresolvedItemName>,
1891    filter_constraints: FilterConstraints,
1892    exclude_all_not_null_constraints: bool,
1893    unresolved_source_name: UnresolvedItemName,
1894}
1895
1896async fn execute_table_from_source_preparation(
1897    request: TableFromSourcePreparation,
1898    storage_configuration: &StorageConfiguration,
1899) -> Result<(PurifiedSourceExport, SourceFormatOptions), PlanError> {
1900    let TableFromSourcePreparation {
1901        connection,
1902        requested_references,
1903        qualified_text_columns,
1904        qualified_exclude_columns,
1905        filter_constraints,
1906        exclude_all_not_null_constraints,
1907        unresolved_source_name,
1908    } = request;
1909    // Should be overriden below if a source-specific format is required.
1910    let mut format_options = SourceFormatOptions::Default;
1911
1912    let retrieved_source_references: RetrievedSourceReferences;
1913
1914    if (filter_constraints.excludes_any() || exclude_all_not_null_constraints)
1915        && !matches!(
1916            connection,
1917            GenericSourceConnection::Postgres(_) | GenericSourceConnection::MySql(_)
1918        )
1919    {
1920        sql_bail!(
1921            "EXCLUDE CONSTRAINTS is not supported for {} sources",
1922            connection.name()
1923        );
1924    }
1925
1926    // Run purification work specific to each source type: resolve the external reference to
1927    // a fully qualified name and obtain the appropriate details for the source-export statement
1928    let purified_export = match connection {
1929        GenericSourceConnection::Postgres(pg_source_connection) => {
1930            // Get PostgresConnection for generating subsources.
1931            let pg_connection = &pg_source_connection.connection;
1932
1933            let client = pg_connection
1934                .validate(pg_source_connection.connection_id, storage_configuration)
1935                .await?;
1936
1937            let reference_client = SourceReferenceClient::Postgres {
1938                client: &client,
1939                publication: &pg_source_connection.publication,
1940                database: &pg_connection.database,
1941            };
1942            retrieved_source_references = reference_client.get_source_references().await?;
1943
1944            // Read after the references above, so it bounds the LSN their schemas belong to.
1945            let initial_lsn = MzOffset::from(
1946                mz_postgres_util::fetch_max_lsn(
1947                    &client,
1948                    pg_source_connection
1949                        .publication_details
1950                        .get_is_physical_replica(),
1951                )
1952                .await?,
1953            );
1954
1955            let postgres::PurifiedSourceExports {
1956                source_exports,
1957                // TODO(database-issues#8620): Remove once subsources are removed
1958                // This `normalized_text_columns` is not relevant for us and is only returned for
1959                // `CREATE SOURCE` statements that automatically generate subsources
1960                normalized_text_columns: _,
1961            } = postgres::purify_source_exports(
1962                &client,
1963                &retrieved_source_references,
1964                &requested_references,
1965                qualified_text_columns,
1966                qualified_exclude_columns,
1967                &filter_constraints,
1968                exclude_all_not_null_constraints,
1969                &unresolved_source_name,
1970                &SourceReferencePolicy::Required,
1971                initial_lsn,
1972            )
1973            .await?;
1974            // There should be exactly one source_export returned for this statement
1975            let (_, purified_export) = source_exports.into_element();
1976            purified_export
1977        }
1978        GenericSourceConnection::MySql(mysql_source_connection) => {
1979            let mysql_connection = &mysql_source_connection.connection;
1980            let config = mysql_connection
1981                .config(
1982                    &storage_configuration.connection_context.secrets_reader,
1983                    storage_configuration,
1984                    InTask::No,
1985                )
1986                .await?;
1987
1988            let mut conn = config
1989                .connect(
1990                    "mysql purification",
1991                    &storage_configuration.connection_context.ssh_tunnel_manager,
1992                )
1993                .await?;
1994
1995            ensure_binlog_full_metadata(&mut conn).await?;
1996            let binlog_full_metadata = true;
1997
1998            // Retrieve the current @gtid_executed value of the server to mark as the effective
1999            // initial snapshot point for this table.
2000            let initial_gtid_set =
2001                mz_mysql_util::query_sys_var(&mut conn, "global.gtid_executed").await?;
2002
2003            let reference_client = SourceReferenceClient::MySql {
2004                conn: &mut conn,
2005                include_system_schemas: mysql::references_system_schemas(&requested_references),
2006            };
2007            retrieved_source_references = reference_client.get_source_references().await?;
2008
2009            let mysql::PurifiedSourceExports {
2010                source_exports,
2011                // TODO(database-issues#8620): Remove once subsources are removed
2012                // `normalized_text/exclude_columns` is not relevant for us and is only returned for
2013                // `CREATE SOURCE` statements that automatically generate subsources
2014                normalized_text_columns: _,
2015                normalized_exclude_columns: _,
2016            } = mysql::purify_source_exports(
2017                &mut conn,
2018                &retrieved_source_references,
2019                &requested_references,
2020                qualified_text_columns,
2021                qualified_exclude_columns,
2022                &filter_constraints,
2023                exclude_all_not_null_constraints,
2024                &unresolved_source_name,
2025                initial_gtid_set,
2026                &SourceReferencePolicy::Required,
2027                binlog_full_metadata,
2028            )
2029            .await?;
2030            // There should be exactly one source_export returned for this statement
2031            let (_, purified_export) = source_exports.into_element();
2032            purified_export
2033        }
2034        GenericSourceConnection::SqlServer(sql_server_source) => {
2035            let connection = sql_server_source.connection;
2036            let config = connection
2037                .resolve_config(
2038                    &storage_configuration.connection_context.secrets_reader,
2039                    storage_configuration,
2040                    InTask::No,
2041                )
2042                .await?;
2043            let mut client = mz_sql_server_util::Client::connect(config).await?;
2044
2045            let database: Arc<str> = connection.database.into();
2046            let reference_client = SourceReferenceClient::SqlServer {
2047                client: &mut client,
2048                database: Arc::clone(&database),
2049            };
2050            retrieved_source_references = reference_client.get_source_references().await?;
2051            tracing::debug!(?retrieved_source_references, "got source references");
2052
2053            let timeout = mz_storage_types::sources::sql_server::MAX_LSN_WAIT
2054                .get(storage_configuration.config_set());
2055
2056            let purified_source_exports = sql_server::purify_source_exports(
2057                &*database,
2058                &mut client,
2059                &retrieved_source_references,
2060                &requested_references,
2061                &qualified_text_columns,
2062                &qualified_exclude_columns,
2063                &unresolved_source_name,
2064                timeout,
2065                &SourceReferencePolicy::Required,
2066            )
2067            .await?;
2068
2069            // There should be exactly one source_export returned for this statement
2070            let (_, purified_export) = purified_source_exports.source_exports.into_element();
2071            purified_export
2072        }
2073        GenericSourceConnection::LoadGenerator(load_gen_connection) => {
2074            let reference_client = SourceReferenceClient::LoadGenerator {
2075                generator: &load_gen_connection.load_generator,
2076            };
2077            retrieved_source_references = reference_client.get_source_references().await?;
2078
2079            let requested_exports = retrieved_source_references
2080                .requested_source_exports(requested_references.as_ref(), &unresolved_source_name)?;
2081            // There should be exactly one source_export returned
2082            let export = requested_exports.into_element();
2083            PurifiedSourceExport {
2084                external_reference: export.external_reference,
2085                details: PurifiedExportDetails::LoadGenerator {
2086                    table: export
2087                        .meta
2088                        .load_generator_desc()
2089                        .ok_or_else(|| internal_err!("expected load generator source reference"))?
2090                        .clone(),
2091                    output: export
2092                        .meta
2093                        .load_generator_output()
2094                        .ok_or_else(|| internal_err!("expected load generator source reference"))?
2095                        .clone(),
2096                },
2097            }
2098        }
2099        GenericSourceConnection::Kafka(kafka_conn) => {
2100            let reference_client = SourceReferenceClient::Kafka {
2101                topic: &kafka_conn.topic,
2102            };
2103            retrieved_source_references = reference_client.get_source_references().await?;
2104            let requested_exports = retrieved_source_references
2105                .requested_source_exports(requested_references.as_ref(), &unresolved_source_name)?;
2106            // There should be exactly one source_export returned
2107            let export = requested_exports.into_element();
2108
2109            format_options = SourceFormatOptions::Kafka {
2110                topic: kafka_conn.topic.clone(),
2111            };
2112            PurifiedSourceExport {
2113                external_reference: export.external_reference,
2114                details: PurifiedExportDetails::Kafka {},
2115            }
2116        }
2117    };
2118
2119    Ok((purified_export, format_options))
2120}
2121
2122fn prepare_table_from_source(
2123    catalog: &dyn SessionCatalog,
2124    stmt: &CreateTableFromSourceStatement<Aug>,
2125) -> Result<TableFromSourcePreparation, PlanError> {
2126    let scx = StatementContext::new(None, catalog);
2127    let CreateTableFromSourceStatement {
2128        name: _,
2129        columns,
2130        constraints,
2131        source: source_name,
2132        if_not_exists: _,
2133        external_reference,
2134        format: _,
2135        envelope: _,
2136        include_metadata: _,
2137        with_options,
2138    } = stmt;
2139
2140    // Columns and constraints cannot be specified by the user but will be populated below.
2141    if matches!(columns, TableFromSourceColumns::Defined(_)) {
2142        sql_bail!("CREATE TABLE .. FROM SOURCE column definitions cannot be specified directly");
2143    }
2144    if !constraints.is_empty() {
2145        sql_bail!(
2146            "CREATE TABLE .. FROM SOURCE constraint definitions cannot be specified directly"
2147        );
2148    }
2149
2150    // Get the source item
2151    let item = match scx.get_item_by_resolved_name(source_name) {
2152        Ok(item) => item,
2153        Err(e) => return Err(e),
2154    };
2155
2156    // Ensure it's an ingestion-based and alterable source.
2157    let desc = match item.source_desc()? {
2158        Some(desc) => desc.clone().into_inline_connection(scx.catalog),
2159        None => {
2160            sql_bail!("cannot ALTER this type of source")
2161        }
2162    };
2163    let unresolved_source_name: UnresolvedItemName = source_name.full_item_name().clone().into();
2164
2165    let crate::plan::statement::ddl::TableFromSourceOptionExtracted {
2166        text_columns,
2167        exclude_columns,
2168        exclude_constraints,
2169        exclude_all_constraints,
2170        exclude_all_not_null_constraints,
2171        retain_history: _,
2172        details,
2173        partition_by: _,
2174        seen: _,
2175    } = with_options.clone().try_into()?;
2176    if details.is_some() {
2177        sql_bail!("DETAILS option cannot be explicitly set");
2178    }
2179
2180    let filter_constraints =
2181        FilterConstraints::from_options(exclude_constraints, exclude_all_constraints)?;
2182    if filter_constraints.excludes_any() || exclude_all_not_null_constraints {
2183        scx.require_feature_flag(&crate::session::vars::ENABLE_EXCLUDE_CONSTRAINTS_OPTION)?;
2184    }
2185
2186    // Our text column values are unqualified (just column names), but the purification methods below
2187    // expect to match the fully-qualified names against the full set of tables in upstream, so we
2188    // need to qualify them using the external reference first.
2189    let qualified_text_columns = text_columns
2190        .iter()
2191        .map(|col| {
2192            UnresolvedItemName(
2193                external_reference
2194                    .as_ref()
2195                    .map(|er| er.0.iter().chain_one(col).map(|i| i.clone()).collect())
2196                    .unwrap_or_else(|| vec![col.clone()]),
2197            )
2198        })
2199        .collect_vec();
2200    let qualified_exclude_columns = exclude_columns
2201        .iter()
2202        .map(|col| {
2203            UnresolvedItemName(
2204                external_reference
2205                    .as_ref()
2206                    .map(|er| er.0.iter().chain_one(col).map(|i| i.clone()).collect())
2207                    .unwrap_or_else(|| vec![col.clone()]),
2208            )
2209        })
2210        .collect_vec();
2211
2212    let requested_references = external_reference.as_ref().map(|ref_name| {
2213        ExternalReferences::SubsetTables(vec![ExternalReferenceExport {
2214            reference: ref_name.clone(),
2215            alias: None,
2216        }])
2217    });
2218
2219    Ok(TableFromSourcePreparation {
2220        connection: desc.connection,
2221        requested_references,
2222        qualified_text_columns,
2223        qualified_exclude_columns,
2224        filter_constraints,
2225        exclude_all_not_null_constraints,
2226        unresolved_source_name,
2227    })
2228}
2229
2230fn finish_table_from_source(
2231    catalog: &dyn SessionCatalog,
2232    mut stmt: CreateTableFromSourceStatement<Aug>,
2233    purified_export: PurifiedSourceExport,
2234) -> Result<PurifiedStatement, PlanError> {
2235    let scx = StatementContext::new(None, catalog);
2236    let CreateTableFromSourceStatement {
2237        columns,
2238        constraints,
2239        external_reference,
2240        with_options,
2241        ..
2242    } = &mut stmt;
2243    // Update the external reference in the statement to the resolved fully-qualified
2244    // external reference
2245    *external_reference = Some(purified_export.external_reference.clone());
2246
2247    // Update options in the statement using the purified export details
2248    match &purified_export.details {
2249        PurifiedExportDetails::Postgres { .. } => {
2250            let mut unsupported_cols = vec![];
2251            let postgres::PostgresExportStatementValues {
2252                columns: gen_columns,
2253                constraints: gen_constraints,
2254                text_columns: gen_text_columns,
2255                exclude_columns: gen_exclude_columns,
2256                details: gen_details,
2257                external_reference: _,
2258            } = postgres::generate_source_export_statement_values(
2259                &scx,
2260                purified_export,
2261                &mut unsupported_cols,
2262            )?;
2263            if !unsupported_cols.is_empty() {
2264                unsupported_cols.sort();
2265                Err(PgSourcePurificationError::UnrecognizedTypes {
2266                    cols: unsupported_cols,
2267                })?;
2268            }
2269
2270            if let Some(text_cols_option) = with_options
2271                .iter_mut()
2272                .find(|option| option.name == TableFromSourceOptionName::TextColumns)
2273            {
2274                if let Some(gen_text_columns) = gen_text_columns {
2275                    text_cols_option.value = Some(WithOptionValue::Sequence(gen_text_columns));
2276                }
2277            }
2278            if let Some(exclude_cols_option) = with_options
2279                .iter_mut()
2280                .find(|option| option.name == TableFromSourceOptionName::ExcludeColumns)
2281            {
2282                if let Some(gen_exclude_columns) = gen_exclude_columns {
2283                    exclude_cols_option.value =
2284                        Some(WithOptionValue::Sequence(gen_exclude_columns));
2285                }
2286            }
2287            match columns {
2288                TableFromSourceColumns::Defined(_) => {
2289                    // Purification is what populates the `Defined` variant;
2290                    // reaching this is an internal invariant violation.
2291                    bail_internal!(
2292                        "column definitions cannot be explicitly set for this source type"
2293                    )
2294                }
2295                TableFromSourceColumns::NotSpecified => {
2296                    *columns = TableFromSourceColumns::Defined(gen_columns);
2297                    *constraints = gen_constraints;
2298                }
2299                TableFromSourceColumns::Named(_) => {
2300                    sql_bail!("columns cannot be named for Postgres sources")
2301                }
2302            }
2303            with_options.push(TableFromSourceOption {
2304                name: TableFromSourceOptionName::Details,
2305                value: Some(WithOptionValue::Value(Value::String(hex::encode(
2306                    gen_details.into_proto().encode_to_vec(),
2307                )))),
2308            })
2309        }
2310        PurifiedExportDetails::MySql { .. } => {
2311            let mysql::MySqlExportStatementValues {
2312                columns: gen_columns,
2313                constraints: gen_constraints,
2314                text_columns: gen_text_columns,
2315                exclude_columns: gen_exclude_columns,
2316                details: gen_details,
2317                external_reference: _,
2318            } = mysql::generate_source_export_statement_values(&scx, purified_export)?;
2319
2320            if let Some(text_cols_option) = with_options
2321                .iter_mut()
2322                .find(|option| option.name == TableFromSourceOptionName::TextColumns)
2323            {
2324                if let Some(gen_text_columns) = gen_text_columns {
2325                    text_cols_option.value = Some(WithOptionValue::Sequence(gen_text_columns));
2326                }
2327            }
2328            if let Some(exclude_cols_option) = with_options
2329                .iter_mut()
2330                .find(|option| option.name == TableFromSourceOptionName::ExcludeColumns)
2331            {
2332                if let Some(gen_exclude_columns) = gen_exclude_columns {
2333                    exclude_cols_option.value =
2334                        Some(WithOptionValue::Sequence(gen_exclude_columns));
2335                }
2336            }
2337            match columns {
2338                TableFromSourceColumns::Defined(_) => {
2339                    // Purification is what populates the `Defined` variant;
2340                    // reaching this is an internal invariant violation.
2341                    bail_internal!(
2342                        "column definitions cannot be explicitly set for this source type"
2343                    )
2344                }
2345                TableFromSourceColumns::NotSpecified => {
2346                    *columns = TableFromSourceColumns::Defined(gen_columns);
2347                    *constraints = gen_constraints;
2348                }
2349                TableFromSourceColumns::Named(_) => {
2350                    sql_bail!("columns cannot be named for MySQL sources")
2351                }
2352            }
2353            with_options.push(TableFromSourceOption {
2354                name: TableFromSourceOptionName::Details,
2355                value: Some(WithOptionValue::Value(Value::String(hex::encode(
2356                    gen_details.into_proto().encode_to_vec(),
2357                )))),
2358            })
2359        }
2360        PurifiedExportDetails::SqlServer { .. } => {
2361            let sql_server::SqlServerExportStatementValues {
2362                columns: gen_columns,
2363                constraints: gen_constraints,
2364                text_columns: gen_text_columns,
2365                excl_columns: gen_excl_columns,
2366                details: gen_details,
2367                external_reference: _,
2368            } = sql_server::generate_source_export_statement_values(&scx, purified_export)?;
2369
2370            if let Some(text_cols_option) = with_options
2371                .iter_mut()
2372                .find(|opt| opt.name == TableFromSourceOptionName::TextColumns)
2373            {
2374                if let Some(gen_text_columns) = gen_text_columns {
2375                    text_cols_option.value = Some(WithOptionValue::Sequence(gen_text_columns));
2376                }
2377            }
2378            if let Some(exclude_cols_option) = with_options
2379                .iter_mut()
2380                .find(|opt| opt.name == TableFromSourceOptionName::ExcludeColumns)
2381            {
2382                if let Some(gen_excl_columns) = gen_excl_columns {
2383                    exclude_cols_option.value = Some(WithOptionValue::Sequence(gen_excl_columns));
2384                }
2385            }
2386
2387            match columns {
2388                TableFromSourceColumns::NotSpecified => {
2389                    *columns = TableFromSourceColumns::Defined(gen_columns);
2390                    *constraints = gen_constraints;
2391                }
2392                TableFromSourceColumns::Named(_) => {
2393                    sql_bail!("columns cannot be named for SQL Server sources")
2394                }
2395                TableFromSourceColumns::Defined(_) => {
2396                    // Purification is what populates the `Defined` variant;
2397                    // reaching this is an internal invariant violation.
2398                    bail_internal!(
2399                        "column definitions cannot be explicitly set for this source type"
2400                    )
2401                }
2402            }
2403
2404            with_options.push(TableFromSourceOption {
2405                name: TableFromSourceOptionName::Details,
2406                value: Some(WithOptionValue::Value(Value::String(hex::encode(
2407                    gen_details.into_proto().encode_to_vec(),
2408                )))),
2409            })
2410        }
2411        PurifiedExportDetails::LoadGenerator { .. } => {
2412            let (desc, output) = match purified_export.details {
2413                PurifiedExportDetails::LoadGenerator { table, output } => (table, output),
2414                _ => bail_internal!("purified export details must be load generator"),
2415            };
2416            // We only determine the table description for multi-output load generator sources here,
2417            // whereas single-output load generators will have their relation description
2418            // determined during statment planning as envelope and format options may affect their
2419            // schema.
2420            if let Some(desc) = desc {
2421                let (gen_columns, gen_constraints) = scx.relation_desc_into_table_defs(&desc)?;
2422                match columns {
2423                    // Purification is what populates the `Defined` variant;
2424                    // reaching this is an internal invariant violation.
2425                    TableFromSourceColumns::Defined(_) => bail_internal!(
2426                        "column definitions cannot be explicitly set for this source type"
2427                    ),
2428                    TableFromSourceColumns::NotSpecified => {
2429                        *columns = TableFromSourceColumns::Defined(gen_columns);
2430                        *constraints = gen_constraints;
2431                    }
2432                    TableFromSourceColumns::Named(_) => {
2433                        sql_bail!("columns cannot be named for multi-output load generator sources")
2434                    }
2435                }
2436            }
2437            let details = SourceExportStatementDetails::LoadGenerator { output };
2438            with_options.push(TableFromSourceOption {
2439                name: TableFromSourceOptionName::Details,
2440                value: Some(WithOptionValue::Value(Value::String(hex::encode(
2441                    details.into_proto().encode_to_vec(),
2442                )))),
2443            })
2444        }
2445        PurifiedExportDetails::Kafka {} => {
2446            // NOTE: Kafka tables have their 'schemas' purified into the statement inside the
2447            // format field, so we don't specify any columns or constraints to be stored
2448            // on the statement here. The RelationDesc will be determined during planning.
2449            let details = SourceExportStatementDetails::Kafka {};
2450            with_options.push(TableFromSourceOption {
2451                name: TableFromSourceOptionName::Details,
2452                value: Some(WithOptionValue::Value(Value::String(hex::encode(
2453                    details.into_proto().encode_to_vec(),
2454                )))),
2455            })
2456        }
2457    };
2458
2459    // TODO: We might as well use the retrieved available references to update the source
2460    // available references table in the catalog, so plumb this through.
2461    // available_source_references: retrieved_source_references.available_source_references(),
2462    Ok(PurifiedStatement::PurifiedCreateTableFromSource { stmt })
2463}
2464
2465#[derive(Clone, Debug)]
2466enum SourceFormatOptions {
2467    Default,
2468    Kafka { topic: String },
2469}
2470
2471/// Generates a progress subsource, choosing an available name in the supplied catalog.
2472pub fn generate_progress_subsource_statement(
2473    scx: &StatementContext,
2474    source: &CreateSourceStatement<Aug>,
2475) -> Result<CreateSubsourceStatement<Aug>, PlanError> {
2476    let CreateSourceStatement {
2477        name: source_name,
2478        connection: source_connection,
2479        progress_subsource,
2480        with_options,
2481        ..
2482    } = source;
2483
2484    let progress_desc = match source_connection {
2485        CreateSourceConnection::Kafka { .. } => {
2486            &mz_storage_types::sources::kafka::KAFKA_PROGRESS_DESC
2487        }
2488        CreateSourceConnection::Postgres { .. } => {
2489            &mz_storage_types::sources::postgres::PG_PROGRESS_DESC
2490        }
2491        CreateSourceConnection::SqlServer { .. } => {
2492            &mz_storage_types::sources::sql_server::SQL_SERVER_PROGRESS_DESC
2493        }
2494        CreateSourceConnection::MySql { .. } => {
2495            &mz_storage_types::sources::mysql::MYSQL_PROGRESS_DESC
2496        }
2497        CreateSourceConnection::LoadGenerator { .. } => {
2498            &mz_storage_types::sources::load_generator::LOAD_GEN_PROGRESS_DESC
2499        }
2500    };
2501
2502    let name = match progress_subsource {
2503        Some(name) => match name {
2504            DeferredItemName::Deferred(name) => name.clone(),
2505            DeferredItemName::Named(_) => {
2506                sql_bail!("progress subsource name cannot be a resolved name")
2507            }
2508        },
2509        None => {
2510            let (item, prefix) = source_name
2511                .0
2512                .split_last()
2513                .ok_or_else(|| sql_err!("source name must have at least one component"))?;
2514            let item_name = Ident::try_generate_name(item.to_string(), "_progress", |candidate| {
2515                let mut suggested_name = prefix.to_vec();
2516                suggested_name.push(candidate.clone());
2517
2518                let partial = normalize::unresolved_item_name(UnresolvedItemName(suggested_name))?;
2519                let qualified = scx.allocate_qualified_name(partial)?;
2520                let item_exists = scx.catalog.get_item_by_name(&qualified).is_some();
2521                let type_exists = scx.catalog.get_type_by_name(&qualified).is_some();
2522                Ok::<_, PlanError>(!item_exists && !type_exists)
2523            })?;
2524
2525            let mut full_name = prefix.to_vec();
2526            full_name.push(item_name);
2527            let full_name = normalize::unresolved_item_name(UnresolvedItemName(full_name))?;
2528            let qualified_name = scx.allocate_qualified_name(full_name)?;
2529            let full_name = scx.catalog.resolve_full_name(&qualified_name);
2530
2531            UnresolvedItemName::from(full_name.clone())
2532        }
2533    };
2534
2535    let (columns, constraints) = scx.relation_desc_into_table_defs(progress_desc)?;
2536
2537    let mut progress_with_options: Vec<_> = with_options
2538        .iter()
2539        .filter_map(|opt| match opt.name {
2540            CreateSourceOptionName::TimestampInterval => None,
2541            CreateSourceOptionName::RetainHistory => Some(CreateSubsourceOption {
2542                name: CreateSubsourceOptionName::RetainHistory,
2543                value: opt.value.clone(),
2544            }),
2545        })
2546        .collect();
2547    progress_with_options.push(CreateSubsourceOption {
2548        name: CreateSubsourceOptionName::Progress,
2549        value: Some(WithOptionValue::Value(Value::Boolean(true))),
2550    });
2551
2552    Ok(CreateSubsourceStatement {
2553        name,
2554        columns,
2555        // Progress subsources do not refer to the source to which they belong.
2556        // Instead the primary source depends on it (the opposite is true of
2557        // ingestion exports, which depend on the primary source).
2558        of_source: None,
2559        constraints,
2560        if_not_exists: false,
2561        with_options: progress_with_options,
2562    })
2563}
2564
2565pub fn generate_subsource_statements(
2566    scx: &StatementContext,
2567    source_name: ResolvedItemName,
2568    subsources: BTreeMap<UnresolvedItemName, PurifiedSourceExport>,
2569) -> Result<Vec<CreateSubsourceStatement<Aug>>, PlanError> {
2570    // get the first subsource to determine the connection type
2571    if subsources.is_empty() {
2572        return Ok(vec![]);
2573    }
2574    let (_, purified_export) = subsources
2575        .iter()
2576        .next()
2577        .ok_or_else(|| internal_err!("expected at least one subsource"))?;
2578
2579    let statements = match &purified_export.details {
2580        PurifiedExportDetails::Postgres { .. } => {
2581            crate::pure::postgres::generate_create_subsource_statements(
2582                scx,
2583                source_name,
2584                subsources,
2585            )?
2586        }
2587        PurifiedExportDetails::MySql { .. } => {
2588            crate::pure::mysql::generate_create_subsource_statements(scx, source_name, subsources)?
2589        }
2590        PurifiedExportDetails::SqlServer { .. } => {
2591            crate::pure::sql_server::generate_create_subsource_statements(
2592                scx,
2593                source_name,
2594                subsources,
2595            )?
2596        }
2597        PurifiedExportDetails::LoadGenerator { .. } => {
2598            let mut subsource_stmts = Vec::with_capacity(subsources.len());
2599            for (subsource_name, purified_export) in subsources {
2600                let (desc, output) = match purified_export.details {
2601                    PurifiedExportDetails::LoadGenerator { table, output } => (table, output),
2602                    _ => {
2603                        bail_internal!("purified export details must be load generator")
2604                    }
2605                };
2606                let desc = desc.ok_or_else(|| {
2607                    internal_err!(
2608                        "subsources cannot be generated for single-output load generators"
2609                    )
2610                })?;
2611
2612                let (columns, table_constraints) = scx.relation_desc_into_table_defs(&desc)?;
2613                let details = SourceExportStatementDetails::LoadGenerator { output };
2614                // Create the subsource statement
2615                let subsource = CreateSubsourceStatement {
2616                    name: subsource_name,
2617                    columns,
2618                    of_source: Some(source_name.clone()),
2619                    // unlike sources that come from an external upstream, we
2620                    // have more leniency to introduce different constraints
2621                    // every time the load generator is run; i.e. we are not as
2622                    // worried about introducing junk data.
2623                    constraints: table_constraints,
2624                    if_not_exists: false,
2625                    with_options: vec![
2626                        CreateSubsourceOption {
2627                            name: CreateSubsourceOptionName::ExternalReference,
2628                            value: Some(WithOptionValue::UnresolvedItemName(
2629                                purified_export.external_reference,
2630                            )),
2631                        },
2632                        CreateSubsourceOption {
2633                            name: CreateSubsourceOptionName::Details,
2634                            value: Some(WithOptionValue::Value(Value::String(hex::encode(
2635                                details.into_proto().encode_to_vec(),
2636                            )))),
2637                        },
2638                    ],
2639                };
2640                subsource_stmts.push(subsource);
2641            }
2642
2643            subsource_stmts
2644        }
2645        PurifiedExportDetails::Kafka { .. } => {
2646            // TODO: as part of database-issues#8322, Kafka sources will begin
2647            // producing data––we'll need to understand the schema
2648            // of the output here.
2649            if !subsources.is_empty() {
2650                bail_internal!("Kafka sources do not produce data-bearing subsources");
2651            }
2652            vec![]
2653        }
2654    };
2655    Ok(statements)
2656}
2657
2658#[derive(Clone, Debug, PartialEq, Eq)]
2659struct CsrProtobufPreparation {
2660    connection: mz_storage_types::connections::CsrConnection,
2661    topic: String,
2662    require_key: bool,
2663}
2664
2665impl CsrProtobufPreparation {
2666    async fn execute(
2667        self,
2668        storage_configuration: &StorageConfiguration,
2669    ) -> Result<CsrSeedProtobuf, PlanError> {
2670        let Self {
2671            connection: ccsr_connection,
2672            topic,
2673            require_key,
2674        } = self;
2675        let ccsr_client = ccsr_connection
2676            .connect(storage_configuration, InTask::No)
2677            .await
2678            .map_err(|e| CsrPurificationError::ClientError(Arc::new(e)))?;
2679
2680        let value = compile_proto(&format!("{}-value", topic), &ccsr_client).await?;
2681        let key = compile_proto(&format!("{}-key", topic), &ccsr_client)
2682            .await
2683            .ok();
2684
2685        if require_key && key.is_none() {
2686            sql_bail!("Key schema is required for ENVELOPE DEBEZIUM");
2687        }
2688
2689        Ok(CsrSeedProtobuf { value, key })
2690    }
2691}
2692
2693#[derive(Clone, Debug, PartialEq, Eq)]
2694struct CsrAvroPreparation {
2695    connection: mz_storage_types::connections::CsrConnection,
2696    topic: String,
2697    key_strategy: ReaderSchemaSelectionStrategy,
2698    value_strategy: ReaderSchemaSelectionStrategy,
2699    require_key: bool,
2700}
2701
2702impl CsrAvroPreparation {
2703    async fn execute(
2704        self,
2705        storage_configuration: &StorageConfiguration,
2706    ) -> Result<CsrSeedAvro, PlanError> {
2707        let Self {
2708            connection: csr_connection,
2709            topic,
2710            key_strategy,
2711            value_strategy,
2712            require_key,
2713        } = self;
2714        let ccsr_client = csr_connection
2715            .connect(storage_configuration, InTask::No)
2716            .await
2717            .map_err(|e| CsrPurificationError::ClientError(Arc::new(e)))?;
2718
2719        let Schema {
2720            key_schema,
2721            value_schema,
2722            key_reference_schemas,
2723            value_reference_schemas,
2724        } = get_remote_csr_schema(&ccsr_client, key_strategy, value_strategy, &topic).await?;
2725        if require_key && key_schema.is_none() {
2726            sql_bail!("Key schema is required for ENVELOPE DEBEZIUM");
2727        }
2728
2729        Ok(CsrSeedAvro {
2730            key_schema,
2731            value_schema,
2732            key_reference_schemas,
2733            value_reference_schemas,
2734        })
2735    }
2736}
2737
2738#[derive(Clone, Debug, PartialEq, Eq)]
2739struct GlueAvroPreparation {
2740    connection: mz_storage_types::connections::GlueSchemaRegistryConnection,
2741    schema_name: String,
2742}
2743
2744impl GlueAvroPreparation {
2745    async fn execute(
2746        self,
2747        storage_configuration: &StorageConfiguration,
2748    ) -> Result<GlueAvroSeed, PlanError> {
2749        use crate::pure::error::GluePurificationError;
2750        let Self {
2751            connection: gsr_connection,
2752            schema_name,
2753        } = self;
2754        // Use storage decoding's authentication and address policy.
2755        let enforce_external_addresses = mz_storage_types::dyncfgs::ENFORCE_EXTERNAL_ADDRESSES
2756            .get(storage_configuration.config_set());
2757        let sdk_config = gsr_connection
2758            .aws_connection
2759            .connection
2760            .load_sdk_config(
2761                &storage_configuration.connection_context,
2762                gsr_connection.aws_connection.connection_id,
2763                // We are in a normal tokio context during purification.
2764                InTask::No,
2765                enforce_external_addresses,
2766            )
2767            .await
2768            .map_err(|e| GluePurificationError::LoadSdkConfigError(Arc::new(e)))?;
2769        let glue_client = mz_aws_glue_schema_registry::ClientConfig::new(sdk_config).build();
2770
2771        let version = glue_client
2772            .get_schema_version_latest_by_name(&gsr_connection.registry_name, &schema_name)
2773            .await
2774            .map_err(|e| GluePurificationError::SchemaLookupError {
2775                registry: gsr_connection.registry_name.clone(),
2776                schema: schema_name.clone(),
2777                cause: Arc::new(e),
2778            })?;
2779        // The runtime decode path only handles Avro; reject other formats at
2780        // planning time so the failure mode is a clear SQL error, not a
2781        // permanent decode error on every record.
2782        match &version.data_format {
2783            Some(mz_aws_glue_schema_registry::DataFormat::Avro) => {}
2784            other => {
2785                return Err(GluePurificationError::UnsupportedDataFormat {
2786                    registry: gsr_connection.registry_name.clone(),
2787                    schema: schema_name.clone(),
2788                    format: other
2789                        .as_ref()
2790                        .map(|f| f.as_str().to_string())
2791                        .unwrap_or_else(|| "<unspecified>".to_string()),
2792                }
2793                .into());
2794            }
2795        }
2796        let value_schema =
2797            version
2798                .definition
2799                .ok_or_else(|| GluePurificationError::EmptyDefinition {
2800                    registry: gsr_connection.registry_name.clone(),
2801                    schema: schema_name.clone(),
2802                })?;
2803
2804        Ok(GlueAvroSeed { value_schema })
2805    }
2806}
2807
2808fn prepare_csr_connection_proto(
2809    catalog: &dyn SessionCatalog,
2810    options: &SourceFormatOptions,
2811    csr_connection: &CsrConnectionProtobuf<Aug>,
2812    envelope: &Option<SourceEnvelope>,
2813) -> Result<Option<CsrProtobufPreparation>, PlanError> {
2814    let SourceFormatOptions::Kafka { topic } = options else {
2815        sql_bail!("Confluent Schema Registry is only supported with Kafka sources")
2816    };
2817
2818    let CsrConnectionProtobuf {
2819        seed,
2820        connection: CsrConnection {
2821            connection,
2822            options: _,
2823        },
2824    } = csr_connection;
2825    match seed {
2826        None => {
2827            let scx = StatementContext::new(None, &*catalog);
2828
2829            let ccsr_connection = match scx.get_item_by_resolved_name(connection)?.connection()? {
2830                Connection::Csr(connection) => connection.clone().into_inline_connection(catalog),
2831                _ => sql_bail!("{} is not a schema registry connection", connection),
2832            };
2833
2834            return Ok(Some(CsrProtobufPreparation {
2835                connection: ccsr_connection,
2836                topic: topic.clone(),
2837                require_key: matches!(envelope, Some(SourceEnvelope::Debezium)),
2838            }));
2839        }
2840        Some(_) => (),
2841    }
2842
2843    Ok(None)
2844}
2845
2846fn prepare_csr_connection_avro(
2847    catalog: &dyn SessionCatalog,
2848    options: &SourceFormatOptions,
2849    csr_connection: &CsrConnectionAvro<Aug>,
2850    envelope: &Option<SourceEnvelope>,
2851) -> Result<Option<CsrAvroPreparation>, PlanError> {
2852    let SourceFormatOptions::Kafka { topic } = options else {
2853        sql_bail!("Confluent Schema Registry is only supported with Kafka sources")
2854    };
2855
2856    let CsrConnectionAvro {
2857        connection: CsrConnection { connection, .. },
2858        seed,
2859        key_strategy,
2860        value_strategy,
2861    } = csr_connection;
2862    if seed.is_none() {
2863        let scx = StatementContext::new(None, &*catalog);
2864        let csr_connection = match scx.get_item_by_resolved_name(connection)?.connection()? {
2865            Connection::Csr(connection) => connection.clone().into_inline_connection(catalog),
2866            _ => sql_bail!("{} is not a schema registry connection", connection),
2867        };
2868        return Ok(Some(CsrAvroPreparation {
2869            connection: csr_connection,
2870            topic: topic.clone(),
2871            key_strategy: key_strategy.clone().unwrap_or_default(),
2872            value_strategy: value_strategy.clone().unwrap_or_default(),
2873            require_key: matches!(envelope, Some(SourceEnvelope::Debezium)),
2874        }));
2875    }
2876
2877    Ok(None)
2878}
2879
2880fn prepare_glue_connection_avro(
2881    catalog: &dyn SessionCatalog,
2882    options: &SourceFormatOptions,
2883    connection: &ResolvedItemName,
2884    with_options: &[GlueAvroOption<Aug>],
2885    seed: &Option<GlueAvroSeed>,
2886) -> Result<Option<GlueAvroPreparation>, PlanError> {
2887    use crate::pure::error::GluePurificationError;
2888    let SourceFormatOptions::Kafka { .. } = options else {
2889        sql_bail!("AWS Glue Schema Registry is only supported with Kafka sources")
2890    };
2891
2892    let scx = StatementContext::new(None, &*catalog);
2893    let item = scx.get_item_by_resolved_name(connection)?;
2894    let full_name = scx.catalog.resolve_full_name(item.name());
2895    // Structural validation below runs unconditionally — even when a seed is
2896    // already present. The seed is not proof the rest of the statement is
2897    // well-formed: the grammar accepts a user-written `SEED VALUE SCHEMA`, so a
2898    // seeded statement may have arrived straight from a user rather than from
2899    // our own re-parsed `create_sql`. Only the Glue *fetch* is safe to skip
2900    // when seeded; the connection-type and `SCHEMA NAME` checks are not.
2901    let Connection::GlueSchemaRegistry(gsr_connection) = item.connection()? else {
2902        return Err(GluePurificationError::NotGlueConnection(full_name).into());
2903    };
2904
2905    // Pull `SCHEMA NAME` out of the option bag. Required. Use the shared
2906    // extractor (rather than a hand-rolled match) so non-string values get a
2907    // proper "invalid value" error instead of being silently dropped.
2908    let crate::plan::statement::ddl::GlueAvroOptionExtracted {
2909        schema_name,
2910        key_schema_name,
2911        value_schema_name,
2912        key_compatibility_level,
2913        value_compatibility_level,
2914        seen: _,
2915    } = with_options.to_vec().try_into()?;
2916    // The per-side and compatibility options are sink-only. A source reads a
2917    // single writer schema, so reject them here rather than silently ignoring.
2918    if key_schema_name.is_some()
2919        || value_schema_name.is_some()
2920        || key_compatibility_level.is_some()
2921        || value_compatibility_level.is_some()
2922    {
2923        sql_bail!(
2924            "KEY SCHEMA NAME, VALUE SCHEMA NAME, KEY COMPATIBILITY LEVEL, and VALUE \
2925             COMPATIBILITY LEVEL are not supported for AWS Glue Schema Registry sources, \
2926             use SCHEMA NAME instead"
2927        );
2928    }
2929    let schema_name = schema_name.ok_or(GluePurificationError::MissingSchemaName)?;
2930
2931    if seed.is_some() {
2932        // The schema is already pinned — either we purified this statement
2933        // before and re-parsed it from persisted create_sql, or a user supplied
2934        // the seed directly. Either way the reader schema is fixed, so skip the
2935        // Glue lookup. (The structural checks above still ran.)
2936        return Ok(None);
2937    }
2938    let gsr_connection = gsr_connection.into_inline_connection(catalog);
2939
2940    Ok(Some(GlueAvroPreparation {
2941        connection: gsr_connection,
2942        schema_name,
2943    }))
2944}
2945
2946#[derive(Debug)]
2947pub struct Schema {
2948    pub key_schema: Option<String>,
2949    pub value_schema: String,
2950    /// Reference schemas for the key schema, in dependency order.
2951    pub key_reference_schemas: Vec<String>,
2952    /// Reference schemas for the value schema, in dependency order.
2953    pub value_reference_schemas: Vec<String>,
2954}
2955
2956/// Result of fetching a schema, including any referenced schemas.
2957struct SchemaWithReferences {
2958    /// The primary schema.
2959    schema: String,
2960    /// Reference schemas in dependency order.
2961    references: Vec<String>,
2962}
2963
2964async fn get_schema_with_strategy(
2965    client: &Client,
2966    strategy: ReaderSchemaSelectionStrategy,
2967    subject: &str,
2968) -> Result<Option<SchemaWithReferences>, PlanError> {
2969    match strategy {
2970        ReaderSchemaSelectionStrategy::Latest => {
2971            // Use get_subject_and_references to also fetch referenced schemas
2972            match client.get_subject_and_references(subject).await {
2973                Ok((primary, dependencies)) => Ok(Some(SchemaWithReferences {
2974                    schema: primary.schema.raw,
2975                    references: dependencies.into_iter().map(|s| s.schema.raw).collect(),
2976                })),
2977                Err(GetBySubjectError::SubjectNotFound)
2978                | Err(GetBySubjectError::VersionNotFound(_)) => Ok(None),
2979                Err(e) => Err(PlanError::FetchingCsrSchemaFailed {
2980                    schema_lookup: format!("subject {}", subject.quoted()),
2981                    cause: Arc::new(e),
2982                }),
2983            }
2984        }
2985        // It's possible that CSR was provided with an inline strategy, but without a subject
2986        // or schema id to look up, there isn't a clean way to lookup references for the schema.
2987        // This could be done with extra work, but at this point, it isn't clear this is needed.
2988        ReaderSchemaSelectionStrategy::Inline(raw) => Ok(Some(SchemaWithReferences {
2989            schema: raw,
2990            references: vec![],
2991        })),
2992        ReaderSchemaSelectionStrategy::ById(id) => {
2993            match client.get_subject_and_references_by_id(id).await {
2994                Ok((primary, dependencies)) => Ok(Some(SchemaWithReferences {
2995                    schema: primary.schema.raw,
2996                    references: dependencies.into_iter().map(|s| s.schema.raw).collect(),
2997                })),
2998                Err(GetBySubjectError::SubjectNotFound)
2999                | Err(GetBySubjectError::VersionNotFound(_)) => Ok(None),
3000                Err(e) => Err(PlanError::FetchingCsrSchemaFailed {
3001                    schema_lookup: format!("subject {}", subject.quoted()),
3002                    cause: Arc::new(e),
3003                }),
3004            }
3005        }
3006    }
3007}
3008
3009async fn get_remote_csr_schema(
3010    ccsr_client: &mz_ccsr::Client,
3011    key_strategy: ReaderSchemaSelectionStrategy,
3012    value_strategy: ReaderSchemaSelectionStrategy,
3013    topic: &str,
3014) -> Result<Schema, PlanError> {
3015    let value_schema_name = format!("{}-value", topic);
3016    let value_result =
3017        get_schema_with_strategy(ccsr_client, value_strategy, &value_schema_name).await?;
3018    let value_result = value_result.ok_or_else(|| anyhow!("No value schema found"))?;
3019
3020    let key_subject = format!("{}-key", topic);
3021    let key_result = get_schema_with_strategy(ccsr_client, key_strategy, &key_subject).await?;
3022    Ok(Schema {
3023        key_schema: key_result.as_ref().map(|r| r.schema.clone()),
3024        value_schema: value_result.schema,
3025        key_reference_schemas: key_result.map(|r| r.references).unwrap_or_default(),
3026        value_reference_schemas: value_result.references,
3027    })
3028}
3029
3030/// Collect protobuf message descriptor from CSR and compile the descriptor.
3031async fn compile_proto(
3032    subject_name: &String,
3033    ccsr_client: &Client,
3034) -> Result<CsrSeedProtobufSchema, PlanError> {
3035    let (primary_subject, dependency_subjects) = ccsr_client
3036        .get_subject_and_references(subject_name)
3037        .await
3038        .map_err(|e| PlanError::FetchingCsrSchemaFailed {
3039            schema_lookup: format!("subject {}", subject_name.quoted()),
3040            cause: Arc::new(e),
3041        })?;
3042
3043    // Compile .proto files into a file descriptor set.
3044    let mut source_tree = VirtualSourceTree::new();
3045
3046    // Add well-known types (e.g., google/protobuf/timestamp.proto) to the source
3047    // tree. These are implicitly available to protoc but are typically not
3048    // registered in the schema registry.
3049    source_tree.as_mut().map_well_known_types();
3050
3051    for subject in iter::once(&primary_subject).chain(dependency_subjects.iter()) {
3052        source_tree.as_mut().add_file(
3053            Path::new(&subject.name),
3054            subject.schema.raw.as_bytes().to_vec(),
3055        );
3056    }
3057    let mut db = SourceTreeDescriptorDatabase::new(source_tree.as_mut());
3058    let fds = db
3059        .as_mut()
3060        .build_file_descriptor_set(&[Path::new(&primary_subject.name)])
3061        .map_err(|cause| PlanError::InvalidProtobufSchema { cause })?;
3062
3063    // Ensure there is exactly one message in the file.
3064    let primary_fd = fds.file(0);
3065    let message_name = match primary_fd.message_type_size() {
3066        1 => String::from_utf8_lossy(primary_fd.message_type(0).name()).into_owned(),
3067        0 => bail_unsupported!(29603, "Protobuf schemas with no messages"),
3068        _ => bail_unsupported!(29603, "Protobuf schemas with multiple messages"),
3069    };
3070
3071    // Encode the file descriptor set into a SQL byte string.
3072    let bytes = &fds
3073        .serialize()
3074        .map_err(|cause| PlanError::InvalidProtobufSchema { cause })?;
3075    let mut schema = String::new();
3076    strconv::format_bytes(&mut schema, bytes);
3077
3078    Ok(CsrSeedProtobufSchema {
3079        schema,
3080        message_name,
3081    })
3082}
3083
3084const MZ_NOW_NAME: &str = "mz_now";
3085const MZ_NOW_SCHEMA: &str = "mz_catalog";
3086
3087/// Purifies a CREATE MATERIALIZED VIEW statement. Additionally, it adjusts `resolved_ids` if
3088/// references to ids appear or disappear during the purification.
3089///
3090/// Note that in contrast with [`purify_statement`], this doesn't need to be async, because
3091/// this function is not making any network calls.
3092pub fn purify_create_materialized_view_options(
3093    catalog: impl SessionCatalog,
3094    mz_now: Option<Timestamp>,
3095    cmvs: &mut CreateMaterializedViewStatement<Aug>,
3096    resolved_ids: &mut ResolvedIds,
3097) {
3098    // 0. Preparations:
3099    // Prepare an expression that calls `mz_now()`, which we can insert in various later steps.
3100    let (mz_now_id, mz_now_expr) = {
3101        let item = catalog
3102            .resolve_function(&PartialItemName {
3103                database: None,
3104                schema: Some(MZ_NOW_SCHEMA.to_string()),
3105                item: MZ_NOW_NAME.to_string(),
3106            })
3107            .expect("we should be able to resolve mz_now");
3108        (
3109            item.id(),
3110            Expr::Function(Function {
3111                name: ResolvedItemName::Item {
3112                    id: item.id(),
3113                    qualifiers: item.name().qualifiers.clone(),
3114                    full_name: catalog.resolve_full_name(item.name()),
3115                    print_id: false,
3116                    version: RelationVersionSelector::Latest,
3117                },
3118                args: FunctionArgs::Args {
3119                    args: Vec::new(),
3120                    order_by: Vec::new(),
3121                },
3122                filter: None,
3123                over: None,
3124                distinct: false,
3125            }),
3126        )
3127    };
3128    // Prepare the `mz_timestamp` type.
3129    let (mz_timestamp_id, mz_timestamp_type) = {
3130        let item = catalog.get_system_type("mz_timestamp");
3131        let full_name = catalog.resolve_full_name(item.name());
3132        (
3133            item.id(),
3134            ResolvedDataType::Named {
3135                id: item.id(),
3136                qualifiers: item.name().qualifiers.clone(),
3137                full_name,
3138                modifiers: vec![],
3139                print_id: true,
3140            },
3141        )
3142    };
3143
3144    let mut introduced_mz_timestamp = false;
3145
3146    for option in cmvs.with_options.iter_mut() {
3147        // 1. Purify `REFRESH AT CREATION` to `REFRESH AT mz_now()`.
3148        if matches!(
3149            option.value,
3150            Some(WithOptionValue::Refresh(RefreshOptionValue::AtCreation))
3151        ) {
3152            option.value = Some(WithOptionValue::Refresh(RefreshOptionValue::At(
3153                RefreshAtOptionValue {
3154                    time: mz_now_expr.clone(),
3155                },
3156            )));
3157        }
3158
3159        // 2. If `REFRESH EVERY` doesn't have an `ALIGNED TO`, then add `ALIGNED TO mz_now()`.
3160        if let Some(WithOptionValue::Refresh(RefreshOptionValue::Every(
3161            RefreshEveryOptionValue { aligned_to, .. },
3162        ))) = &mut option.value
3163        {
3164            if aligned_to.is_none() {
3165                *aligned_to = Some(mz_now_expr.clone());
3166            }
3167        }
3168
3169        // 3. Substitute `mz_now()` with the timestamp chosen for the CREATE MATERIALIZED VIEW
3170        // statement. (This has to happen after the above steps, which might introduce `mz_now()`.)
3171        match &mut option.value {
3172            Some(WithOptionValue::Refresh(RefreshOptionValue::At(RefreshAtOptionValue {
3173                time,
3174            }))) => {
3175                let mut visitor = MzNowPurifierVisitor::new(mz_now, mz_timestamp_type.clone());
3176                visitor.visit_expr_mut(time);
3177                introduced_mz_timestamp |= visitor.introduced_mz_timestamp;
3178            }
3179            Some(WithOptionValue::Refresh(RefreshOptionValue::Every(
3180                RefreshEveryOptionValue {
3181                    interval: _,
3182                    aligned_to: Some(aligned_to),
3183                },
3184            ))) => {
3185                let mut visitor = MzNowPurifierVisitor::new(mz_now, mz_timestamp_type.clone());
3186                visitor.visit_expr_mut(aligned_to);
3187                introduced_mz_timestamp |= visitor.introduced_mz_timestamp;
3188            }
3189            _ => {}
3190        }
3191    }
3192
3193    // 4. If the user didn't give any REFRESH option, then default to ON COMMIT.
3194    if !cmvs.with_options.iter().any(|o| {
3195        matches!(
3196            o,
3197            MaterializedViewOption {
3198                value: Some(WithOptionValue::Refresh(..)),
3199                ..
3200            }
3201        )
3202    }) {
3203        cmvs.with_options.push(MaterializedViewOption {
3204            name: MaterializedViewOptionName::Refresh,
3205            value: Some(WithOptionValue::Refresh(RefreshOptionValue::OnCommit)),
3206        })
3207    }
3208
3209    // 5. Attend to `resolved_ids`: The purification might have
3210    // - added references to `mz_timestamp`;
3211    // - removed references to `mz_now`.
3212    if introduced_mz_timestamp {
3213        resolved_ids.add_item(mz_timestamp_id);
3214    }
3215    // Even though we always remove `mz_now()` from the `with_options`, there might be `mz_now()`
3216    // remaining in the main query expression of the MV, so let's visit the entire statement to look
3217    // for `mz_now()` everywhere.
3218    let mut visitor = ExprContainsTemporalVisitor::new();
3219    visitor.visit_create_materialized_view_statement(cmvs);
3220    if !visitor.contains_temporal {
3221        resolved_ids.remove_item(&mz_now_id);
3222    }
3223}
3224
3225/// Returns true if the [MaterializedViewOption] either already involves `mz_now()` or will involve
3226/// after purification.
3227pub fn materialized_view_option_contains_temporal(mvo: &MaterializedViewOption<Aug>) -> bool {
3228    match &mvo.value {
3229        Some(WithOptionValue::Refresh(RefreshOptionValue::At(RefreshAtOptionValue { time }))) => {
3230            let mut visitor = ExprContainsTemporalVisitor::new();
3231            visitor.visit_expr(time);
3232            visitor.contains_temporal
3233        }
3234        Some(WithOptionValue::Refresh(RefreshOptionValue::Every(RefreshEveryOptionValue {
3235            interval: _,
3236            aligned_to: Some(aligned_to),
3237        }))) => {
3238            let mut visitor = ExprContainsTemporalVisitor::new();
3239            visitor.visit_expr(aligned_to);
3240            visitor.contains_temporal
3241        }
3242        Some(WithOptionValue::Refresh(RefreshOptionValue::Every(RefreshEveryOptionValue {
3243            interval: _,
3244            aligned_to: None,
3245        }))) => {
3246            // For a `REFRESH EVERY` without an `ALIGNED TO`, purification will default the
3247            // `ALIGNED TO` to `mz_now()`.
3248            true
3249        }
3250        Some(WithOptionValue::Refresh(RefreshOptionValue::AtCreation)) => {
3251            // `REFRESH AT CREATION` will be purified to `REFRESH AT mz_now()`.
3252            true
3253        }
3254        _ => false,
3255    }
3256}
3257
3258/// Determines whether the AST involves `mz_now()`.
3259struct ExprContainsTemporalVisitor {
3260    pub contains_temporal: bool,
3261}
3262
3263impl ExprContainsTemporalVisitor {
3264    pub fn new() -> ExprContainsTemporalVisitor {
3265        ExprContainsTemporalVisitor {
3266            contains_temporal: false,
3267        }
3268    }
3269}
3270
3271impl Visit<'_, Aug> for ExprContainsTemporalVisitor {
3272    fn visit_function(&mut self, func: &Function<Aug>) {
3273        self.contains_temporal |= func.name.full_item_name().item == MZ_NOW_NAME;
3274        visit_function(self, func);
3275    }
3276}
3277
3278struct MzNowPurifierVisitor {
3279    pub mz_now: Option<Timestamp>,
3280    pub mz_timestamp_type: ResolvedDataType,
3281    pub introduced_mz_timestamp: bool,
3282}
3283
3284impl MzNowPurifierVisitor {
3285    pub fn new(
3286        mz_now: Option<Timestamp>,
3287        mz_timestamp_type: ResolvedDataType,
3288    ) -> MzNowPurifierVisitor {
3289        MzNowPurifierVisitor {
3290            mz_now,
3291            mz_timestamp_type,
3292            introduced_mz_timestamp: false,
3293        }
3294    }
3295}
3296
3297impl VisitMut<'_, Aug> for MzNowPurifierVisitor {
3298    fn visit_expr_mut(&mut self, expr: &'_ mut Expr<Aug>) {
3299        match expr {
3300            Expr::Function(Function {
3301                name:
3302                    ResolvedItemName::Item {
3303                        full_name: FullItemName { item, .. },
3304                        ..
3305                    },
3306                ..
3307            }) if item == &MZ_NOW_NAME.to_string() => {
3308                let mz_now = self.mz_now.expect(
3309                    "we should have chosen a timestamp if the expression contains mz_now()",
3310                );
3311                // We substitute `mz_now()` with number + a cast to `mz_timestamp`. The cast is to
3312                // not alter the type of the expression.
3313                *expr = Expr::Cast {
3314                    expr: Box::new(Expr::Value(Value::Number(mz_now.to_string()))),
3315                    data_type: self.mz_timestamp_type.clone(),
3316                    failure_mode: CastFailureMode::Error,
3317                };
3318                self.introduced_mz_timestamp = true;
3319            }
3320            _ => visit_expr_mut(self, expr),
3321        }
3322    }
3323}