risingwave_connector/connector_common/
connection.rs

1// Copyright 2024 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::{BTreeMap, HashMap};
16use std::time::Duration;
17
18use anyhow::Context;
19use opendal::Operator;
20use opendal::services::{Azblob, Gcs, S3};
21use phf::{Set, phf_set};
22use rdkafka::ClientConfig;
23use rdkafka::config::RDKafkaLogLevel;
24use rdkafka::consumer::{BaseConsumer, Consumer};
25use risingwave_common::bail;
26use risingwave_common::secret::LocalSecretManager;
27use risingwave_common::util::env_var::env_var_is_true;
28use risingwave_pb::catalog::PbConnection;
29use serde::Deserialize;
30use serde_with::serde_as;
31use tonic::async_trait;
32use url::Url;
33use with_options::WithOptions;
34
35use crate::connector_common::common::DISABLE_DEFAULT_CREDENTIAL;
36use crate::connector_common::{
37    AwsAuthProps, IcebergCommon, IcebergTableIdentifier, KafkaConnectionProps,
38    KafkaPrivateLinkCommon,
39};
40use crate::enforce_secret::EnforceSecret;
41use crate::error::ConnectorResult;
42use crate::schema::schema_registry::Client as ConfluentSchemaRegistryClient;
43use crate::sink::elasticsearch_opensearch::elasticsearch_opensearch_config::ElasticSearchOpenSearchConfig;
44use crate::source::build_connection;
45use crate::source::kafka::{KafkaContextCommon, RwConsumerContext};
46
47pub const SCHEMA_REGISTRY_CONNECTION_TYPE: &str = "schema_registry";
48
49// All XxxConnection structs should implement this trait as well as EnforceSecretOnCloud trait.
50#[async_trait]
51pub trait Connection: Send {
52    async fn validate_connection(&self) -> ConnectorResult<()>;
53}
54
55#[serde_as]
56#[derive(Debug, Clone, Deserialize, WithOptions, PartialEq)]
57#[serde(deny_unknown_fields)]
58pub struct KafkaConnection {
59    #[serde(flatten)]
60    pub inner: KafkaConnectionProps,
61    #[serde(flatten)]
62    pub kafka_private_link_common: KafkaPrivateLinkCommon,
63    #[serde(flatten)]
64    pub aws_auth_props: AwsAuthProps,
65}
66
67impl EnforceSecret for KafkaConnection {
68    fn enforce_secret<'a>(prop_iter: impl Iterator<Item = &'a str>) -> ConnectorResult<()> {
69        for prop in prop_iter {
70            KafkaConnectionProps::enforce_one(prop)?;
71            AwsAuthProps::enforce_one(prop)?;
72        }
73        Ok(())
74    }
75}
76
77pub async fn validate_connection(connection: &PbConnection) -> ConnectorResult<()> {
78    if let Some(ref info) = connection.info {
79        match info {
80            risingwave_pb::catalog::connection::Info::ConnectionParams(cp) => {
81                let options = cp.properties.clone().into_iter().collect();
82                let secret_refs = cp.secret_refs.clone().into_iter().collect();
83                let props_secret_resolved =
84                    LocalSecretManager::global().fill_secrets(options, secret_refs)?;
85                let connection = build_connection(cp.connection_type(), props_secret_resolved)?;
86                connection.validate_connection().await?
87            }
88            #[expect(deprecated)]
89            risingwave_pb::catalog::connection::Info::PrivateLinkService(_) => unreachable!(),
90        }
91    }
92    Ok(())
93}
94
95#[async_trait]
96impl Connection for KafkaConnection {
97    async fn validate_connection(&self) -> ConnectorResult<()> {
98        let client = self.build_client().await?;
99        // describe cluster here
100        client.fetch_metadata(None, Duration::from_secs(10)).await?;
101        Ok(())
102    }
103}
104
105pub fn read_kafka_log_level() -> Option<RDKafkaLogLevel> {
106    let log_level = std::env::var("RISINGWAVE_KAFKA_LOG_LEVEL").ok()?;
107    match log_level.to_uppercase().as_str() {
108        "DEBUG" => Some(RDKafkaLogLevel::Debug),
109        "INFO" => Some(RDKafkaLogLevel::Info),
110        "WARN" => Some(RDKafkaLogLevel::Warning),
111        "ERROR" => Some(RDKafkaLogLevel::Error),
112        "CRITICAL" => Some(RDKafkaLogLevel::Critical),
113        "EMERG" => Some(RDKafkaLogLevel::Emerg),
114        "ALERT" => Some(RDKafkaLogLevel::Alert),
115        "NOTICE" => Some(RDKafkaLogLevel::Notice),
116        _ => None,
117    }
118}
119
120impl KafkaConnection {
121    async fn build_client(&self) -> ConnectorResult<BaseConsumer<RwConsumerContext>> {
122        let mut config = ClientConfig::new();
123        let bootstrap_servers = &self.inner.brokers;
124        let broker_rewrite_map = self.kafka_private_link_common.broker_rewrite_map.clone();
125        config.set("bootstrap.servers", bootstrap_servers);
126        self.inner.set_security_properties(&mut config);
127
128        // dup with Kafka Enumerator
129        let ctx_common = KafkaContextCommon::new(
130            broker_rewrite_map,
131            None,
132            None,
133            self.aws_auth_props.clone(),
134            self.inner.is_aws_msk_iam(),
135        )
136        .await?;
137        let client_ctx = RwConsumerContext::new(ctx_common);
138
139        if let Some(log_level) = read_kafka_log_level() {
140            config.set_log_level(log_level);
141        }
142        let client: BaseConsumer<RwConsumerContext> =
143            config.create_with_context(client_ctx).await?;
144        if self.inner.is_aws_msk_iam() {
145            #[cfg(not(madsim))]
146            client.poll(Duration::from_secs(10)); // note: this is a blocking call
147            #[cfg(madsim)]
148            client.poll(Duration::from_secs(10)).await;
149        }
150        Ok(client)
151    }
152}
153
154#[serde_as]
155#[derive(Debug, Clone, PartialEq, Eq, Deserialize, WithOptions)]
156#[serde(deny_unknown_fields)]
157pub struct IcebergConnection {
158    #[serde(flatten)]
159    pub common: IcebergCommon,
160
161    #[serde(rename = "catalog.jdbc.user")]
162    pub jdbc_user: Option<String>,
163
164    #[serde(rename = "catalog.jdbc.password")]
165    pub jdbc_password: Option<String>,
166}
167
168impl EnforceSecret for IcebergConnection {
169    const ENFORCE_SECRET_PROPERTIES: Set<&'static str> = phf_set! {
170        "s3.access.key",
171        "s3.secret.key",
172        "gcs.credential",
173        "catalog.token",
174    };
175}
176
177#[async_trait]
178impl Connection for IcebergConnection {
179    async fn validate_connection(&self) -> ConnectorResult<()> {
180        let common = &self.common;
181        let is_rest_catalog = common.is_rest_catalog()?;
182
183        let info = match &common.warehouse_path {
184            Some(warehouse_path) => {
185                let is_s3_tables = warehouse_path.starts_with("arn:aws:s3tables");
186                let url = Url::parse(warehouse_path);
187                if (url.is_err() || is_s3_tables) && is_rest_catalog {
188                    // If the warehouse path is not a valid URL, it could be a warehouse name in rest catalog,
189                    // Or it could be a s3tables path, which is not a valid URL but a valid warehouse path,
190                    // so we allow it to pass here.
191                    None
192                } else {
193                    let url =
194                        url.with_context(|| format!("Invalid warehouse path: {}", warehouse_path))?;
195                    let bucket = url
196                        .host_str()
197                        .with_context(|| {
198                            format!("Invalid s3 path: {}, bucket is missing", warehouse_path)
199                        })?
200                        .to_owned();
201                    let root = url.path().trim_start_matches('/').to_owned();
202                    Some((url.scheme().to_owned(), bucket, root))
203                }
204            }
205            None => {
206                if is_rest_catalog {
207                    None
208                } else {
209                    bail!("`warehouse.path` must be set");
210                }
211            }
212        };
213
214        // Test warehouse
215        if let Some((scheme, bucket, root)) = info {
216            match scheme.as_str() {
217                "s3" | "s3a" => {
218                    let mut builder = S3::default();
219                    if let Some(region) = &common.s3_region {
220                        builder = builder.region(region);
221                    }
222                    if let Some(endpoint) = &common.s3_endpoint {
223                        builder = builder.endpoint(endpoint);
224                    }
225                    if let Some(access_key) = &common.s3_access_key {
226                        builder = builder.access_key_id(access_key);
227                    }
228                    if let Some(secret_key) = &common.s3_secret_key {
229                        builder = builder.secret_access_key(secret_key);
230                    }
231                    builder = builder.root(root.as_str()).bucket(bucket.as_str());
232                    let op = Operator::new(builder)?.finish();
233                    op.check().await?;
234                }
235                "gs" | "gcs" => {
236                    let mut builder = Gcs::default();
237                    if let Some(credential) = &common.gcs_credential {
238                        builder = builder.credential(credential);
239                    }
240                    builder = builder.root(root.as_str()).bucket(bucket.as_str());
241                    let op = Operator::new(builder)?.finish();
242                    op.check().await?;
243                }
244                "azblob" => {
245                    let mut builder = Azblob::default();
246                    if let Some(account_name) = &common.azblob_account_name {
247                        builder = builder.account_name(account_name);
248                    }
249                    if let Some(azblob_account_key) = &common.azblob_account_key {
250                        builder = builder.account_key(azblob_account_key);
251                    }
252                    if let Some(azblob_endpoint_url) = &common.azblob_endpoint_url {
253                        builder = builder.endpoint(azblob_endpoint_url);
254                    }
255                    builder = builder.root(root.as_str()).container(bucket.as_str());
256                    let op = Operator::new(builder)?.finish();
257                    op.check().await?;
258                }
259                _ => {
260                    bail!("Unsupported scheme: {}", scheme);
261                }
262            }
263        }
264
265        if env_var_is_true(DISABLE_DEFAULT_CREDENTIAL)
266            && matches!(common.enable_config_load, Some(true))
267        {
268            bail!("`enable_config_load` can't be enabled in this environment");
269        }
270
271        if common.hosted_catalog.unwrap_or(false) {
272            // If `hosted_catalog` is set, we don't need to test the catalog, but just ensure no catalog fields are set.
273            if common.catalog_type.is_some() {
274                bail!("`catalog.type` must not be set when `hosted_catalog` is set");
275            }
276            if common.catalog_uri.is_some() {
277                bail!("`catalog.uri` must not be set when `hosted_catalog` is set");
278            }
279            if common.catalog_name.is_some() {
280                bail!("`catalog.name` must not be set when `hosted_catalog` is set");
281            }
282            if self.jdbc_user.is_some() {
283                bail!("`catalog.jdbc.user` must not be set when `hosted_catalog` is set");
284            }
285            if self.jdbc_password.is_some() {
286                bail!("`catalog.jdbc.password` must not be set when `hosted_catalog` is set");
287            }
288            return Ok(());
289        }
290
291        if common.catalog_type.is_none() {
292            bail!("`catalog.type` must be set");
293        }
294
295        // Test catalog
296        let iceberg_common = common.clone();
297
298        let mut java_map = HashMap::new();
299        if let Some(jdbc_user) = &self.jdbc_user {
300            java_map.insert("jdbc.user".to_owned(), jdbc_user.to_owned());
301        }
302        if let Some(jdbc_password) = &self.jdbc_password {
303            java_map.insert("jdbc.password".to_owned(), jdbc_password.to_owned());
304        }
305        let catalog = iceberg_common
306            .resolve_catalog_config(java_map)?
307            .create_catalog()
308            .await?;
309        // test catalog by `table_exists` api
310        let test_table_ident = IcebergTableIdentifier {
311            database_name: Some("test_database".to_owned()),
312            table_name: "test_table".to_owned(),
313        }
314        .to_table_ident()?;
315        catalog.table_exists(&test_table_ident).await?;
316        Ok(())
317    }
318}
319
320#[serde_as]
321#[derive(Debug, Clone, Deserialize, WithOptions, PartialEq, Hash, Eq)]
322#[serde(deny_unknown_fields)]
323pub struct ConfluentSchemaRegistryConnection {
324    #[serde(rename = "schema.registry")]
325    pub url: String,
326    // ref `SchemaRegistryAuth`
327    #[serde(rename = "schema.registry.username")]
328    pub username: Option<String>,
329    #[serde(rename = "schema.registry.password")]
330    pub password: Option<String>,
331}
332
333#[async_trait]
334impl Connection for ConfluentSchemaRegistryConnection {
335    async fn validate_connection(&self) -> ConnectorResult<()> {
336        // GET /config to validate the connection
337        let client = ConfluentSchemaRegistryClient::try_from(self)?;
338        client.validate_connection().await?;
339        Ok(())
340    }
341}
342
343impl EnforceSecret for ConfluentSchemaRegistryConnection {
344    const ENFORCE_SECRET_PROPERTIES: Set<&'static str> = phf_set! {
345        "schema.registry.password",
346    };
347}
348
349#[derive(Debug, Clone, Deserialize, PartialEq, Hash, Eq)]
350pub struct ElasticsearchConnection(pub BTreeMap<String, String>);
351
352#[async_trait]
353impl Connection for ElasticsearchConnection {
354    async fn validate_connection(&self) -> ConnectorResult<()> {
355        const CONNECTOR: &str = "elasticsearch";
356
357        let config = ElasticSearchOpenSearchConfig::try_from(self)?;
358        let client = config.build_client(CONNECTOR)?;
359        client.ping().await?;
360        Ok(())
361    }
362}
363
364impl EnforceSecret for ElasticsearchConnection {
365    const ENFORCE_SECRET_PROPERTIES: Set<&'static str> = phf_set! {
366        "elasticsearch.password",
367    };
368}
risingwave_connector/connector_common/connection.rs

risingwave_connector/connector_common/
connection.rs
