Skip to main content

mz_storage_types/
sources.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//! Types and traits related to the introduction of changing collections into `dataflow`.
11
12use std::collections::BTreeMap;
13use std::fmt::Debug;
14use std::hash::Hash;
15use std::ops::{Add, AddAssign, Deref, DerefMut};
16use std::str::FromStr;
17use std::sync::Arc;
18use std::time::Duration;
19
20use arrow::array::{Array, ArrayRef, BinaryArray, BinaryBuilder, NullArray, StructArray};
21use arrow::datatypes::{Field, Fields};
22use bytes::{BufMut, Bytes};
23use columnation::Columnation;
24use itertools::EitherOrBoth::Both;
25use itertools::Itertools;
26use kafka::KafkaSourceExportDetails;
27use load_generator::{LoadGeneratorOutput, LoadGeneratorSourceExportDetails};
28use mz_ore::assert_none;
29use mz_persist_types::Codec;
30use mz_persist_types::arrow::ArrayOrd;
31use mz_persist_types::columnar::{ColumnDecoder, ColumnEncoder, Schema};
32use mz_persist_types::stats::{
33    ColumnNullStats, ColumnStatKinds, ColumnarStats, ColumnarStatsBuilder, PrimitiveStats,
34    StructStats,
35};
36use mz_proto::{IntoRustIfSome, ProtoType, RustType, TryFromProtoError};
37#[cfg(any(test, feature = "proptest"))]
38use mz_repr::arb_row_for_relation;
39use mz_repr::{
40    CatalogItemId, Datum, GlobalId, ProtoRelationDesc, ProtoRow, RelationDesc, Row,
41    RowColumnarDecoder, RowColumnarEncoder,
42};
43use mz_sql_parser::ast::{Ident, IdentError, UnresolvedItemName};
44#[cfg(any(test, feature = "proptest"))]
45use proptest::prelude::any;
46#[cfg(any(test, feature = "proptest"))]
47use proptest::strategy::Strategy;
48use prost::Message;
49use serde::{Deserialize, Serialize};
50use timely::order::{PartialOrder, TotalOrder};
51use timely::progress::timestamp::Refines;
52use timely::progress::{PathSummary, Timestamp};
53
54use crate::AlterCompatible;
55use crate::connections::inline::{
56    ConnectionAccess, ConnectionResolver, InlinedConnection, IntoInlineConnection,
57    ReferencedConnection,
58};
59use crate::controller::AlterError;
60use crate::errors::{DataflowError, ProtoDataflowError};
61use crate::instances::StorageInstanceId;
62use crate::sources::sql_server::SqlServerSourceExportDetails;
63
64pub mod casts;
65pub mod encoding;
66pub mod envelope;
67pub mod kafka;
68pub mod load_generator;
69pub mod mysql;
70pub mod postgres;
71pub mod sql_server;
72
73pub use crate::sources::envelope::SourceEnvelope;
74pub use crate::sources::kafka::KafkaSourceConnection;
75pub use crate::sources::load_generator::LoadGeneratorSourceConnection;
76pub use crate::sources::mysql::{MySqlSourceConnection, MySqlSourceExportDetails};
77pub use crate::sources::postgres::{PostgresSourceConnection, PostgresSourceExportDetails};
78pub use crate::sources::sql_server::{SqlServerSourceConnection, SqlServerSourceExtras};
79
80include!(concat!(env!("OUT_DIR"), "/mz_storage_types.sources.rs"));
81
82/// A description of a source ingestion
83#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
84pub struct IngestionDescription<S: 'static = (), C: ConnectionAccess = InlinedConnection> {
85    /// The source description.
86    pub desc: SourceDesc<C>,
87    /// Collections to be exported by this ingestion.
88    ///
89    /// # Notes
90    /// - For multi-output sources:
91    ///     - Add exports by adding a new [`SourceExport`].
92    ///     - Remove exports by removing the [`SourceExport`].
93    ///
94    ///   Re-rendering/executing the source after making these modifications
95    ///   adds and drops the subsource, respectively.
96    /// - This field includes the primary source's ID, which might need to be
97    ///   filtered out to understand which exports are explicit ingestion exports.
98    /// - This field does _not_ include the remap collection, which is tracked
99    ///   in its own field.
100    pub source_exports: BTreeMap<GlobalId, SourceExport<S>>,
101    /// The ID of the instance in which to install the source.
102    pub instance_id: StorageInstanceId,
103    /// The ID of this ingestion's remap/progress collection.
104    pub remap_collection_id: GlobalId,
105    /// The storage metadata for the remap/progress collection
106    pub remap_metadata: S,
107}
108
109impl IngestionDescription {
110    pub fn new(
111        desc: SourceDesc,
112        instance_id: StorageInstanceId,
113        remap_collection_id: GlobalId,
114    ) -> Self {
115        Self {
116            desc,
117            remap_metadata: (),
118            source_exports: BTreeMap::new(),
119            instance_id,
120            remap_collection_id,
121        }
122    }
123}
124
125impl<S> IngestionDescription<S> {
126    /// Return an iterator over the `GlobalId`s of `self`'s collections.
127    /// This will contain ids for the remap collection, subsources,
128    /// tables for this source, and the primary collection ID, even if
129    /// no data will be exported to the primary collection.
130    pub fn collection_ids(&self) -> impl Iterator<Item = GlobalId> + '_ {
131        // Expand self so that any new fields added generate a compiler error to
132        // increase the likelihood of developers seeing this function.
133        let IngestionDescription {
134            desc: _,
135            remap_metadata: _,
136            source_exports,
137            instance_id: _,
138            remap_collection_id,
139        } = &self;
140
141        source_exports
142            .keys()
143            .copied()
144            .chain(std::iter::once(*remap_collection_id))
145    }
146}
147
148impl<S: Debug + Eq + PartialEq + AlterCompatible> AlterCompatible for IngestionDescription<S> {
149    fn alter_compatible(
150        &self,
151        id: GlobalId,
152        other: &IngestionDescription<S>,
153    ) -> Result<(), AlterError> {
154        if self == other {
155            return Ok(());
156        }
157        let IngestionDescription {
158            desc,
159            remap_metadata,
160            source_exports,
161            instance_id,
162            remap_collection_id,
163        } = self;
164
165        let compatibility_checks = [
166            (desc.alter_compatible(id, &other.desc).is_ok(), "desc"),
167            (remap_metadata == &other.remap_metadata, "remap_metadata"),
168            (
169                source_exports
170                    .iter()
171                    .merge_join_by(&other.source_exports, |(l_key, _), (r_key, _)| {
172                        l_key.cmp(r_key)
173                    })
174                    .all(|r| match r {
175                        Both(
176                            (
177                                _,
178                                SourceExport {
179                                    storage_metadata: l_metadata,
180                                    details: l_details,
181                                    data_config: l_data_config,
182                                },
183                            ),
184                            (
185                                _,
186                                SourceExport {
187                                    storage_metadata: r_metadata,
188                                    details: r_details,
189                                    data_config: r_data_config,
190                                },
191                            ),
192                        ) => {
193                            l_metadata.alter_compatible(id, r_metadata).is_ok()
194                                && l_details.alter_compatible(id, r_details).is_ok()
195                                && l_data_config.alter_compatible(id, r_data_config).is_ok()
196                        }
197                        _ => true,
198                    }),
199                "source_exports",
200            ),
201            (instance_id == &other.instance_id, "instance_id"),
202            (
203                remap_collection_id == &other.remap_collection_id,
204                "remap_collection_id",
205            ),
206        ];
207        for (compatible, field) in compatibility_checks {
208            if !compatible {
209                tracing::warn!(
210                    "IngestionDescription incompatible at {field}:\nself:\n{:#?}\n\nother\n{:#?}",
211                    self,
212                    other
213                );
214
215                return Err(AlterError { id });
216            }
217        }
218
219        Ok(())
220    }
221}
222
223impl<R: ConnectionResolver> IntoInlineConnection<IngestionDescription, R>
224    for IngestionDescription<(), ReferencedConnection>
225{
226    fn into_inline_connection(self, r: R) -> IngestionDescription {
227        let IngestionDescription {
228            desc,
229            remap_metadata,
230            source_exports,
231            instance_id,
232            remap_collection_id,
233        } = self;
234
235        IngestionDescription {
236            desc: desc.into_inline_connection(r),
237            remap_metadata,
238            source_exports,
239            instance_id,
240            remap_collection_id,
241        }
242    }
243}
244
245#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
246pub struct SourceExport<S = (), C: ConnectionAccess = InlinedConnection> {
247    /// The collection metadata needed to write the exported data
248    pub storage_metadata: S,
249    /// Details necessary for the source to export data to this export's collection.
250    pub details: SourceExportDetails,
251    /// Config necessary to handle (e.g. decode and envelope) the data for this export.
252    pub data_config: SourceExportDataConfig<C>,
253}
254
255pub trait SourceTimestamp:
256    Timestamp + Columnation + Refines<()> + std::fmt::Display + Sync
257{
258    fn encode_row(&self) -> Row;
259    fn decode_row(row: &Row) -> Self;
260}
261
262impl SourceTimestamp for MzOffset {
263    fn encode_row(&self) -> Row {
264        Row::pack([Datum::UInt64(self.offset)])
265    }
266
267    fn decode_row(row: &Row) -> Self {
268        let mut datums = row.iter();
269        match (datums.next(), datums.next()) {
270            (Some(Datum::UInt64(offset)), None) => MzOffset::from(offset),
271            _ => panic!("invalid row {row:?}"),
272        }
273    }
274}
275
276/// Universal language for describing message positions in Materialize, in a source independent
277/// way. Individual sources like Kafka or File sources should explicitly implement their own offset
278/// type that converts to/From MzOffsets. A 0-MzOffset denotes an empty stream.
279#[derive(
280    Copy,
281    Clone,
282    Default,
283    Debug,
284    PartialEq,
285    PartialOrd,
286    Eq,
287    Ord,
288    Hash,
289    Serialize,
290    Deserialize
291)]
292pub struct MzOffset {
293    pub offset: u64,
294}
295
296impl differential_dataflow::difference::Semigroup for MzOffset {
297    fn plus_equals(&mut self, rhs: &Self) {
298        self.offset.plus_equals(&rhs.offset)
299    }
300}
301
302impl differential_dataflow::difference::IsZero for MzOffset {
303    fn is_zero(&self) -> bool {
304        self.offset.is_zero()
305    }
306}
307
308impl mz_persist_types::Codec64 for MzOffset {
309    fn codec_name() -> String {
310        "MzOffset".to_string()
311    }
312
313    fn encode(&self) -> [u8; 8] {
314        mz_persist_types::Codec64::encode(&self.offset)
315    }
316
317    fn decode(buf: [u8; 8]) -> Self {
318        Self {
319            offset: mz_persist_types::Codec64::decode(buf),
320        }
321    }
322}
323
324impl columnation::Columnation for MzOffset {
325    type InnerRegion = columnation::CopyRegion<MzOffset>;
326}
327
328impl MzOffset {
329    pub fn checked_sub(self, other: Self) -> Option<Self> {
330        self.offset
331            .checked_sub(other.offset)
332            .map(|offset| Self { offset })
333    }
334}
335
336/// Convert from MzOffset to Kafka::Offset as long as
337/// the offset is not negative
338impl From<u64> for MzOffset {
339    fn from(offset: u64) -> Self {
340        Self { offset }
341    }
342}
343
344impl std::fmt::Display for MzOffset {
345    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
346        write!(f, "{}", self.offset)
347    }
348}
349
350// Assume overflow does not occur for addition
351impl Add<u64> for MzOffset {
352    type Output = MzOffset;
353
354    fn add(self, x: u64) -> MzOffset {
355        MzOffset {
356            offset: self.offset + x,
357        }
358    }
359}
360impl Add<Self> for MzOffset {
361    type Output = Self;
362
363    fn add(self, x: Self) -> Self {
364        MzOffset {
365            offset: self.offset + x.offset,
366        }
367    }
368}
369impl AddAssign<u64> for MzOffset {
370    fn add_assign(&mut self, x: u64) {
371        self.offset += x;
372    }
373}
374impl AddAssign<Self> for MzOffset {
375    fn add_assign(&mut self, x: Self) {
376        self.offset += x.offset;
377    }
378}
379
380/// Convert from `PgLsn` to MzOffset
381impl From<tokio_postgres::types::PgLsn> for MzOffset {
382    fn from(lsn: tokio_postgres::types::PgLsn) -> Self {
383        MzOffset { offset: lsn.into() }
384    }
385}
386
387impl Timestamp for MzOffset {
388    type Summary = MzOffset;
389
390    fn minimum() -> Self {
391        MzOffset {
392            offset: Timestamp::minimum(),
393        }
394    }
395}
396
397impl PathSummary<MzOffset> for MzOffset {
398    fn results_in(&self, src: &MzOffset) -> Option<MzOffset> {
399        Some(MzOffset {
400            offset: self.offset.results_in(&src.offset)?,
401        })
402    }
403
404    fn followed_by(&self, other: &Self) -> Option<Self> {
405        Some(MzOffset {
406            offset: PathSummary::<u64>::followed_by(&self.offset, &other.offset)?,
407        })
408    }
409}
410
411impl Refines<()> for MzOffset {
412    fn to_inner(_: ()) -> Self {
413        MzOffset::minimum()
414    }
415    fn to_outer(self) {}
416    fn summarize(_: Self::Summary) {}
417}
418
419impl PartialOrder for MzOffset {
420    #[inline]
421    fn less_equal(&self, other: &Self) -> bool {
422        self.offset.less_equal(&other.offset)
423    }
424}
425
426impl TotalOrder for MzOffset {}
427
428/// The meaning of the timestamp number produced by data sources. This type
429/// is not concerned with the source of the timestamp (like if the data came
430/// from a Debezium consistency topic or a CDCv2 stream), instead only what the
431/// timestamp number means.
432///
433/// Some variants here have attached data used to differentiate incomparable
434/// instantiations. These attached data types should be expanded in the future
435/// if we need to tell apart more kinds of sources.
436#[derive(
437    Clone,
438    Debug,
439    Ord,
440    PartialOrd,
441    Eq,
442    PartialEq,
443    Serialize,
444    Deserialize,
445    Hash
446)]
447pub enum Timeline {
448    /// EpochMilliseconds means the timestamp is the number of milliseconds since
449    /// the Unix epoch.
450    EpochMilliseconds,
451    /// External means the timestamp comes from an external data source and we
452    /// don't know what the number means. The attached String is the source's name,
453    /// which will result in different sources being incomparable.
454    External(String),
455    /// User means the user has manually specified a timeline. The attached
456    /// String is specified by the user, allowing them to decide sources that are
457    /// joinable.
458    User(String),
459}
460
461impl Timeline {
462    const EPOCH_MILLISECOND_ID_CHAR: char = 'M';
463    const EXTERNAL_ID_CHAR: char = 'E';
464    const USER_ID_CHAR: char = 'U';
465
466    fn id_char(&self) -> char {
467        match self {
468            Self::EpochMilliseconds => Self::EPOCH_MILLISECOND_ID_CHAR,
469            Self::External(_) => Self::EXTERNAL_ID_CHAR,
470            Self::User(_) => Self::USER_ID_CHAR,
471        }
472    }
473}
474
475impl ToString for Timeline {
476    fn to_string(&self) -> String {
477        match self {
478            Self::EpochMilliseconds => format!("{}", self.id_char()),
479            Self::External(id) => format!("{}.{id}", self.id_char()),
480            Self::User(id) => format!("{}.{id}", self.id_char()),
481        }
482    }
483}
484
485impl FromStr for Timeline {
486    type Err = String;
487
488    fn from_str(s: &str) -> Result<Self, Self::Err> {
489        if s.is_empty() {
490            return Err("empty timeline".to_string());
491        }
492        let mut chars = s.chars();
493        match chars.next().expect("non-empty string") {
494            Self::EPOCH_MILLISECOND_ID_CHAR => match chars.next() {
495                None => Ok(Self::EpochMilliseconds),
496                Some(_) => Err(format!("unknown timeline: {s}")),
497            },
498            Self::EXTERNAL_ID_CHAR => match chars.next() {
499                Some('.') => Ok(Self::External(chars.as_str().to_string())),
500                _ => Err(format!("unknown timeline: {s}")),
501            },
502            Self::USER_ID_CHAR => match chars.next() {
503                Some('.') => Ok(Self::User(chars.as_str().to_string())),
504                _ => Err(format!("unknown timeline: {s}")),
505            },
506            _ => Err(format!("unknown timeline: {s}")),
507        }
508    }
509}
510
511/// A connection to an external system
512pub trait SourceConnection: Debug + Clone + PartialEq + AlterCompatible {
513    /// The name of the external system (e.g kafka, postgres, etc).
514    fn name(&self) -> &'static str;
515
516    /// The name of the resource in the external system (e.g kafka topic) if any
517    fn external_reference(&self) -> Option<&str>;
518
519    /// Defines the key schema to use by default for this source connection type.
520    /// This will be used for the primary export of the source and as the default
521    /// pre-encoding key schema for the source.
522    fn default_key_desc(&self) -> RelationDesc;
523
524    /// Defines the value schema to use by default for this source connection type.
525    /// This will be used for the primary export of the source and as the default
526    /// pre-encoding value schema for the source.
527    fn default_value_desc(&self) -> RelationDesc;
528
529    /// The schema of this connection's timestamp type. This will also be the schema of the
530    /// progress relation.
531    fn timestamp_desc(&self) -> RelationDesc;
532
533    /// The id of the connection object (i.e the one obtained from running `CREATE CONNECTION`) in
534    /// the catalog, if any.
535    fn connection_id(&self) -> Option<CatalogItemId>;
536
537    /// Whether the source type supports read only mode.
538    fn supports_read_only(&self) -> bool;
539
540    /// Whether the source type prefers to run on only one replica of a multi-replica cluster.
541    fn prefers_single_replica(&self) -> bool;
542}
543
544#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
545pub enum Compression {
546    Gzip,
547    None,
548}
549
550/// Defines the configuration for how to handle data that is exported for a given
551/// Source Export.
552#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
553pub struct SourceExportDataConfig<C: ConnectionAccess = InlinedConnection> {
554    pub encoding: Option<encoding::SourceDataEncoding<C>>,
555    pub envelope: SourceEnvelope,
556}
557
558impl<R: ConnectionResolver> IntoInlineConnection<SourceExportDataConfig, R>
559    for SourceExportDataConfig<ReferencedConnection>
560{
561    fn into_inline_connection(self, r: R) -> SourceExportDataConfig {
562        let SourceExportDataConfig { encoding, envelope } = self;
563
564        SourceExportDataConfig {
565            encoding: encoding.map(|e| e.into_inline_connection(r)),
566            envelope,
567        }
568    }
569}
570
571impl<C: ConnectionAccess> AlterCompatible for SourceExportDataConfig<C> {
572    fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
573        if self == other {
574            return Ok(());
575        }
576        let Self { encoding, envelope } = &self;
577
578        let compatibility_checks = [
579            (
580                match (encoding, &other.encoding) {
581                    (Some(s), Some(o)) => s.alter_compatible(id, o).is_ok(),
582                    (s, o) => s == o,
583                },
584                "encoding",
585            ),
586            (envelope == &other.envelope, "envelope"),
587        ];
588
589        for (compatible, field) in compatibility_checks {
590            if !compatible {
591                tracing::warn!(
592                    "SourceDesc incompatible {field}:\nself:\n{:#?}\n\nother\n{:#?}",
593                    self,
594                    other
595                );
596
597                return Err(AlterError { id });
598            }
599        }
600        Ok(())
601    }
602}
603
604impl<C: ConnectionAccess> SourceExportDataConfig<C> {
605    /// Returns `true` if this connection yields data that is
606    /// append-only/monotonic. Append-monly means the source
607    /// never produces retractions.
608    // TODO(guswynn): consider enforcing this more completely at the
609    // parsing/typechecking level, by not using an `envelope`
610    // for sources like pg
611    pub fn monotonic(&self, connection: &GenericSourceConnection<C>) -> bool {
612        match &self.envelope {
613            // Upsert and CdcV2 may produce retractions.
614            SourceEnvelope::Upsert(_) | SourceEnvelope::CdcV2 => false,
615            SourceEnvelope::None(_) => {
616                match connection {
617                    // Postgres can produce retractions (deletes).
618                    GenericSourceConnection::Postgres(_) => false,
619                    // MySQL can produce retractions (deletes).
620                    GenericSourceConnection::MySql(_) => false,
621                    // SQL Server can produce retractions (deletes).
622                    GenericSourceConnection::SqlServer(_) => false,
623                    // Whether or not a Loadgen source can produce retractions varies.
624                    GenericSourceConnection::LoadGenerator(g) => g.load_generator.is_monotonic(),
625                    // Kafka exports with `None` envelope are append-only.
626                    GenericSourceConnection::Kafka(_) => true,
627                }
628            }
629        }
630    }
631}
632
633/// An external source of updates for a relational collection.
634#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
635pub struct SourceDesc<C: ConnectionAccess = InlinedConnection> {
636    pub connection: GenericSourceConnection<C>,
637    pub timestamp_interval: Duration,
638}
639
640impl<R: ConnectionResolver> IntoInlineConnection<SourceDesc, R>
641    for SourceDesc<ReferencedConnection>
642{
643    fn into_inline_connection(self, r: R) -> SourceDesc {
644        let SourceDesc {
645            connection,
646            timestamp_interval,
647        } = self;
648
649        SourceDesc {
650            connection: connection.into_inline_connection(&r),
651            timestamp_interval,
652        }
653    }
654}
655
656impl<C: ConnectionAccess> AlterCompatible for SourceDesc<C> {
657    /// Determines if `self` is compatible with another `SourceDesc`, in such a
658    /// way that it is possible to turn `self` into `other` through a valid
659    /// series of transformations (e.g. no transformation or `ALTER SOURCE`).
660    fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
661        if self == other {
662            return Ok(());
663        }
664        let Self {
665            connection,
666            // timestamp_interval is allowed to change via ALTER SOURCE
667            timestamp_interval: _,
668        } = &self;
669
670        let compatibility_checks = [(
671            connection.alter_compatible(id, &other.connection).is_ok(),
672            "connection",
673        )];
674
675        for (compatible, field) in compatibility_checks {
676            if !compatible {
677                tracing::warn!(
678                    "SourceDesc incompatible {field}:\nself:\n{:#?}\n\nother\n{:#?}",
679                    self,
680                    other
681                );
682
683                return Err(AlterError { id });
684            }
685        }
686
687        Ok(())
688    }
689}
690
691#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
692pub enum GenericSourceConnection<C: ConnectionAccess = InlinedConnection> {
693    Kafka(KafkaSourceConnection<C>),
694    Postgres(PostgresSourceConnection<C>),
695    MySql(MySqlSourceConnection<C>),
696    SqlServer(SqlServerSourceConnection<C>),
697    LoadGenerator(LoadGeneratorSourceConnection),
698}
699
700impl<C: ConnectionAccess> From<KafkaSourceConnection<C>> for GenericSourceConnection<C> {
701    fn from(conn: KafkaSourceConnection<C>) -> Self {
702        Self::Kafka(conn)
703    }
704}
705
706impl<C: ConnectionAccess> From<PostgresSourceConnection<C>> for GenericSourceConnection<C> {
707    fn from(conn: PostgresSourceConnection<C>) -> Self {
708        Self::Postgres(conn)
709    }
710}
711
712impl<C: ConnectionAccess> From<MySqlSourceConnection<C>> for GenericSourceConnection<C> {
713    fn from(conn: MySqlSourceConnection<C>) -> Self {
714        Self::MySql(conn)
715    }
716}
717
718impl<C: ConnectionAccess> From<SqlServerSourceConnection<C>> for GenericSourceConnection<C> {
719    fn from(conn: SqlServerSourceConnection<C>) -> Self {
720        Self::SqlServer(conn)
721    }
722}
723
724impl<C: ConnectionAccess> From<LoadGeneratorSourceConnection> for GenericSourceConnection<C> {
725    fn from(conn: LoadGeneratorSourceConnection) -> Self {
726        Self::LoadGenerator(conn)
727    }
728}
729
730impl<R: ConnectionResolver> IntoInlineConnection<GenericSourceConnection, R>
731    for GenericSourceConnection<ReferencedConnection>
732{
733    fn into_inline_connection(self, r: R) -> GenericSourceConnection {
734        match self {
735            GenericSourceConnection::Kafka(kafka) => {
736                GenericSourceConnection::Kafka(kafka.into_inline_connection(r))
737            }
738            GenericSourceConnection::Postgres(pg) => {
739                GenericSourceConnection::Postgres(pg.into_inline_connection(r))
740            }
741            GenericSourceConnection::MySql(mysql) => {
742                GenericSourceConnection::MySql(mysql.into_inline_connection(r))
743            }
744            GenericSourceConnection::SqlServer(sql_server) => {
745                GenericSourceConnection::SqlServer(sql_server.into_inline_connection(r))
746            }
747            GenericSourceConnection::LoadGenerator(lg) => {
748                GenericSourceConnection::LoadGenerator(lg)
749            }
750        }
751    }
752}
753
754impl<C: ConnectionAccess> SourceConnection for GenericSourceConnection<C> {
755    fn name(&self) -> &'static str {
756        match self {
757            Self::Kafka(conn) => conn.name(),
758            Self::Postgres(conn) => conn.name(),
759            Self::MySql(conn) => conn.name(),
760            Self::SqlServer(conn) => conn.name(),
761            Self::LoadGenerator(conn) => conn.name(),
762        }
763    }
764
765    fn external_reference(&self) -> Option<&str> {
766        match self {
767            Self::Kafka(conn) => conn.external_reference(),
768            Self::Postgres(conn) => conn.external_reference(),
769            Self::MySql(conn) => conn.external_reference(),
770            Self::SqlServer(conn) => conn.external_reference(),
771            Self::LoadGenerator(conn) => conn.external_reference(),
772        }
773    }
774
775    fn default_key_desc(&self) -> RelationDesc {
776        match self {
777            Self::Kafka(conn) => conn.default_key_desc(),
778            Self::Postgres(conn) => conn.default_key_desc(),
779            Self::MySql(conn) => conn.default_key_desc(),
780            Self::SqlServer(conn) => conn.default_key_desc(),
781            Self::LoadGenerator(conn) => conn.default_key_desc(),
782        }
783    }
784
785    fn default_value_desc(&self) -> RelationDesc {
786        match self {
787            Self::Kafka(conn) => conn.default_value_desc(),
788            Self::Postgres(conn) => conn.default_value_desc(),
789            Self::MySql(conn) => conn.default_value_desc(),
790            Self::SqlServer(conn) => conn.default_value_desc(),
791            Self::LoadGenerator(conn) => conn.default_value_desc(),
792        }
793    }
794
795    fn timestamp_desc(&self) -> RelationDesc {
796        match self {
797            Self::Kafka(conn) => conn.timestamp_desc(),
798            Self::Postgres(conn) => conn.timestamp_desc(),
799            Self::MySql(conn) => conn.timestamp_desc(),
800            Self::SqlServer(conn) => conn.timestamp_desc(),
801            Self::LoadGenerator(conn) => conn.timestamp_desc(),
802        }
803    }
804
805    fn connection_id(&self) -> Option<CatalogItemId> {
806        match self {
807            Self::Kafka(conn) => conn.connection_id(),
808            Self::Postgres(conn) => conn.connection_id(),
809            Self::MySql(conn) => conn.connection_id(),
810            Self::SqlServer(conn) => conn.connection_id(),
811            Self::LoadGenerator(conn) => conn.connection_id(),
812        }
813    }
814
815    fn supports_read_only(&self) -> bool {
816        match self {
817            GenericSourceConnection::Kafka(conn) => conn.supports_read_only(),
818            GenericSourceConnection::Postgres(conn) => conn.supports_read_only(),
819            GenericSourceConnection::MySql(conn) => conn.supports_read_only(),
820            GenericSourceConnection::SqlServer(conn) => conn.supports_read_only(),
821            GenericSourceConnection::LoadGenerator(conn) => conn.supports_read_only(),
822        }
823    }
824
825    fn prefers_single_replica(&self) -> bool {
826        match self {
827            GenericSourceConnection::Kafka(conn) => conn.prefers_single_replica(),
828            GenericSourceConnection::Postgres(conn) => conn.prefers_single_replica(),
829            GenericSourceConnection::MySql(conn) => conn.prefers_single_replica(),
830            GenericSourceConnection::SqlServer(conn) => conn.prefers_single_replica(),
831            GenericSourceConnection::LoadGenerator(conn) => conn.prefers_single_replica(),
832        }
833    }
834}
835impl<C: ConnectionAccess> crate::AlterCompatible for GenericSourceConnection<C> {
836    fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
837        if self == other {
838            return Ok(());
839        }
840        let r = match (self, other) {
841            (Self::Kafka(conn), Self::Kafka(other)) => conn.alter_compatible(id, other),
842            (Self::Postgres(conn), Self::Postgres(other)) => conn.alter_compatible(id, other),
843            (Self::MySql(conn), Self::MySql(other)) => conn.alter_compatible(id, other),
844            (Self::SqlServer(conn), Self::SqlServer(other)) => conn.alter_compatible(id, other),
845            (Self::LoadGenerator(conn), Self::LoadGenerator(other)) => {
846                conn.alter_compatible(id, other)
847            }
848            _ => Err(AlterError { id }),
849        };
850
851        if r.is_err() {
852            tracing::warn!(
853                "GenericSourceConnection incompatible:\nself:\n{:#?}\n\nother\n{:#?}",
854                self,
855                other
856            );
857        }
858
859        r
860    }
861}
862
863/// Details necessary for each source export to allow the source implementations
864/// to export data to the export's collection.
865#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
866pub enum SourceExportDetails {
867    /// Used when the primary collection of a source isn't an export to
868    /// output to.
869    None,
870    Kafka(KafkaSourceExportDetails),
871    Postgres(PostgresSourceExportDetails),
872    MySql(MySqlSourceExportDetails),
873    SqlServer(SqlServerSourceExportDetails),
874    LoadGenerator(LoadGeneratorSourceExportDetails),
875}
876
877impl crate::AlterCompatible for SourceExportDetails {
878    fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
879        if self == other {
880            return Ok(());
881        }
882        let r = match (self, other) {
883            (Self::None, Self::None) => Ok(()),
884            (Self::Kafka(s), Self::Kafka(o)) => s.alter_compatible(id, o),
885            (Self::Postgres(s), Self::Postgres(o)) => s.alter_compatible(id, o),
886            (Self::MySql(s), Self::MySql(o)) => s.alter_compatible(id, o),
887            (Self::SqlServer(s), Self::SqlServer(o)) => s.alter_compatible(id, o),
888            (Self::LoadGenerator(s), Self::LoadGenerator(o)) => s.alter_compatible(id, o),
889            _ => Err(AlterError { id }),
890        };
891
892        if r.is_err() {
893            tracing::warn!(
894                "SourceExportDetails incompatible:\nself:\n{:#?}\n\nother\n{:#?}",
895                self,
896                other
897            );
898        }
899
900        r
901    }
902}
903
904/// Details necessary to store in the `Details` option of a source export
905/// statement (`CREATE SUBSOURCE` and `CREATE TABLE .. FROM SOURCE` statements),
906/// to generate the appropriate `SourceExportDetails` struct during planning.
907/// NOTE that this is serialized as proto to the catalog, so any changes here
908/// must be backwards compatible or will require a migration.
909pub enum SourceExportStatementDetails {
910    Postgres {
911        table: mz_postgres_util::desc::PostgresTableDesc,
912        /// Whether the text-to-oid cast for this export accepts the full `u32`
913        /// range. Exports created before the cast was widened decode as
914        /// `false` and must keep the legacy `i32`-range cast forever, because
915        /// replication re-casts old tuples on delete and the persisted rows
916        /// were ingested under the legacy semantics.
917        cast_oid_full_range: bool,
918    },
919    MySql {
920        table: mz_mysql_util::MySqlTableDesc,
921        initial_gtid_set: String,
922        binlog_full_metadata: bool,
923    },
924    SqlServer {
925        table: mz_sql_server_util::desc::SqlServerTableDesc,
926        capture_instance: Arc<str>,
927        initial_lsn: mz_sql_server_util::cdc::Lsn,
928    },
929    LoadGenerator {
930        output: LoadGeneratorOutput,
931    },
932    Kafka {},
933}
934
935impl RustType<ProtoSourceExportStatementDetails> for SourceExportStatementDetails {
936    fn into_proto(&self) -> ProtoSourceExportStatementDetails {
937        match self {
938            SourceExportStatementDetails::Postgres {
939                table,
940                cast_oid_full_range,
941            } => ProtoSourceExportStatementDetails {
942                kind: Some(proto_source_export_statement_details::Kind::Postgres(
943                    postgres::ProtoPostgresSourceExportStatementDetails {
944                        table: Some(table.into_proto()),
945                        cast_oid_full_range: *cast_oid_full_range,
946                    },
947                )),
948            },
949            SourceExportStatementDetails::MySql {
950                table,
951                initial_gtid_set,
952                binlog_full_metadata,
953            } => ProtoSourceExportStatementDetails {
954                kind: Some(proto_source_export_statement_details::Kind::Mysql(
955                    mysql::ProtoMySqlSourceExportStatementDetails {
956                        table: Some(table.into_proto()),
957                        initial_gtid_set: initial_gtid_set.clone(),
958                        binlog_full_metadata: *binlog_full_metadata,
959                    },
960                )),
961            },
962            SourceExportStatementDetails::SqlServer {
963                table,
964                capture_instance,
965                initial_lsn,
966            } => ProtoSourceExportStatementDetails {
967                kind: Some(proto_source_export_statement_details::Kind::SqlServer(
968                    sql_server::ProtoSqlServerSourceExportStatementDetails {
969                        table: Some(table.into_proto()),
970                        capture_instance: capture_instance.to_string(),
971                        initial_lsn: initial_lsn.as_bytes().to_vec(),
972                    },
973                )),
974            },
975            SourceExportStatementDetails::LoadGenerator { output } => {
976                ProtoSourceExportStatementDetails {
977                    kind: Some(proto_source_export_statement_details::Kind::Loadgen(
978                        load_generator::ProtoLoadGeneratorSourceExportStatementDetails {
979                            output: output.into_proto().into(),
980                        },
981                    )),
982                }
983            }
984            SourceExportStatementDetails::Kafka {} => ProtoSourceExportStatementDetails {
985                kind: Some(proto_source_export_statement_details::Kind::Kafka(
986                    kafka::ProtoKafkaSourceExportStatementDetails {},
987                )),
988            },
989        }
990    }
991
992    fn from_proto(proto: ProtoSourceExportStatementDetails) -> Result<Self, TryFromProtoError> {
993        use proto_source_export_statement_details::Kind;
994        Ok(match proto.kind {
995            Some(Kind::Postgres(details)) => SourceExportStatementDetails::Postgres {
996                table: details
997                    .table
998                    .into_rust_if_some("ProtoPostgresSourceExportStatementDetails::table")?,
999                cast_oid_full_range: details.cast_oid_full_range,
1000            },
1001            Some(Kind::Mysql(details)) => SourceExportStatementDetails::MySql {
1002                table: details
1003                    .table
1004                    .into_rust_if_some("ProtoMySqlSourceExportStatementDetails::table")?,
1005
1006                initial_gtid_set: details.initial_gtid_set,
1007                binlog_full_metadata: details.binlog_full_metadata,
1008            },
1009            Some(Kind::SqlServer(details)) => SourceExportStatementDetails::SqlServer {
1010                table: details
1011                    .table
1012                    .into_rust_if_some("ProtoSqlServerSourceExportStatementDetails::table")?,
1013                capture_instance: details.capture_instance.into(),
1014                initial_lsn: mz_sql_server_util::cdc::Lsn::try_from(details.initial_lsn.as_slice())
1015                    .map_err(|e| TryFromProtoError::InvalidFieldError(e.to_string()))?,
1016            },
1017            Some(Kind::Loadgen(details)) => SourceExportStatementDetails::LoadGenerator {
1018                output: details
1019                    .output
1020                    .into_rust_if_some("ProtoLoadGeneratorSourceExportStatementDetails::output")?,
1021            },
1022            Some(Kind::Kafka(_details)) => SourceExportStatementDetails::Kafka {},
1023            None => {
1024                return Err(TryFromProtoError::missing_field(
1025                    "ProtoSourceExportStatementDetails::kind",
1026                ));
1027            }
1028        })
1029    }
1030}
1031
1032#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
1033#[repr(transparent)]
1034pub struct SourceData(pub Result<Row, DataflowError>);
1035
1036impl Default for SourceData {
1037    fn default() -> Self {
1038        SourceData(Ok(Row::default()))
1039    }
1040}
1041
1042impl Deref for SourceData {
1043    type Target = Result<Row, DataflowError>;
1044
1045    fn deref(&self) -> &Self::Target {
1046        &self.0
1047    }
1048}
1049
1050impl DerefMut for SourceData {
1051    fn deref_mut(&mut self) -> &mut Self::Target {
1052        &mut self.0
1053    }
1054}
1055
1056impl RustType<ProtoSourceData> for SourceData {
1057    fn into_proto(&self) -> ProtoSourceData {
1058        use proto_source_data::Kind;
1059        ProtoSourceData {
1060            kind: Some(match &**self {
1061                Ok(row) => Kind::Ok(row.into_proto()),
1062                Err(err) => Kind::Err(err.into_proto()),
1063            }),
1064        }
1065    }
1066
1067    fn from_proto(proto: ProtoSourceData) -> Result<Self, TryFromProtoError> {
1068        use proto_source_data::Kind;
1069        match proto.kind {
1070            Some(kind) => match kind {
1071                Kind::Ok(row) => Ok(SourceData(Ok(row.into_rust()?))),
1072                Kind::Err(err) => Ok(SourceData(Err(err.into_rust()?))),
1073            },
1074            None => Result::Err(TryFromProtoError::missing_field("ProtoSourceData::kind")),
1075        }
1076    }
1077}
1078
1079impl Codec for SourceData {
1080    type Storage = ProtoRow;
1081    type Schema = RelationDesc;
1082
1083    fn codec_name() -> String {
1084        "protobuf[SourceData]".into()
1085    }
1086
1087    fn encode<B: BufMut>(&self, buf: &mut B) {
1088        self.into_proto()
1089            .encode(buf)
1090            .expect("no required fields means no initialization errors");
1091    }
1092
1093    fn decode(buf: &[u8], schema: &RelationDesc) -> Result<Self, String> {
1094        let mut val = SourceData::default();
1095        <Self as Codec>::decode_from(&mut val, buf, &mut None, schema)?;
1096        Ok(val)
1097    }
1098
1099    fn decode_from<'a>(
1100        &mut self,
1101        buf: &'a [u8],
1102        storage: &mut Option<ProtoRow>,
1103        schema: &RelationDesc,
1104    ) -> Result<(), String> {
1105        // Optimize for common case of `Ok` by leaving a (cleared) `ProtoRow` in
1106        // the `Ok` variant of `ProtoSourceData`. prost's `Message::merge` impl
1107        // is smart about reusing the `Vec<Datum>` when it can.
1108        let mut proto = storage.take().unwrap_or_default();
1109        proto.clear();
1110        let mut proto = ProtoSourceData {
1111            kind: Some(proto_source_data::Kind::Ok(proto)),
1112        };
1113        proto.merge(buf).map_err(|err| err.to_string())?;
1114        match (proto.kind, &mut self.0) {
1115            // Again, optimize for the common case...
1116            (Some(proto_source_data::Kind::Ok(proto)), Ok(row)) => {
1117                let ret = row.decode_from_proto(&proto, schema);
1118                storage.replace(proto);
1119                ret
1120            }
1121            // ...otherwise fall back to the obvious thing.
1122            (kind, _) => {
1123                let proto = ProtoSourceData { kind };
1124                *self = proto.into_rust().map_err(|err| err.to_string())?;
1125                // Nothing to put back in storage.
1126                Ok(())
1127            }
1128        }
1129    }
1130
1131    fn validate(val: &Self, desc: &Self::Schema) -> Result<(), String> {
1132        match &val.0 {
1133            Ok(row) => Row::validate(row, desc),
1134            Err(_) => Ok(()),
1135        }
1136    }
1137
1138    fn encode_schema(schema: &Self::Schema) -> Bytes {
1139        schema.into_proto().encode_to_vec().into()
1140    }
1141
1142    fn decode_schema(buf: &Bytes) -> Self::Schema {
1143        let proto = ProtoRelationDesc::decode(buf.as_ref()).expect("valid schema");
1144        proto.into_rust().expect("valid schema")
1145    }
1146}
1147
1148/// Given a [`RelationDesc`] returns an arbitrary [`SourceData`].
1149#[cfg(any(test, feature = "proptest"))]
1150pub fn arb_source_data_for_relation_desc(
1151    desc: &RelationDesc,
1152) -> impl Strategy<Value = SourceData> + use<> {
1153    let row_strat = arb_row_for_relation(desc).no_shrink();
1154
1155    proptest::strategy::Union::new_weighted(vec![
1156        (50, row_strat.prop_map(|row| SourceData(Ok(row))).boxed()),
1157        (
1158            1,
1159            any::<DataflowError>()
1160                .prop_map(|err| SourceData(Err(err)))
1161                .no_shrink()
1162                .boxed(),
1163        ),
1164    ])
1165}
1166
1167/// Describes how external references should be organized in a multi-level
1168/// hierarchy.
1169///
1170/// For both PostgreSQL and MySQL sources, these levels of reference are
1171/// intrinsic to the items which we're referencing. If there are other naming
1172/// schemas for other types of sources we discover, we might need to revisit
1173/// this.
1174pub trait ExternalCatalogReference {
1175    /// The "second" level of namespacing for the reference.
1176    fn schema_name(&self) -> &str;
1177    /// The lowest level of namespacing for the reference.
1178    fn item_name(&self) -> &str;
1179}
1180
1181impl ExternalCatalogReference for &mz_mysql_util::MySqlTableDesc {
1182    fn schema_name(&self) -> &str {
1183        &self.schema_name
1184    }
1185
1186    fn item_name(&self) -> &str {
1187        &self.name
1188    }
1189}
1190
1191impl ExternalCatalogReference for mz_postgres_util::desc::PostgresTableDesc {
1192    fn schema_name(&self) -> &str {
1193        &self.namespace
1194    }
1195
1196    fn item_name(&self) -> &str {
1197        &self.name
1198    }
1199}
1200
1201impl ExternalCatalogReference for &mz_sql_server_util::desc::SqlServerTableDesc {
1202    fn schema_name(&self) -> &str {
1203        &*self.schema_name
1204    }
1205
1206    fn item_name(&self) -> &str {
1207        &*self.name
1208    }
1209}
1210
1211// This implementation provides a means of converting arbitrary objects into a
1212// `SubsourceCatalogReference`, e.g. load generator view names.
1213impl<'a> ExternalCatalogReference for (&'a str, &'a str) {
1214    fn schema_name(&self) -> &str {
1215        self.0
1216    }
1217
1218    fn item_name(&self) -> &str {
1219        self.1
1220    }
1221}
1222
1223/// Stores and resolves references to a `&[T: ExternalCatalogReference]`.
1224///
1225/// This is meant to provide an API to quickly look up a source's subsources.
1226///
1227/// For sources that do not provide any subsources, use the `Default`
1228/// implementation, which is empty and will not be able to resolve any
1229/// references.
1230#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
1231pub struct SourceReferenceResolver {
1232    inner: BTreeMap<Ident, BTreeMap<Ident, BTreeMap<Ident, usize>>>,
1233}
1234
1235#[derive(Debug, Clone, thiserror::Error)]
1236pub enum ExternalReferenceResolutionError {
1237    #[error("reference to {name} not found in source")]
1238    DoesNotExist { name: String },
1239    #[error(
1240        "reference {name} is ambiguous, consider specifying an additional \
1241    layer of qualification"
1242    )]
1243    Ambiguous { name: String },
1244    #[error("invalid identifier: {0}")]
1245    Ident(#[from] IdentError),
1246}
1247
1248impl<'a> SourceReferenceResolver {
1249    /// Constructs a new `SourceReferenceResolver` from a slice of `T:
1250    /// SubsourceCatalogReference`.
1251    ///
1252    /// # Errors
1253    /// - If any `&str` provided cannot be taken to an [`Ident`].
1254    pub fn new<T: ExternalCatalogReference>(
1255        database: &str,
1256        referenceable_items: &'a [T],
1257    ) -> Result<SourceReferenceResolver, ExternalReferenceResolutionError> {
1258        // An index from table name -> schema name -> database name -> index in
1259        // `referenceable_items`.
1260        let mut inner = BTreeMap::new();
1261
1262        let database = Ident::new(database)?;
1263
1264        for (reference_idx, item) in referenceable_items.iter().enumerate() {
1265            let item_name = Ident::new(item.item_name())?;
1266            let schema_name = Ident::new(item.schema_name())?;
1267
1268            inner
1269                .entry(item_name)
1270                .or_insert_with(BTreeMap::new)
1271                .entry(schema_name)
1272                .or_insert_with(BTreeMap::new)
1273                .entry(database.clone())
1274                .or_insert(reference_idx);
1275        }
1276
1277        Ok(SourceReferenceResolver { inner })
1278    }
1279
1280    /// Returns the canonical reference and index from which it originated in
1281    /// the `referenceable_items` provided to [`Self::new`].
1282    ///
1283    /// # Args
1284    /// - `name` is `&[Ident]` to let users provide the inner element of
1285    ///   [`UnresolvedItemName`].
1286    /// - `canonicalize_to_width` limits the number of elements in the returned
1287    ///   [`UnresolvedItemName`];this is useful if the source type requires
1288    ///   contriving database and schema names that a subsource should not
1289    ///   persist as its reference.
1290    ///
1291    /// # Errors
1292    /// - If `name` does not resolve to an item in `self.inner`.
1293    ///
1294    /// # Panics
1295    /// - If `canonicalize_to_width`` is not in `1..=3`.
1296    pub fn resolve(
1297        &self,
1298        name: &[Ident],
1299        canonicalize_to_width: usize,
1300    ) -> Result<(UnresolvedItemName, usize), ExternalReferenceResolutionError> {
1301        let (db, schema, idx) = self.resolve_inner(name)?;
1302
1303        let item = name.last().expect("must have provided at least 1 element");
1304
1305        let canonical_name = match canonicalize_to_width {
1306            1 => vec![item.clone()],
1307            2 => vec![schema.clone(), item.clone()],
1308            3 => vec![db.clone(), schema.clone(), item.clone()],
1309            o => panic!("canonicalize_to_width values must be 1..=3, but got {}", o),
1310        };
1311
1312        Ok((UnresolvedItemName(canonical_name), idx))
1313    }
1314
1315    /// Returns the index from which it originated in the `referenceable_items`
1316    /// provided to [`Self::new`].
1317    ///
1318    /// # Args
1319    /// `name` is `&[Ident]` to let users provide the inner element of
1320    /// [`UnresolvedItemName`].
1321    ///
1322    /// # Errors
1323    /// - If `name` does not resolve to an item in `self.inner`.
1324    pub fn resolve_idx(&self, name: &[Ident]) -> Result<usize, ExternalReferenceResolutionError> {
1325        let (_db, _schema, idx) = self.resolve_inner(name)?;
1326        Ok(idx)
1327    }
1328
1329    /// Returns the index from which it originated in the `referenceable_items`
1330    /// provided to [`Self::new`].
1331    ///
1332    /// # Args
1333    /// `name` is `&[Ident]` to let users provide the inner element of
1334    /// [`UnresolvedItemName`].
1335    ///
1336    /// # Return
1337    /// Returns a tuple whose elements are:
1338    /// 1. The "database"- or top-level namespace of the reference.
1339    /// 2. The "schema"- or second-level namespace of the reference.
1340    /// 3. The index to find the item in `referenceable_items` argument provided
1341    ///    to `SourceReferenceResolver::new`.
1342    ///
1343    /// # Errors
1344    /// - If `name` does not resolve to an item in `self.inner`.
1345    fn resolve_inner<'name: 'a>(
1346        &'a self,
1347        name: &'name [Ident],
1348    ) -> Result<(&'a Ident, &'a Ident, usize), ExternalReferenceResolutionError> {
1349        let get_provided_name = || UnresolvedItemName(name.to_vec()).to_string();
1350
1351        // Names must be composed of 1..=3 elements.
1352        if !(1..=3).contains(&name.len()) {
1353            Err(ExternalReferenceResolutionError::DoesNotExist {
1354                name: get_provided_name(),
1355            })?;
1356        }
1357
1358        // Fill on the leading elements with `None` if they aren't present.
1359        let mut names = std::iter::repeat(None)
1360            .take(3 - name.len())
1361            .chain(name.iter().map(Some));
1362
1363        let database = names.next().flatten();
1364        let schema = names.next().flatten();
1365        let item = names
1366            .next()
1367            .flatten()
1368            .expect("must have provided the item name");
1369
1370        assert_none!(names.next(), "expected a 3-element iterator");
1371
1372        let schemas =
1373            self.inner
1374                .get(item)
1375                .ok_or_else(|| ExternalReferenceResolutionError::DoesNotExist {
1376                    name: get_provided_name(),
1377                })?;
1378
1379        let schema = match schema {
1380            Some(schema) => schema,
1381            None => schemas.keys().exactly_one().map_err(|_e| {
1382                ExternalReferenceResolutionError::Ambiguous {
1383                    name: get_provided_name(),
1384                }
1385            })?,
1386        };
1387
1388        let databases =
1389            schemas
1390                .get(schema)
1391                .ok_or_else(|| ExternalReferenceResolutionError::DoesNotExist {
1392                    name: get_provided_name(),
1393                })?;
1394
1395        let database = match database {
1396            Some(database) => database,
1397            None => databases.keys().exactly_one().map_err(|_e| {
1398                ExternalReferenceResolutionError::Ambiguous {
1399                    name: get_provided_name(),
1400                }
1401            })?,
1402        };
1403
1404        let reference_idx = databases.get(database).ok_or_else(|| {
1405            ExternalReferenceResolutionError::DoesNotExist {
1406                name: get_provided_name(),
1407            }
1408        })?;
1409
1410        Ok((database, schema, *reference_idx))
1411    }
1412}
1413
1414/// A decoder for [`Row`]s within [`SourceData`].
1415///
1416/// This type exists as a wrapper around [`RowColumnarDecoder`] to handle the
1417/// case where the [`RelationDesc`] we're encoding with has no columns. See
1418/// [`SourceDataRowColumnarEncoder`] for more details.
1419#[derive(Debug)]
1420pub enum SourceDataRowColumnarDecoder {
1421    Row(RowColumnarDecoder),
1422    EmptyRow,
1423}
1424
1425impl SourceDataRowColumnarDecoder {
1426    pub fn decode(&self, idx: usize, row: &mut Row) {
1427        match self {
1428            SourceDataRowColumnarDecoder::Row(decoder) => decoder.decode(idx, row),
1429            SourceDataRowColumnarDecoder::EmptyRow => {
1430                // Create a packer just to clear the Row.
1431                row.packer();
1432            }
1433        }
1434    }
1435
1436    pub fn goodbytes(&self) -> usize {
1437        match self {
1438            SourceDataRowColumnarDecoder::Row(decoder) => decoder.goodbytes(),
1439            SourceDataRowColumnarDecoder::EmptyRow => 0,
1440        }
1441    }
1442}
1443
1444#[derive(Debug)]
1445pub struct SourceDataColumnarDecoder {
1446    row_decoder: SourceDataRowColumnarDecoder,
1447    err_decoder: BinaryArray,
1448}
1449
1450impl SourceDataColumnarDecoder {
1451    pub fn new(col: StructArray, desc: &RelationDesc) -> Result<Self, anyhow::Error> {
1452        // TODO(parkmcar): We should validate the fields here.
1453        let (_fields, arrays, nullability) = col.into_parts();
1454
1455        if nullability.is_some() {
1456            anyhow::bail!("SourceData is not nullable, but found {nullability:?}");
1457        }
1458        if arrays.len() != 2 {
1459            anyhow::bail!("SourceData should only have two fields, found {arrays:?}");
1460        }
1461
1462        let errs = arrays[1]
1463            .as_any()
1464            .downcast_ref::<BinaryArray>()
1465            .ok_or_else(|| anyhow::anyhow!("expected BinaryArray, found {:?}", arrays[1]))?;
1466
1467        let row_decoder = match arrays[0].data_type() {
1468            arrow::datatypes::DataType::Struct(_) => {
1469                let rows = arrays[0]
1470                    .as_any()
1471                    .downcast_ref::<StructArray>()
1472                    .ok_or_else(|| {
1473                        anyhow::anyhow!("expected StructArray, found {:?}", arrays[0])
1474                    })?;
1475                let decoder = RowColumnarDecoder::new(rows.clone(), desc)?;
1476                SourceDataRowColumnarDecoder::Row(decoder)
1477            }
1478            arrow::datatypes::DataType::Null => SourceDataRowColumnarDecoder::EmptyRow,
1479            other => anyhow::bail!("expected Struct or Null Array, found {other:?}"),
1480        };
1481
1482        Ok(SourceDataColumnarDecoder {
1483            row_decoder,
1484            err_decoder: errs.clone(),
1485        })
1486    }
1487}
1488
1489impl ColumnDecoder<SourceData> for SourceDataColumnarDecoder {
1490    fn decode(&self, idx: usize, val: &mut SourceData) {
1491        let err_null = self.err_decoder.is_null(idx);
1492        let row_null = match &self.row_decoder {
1493            SourceDataRowColumnarDecoder::Row(decoder) => decoder.is_null(idx),
1494            SourceDataRowColumnarDecoder::EmptyRow => !err_null,
1495        };
1496
1497        match (row_null, err_null) {
1498            (true, false) => {
1499                let err = self.err_decoder.value(idx);
1500                let err = ProtoDataflowError::decode(err)
1501                    .expect("proto should be valid")
1502                    .into_rust()
1503                    .expect("error should be valid");
1504                val.0 = Err(err);
1505            }
1506            (false, true) => {
1507                let row = match val.0.as_mut() {
1508                    Ok(row) => row,
1509                    Err(_) => {
1510                        val.0 = Ok(Row::default());
1511                        val.0.as_mut().unwrap()
1512                    }
1513                };
1514                self.row_decoder.decode(idx, row);
1515            }
1516            (true, true) => panic!("should have one of 'ok' or 'err'"),
1517            (false, false) => panic!("cannot have both 'ok' and 'err'"),
1518        }
1519    }
1520
1521    fn is_null(&self, idx: usize) -> bool {
1522        let err_null = self.err_decoder.is_null(idx);
1523        let row_null = match &self.row_decoder {
1524            SourceDataRowColumnarDecoder::Row(decoder) => decoder.is_null(idx),
1525            SourceDataRowColumnarDecoder::EmptyRow => !err_null,
1526        };
1527        assert!(!err_null || !row_null, "SourceData should never be null!");
1528
1529        false
1530    }
1531
1532    fn goodbytes(&self) -> usize {
1533        self.row_decoder.goodbytes() + ArrayOrd::Binary(self.err_decoder.clone()).goodbytes()
1534    }
1535
1536    fn stats(&self) -> StructStats {
1537        let len = self.err_decoder.len();
1538        let err_stats = ColumnarStats {
1539            nulls: Some(ColumnNullStats {
1540                count: self.err_decoder.null_count(),
1541            }),
1542            values: PrimitiveStats::<Vec<u8>>::from_column(&self.err_decoder).into(),
1543        };
1544        // The top level struct is non-nullable and every entry is either an
1545        // `Ok(Row)` or an `Err(String)`. As a result, we can compute the number
1546        // of `Ok` entries by subtracting the number of `Err` entries from the
1547        // total count.
1548        let row_null_count = len - self.err_decoder.null_count();
1549        let row_stats = match &self.row_decoder {
1550            SourceDataRowColumnarDecoder::Row(encoder) => {
1551                // Sanity check that the number of row nulls/nones we calculated
1552                // using the error column matches what the row column thinks it
1553                // has.
1554                assert_eq!(encoder.null_count(), row_null_count);
1555                encoder.stats()
1556            }
1557            SourceDataRowColumnarDecoder::EmptyRow => StructStats {
1558                len,
1559                cols: BTreeMap::default(),
1560            },
1561        };
1562        let row_stats = ColumnarStats {
1563            nulls: Some(ColumnNullStats {
1564                count: row_null_count,
1565            }),
1566            values: ColumnStatKinds::Struct(row_stats),
1567        };
1568
1569        let stats = [
1570            (
1571                SourceDataColumnarEncoder::OK_COLUMN_NAME.to_string(),
1572                row_stats,
1573            ),
1574            (
1575                SourceDataColumnarEncoder::ERR_COLUMN_NAME.to_string(),
1576                err_stats,
1577            ),
1578        ];
1579        StructStats {
1580            len,
1581            cols: stats.into_iter().map(|(name, s)| (name, s)).collect(),
1582        }
1583    }
1584}
1585
1586/// An encoder for [`Row`]s within [`SourceData`].
1587///
1588/// This type exists as a wrapper around [`RowColumnarEncoder`] to support
1589/// encoding empty [`Row`]s. A [`RowColumnarEncoder`] finishes as a
1590/// [`StructArray`] which is required to have at least one column, and thus
1591/// cannot support empty [`Row`]s.
1592#[derive(Debug)]
1593pub enum SourceDataRowColumnarEncoder {
1594    Row(RowColumnarEncoder),
1595    EmptyRow,
1596}
1597
1598impl SourceDataRowColumnarEncoder {
1599    pub(crate) fn goodbytes(&self) -> usize {
1600        match self {
1601            SourceDataRowColumnarEncoder::Row(e) => e.goodbytes(),
1602            SourceDataRowColumnarEncoder::EmptyRow => 0,
1603        }
1604    }
1605
1606    pub fn append(&mut self, row: &Row) {
1607        match self {
1608            SourceDataRowColumnarEncoder::Row(encoder) => encoder.append(row),
1609            SourceDataRowColumnarEncoder::EmptyRow => {
1610                assert_eq!(row.iter().count(), 0)
1611            }
1612        }
1613    }
1614
1615    pub fn append_null(&mut self) {
1616        match self {
1617            SourceDataRowColumnarEncoder::Row(encoder) => encoder.append_null(),
1618            SourceDataRowColumnarEncoder::EmptyRow => (),
1619        }
1620    }
1621}
1622
1623#[derive(Debug)]
1624pub struct SourceDataColumnarEncoder {
1625    row_encoder: SourceDataRowColumnarEncoder,
1626    err_encoder: BinaryBuilder,
1627}
1628
1629impl SourceDataColumnarEncoder {
1630    const OK_COLUMN_NAME: &'static str = "ok";
1631    const ERR_COLUMN_NAME: &'static str = "err";
1632
1633    pub fn new(desc: &RelationDesc) -> Self {
1634        let row_encoder = match RowColumnarEncoder::new(desc) {
1635            Some(encoder) => SourceDataRowColumnarEncoder::Row(encoder),
1636            None => {
1637                assert!(desc.typ().columns().is_empty());
1638                SourceDataRowColumnarEncoder::EmptyRow
1639            }
1640        };
1641        let err_encoder = BinaryBuilder::new();
1642
1643        SourceDataColumnarEncoder {
1644            row_encoder,
1645            err_encoder,
1646        }
1647    }
1648}
1649
1650impl ColumnEncoder<SourceData> for SourceDataColumnarEncoder {
1651    type FinishedColumn = StructArray;
1652
1653    fn goodbytes(&self) -> usize {
1654        self.row_encoder.goodbytes() + self.err_encoder.values_slice().len()
1655    }
1656
1657    #[inline]
1658    fn append(&mut self, val: &SourceData) {
1659        match val.0.as_ref() {
1660            Ok(row) => {
1661                self.row_encoder.append(row);
1662                self.err_encoder.append_null();
1663            }
1664            Err(err) => {
1665                self.row_encoder.append_null();
1666                self.err_encoder
1667                    .append_value(err.into_proto().encode_to_vec());
1668            }
1669        }
1670    }
1671
1672    #[inline]
1673    fn append_null(&mut self) {
1674        panic!("appending a null into SourceDataColumnarEncoder is not supported");
1675    }
1676
1677    fn finish(self) -> Self::FinishedColumn {
1678        let SourceDataColumnarEncoder {
1679            row_encoder,
1680            mut err_encoder,
1681        } = self;
1682
1683        let err_column = BinaryBuilder::finish(&mut err_encoder);
1684        let row_column: ArrayRef = match row_encoder {
1685            SourceDataRowColumnarEncoder::Row(encoder) => {
1686                let column = encoder.finish();
1687                Arc::new(column)
1688            }
1689            SourceDataRowColumnarEncoder::EmptyRow => Arc::new(NullArray::new(err_column.len())),
1690        };
1691
1692        assert_eq!(row_column.len(), err_column.len());
1693
1694        let fields = vec![
1695            Field::new(Self::OK_COLUMN_NAME, row_column.data_type().clone(), true),
1696            Field::new(Self::ERR_COLUMN_NAME, err_column.data_type().clone(), true),
1697        ];
1698        let arrays: Vec<Arc<dyn Array>> = vec![row_column, Arc::new(err_column)];
1699        StructArray::new(Fields::from(fields), arrays, None)
1700    }
1701}
1702
1703impl Schema<SourceData> for RelationDesc {
1704    type ArrowColumn = StructArray;
1705    type Statistics = StructStats;
1706
1707    type Decoder = SourceDataColumnarDecoder;
1708    type Encoder = SourceDataColumnarEncoder;
1709
1710    fn decoder(&self, col: Self::ArrowColumn) -> Result<Self::Decoder, anyhow::Error> {
1711        SourceDataColumnarDecoder::new(col, self)
1712    }
1713
1714    fn encoder(&self) -> Result<Self::Encoder, anyhow::Error> {
1715        Ok(SourceDataColumnarEncoder::new(self))
1716    }
1717}
1718
1719#[cfg(test)]
1720mod tests {
1721    use arrow::array::{ArrayData, make_comparator};
1722    use base64::Engine;
1723    use bytes::Bytes;
1724    use mz_expr::EvalError;
1725    use mz_ore::assert_err;
1726    use mz_ore::metrics::MetricsRegistry;
1727    use mz_persist::indexed::columnar::arrow::{realloc_any, realloc_array};
1728    use mz_persist::metrics::ColumnarMetrics;
1729    use mz_persist_types::parquet::EncodingConfig;
1730    use mz_persist_types::schema::{Migration, backward_compatible};
1731    use mz_persist_types::stats::{PartStats, PartStatsMetrics};
1732    use mz_repr::{
1733        ColumnIndex, DatumVec, PropRelationDescDiff, ProtoRelationDesc, RelationDescBuilder,
1734        RowArena, SqlScalarType, arb_relation_desc_diff, arb_relation_desc_projection,
1735    };
1736    use proptest::prelude::*;
1737    use proptest::strategy::{Union, ValueTree};
1738
1739    use crate::stats::RelationPartStats;
1740
1741    use super::*;
1742
1743    #[mz_ore::test]
1744    fn test_timeline_parsing() {
1745        assert_eq!(Ok(Timeline::EpochMilliseconds), "M".parse());
1746        assert_eq!(Ok(Timeline::External("JOE".to_string())), "E.JOE".parse());
1747        assert_eq!(Ok(Timeline::User("MIKE".to_string())), "U.MIKE".parse());
1748
1749        assert_err!("Materialize".parse::<Timeline>());
1750        assert_err!("Ejoe".parse::<Timeline>());
1751        assert_err!("Umike".parse::<Timeline>());
1752        assert_err!("Dance".parse::<Timeline>());
1753        assert_err!("".parse::<Timeline>());
1754    }
1755
1756    #[track_caller]
1757    fn roundtrip_source_data(
1758        desc: &RelationDesc,
1759        datas: Vec<SourceData>,
1760        read_desc: &RelationDesc,
1761        config: &EncodingConfig,
1762    ) {
1763        let metrics = ColumnarMetrics::disconnected();
1764        let mut encoder = <RelationDesc as Schema<SourceData>>::encoder(desc).unwrap();
1765        for data in &datas {
1766            encoder.append(data);
1767        }
1768        let col = encoder.finish();
1769
1770        // The top-level StructArray for SourceData should always be non-nullable.
1771        assert!(!col.is_nullable());
1772
1773        // Reallocate our arrays with lgalloc.
1774        let col = realloc_array(&col, &metrics);
1775
1776        // Roundtrip through ProtoArray format.
1777        {
1778            let proto = col.to_data().into_proto();
1779            let bytes = proto.encode_to_vec();
1780            let proto = mz_persist_types::arrow::ProtoArrayData::decode(&bytes[..]).unwrap();
1781            let array_data: ArrayData = proto.into_rust().unwrap();
1782
1783            let col_rnd = StructArray::from(array_data.clone());
1784            assert_eq!(col, col_rnd);
1785
1786            let col_dyn = arrow::array::make_array(array_data);
1787            let col_dyn = col_dyn.as_any().downcast_ref::<StructArray>().unwrap();
1788            assert_eq!(&col, col_dyn);
1789        }
1790
1791        // Encode to Parquet.
1792        let mut buf = Vec::new();
1793        let fields = Fields::from(vec![Field::new("k", col.data_type().clone(), false)]);
1794        let arrays: Vec<Arc<dyn Array>> = vec![Arc::new(col.clone())];
1795        mz_persist_types::parquet::encode_arrays(&mut buf, fields, arrays, config).unwrap();
1796
1797        // Decode from Parquet.
1798        let buf = Bytes::from(buf);
1799        let mut reader = mz_persist_types::parquet::decode_arrays(buf).unwrap();
1800        let maybe_batch = reader.next();
1801
1802        // If we didn't encode any data then our record_batch will be empty.
1803        let Some(record_batch) = maybe_batch else {
1804            assert!(datas.is_empty());
1805            return;
1806        };
1807        let record_batch = record_batch.unwrap();
1808
1809        assert_eq!(record_batch.columns().len(), 1);
1810        let rnd_col = &record_batch.columns()[0];
1811        let rnd_col = realloc_any(Arc::clone(rnd_col), &metrics);
1812        let rnd_col = rnd_col
1813            .as_any()
1814            .downcast_ref::<StructArray>()
1815            .unwrap()
1816            .clone();
1817
1818        // Try generating stats for the data, just to make sure we don't panic.
1819        let stats = <RelationDesc as Schema<SourceData>>::decoder_any(desc, &rnd_col)
1820            .expect("valid decoder")
1821            .stats();
1822
1823        // Read back all of our data and assert it roundtrips.
1824        let mut rnd_data = SourceData(Ok(Row::default()));
1825        let decoder = <RelationDesc as Schema<SourceData>>::decoder(desc, rnd_col.clone()).unwrap();
1826        for (idx, og_data) in datas.iter().enumerate() {
1827            decoder.decode(idx, &mut rnd_data);
1828            assert_eq!(og_data, &rnd_data);
1829        }
1830
1831        // Read back all of our data a second time with a projection applied, and make sure the
1832        // stats are valid.
1833        let stats_metrics = PartStatsMetrics::new(&MetricsRegistry::new());
1834        let stats = RelationPartStats {
1835            name: "test",
1836            metrics: &stats_metrics,
1837            stats: &PartStats { key: stats },
1838            desc: read_desc,
1839        };
1840        let mut datum_vec = DatumVec::new();
1841        let arena = RowArena::default();
1842        let decoder = <RelationDesc as Schema<SourceData>>::decoder(read_desc, rnd_col).unwrap();
1843
1844        for (idx, og_data) in datas.iter().enumerate() {
1845            decoder.decode(idx, &mut rnd_data);
1846            match (&og_data.0, &rnd_data.0) {
1847                (Ok(og_row), Ok(rnd_row)) => {
1848                    // Filter down to just the Datums in the projection schema.
1849                    {
1850                        let datums = datum_vec.borrow_with(og_row);
1851                        let projected_datums =
1852                            datums.iter().enumerate().filter_map(|(idx, datum)| {
1853                                read_desc
1854                                    .contains_index(&ColumnIndex::from_raw(idx))
1855                                    .then_some(datum)
1856                            });
1857                        let og_projected_row = Row::pack(projected_datums);
1858                        assert_eq!(&og_projected_row, rnd_row);
1859                    }
1860
1861                    // Validate the stats for all of our projected columns.
1862                    {
1863                        let proj_datums = datum_vec.borrow_with(rnd_row);
1864                        for (pos, (idx, _, _)) in read_desc.iter_all().enumerate() {
1865                            let spec = stats.col_stats(idx, &arena);
1866                            assert!(spec.may_contain(proj_datums[pos]));
1867                        }
1868                    }
1869                }
1870                (Err(_), Err(_)) => assert_eq!(og_data, &rnd_data),
1871                (_, _) => panic!("decoded to a different type? {og_data:?} {rnd_data:?}"),
1872            }
1873        }
1874
1875        // Verify that the RelationDesc itself roundtrips through
1876        // {encode,decode}_schema.
1877        let encoded_schema = SourceData::encode_schema(desc);
1878        let roundtrip_desc = SourceData::decode_schema(&encoded_schema);
1879        assert_eq!(desc, &roundtrip_desc);
1880
1881        // Verify that the RelationDesc is backward compatible with itself (this
1882        // mostly checks for `unimplemented!` type panics).
1883        let migration =
1884            mz_persist_types::schema::backward_compatible(col.data_type(), col.data_type());
1885        let migration = migration.expect("should be backward compatible with self");
1886        // Also verify that the Fn doesn't do anything wonky.
1887        let migrated = migration.migrate(Arc::new(col.clone()));
1888        assert_eq!(col.data_type(), migrated.data_type());
1889    }
1890
1891    #[mz_ore::test]
1892    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decContextDefault` on OS `linux`
1893    fn all_source_data_roundtrips() {
1894        let mut weights = vec![(500, Just(0..8)), (50, Just(8..32))];
1895        if std::env::var("PROPTEST_LARGE_DATA").is_ok() {
1896            weights.extend([
1897                (10, Just(32..128)),
1898                (5, Just(128..512)),
1899                (3, Just(512..2048)),
1900                (1, Just(2048..8192)),
1901            ]);
1902        }
1903        let num_rows = Union::new_weighted(weights);
1904
1905        // TODO(parkmycar): There are so many clones going on here, and maybe we can avoid them?
1906        let strat = (any::<RelationDesc>(), num_rows)
1907            .prop_flat_map(|(desc, num_rows)| {
1908                arb_relation_desc_projection(desc.clone())
1909                    .prop_map(move |read_desc| (desc.clone(), read_desc, num_rows.clone()))
1910            })
1911            .prop_flat_map(|(desc, read_desc, num_rows)| {
1912                proptest::collection::vec(arb_source_data_for_relation_desc(&desc), num_rows)
1913                    .prop_map(move |datas| (desc.clone(), datas, read_desc.clone()))
1914            });
1915
1916        let combined_strat = (any::<EncodingConfig>(), strat);
1917        proptest!(|((config, (desc, source_datas, read_desc)) in combined_strat)| {
1918            roundtrip_source_data(&desc, source_datas, &read_desc, &config);
1919        });
1920    }
1921
1922    #[mz_ore::test]
1923    fn roundtrip_error_nulls() {
1924        let desc = RelationDescBuilder::default()
1925            .with_column(
1926                "ts",
1927                SqlScalarType::TimestampTz { precision: None }.nullable(false),
1928            )
1929            .finish();
1930        let source_datas = vec![SourceData(Err(DataflowError::EvalError(
1931            EvalError::DateOutOfRange.into(),
1932        )))];
1933        let config = EncodingConfig::default();
1934        roundtrip_source_data(&desc, source_datas, &desc, &config);
1935    }
1936
1937    fn is_sorted(array: &dyn Array) -> bool {
1938        let sort_options = arrow::compute::SortOptions::default();
1939        let Ok(cmp) = make_comparator(array, array, sort_options) else {
1940            // TODO: arrow v51.0.0 doesn't support comparing structs. When
1941            // we migrate to v52+, the `build_compare` function is
1942            // deprecated and replaced by `make_comparator`, which does
1943            // support structs. At which point, this will work (and we
1944            // should switch this early return to an expect, if possible).
1945            return false;
1946        };
1947        (0..array.len())
1948            .tuple_windows()
1949            .all(|(i, j)| cmp(i, j).is_le())
1950    }
1951
1952    fn get_data_type(schema: &impl Schema<SourceData>) -> arrow::datatypes::DataType {
1953        use mz_persist_types::columnar::ColumnEncoder;
1954        let array = Schema::encoder(schema).expect("valid schema").finish();
1955        Array::data_type(&array).clone()
1956    }
1957
1958    #[track_caller]
1959    fn backward_compatible_testcase(
1960        old: &RelationDesc,
1961        new: &RelationDesc,
1962        migration: Migration,
1963        datas: &[SourceData],
1964    ) {
1965        let mut encoder = Schema::<SourceData>::encoder(old).expect("valid schema");
1966        for data in datas {
1967            encoder.append(data);
1968        }
1969        let old = encoder.finish();
1970        let new = Schema::<SourceData>::encoder(new)
1971            .expect("valid schema")
1972            .finish();
1973        let old: Arc<dyn Array> = Arc::new(old);
1974        let new: Arc<dyn Array> = Arc::new(new);
1975        let migrated = migration.migrate(Arc::clone(&old));
1976        assert_eq!(migrated.data_type(), new.data_type());
1977
1978        // Check the sortedness preservation, if we can.
1979        if migration.preserves_order() && is_sorted(&old) {
1980            assert!(is_sorted(&new))
1981        }
1982    }
1983
1984    #[mz_ore::test]
1985    fn backward_compatible_empty_add_column() {
1986        let old = RelationDesc::empty();
1987        let new = RelationDesc::from_names_and_types([("a", SqlScalarType::Bool.nullable(true))]);
1988
1989        let old_data_type = get_data_type(&old);
1990        let new_data_type = get_data_type(&new);
1991
1992        let migration = backward_compatible(&old_data_type, &new_data_type);
1993        assert!(migration.is_some());
1994    }
1995
1996    #[mz_ore::test]
1997    fn backward_compatible_project_away_all() {
1998        let old = RelationDesc::from_names_and_types([("a", SqlScalarType::Bool.nullable(true))]);
1999        let new = RelationDesc::empty();
2000
2001        let old_data_type = get_data_type(&old);
2002        let new_data_type = get_data_type(&new);
2003
2004        let migration = backward_compatible(&old_data_type, &new_data_type);
2005        assert!(migration.is_some());
2006    }
2007
2008    #[mz_ore::test]
2009    #[cfg_attr(miri, ignore)]
2010    fn backward_compatible_migrate() {
2011        let strat = (any::<RelationDesc>(), any::<RelationDesc>()).prop_flat_map(|(old, new)| {
2012            proptest::collection::vec(arb_source_data_for_relation_desc(&old), 2)
2013                .prop_map(move |datas| (old.clone(), new.clone(), datas))
2014        });
2015
2016        proptest!(|((old, new, datas) in strat)| {
2017            let old_data_type = get_data_type(&old);
2018            let new_data_type = get_data_type(&new);
2019
2020            if let Some(migration) = backward_compatible(&old_data_type, &new_data_type) {
2021                backward_compatible_testcase(&old, &new, migration, &datas);
2022            };
2023        });
2024    }
2025
2026    #[mz_ore::test]
2027    #[cfg_attr(miri, ignore)]
2028    fn backward_compatible_migrate_from_common() {
2029        use mz_repr::SqlColumnType;
2030        fn test_case(old: RelationDesc, diffs: Vec<PropRelationDescDiff>, datas: Vec<SourceData>) {
2031            // TODO(parkmycar): As we iterate on schema migrations more things should become compatible.
2032            let should_be_compatible = diffs.iter().all(|diff| match diff {
2033                // We only support adding nullable columns.
2034                PropRelationDescDiff::AddColumn {
2035                    typ: SqlColumnType { nullable, .. },
2036                    ..
2037                } => *nullable,
2038                PropRelationDescDiff::DropColumn { .. } => true,
2039                _ => false,
2040            });
2041
2042            let mut new = old.clone();
2043            for diff in diffs.into_iter() {
2044                diff.apply(&mut new)
2045            }
2046
2047            let old_data_type = get_data_type(&old);
2048            let new_data_type = get_data_type(&new);
2049
2050            if let Some(migration) = backward_compatible(&old_data_type, &new_data_type) {
2051                backward_compatible_testcase(&old, &new, migration, &datas);
2052            } else if should_be_compatible {
2053                panic!("new DataType was not compatible when it should have been!");
2054            }
2055        }
2056
2057        let strat = any::<RelationDesc>()
2058            .prop_flat_map(|desc| {
2059                proptest::collection::vec(arb_source_data_for_relation_desc(&desc), 2)
2060                    .no_shrink()
2061                    .prop_map(move |datas| (desc.clone(), datas))
2062            })
2063            .prop_flat_map(|(desc, datas)| {
2064                arb_relation_desc_diff(&desc)
2065                    .prop_map(move |diffs| (desc.clone(), diffs, datas.clone()))
2066            });
2067
2068        proptest!(|((old, diffs, datas) in strat)| {
2069            test_case(old, diffs, datas);
2070        });
2071    }
2072
2073    #[mz_ore::test]
2074    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decContextDefault` on OS `linux`
2075    fn empty_relation_desc_roundtrips() {
2076        let empty = RelationDesc::empty();
2077        let rows = proptest::collection::vec(arb_source_data_for_relation_desc(&empty), 0..8)
2078            .prop_map(move |datas| (empty.clone(), datas));
2079
2080        // Note: This case should be covered by the `all_source_data_roundtrips` test above, but
2081        // it's a special case that we explicitly want to exercise.
2082        proptest!(|((config, (desc, source_datas)) in (any::<EncodingConfig>(), rows))| {
2083            roundtrip_source_data(&desc, source_datas, &desc, &config);
2084        });
2085    }
2086
2087    #[mz_ore::test]
2088    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decContextDefault` on OS `linux`
2089    fn arrow_datatype_consistent() {
2090        fn test_case(desc: RelationDesc, datas: Vec<SourceData>) {
2091            let half = datas.len() / 2;
2092
2093            let mut encoder_a = <RelationDesc as Schema<SourceData>>::encoder(&desc).unwrap();
2094            for data in &datas[..half] {
2095                encoder_a.append(data);
2096            }
2097            let col_a = encoder_a.finish();
2098
2099            let mut encoder_b = <RelationDesc as Schema<SourceData>>::encoder(&desc).unwrap();
2100            for data in &datas[half..] {
2101                encoder_b.append(data);
2102            }
2103            let col_b = encoder_b.finish();
2104
2105            // The DataType of the resulting column should not change based on what data was
2106            // encoded.
2107            assert_eq!(col_a.data_type(), col_b.data_type());
2108        }
2109
2110        let num_rows = 12;
2111        let strat = any::<RelationDesc>().prop_flat_map(|desc| {
2112            proptest::collection::vec(arb_source_data_for_relation_desc(&desc), num_rows)
2113                .prop_map(move |datas| (desc.clone(), datas))
2114        });
2115
2116        proptest!(|((desc, data) in strat)| {
2117            test_case(desc, data);
2118        });
2119    }
2120
2121    #[mz_ore::test]
2122    #[cfg_attr(miri, ignore)] // too slow
2123    fn source_proto_serialization_stability() {
2124        let min_protos = 10;
2125        let encoded = include_str!("snapshots/source-datas.txt");
2126
2127        // Decode the pre-generated source datas
2128        let mut decoded: Vec<(RelationDesc, SourceData)> = encoded
2129            .lines()
2130            .map(|s| {
2131                let (desc, data) = s.split_once(',').expect("comma separated data");
2132                let desc = base64::engine::general_purpose::STANDARD
2133                    .decode(desc)
2134                    .expect("valid base64");
2135                let data = base64::engine::general_purpose::STANDARD
2136                    .decode(data)
2137                    .expect("valid base64");
2138                (desc, data)
2139            })
2140            .map(|(desc, data)| {
2141                let desc = ProtoRelationDesc::decode(&desc[..]).expect("valid proto");
2142                let desc = desc.into_rust().expect("valid proto");
2143                let data = SourceData::decode(&data, &desc).expect("valid proto");
2144                (desc, data)
2145            })
2146            .collect();
2147
2148        // If there are fewer than the minimum examples, generate some new ones arbitrarily
2149        let mut runner = proptest::test_runner::TestRunner::deterministic();
2150        let strategy = RelationDesc::arbitrary().prop_flat_map(|desc| {
2151            arb_source_data_for_relation_desc(&desc).prop_map(move |data| (desc.clone(), data))
2152        });
2153        while decoded.len() < min_protos {
2154            let arbitrary_data = strategy
2155                .new_tree(&mut runner)
2156                .expect("source data")
2157                .current();
2158            decoded.push(arbitrary_data);
2159        }
2160
2161        // Reencode and compare the strings
2162        let mut reencoded = String::new();
2163        let mut buf = vec![];
2164        for (desc, data) in decoded {
2165            buf.clear();
2166            desc.into_proto().encode(&mut buf).expect("success");
2167            base64::engine::general_purpose::STANDARD.encode_string(buf.as_slice(), &mut reencoded);
2168            reencoded.push(',');
2169
2170            buf.clear();
2171            data.encode(&mut buf);
2172            base64::engine::general_purpose::STANDARD.encode_string(buf.as_slice(), &mut reencoded);
2173            reencoded.push('\n');
2174        }
2175
2176        // Optimizations in Persist, particularly consolidation on read,
2177        // depend on a stable serialization for the serialized data.
2178        // For example, reordering proto fields could cause us
2179        // to generate a different (equivalent) serialization for a record,
2180        // and the two versions would not consolidate out.
2181        // This can impact correctness!
2182        //
2183        // If you need to change how SourceDatas are encoded, that's still fine...
2184        // but we'll also need to increase
2185        // the MINIMUM_CONSOLIDATED_VERSION as part of the same release.
2186        assert_eq!(
2187            encoded,
2188            reencoded.as_str(),
2189            "SourceData serde should be stable"
2190        )
2191    }
2192}