1use std::borrow::Cow;
13use std::fmt::Debug;
14use std::time::Duration;
15
16use mz_dyncfg::ConfigSet;
17use mz_expr::MirScalarExpr;
18use mz_pgcopy::CopyFormatParams;
19use mz_repr::bytes::ByteSize;
20use mz_repr::{CatalogItemId, GlobalId, RelationDesc};
21#[cfg(any(test, feature = "proptest"))]
22use proptest_derive::Arbitrary;
23use serde::{Deserialize, Serialize};
24use timely::PartialOrder;
25use timely::progress::frontier::Antichain;
26
27use crate::AlterCompatible;
28use crate::connections::inline::{
29 ConnectionAccess, ConnectionResolver, InlinedConnection, IntoInlineConnection,
30 ReferencedConnection,
31};
32use crate::connections::{ConnectionContext, KafkaConnection, KafkaTopicOptions};
33use crate::controller::AlterError;
34use crate::wire_format::WireFormat;
35
36pub mod s3_oneshot_sink;
37
38#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
40pub struct StorageSinkDesc<S, T = mz_repr::Timestamp> {
41 pub from: GlobalId,
42 pub from_desc: RelationDesc,
43 pub connection: StorageSinkConnection,
44 pub with_snapshot: bool,
45 pub version: u64,
46 pub envelope: SinkEnvelope,
47 pub as_of: Antichain<T>,
48 pub from_storage_metadata: S,
49 pub to_storage_metadata: S,
50 pub commit_interval: Option<Duration>,
56}
57
58impl<S: Debug + PartialEq, T: Debug + PartialEq + PartialOrder> AlterCompatible
59 for StorageSinkDesc<S, T>
60{
61 fn alter_compatible(
69 &self,
70 id: GlobalId,
71 other: &StorageSinkDesc<S, T>,
72 ) -> Result<(), AlterError> {
73 if self == other {
74 return Ok(());
75 }
76 let StorageSinkDesc {
77 from,
78 from_desc,
79 connection,
80 envelope,
81 version: _,
82 as_of: _,
84 from_storage_metadata,
85 with_snapshot,
86 to_storage_metadata,
87 commit_interval: _,
88 } = self;
89
90 let compatibility_checks = [
91 (from == &other.from, "from"),
92 (from_desc == &other.from_desc, "from_desc"),
93 (
94 connection.alter_compatible(id, &other.connection).is_ok(),
95 "connection",
96 ),
97 (envelope == &other.envelope, "envelope"),
98 (*with_snapshot || !other.with_snapshot, "with_snapshot"),
101 (
102 from_storage_metadata == &other.from_storage_metadata,
103 "from_storage_metadata",
104 ),
105 (
106 to_storage_metadata == &other.to_storage_metadata,
107 "to_storage_metadata",
108 ),
109 ];
110
111 for (compatible, field) in compatibility_checks {
112 if !compatible {
113 tracing::warn!(
114 "StorageSinkDesc incompatible at {field}:\nself:\n{:#?}\n\nother\n{:#?}",
115 self,
116 other
117 );
118
119 return Err(AlterError { id });
120 }
121 }
122
123 Ok(())
124 }
125}
126
127#[derive(Copy, Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
128pub enum SinkEnvelope {
129 Debezium,
131 Upsert,
132 Append,
134}
135
136#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
137pub enum StorageSinkConnection<C: ConnectionAccess = InlinedConnection> {
138 Kafka(KafkaSinkConnection<C>),
139 Iceberg(IcebergSinkConnection<C>),
140}
141
142impl<C: ConnectionAccess> StorageSinkConnection<C> {
143 pub fn alter_compatible(
148 &self,
149 id: GlobalId,
150 other: &StorageSinkConnection<C>,
151 ) -> Result<(), AlterError> {
152 if self == other {
153 return Ok(());
154 }
155 match (self, other) {
156 (StorageSinkConnection::Kafka(s), StorageSinkConnection::Kafka(o)) => {
157 s.alter_compatible(id, o)?
158 }
159 (StorageSinkConnection::Iceberg(s), StorageSinkConnection::Iceberg(o)) => {
160 s.alter_compatible(id, o)?
161 }
162 _ => {
163 tracing::warn!(
164 "StorageSinkConnection incompatible:\nself:\n{:#?}\n\nother\n{:#?}",
165 self,
166 other
167 );
168 return Err(AlterError { id });
169 }
170 }
171
172 Ok(())
173 }
174}
175
176impl<R: ConnectionResolver> IntoInlineConnection<StorageSinkConnection, R>
177 for StorageSinkConnection<ReferencedConnection>
178{
179 fn into_inline_connection(self, r: R) -> StorageSinkConnection {
180 match self {
181 Self::Kafka(conn) => StorageSinkConnection::Kafka(conn.into_inline_connection(r)),
182 Self::Iceberg(conn) => StorageSinkConnection::Iceberg(conn.into_inline_connection(r)),
183 }
184 }
185}
186
187impl<C: ConnectionAccess> StorageSinkConnection<C> {
188 pub fn connection_id(&self) -> Option<CatalogItemId> {
193 use StorageSinkConnection::*;
194 match self {
195 Kafka(KafkaSinkConnection { connection_id, .. }) => Some(*connection_id),
196 Iceberg(IcebergSinkConnection {
197 catalog_connection_id: connection_id,
198 ..
199 }) => Some(*connection_id),
200 }
201 }
202
203 pub fn name(&self) -> &'static str {
208 use StorageSinkConnection::*;
209 match self {
210 Kafka(_) => "kafka",
211 Iceberg(_) => "iceberg",
212 }
213 }
214}
215
216#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
217pub enum KafkaSinkCompressionType {
218 None,
219 Gzip,
220 Snappy,
221 Lz4,
222 Zstd,
223}
224
225impl KafkaSinkCompressionType {
226 pub fn to_librdkafka_option(&self) -> &'static str {
229 match self {
230 KafkaSinkCompressionType::None => "none",
231 KafkaSinkCompressionType::Gzip => "gzip",
232 KafkaSinkCompressionType::Snappy => "snappy",
233 KafkaSinkCompressionType::Lz4 => "lz4",
234 KafkaSinkCompressionType::Zstd => "zstd",
235 }
236 }
237}
238
239#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
240pub struct KafkaSinkConnection<C: ConnectionAccess = InlinedConnection> {
241 pub connection_id: CatalogItemId,
242 pub connection: C::Kafka,
243 pub format: KafkaSinkFormat<C>,
244 pub relation_key_indices: Option<Vec<usize>>,
246 pub key_desc_and_indices: Option<(RelationDesc, Vec<usize>)>,
248 pub headers_index: Option<usize>,
250 pub value_desc: RelationDesc,
251 pub partition_by: Option<MirScalarExpr>,
254 pub topic: String,
255 pub topic_options: KafkaTopicOptions,
257 pub compression_type: KafkaSinkCompressionType,
258 pub progress_group_id: KafkaIdStyle,
259 pub transactional_id: KafkaIdStyle,
260 pub topic_metadata_refresh_interval: Duration,
261}
262
263impl KafkaSinkConnection {
264 pub fn client_id(
269 &self,
270 configs: &ConfigSet,
271 connection_context: &ConnectionContext,
272 sink_id: GlobalId,
273 ) -> String {
274 let mut client_id =
275 KafkaConnection::id_base(connection_context, self.connection_id, sink_id);
276 self.connection.enrich_client_id(configs, &mut client_id);
277 client_id
278 }
279
280 pub fn progress_topic(&self, connection_context: &ConnectionContext) -> Cow<'_, str> {
282 self.connection
283 .progress_topic(connection_context, self.connection_id)
284 }
285
286 pub fn progress_group_id(
292 &self,
293 connection_context: &ConnectionContext,
294 sink_id: GlobalId,
295 ) -> String {
296 match self.progress_group_id {
297 KafkaIdStyle::Prefix(ref prefix) => format!(
298 "{}{}",
299 prefix.as_deref().unwrap_or(""),
300 KafkaConnection::id_base(connection_context, self.connection_id, sink_id),
301 ),
302 KafkaIdStyle::Legacy => format!("materialize-bootstrap-sink-{sink_id}"),
303 }
304 }
305
306 pub fn transactional_id(
311 &self,
312 connection_context: &ConnectionContext,
313 sink_id: GlobalId,
314 ) -> String {
315 match self.transactional_id {
316 KafkaIdStyle::Prefix(ref prefix) => format!(
317 "{}{}",
318 prefix.as_deref().unwrap_or(""),
319 KafkaConnection::id_base(connection_context, self.connection_id, sink_id)
320 ),
321 KafkaIdStyle::Legacy => format!("mz-producer-{sink_id}-0"),
322 }
323 }
324}
325
326impl<C: ConnectionAccess> KafkaSinkConnection<C> {
327 pub fn alter_compatible(
332 &self,
333 id: GlobalId,
334 other: &KafkaSinkConnection<C>,
335 ) -> Result<(), AlterError> {
336 if self == other {
337 return Ok(());
338 }
339 let KafkaSinkConnection {
340 connection_id,
341 connection,
342 format,
343 relation_key_indices,
344 key_desc_and_indices,
345 headers_index,
346 value_desc,
347 partition_by,
348 topic,
349 compression_type,
350 progress_group_id,
351 transactional_id,
352 topic_options,
353 topic_metadata_refresh_interval,
354 } = self;
355
356 let compatibility_checks = [
357 (connection_id == &other.connection_id, "connection_id"),
358 (
359 connection.alter_compatible(id, &other.connection).is_ok(),
360 "connection",
361 ),
362 (format.alter_compatible(id, &other.format).is_ok(), "format"),
363 (
364 relation_key_indices == &other.relation_key_indices,
365 "relation_key_indices",
366 ),
367 (
368 key_desc_and_indices == &other.key_desc_and_indices,
369 "key_desc_and_indices",
370 ),
371 (headers_index == &other.headers_index, "headers_index"),
372 (value_desc == &other.value_desc, "value_desc"),
373 (partition_by == &other.partition_by, "partition_by"),
374 (topic == &other.topic, "topic"),
375 (
376 compression_type == &other.compression_type,
377 "compression_type",
378 ),
379 (
380 progress_group_id == &other.progress_group_id,
381 "progress_group_id",
382 ),
383 (
384 transactional_id == &other.transactional_id,
385 "transactional_id",
386 ),
387 (topic_options == &other.topic_options, "topic_config"),
388 (
389 topic_metadata_refresh_interval == &other.topic_metadata_refresh_interval,
390 "topic_metadata_refresh_interval",
391 ),
392 ];
393 for (compatible, field) in compatibility_checks {
394 if !compatible {
395 tracing::warn!(
396 "KafkaSinkConnection incompatible at {field}:\nself:\n{:#?}\n\nother\n{:#?}",
397 self,
398 other
399 );
400
401 return Err(AlterError { id });
402 }
403 }
404
405 Ok(())
406 }
407}
408
409impl<R: ConnectionResolver> IntoInlineConnection<KafkaSinkConnection, R>
410 for KafkaSinkConnection<ReferencedConnection>
411{
412 fn into_inline_connection(self, r: R) -> KafkaSinkConnection {
413 let KafkaSinkConnection {
414 connection_id,
415 connection,
416 format,
417 relation_key_indices,
418 key_desc_and_indices,
419 headers_index,
420 value_desc,
421 partition_by,
422 topic,
423 compression_type,
424 progress_group_id,
425 transactional_id,
426 topic_options,
427 topic_metadata_refresh_interval,
428 } = self;
429 KafkaSinkConnection {
430 connection_id,
431 connection: r.resolve_connection(connection).unwrap_kafka(),
432 format: format.into_inline_connection(r),
433 relation_key_indices,
434 key_desc_and_indices,
435 headers_index,
436 value_desc,
437 partition_by,
438 topic,
439 compression_type,
440 progress_group_id,
441 transactional_id,
442 topic_options,
443 topic_metadata_refresh_interval,
444 }
445 }
446}
447
448#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
449pub enum KafkaIdStyle {
450 Prefix(Option<String>),
452 Legacy,
454}
455
456#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
457pub struct KafkaSinkFormat<C: ConnectionAccess = InlinedConnection> {
458 pub key_format: Option<KafkaSinkFormatType<C>>,
459 pub value_format: KafkaSinkFormatType<C>,
460}
461
462#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
463pub enum KafkaSinkFormatType<C: ConnectionAccess = InlinedConnection> {
464 Avro {
465 schema: String,
466 compatibility_level: Option<mz_ccsr::CompatibilityLevel>,
467 #[serde(default)]
472 schema_name: Option<String>,
473 wire_format: WireFormat<C>,
476 },
477 Json,
478 Text,
479 Bytes,
480}
481
482impl<C: ConnectionAccess> KafkaSinkFormatType<C> {
483 pub fn get_format_name(&self) -> &str {
484 match self {
485 Self::Avro { .. } => "avro",
486 Self::Json => "json",
487 Self::Text => "text",
488 Self::Bytes => "bytes",
489 }
490 }
491}
492
493impl<C: ConnectionAccess> KafkaSinkFormat<C> {
494 pub fn get_format_name<'a>(&'a self) -> Cow<'a, str> {
497 match &self.key_format {
501 None => self.value_format.get_format_name().into(),
502 Some(key_format) => match (key_format, &self.value_format) {
503 (KafkaSinkFormatType::Avro { .. }, KafkaSinkFormatType::Avro { .. }) => {
504 "avro".into()
505 }
506 (KafkaSinkFormatType::Json, KafkaSinkFormatType::Json) => "json".into(),
507 (keyf, valuef) => format!(
508 "key-{}-value-{}",
509 keyf.get_format_name(),
510 valuef.get_format_name()
511 )
512 .into(),
513 },
514 }
515 }
516
517 fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
518 if self == other {
519 return Ok(());
520 }
521
522 match (&self.value_format, &other.value_format) {
523 (
524 KafkaSinkFormatType::Avro {
525 schema,
526 compatibility_level: _,
527 schema_name,
528 wire_format,
529 },
530 KafkaSinkFormatType::Avro {
531 schema: other_schema,
532 compatibility_level: _,
533 schema_name: other_schema_name,
534 wire_format: other_wire_format,
535 },
536 ) => {
537 if schema != other_schema
538 || schema_name != other_schema_name
539 || wire_format.alter_compatible(id, other_wire_format).is_err()
540 {
541 tracing::warn!(
542 "KafkaSinkFormat::Avro incompatible at value_format:\nself:\n{:#?}\n\nother\n{:#?}",
543 self,
544 other
545 );
546
547 return Err(AlterError { id });
548 }
549 }
550 (s, o) => {
551 if s != o {
552 tracing::warn!(
553 "KafkaSinkFormat incompatible at value_format:\nself:\n{:#?}\n\nother:{:#?}",
554 s,
555 o
556 );
557 return Err(AlterError { id });
558 }
559 }
560 }
561
562 match (&self.key_format, &other.key_format) {
563 (
564 Some(KafkaSinkFormatType::Avro {
565 schema,
566 compatibility_level: _,
567 schema_name,
568 wire_format,
569 }),
570 Some(KafkaSinkFormatType::Avro {
571 schema: other_schema,
572 compatibility_level: _,
573 schema_name: other_schema_name,
574 wire_format: other_wire_format,
575 }),
576 ) => {
577 if schema != other_schema
578 || schema_name != other_schema_name
579 || wire_format.alter_compatible(id, other_wire_format).is_err()
580 {
581 tracing::warn!(
582 "KafkaSinkFormat::Avro incompatible at key_format:\nself:\n{:#?}\n\nother\n{:#?}",
583 self,
584 other
585 );
586
587 return Err(AlterError { id });
588 }
589 }
590 (s, o) => {
591 if s != o {
592 tracing::warn!(
593 "KafkaSinkFormat incompatible at key_format\nself:\n{:#?}\n\nother:{:#?}",
594 s,
595 o
596 );
597 return Err(AlterError { id });
598 }
599 }
600 }
601
602 Ok(())
603 }
604}
605
606impl<R: ConnectionResolver> IntoInlineConnection<KafkaSinkFormat, R>
607 for KafkaSinkFormat<ReferencedConnection>
608{
609 fn into_inline_connection(self, r: R) -> KafkaSinkFormat {
610 KafkaSinkFormat {
611 key_format: self.key_format.map(|f| f.into_inline_connection(&r)),
612 value_format: self.value_format.into_inline_connection(&r),
613 }
614 }
615}
616
617impl<R: ConnectionResolver> IntoInlineConnection<KafkaSinkFormatType, R>
618 for KafkaSinkFormatType<ReferencedConnection>
619{
620 fn into_inline_connection(self, r: R) -> KafkaSinkFormatType {
621 match self {
622 KafkaSinkFormatType::Avro {
623 schema,
624 compatibility_level,
625 schema_name,
626 wire_format,
627 } => KafkaSinkFormatType::Avro {
628 schema,
629 compatibility_level,
630 schema_name,
631 wire_format: wire_format.into_inline_connection(r),
632 },
633 KafkaSinkFormatType::Json => KafkaSinkFormatType::Json,
634 KafkaSinkFormatType::Text => KafkaSinkFormatType::Text,
635 KafkaSinkFormatType::Bytes => KafkaSinkFormatType::Bytes,
636 }
637 }
638}
639
640#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
641pub enum S3SinkFormat {
642 PgCopy(CopyFormatParams<'static>),
644 Parquet,
646}
647
648#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
650pub struct S3UploadInfo {
651 pub uri: String,
653 pub max_file_size: u64,
655 pub desc: RelationDesc,
657 pub format: S3SinkFormat,
659}
660
661pub const MIN_S3_SINK_FILE_SIZE: ByteSize = ByteSize::mb(16);
662pub const MAX_S3_SINK_FILE_SIZE: ByteSize = ByteSize::gb(4);
663
664pub const ICEBERG_APPEND_DIFF_COLUMN: &str = "_mz_diff";
666pub const ICEBERG_APPEND_TIMESTAMP_COLUMN: &str = "_mz_timestamp";
668
669pub const ICEBERG_UINT64_DECIMAL_PRECISION: u8 = 20;
672
673pub fn iceberg_type_overrides(
688 scalar_type: &mz_repr::SqlScalarType,
689) -> Option<(arrow::datatypes::DataType, String)> {
690 use arrow::datatypes::DataType;
691 use mz_repr::SqlScalarType;
692 match scalar_type {
693 SqlScalarType::UInt16 => Some((DataType::Int32, "uint2".to_string())),
694 SqlScalarType::UInt32 => Some((DataType::Int64, "uint4".to_string())),
695 SqlScalarType::UInt64 => Some((
696 DataType::Decimal128(ICEBERG_UINT64_DECIMAL_PRECISION, 0),
697 "uint8".to_string(),
698 )),
699 SqlScalarType::MzTimestamp => Some((
700 DataType::Decimal128(ICEBERG_UINT64_DECIMAL_PRECISION, 0),
701 "mz_timestamp".to_string(),
702 )),
703 SqlScalarType::Interval => Some((DataType::LargeUtf8, "interval".to_string())),
704 _ => None,
705 }
706}
707
708#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
709#[cfg_attr(any(test, feature = "proptest"), derive(Arbitrary))]
710pub struct IcebergSinkConnection<C: ConnectionAccess = InlinedConnection> {
711 pub catalog_connection_id: CatalogItemId,
712 pub catalog_connection: C::IcebergCatalog,
713
714 pub storage_connection_id: Option<CatalogItemId>,
723 pub storage_connection: Option<C::Aws>,
724
725 pub relation_key_indices: Option<Vec<usize>>,
727 pub key_desc_and_indices: Option<(RelationDesc, Vec<usize>)>,
729 pub namespace: String,
730 pub table: String,
731}
732
733impl<C: ConnectionAccess> IcebergSinkConnection<C> {
734 pub fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
739 if self == other {
740 return Ok(());
741 }
742 let IcebergSinkConnection {
743 catalog_connection_id: connection_id,
744 catalog_connection,
745 storage_connection_id,
746 storage_connection,
747 relation_key_indices,
748 key_desc_and_indices,
749 namespace,
750 table,
751 } = self;
752
753 let compatibility_checks = [
754 (
755 connection_id == &other.catalog_connection_id,
756 "connection_id",
757 ),
758 (
759 catalog_connection
760 .alter_compatible(id, &other.catalog_connection)
761 .is_ok(),
762 "catalog_connection",
763 ),
764 (
767 other.storage_connection_id.is_none()
768 || storage_connection_id == &other.storage_connection_id,
769 "storage_connection_id",
770 ),
771 (
772 match &other.storage_connection {
773 None => true, Some(after) => {
775 match storage_connection {
776 None => false, Some(before) => before.alter_compatible(id, after).is_ok(),
778 }
779 }
780 },
781 "storage_connection",
782 ),
783 (
784 relation_key_indices == &other.relation_key_indices,
785 "relation_key_indices",
786 ),
787 (
788 key_desc_and_indices == &other.key_desc_and_indices,
789 "key_desc_and_indices",
790 ),
791 (namespace == &other.namespace, "namespace"),
792 (table == &other.table, "table"),
793 ];
794 for (compatible, field) in compatibility_checks {
795 if !compatible {
796 tracing::warn!(
797 "IcebergSinkConnection incompatible at {field}:\nself:\n{:#?}\n\nother\n{:#?}",
798 self,
799 other
800 );
801
802 return Err(AlterError { id });
803 }
804 }
805
806 Ok(())
807 }
808}
809
810impl<R: ConnectionResolver> IntoInlineConnection<IcebergSinkConnection, R>
811 for IcebergSinkConnection<ReferencedConnection>
812{
813 fn into_inline_connection(self, r: R) -> IcebergSinkConnection {
814 let IcebergSinkConnection {
815 catalog_connection_id,
816 catalog_connection,
817 storage_connection_id,
818 storage_connection,
819 relation_key_indices,
820 key_desc_and_indices,
821 namespace,
822 table,
823 } = self;
824 IcebergSinkConnection {
825 catalog_connection_id,
826 catalog_connection: r
827 .resolve_connection(catalog_connection)
828 .unwrap_iceberg_catalog(),
829 storage_connection_id,
830 storage_connection: storage_connection.map(|c| r.resolve_connection(c).unwrap_aws()),
831 relation_key_indices,
832 key_desc_and_indices,
833 namespace,
834 table,
835 }
836 }
837}