Skip to main content

mz_storage_types/
connections.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//! Connection types.
11
12use std::borrow::Cow;
13use std::collections::{BTreeMap, BTreeSet};
14use std::net::SocketAddr;
15use std::sync::Arc;
16use std::time::SystemTime;
17
18use anyhow::{Context, anyhow};
19use async_trait::async_trait;
20use aws_credential_types::provider::{ProvideCredentials, SharedCredentialsProvider};
21use aws_sigv4::http_request::{SignableBody, SignableRequest, SigningSettings, sign};
22use aws_sigv4::sign::v4;
23// Aliased to avoid colliding with `mz_ccsr::tls::Identity`.
24use aws_smithy_runtime_api::client::identity::Identity as AwsIdentity;
25use base64::Engine;
26use http::{HeaderName, HeaderValue};
27use iceberg::Catalog;
28use iceberg::CatalogBuilder;
29use iceberg::io::{
30    GCS_CREDENTIALS_JSON, GCS_DISABLE_CONFIG_LOAD, GCS_DISABLE_VM_METADATA, GCS_USER_PROJECT,
31    S3_ACCESS_KEY_ID, S3_DISABLE_EC2_METADATA, S3_REGION, S3_SECRET_ACCESS_KEY,
32};
33use iceberg_catalog_rest::{
34    REST_CATALOG_PROP_URI, REST_CATALOG_PROP_WAREHOUSE, RequestAuthenticator, RestCatalogBuilder,
35};
36use iceberg_storage_opendal::{
37    AwsCredential, AwsCredentialLoad, CustomAwsCredentialLoader, OpenDalStorageFactory,
38};
39use itertools::Itertools;
40use mz_ccsr::tls::{Certificate, Identity};
41use mz_cloud_resources::{AwsExternalIdPrefix, CloudResourceReader, vpc_endpoint_host};
42use mz_dyncfg::ConfigSet;
43use mz_kafka_util::client::{
44    BrokerAddr, BrokerRewrite, HostMappingRules, MzClientContext, MzKafkaError, TunnelConfig,
45    TunnelingClientContext,
46};
47use mz_mysql_util::{MySqlConn, MySqlError};
48use mz_ore::assert_none;
49use mz_ore::error::ErrorExt;
50use mz_ore::future::{InTask, OreFutureExt};
51use mz_ore::netio::resolve_address;
52use mz_ore::num::NonNeg;
53use mz_repr::{CatalogItemId, GlobalId};
54use mz_secrets::SecretsReader;
55use mz_sql_parser::ast::ConnectionRulePattern;
56use mz_ssh_util::keys::SshKeyPair;
57use mz_ssh_util::tunnel::SshTunnelConfig;
58use mz_ssh_util::tunnel_manager::{ManagedSshTunnelHandle, SshTunnelManager};
59use mz_tracing::CloneableEnvFilter;
60use rdkafka::ClientContext;
61use rdkafka::config::FromClientConfigAndContext;
62use rdkafka::consumer::{BaseConsumer, Consumer};
63use regex::Regex;
64use reqwest::Request;
65use serde::{Deserialize, Deserializer, Serialize};
66use tokio::net;
67use tokio::runtime::Handle;
68use tokio_postgres::config::SslMode;
69use tracing::{debug, info, warn};
70use url::Url;
71
72use crate::AlterCompatible;
73use crate::configuration::StorageConfiguration;
74use crate::connections::aws::{
75    AwsAuth, AwsConnection, AwsConnectionReference, AwsConnectionValidationError,
76};
77use crate::connections::gcp::{GcpConnectionReference, GcpTokenProvider};
78use crate::connections::string_or_secret::StringOrSecret;
79use crate::controller::AlterError;
80use crate::dyncfgs::{
81    ENFORCE_EXTERNAL_ADDRESSES, KAFKA_CLIENT_ID_ENRICHMENT_RULES,
82    KAFKA_DEFAULT_AWS_PRIVATELINK_ENDPOINT_IDENTIFICATION_ALGORITHM, KAFKA_RECONNECT_BACKOFF,
83    KAFKA_RECONNECT_BACKOFF_MAX, KAFKA_RETRY_BACKOFF, KAFKA_RETRY_BACKOFF_MAX,
84};
85use crate::errors::{ContextCreationError, CsrConnectError};
86
87pub mod aws;
88pub mod gcp;
89pub mod inline;
90pub mod string_or_secret;
91
92const REST_CATALOG_PROP_SCOPE: &str = "scope";
93const REST_CATALOG_PROP_CREDENTIAL: &str = "credential";
94
95/// A credential loader that wraps an aws-sdk-rust credentials provider for use with
96/// iceberg/OpenDAL. This allows us to provide refreshable credentials from the AWS SDK
97/// credential chain (including the full assume role chain) to OpenDAL's S3 implementation.
98///
99/// We use this instead of OpenDAL's built-in assume role support because Materialize
100/// has a runtime-defined credential chain (ambient → jump role → user role with external ID)
101/// that can't be expressed via OpenDAL's static configuration properties.
102struct AwsSdkCredentialLoader {
103    /// The underlying AWS SDK credentials provider. For assume role auth, this provider
104    /// already handles the full chain: ambient creds -> jump role -> user role.
105    provider: SharedCredentialsProvider,
106}
107
108impl AwsSdkCredentialLoader {
109    fn new(provider: SharedCredentialsProvider) -> Self {
110        Self { provider }
111    }
112}
113
114#[async_trait]
115impl AwsCredentialLoad for AwsSdkCredentialLoader {
116    async fn load_credential(
117        &self,
118        _client: reqwest::Client,
119    ) -> anyhow::Result<Option<AwsCredential>> {
120        let creds = self
121            .provider
122            .provide_credentials()
123            .await
124            .map_err(|e| {
125                warn!(
126                    error = %e.display_with_causes(),
127                    "failed to load AWS credentials for Iceberg FileIO from SDK provider"
128                );
129                e
130            })
131            .context(
132                "failed to load AWS credentials from SDK provider for Iceberg FileIO \
133                 (credential source may be temporarily unavailable)",
134            )?;
135
136        Ok(Some(AwsCredential {
137            access_key_id: creds.access_key_id().to_string(),
138            secret_access_key: creds.secret_access_key().to_string(),
139            session_token: creds.session_token().map(|s| s.to_string()),
140            expires_in: creds.expiry().map(|t| t.into()),
141        }))
142    }
143}
144
145/// Signs each outgoing REST-catalog request with AWS SigV4.
146///
147/// Holds a [`SharedCredentialsProvider`] (not static `Credentials`) so each
148/// request signs with refreshable creds from Materialize's chain
149/// (ambient -> jump role -> user role w/ external ID).
150struct Sigv4Authenticator {
151    provider: SharedCredentialsProvider,
152    region: String,
153    /// The AWS signing name. `"s3tables"` for AWS S3 Tables REST catalog.
154    signing_name: String,
155}
156
157impl std::fmt::Debug for Sigv4Authenticator {
158    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
159        f.debug_struct("Sigv4Authenticator")
160            .field("region", &self.region)
161            .field("signing_name", &self.signing_name)
162            .finish_non_exhaustive()
163    }
164}
165
166fn sigv4_err(e: impl Into<anyhow::Error>) -> iceberg::Error {
167    iceberg::Error::new(iceberg::ErrorKind::DataInvalid, "AWS SigV4").with_source(e)
168}
169
170#[async_trait]
171impl RequestAuthenticator for Sigv4Authenticator {
172    async fn authenticate_request(&self, req: &mut Request) -> iceberg::Result<()> {
173        let creds = self
174            .provider
175            .provide_credentials()
176            .await
177            .map_err(sigv4_err)?;
178        let identity: AwsIdentity = creds.into();
179        let params = v4::SigningParams::builder()
180            .identity(&identity)
181            .region(&self.region)
182            .name(&self.signing_name)
183            .time(SystemTime::now())
184            .settings(SigningSettings::default())
185            .build()
186            .map_err(sigv4_err)?
187            .into();
188        let body: &[u8] = req
189            .body()
190            .map(|b| match b.as_bytes() {
191                Some(b) => Ok(b),
192                None => Err(iceberg::Error::new(
193                    iceberg::ErrorKind::FeatureUnsupported,
194                    "SigV4 Authenticator cannot sign a streaming request body.",
195                )),
196            })
197            .transpose()?
198            .unwrap_or_default();
199        let headers = req
200            .headers()
201            .iter()
202            .map(|(k, v)| {
203                Ok((
204                    k.as_str(),
205                    v.to_str().map_err(|_| {
206                        iceberg::Error::new(
207                            iceberg::ErrorKind::DataInvalid,
208                            format!("header '{}' value is not all visible ASCII", k),
209                        )
210                    })?,
211                ))
212            })
213            .collect::<iceberg::Result<Vec<(&str, &str)>>>()?;
214        let signable = SignableRequest::new(
215            req.method().as_str(),
216            req.url().as_str(),
217            headers.into_iter(),
218            SignableBody::Bytes(body),
219        )
220        .map_err(sigv4_err)?;
221        let (instructions, _sig) = sign(signable, &params).map_err(sigv4_err)?.into_parts();
222        let (new_headers, new_query) = instructions.into_parts();
223        for header in new_headers {
224            let mut value = HeaderValue::from_str(header.value()).map_err(sigv4_err)?;
225            value.set_sensitive(header.sensitive());
226            req.headers_mut()
227                .insert(HeaderName::from_static(header.name()), value);
228        }
229        if !new_query.is_empty() {
230            let url = req.url_mut();
231            let mut pairs = url.query_pairs_mut();
232            for (name, value) in new_query {
233                pairs.append_pair(name, &value);
234            }
235        }
236        Ok(())
237    }
238
239    // SigV4 is stateless: nothing to cache, invalidate, or refresh.
240    async fn invalidate_cache(&self) -> iceberg::Result<()> {
241        Ok(())
242    }
243    async fn regenerate_cache(&self) -> iceberg::Result<()> {
244        Ok(())
245    }
246}
247
248/// An extension trait for [`SecretsReader`]
249#[async_trait::async_trait]
250trait SecretsReaderExt {
251    /// `SecretsReader::read`, but optionally run in a task.
252    async fn read_in_task_if(
253        &self,
254        in_task: InTask,
255        id: CatalogItemId,
256    ) -> Result<Vec<u8>, anyhow::Error>;
257
258    /// `SecretsReader::read_string`, but optionally run in a task.
259    async fn read_string_in_task_if(
260        &self,
261        in_task: InTask,
262        id: CatalogItemId,
263    ) -> Result<String, anyhow::Error>;
264}
265
266#[async_trait::async_trait]
267impl SecretsReaderExt for Arc<dyn SecretsReader> {
268    async fn read_in_task_if(
269        &self,
270        in_task: InTask,
271        id: CatalogItemId,
272    ) -> Result<Vec<u8>, anyhow::Error> {
273        let sr = Arc::clone(self);
274        async move { sr.read(id).await }
275            .run_in_task_if(in_task, || "secrets_reader_read".to_string())
276            .await
277    }
278    async fn read_string_in_task_if(
279        &self,
280        in_task: InTask,
281        id: CatalogItemId,
282    ) -> Result<String, anyhow::Error> {
283        let sr = Arc::clone(self);
284        async move { sr.read_string(id).await }
285            .run_in_task_if(in_task, || "secrets_reader_read".to_string())
286            .await
287    }
288}
289
290/// Extra context to pass through when instantiating a connection for a source
291/// or sink.
292///
293/// Should be kept cheaply cloneable.
294#[derive(Debug, Clone)]
295pub struct ConnectionContext {
296    /// An opaque identifier for the environment in which this process is
297    /// running.
298    ///
299    /// The storage layer is intentionally unaware of the structure within this
300    /// identifier. Higher layers of the stack can make use of that structure,
301    /// but the storage layer should be oblivious to it.
302    pub environment_id: String,
303    /// The level for librdkafka's logs.
304    pub librdkafka_log_level: tracing::Level,
305    /// A prefix for an external ID to use for all AWS AssumeRole operations.
306    pub aws_external_id_prefix: Option<AwsExternalIdPrefix>,
307    /// The ARN for a Materialize-controlled role to assume before assuming
308    /// a customer's requested role for an AWS connection.
309    pub aws_connection_role_arn: Option<String>,
310    /// A secrets reader.
311    pub secrets_reader: Arc<dyn SecretsReader>,
312    /// A cloud resource reader, if supported in this configuration.
313    pub cloud_resource_reader: Option<Arc<dyn CloudResourceReader>>,
314    /// A manager for SSH tunnels.
315    pub ssh_tunnel_manager: SshTunnelManager,
316}
317
318impl ConnectionContext {
319    /// Constructs a new connection context from command line arguments.
320    ///
321    /// **WARNING:** it is critical for security that the `aws_external_id` be
322    /// provided by the operator of the Materialize service (i.e., via a CLI
323    /// argument or environment variable) and not the end user of Materialize
324    /// (e.g., via a configuration option in a SQL statement). See
325    /// [`AwsExternalIdPrefix`] for details.
326    pub fn from_cli_args(
327        environment_id: String,
328        startup_log_level: &CloneableEnvFilter,
329        aws_external_id_prefix: Option<AwsExternalIdPrefix>,
330        aws_connection_role_arn: Option<String>,
331        secrets_reader: Arc<dyn SecretsReader>,
332        cloud_resource_reader: Option<Arc<dyn CloudResourceReader>>,
333    ) -> ConnectionContext {
334        ConnectionContext {
335            environment_id,
336            librdkafka_log_level: mz_ore::tracing::crate_level(
337                &startup_log_level.clone().into(),
338                "librdkafka",
339            ),
340            aws_external_id_prefix,
341            aws_connection_role_arn,
342            secrets_reader,
343            cloud_resource_reader,
344            ssh_tunnel_manager: SshTunnelManager::default(),
345        }
346    }
347
348    /// Constructs a new connection context for usage in tests.
349    pub fn for_tests(secrets_reader: Arc<dyn SecretsReader>) -> ConnectionContext {
350        ConnectionContext {
351            environment_id: "test-environment-id".into(),
352            librdkafka_log_level: tracing::Level::INFO,
353            aws_external_id_prefix: Some(
354                AwsExternalIdPrefix::new_from_cli_argument_or_environment_variable(
355                    "test-aws-external-id-prefix",
356                )
357                .expect("infallible"),
358            ),
359            aws_connection_role_arn: Some(
360                "arn:aws:iam::123456789000:role/MaterializeConnection".into(),
361            ),
362            secrets_reader,
363            cloud_resource_reader: None,
364            ssh_tunnel_manager: SshTunnelManager::default(),
365        }
366    }
367}
368
369#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
370pub enum Connection<C: ConnectionAccess = InlinedConnection> {
371    Kafka(KafkaConnection<C>),
372    Csr(CsrConnection<C>),
373    GlueSchemaRegistry(GlueSchemaRegistryConnection<C>),
374    Postgres(PostgresConnection<C>),
375    Ssh(SshConnection),
376    Aws(AwsConnection),
377    AwsPrivatelink(AwsPrivatelinkConnection),
378    Gcp(gcp::GcpConnection),
379    MySql(MySqlConnection<C>),
380    SqlServer(SqlServerConnectionDetails<C>),
381    IcebergCatalog(IcebergCatalogConnection<C>),
382}
383
384impl<R: ConnectionResolver> IntoInlineConnection<Connection, R>
385    for Connection<ReferencedConnection>
386{
387    fn into_inline_connection(self, r: R) -> Connection {
388        match self {
389            Connection::Kafka(kafka) => Connection::Kafka(kafka.into_inline_connection(r)),
390            Connection::Csr(csr) => Connection::Csr(csr.into_inline_connection(r)),
391            Connection::GlueSchemaRegistry(glue) => {
392                Connection::GlueSchemaRegistry(glue.into_inline_connection(r))
393            }
394            Connection::Postgres(pg) => Connection::Postgres(pg.into_inline_connection(r)),
395            Connection::Ssh(ssh) => Connection::Ssh(ssh),
396            Connection::Aws(aws) => Connection::Aws(aws),
397            Connection::AwsPrivatelink(awspl) => Connection::AwsPrivatelink(awspl),
398            Connection::Gcp(gcp) => Connection::Gcp(gcp),
399            Connection::MySql(mysql) => Connection::MySql(mysql.into_inline_connection(r)),
400            Connection::SqlServer(sql_server) => {
401                Connection::SqlServer(sql_server.into_inline_connection(r))
402            }
403            Connection::IcebergCatalog(iceberg) => {
404                Connection::IcebergCatalog(iceberg.into_inline_connection(r))
405            }
406        }
407    }
408}
409
410impl<C: ConnectionAccess> Connection<C> {
411    /// Whether this connection should be validated by default on creation.
412    pub fn validate_by_default(&self) -> bool {
413        match self {
414            Connection::Kafka(conn) => conn.validate_by_default(),
415            Connection::Csr(conn) => conn.validate_by_default(),
416            Connection::GlueSchemaRegistry(conn) => conn.validate_by_default(),
417            Connection::Postgres(conn) => conn.validate_by_default(),
418            Connection::Ssh(conn) => conn.validate_by_default(),
419            Connection::Aws(conn) => conn.validate_by_default(),
420            Connection::AwsPrivatelink(conn) => conn.validate_by_default(),
421            Connection::Gcp(conn) => conn.validate_by_default(),
422            Connection::MySql(conn) => conn.validate_by_default(),
423            Connection::SqlServer(conn) => conn.validate_by_default(),
424            Connection::IcebergCatalog(conn) => conn.validate_by_default(),
425        }
426    }
427}
428
429impl Connection<InlinedConnection> {
430    /// Validates this connection by attempting to connect to the upstream system.
431    pub async fn validate(
432        &self,
433        id: CatalogItemId,
434        storage_configuration: &StorageConfiguration,
435    ) -> Result<(), ConnectionValidationError> {
436        match self {
437            Connection::Kafka(conn) => conn.validate(id, storage_configuration).await?,
438            Connection::Csr(conn) => conn.validate(id, storage_configuration).await?,
439            Connection::GlueSchemaRegistry(conn) => {
440                conn.validate(id, storage_configuration).await?
441            }
442            Connection::Postgres(conn) => {
443                conn.validate(id, storage_configuration).await?;
444            }
445            Connection::Ssh(conn) => conn.validate(id, storage_configuration).await?,
446            Connection::Aws(conn) => conn.validate(id, storage_configuration).await?,
447            Connection::AwsPrivatelink(conn) => conn.validate(id, storage_configuration).await?,
448            Connection::Gcp(conn) => conn.validate(id, storage_configuration).await?,
449            Connection::MySql(conn) => {
450                conn.validate(id, storage_configuration).await?;
451            }
452            Connection::SqlServer(conn) => {
453                conn.validate(id, storage_configuration).await?;
454            }
455            Connection::IcebergCatalog(conn) => conn.validate(id, storage_configuration).await?,
456        }
457        Ok(())
458    }
459
460    pub fn unwrap_kafka(self) -> <InlinedConnection as ConnectionAccess>::Kafka {
461        match self {
462            Self::Kafka(conn) => conn,
463            o => unreachable!("{o:?} is not a Kafka connection"),
464        }
465    }
466
467    pub fn unwrap_pg(self) -> <InlinedConnection as ConnectionAccess>::Pg {
468        match self {
469            Self::Postgres(conn) => conn,
470            o => unreachable!("{o:?} is not a Postgres connection"),
471        }
472    }
473
474    pub fn unwrap_mysql(self) -> <InlinedConnection as ConnectionAccess>::MySql {
475        match self {
476            Self::MySql(conn) => conn,
477            o => unreachable!("{o:?} is not a MySQL connection"),
478        }
479    }
480
481    pub fn unwrap_sql_server(self) -> <InlinedConnection as ConnectionAccess>::SqlServer {
482        match self {
483            Self::SqlServer(conn) => conn,
484            o => unreachable!("{o:?} is not a SQL Server connection"),
485        }
486    }
487
488    pub fn unwrap_aws(self) -> <InlinedConnection as ConnectionAccess>::Aws {
489        match self {
490            Self::Aws(conn) => conn,
491            o => unreachable!("{o:?} is not an AWS connection"),
492        }
493    }
494
495    pub fn unwrap_gcp(self) -> <InlinedConnection as ConnectionAccess>::Gcp {
496        match self {
497            Self::Gcp(conn) => conn,
498            o => unreachable!("{o:?} is not a GCP connection"),
499        }
500    }
501
502    pub fn unwrap_ssh(self) -> <InlinedConnection as ConnectionAccess>::Ssh {
503        match self {
504            Self::Ssh(conn) => conn,
505            o => unreachable!("{o:?} is not an SSH connection"),
506        }
507    }
508
509    pub fn unwrap_csr(self) -> <InlinedConnection as ConnectionAccess>::Csr {
510        match self {
511            Self::Csr(conn) => conn,
512            o => unreachable!("{o:?} is not a Kafka connection"),
513        }
514    }
515
516    pub fn unwrap_glue_schema_registry(
517        self,
518    ) -> <InlinedConnection as ConnectionAccess>::GlueSchemaRegistry {
519        match self {
520            Self::GlueSchemaRegistry(conn) => conn,
521            o => unreachable!("{o:?} is not an AWS Glue Schema Registry connection"),
522        }
523    }
524
525    pub fn unwrap_iceberg_catalog(self) -> <InlinedConnection as ConnectionAccess>::IcebergCatalog {
526        match self {
527            Self::IcebergCatalog(conn) => conn,
528            o => unreachable!("{o:?} is not an Iceberg catalog connection"),
529        }
530    }
531}
532
533/// An error returned by [`Connection::validate`].
534#[derive(thiserror::Error, Debug)]
535pub enum ConnectionValidationError {
536    #[error(transparent)]
537    Postgres(#[from] PostgresConnectionValidationError),
538    #[error(transparent)]
539    MySql(#[from] MySqlConnectionValidationError),
540    #[error(transparent)]
541    SqlServer(#[from] SqlServerConnectionValidationError),
542    #[error(transparent)]
543    Aws(#[from] AwsConnectionValidationError),
544    #[error(transparent)]
545    Gcp(#[from] gcp::GcpConnectionValidationError),
546    #[error("{}", .0.display_with_causes())]
547    Other(#[from] anyhow::Error),
548}
549
550impl ConnectionValidationError {
551    /// Reports additional details about the error, if any are available.
552    pub fn detail(&self) -> Option<String> {
553        match self {
554            ConnectionValidationError::Postgres(e) => e.detail(),
555            ConnectionValidationError::MySql(e) => e.detail(),
556            ConnectionValidationError::SqlServer(e) => e.detail(),
557            ConnectionValidationError::Aws(e) => e.detail(),
558            ConnectionValidationError::Gcp(e) => e.detail(),
559            ConnectionValidationError::Other(_) => None,
560        }
561    }
562
563    /// Reports a hint for the user about how the error could be fixed.
564    pub fn hint(&self) -> Option<String> {
565        match self {
566            ConnectionValidationError::Postgres(e) => e.hint(),
567            ConnectionValidationError::MySql(e) => e.hint(),
568            ConnectionValidationError::SqlServer(e) => e.hint(),
569            ConnectionValidationError::Aws(e) => e.hint(),
570            ConnectionValidationError::Gcp(e) => e.hint(),
571            ConnectionValidationError::Other(_) => None,
572        }
573    }
574}
575
576impl<C: ConnectionAccess> AlterCompatible for Connection<C> {
577    fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
578        match (self, other) {
579            (Self::Aws(s), Self::Aws(o)) => s.alter_compatible(id, o),
580            (Self::AwsPrivatelink(s), Self::AwsPrivatelink(o)) => s.alter_compatible(id, o),
581            (Self::Gcp(s), Self::Gcp(o)) => s.alter_compatible(id, o),
582            (Self::Ssh(s), Self::Ssh(o)) => s.alter_compatible(id, o),
583            (Self::Csr(s), Self::Csr(o)) => s.alter_compatible(id, o),
584            (Self::Kafka(s), Self::Kafka(o)) => s.alter_compatible(id, o),
585            (Self::Postgres(s), Self::Postgres(o)) => s.alter_compatible(id, o),
586            (Self::MySql(s), Self::MySql(o)) => s.alter_compatible(id, o),
587            _ => {
588                tracing::warn!(
589                    "Connection incompatible:\nself:\n{:#?}\n\nother\n{:#?}",
590                    self,
591                    other
592                );
593                Err(AlterError { id })
594            }
595        }
596    }
597}
598
599/// Auth mechanism for Iceberg REST catalogs.
600#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
601pub enum IcebergCatalogAuth<C: ConnectionAccess = InlinedConnection> {
602    /// Use Iceberg catalog REST API's standard OAuth flow.
603    OAuth {
604        /// client_id:client_secret
605        credential: StringOrSecret,
606        /// OAuth2 scope
607        scope: Option<String>,
608    },
609    Gcp(GcpConnectionReference<C>),
610}
611
612#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
613pub struct RestIcebergCatalog<C: ConnectionAccess = InlinedConnection> {
614    pub auth: IcebergCatalogAuth<C>,
615    /// The warehouse for REST catalogs
616    pub warehouse: Option<String>,
617}
618
619#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
620pub struct S3TablesRestIcebergCatalog<C: ConnectionAccess = InlinedConnection> {
621    /// The AWS connection details, for s3tables
622    pub aws_connection: AwsConnectionReference<C>,
623    /// The warehouse for s3tables
624    pub warehouse: String,
625}
626
627impl<R: ConnectionResolver> IntoInlineConnection<IcebergCatalogAuth, R>
628    for IcebergCatalogAuth<ReferencedConnection>
629{
630    fn into_inline_connection(self, r: R) -> IcebergCatalogAuth {
631        match self {
632            IcebergCatalogAuth::Gcp(x) => IcebergCatalogAuth::Gcp(x.into_inline_connection(&r)),
633            IcebergCatalogAuth::OAuth { credential, scope } => {
634                IcebergCatalogAuth::OAuth { credential, scope }
635            }
636        }
637    }
638}
639
640impl<R: ConnectionResolver> IntoInlineConnection<RestIcebergCatalog, R>
641    for RestIcebergCatalog<ReferencedConnection>
642{
643    fn into_inline_connection(self, r: R) -> RestIcebergCatalog {
644        RestIcebergCatalog {
645            auth: self.auth.into_inline_connection(&r),
646            warehouse: self.warehouse,
647        }
648    }
649}
650
651impl<R: ConnectionResolver> IntoInlineConnection<S3TablesRestIcebergCatalog, R>
652    for S3TablesRestIcebergCatalog<ReferencedConnection>
653{
654    fn into_inline_connection(self, r: R) -> S3TablesRestIcebergCatalog {
655        S3TablesRestIcebergCatalog {
656            aws_connection: self.aws_connection.into_inline_connection(&r),
657            warehouse: self.warehouse,
658        }
659    }
660}
661
662#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
663pub enum IcebergCatalogType {
664    Rest,
665    S3TablesRest,
666}
667
668#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
669pub enum IcebergCatalogImpl<C: ConnectionAccess = InlinedConnection> {
670    Rest(RestIcebergCatalog<C>),
671    S3TablesRest(S3TablesRestIcebergCatalog<C>),
672}
673
674impl<R: ConnectionResolver> IntoInlineConnection<IcebergCatalogImpl, R>
675    for IcebergCatalogImpl<ReferencedConnection>
676{
677    fn into_inline_connection(self, r: R) -> IcebergCatalogImpl {
678        match self {
679            IcebergCatalogImpl::Rest(rest) => {
680                IcebergCatalogImpl::Rest(rest.into_inline_connection(r))
681            }
682            IcebergCatalogImpl::S3TablesRest(s3tables) => {
683                IcebergCatalogImpl::S3TablesRest(s3tables.into_inline_connection(r))
684            }
685        }
686    }
687}
688
689#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
690pub struct IcebergCatalogConnection<C: ConnectionAccess = InlinedConnection> {
691    /// The catalog impl impl of that catalog
692    pub catalog: IcebergCatalogImpl<C>,
693    /// Where the catalog is located
694    pub uri: reqwest::Url,
695}
696
697impl AlterCompatible for IcebergCatalogConnection {
698    fn alter_compatible(&self, id: GlobalId, _other: &Self) -> Result<(), AlterError> {
699        Err(AlterError { id })
700    }
701}
702
703impl<R: ConnectionResolver> IntoInlineConnection<IcebergCatalogConnection, R>
704    for IcebergCatalogConnection<ReferencedConnection>
705{
706    fn into_inline_connection(self, r: R) -> IcebergCatalogConnection {
707        IcebergCatalogConnection {
708            catalog: self.catalog.into_inline_connection(&r),
709            uri: self.uri,
710        }
711    }
712}
713
714impl<C: ConnectionAccess> IcebergCatalogConnection<C> {
715    fn validate_by_default(&self) -> bool {
716        true
717    }
718}
719
720impl IcebergCatalogConnection<InlinedConnection> {
721    pub async fn connect(
722        &self,
723        storage_configuration: &StorageConfiguration,
724        in_task: InTask,
725    ) -> Result<Arc<dyn Catalog>, anyhow::Error> {
726        match self.catalog {
727            IcebergCatalogImpl::S3TablesRest(ref s3tables) => {
728                self.connect_s3tables(s3tables, storage_configuration, in_task)
729                    .await
730            }
731            IcebergCatalogImpl::Rest(ref rest) => {
732                self.connect_rest(rest, storage_configuration, in_task)
733                    .await
734            }
735        }
736    }
737
738    pub fn catalog_type(&self) -> IcebergCatalogType {
739        match self.catalog {
740            IcebergCatalogImpl::S3TablesRest(_) => IcebergCatalogType::S3TablesRest,
741            IcebergCatalogImpl::Rest(_) => IcebergCatalogType::Rest,
742        }
743    }
744
745    pub fn s3tables_catalog(&self) -> Option<&S3TablesRestIcebergCatalog> {
746        match &self.catalog {
747            IcebergCatalogImpl::S3TablesRest(s3tables) => Some(s3tables),
748            IcebergCatalogImpl::Rest(_) => None,
749        }
750    }
751
752    pub fn rest_catalog(&self) -> Option<&RestIcebergCatalog> {
753        match &self.catalog {
754            IcebergCatalogImpl::Rest(rest) => Some(rest),
755            IcebergCatalogImpl::S3TablesRest(_) => None,
756        }
757    }
758
759    async fn connect_s3tables(
760        &self,
761        s3tables: &S3TablesRestIcebergCatalog,
762        storage_configuration: &StorageConfiguration,
763        in_task: InTask,
764    ) -> Result<Arc<dyn Catalog>, anyhow::Error> {
765        let secret_reader = &storage_configuration.connection_context.secrets_reader;
766        let aws_ref = &s3tables.aws_connection;
767
768        let aws_region = aws_ref
769            .connection
770            .region
771            .clone()
772            .unwrap_or_else(|| "us-east-1".to_string());
773
774        let mut props = vec![
775            (S3_REGION.to_string(), aws_region.clone()),
776            (S3_DISABLE_EC2_METADATA.to_string(), "true".to_string()),
777            (
778                REST_CATALOG_PROP_WAREHOUSE.to_string(),
779                s3tables.warehouse.clone(),
780            ),
781            (REST_CATALOG_PROP_URI.to_string(), self.uri.to_string()),
782        ];
783
784        let aws_auth = aws_ref.connection.auth.clone();
785
786        if let AwsAuth::Credentials(creds) = &aws_auth {
787            props.push((
788                S3_ACCESS_KEY_ID.to_string(),
789                creds
790                    .access_key_id
791                    .get_string(in_task, secret_reader)
792                    .await?,
793            ));
794            props.push((
795                S3_SECRET_ACCESS_KEY.to_string(),
796                secret_reader.read_string(creds.secret_access_key).await?,
797            ));
798        }
799
800        // Sign REST catalog requests with the Materialize AWS credential chain
801        // via a custom `RequestAuthenticator`. For AssumeRole auth, also feed
802        // the chain to OpenDAL's S3 loader so data-file IO uses the same creds.
803        //
804        // For AssumeRole auth, the provider serves cached credentials that a
805        // background task keeps fresh, so no request through this catalog ever
806        // waits on STS. The task lives as long as the catalog holds the
807        // provider.
808        let credentials_provider = match &aws_auth {
809            // NOTE: This branch never contacts the connection's ENDPOINT.
810            // REST requests go to the catalog URI and the STS calls use the
811            // SDK defaults. The endpoint is still validated so a forbidden
812            // one is rejected rather than silently ignored.
813            AwsAuth::AssumeRole(assume_role) => {
814                aws_ref.connection.validate_endpoint(
815                    ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
816                )?;
817                assume_role
818                    .prefetch_credentials(
819                        &storage_configuration.connection_context,
820                        aws_ref.connection_id,
821                        storage_configuration.config_set(),
822                        format!("aws-connection-{}", aws_ref.connection_id),
823                    )
824                    .await
825                    .with_context(|| {
826                        format!(
827                            "failed to initialize AssumeRole credentials for S3 Tables Iceberg \
828                             catalog (catalog uri: {}, warehouse: {})",
829                            self.uri, s3tables.warehouse
830                        )
831                    })?
832            }
833            AwsAuth::Credentials(_) => {
834                let aws_config = aws_ref
835                    .connection
836                    .load_sdk_config(
837                        &storage_configuration.connection_context,
838                        aws_ref.connection_id,
839                        in_task,
840                        ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
841                    )
842                    .await
843                    .with_context(|| {
844                        format!(
845                            "failed to load AWS SDK config for S3 Tables Iceberg catalog \
846                             (connection id: {}, auth method: {}, catalog uri: {}, warehouse: {})",
847                            aws_ref.connection_id,
848                            aws_ref.connection.auth_method(),
849                            self.uri,
850                            s3tables.warehouse
851                        )
852                    })?;
853                aws_config
854                    .credentials_provider()
855                    .ok_or_else(|| anyhow!("aws_config missing credentials provider"))?
856            }
857        };
858
859        let authenticator = Arc::new(Sigv4Authenticator {
860            provider: credentials_provider.clone(),
861            region: aws_region.clone(),
862            signing_name: "s3tables".to_string(),
863        });
864
865        // N.B. We're using the AWS credentials from the catalog connection for the storage layer
866        //   even though the sink comes with its own (unused) AWS credentials for storage.
867        let customized_credential_load = if matches!(aws_auth, AwsAuth::AssumeRole(_)) {
868            Some(CustomAwsCredentialLoader::new(Arc::new(
869                AwsSdkCredentialLoader::new(credentials_provider),
870            )))
871        } else {
872            None
873        };
874
875        let storage_factory = Arc::new(OpenDalStorageFactory::S3 {
876            configured_scheme: "s3".to_string(),
877            customized_credential_load,
878        });
879
880        let catalog = RestCatalogBuilder::default()
881            .with_storage_factory(storage_factory)
882            .with_authenticator(authenticator)
883            .load("IcebergCatalog", props.into_iter().collect())
884            .await
885            .with_context(|| {
886                format!(
887                    "failed to create S3 Tables Iceberg catalog \
888                     (connection id: {}, catalog uri: {}, warehouse: {})",
889                    aws_ref.connection_id, self.uri, s3tables.warehouse
890                )
891            })?;
892
893        Ok(Arc::new(catalog))
894    }
895
896    async fn connect_rest(
897        &self,
898        rest: &RestIcebergCatalog,
899        storage_configuration: &StorageConfiguration,
900        in_task: InTask,
901    ) -> Result<Arc<dyn Catalog>, anyhow::Error> {
902        let mut props = BTreeMap::from([(
903            REST_CATALOG_PROP_URI.to_string(),
904            self.uri.to_string().clone(),
905        )]);
906
907        if let Some(warehouse) = &rest.warehouse {
908            props.insert(REST_CATALOG_PROP_WAREHOUSE.to_string(), warehouse.clone());
909        }
910
911        // Catalog auth is configured through a combination of `props` and `.with_authenticator(...)`,
912        // which happen at different stages of the [`RestCatalogBuilder`] -> [`RestCatalog`]
913        // construction pipeline.
914        let (storage_factory, custom_authenticator) = match &rest.auth {
915            IcebergCatalogAuth::OAuth { credential, scope } => {
916                let credential = credential
917                    .get_string(
918                        in_task,
919                        &storage_configuration.connection_context.secrets_reader,
920                    )
921                    .await
922                    .map_err(|e| anyhow!("failed to read Iceberg catalog credential: {e}"))?;
923                props.insert(REST_CATALOG_PROP_CREDENTIAL.to_string(), credential);
924
925                if let Some(scope) = scope {
926                    props.insert(REST_CATALOG_PROP_SCOPE.to_string(), scope.clone());
927                }
928                (
929                    OpenDalStorageFactory::S3 {
930                        configured_scheme: "s3".to_string(),
931                        // When used with MinIO, Polaris returns a config with:
932                        //   s3.access-key-id, s3.secret-access-key, s3.endpoint, ...
933                        // `iceberg-rust` forwards these props to `opendal`.
934                        // N.B. This is not confirmed to work with other catalog & storage implementations.
935                        customized_credential_load: None,
936                    },
937                    None,
938                )
939            }
940            IcebergCatalogAuth::Gcp(gcp_connection_reference) => {
941                let (creds_json, service_account) = gcp_connection_reference
942                    .connection
943                    .read_credentials(storage_configuration)
944                    .await
945                    .map_err(|e| anyhow!("failed to parse GCP service account JSON: {e}"))?;
946
947                props.insert(
948                    GCS_CREDENTIALS_JSON.to_owned(),
949                    base64::engine::general_purpose::STANDARD.encode(creds_json),
950                );
951                // We supplied a service account key. Don't look elsewhere for GCP credentials.
952                props.insert(GCS_DISABLE_VM_METADATA.to_owned(), "true".to_owned());
953                props.insert(GCS_DISABLE_CONFIG_LOAD.to_owned(), "true".to_owned());
954                if let Some(project_id) = service_account.project_id() {
955                    props.insert(GCS_USER_PROJECT.to_owned(), project_id.to_owned());
956                    props.insert(
957                        "header.x-goog-user-project".to_owned(),
958                        project_id.to_owned(),
959                    );
960                }
961
962                (
963                    OpenDalStorageFactory::Gcs,
964                    Some(iceberg_catalog_rest::BearerTokenAuthenticator::new(
965                        Arc::new(GcpTokenProvider { service_account }),
966                    )),
967                )
968            }
969        };
970
971        let mut catalog =
972            RestCatalogBuilder::default().with_storage_factory(Arc::new(storage_factory));
973        if let Some(auth) = custom_authenticator {
974            catalog = catalog.with_authenticator(Arc::new(auth));
975        }
976        let catalog = catalog
977            .load("IcebergCatalog", props.into_iter().collect())
978            .await
979            .map_err(|e| anyhow!("failed to create Iceberg catalog: {e}"))?;
980        Ok(Arc::new(catalog))
981    }
982
983    async fn validate(
984        &self,
985        _id: CatalogItemId,
986        storage_configuration: &StorageConfiguration,
987    ) -> Result<(), ConnectionValidationError> {
988        let catalog = self
989            .connect(storage_configuration, InTask::No)
990            .await
991            .map_err(|e| {
992                ConnectionValidationError::Other(anyhow!("failed to connect to catalog: {e}"))
993            })?;
994
995        // If we can list namespaces, the connection is valid.
996        catalog.list_namespaces(None).await.map_err(|e| {
997            ConnectionValidationError::Other(anyhow!("failed to list namespaces: {e}"))
998        })?;
999
1000        Ok(())
1001    }
1002}
1003
1004#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
1005pub struct AwsPrivatelinkConnection {
1006    pub service_name: String,
1007    pub availability_zones: Vec<String>,
1008}
1009
1010impl AlterCompatible for AwsPrivatelinkConnection {
1011    fn alter_compatible(&self, _id: GlobalId, _other: &Self) -> Result<(), AlterError> {
1012        // Every element of the AwsPrivatelinkConnection connection is configurable.
1013        Ok(())
1014    }
1015}
1016
1017#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
1018pub struct KafkaTlsConfig {
1019    pub identity: Option<TlsIdentity>,
1020    pub root_cert: Option<StringOrSecret>,
1021}
1022
1023#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
1024pub struct KafkaSaslConfig<C: ConnectionAccess = InlinedConnection> {
1025    pub mechanism: String,
1026    pub username: StringOrSecret,
1027    pub password: Option<CatalogItemId>,
1028    pub aws: Option<AwsConnectionReference<C>>,
1029}
1030
1031impl<R: ConnectionResolver> IntoInlineConnection<KafkaSaslConfig, R>
1032    for KafkaSaslConfig<ReferencedConnection>
1033{
1034    fn into_inline_connection(self, r: R) -> KafkaSaslConfig {
1035        KafkaSaslConfig {
1036            mechanism: self.mechanism,
1037            username: self.username,
1038            password: self.password,
1039            aws: self.aws.map(|aws| aws.into_inline_connection(&r)),
1040        }
1041    }
1042}
1043
1044/// Specifies a Kafka broker in a [`KafkaConnection`].
1045#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
1046pub struct KafkaBroker<C: ConnectionAccess = InlinedConnection> {
1047    /// The address of the Kafka broker.
1048    pub address: String,
1049    /// An optional tunnel to use when connecting to the broker.
1050    pub tunnel: Tunnel<C>,
1051}
1052
1053impl<R: ConnectionResolver> IntoInlineConnection<KafkaBroker, R>
1054    for KafkaBroker<ReferencedConnection>
1055{
1056    fn into_inline_connection(self, r: R) -> KafkaBroker {
1057        let KafkaBroker { address, tunnel } = self;
1058        KafkaBroker {
1059            address,
1060            tunnel: tunnel.into_inline_connection(r),
1061        }
1062    }
1063}
1064
1065#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize, Default)]
1066pub struct KafkaTopicOptions {
1067    /// The replication factor for the topic.
1068    /// If `None`, the broker default will be used.
1069    pub replication_factor: Option<NonNeg<i32>>,
1070    /// The number of partitions to create.
1071    /// If `None`, the broker default will be used.
1072    pub partition_count: Option<NonNeg<i32>>,
1073    /// The initial configuration parameters for the topic.
1074    pub topic_config: BTreeMap<String, String>,
1075}
1076
1077#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
1078pub struct KafkaConnection<C: ConnectionAccess = InlinedConnection> {
1079    pub brokers: Vec<KafkaBroker<C>>,
1080    /// A tunnel through which to route traffic,
1081    /// that can be overridden for individual brokers
1082    /// in `brokers`.
1083    pub default_tunnel: Tunnel<C>,
1084    pub progress_topic: Option<String>,
1085    pub progress_topic_options: KafkaTopicOptions,
1086    pub options: BTreeMap<String, StringOrSecret>,
1087    pub tls: Option<KafkaTlsConfig>,
1088    pub sasl: Option<KafkaSaslConfig<C>>,
1089}
1090
1091impl<R: ConnectionResolver> IntoInlineConnection<KafkaConnection, R>
1092    for KafkaConnection<ReferencedConnection>
1093{
1094    fn into_inline_connection(self, r: R) -> KafkaConnection {
1095        let KafkaConnection {
1096            brokers,
1097            progress_topic,
1098            progress_topic_options,
1099            default_tunnel,
1100            options,
1101            tls,
1102            sasl,
1103        } = self;
1104
1105        let brokers = brokers
1106            .into_iter()
1107            .map(|broker| broker.into_inline_connection(&r))
1108            .collect();
1109
1110        KafkaConnection {
1111            brokers,
1112            progress_topic,
1113            progress_topic_options,
1114            default_tunnel: default_tunnel.into_inline_connection(&r),
1115            options,
1116            tls,
1117            sasl: sasl.map(|sasl| sasl.into_inline_connection(&r)),
1118        }
1119    }
1120}
1121
1122impl<C: ConnectionAccess> KafkaConnection<C> {
1123    /// Returns the name of the progress topic to use for the connection.
1124    ///
1125    /// The caller is responsible for providing the connection ID as it is not
1126    /// known to `KafkaConnection`.
1127    pub fn progress_topic(
1128        &self,
1129        connection_context: &ConnectionContext,
1130        connection_id: CatalogItemId,
1131    ) -> Cow<'_, str> {
1132        if let Some(progress_topic) = &self.progress_topic {
1133            Cow::Borrowed(progress_topic)
1134        } else {
1135            Cow::Owned(format!(
1136                "_materialize-progress-{}-{}",
1137                connection_context.environment_id, connection_id,
1138            ))
1139        }
1140    }
1141
1142    fn validate_by_default(&self) -> bool {
1143        true
1144    }
1145}
1146
1147impl KafkaConnection {
1148    /// Generates a string that can be used as the base for a configuration ID
1149    /// (e.g., `client.id`, `group.id`, `transactional.id`) for a Kafka source
1150    /// or sink.
1151    pub fn id_base(
1152        connection_context: &ConnectionContext,
1153        connection_id: CatalogItemId,
1154        object_id: GlobalId,
1155    ) -> String {
1156        format!(
1157            "materialize-{}-{}-{}",
1158            connection_context.environment_id, connection_id, object_id,
1159        )
1160    }
1161
1162    /// Enriches the provided `client_id` according to any enrichment rules in
1163    /// the `kafka_client_id_enrichment_rules` configuration parameter.
1164    pub fn enrich_client_id(&self, configs: &ConfigSet, client_id: &mut String) {
1165        #[derive(Debug, Deserialize)]
1166        struct EnrichmentRule {
1167            #[serde(deserialize_with = "deserialize_regex")]
1168            pattern: Regex,
1169            payload: String,
1170        }
1171
1172        fn deserialize_regex<'de, D>(deserializer: D) -> Result<Regex, D::Error>
1173        where
1174            D: Deserializer<'de>,
1175        {
1176            let buf = String::deserialize(deserializer)?;
1177            Regex::new(&buf).map_err(serde::de::Error::custom)
1178        }
1179
1180        let rules = KAFKA_CLIENT_ID_ENRICHMENT_RULES.get(configs);
1181        let rules = match serde_json::from_value::<Vec<EnrichmentRule>>(rules) {
1182            Ok(rules) => rules,
1183            Err(e) => {
1184                warn!(%e, "failed to decode kafka_client_id_enrichment_rules");
1185                return;
1186            }
1187        };
1188
1189        // Check every rule against every broker. Rules are matched in the order
1190        // that they are specified. It is usually a configuration error if
1191        // multiple rules match the same list of Kafka brokers, but we
1192        // nonetheless want to provide well defined semantics.
1193        debug!(?self.brokers, "evaluating client ID enrichment rules");
1194        for rule in rules {
1195            let is_match = self
1196                .brokers
1197                .iter()
1198                .any(|b| rule.pattern.is_match(&b.address));
1199            debug!(?rule, is_match, "evaluated client ID enrichment rule");
1200            if is_match {
1201                client_id.push('-');
1202                client_id.push_str(&rule.payload);
1203            }
1204        }
1205    }
1206
1207    /// Creates a Kafka client for the connection.
1208    pub async fn create_with_context<C, T>(
1209        &self,
1210        storage_configuration: &StorageConfiguration,
1211        context: C,
1212        extra_options: &BTreeMap<&str, String>,
1213        in_task: InTask,
1214    ) -> Result<T, ContextCreationError>
1215    where
1216        C: ClientContext,
1217        T: FromClientConfigAndContext<TunnelingClientContext<C>>,
1218    {
1219        let mut options = self.options.clone();
1220
1221        // Ensure that Kafka topics are *not* automatically created when
1222        // consuming, producing, or fetching metadata for a topic. This ensures
1223        // that we don't accidentally create topics with the wrong number of
1224        // partitions.
1225        options.insert("allow.auto.create.topics".into(), "false".into());
1226
1227        let brokers = match &self.default_tunnel {
1228            Tunnel::AwsPrivatelink(t) => {
1229                assert!(&self.brokers.is_empty());
1230
1231                let algo = KAFKA_DEFAULT_AWS_PRIVATELINK_ENDPOINT_IDENTIFICATION_ALGORITHM
1232                    .get(storage_configuration.config_set());
1233                options.insert("ssl.endpoint.identification.algorithm".into(), algo.into());
1234
1235                // When using a default privatelink tunnel broker/brokers cannot be specified
1236                // instead the tunnel connection_id and port are used for the initial connection.
1237                format!(
1238                    "{}:{}",
1239                    vpc_endpoint_host(
1240                        t.connection_id,
1241                        None, // Default tunnel does not support availability zones.
1242                    ),
1243                    t.port.unwrap_or(9092)
1244                )
1245            }
1246            Tunnel::AwsPrivatelinks(_pl) => {
1247                let algo = KAFKA_DEFAULT_AWS_PRIVATELINK_ENDPOINT_IDENTIFICATION_ALGORITHM
1248                    .get(storage_configuration.config_set());
1249                options.insert("ssl.endpoint.identification.algorithm".into(), algo.into());
1250
1251                if self.brokers.is_empty() {
1252                    return Err(ContextCreationError::Other(anyhow::anyhow!(
1253                        "at least one static broker is required when using BROKER or BROKERS"
1254                    )));
1255                }
1256                self.brokers.iter().map(|b| &b.address).join(",")
1257            }
1258            _ => self.brokers.iter().map(|b| &b.address).join(","),
1259        };
1260        options.insert("bootstrap.servers".into(), brokers.clone().into());
1261        let security_protocol = match (self.tls.is_some(), self.sasl.is_some()) {
1262            (false, false) => "PLAINTEXT",
1263            (true, false) => "SSL",
1264            (false, true) => "SASL_PLAINTEXT",
1265            (true, true) => "SASL_SSL",
1266        };
1267        info!(
1268            "kafka: create_with_context bootstrap.servers={brokers}, security_protocol={security_protocol}"
1269        );
1270        options.insert("security.protocol".into(), security_protocol.into());
1271        if let Some(tls) = &self.tls {
1272            if let Some(root_cert) = &tls.root_cert {
1273                options.insert("ssl.ca.pem".into(), root_cert.clone());
1274            }
1275            if let Some(identity) = &tls.identity {
1276                options.insert("ssl.key.pem".into(), StringOrSecret::Secret(identity.key));
1277                options.insert("ssl.certificate.pem".into(), identity.cert.clone());
1278            }
1279        }
1280        if let Some(sasl) = &self.sasl {
1281            options.insert("sasl.mechanisms".into(), (&sasl.mechanism).into());
1282            options.insert("sasl.username".into(), sasl.username.clone());
1283            if let Some(password) = sasl.password {
1284                options.insert("sasl.password".into(), StringOrSecret::Secret(password));
1285            }
1286        }
1287
1288        options.insert(
1289            "retry.backoff.ms".into(),
1290            KAFKA_RETRY_BACKOFF
1291                .get(storage_configuration.config_set())
1292                .as_millis()
1293                .into(),
1294        );
1295        options.insert(
1296            "retry.backoff.max.ms".into(),
1297            KAFKA_RETRY_BACKOFF_MAX
1298                .get(storage_configuration.config_set())
1299                .as_millis()
1300                .into(),
1301        );
1302        options.insert(
1303            "reconnect.backoff.ms".into(),
1304            KAFKA_RECONNECT_BACKOFF
1305                .get(storage_configuration.config_set())
1306                .as_millis()
1307                .into(),
1308        );
1309        options.insert(
1310            "reconnect.backoff.max.ms".into(),
1311            KAFKA_RECONNECT_BACKOFF_MAX
1312                .get(storage_configuration.config_set())
1313                .as_millis()
1314                .into(),
1315        );
1316
1317        let mut config = mz_kafka_util::client::create_new_client_config(
1318            storage_configuration
1319                .connection_context
1320                .librdkafka_log_level,
1321            storage_configuration.parameters.kafka_timeout_config,
1322        );
1323        for (k, v) in options {
1324            config.set(
1325                k,
1326                v.get_string(
1327                    in_task,
1328                    &storage_configuration.connection_context.secrets_reader,
1329                )
1330                .await
1331                .context("reading kafka secret")?,
1332            );
1333        }
1334        for (k, v) in extra_options {
1335            config.set(*k, v);
1336        }
1337
1338        let aws_config = match self.sasl.as_ref().and_then(|sasl| sasl.aws.as_ref()) {
1339            None => None,
1340            Some(aws) => Some(
1341                aws.connection
1342                    .load_sdk_config(
1343                        &storage_configuration.connection_context,
1344                        aws.connection_id,
1345                        in_task,
1346                        ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
1347                    )
1348                    .await?,
1349            ),
1350        };
1351
1352        // TODO(roshan): Implement enforcement of external address validation once
1353        // rdkafka client has been updated to support providing multiple resolved
1354        // addresses for brokers
1355        let mut context = TunnelingClientContext::new(
1356            context,
1357            Handle::current(),
1358            storage_configuration
1359                .connection_context
1360                .ssh_tunnel_manager
1361                .clone(),
1362            storage_configuration.parameters.ssh_timeout_config,
1363            aws_config,
1364            in_task,
1365        );
1366
1367        match &self.default_tunnel {
1368            Tunnel::Direct => {
1369                // By default, don't offer a default override for broker address lookup.
1370            }
1371            Tunnel::AwsPrivatelink(pl) => {
1372                context.set_default_tunnel(TunnelConfig::StaticHost(
1373                    // Possible bug: We have been ignoring the configured port.
1374                    KafkaConnection::from_default_aws_privatelink(pl).host,
1375                ));
1376            }
1377            Tunnel::AwsPrivatelinks(pl) => {
1378                context.set_default_tunnel(TunnelConfig::Rules(
1379                    KafkaConnection::from_aws_privatelinks(pl),
1380                ));
1381            }
1382            Tunnel::Ssh(ssh_tunnel) => {
1383                let secret = storage_configuration
1384                    .connection_context
1385                    .secrets_reader
1386                    .read_in_task_if(in_task, ssh_tunnel.connection_id)
1387                    .await?;
1388                let key_pair = SshKeyPair::from_bytes(&secret)?;
1389
1390                // Ensure any ssh-bastion address we connect to is resolved to an external address.
1391                let resolved = resolve_address(
1392                    &ssh_tunnel.connection.host,
1393                    ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
1394                )
1395                .await?;
1396                context.set_default_tunnel(TunnelConfig::Ssh(SshTunnelConfig {
1397                    host: resolved
1398                        .iter()
1399                        .map(|a| a.to_string())
1400                        .collect::<BTreeSet<_>>(),
1401                    port: ssh_tunnel.connection.port,
1402                    user: ssh_tunnel.connection.user.clone(),
1403                    key_pair,
1404                }));
1405            }
1406        }
1407        info!(
1408            "kafka: tunnel config set to {}",
1409            match &self.default_tunnel {
1410                Tunnel::Direct => "Direct".to_string(),
1411                Tunnel::AwsPrivatelink(_) => "AwsPrivatelink (static host)".to_string(),
1412                Tunnel::AwsPrivatelinks(pl) =>
1413                    format!("AwsPrivatelinks ({} rules)", pl.rules.len()),
1414                Tunnel::Ssh(_) => "Ssh".to_string(),
1415            }
1416        );
1417
1418        // Here, we preemptively rewrite broker addresses.
1419        // In concept, this overlaps with 'TunnelingClientContext::resolve_broker_addr'.
1420        for broker in &self.brokers {
1421            let mut addr_parts = broker.address.splitn(2, ':');
1422            let addr = BrokerAddr {
1423                host: addr_parts
1424                    .next()
1425                    .context("BROKER is not address:port")?
1426                    .into(),
1427                port: addr_parts
1428                    .next()
1429                    .unwrap_or("9092")
1430                    .parse()
1431                    .context("parsing BROKER port")?,
1432            };
1433            match &broker.tunnel {
1434                Tunnel::Direct => {
1435                    // By default, don't override broker address lookup.
1436                    //
1437                    // N.B.
1438                    //
1439                    // We _could_ pre-setup the default ssh tunnel for all known brokers here, but
1440                    // we avoid doing because:
1441                    // - Its not necessary.
1442                    // - Not doing so makes it easier to test the `FailedDefaultSshTunnel` path
1443                    // in the `TunnelingClientContext`.
1444                }
1445                Tunnel::AwsPrivatelink(aws_privatelink) => {
1446                    context.add_broker_rewrite(
1447                        addr,
1448                        KafkaConnection::from_aws_privatelink(aws_privatelink),
1449                    );
1450                }
1451                Tunnel::AwsPrivatelinks(_) => unreachable!(
1452                    "Individually predefined brokers do not use rule-based PrivateLinks routing."
1453                ),
1454                Tunnel::Ssh(ssh_tunnel) => {
1455                    // Ensure any SSH bastion address we connect to is resolved to an external address.
1456                    let ssh_host_resolved = resolve_address(
1457                        &ssh_tunnel.connection.host,
1458                        ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
1459                    )
1460                    .await?;
1461                    context
1462                        .add_ssh_tunnel(
1463                            addr,
1464                            SshTunnelConfig {
1465                                host: ssh_host_resolved
1466                                    .iter()
1467                                    .map(|a| a.to_string())
1468                                    .collect::<BTreeSet<_>>(),
1469                                port: ssh_tunnel.connection.port,
1470                                user: ssh_tunnel.connection.user.clone(),
1471                                key_pair: SshKeyPair::from_bytes(
1472                                    &storage_configuration
1473                                        .connection_context
1474                                        .secrets_reader
1475                                        .read_in_task_if(in_task, ssh_tunnel.connection_id)
1476                                        .await?,
1477                                )?,
1478                            },
1479                        )
1480                        .await
1481                        .map_err(ContextCreationError::Ssh)?;
1482                }
1483            }
1484        }
1485
1486        Ok(config.create_with_context(context)?)
1487    }
1488
1489    async fn validate(
1490        &self,
1491        _id: CatalogItemId,
1492        storage_configuration: &StorageConfiguration,
1493    ) -> Result<(), anyhow::Error> {
1494        let (context, error_rx) = MzClientContext::with_errors();
1495        let consumer: BaseConsumer<_> = self
1496            .create_with_context(
1497                storage_configuration,
1498                context,
1499                &BTreeMap::new(),
1500                // We are in a normal tokio context during validation, already.
1501                InTask::No,
1502            )
1503            .await?;
1504        let consumer = Arc::new(consumer);
1505
1506        let timeout = storage_configuration
1507            .parameters
1508            .kafka_timeout_config
1509            .fetch_metadata_timeout;
1510
1511        // librdkafka doesn't expose an API for determining whether a connection to
1512        // the Kafka cluster has been successfully established. So we make a
1513        // metadata request, though we don't care about the results, so that we can
1514        // report any errors making that request. If the request succeeds, we know
1515        // we were able to contact at least one broker, and that's a good proxy for
1516        // being able to contact all the brokers in the cluster.
1517        //
1518        // The downside of this approach is it produces a generic error message like
1519        // "metadata fetch error" with no additional details. The real networking
1520        // error is buried in the librdkafka logs, which are not visible to users.
1521        info!("kafka: starting connection validation via fetch_metadata (timeout={timeout:?})");
1522        let result = mz_ore::task::spawn_blocking(|| "kafka_get_metadata", {
1523            let consumer = Arc::clone(&consumer);
1524            move || consumer.fetch_metadata(None, timeout)
1525        })
1526        .await;
1527        info!(
1528            "kafka: connection validation result: {}",
1529            if result.is_ok() { "success" } else { "failed" },
1530        );
1531        match result {
1532            Ok(_) => Ok(()),
1533            // The error returned by `fetch_metadata` does not provide any details which makes for
1534            // a crappy user facing error message. For this reason we attempt to grab a better
1535            // error message from the client context, which should contain any error logs emitted
1536            // by librdkafka, and fallback to the generic error if there is nothing there.
1537            Err(err) => {
1538                // Multiple errors might have been logged during this validation but some are more
1539                // relevant than others. Specifically, we prefer non-internal errors over internal
1540                // errors since those give much more useful information to the users.
1541                let main_err = error_rx.try_iter().reduce(|cur, new| match cur {
1542                    MzKafkaError::Internal(_) => new,
1543                    _ => cur,
1544                });
1545
1546                // Don't drop the consumer until after we've drained the errors
1547                // channel. Dropping the consumer can introduce spurious errors.
1548                // See database-issues#7432.
1549                drop(consumer);
1550
1551                match main_err {
1552                    Some(err) => Err(err.into()),
1553                    None => Err(err.into()),
1554                }
1555            }
1556        }
1557    }
1558
1559    /// The "default" PrivateLink connection is used for bootstrapping Kafka.
1560    fn from_default_aws_privatelink(pl: &AwsPrivatelink) -> BrokerRewrite {
1561        BrokerRewrite {
1562            host: vpc_endpoint_host(
1563                pl.connection_id,
1564                None, // Default tunnel does not support availability zones.
1565            ),
1566            port: pl.port,
1567        }
1568    }
1569
1570    /// The "not default" PrivateLink connections are used for routing to specific Kafka brokers.
1571    fn from_aws_privatelink(pl: &AwsPrivatelink) -> BrokerRewrite {
1572        BrokerRewrite {
1573            host: vpc_endpoint_host(pl.connection_id, pl.availability_zone.as_deref()),
1574            port: pl.port,
1575        }
1576    }
1577
1578    fn from_aws_privatelink_rule(
1579        AwsPrivatelinkRule { pattern, to }: &AwsPrivatelinkRule,
1580    ) -> (mz_kafka_util::client::ConnectionRulePattern, BrokerRewrite) {
1581        (
1582            mz_kafka_util::client::ConnectionRulePattern {
1583                prefix_wildcard: pattern.prefix_wildcard,
1584                literal_match: pattern.literal_match.clone(),
1585                suffix_wildcard: pattern.suffix_wildcard,
1586            },
1587            KafkaConnection::from_aws_privatelink(to),
1588        )
1589    }
1590
1591    fn from_aws_privatelinks(pl: &AwsPrivatelinks) -> HostMappingRules {
1592        HostMappingRules {
1593            rules: pl
1594                .rules
1595                .iter()
1596                .map(KafkaConnection::from_aws_privatelink_rule)
1597                .collect_vec(),
1598        }
1599    }
1600}
1601
1602impl<C: ConnectionAccess> AlterCompatible for KafkaConnection<C> {
1603    fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
1604        let KafkaConnection {
1605            brokers: _,
1606            default_tunnel: _,
1607            progress_topic,
1608            progress_topic_options,
1609            options: _,
1610            tls: _,
1611            sasl: _,
1612        } = self;
1613
1614        let compatibility_checks = [
1615            (progress_topic == &other.progress_topic, "progress_topic"),
1616            (
1617                progress_topic_options == &other.progress_topic_options,
1618                "progress_topic_options",
1619            ),
1620        ];
1621
1622        for (compatible, field) in compatibility_checks {
1623            if !compatible {
1624                tracing::warn!(
1625                    "KafkaConnection incompatible at {field}:\nself:\n{:#?}\n\nother\n{:#?}",
1626                    self,
1627                    other
1628                );
1629
1630                return Err(AlterError { id });
1631            }
1632        }
1633
1634        Ok(())
1635    }
1636}
1637
1638/// A connection to a Confluent Schema Registry.
1639#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
1640pub struct CsrConnection<C: ConnectionAccess = InlinedConnection> {
1641    /// The URL of the schema registry.
1642    pub url: Url,
1643    /// Trusted root TLS certificate in PEM format.
1644    pub tls_root_cert: Option<StringOrSecret>,
1645    /// An optional TLS client certificate for authentication with the schema
1646    /// registry.
1647    pub tls_identity: Option<TlsIdentity>,
1648    /// Optional HTTP authentication credentials for the schema registry.
1649    pub http_auth: Option<CsrConnectionHttpAuth>,
1650    /// A tunnel through which to route traffic.
1651    pub tunnel: Tunnel<C>,
1652}
1653
1654impl<R: ConnectionResolver> IntoInlineConnection<CsrConnection, R>
1655    for CsrConnection<ReferencedConnection>
1656{
1657    fn into_inline_connection(self, r: R) -> CsrConnection {
1658        let CsrConnection {
1659            url,
1660            tls_root_cert,
1661            tls_identity,
1662            http_auth,
1663            tunnel,
1664        } = self;
1665        CsrConnection {
1666            url,
1667            tls_root_cert,
1668            tls_identity,
1669            http_auth,
1670            tunnel: tunnel.into_inline_connection(r),
1671        }
1672    }
1673}
1674
1675impl<C: ConnectionAccess> CsrConnection<C> {
1676    fn validate_by_default(&self) -> bool {
1677        true
1678    }
1679}
1680
1681impl CsrConnection {
1682    /// Constructs a schema registry client from the connection.
1683    pub async fn connect(
1684        &self,
1685        storage_configuration: &StorageConfiguration,
1686        in_task: InTask,
1687    ) -> Result<mz_ccsr::Client, CsrConnectError> {
1688        let mut client_config = mz_ccsr::ClientConfig::new(self.url.clone());
1689        if let Some(root_cert) = &self.tls_root_cert {
1690            let root_cert = root_cert
1691                .get_string(
1692                    in_task,
1693                    &storage_configuration.connection_context.secrets_reader,
1694                )
1695                .await?;
1696            let root_cert = Certificate::from_pem(root_cert.as_bytes())?;
1697            client_config = client_config.add_root_certificate(root_cert);
1698        }
1699
1700        if let Some(tls_identity) = &self.tls_identity {
1701            let key = &storage_configuration
1702                .connection_context
1703                .secrets_reader
1704                .read_string_in_task_if(in_task, tls_identity.key)
1705                .await?;
1706            let cert = tls_identity
1707                .cert
1708                .get_string(
1709                    in_task,
1710                    &storage_configuration.connection_context.secrets_reader,
1711                )
1712                .await?;
1713            let ident = Identity::from_pem(key.as_bytes(), cert.as_bytes())?;
1714            client_config = client_config.identity(ident);
1715        }
1716
1717        if let Some(http_auth) = &self.http_auth {
1718            let username = http_auth
1719                .username
1720                .get_string(
1721                    in_task,
1722                    &storage_configuration.connection_context.secrets_reader,
1723                )
1724                .await?;
1725            let password = match http_auth.password {
1726                None => None,
1727                Some(password) => Some(
1728                    storage_configuration
1729                        .connection_context
1730                        .secrets_reader
1731                        .read_string_in_task_if(in_task, password)
1732                        .await?,
1733                ),
1734            };
1735            client_config = client_config.auth(username, password);
1736        }
1737
1738        // TODO: use types to enforce that the URL has a string hostname.
1739        let host = self
1740            .url
1741            .host_str()
1742            .ok_or_else(|| anyhow!("url missing host"))?;
1743        match &self.tunnel {
1744            Tunnel::Direct => {
1745                // Ensure any host we connect to is resolved to an external address.
1746                let resolved = resolve_address(
1747                    host,
1748                    ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
1749                )
1750                .await?;
1751                client_config = client_config.resolve_to_addrs(
1752                    host,
1753                    &resolved
1754                        .iter()
1755                        .map(|addr| SocketAddr::new(*addr, 0))
1756                        .collect::<Vec<_>>(),
1757                )
1758            }
1759            Tunnel::Ssh(ssh_tunnel) => {
1760                let ssh_tunnel = ssh_tunnel
1761                    .connect(
1762                        storage_configuration,
1763                        host,
1764                        // Honor the URL scheme's default port (443 for https,
1765                        // 80 for http) if no explicit port was provided.
1766                        self.url.port_or_known_default().unwrap_or(80),
1767                        in_task,
1768                    )
1769                    .await
1770                    .map_err(CsrConnectError::Ssh)?;
1771
1772                // Carefully inject the SSH tunnel into the client
1773                // configuration. This is delicate because we need TLS
1774                // verification to continue to use the remote hostname rather
1775                // than the tunnel hostname.
1776
1777                client_config = client_config
1778                    // `resolve_to_addrs` allows us to rewrite the hostname
1779                    // at the DNS level, which means the TCP connection is
1780                    // correctly routed through the tunnel, but TLS verification
1781                    // is still performed against the remote hostname.
1782                    // Unfortunately the port here is ignored if the URL also
1783                    // specifies a port...
1784                    .resolve_to_addrs(host, &[SocketAddr::new(ssh_tunnel.local_addr().ip(), 0)])
1785                    // ...so we also dynamically rewrite the URL to use the
1786                    // current port for the SSH tunnel.
1787                    //
1788                    // WARNING: this is brittle, because we only dynamically
1789                    // update the client configuration with the tunnel *port*,
1790                    // and not the hostname This works fine in practice, because
1791                    // only the SSH tunnel port will change if the tunnel fails
1792                    // and has to be restarted (the hostname is always
1793                    // 127.0.0.1)--but this is an an implementation detail of
1794                    // the SSH tunnel code that we're relying on.
1795                    .dynamic_url({
1796                        let remote_url = self.url.clone();
1797                        move || {
1798                            let mut url = remote_url.clone();
1799                            url.set_port(Some(ssh_tunnel.local_addr().port()))
1800                                .expect("cannot fail");
1801                            url
1802                        }
1803                    });
1804            }
1805            Tunnel::AwsPrivatelink(connection) => {
1806                assert_none!(connection.port);
1807
1808                let privatelink_host = mz_cloud_resources::vpc_endpoint_host(
1809                    connection.connection_id,
1810                    connection.availability_zone.as_deref(),
1811                );
1812                let addrs: Vec<_> = net::lookup_host((privatelink_host, 0))
1813                    .await
1814                    .context("resolving PrivateLink host")?
1815                    .collect();
1816                client_config = client_config.resolve_to_addrs(host, &addrs)
1817            }
1818            Tunnel::AwsPrivatelinks(_) => {
1819                unreachable!("MATCHING broker rules are only available for Kafka connections.");
1820            }
1821        }
1822
1823        Ok(client_config.build()?)
1824    }
1825
1826    async fn validate(
1827        &self,
1828        _id: CatalogItemId,
1829        storage_configuration: &StorageConfiguration,
1830    ) -> Result<(), anyhow::Error> {
1831        let client = self
1832            .connect(
1833                storage_configuration,
1834                // We are in a normal tokio context during validation, already.
1835                InTask::No,
1836            )
1837            .await?;
1838        client.list_subjects().await?;
1839        Ok(())
1840    }
1841}
1842
1843impl<C: ConnectionAccess> AlterCompatible for CsrConnection<C> {
1844    fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
1845        let CsrConnection {
1846            tunnel,
1847            // All non-tunnel fields may change
1848            url: _,
1849            tls_root_cert: _,
1850            tls_identity: _,
1851            http_auth: _,
1852        } = self;
1853
1854        let compatibility_checks = [(tunnel.alter_compatible(id, &other.tunnel).is_ok(), "tunnel")];
1855
1856        for (compatible, field) in compatibility_checks {
1857            if !compatible {
1858                tracing::warn!(
1859                    "CsrConnection incompatible at {field}:\nself:\n{:#?}\n\nother\n{:#?}",
1860                    self,
1861                    other
1862                );
1863
1864                return Err(AlterError { id });
1865            }
1866        }
1867        Ok(())
1868    }
1869}
1870
1871/// A connection to an AWS Glue Schema Registry.
1872///
1873/// AWS credentials, region, and endpoint are inherited from the referenced
1874/// [`AwsConnection`]; this struct only carries the per-registry settings.
1875#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
1876pub struct GlueSchemaRegistryConnection<C: ConnectionAccess = InlinedConnection> {
1877    /// The referenced AWS connection that supplies credentials, region, and
1878    /// (optional) endpoint.
1879    pub aws_connection: AwsConnectionReference<C>,
1880    /// The Glue Schema Registry name within the AWS account/region.
1881    pub registry_name: String,
1882}
1883
1884impl<R: ConnectionResolver> IntoInlineConnection<GlueSchemaRegistryConnection, R>
1885    for GlueSchemaRegistryConnection<ReferencedConnection>
1886{
1887    fn into_inline_connection(self, r: R) -> GlueSchemaRegistryConnection {
1888        let GlueSchemaRegistryConnection {
1889            aws_connection,
1890            registry_name,
1891        } = self;
1892        GlueSchemaRegistryConnection {
1893            aws_connection: aws_connection.into_inline_connection(&r),
1894            registry_name,
1895        }
1896    }
1897}
1898
1899impl<C: ConnectionAccess> GlueSchemaRegistryConnection<C> {
1900    fn validate_by_default(&self) -> bool {
1901        // Matches CSR: default-validate so a bad registry name fails at
1902        // `CREATE CONNECTION` rather than surfacing later on first use.
1903        // Users can still opt out with `WITH (VALIDATE = false)`.
1904        true
1905    }
1906}
1907
1908impl GlueSchemaRegistryConnection {
1909    async fn validate(
1910        &self,
1911        _id: CatalogItemId,
1912        storage_configuration: &StorageConfiguration,
1913    ) -> Result<(), anyhow::Error> {
1914        let enforce_external_addresses =
1915            crate::dyncfgs::ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set());
1916        let sdk_config = self
1917            .aws_connection
1918            .connection
1919            .load_sdk_config(
1920                &storage_configuration.connection_context,
1921                self.aws_connection.connection_id,
1922                // We are in a normal tokio context during validation.
1923                InTask::No,
1924                enforce_external_addresses,
1925            )
1926            .await?;
1927        let client = mz_aws_glue_schema_registry::ClientConfig::new(sdk_config).build();
1928        match client.get_registry(&self.registry_name).await {
1929            Ok(_) => Ok(()),
1930            Err(mz_aws_glue_schema_registry::GetRegistryError::NotFound) => Err(anyhow!(
1931                "AWS Glue Schema Registry {:?} does not exist in the configured account/region",
1932                self.registry_name
1933            )),
1934            Err(err) => Err(anyhow::Error::new(err).context(format!(
1935                "failed to validate AWS Glue Schema Registry connection (registry={:?})",
1936                self.registry_name
1937            ))),
1938        }
1939    }
1940}
1941
1942impl<C: ConnectionAccess> AlterCompatible for GlueSchemaRegistryConnection<C> {
1943    fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
1944        let GlueSchemaRegistryConnection {
1945            registry_name,
1946            // The referenced AWS connection itself may be swapped; matches
1947            // the permissive policy of MySqlConnection / SqlServerConnection.
1948            aws_connection: _,
1949        } = self;
1950
1951        let compatibility_checks = [(registry_name == &other.registry_name, "registry_name")];
1952
1953        for (compatible, field) in compatibility_checks {
1954            if !compatible {
1955                tracing::warn!(
1956                    "GlueSchemaRegistryConnection incompatible at {field}:\nself:\n{:#?}\n\nother\n{:#?}",
1957                    self,
1958                    other
1959                );
1960
1961                return Err(AlterError { id });
1962            }
1963        }
1964        Ok(())
1965    }
1966}
1967
1968/// A TLS key pair used for client identity.
1969#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
1970pub struct TlsIdentity {
1971    /// The client's TLS public certificate in PEM format.
1972    pub cert: StringOrSecret,
1973    /// The ID of the secret containing the client's TLS private key in PEM
1974    /// format.
1975    pub key: CatalogItemId,
1976}
1977
1978/// HTTP authentication credentials in a [`CsrConnection`].
1979#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
1980pub struct CsrConnectionHttpAuth {
1981    /// The username.
1982    pub username: StringOrSecret,
1983    /// The ID of the secret containing the password, if any.
1984    pub password: Option<CatalogItemId>,
1985}
1986
1987/// A connection to a PostgreSQL server.
1988#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
1989pub struct PostgresConnection<C: ConnectionAccess = InlinedConnection> {
1990    /// The hostname of the server.
1991    pub host: String,
1992    /// The port of the server.
1993    pub port: u16,
1994    /// The name of the database to connect to.
1995    pub database: String,
1996    /// The username to authenticate as.
1997    pub user: StringOrSecret,
1998    /// An optional password for authentication.
1999    pub password: Option<CatalogItemId>,
2000    /// A tunnel through which to route traffic.
2001    pub tunnel: Tunnel<C>,
2002    /// Whether to use TLS for encryption, authentication, or both.
2003    pub tls_mode: SslMode,
2004    /// An optional root TLS certificate in PEM format, to verify the server's
2005    /// identity.
2006    pub tls_root_cert: Option<StringOrSecret>,
2007    /// An optional TLS client certificate for authentication.
2008    pub tls_identity: Option<TlsIdentity>,
2009}
2010
2011impl<R: ConnectionResolver> IntoInlineConnection<PostgresConnection, R>
2012    for PostgresConnection<ReferencedConnection>
2013{
2014    fn into_inline_connection(self, r: R) -> PostgresConnection {
2015        let PostgresConnection {
2016            host,
2017            port,
2018            database,
2019            user,
2020            password,
2021            tunnel,
2022            tls_mode,
2023            tls_root_cert,
2024            tls_identity,
2025        } = self;
2026
2027        PostgresConnection {
2028            host,
2029            port,
2030            database,
2031            user,
2032            password,
2033            tunnel: tunnel.into_inline_connection(r),
2034            tls_mode,
2035            tls_root_cert,
2036            tls_identity,
2037        }
2038    }
2039}
2040
2041impl<C: ConnectionAccess> PostgresConnection<C> {
2042    fn validate_by_default(&self) -> bool {
2043        true
2044    }
2045}
2046
2047impl PostgresConnection<InlinedConnection> {
2048    pub async fn config(
2049        &self,
2050        secrets_reader: &Arc<dyn mz_secrets::SecretsReader>,
2051        storage_configuration: &StorageConfiguration,
2052        in_task: InTask,
2053    ) -> Result<mz_postgres_util::Config, anyhow::Error> {
2054        let params = &storage_configuration.parameters;
2055
2056        let mut config = tokio_postgres::Config::new();
2057        config
2058            .host(&self.host)
2059            .port(self.port)
2060            .dbname(&self.database)
2061            .user(&self.user.get_string(in_task, secrets_reader).await?)
2062            .ssl_mode(self.tls_mode);
2063        if let Some(password) = self.password {
2064            let password = secrets_reader
2065                .read_string_in_task_if(in_task, password)
2066                .await?;
2067            config.password(password);
2068        }
2069        if let Some(tls_root_cert) = &self.tls_root_cert {
2070            let tls_root_cert = tls_root_cert.get_string(in_task, secrets_reader).await?;
2071            config.ssl_root_cert(tls_root_cert.as_bytes());
2072        }
2073        if let Some(tls_identity) = &self.tls_identity {
2074            let cert = tls_identity
2075                .cert
2076                .get_string(in_task, secrets_reader)
2077                .await?;
2078            let key = secrets_reader
2079                .read_string_in_task_if(in_task, tls_identity.key)
2080                .await?;
2081            config.ssl_cert(cert.as_bytes()).ssl_key(key.as_bytes());
2082        }
2083
2084        if let Some(connect_timeout) = params.pg_source_connect_timeout {
2085            config.connect_timeout(connect_timeout);
2086        }
2087        if let Some(keepalives_retries) = params.pg_source_tcp_keepalives_retries {
2088            config.keepalives_retries(keepalives_retries);
2089        }
2090        if let Some(keepalives_idle) = params.pg_source_tcp_keepalives_idle {
2091            config.keepalives_idle(keepalives_idle);
2092        }
2093        if let Some(keepalives_interval) = params.pg_source_tcp_keepalives_interval {
2094            config.keepalives_interval(keepalives_interval);
2095        }
2096        if let Some(tcp_user_timeout) = params.pg_source_tcp_user_timeout {
2097            config.tcp_user_timeout(tcp_user_timeout);
2098        }
2099
2100        let mut options = vec![];
2101        if let Some(wal_sender_timeout) = params.pg_source_wal_sender_timeout {
2102            options.push(format!(
2103                "--wal_sender_timeout={}",
2104                wal_sender_timeout.as_millis()
2105            ));
2106        };
2107        if params.pg_source_tcp_configure_server {
2108            if let Some(keepalives_retries) = params.pg_source_tcp_keepalives_retries {
2109                options.push(format!("--tcp_keepalives_count={}", keepalives_retries));
2110            }
2111            if let Some(keepalives_idle) = params.pg_source_tcp_keepalives_idle {
2112                options.push(format!(
2113                    "--tcp_keepalives_idle={}",
2114                    keepalives_idle.as_secs()
2115                ));
2116            }
2117            if let Some(keepalives_interval) = params.pg_source_tcp_keepalives_interval {
2118                options.push(format!(
2119                    "--tcp_keepalives_interval={}",
2120                    keepalives_interval.as_secs()
2121                ));
2122            }
2123            if let Some(tcp_user_timeout) = params.pg_source_tcp_user_timeout {
2124                options.push(format!(
2125                    "--tcp_user_timeout={}",
2126                    tcp_user_timeout.as_millis()
2127                ));
2128            }
2129        }
2130        config.options(options.join(" ").as_str());
2131
2132        let tunnel = match &self.tunnel {
2133            Tunnel::Direct => {
2134                // Ensure any host we connect to is resolved to an external address.
2135                let resolved = resolve_address(
2136                    &self.host,
2137                    ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
2138                )
2139                .await?;
2140                mz_postgres_util::TunnelConfig::Direct {
2141                    resolved_ips: Some(resolved),
2142                }
2143            }
2144            Tunnel::Ssh(SshTunnel {
2145                connection_id,
2146                connection,
2147            }) => {
2148                let secret = secrets_reader
2149                    .read_in_task_if(in_task, *connection_id)
2150                    .await?;
2151                let key_pair = SshKeyPair::from_bytes(&secret)?;
2152                // Ensure any ssh-bastion host we connect to is resolved to an external address.
2153                let resolved = resolve_address(
2154                    &connection.host,
2155                    ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
2156                )
2157                .await?;
2158                mz_postgres_util::TunnelConfig::Ssh {
2159                    config: SshTunnelConfig {
2160                        host: resolved
2161                            .iter()
2162                            .map(|a| a.to_string())
2163                            .collect::<BTreeSet<_>>(),
2164                        port: connection.port,
2165                        user: connection.user.clone(),
2166                        key_pair,
2167                    },
2168                }
2169            }
2170            Tunnel::AwsPrivatelink(connection) => {
2171                assert_none!(connection.port);
2172                mz_postgres_util::TunnelConfig::AwsPrivatelink {
2173                    connection_id: connection.connection_id,
2174                }
2175            }
2176            Tunnel::AwsPrivatelinks(_) => {
2177                unreachable!("MATCHING broker rules are only available for Kafka connections.");
2178            }
2179        };
2180
2181        Ok(mz_postgres_util::Config::new(
2182            config,
2183            tunnel,
2184            params.ssh_timeout_config,
2185            in_task,
2186        )?)
2187    }
2188
2189    pub async fn validate(
2190        &self,
2191        _id: CatalogItemId,
2192        storage_configuration: &StorageConfiguration,
2193    ) -> Result<mz_postgres_util::Client, anyhow::Error> {
2194        let config = self
2195            .config(
2196                &storage_configuration.connection_context.secrets_reader,
2197                storage_configuration,
2198                // We are in a normal tokio context during validation, already.
2199                InTask::No,
2200            )
2201            .await?;
2202        let client = config
2203            .connect(
2204                "connection validation",
2205                &storage_configuration.connection_context.ssh_tunnel_manager,
2206            )
2207            .await?;
2208
2209        let wal_level = mz_postgres_util::get_wal_level(&client).await?;
2210
2211        if wal_level < mz_postgres_util::replication::WalLevel::Logical {
2212            Err(PostgresConnectionValidationError::InsufficientWalLevel { wal_level })?;
2213        }
2214
2215        let max_wal_senders = mz_postgres_util::get_max_wal_senders(&client).await?;
2216
2217        if max_wal_senders < 1 {
2218            Err(PostgresConnectionValidationError::ReplicationDisabled)?;
2219        }
2220
2221        let available_replication_slots =
2222            mz_postgres_util::available_replication_slots(&client).await?;
2223
2224        // We need 1 replication slot for the snapshots and 1 for the continuing replication
2225        if available_replication_slots < 2 {
2226            Err(
2227                PostgresConnectionValidationError::InsufficientReplicationSlotsAvailable {
2228                    count: 2,
2229                },
2230            )?;
2231        }
2232
2233        Ok(client)
2234    }
2235}
2236
2237#[derive(Debug, Clone, thiserror::Error)]
2238pub enum PostgresConnectionValidationError {
2239    #[error("PostgreSQL server has insufficient number of replication slots available")]
2240    InsufficientReplicationSlotsAvailable { count: usize },
2241    #[error("server must have wal_level >= logical, but has {wal_level}")]
2242    InsufficientWalLevel {
2243        wal_level: mz_postgres_util::replication::WalLevel,
2244    },
2245    #[error("replication disabled on server")]
2246    ReplicationDisabled,
2247}
2248
2249impl PostgresConnectionValidationError {
2250    pub fn detail(&self) -> Option<String> {
2251        match self {
2252            Self::InsufficientReplicationSlotsAvailable { count } => Some(format!(
2253                "executing this statement requires {} replication slot{}",
2254                count,
2255                if *count == 1 { "" } else { "s" }
2256            )),
2257            _ => None,
2258        }
2259    }
2260
2261    pub fn hint(&self) -> Option<String> {
2262        match self {
2263            Self::InsufficientReplicationSlotsAvailable { .. } => Some(
2264                "you might be able to wait for other sources to finish snapshotting and try again"
2265                    .into(),
2266            ),
2267            Self::ReplicationDisabled => Some("set max_wal_senders to a value > 0".into()),
2268            Self::InsufficientWalLevel { .. } => None,
2269        }
2270    }
2271}
2272
2273impl<C: ConnectionAccess> AlterCompatible for PostgresConnection<C> {
2274    fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
2275        let PostgresConnection {
2276            tunnel,
2277            // All non-tunnel options may change arbitrarily
2278            host: _,
2279            port: _,
2280            database: _,
2281            user: _,
2282            password: _,
2283            tls_mode: _,
2284            tls_root_cert: _,
2285            tls_identity: _,
2286        } = self;
2287
2288        let compatibility_checks = [(tunnel.alter_compatible(id, &other.tunnel).is_ok(), "tunnel")];
2289
2290        for (compatible, field) in compatibility_checks {
2291            if !compatible {
2292                tracing::warn!(
2293                    "PostgresConnection incompatible at {field}:\nself:\n{:#?}\n\nother\n{:#?}",
2294                    self,
2295                    other
2296                );
2297
2298                return Err(AlterError { id });
2299            }
2300        }
2301        Ok(())
2302    }
2303}
2304
2305/// Specifies how to tunnel a connection.
2306#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
2307pub enum Tunnel<C: ConnectionAccess = InlinedConnection> {
2308    /// No tunneling.
2309    Direct,
2310    /// Via the specified SSH tunnel connection.
2311    Ssh(SshTunnel<C>),
2312    /// Via the specified AWS PrivateLink connection.
2313    AwsPrivatelink(AwsPrivatelink),
2314    AwsPrivatelinks(AwsPrivatelinks),
2315}
2316
2317impl<R: ConnectionResolver> IntoInlineConnection<Tunnel, R> for Tunnel<ReferencedConnection> {
2318    fn into_inline_connection(self, r: R) -> Tunnel {
2319        match self {
2320            Tunnel::Direct => Tunnel::Direct,
2321            Tunnel::Ssh(ssh) => Tunnel::Ssh(ssh.into_inline_connection(r)),
2322            Tunnel::AwsPrivatelink(awspl) => Tunnel::AwsPrivatelink(awspl),
2323            Tunnel::AwsPrivatelinks(x) => Tunnel::AwsPrivatelinks(x),
2324        }
2325    }
2326}
2327
2328impl<C: ConnectionAccess> AlterCompatible for Tunnel<C> {
2329    fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
2330        let compatible = match (self, other) {
2331            (Self::Ssh(s), Self::Ssh(o)) => s.alter_compatible(id, o).is_ok(),
2332            (s, o) => s == o,
2333        };
2334
2335        if !compatible {
2336            tracing::warn!(
2337                "Tunnel incompatible:\nself:\n{:#?}\n\nother\n{:#?}",
2338                self,
2339                other
2340            );
2341
2342            return Err(AlterError { id });
2343        }
2344
2345        Ok(())
2346    }
2347}
2348
2349/// Specifies which MySQL SSL Mode to use:
2350/// <https://dev.mysql.com/doc/refman/8.0/en/connection-options.html#option_general_ssl-mode>
2351/// This is not available as an enum in the mysql-async crate, so we define our own.
2352#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
2353pub enum MySqlSslMode {
2354    Disabled,
2355    Required,
2356    VerifyCa,
2357    VerifyIdentity,
2358}
2359
2360/// A connection to a MySQL server.
2361#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
2362pub struct MySqlConnection<C: ConnectionAccess = InlinedConnection> {
2363    /// The hostname of the server.
2364    pub host: String,
2365    /// The port of the server.
2366    pub port: u16,
2367    /// The username to authenticate as.
2368    pub user: StringOrSecret,
2369    /// An optional password for authentication.
2370    pub password: Option<CatalogItemId>,
2371    /// A tunnel through which to route traffic.
2372    pub tunnel: Tunnel<C>,
2373    /// Whether to use TLS for encryption, verify the server's certificate, and identity.
2374    pub tls_mode: MySqlSslMode,
2375    /// An optional root TLS certificate in PEM format, to verify the server's
2376    /// identity.
2377    pub tls_root_cert: Option<StringOrSecret>,
2378    /// An optional TLS client certificate for authentication.
2379    pub tls_identity: Option<TlsIdentity>,
2380    /// Reference to the AWS connection information to be used for IAM authenitcation and
2381    /// assuming AWS roles.
2382    pub aws_connection: Option<AwsConnectionReference<C>>,
2383}
2384
2385impl<R: ConnectionResolver> IntoInlineConnection<MySqlConnection, R>
2386    for MySqlConnection<ReferencedConnection>
2387{
2388    fn into_inline_connection(self, r: R) -> MySqlConnection {
2389        let MySqlConnection {
2390            host,
2391            port,
2392            user,
2393            password,
2394            tunnel,
2395            tls_mode,
2396            tls_root_cert,
2397            tls_identity,
2398            aws_connection,
2399        } = self;
2400
2401        MySqlConnection {
2402            host,
2403            port,
2404            user,
2405            password,
2406            tunnel: tunnel.into_inline_connection(&r),
2407            tls_mode,
2408            tls_root_cert,
2409            tls_identity,
2410            aws_connection: aws_connection.map(|aws| aws.into_inline_connection(&r)),
2411        }
2412    }
2413}
2414
2415impl<C: ConnectionAccess> MySqlConnection<C> {
2416    fn validate_by_default(&self) -> bool {
2417        true
2418    }
2419}
2420
2421impl MySqlConnection<InlinedConnection> {
2422    pub async fn config(
2423        &self,
2424        secrets_reader: &Arc<dyn mz_secrets::SecretsReader>,
2425        storage_configuration: &StorageConfiguration,
2426        in_task: InTask,
2427    ) -> Result<mz_mysql_util::Config, anyhow::Error> {
2428        // TODO(roshan): Set appropriate connection timeouts
2429        let mut opts = mysql_async::OptsBuilder::default()
2430            .ip_or_hostname(&self.host)
2431            .tcp_port(self.port)
2432            .user(Some(&self.user.get_string(in_task, secrets_reader).await?));
2433
2434        if let Some(password) = self.password {
2435            let password = secrets_reader
2436                .read_string_in_task_if(in_task, password)
2437                .await?;
2438            opts = opts.pass(Some(password));
2439        }
2440
2441        // Our `MySqlSslMode` enum matches the official MySQL Client `--ssl-mode` parameter values
2442        // which uses opt-in security features (SSL, CA verification, & Identity verification).
2443        // The mysql_async crate `SslOpts` struct uses an opt-out mechanism for each of these, so
2444        // we need to appropriately disable features to match the intent of each enum value.
2445        let mut ssl_opts = match self.tls_mode {
2446            MySqlSslMode::Disabled => None,
2447            MySqlSslMode::Required => Some(
2448                mysql_async::SslOpts::default()
2449                    .with_danger_accept_invalid_certs(true)
2450                    .with_danger_skip_domain_validation(true),
2451            ),
2452            MySqlSslMode::VerifyCa => {
2453                Some(mysql_async::SslOpts::default().with_danger_skip_domain_validation(true))
2454            }
2455            MySqlSslMode::VerifyIdentity => Some(mysql_async::SslOpts::default()),
2456        };
2457
2458        if matches!(
2459            self.tls_mode,
2460            MySqlSslMode::VerifyCa | MySqlSslMode::VerifyIdentity
2461        ) {
2462            if let Some(tls_root_cert) = &self.tls_root_cert {
2463                let tls_root_cert = tls_root_cert.get_string(in_task, secrets_reader).await?;
2464                ssl_opts = ssl_opts.map(|opts| {
2465                    opts.with_root_certs(vec![tls_root_cert.as_bytes().to_vec().into()])
2466                });
2467            }
2468        }
2469
2470        if let Some(identity) = &self.tls_identity {
2471            let key = secrets_reader
2472                .read_string_in_task_if(in_task, identity.key)
2473                .await?;
2474            let cert = identity.cert.get_string(in_task, secrets_reader).await?;
2475            let (der, pass) =
2476                mz_tls_util::pkcs12der_from_pem(key.as_bytes(), cert.as_bytes())?.into_parts();
2477
2478            // Add client identity to SSLOpts
2479            ssl_opts = ssl_opts.map(|opts| {
2480                opts.with_client_identity(Some(
2481                    mysql_async::ClientIdentity::new(der.into()).with_password(pass),
2482                ))
2483            });
2484        }
2485
2486        opts = opts.ssl_opts(ssl_opts);
2487
2488        let tunnel = match &self.tunnel {
2489            Tunnel::Direct => {
2490                // Ensure any host we connect to is resolved to an external address.
2491                let resolved = resolve_address(
2492                    &self.host,
2493                    ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
2494                )
2495                .await?;
2496                mz_mysql_util::TunnelConfig::Direct {
2497                    resolved_ips: Some(resolved),
2498                }
2499            }
2500            Tunnel::Ssh(SshTunnel {
2501                connection_id,
2502                connection,
2503            }) => {
2504                let secret = secrets_reader
2505                    .read_in_task_if(in_task, *connection_id)
2506                    .await?;
2507                let key_pair = SshKeyPair::from_bytes(&secret)?;
2508                // Ensure any ssh-bastion host we connect to is resolved to an external address.
2509                let resolved = resolve_address(
2510                    &connection.host,
2511                    ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
2512                )
2513                .await?;
2514                mz_mysql_util::TunnelConfig::Ssh {
2515                    config: SshTunnelConfig {
2516                        host: resolved
2517                            .iter()
2518                            .map(|a| a.to_string())
2519                            .collect::<BTreeSet<_>>(),
2520                        port: connection.port,
2521                        user: connection.user.clone(),
2522                        key_pair,
2523                    },
2524                }
2525            }
2526            Tunnel::AwsPrivatelink(connection) => {
2527                assert_none!(connection.port);
2528                mz_mysql_util::TunnelConfig::AwsPrivatelink {
2529                    connection_id: connection.connection_id,
2530                }
2531            }
2532            Tunnel::AwsPrivatelinks(_) => {
2533                unreachable!("MATCHING broker rules are only available for Kafka connections.");
2534            }
2535        };
2536
2537        let aws_config = match self.aws_connection.as_ref() {
2538            None => None,
2539            Some(aws_ref) => Some(
2540                aws_ref
2541                    .connection
2542                    .load_sdk_config(
2543                        &storage_configuration.connection_context,
2544                        aws_ref.connection_id,
2545                        in_task,
2546                        ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
2547                    )
2548                    .await?,
2549            ),
2550        };
2551
2552        Ok(mz_mysql_util::Config::new(
2553            opts,
2554            tunnel,
2555            storage_configuration.parameters.ssh_timeout_config,
2556            in_task,
2557            storage_configuration
2558                .parameters
2559                .mysql_source_timeouts
2560                .clone(),
2561            aws_config,
2562        )?)
2563    }
2564
2565    pub async fn validate(
2566        &self,
2567        _id: CatalogItemId,
2568        storage_configuration: &StorageConfiguration,
2569    ) -> Result<MySqlConn, MySqlConnectionValidationError> {
2570        let config = self
2571            .config(
2572                &storage_configuration.connection_context.secrets_reader,
2573                storage_configuration,
2574                // We are in a normal tokio context during validation, already.
2575                InTask::No,
2576            )
2577            .await?;
2578        let mut conn = config
2579            .connect(
2580                "connection validation",
2581                &storage_configuration.connection_context.ssh_tunnel_manager,
2582            )
2583            .await?;
2584
2585        // Check if the MySQL database is configured to allow row-based consistent GTID replication
2586        let mut setting_errors = vec![];
2587        let gtid_res = mz_mysql_util::ensure_gtid_consistency(&mut conn).await;
2588        let binlog_res = mz_mysql_util::ensure_full_row_binlog_format(&mut conn).await;
2589        let order_res = mz_mysql_util::ensure_replication_commit_order(&mut conn).await;
2590        for res in [gtid_res, binlog_res, order_res] {
2591            match res {
2592                Err(MySqlError::InvalidSystemSetting {
2593                    setting,
2594                    expected,
2595                    actual,
2596                }) => {
2597                    setting_errors.push((setting, expected, actual));
2598                }
2599                Err(err) => Err(err)?,
2600                Ok(()) => {}
2601            }
2602        }
2603        if !setting_errors.is_empty() {
2604            Err(MySqlConnectionValidationError::ReplicationSettingsError(
2605                setting_errors,
2606            ))?;
2607        }
2608
2609        Ok(conn)
2610    }
2611}
2612
2613#[derive(Debug, thiserror::Error)]
2614pub enum MySqlConnectionValidationError {
2615    #[error("Invalid MySQL system replication settings")]
2616    ReplicationSettingsError(Vec<(String, String, String)>),
2617    #[error(transparent)]
2618    Client(#[from] MySqlError),
2619    #[error("{}", .0.display_with_causes())]
2620    Other(#[from] anyhow::Error),
2621}
2622
2623impl MySqlConnectionValidationError {
2624    pub fn detail(&self) -> Option<String> {
2625        match self {
2626            Self::ReplicationSettingsError(settings) => Some(format!(
2627                "Invalid MySQL system replication settings: {}",
2628                itertools::join(
2629                    settings.iter().map(|(setting, expected, actual)| format!(
2630                        "{}: expected {}, got {}",
2631                        setting, expected, actual
2632                    )),
2633                    "; "
2634                )
2635            )),
2636            _ => None,
2637        }
2638    }
2639
2640    pub fn hint(&self) -> Option<String> {
2641        match self {
2642            Self::ReplicationSettingsError(_) => {
2643                Some("Set the necessary MySQL database system settings.".into())
2644            }
2645            _ => None,
2646        }
2647    }
2648}
2649
2650impl<C: ConnectionAccess> AlterCompatible for MySqlConnection<C> {
2651    fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
2652        let MySqlConnection {
2653            tunnel,
2654            // All non-tunnel options may change arbitrarily
2655            host: _,
2656            port: _,
2657            user: _,
2658            password: _,
2659            tls_mode: _,
2660            tls_root_cert: _,
2661            tls_identity: _,
2662            aws_connection: _,
2663        } = self;
2664
2665        let compatibility_checks = [(tunnel.alter_compatible(id, &other.tunnel).is_ok(), "tunnel")];
2666
2667        for (compatible, field) in compatibility_checks {
2668            if !compatible {
2669                tracing::warn!(
2670                    "MySqlConnection incompatible at {field}:\nself:\n{:#?}\n\nother\n{:#?}",
2671                    self,
2672                    other
2673                );
2674
2675                return Err(AlterError { id });
2676            }
2677        }
2678        Ok(())
2679    }
2680}
2681
2682/// Details how to connect to an instance of Microsoft SQL Server.
2683///
2684/// For specifics of connecting to SQL Server for purposes of creating a
2685/// Materialize Source, see [`SqlServerSourceConnection`] which wraps this type.
2686///
2687/// [`SqlServerSourceConnection`]: crate::sources::SqlServerSourceConnection
2688#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
2689pub struct SqlServerConnectionDetails<C: ConnectionAccess = InlinedConnection> {
2690    /// The hostname of the server.
2691    pub host: String,
2692    /// The port of the server.
2693    pub port: u16,
2694    /// Database we should connect to.
2695    pub database: String,
2696    /// The username to authenticate as.
2697    pub user: StringOrSecret,
2698    /// Password used for authentication.
2699    pub password: CatalogItemId,
2700    /// A tunnel through which to route traffic.
2701    pub tunnel: Tunnel<C>,
2702    /// Level of encryption to use for the connection.
2703    pub encryption: mz_sql_server_util::config::EncryptionLevel,
2704    /// Certificate validation policy
2705    pub certificate_validation_policy: mz_sql_server_util::config::CertificateValidationPolicy,
2706    /// TLS CA Certifiecate in PEM format
2707    pub tls_root_cert: Option<StringOrSecret>,
2708}
2709
2710impl<C: ConnectionAccess> SqlServerConnectionDetails<C> {
2711    fn validate_by_default(&self) -> bool {
2712        true
2713    }
2714}
2715
2716impl SqlServerConnectionDetails<InlinedConnection> {
2717    /// Attempts to open a connection to the upstream SQL Server instance.
2718    pub async fn validate(
2719        &self,
2720        _id: CatalogItemId,
2721        storage_configuration: &StorageConfiguration,
2722    ) -> Result<mz_sql_server_util::Client, anyhow::Error> {
2723        let config = self
2724            .resolve_config(
2725                &storage_configuration.connection_context.secrets_reader,
2726                storage_configuration,
2727                InTask::No,
2728            )
2729            .await?;
2730        tracing::debug!(?config, "Validating SQL Server connection");
2731
2732        let mut client = mz_sql_server_util::Client::connect(config).await?;
2733
2734        // Ensure the upstream SQL Server instance is configured to allow CDC.
2735        //
2736        // Run all of the checks necessary and collect the errors to provide the best
2737        // guidance as to which system settings need to be enabled.
2738        let mut replication_errors = vec![];
2739        for error in [
2740            mz_sql_server_util::inspect::ensure_database_cdc_enabled(&mut client).await,
2741            mz_sql_server_util::inspect::ensure_snapshot_isolation_enabled(&mut client).await,
2742            mz_sql_server_util::inspect::ensure_sql_server_agent_running(&mut client).await,
2743        ] {
2744            match error {
2745                Err(mz_sql_server_util::SqlServerError::InvalidSystemSetting {
2746                    name,
2747                    expected,
2748                    actual,
2749                }) => replication_errors.push((name, expected, actual)),
2750                Err(other) => Err(other)?,
2751                Ok(()) => (),
2752            }
2753        }
2754        if !replication_errors.is_empty() {
2755            Err(SqlServerConnectionValidationError::ReplicationSettingsError(replication_errors))?;
2756        }
2757
2758        Ok(client)
2759    }
2760
2761    /// Resolve all of the connection details (e.g. read from the [`SecretsReader`])
2762    /// so the returned [`Config`] can be used to open a connection with the
2763    /// upstream system.
2764    ///
2765    /// The provided [`InTask`] argument determines whether any I/O is run in an
2766    /// [`mz_ore::task`] (i.e. a different thread) or directly in the returned
2767    /// future. The main goal here is to prevent running I/O in timely threads.
2768    ///
2769    /// [`Config`]: mz_sql_server_util::Config
2770    pub async fn resolve_config(
2771        &self,
2772        secrets_reader: &Arc<dyn mz_secrets::SecretsReader>,
2773        storage_configuration: &StorageConfiguration,
2774        in_task: InTask,
2775    ) -> Result<mz_sql_server_util::Config, anyhow::Error> {
2776        let dyncfg = storage_configuration.config_set();
2777        let mut inner_config = tiberius::Config::new();
2778
2779        // Setup default connection params.
2780        inner_config.host(&self.host);
2781        inner_config.port(self.port);
2782        inner_config.database(self.database.clone());
2783        inner_config.encryption(self.encryption.into());
2784        match self.certificate_validation_policy {
2785            mz_sql_server_util::config::CertificateValidationPolicy::TrustAll => {
2786                inner_config.trust_cert()
2787            }
2788            mz_sql_server_util::config::CertificateValidationPolicy::VerifyCA => {
2789                inner_config.trust_cert_ca_pem(
2790                    self.tls_root_cert
2791                        .as_ref()
2792                        .unwrap()
2793                        .get_string(in_task, secrets_reader)
2794                        .await
2795                        .context("ca certificate")?,
2796                );
2797            }
2798            mz_sql_server_util::config::CertificateValidationPolicy::VerifySystem => (), // no-op
2799        }
2800
2801        inner_config.application_name("materialize");
2802
2803        // Read our auth settings from
2804        let user = self
2805            .user
2806            .get_string(in_task, secrets_reader)
2807            .await
2808            .context("username")?;
2809        let password = secrets_reader
2810            .read_string_in_task_if(in_task, self.password)
2811            .await
2812            .context("password")?;
2813        // TODO(sql_server3): Support other methods of authentication besides
2814        // username and password.
2815        inner_config.authentication(tiberius::AuthMethod::sql_server(user, password));
2816
2817        // Prevent users from probing our internal network ports by trying to
2818        // connect to localhost, or another non-external IP.
2819        let enforce_external_addresses = ENFORCE_EXTERNAL_ADDRESSES.get(dyncfg);
2820
2821        let tunnel = match &self.tunnel {
2822            Tunnel::Direct => {
2823                let resolved_addresses: Vec<SocketAddr> =
2824                    resolve_address(&self.host, enforce_external_addresses)
2825                        .await?
2826                        .into_iter()
2827                        .map(|ip| SocketAddr::new(ip, self.port))
2828                        .collect();
2829                mz_sql_server_util::config::TunnelConfig::Direct {
2830                    resolved_addresses: resolved_addresses.into_boxed_slice(),
2831                }
2832            }
2833            Tunnel::Ssh(SshTunnel {
2834                connection_id,
2835                connection: ssh_connection,
2836            }) => {
2837                let secret = secrets_reader
2838                    .read_in_task_if(in_task, *connection_id)
2839                    .await
2840                    .context("ssh secret")?;
2841                let key_pair = SshKeyPair::from_bytes(&secret).context("ssh key pair")?;
2842                // Ensure any SSH-bastion host we connect to is resolved to an
2843                // external address.
2844                let addresses = resolve_address(&ssh_connection.host, enforce_external_addresses)
2845                    .await
2846                    .context("ssh tunnel")?;
2847
2848                let config = SshTunnelConfig {
2849                    host: addresses.into_iter().map(|a| a.to_string()).collect(),
2850                    port: ssh_connection.port,
2851                    user: ssh_connection.user.clone(),
2852                    key_pair,
2853                };
2854                mz_sql_server_util::config::TunnelConfig::Ssh {
2855                    config,
2856                    manager: storage_configuration
2857                        .connection_context
2858                        .ssh_tunnel_manager
2859                        .clone(),
2860                    timeout: storage_configuration.parameters.ssh_timeout_config.clone(),
2861                    host: self.host.clone(),
2862                    port: self.port,
2863                }
2864            }
2865            Tunnel::AwsPrivatelink(private_link_connection) => {
2866                assert_none!(private_link_connection.port);
2867                mz_sql_server_util::config::TunnelConfig::AwsPrivatelink {
2868                    connection_id: private_link_connection.connection_id,
2869                    port: self.port,
2870                }
2871            }
2872            Tunnel::AwsPrivatelinks(_) => {
2873                unreachable!("MATCHING broker rules are only available for Kafka connections.");
2874            }
2875        };
2876
2877        Ok(mz_sql_server_util::Config::new(
2878            inner_config,
2879            tunnel,
2880            in_task,
2881        ))
2882    }
2883}
2884
2885#[derive(Debug, Clone, thiserror::Error)]
2886pub enum SqlServerConnectionValidationError {
2887    #[error("Invalid SQL Server system replication settings")]
2888    ReplicationSettingsError(Vec<(String, String, String)>),
2889}
2890
2891impl SqlServerConnectionValidationError {
2892    pub fn detail(&self) -> Option<String> {
2893        match self {
2894            Self::ReplicationSettingsError(settings) => Some(format!(
2895                "Invalid SQL Server system replication settings: {}",
2896                itertools::join(
2897                    settings.iter().map(|(setting, expected, actual)| format!(
2898                        "{}: expected {}, got {}",
2899                        setting, expected, actual
2900                    )),
2901                    "; "
2902                )
2903            )),
2904        }
2905    }
2906
2907    pub fn hint(&self) -> Option<String> {
2908        match self {
2909            _ => None,
2910        }
2911    }
2912}
2913
2914impl<R: ConnectionResolver> IntoInlineConnection<SqlServerConnectionDetails, R>
2915    for SqlServerConnectionDetails<ReferencedConnection>
2916{
2917    fn into_inline_connection(self, r: R) -> SqlServerConnectionDetails {
2918        let SqlServerConnectionDetails {
2919            host,
2920            port,
2921            database,
2922            user,
2923            password,
2924            tunnel,
2925            encryption,
2926            certificate_validation_policy,
2927            tls_root_cert,
2928        } = self;
2929
2930        SqlServerConnectionDetails {
2931            host,
2932            port,
2933            database,
2934            user,
2935            password,
2936            tunnel: tunnel.into_inline_connection(&r),
2937            encryption,
2938            certificate_validation_policy,
2939            tls_root_cert,
2940        }
2941    }
2942}
2943
2944impl<C: ConnectionAccess> AlterCompatible for SqlServerConnectionDetails<C> {
2945    fn alter_compatible(
2946        &self,
2947        id: mz_repr::GlobalId,
2948        other: &Self,
2949    ) -> Result<(), crate::controller::AlterError> {
2950        let SqlServerConnectionDetails {
2951            tunnel,
2952            // TODO(sql_server2): Figure out how these variables are allowed to change.
2953            host: _,
2954            port: _,
2955            database: _,
2956            user: _,
2957            password: _,
2958            encryption: _,
2959            certificate_validation_policy: _,
2960            tls_root_cert: _,
2961        } = self;
2962
2963        let compatibility_checks = [(tunnel.alter_compatible(id, &other.tunnel).is_ok(), "tunnel")];
2964
2965        for (compatible, field) in compatibility_checks {
2966            if !compatible {
2967                tracing::warn!(
2968                    "SqlServerConnectionDetails incompatible at {field}:\nself:\n{:#?}\n\nother\n{:#?}",
2969                    self,
2970                    other
2971                );
2972
2973                return Err(AlterError { id });
2974            }
2975        }
2976        Ok(())
2977    }
2978}
2979
2980/// A connection to an SSH tunnel.
2981#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
2982pub struct SshConnection {
2983    pub host: String,
2984    pub port: u16,
2985    pub user: String,
2986}
2987
2988use self::inline::{
2989    ConnectionAccess, ConnectionResolver, InlinedConnection, IntoInlineConnection,
2990    ReferencedConnection,
2991};
2992
2993impl AlterCompatible for SshConnection {
2994    fn alter_compatible(&self, _id: GlobalId, _other: &Self) -> Result<(), AlterError> {
2995        // Every element of the SSH connection is configurable.
2996        Ok(())
2997    }
2998}
2999
3000/// Specifies an AWS PrivateLink service for a [`Tunnel`].
3001#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
3002pub struct AwsPrivatelink {
3003    /// The ID of the connection to the AWS PrivateLink service.
3004    pub connection_id: CatalogItemId,
3005    // The availability zone to use when connecting to the AWS PrivateLink service.
3006    pub availability_zone: Option<String>,
3007    /// The port to use when connecting to the AWS PrivateLink service, if
3008    /// different from the port in [`KafkaBroker::address`].
3009    pub port: Option<u16>,
3010}
3011
3012impl AlterCompatible for AwsPrivatelink {
3013    fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
3014        let AwsPrivatelink {
3015            connection_id,
3016            availability_zone: _,
3017            port: _,
3018        } = self;
3019
3020        let compatibility_checks = [(connection_id == &other.connection_id, "connection_id")];
3021
3022        for (compatible, field) in compatibility_checks {
3023            if !compatible {
3024                tracing::warn!(
3025                    "AwsPrivatelink incompatible at {field}:\nself:\n{:#?}\n\nother\n{:#?}",
3026                    self,
3027                    other
3028                );
3029
3030                return Err(AlterError { id });
3031            }
3032        }
3033
3034        Ok(())
3035    }
3036}
3037
3038#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
3039pub struct AwsPrivatelinks {
3040    /// Route to brokers through PrivateLink connections according to these rules.
3041    /// Exact-match rules (no wildcards) are used as bootstrap brokers.
3042    /// Wildcard rules are applied dynamically to discovered brokers.
3043    pub rules: Vec<AwsPrivatelinkRule>,
3044}
3045
3046#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
3047pub struct AwsPrivatelinkRule {
3048    /// Given a broker's host:port, should we use this route?
3049    pub pattern: ConnectionRulePattern,
3050    /// Route to the broker through this PrivateLink connection.
3051    pub to: AwsPrivatelink,
3052}
3053
3054/// Specifies an SSH tunnel connection.
3055#[derive(Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
3056pub struct SshTunnel<C: ConnectionAccess = InlinedConnection> {
3057    /// id of the ssh connection
3058    pub connection_id: CatalogItemId,
3059    /// ssh connection object
3060    pub connection: C::Ssh,
3061}
3062
3063impl<R: ConnectionResolver> IntoInlineConnection<SshTunnel, R> for SshTunnel<ReferencedConnection> {
3064    fn into_inline_connection(self, r: R) -> SshTunnel {
3065        let SshTunnel {
3066            connection,
3067            connection_id,
3068        } = self;
3069
3070        SshTunnel {
3071            connection: r.resolve_connection(connection).unwrap_ssh(),
3072            connection_id,
3073        }
3074    }
3075}
3076
3077impl SshTunnel<InlinedConnection> {
3078    /// Like [`SshTunnelConfig::connect`], but the SSH key is loaded from a
3079    /// secret.
3080    async fn connect(
3081        &self,
3082        storage_configuration: &StorageConfiguration,
3083        remote_host: &str,
3084        remote_port: u16,
3085        in_task: InTask,
3086    ) -> Result<ManagedSshTunnelHandle, anyhow::Error> {
3087        // Ensure any ssh-bastion host we connect to is resolved to an external address.
3088        let resolved = resolve_address(
3089            &self.connection.host,
3090            ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
3091        )
3092        .await?;
3093        storage_configuration
3094            .connection_context
3095            .ssh_tunnel_manager
3096            .connect(
3097                SshTunnelConfig {
3098                    host: resolved
3099                        .iter()
3100                        .map(|a| a.to_string())
3101                        .collect::<BTreeSet<_>>(),
3102                    port: self.connection.port,
3103                    user: self.connection.user.clone(),
3104                    key_pair: SshKeyPair::from_bytes(
3105                        &storage_configuration
3106                            .connection_context
3107                            .secrets_reader
3108                            .read_in_task_if(in_task, self.connection_id)
3109                            .await?,
3110                    )?,
3111                },
3112                remote_host,
3113                remote_port,
3114                storage_configuration.parameters.ssh_timeout_config,
3115                in_task,
3116            )
3117            .await
3118    }
3119}
3120
3121impl<C: ConnectionAccess> AlterCompatible for SshTunnel<C> {
3122    fn alter_compatible(&self, id: GlobalId, other: &Self) -> Result<(), AlterError> {
3123        let SshTunnel {
3124            connection_id,
3125            connection,
3126        } = self;
3127
3128        let compatibility_checks = [
3129            (connection_id == &other.connection_id, "connection_id"),
3130            (
3131                connection.alter_compatible(id, &other.connection).is_ok(),
3132                "connection",
3133            ),
3134        ];
3135
3136        for (compatible, field) in compatibility_checks {
3137            if !compatible {
3138                tracing::warn!(
3139                    "SshTunnel incompatible at {field}:\nself:\n{:#?}\n\nother\n{:#?}",
3140                    self,
3141                    other
3142                );
3143
3144                return Err(AlterError { id });
3145            }
3146        }
3147
3148        Ok(())
3149    }
3150}
3151
3152impl SshConnection {
3153    #[allow(clippy::unused_async)]
3154    async fn validate(
3155        &self,
3156        id: CatalogItemId,
3157        storage_configuration: &StorageConfiguration,
3158    ) -> Result<(), anyhow::Error> {
3159        let secret = storage_configuration
3160            .connection_context
3161            .secrets_reader
3162            .read_in_task_if(
3163                // We are in a normal tokio context during validation, already.
3164                InTask::No,
3165                id,
3166            )
3167            .await?;
3168        let key_pair = SshKeyPair::from_bytes(&secret)?;
3169
3170        // Ensure any ssh-bastion host we connect to is resolved to an external address.
3171        let resolved = resolve_address(
3172            &self.host,
3173            ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
3174        )
3175        .await?;
3176
3177        let config = SshTunnelConfig {
3178            host: resolved
3179                .iter()
3180                .map(|a| a.to_string())
3181                .collect::<BTreeSet<_>>(),
3182            port: self.port,
3183            user: self.user.clone(),
3184            key_pair,
3185        };
3186        // Note that we do NOT use the `SshTunnelManager` here, as we want to validate that we
3187        // can actually create a new connection to the ssh bastion, without tunneling.
3188        config
3189            .validate(storage_configuration.parameters.ssh_timeout_config)
3190            .await
3191    }
3192
3193    fn validate_by_default(&self) -> bool {
3194        false
3195    }
3196}
3197
3198impl AwsPrivatelinkConnection {
3199    #[allow(clippy::unused_async)]
3200    async fn validate(
3201        &self,
3202        id: CatalogItemId,
3203        storage_configuration: &StorageConfiguration,
3204    ) -> Result<(), anyhow::Error> {
3205        let Some(ref cloud_resource_reader) = storage_configuration
3206            .connection_context
3207            .cloud_resource_reader
3208        else {
3209            return Err(anyhow!("AWS PrivateLink connections are unsupported"));
3210        };
3211
3212        // No need to optionally run this in a task, as we are just validating from envd.
3213        let status = cloud_resource_reader.read(id).await?;
3214
3215        let availability = status
3216            .conditions
3217            .as_ref()
3218            .and_then(|conditions| conditions.iter().find(|c| c.type_ == "Available"));
3219
3220        match availability {
3221            Some(condition) if condition.status == "True" => Ok(()),
3222            Some(condition) => Err(anyhow!("{}", condition.message)),
3223            None => Err(anyhow!("Endpoint availability is unknown")),
3224        }
3225    }
3226
3227    fn validate_by_default(&self) -> bool {
3228        false
3229    }
3230}