feat: Connection for Kafka source & sink

Cargo.toml

@@ -160,7 +160,11 @@ arrow-udf-flight = "0.4"
 clap = { version = "4", features = ["cargo", "derive", "env"] }
 # Use a forked version which removes the dependencies on dynamo db to reduce
 # compile time and binary size.
-deltalake = { version = "0.20.1", features = ["s3", "gcs", "datafusion"] }
+deltalake = { version = "0.20.1", features = [
+    "s3",
+    "gcs",
+    "datafusion",
+] }
 itertools = "0.13.0"
 jsonbb = "0.1.4"
 lru = { git = "https://github.com/risingwavelabs/lru-rs.git", rev = "2682b85" }

e2e_test/source_inline/connection/ddl.slt

@@ -0,0 +1,107 @@
+control substitution on
+
+# for non-shared source
+statement ok
+set streaming_use_shared_source to false;
+
+statement ok
+create secret sec_broker with (backend = 'meta') as '${RISEDEV_KAFKA_BOOTSTRAP_SERVERS}';
+
+statement error unknown field `foo`
+create connection conn with (type = 'kafka', properties.bootstrap.server = secret sec_broker, foo = 'bar');
+
+statement error Connection type "kinesis" is not supported
+create connection conn with (type = 'kinesis');
+
+statement ok
+create connection conn with (type = 'kafka', properties.bootstrap.server = secret sec_broker, properties.security.protocol = 'plaintext');
+
+query TTT
+select "name", "type_" from rw_catalog.rw_connections;
+----
+conn CONNECTION_TYPE_KAFKA
+
+# unstable test serialization due to iter on hashmap
+# the "connectiond_params" looks like:
+# {"properties.bootstrap.server":"SECRET sec_broker AS TEXT","properties.security.protocol":"plaintext"}
+
+statement error Permission denied: PermissionDenied: secret used by 1 other objects.
+drop secret sec_broker;
+
+statement error Duplicated key in both WITH clause and Connection catalog: properties.security.protocol
+create table t1 (a int, b varchar) with (
+  connector = 'kafka',
+  connection = conn,
+  topic = 'connection_ddl_1',
+  properties.security.protocol = 'plaintext')
+format plain encode json;
+
+statement error connector kinesis and connection type Kafka are not compatible
+create table t1 (a int, b varchar) with (
+  connector = 'kinesis',
+  connection = conn,
+  stream = 'connection_ddl_1',
+  region = 'us-east-1')
+format plain encode json;
+
+system ok
+rpk topic create connection_ddl_1 -p 1
+
+statement ok
+create table t1 (a int, b varchar) with (
+  connector = 'kafka',
+  connection = conn,
+  topic = 'connection_ddl_1')
+format plain encode json;
+
+statement error Permission denied: PermissionDenied: connection used by 1 other objects.
+drop connection conn;
+
+# Connection & Source & Sink will have independent rely on the secret
+statement error Permission denied: PermissionDenied: secret used by 2 other objects.
+drop secret sec_broker;
+
+statement ok
+create table data_table (a int, b varchar);
+
+statement ok
+insert into data_table values (1, 'a'), (2, 'b'), (3, 'c');
+
+statement ok
+flush;
+
+statement ok
+create sink sink_kafka from data_table with (
+  connector = 'kafka',
+  connection = conn,
+  topic = 'connection_ddl_1'
+) format plain encode json (
+  force_append_only='true'
+);
+
+sleep 3s
+
+query IT rowsort
+select a, b from t1;
+----
+1 a
+2 b
+3 c
+
+statement ok
+drop sink sink_kafka
+
+statement ok
+drop table data_table;
+
+statement ok
+drop table t1;
+
+statement ok
+drop connection conn;
+
+statement ok
+drop secret sec_broker;
+
+statement ok
+set streaming_use_shared_source to true;

proto/catalog.proto

@@ -92,6 +92,9 @@ message StreamSourceInfo {
   // Handle the source relies on any sceret. The key is the propertity name and the value is the secret id and type.
   // For format and encode options.
   map<string, secret.SecretRef> format_encode_secret_refs = 16;
+
+  // ref connection for schema registry
+  optional uint32 connection_id = 17;
 }
 
 message WebhookSourceInfo {
@@ -128,6 +131,8 @@ message Source {
     uint32 associated_table_id = 12;
   }
   string definition = 13;
+
+  // ref connection for connector
   optional uint32 connection_id = 14;
 
   optional uint64 initialized_at_epoch = 15;
@@ -161,6 +166,9 @@ message SinkFormatDesc {
   optional plan_common.EncodeType key_encode = 4;
   // Secret used for format encode options.
   map<string, secret.SecretRef> secret_refs = 5;
+
+  // ref connection for schema registry
+  optional uint32 connection_id = 6;
 }
 
 // the catalog of the sink. There are two kind of schema here. The full schema is all columns
@@ -183,6 +191,8 @@ message Sink {
   uint32 owner = 11;
   map<string, string> properties = 12;
   string definition = 13;
+
+  // ref connection for connector
   optional uint32 connection_id = 14;
   optional uint64 initialized_at_epoch = 15;
   optional uint64 created_at_epoch = 16;
@@ -231,6 +241,19 @@ message Subscription {
   SubscriptionState subscription_state = 19;
 }
 
+message ConnectionParams {
+  enum ConnectionType {
+    CONNECTION_TYPE_UNSPECIFIED = 0;
+    CONNECTION_TYPE_KAFKA = 1;
+    CONNECTION_TYPE_ICEBERG = 2;
+    CONNECTION_TYPE_SCHEMA_REGISTRY = 3;
+  }
+
+  ConnectionType connection_type = 1;
+  map<string, string> properties = 2;
+  map<string, secret.SecretRef> secret_refs = 3;
+}
+
 message Connection {
   message PrivateLinkService {
     enum PrivateLinkProvider {
@@ -251,6 +274,7 @@ message Connection {
   string name = 4;
   oneof info {
     PrivateLinkService private_link_service = 5 [deprecated = true];
+    ConnectionParams connection_params = 7;
   }
   uint32 owner = 6;
 }

proto/ddl_service.proto

@@ -442,7 +442,8 @@ message CreateConnectionRequest {
   uint32 database_id = 2;
   uint32 schema_id = 3;
   oneof payload {
-    PrivateLink private_link = 4;
+    PrivateLink private_link = 4 [deprecated = true];
+    catalog.ConnectionParams connection_params = 6;
   }
   uint32 owner_id = 5;
 }

src/bench/sink_bench/main.rs

@@ -487,6 +487,7 @@ fn mock_from_legacy_type(
             options: Default::default(),
             secret_refs: Default::default(),
             key_encode: None,
+            connection_id: None,
         }))
     } else {
         SinkFormatDesc::from_legacy_type(connector, r#type)

src/connector/src/connector_common/common.rs

@@ -163,7 +163,7 @@ impl AwsAuthProps {
 
 #[serde_as]
 #[derive(Debug, Clone, Deserialize, WithOptions, PartialEq, Hash, Eq)]
-pub struct KafkaConnection {
+pub struct KafkaConnectionProps {
     #[serde(rename = "properties.bootstrap.server", alias = "kafka.brokers")]
     pub brokers: String,
 
@@ -244,6 +244,7 @@ pub struct KafkaConnection {
 #[serde_as]
 #[derive(Debug, Clone, Deserialize, WithOptions)]
 pub struct KafkaCommon {
+    // connection related props are moved to `KafkaConnection`
     #[serde(rename = "topic", alias = "kafka.topic")]
     pub topic: String,
 
@@ -256,7 +257,7 @@ pub struct KafkaCommon {
 }
 
 #[serde_as]
-#[derive(Debug, Clone, Deserialize, WithOptions, PartialEq)]
+#[derive(Debug, Clone, Deserialize, WithOptions, PartialEq, Hash, Eq)]
 pub struct KafkaPrivateLinkCommon {
     /// This is generated from `private_link_targets` and `private_link_endpoint` in frontend, instead of given by users.
     #[serde(rename = "broker.rewrite.endpoints")]
@@ -321,7 +322,32 @@ impl RdKafkaPropertiesCommon {
     }
 }
 
-impl KafkaConnection {
+impl KafkaConnectionProps {
+    #[cfg(test)]
+    pub fn test_default() -> Self {
+        Self {
+            brokers: "localhost:9092".to_string(),
+            security_protocol: None,
+            ssl_ca_location: None,
+            ssl_certificate_location: None,
+            ssl_key_location: None,
+            ssl_ca_pem: None,
+            ssl_certificate_pem: None,
+            ssl_key_pem: None,
+            ssl_key_password: None,
+            ssl_endpoint_identification_algorithm: None,
+            sasl_mechanism: None,
+            sasl_username: None,
+            sasl_password: None,
+            sasl_kerberos_service_name: None,
+            sasl_kerberos_keytab: None,
+            sasl_kerberos_principal: None,
+            sasl_kerberos_kinit_cmd: None,
+            sasl_kerberos_min_time_before_relogin: None,
+            sasl_oathbearer_config: None,
+        }
+    }
+
     pub(crate) fn set_security_properties(&self, config: &mut ClientConfig) {
         // AWS_MSK_IAM
         if self.is_aws_msk_iam() {

src/connector/src/connector_common/connection.rs

@@ -0,0 +1,128 @@
+// Copyright 2024 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::time::Duration;
+
+use rdkafka::consumer::{BaseConsumer, Consumer};
+use rdkafka::ClientConfig;
+use risingwave_common::secret::LocalSecretManager;
+use risingwave_pb::catalog::PbConnection;
+use serde_derive::Deserialize;
+use serde_with::serde_as;
+use tonic::async_trait;
+use with_options::WithOptions;
+
+use crate::connector_common::{AwsAuthProps, KafkaConnectionProps, KafkaPrivateLinkCommon};
+use crate::error::ConnectorResult;
+use crate::source::kafka::{KafkaContextCommon, RwConsumerContext};
+use crate::{dispatch_connection_impl, ConnectionImpl};
+
+#[async_trait]
+pub trait Connection {
+    async fn test_connection(&self) -> ConnectorResult<()>;
+}
+
+#[serde_as]
+#[derive(Debug, Clone, Deserialize, WithOptions, PartialEq)]
+#[serde(deny_unknown_fields)]
+pub struct KafkaConnection {
+    #[serde(flatten)]
+    pub inner: KafkaConnectionProps,
+    #[serde(flatten)]
+    pub kafka_private_link_common: KafkaPrivateLinkCommon,
+    #[serde(flatten)]
+    pub aws_auth_props: AwsAuthProps,
+}
+
+pub async fn validate_connection(connection: &PbConnection) -> ConnectorResult<()> {
+    if let Some(ref info) = connection.info {
+        match info {
+            risingwave_pb::catalog::connection::Info::ConnectionParams(cp) => {
+                let options = cp.properties.clone().into_iter().collect();
+                let secret_refs = cp.secret_refs.clone().into_iter().collect();
+                let props_secret_resolved =
+                    LocalSecretManager::global().fill_secrets(options, secret_refs)?;
+                let connection_impl =
+                    ConnectionImpl::from_proto(cp.connection_type(), props_secret_resolved)?;
+                dispatch_connection_impl!(connection_impl, inner, inner.test_connection().await?)
+            }
+            risingwave_pb::catalog::connection::Info::PrivateLinkService(_) => unreachable!(),
+        }
+    }
+    Ok(())
+}
+
+#[async_trait]
+impl Connection for KafkaConnection {
+    async fn test_connection(&self) -> ConnectorResult<()> {
+        let client = self.build_client().await?;
+        // describe cluster here
+        client.fetch_metadata(None, Duration::from_secs(10)).await?;
+        Ok(())
+    }
+}
+
+impl KafkaConnection {
+    async fn build_client(&self) -> ConnectorResult<BaseConsumer<RwConsumerContext>> {
+        let mut config = ClientConfig::new();
+        let bootstrap_servers = &self.inner.brokers;
+        let broker_rewrite_map = self.kafka_private_link_common.broker_rewrite_map.clone();
+        config.set("bootstrap.servers", bootstrap_servers);
+        self.inner.set_security_properties(&mut config);
+
+        // dup with Kafka Enumerator
+        let ctx_common = KafkaContextCommon::new(
+            broker_rewrite_map,
+            None,
+            None,
+            self.aws_auth_props.clone(),
+            self.inner.is_aws_msk_iam(),
+        )
+        .await?;
+        let client_ctx = RwConsumerContext::new(ctx_common);
+        let client: BaseConsumer<RwConsumerContext> =
+            config.create_with_context(client_ctx).await?;
+        if self.inner.is_aws_msk_iam() {
+            #[cfg(not(madsim))]
+            client.poll(Duration::from_secs(10)); // note: this is a blocking call
+            #[cfg(madsim)]
+            client.poll(Duration::from_secs(10)).await;
+        }
+        Ok(client)
+    }
+}
+
+#[serde_as]
+#[derive(Debug, Clone, PartialEq, Eq, Deserialize, WithOptions)]
+#[serde(deny_unknown_fields)]
+pub struct IcebergConnection {}
+
+#[async_trait]
+impl Connection for IcebergConnection {
+    async fn test_connection(&self) -> ConnectorResult<()> {
+        todo!()
+    }
+}
+
+#[serde_as]
+#[derive(Debug, Clone, Deserialize, WithOptions, PartialEq, Hash, Eq)]
+#[serde(deny_unknown_fields)]
+pub struct SchemaRegistryConnection {}
+
+#[async_trait]
+impl Connection for SchemaRegistryConnection {
+    async fn test_connection(&self) -> ConnectorResult<()> {
+        todo!()
+    }
+}