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