risingwave_connector/sink/elasticsearch_opensearch/
elasticsearch_opensearch_config.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, HashSet};
16
17use anyhow::anyhow;
18use maplit::hashset;
19use risingwave_common::catalog::Schema;
20use risingwave_common::types::DataType;
21use serde::Deserialize;
22use serde_with::{DisplayFromStr, serde_as};
23use url::Url;
24use with_options::WithOptions;
25
26use super::super::SinkError;
27use super::elasticsearch::ES_SINK;
28use super::elasticsearch_opensearch_client::ElasticSearchOpenSearchClient;
29use super::opensearch::OPENSEARCH_SINK;
30use crate::connector_common::ElasticsearchConnection;
31use crate::enforce_secret::EnforceSecret;
32use crate::error::ConnectorError;
33use crate::sink::Result;
34
35pub const ES_OPTION_DELIMITER: &str = "delimiter";
36pub const ES_OPTION_INDEX_COLUMN: &str = "index_column";
37pub const ES_OPTION_INDEX: &str = "index";
38pub const ES_OPTION_ROUTING_COLUMN: &str = "routing_column";
39
40#[serde_as]
41#[derive(Deserialize, Debug, Clone, WithOptions)]
42pub struct ElasticSearchConfig {
43    #[serde(flatten)]
44    pub inner: ElasticSearchOpenSearchConfig,
45}
46
47#[serde_as]
48#[derive(Deserialize, Debug, Clone, WithOptions)]
49pub struct OpenSearchConfig {
50    #[serde(flatten)]
51    pub inner: ElasticSearchOpenSearchConfig,
52}
53
54#[serde_as]
55#[derive(Deserialize, Debug, Clone, WithOptions)]
56pub struct ElasticSearchOpenSearchConfig {
57    #[serde(rename = "url")]
58    pub url: String,
59    /// The index's name of elasticsearch or openserach
60    #[serde(rename = "index")]
61    pub index: Option<String>,
62    /// If pk is set, then "pk1+delimiter+pk2+delimiter..." will be used as the key, if pk is not set, we will just use the first column as the key.
63    #[serde(rename = "delimiter")]
64    pub delimiter: Option<String>,
65    /// The username of elasticsearch or openserach
66    #[serde(rename = "username")]
67    pub username: Option<String>,
68    /// The username of elasticsearch or openserach
69    #[serde(rename = "password")]
70    pub password: Option<String>,
71    /// It is used for dynamic index, if it is be set, the value of this column will be used as the index. It and `index` can only set one
72    #[serde(rename = "index_column")]
73    pub index_column: Option<String>,
74
75    /// It is used for dynamic route, if it is be set, the value of this column will be used as the route
76    #[serde(rename = "routing_column")]
77    pub routing_column: Option<String>,
78
79    #[serde(flatten)]
80    pub unknown_fields: std::collections::HashMap<String, String>,
81
82    #[serde(rename = "retry_on_conflict")]
83    #[serde_as(as = "DisplayFromStr")]
84    #[serde(default = "default_retry_on_conflict")]
85    pub retry_on_conflict: i32,
86
87    #[serde(rename = "batch_num_messages")]
88    #[serde_as(as = "DisplayFromStr")]
89    #[serde(default = "default_batch_num_messages")]
90    #[with_option(allow_alter_on_fly)]
91    pub batch_num_messages: usize,
92
93    #[serde(rename = "batch_size_kb")]
94    #[serde_as(as = "DisplayFromStr")]
95    #[serde(default = "default_batch_size_kb")]
96    #[with_option(allow_alter_on_fly)]
97    pub batch_size_kb: usize,
98
99    #[serde(rename = "concurrent_requests")]
100    #[serde_as(as = "DisplayFromStr")]
101    #[serde(default = "default_concurrent_requests")]
102    #[with_option(allow_alter_on_fly)]
103    pub concurrent_requests: usize,
104
105    #[serde(default = "default_type")]
106    pub r#type: String,
107}
108
109crate::impl_sink_unknown_fields!(ElasticSearchOpenSearchConfig);
110
111impl EnforceSecret for ElasticSearchOpenSearchConfig {
112    const ENFORCE_SECRET_PROPERTIES: phf::Set<&'static str> = phf::phf_set! {
113        "username",
114        "password",
115    };
116}
117
118fn default_type() -> String {
119    "upsert".to_owned()
120}
121
122fn default_retry_on_conflict() -> i32 {
123    3
124}
125
126fn default_batch_num_messages() -> usize {
127    512
128}
129
130fn default_batch_size_kb() -> usize {
131    5 * 1024
132}
133
134fn default_concurrent_requests() -> usize {
135    1024
136}
137
138impl TryFrom<&ElasticsearchConnection> for ElasticSearchOpenSearchConfig {
139    type Error = ConnectorError;
140
141    fn try_from(value: &ElasticsearchConnection) -> std::result::Result<Self, Self::Error> {
142        let allowed_fields: HashSet<&str> = hashset!["url", "username", "password"]; // from ElasticsearchOpenSearchConfig
143
144        for k in value.0.keys() {
145            if !allowed_fields.contains(k.as_str()) {
146                return Err(ConnectorError::from(anyhow!(
147                    "Invalid field: {}, allowed fields: {:?}",
148                    k,
149                    allowed_fields
150                )));
151            }
152        }
153
154        let config = serde_json::from_value::<ElasticSearchOpenSearchConfig>(
155            serde_json::to_value(value.0.clone()).unwrap(),
156        )
157        .map_err(|e| SinkError::Config(anyhow!(e)))?;
158        Ok(config)
159    }
160}
161
162impl ElasticSearchConfig {
163    pub fn from_btreemap(properties: BTreeMap<String, String>) -> Result<Self> {
164        let config = serde_json::from_value::<ElasticSearchConfig>(
165            serde_json::to_value(properties).unwrap(),
166        )
167        .map_err(|e| SinkError::Config(anyhow!(e)))?;
168        Ok(config)
169    }
170}
171
172impl OpenSearchConfig {
173    pub fn from_btreemap(properties: BTreeMap<String, String>) -> Result<Self> {
174        let config =
175            serde_json::from_value::<OpenSearchConfig>(serde_json::to_value(properties).unwrap())
176                .map_err(|e| SinkError::Config(anyhow!(e)))?;
177        Ok(config)
178    }
179}
180
181impl ElasticSearchOpenSearchConfig {
182    pub fn from_btreemap(properties: BTreeMap<String, String>) -> Result<Self> {
183        let config = serde_json::from_value::<ElasticSearchOpenSearchConfig>(
184            serde_json::to_value(properties).unwrap(),
185        )
186        .map_err(|e| SinkError::Config(anyhow!(e)))?;
187        Ok(config)
188    }
189
190    pub fn build_client(&self, connector: &str) -> Result<ElasticSearchOpenSearchClient> {
191        let check_username_password = || -> Result<()> {
192            if self.username.is_some() && self.password.is_none() {
193                return Err(SinkError::Config(anyhow!(
194                    "please set the password when the username is set."
195                )));
196            }
197            if self.username.is_none() && self.password.is_some() {
198                return Err(SinkError::Config(anyhow!(
199                    "please set the username when the password is set."
200                )));
201            }
202            Ok(())
203        };
204        let url =
205            Url::parse(&self.url).map_err(|e| SinkError::ElasticSearchOpenSearch(anyhow!(e)))?;
206        if connector.eq(ES_SINK) {
207            let mut transport_builder = elasticsearch::http::transport::TransportBuilder::new(
208                elasticsearch::http::transport::SingleNodeConnectionPool::new(url),
209            );
210            if let Some(username) = &self.username
211                && let Some(password) = &self.password
212            {
213                transport_builder = transport_builder.auth(
214                    elasticsearch::auth::Credentials::Basic(username.clone(), password.clone()),
215                );
216            }
217            check_username_password()?;
218            let transport = transport_builder
219                .build()
220                .map_err(|e| SinkError::ElasticSearchOpenSearch(anyhow!(e)))?;
221            let client = elasticsearch::Elasticsearch::new(transport);
222            Ok(ElasticSearchOpenSearchClient::ElasticSearch(client))
223        } else if connector.eq(OPENSEARCH_SINK) {
224            let mut transport_builder = opensearch::http::transport::TransportBuilder::new(
225                opensearch::http::transport::SingleNodeConnectionPool::new(url),
226            );
227            if let Some(username) = &self.username
228                && let Some(password) = &self.password
229            {
230                transport_builder = transport_builder.auth(opensearch::auth::Credentials::Basic(
231                    username.clone(),
232                    password.clone(),
233                ));
234            }
235            check_username_password()?;
236            let transport = transport_builder
237                .build()
238                .map_err(|e| SinkError::ElasticSearchOpenSearch(anyhow!(e)))?;
239            let client = opensearch::OpenSearch::new(transport);
240            Ok(ElasticSearchOpenSearchClient::OpenSearch(client))
241        } else {
242            panic!(
243                "connector type must be {} or {}, but get {}",
244                ES_SINK, OPENSEARCH_SINK, connector
245            );
246        }
247    }
248
249    pub fn validate_config(&self, schema: &Schema) -> Result<()> {
250        if self.index_column.is_some() && self.index.is_some()
251            || self.index_column.is_none() && self.index.is_none()
252        {
253            return Err(SinkError::Config(anyhow!(
254                "please set only one of the 'index_column' or 'index' properties."
255            )));
256        }
257
258        if let Some(index_column) = &self.index_column {
259            let filed = schema
260                .fields()
261                .iter()
262                .find(|f| &f.name == index_column)
263                .unwrap();
264            if filed.data_type() != DataType::Varchar {
265                return Err(SinkError::Config(anyhow!(
266                    "please ensure the data type of {} is varchar.",
267                    index_column
268                )));
269            }
270        }
271
272        if let Some(routing_column) = &self.routing_column {
273            let filed = schema
274                .fields()
275                .iter()
276                .find(|f| &f.name == routing_column)
277                .unwrap();
278            if filed.data_type() != DataType::Varchar {
279                return Err(SinkError::Config(anyhow!(
280                    "please ensure the data type of {} is varchar.",
281                    routing_column
282                )));
283            }
284        }
285        Ok(())
286    }
287
288    pub fn get_index_column_index(&self, schema: &Schema) -> Result<Option<usize>> {
289        let index_column_idx = self
290            .index_column
291            .as_ref()
292            .map(|n| {
293                schema
294                    .fields()
295                    .iter()
296                    .position(|s| &s.name == n)
297                    .ok_or_else(|| anyhow!("Cannot find {}", ES_OPTION_INDEX_COLUMN))
298            })
299            .transpose()?;
300        Ok(index_column_idx)
301    }
302
303    pub fn get_routing_column_index(&self, schema: &Schema) -> Result<Option<usize>> {
304        let routing_column_idx = self
305            .routing_column
306            .as_ref()
307            .map(|n| {
308                schema
309                    .fields()
310                    .iter()
311                    .position(|s| &s.name == n)
312                    .ok_or_else(|| anyhow!("Cannot find {}", ES_OPTION_ROUTING_COLUMN))
313            })
314            .transpose()?;
315        Ok(routing_column_idx)
316    }
317}
risingwave_connector/sink/elasticsearch_opensearch/elasticsearch_opensearch_config.rs

risingwave_connector/sink/elasticsearch_opensearch/
elasticsearch_opensearch_config.rs
