risingwave_connector/sink/
doris.rs

1// Copyright 2023 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, HashSet};
16
17use anyhow::{Context, anyhow};
18use async_trait::async_trait;
19use base64::Engine;
20use base64::engine::general_purpose;
21use bytes::{BufMut, Bytes, BytesMut};
22use risingwave_common::array::{Op, StreamChunk};
23use risingwave_common::catalog::Schema;
24use risingwave_common::types::DataType;
25use serde::{Deserialize, Serialize};
26use serde_json::Value;
27use serde_with::{DisplayFromStr, serde_as};
28use thiserror_ext::AsReport;
29use with_options::WithOptions;
30
31use super::doris_starrocks_connector::{
32    DORIS_DELETE_SIGN, DORIS_SUCCESS_STATUS, HeaderBuilder, InserterInner, InserterInnerBuilder,
33    POOL_IDLE_TIMEOUT,
34};
35use super::{
36    Result, SINK_TYPE_APPEND_ONLY, SINK_TYPE_OPTION, SINK_TYPE_UPSERT, SinkError, SinkWriterMetrics,
37};
38use crate::enforce_secret::EnforceSecret;
39use crate::sink::encoder::{DorisJsonConfig, JsonEncoder, RowEncoder};
40use crate::sink::writer::{LogSinkerOf, SinkWriter, SinkWriterExt};
41use crate::sink::{Sink, SinkParam, SinkWriterParam};
42
43pub const DORIS_SINK: &str = "doris";
44
45const fn default_stream_load_http_timeout_ms() -> u64 {
46    30 * 1000
47}
48
49#[derive(Deserialize, Debug, Clone, WithOptions)]
50pub struct DorisCommon {
51    #[serde(rename = "doris.url")]
52    pub url: String,
53    #[serde(rename = "doris.user")]
54    pub user: String,
55    #[serde(rename = "doris.password")]
56    pub password: String,
57    #[serde(rename = "doris.database")]
58    pub database: String,
59    #[serde(rename = "doris.table")]
60    pub table: String,
61    #[serde(rename = "doris.partial_update")]
62    pub partial_update: Option<String>,
63}
64
65impl EnforceSecret for DorisCommon {
66    const ENFORCE_SECRET_PROPERTIES: phf::Set<&'static str> = phf::phf_set! {
67        "doris.password", "doris.user"
68    };
69}
70
71impl DorisCommon {
72    pub(crate) fn build_get_client(&self) -> DorisSchemaClient {
73        DorisSchemaClient::new(
74            self.url.clone(),
75            self.table.clone(),
76            self.database.clone(),
77            self.user.clone(),
78            self.password.clone(),
79        )
80    }
81}
82
83#[serde_as]
84#[derive(Clone, Debug, Deserialize, WithOptions)]
85pub struct DorisConfig {
86    #[serde(flatten)]
87    pub common: DorisCommon,
88
89    pub r#type: String, // accept "append-only" or "upsert"
90
91    /// The timeout in milliseconds for stream load http request, defaults to 30 seconds.
92    #[serde(
93        rename = "doris.stream_load.http.timeout.ms",
94        default = "default_stream_load_http_timeout_ms"
95    )]
96    #[serde_as(as = "DisplayFromStr")]
97    #[with_option(allow_alter_on_fly)]
98    pub stream_load_http_timeout_ms: u64,
99}
100
101impl EnforceSecret for DorisConfig {
102    fn enforce_one(prop: &str) -> crate::error::ConnectorResult<()> {
103        DorisCommon::enforce_one(prop)
104    }
105}
106
107impl DorisConfig {
108    pub fn from_btreemap(properties: BTreeMap<String, String>) -> Result<Self> {
109        let config =
110            serde_json::from_value::<DorisConfig>(serde_json::to_value(properties).unwrap())
111                .map_err(|e| SinkError::Config(anyhow!(e)))?;
112        if config.r#type != SINK_TYPE_APPEND_ONLY && config.r#type != SINK_TYPE_UPSERT {
113            return Err(SinkError::Config(anyhow!(
114                "`{}` must be {}, or {}",
115                SINK_TYPE_OPTION,
116                SINK_TYPE_APPEND_ONLY,
117                SINK_TYPE_UPSERT
118            )));
119        }
120        Ok(config)
121    }
122}
123
124#[derive(Debug)]
125pub struct DorisSink {
126    pub config: DorisConfig,
127    schema: Schema,
128    pk_indices: Vec<usize>,
129    is_append_only: bool,
130}
131
132impl EnforceSecret for DorisSink {
133    fn enforce_secret<'a>(
134        prop_iter: impl Iterator<Item = &'a str>,
135    ) -> crate::error::ConnectorResult<()> {
136        for prop in prop_iter {
137            DorisConfig::enforce_one(prop)?;
138        }
139        Ok(())
140    }
141}
142
143impl DorisSink {
144    pub fn new(
145        config: DorisConfig,
146        schema: Schema,
147        pk_indices: Vec<usize>,
148        is_append_only: bool,
149    ) -> Result<Self> {
150        Ok(Self {
151            config,
152            schema,
153            pk_indices,
154            is_append_only,
155        })
156    }
157}
158
159impl DorisSink {
160    fn check_column_name_and_type(&self, doris_column_fields: Vec<DorisField>) -> Result<()> {
161        let doris_columns_desc: HashMap<String, String> = doris_column_fields
162            .iter()
163            .map(|s| (s.name.clone(), s.r#type.clone()))
164            .collect();
165
166        let rw_fields_name = self.schema.fields();
167        if rw_fields_name.len() > doris_columns_desc.len() {
168            return Err(SinkError::Doris(
169                "The columns of the sink must be equal to or a superset of the target table's columns.".to_owned(),
170            ));
171        }
172
173        for i in rw_fields_name {
174            let value = doris_columns_desc.get(&i.name).ok_or_else(|| {
175                SinkError::Doris(format!(
176                    "Column name don't find in doris, risingwave is {:?} ",
177                    i.name
178                ))
179            })?;
180            if !Self::check_and_correct_column_type(&i.data_type, value.clone())? {
181                return Err(SinkError::Doris(format!(
182                    "Column type don't match, column name is {:?}. doris type is {:?} risingwave type is {:?} ",
183                    i.name, value, i.data_type
184                )));
185            }
186        }
187        Ok(())
188    }
189
190    fn check_and_correct_column_type(
191        rw_data_type: &DataType,
192        doris_data_type: String,
193    ) -> Result<bool> {
194        let doris_data_type = doris_data_type.to_ascii_uppercase();
195        let is_variant = doris_data_type.contains("VARIANT");
196        match rw_data_type {
197            risingwave_common::types::DataType::Boolean => Ok(doris_data_type.contains("BOOLEAN")),
198            risingwave_common::types::DataType::Int16 => Ok(doris_data_type.contains("SMALLINT")),
199            risingwave_common::types::DataType::Int32 => Ok(doris_data_type.contains("INT")),
200            risingwave_common::types::DataType::Int64 => Ok(doris_data_type.contains("BIGINT")),
201            risingwave_common::types::DataType::Float32 => Ok(doris_data_type.contains("FLOAT")),
202            risingwave_common::types::DataType::Float64 => Ok(doris_data_type.contains("DOUBLE")),
203            risingwave_common::types::DataType::Decimal => Ok(doris_data_type.contains("DECIMAL")),
204            risingwave_common::types::DataType::Date => Ok(doris_data_type.contains("DATE")),
205            risingwave_common::types::DataType::Varchar => {
206                Ok(
207                    doris_data_type.contains("STRING")
208                        || doris_data_type.contains("VARCHAR")
209                        || is_variant,
210                )
211            }
212            risingwave_common::types::DataType::Time => {
213                Err(SinkError::Doris("TIME is not supported for Doris sink. Please convert to VARCHAR or other supported types.".to_owned()))
214            }
215            risingwave_common::types::DataType::Timestamp => {
216                Ok(doris_data_type.contains("DATETIME"))
217            }
218            risingwave_common::types::DataType::Timestamptz => Err(SinkError::Doris(
219                "TIMESTAMP WITH TIMEZONE is not supported for Doris sink as Doris doesn't store time values with timezone information. Please convert to TIMESTAMP first.".to_owned(),
220            )),
221            risingwave_common::types::DataType::Interval => Err(SinkError::Doris(
222                "INTERVAL is not supported for Doris sink. Please convert to VARCHAR or other supported types.".to_owned(),
223            )),
224            risingwave_common::types::DataType::Struct(_) => Ok(doris_data_type.contains("STRUCT")),
225            risingwave_common::types::DataType::List(_) => Ok(doris_data_type.contains("ARRAY")),
226            risingwave_common::types::DataType::Bytea => {
227                Err(SinkError::Doris("BYTEA is not supported for Doris sink. Please convert to VARCHAR or other supported types.".to_owned()))
228            }
229            risingwave_common::types::DataType::Jsonb => {
230                Ok(doris_data_type.contains("JSON") || is_variant)
231            }
232            risingwave_common::types::DataType::Serial => Ok(doris_data_type.contains("BIGINT")),
233            risingwave_common::types::DataType::Int256 => {
234                Err(SinkError::Doris("INT256 is not supported for Doris sink.".to_owned()))
235            }
236            risingwave_common::types::DataType::Map(_) => {
237                Err(SinkError::Doris("MAP is not supported for Doris sink.".to_owned()))
238            }
239            DataType::Vector(_) => {
240                Err(SinkError::Doris("VECTOR is not supported for Doris sink.".to_owned()))
241            },
242        }
243    }
244}
245
246impl Sink for DorisSink {
247    type LogSinker = LogSinkerOf<DorisSinkWriter>;
248
249    const SINK_NAME: &'static str = DORIS_SINK;
250
251    async fn new_log_sinker(&self, writer_param: SinkWriterParam) -> Result<Self::LogSinker> {
252        Ok(DorisSinkWriter::new(
253            self.config.clone(),
254            self.schema.clone(),
255            self.pk_indices.clone(),
256            self.is_append_only,
257        )
258        .await?
259        .into_log_sinker(SinkWriterMetrics::new(&writer_param)))
260    }
261
262    async fn validate(&self) -> Result<()> {
263        if !self.is_append_only && self.pk_indices.is_empty() {
264            return Err(SinkError::Config(anyhow!(
265                "Primary key not defined for upsert doris sink (please define in `primary_key` field)"
266            )));
267        }
268        // check reachability
269        let client = self.config.common.build_get_client();
270        let doris_schema = client.get_schema_from_doris().await?;
271
272        if !self.is_append_only && doris_schema.keys_type.ne("UNIQUE_KEYS") {
273            return Err(SinkError::Config(anyhow!(
274                "If you want to use upsert, please set the keysType of doris to UNIQUE_KEYS"
275            )));
276        }
277        self.check_column_name_and_type(doris_schema.properties)?;
278        Ok(())
279    }
280}
281
282pub struct DorisSinkWriter {
283    pub config: DorisConfig,
284    #[expect(dead_code)]
285    schema: Schema,
286    #[expect(dead_code)]
287    pk_indices: Vec<usize>,
288    inserter_inner_builder: InserterInnerBuilder,
289    is_append_only: bool,
290    client: Option<DorisClient>,
291    row_encoder: JsonEncoder,
292}
293
294impl TryFrom<SinkParam> for DorisSink {
295    type Error = SinkError;
296
297    fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
298        let schema = param.schema();
299        let pk_indices = param.downstream_pk_or_empty();
300        let config = DorisConfig::from_btreemap(param.properties)?;
301        DorisSink::new(config, schema, pk_indices, param.sink_type.is_append_only())
302    }
303}
304
305impl DorisSinkWriter {
306    pub async fn new(
307        config: DorisConfig,
308        schema: Schema,
309        pk_indices: Vec<usize>,
310        is_append_only: bool,
311    ) -> Result<Self> {
312        let mut decimal_map = HashMap::default();
313        let mut variant_columns = HashSet::default();
314        let doris_schema = config
315            .common
316            .build_get_client()
317            .get_schema_from_doris()
318            .await?;
319        for s in &doris_schema.properties {
320            if let Some(v) = s.get_decimal_pre_scale()? {
321                decimal_map.insert(s.name.clone(), v);
322            }
323            if s.is_variant() {
324                variant_columns.insert(s.name.clone());
325            }
326        }
327
328        let header_builder = HeaderBuilder::new()
329            .add_common_header()
330            .set_user_password(config.common.user.clone(), config.common.password.clone())
331            .add_json_format()
332            .set_partial_columns(config.common.partial_update.clone())
333            .add_read_json_by_line();
334        let header = if !is_append_only {
335            header_builder.add_hidden_column().build()
336        } else {
337            header_builder.build()
338        };
339
340        let doris_insert_builder = InserterInnerBuilder::new(
341            config.common.url.clone(),
342            config.common.database.clone(),
343            config.common.table.clone(),
344            header,
345            config.stream_load_http_timeout_ms,
346        )?;
347        Ok(Self {
348            config,
349            schema: schema.clone(),
350            pk_indices,
351            inserter_inner_builder: doris_insert_builder,
352            is_append_only,
353            client: None,
354            row_encoder: JsonEncoder::new_with_doris(
355                schema,
356                None,
357                DorisJsonConfig {
358                    decimal_scale: decimal_map,
359                    variant_columns,
360                },
361            ),
362        })
363    }
364
365    async fn append_only(&mut self, chunk: StreamChunk) -> Result<()> {
366        for (op, row) in chunk.rows() {
367            if op != Op::Insert {
368                continue;
369            }
370            let row_json_string = Value::Object(self.row_encoder.encode(row)?).to_string();
371            self.client
372                .as_mut()
373                .ok_or_else(|| SinkError::Doris("Can't find doris sink insert".to_owned()))?
374                .write(row_json_string.into())
375                .await?;
376        }
377        Ok(())
378    }
379
380    async fn upsert(&mut self, chunk: StreamChunk) -> Result<()> {
381        for (op, row) in chunk.rows() {
382            match op {
383                Op::Insert => {
384                    let mut row_json_value = self.row_encoder.encode(row)?;
385                    row_json_value
386                        .insert(DORIS_DELETE_SIGN.to_owned(), Value::String("0".to_owned()));
387                    let row_json_string = serde_json::to_string(&row_json_value).map_err(|e| {
388                        SinkError::Doris(format!("Json derialize error: {}", e.as_report()))
389                    })?;
390                    self.client
391                        .as_mut()
392                        .ok_or_else(|| SinkError::Doris("Can't find doris sink insert".to_owned()))?
393                        .write(row_json_string.into())
394                        .await?;
395                }
396                Op::Delete => {
397                    let mut row_json_value = self.row_encoder.encode(row)?;
398                    row_json_value
399                        .insert(DORIS_DELETE_SIGN.to_owned(), Value::String("1".to_owned()));
400                    let row_json_string = serde_json::to_string(&row_json_value).map_err(|e| {
401                        SinkError::Doris(format!("Json derialize error: {}", e.as_report()))
402                    })?;
403                    self.client
404                        .as_mut()
405                        .ok_or_else(|| SinkError::Doris("Can't find doris sink insert".to_owned()))?
406                        .write(row_json_string.into())
407                        .await?;
408                }
409                Op::UpdateDelete => {}
410                Op::UpdateInsert => {
411                    let mut row_json_value = self.row_encoder.encode(row)?;
412                    row_json_value
413                        .insert(DORIS_DELETE_SIGN.to_owned(), Value::String("0".to_owned()));
414                    let row_json_string = serde_json::to_string(&row_json_value).map_err(|e| {
415                        SinkError::Doris(format!("Json derialize error: {}", e.as_report()))
416                    })?;
417                    self.client
418                        .as_mut()
419                        .ok_or_else(|| SinkError::Doris("Can't find doris sink insert".to_owned()))?
420                        .write(row_json_string.into())
421                        .await?;
422                }
423            }
424        }
425        Ok(())
426    }
427}
428
429#[async_trait]
430impl SinkWriter for DorisSinkWriter {
431    async fn write_batch(&mut self, chunk: StreamChunk) -> Result<()> {
432        if self.client.is_none() {
433            self.client = Some(DorisClient::new(self.inserter_inner_builder.build().await?));
434        }
435        if self.is_append_only {
436            self.append_only(chunk).await
437        } else {
438            self.upsert(chunk).await
439        }
440    }
441
442    async fn begin_epoch(&mut self, _epoch: u64) -> Result<()> {
443        Ok(())
444    }
445
446    async fn abort(&mut self) -> Result<()> {
447        Ok(())
448    }
449
450    async fn barrier(&mut self, _is_checkpoint: bool) -> Result<()> {
451        if self.client.is_some() {
452            let client = self
453                .client
454                .take()
455                .ok_or_else(|| SinkError::Doris("Can't find doris inserter".to_owned()))?;
456            client.finish().await?;
457        }
458        Ok(())
459    }
460}
461
462pub struct DorisSchemaClient {
463    url: String,
464    table: String,
465    db: String,
466    user: String,
467    password: String,
468}
469impl DorisSchemaClient {
470    pub fn new(url: String, table: String, db: String, user: String, password: String) -> Self {
471        Self {
472            url,
473            table,
474            db,
475            user,
476            password,
477        }
478    }
479
480    pub async fn get_schema_from_doris(&self) -> Result<DorisSchema> {
481        let uri = format!("{}/api/{}/{}/_schema", self.url, self.db, self.table);
482
483        let client = reqwest::Client::builder()
484            .pool_idle_timeout(POOL_IDLE_TIMEOUT)
485            .build()
486            .map_err(|err| SinkError::DorisStarrocksConnect(err.into()))?;
487
488        let response = client
489            .get(uri)
490            .header(
491                "Authorization",
492                format!(
493                    "Basic {}",
494                    general_purpose::STANDARD.encode(format!("{}:{}", self.user, self.password))
495                ),
496            )
497            .send()
498            .await
499            .map_err(|err| SinkError::DorisStarrocksConnect(err.into()))?;
500
501        let json: Value = response
502            .json()
503            .await
504            .map_err(|err| SinkError::DorisStarrocksConnect(err.into()))?;
505        let json_data = if json.get("code").is_some() && json.get("msg").is_some() {
506            json.get("data")
507                .ok_or_else(|| {
508                    SinkError::DorisStarrocksConnect(anyhow::anyhow!("Can't find data"))
509                })?
510                .clone()
511        } else {
512            json
513        };
514        let schema: DorisSchema = serde_json::from_value(json_data)
515            .context("Can't get schema from json")
516            .map_err(SinkError::DorisStarrocksConnect)?;
517        Ok(schema)
518    }
519}
520#[derive(Debug, Serialize, Deserialize)]
521pub struct DorisSchema {
522    status: i32,
523    #[serde(rename = "keysType")]
524    pub keys_type: String,
525    pub properties: Vec<DorisField>,
526}
527#[derive(Debug, Serialize, Deserialize)]
528pub struct DorisField {
529    pub name: String,
530    pub r#type: String,
531    comment: String,
532    pub precision: Option<String>,
533    pub scale: Option<String>,
534    aggregation_type: String,
535}
536impl DorisField {
537    pub fn get_decimal_pre_scale(&self) -> Result<Option<u8>> {
538        if self.r#type.contains("DECIMAL") {
539            let scale = self
540                .scale
541                .as_ref()
542                .ok_or_else(|| {
543                    SinkError::Doris(format!(
544                        "In doris, the type of {} is DECIMAL, but `scale` is not found",
545                        self.name
546                    ))
547                })?
548                .parse::<u8>()
549                .map_err(|err| {
550                    SinkError::Doris(format!(
551                        "Unable to convert decimal's scale to u8. error: {:?}",
552                        err.kind()
553                    ))
554                })?;
555            Ok(Some(scale))
556        } else {
557            Ok(None)
558        }
559    }
560
561    pub fn is_variant(&self) -> bool {
562        self.r#type.to_ascii_uppercase().contains("VARIANT")
563    }
564}
565
566#[cfg(test)]
567mod tests {
568    use risingwave_common::types::DataType;
569
570    use super::DorisSink;
571
572    #[test]
573    fn test_jsonb_can_write_to_variant() {
574        assert!(DorisSink::check_and_correct_column_type(&DataType::Jsonb, "VARIANT".into())
575            .unwrap());
576    }
577
578    #[test]
579    fn test_varchar_can_write_to_variant() {
580        assert!(DorisSink::check_and_correct_column_type(&DataType::Varchar, "VARIANT".into())
581            .unwrap());
582    }
583}
584
585#[derive(Debug, Serialize, Deserialize)]
586pub struct DorisInsertResultResponse {
587    #[serde(rename = "TxnId")]
588    txn_id: i64,
589    #[serde(rename = "Label")]
590    label: String,
591    #[serde(rename = "Status")]
592    status: String,
593    #[serde(rename = "TwoPhaseCommit")]
594    two_phase_commit: String,
595    #[serde(rename = "Message")]
596    message: String,
597    #[serde(rename = "NumberTotalRows")]
598    number_total_rows: i64,
599    #[serde(rename = "NumberLoadedRows")]
600    number_loaded_rows: i64,
601    #[serde(rename = "NumberFilteredRows")]
602    number_filtered_rows: i32,
603    #[serde(rename = "NumberUnselectedRows")]
604    number_unselected_rows: i32,
605    #[serde(rename = "LoadBytes")]
606    load_bytes: i64,
607    #[serde(rename = "LoadTimeMs")]
608    load_time_ms: i32,
609    #[serde(rename = "BeginTxnTimeMs")]
610    begin_txn_time_ms: i32,
611    #[serde(rename = "StreamLoadPutTimeMs")]
612    stream_load_put_time_ms: i32,
613    #[serde(rename = "ReadDataTimeMs")]
614    read_data_time_ms: i32,
615    #[serde(rename = "WriteDataTimeMs")]
616    write_data_time_ms: i32,
617    #[serde(rename = "CommitAndPublishTimeMs")]
618    commit_and_publish_time_ms: i32,
619    #[serde(rename = "ErrorURL")]
620    err_url: Option<String>,
621}
622
623pub struct DorisClient {
624    insert: InserterInner,
625    is_first_record: bool,
626}
627impl DorisClient {
628    pub fn new(insert: InserterInner) -> Self {
629        Self {
630            insert,
631            is_first_record: true,
632        }
633    }
634
635    pub async fn write(&mut self, data: Bytes) -> Result<()> {
636        let mut data_build = BytesMut::new();
637        if self.is_first_record {
638            self.is_first_record = false;
639        } else {
640            data_build.put_slice("\n".as_bytes());
641        }
642        data_build.put_slice(&data);
643        self.insert.write(data_build.into()).await?;
644        Ok(())
645    }
646
647    pub async fn finish(self) -> Result<DorisInsertResultResponse> {
648        let raw = self.insert.finish().await?;
649        let res: DorisInsertResultResponse = serde_json::from_slice(&raw)
650            .map_err(|err| SinkError::DorisStarrocksConnect(err.into()))?;
651
652        if !DORIS_SUCCESS_STATUS.contains(&res.status.as_str()) {
653            return Err(SinkError::DorisStarrocksConnect(anyhow::anyhow!(
654                "Insert error: {:?}, error url: {:?}",
655                res.message,
656                res.err_url
657            )));
658        };
659        Ok(res)
660    }
661}
risingwave_connector/sink/doris.rs

risingwave_connector/sink/
doris.rs
