1use 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#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
84pub struct IngestionDescription<S: 'static = (), C: ConnectionAccess = InlinedConnection> {
85 pub desc: SourceDesc<C>,
87 pub source_exports: BTreeMap<GlobalId, SourceExport<S>>,
101 pub instance_id: StorageInstanceId,
103 pub remap_collection_id: GlobalId,
105 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 pub fn collection_ids(&self) -> impl Iterator<Item = GlobalId> + '_ {
131 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 pub storage_metadata: S,
249 pub details: SourceExportDetails,
251 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#[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
336impl 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
350impl 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
380impl 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#[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,
451 External(String),
455 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
511pub trait SourceConnection: Debug + Clone + PartialEq + AlterCompatible {
513 fn name(&self) -> &'static str;
515
516 fn external_reference(&self) -> Option<&str>;
518
519 fn default_key_desc(&self) -> RelationDesc;
523
524 fn default_value_desc(&self) -> RelationDesc;
528
529 fn timestamp_desc(&self) -> RelationDesc;
532
533 fn connection_id(&self) -> Option<CatalogItemId>;
536
537 fn supports_read_only(&self) -> bool;
539
540 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#[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 pub fn monotonic(&self, connection: &GenericSourceConnection<C>) -> bool {
612 match &self.envelope {
613 SourceEnvelope::Upsert(_) | SourceEnvelope::CdcV2 => false,
615 SourceEnvelope::None(_) => {
616 match connection {
617 GenericSourceConnection::Postgres(_) => false,
619 GenericSourceConnection::MySql(_) => false,
621 GenericSourceConnection::SqlServer(_) => false,
623 GenericSourceConnection::LoadGenerator(g) => g.load_generator.is_monotonic(),
625 GenericSourceConnection::Kafka(_) => true,
627 }
628 }
629 }
630 }
631}
632
633#[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 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: _,
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#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
866pub enum SourceExportDetails {
867 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
904pub enum SourceExportStatementDetails {
910 Postgres {
911 table: mz_postgres_util::desc::PostgresTableDesc,
912 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 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 (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 (kind, _) => {
1123 let proto = ProtoSourceData { kind };
1124 *self = proto.into_rust().map_err(|err| err.to_string())?;
1125 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#[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
1167pub trait ExternalCatalogReference {
1175 fn schema_name(&self) -> &str;
1177 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
1211impl<'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#[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 pub fn new<T: ExternalCatalogReference>(
1255 database: &str,
1256 referenceable_items: &'a [T],
1257 ) -> Result<SourceReferenceResolver, ExternalReferenceResolutionError> {
1258 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 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 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 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 if !(1..=3).contains(&name.len()) {
1353 Err(ExternalReferenceResolutionError::DoesNotExist {
1354 name: get_provided_name(),
1355 })?;
1356 }
1357
1358 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#[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 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 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 let row_null_count = len - self.err_decoder.null_count();
1549 let row_stats = match &self.row_decoder {
1550 SourceDataRowColumnarDecoder::Row(encoder) => {
1551 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#[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 assert!(!col.is_nullable());
1772
1773 let col = realloc_array(&col, &metrics);
1775
1776 {
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 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 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 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 let stats = <RelationDesc as Schema<SourceData>>::decoder_any(desc, &rnd_col)
1820 .expect("valid decoder")
1821 .stats();
1822
1823 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 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 {
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 {
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 let encoded_schema = SourceData::encode_schema(desc);
1878 let roundtrip_desc = SourceData::decode_schema(&encoded_schema);
1879 assert_eq!(desc, &roundtrip_desc);
1880
1881 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 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)] 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 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 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 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 let should_be_compatible = diffs.iter().all(|diff| match diff {
2033 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)] 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 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)] 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 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)] fn source_proto_serialization_stability() {
2124 let min_protos = 10;
2125 let encoded = include_str!("snapshots/source-datas.txt");
2126
2127 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 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 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 assert_eq!(
2187 encoded,
2188 reencoded.as_str(),
2189 "SourceData serde should be stable"
2190 )
2191 }
2192}