Skip to main content

risingwave_connector/sink/elasticsearch_opensearch/
elasticsearch.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;
16
17use risingwave_common::catalog::Schema;
18use tonic::async_trait;
19
20use super::super::writer::{AsyncTruncateLogSinkerOf, AsyncTruncateSinkWriterExt};
21use super::super::{Sink, SinkError, SinkParam, SinkWriterParam};
22use super::elasticsearch_opensearch_client::ElasticSearchOpenSearchSinkWriter;
23use super::elasticsearch_opensearch_config::{ElasticSearchConfig, ElasticSearchOpenSearchConfig};
24use crate::enforce_secret::EnforceSecret;
25use crate::sink::Result;
26
27pub const ES_SINK: &str = "elasticsearch";
28
29#[derive(Debug)]
30pub struct ElasticSearchSink {
31    config: ElasticSearchOpenSearchConfig,
32    schema: Schema,
33    pk_indices: Vec<usize>,
34    is_append_only: bool,
35}
36
37impl EnforceSecret for ElasticSearchSink {
38    fn enforce_secret<'a>(
39        prop_iter: impl Iterator<Item = &'a str>,
40    ) -> crate::error::ConnectorResult<()> {
41        for prop in prop_iter {
42            ElasticSearchOpenSearchConfig::enforce_one(prop)?;
43        }
44        Ok(())
45    }
46}
47
48#[async_trait]
49impl TryFrom<SinkParam> for ElasticSearchSink {
50    type Error = SinkError;
51
52    fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
53        let schema = param.schema();
54        let pk_indices = param.downstream_pk_or_empty();
55        let config = ElasticSearchConfig::from_btreemap(param.properties)?.inner;
56        Ok(Self {
57            config,
58            schema,
59            pk_indices,
60            is_append_only: param.sink_type.is_append_only(),
61        })
62    }
63}
64
65impl Sink for ElasticSearchSink {
66    type LogSinker = AsyncTruncateLogSinkerOf<ElasticSearchOpenSearchSinkWriter>;
67
68    const SINK_NAME: &'static str = ES_SINK;
69
70    fn validate_unknown_fields(&self) -> Result<()> {
71        crate::sink::validate_sink_unknown_fields(&self.config)
72    }
73
74    fn support_schema_change() -> bool {
75        true
76    }
77
78    async fn validate(&self) -> Result<()> {
79        self.config.validate_config(&self.schema)?;
80        let client = self.config.build_client(Self::SINK_NAME)?;
81        client.ping().await?;
82        Ok(())
83    }
84
85    fn validate_alter_config(config: &BTreeMap<String, String>) -> Result<()> {
86        ElasticSearchConfig::from_btreemap(config.clone())?;
87        Ok(())
88    }
89
90    async fn new_log_sinker(&self, _writer_param: SinkWriterParam) -> Result<Self::LogSinker> {
91        Ok(ElasticSearchOpenSearchSinkWriter::new(
92            self.config.clone(),
93            self.schema.clone(),
94            self.pk_indices.clone(),
95            Self::SINK_NAME,
96            self.is_append_only,
97        )?
98        .into_log_sinker(self.config.concurrent_requests))
99    }
100}