risingwave_connector/sink/
starrocks.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};
16use std::num::NonZeroU64;
17use std::sync::Arc;
18
19use anyhow::anyhow;
20use async_trait::async_trait;
21use bytes::Bytes;
22use chrono_tz::Tz;
23use mysql_async::Opts;
24use mysql_async::prelude::Queryable;
25use risingwave_common::array::{Op, StreamChunk};
26use risingwave_common::catalog::Schema;
27use risingwave_common::types::DataType;
28use serde::{Deserialize, Serialize};
29use serde_json::Value;
30use serde_with::{DisplayFromStr, serde_as};
31use thiserror_ext::AsReport;
32use url::form_urlencoded;
33use uuid::Uuid;
34use with_options::WithOptions;
35
36use super::decouple_checkpoint_log_sink::DEFAULT_COMMIT_CHECKPOINT_INTERVAL_WITH_SINK_DECOUPLE;
37use super::doris_starrocks_connector::{
38    HeaderBuilder, InserterInner, STARROCKS_DELETE_SIGN, STARROCKS_SUCCESS_STATUS,
39    StarrocksTxnRequestBuilder,
40};
41use super::encoder::{JsonEncoder, RowEncoder};
42use super::{
43    SINK_TYPE_APPEND_ONLY, SINK_TYPE_OPTION, SINK_TYPE_UPSERT, SinkError, SinkParam,
44    SinkWriterMetrics,
45};
46use crate::enforce_secret::EnforceSecret;
47use crate::sink::decouple_checkpoint_log_sink::DecoupleCheckpointLogSinkerOf;
48use crate::sink::writer::SinkWriter;
49use crate::sink::{Result, Sink, SinkWriterParam};
50
51pub const STARROCKS_SINK: &str = "starrocks";
52const STARROCK_MYSQL_PREFER_SOCKET: &str = "false";
53const STARROCK_MYSQL_MAX_ALLOWED_PACKET: usize = 1024;
54const STARROCK_MYSQL_WAIT_TIMEOUT: usize = 28800;
55pub const fn _default_stream_load_http_timeout_ms() -> u64 {
56    30 * 1000
57}
58
59const fn default_use_https() -> bool {
60    false
61}
62
63#[serde_as]
64#[derive(Deserialize, Debug, Clone, WithOptions)]
65pub struct StarrocksCommon {
66    /// The `StarRocks` host address.
67    #[serde(rename = "starrocks.host")]
68    pub host: String,
69    /// The port to the MySQL server of `StarRocks` FE.
70    #[serde(rename = "starrocks.mysqlport", alias = "starrocks.query_port")]
71    pub mysql_port: String,
72    /// The port to the HTTP server of `StarRocks` FE.
73    #[serde(rename = "starrocks.httpport", alias = "starrocks.http_port")]
74    pub http_port: String,
75    /// The user name used to access the `StarRocks` database.
76    #[serde(rename = "starrocks.user")]
77    pub user: String,
78    /// The password associated with the user.
79    #[serde(rename = "starrocks.password")]
80    pub password: String,
81    /// The `StarRocks` database where the target table is located
82    #[serde(rename = "starrocks.database")]
83    pub database: String,
84    /// The `StarRocks` table you want to sink data to.
85    #[serde(rename = "starrocks.table")]
86    pub table: String,
87
88    /// Whether to use https to connect to the `StarRocks` server.
89    #[serde(rename = "starrocks.use_https")]
90    #[serde(default = "default_use_https")]
91    #[serde_as(as = "DisplayFromStr")]
92    pub use_https: bool,
93}
94
95impl EnforceSecret for StarrocksCommon {
96    const ENFORCE_SECRET_PROPERTIES: phf::Set<&'static str> = phf::phf_set! {
97        "starrocks.password", "starrocks.user"
98    };
99}
100
101#[serde_as]
102#[derive(Clone, Debug, Deserialize, WithOptions)]
103pub struct StarrocksConfig {
104    #[serde(flatten)]
105    pub common: StarrocksCommon,
106
107    /// The timeout in milliseconds for stream load http request, defaults to 10 seconds.
108    #[serde(
109        rename = "starrocks.stream_load.http.timeout.ms",
110        default = "_default_stream_load_http_timeout_ms"
111    )]
112    #[serde_as(as = "DisplayFromStr")]
113    #[with_option(allow_alter_on_fly)]
114    pub stream_load_http_timeout_ms: u64,
115
116    /// Set this option to a positive integer n, RisingWave will try to commit data
117    /// to Starrocks at every n checkpoints by leveraging the
118    /// [StreamLoad Transaction API](https://docs.starrocks.io/docs/loading/Stream_Load_transaction_interface/),
119    /// also, in this time, the `sink_decouple` option should be enabled as well.
120    /// Defaults to 10 if `commit_checkpoint_interval` <= 0
121    #[serde(default = "default_commit_checkpoint_interval")]
122    #[serde_as(as = "DisplayFromStr")]
123    #[with_option(allow_alter_on_fly)]
124    pub commit_checkpoint_interval: u64,
125
126    /// Enable partial update
127    #[serde(rename = "starrocks.partial_update")]
128    pub partial_update: Option<String>,
129
130    /// The maximum size in bytes for each `StarRocks` stream load request payload.
131    /// Defaults to unlimited.
132    #[serde(rename = "starrocks.max_batch_size_bytes")]
133    #[serde_as(as = "Option<DisplayFromStr>")]
134    #[with_option(allow_alter_on_fly)]
135    pub max_batch_size_bytes: Option<u64>,
136
137    pub r#type: String, // accept "append-only" or "upsert"
138
139    #[serde(flatten)]
140    pub unknown_fields: std::collections::HashMap<String, String>,
141}
142
143crate::impl_sink_unknown_fields!(StarrocksConfig);
144
145impl EnforceSecret for StarrocksConfig {
146    fn enforce_one(prop: &str) -> crate::error::ConnectorResult<()> {
147        StarrocksCommon::enforce_one(prop)
148    }
149}
150
151fn default_commit_checkpoint_interval() -> u64 {
152    DEFAULT_COMMIT_CHECKPOINT_INTERVAL_WITH_SINK_DECOUPLE
153}
154
155impl StarrocksConfig {
156    pub fn from_btreemap(properties: BTreeMap<String, String>) -> Result<Self> {
157        let config =
158            serde_json::from_value::<StarrocksConfig>(serde_json::to_value(properties).unwrap())
159                .map_err(|e| SinkError::Config(anyhow!(e)))?;
160        if config.r#type != SINK_TYPE_APPEND_ONLY && config.r#type != SINK_TYPE_UPSERT {
161            return Err(SinkError::Config(anyhow!(
162                "`{}` must be {}, or {}",
163                SINK_TYPE_OPTION,
164                SINK_TYPE_APPEND_ONLY,
165                SINK_TYPE_UPSERT
166            )));
167        }
168        if config.commit_checkpoint_interval == 0 {
169            return Err(SinkError::Config(anyhow!(
170                "`commit_checkpoint_interval` must be greater than 0"
171            )));
172        }
173        if let Some(0) = config.max_batch_size_bytes {
174            return Err(SinkError::Config(anyhow!(
175                "`starrocks.max_batch_size_bytes` must be greater than 0"
176            )));
177        }
178        Ok(config)
179    }
180}
181
182#[derive(Debug, PartialEq, Eq)]
183struct LoadRequestSizeDecision {
184    finish_current_load: bool,
185    next_batch_size_bytes: u64,
186}
187
188fn decide_load_request_size(
189    current_batch_size_bytes: u64,
190    row_size: u64,
191    max_batch_size_bytes: u64,
192) -> Result<LoadRequestSizeDecision> {
193    if row_size > max_batch_size_bytes {
194        return Err(SinkError::Starrocks(format!(
195            "single row payload size {} bytes exceeds `starrocks.max_batch_size_bytes` limit {} bytes",
196            row_size, max_batch_size_bytes
197        )));
198    }
199
200    if current_batch_size_bytes > 0
201        && current_batch_size_bytes
202            .checked_add(row_size)
203            .is_none_or(|next_batch_size_bytes| next_batch_size_bytes > max_batch_size_bytes)
204    {
205        return Ok(LoadRequestSizeDecision {
206            finish_current_load: true,
207            next_batch_size_bytes: row_size,
208        });
209    }
210
211    Ok(LoadRequestSizeDecision {
212        finish_current_load: false,
213        next_batch_size_bytes: current_batch_size_bytes
214            .checked_add(row_size)
215            .expect("sum is checked against max_batch_size_bytes above"),
216    })
217}
218
219#[derive(Debug)]
220pub struct StarrocksSink {
221    pub config: StarrocksConfig,
222    schema: Schema,
223    pk_indices: Vec<usize>,
224    is_append_only: bool,
225}
226
227impl EnforceSecret for StarrocksSink {
228    fn enforce_secret<'a>(
229        prop_iter: impl Iterator<Item = &'a str>,
230    ) -> crate::error::ConnectorResult<()> {
231        for prop in prop_iter {
232            StarrocksConfig::enforce_one(prop)?;
233        }
234        Ok(())
235    }
236}
237
238impl StarrocksSink {
239    pub fn new(param: SinkParam, config: StarrocksConfig, schema: Schema) -> Result<Self> {
240        let pk_indices = param.downstream_pk_or_empty();
241        let is_append_only = param.sink_type.is_append_only();
242        Ok(Self {
243            config,
244            schema,
245            pk_indices,
246            is_append_only,
247        })
248    }
249}
250
251impl StarrocksSink {
252    fn starrocks_data_type_contains_any(
253        starrocks_data_type: &str,
254        expected_types: &[&str],
255    ) -> bool {
256        expected_types
257            .iter()
258            .any(|expected_type| starrocks_data_type.contains(expected_type))
259    }
260
261    fn check_column_name_and_type(
262        &self,
263        starrocks_columns_desc: HashMap<String, String>,
264    ) -> Result<()> {
265        let rw_fields_name = self.schema.fields();
266        if rw_fields_name.len() > starrocks_columns_desc.len() {
267            return Err(SinkError::Starrocks("The columns of the sink must be equal to or a superset of the target table's columns.".to_owned()));
268        }
269
270        for i in rw_fields_name {
271            let value = starrocks_columns_desc.get(&i.name).ok_or_else(|| {
272                SinkError::Starrocks(format!(
273                    "Column name don't find in starrocks, risingwave is {:?} ",
274                    i.name
275                ))
276            })?;
277            if !Self::check_and_correct_column_type(&i.data_type, value)? {
278                return Err(SinkError::Starrocks(format!(
279                    "Column type don't match, column name is {:?}. starrocks type is {:?} risingwave type is {:?} ",
280                    i.name, value, i.data_type
281                )));
282            }
283        }
284        Ok(())
285    }
286
287    fn check_and_correct_column_type(
288        rw_data_type: &DataType,
289        starrocks_data_type: &str,
290    ) -> Result<bool> {
291        match rw_data_type {
292            risingwave_common::types::DataType::Boolean => {
293                Ok(Self::starrocks_data_type_contains_any(
294                    starrocks_data_type,
295                    &["tinyint", "boolean"],
296                ))
297            }
298            risingwave_common::types::DataType::Int16 => {
299                Ok(starrocks_data_type.contains("smallint"))
300            }
301            risingwave_common::types::DataType::Int32 => Ok(starrocks_data_type.contains("int")),
302            risingwave_common::types::DataType::Int64 => Ok(starrocks_data_type.contains("bigint")),
303            risingwave_common::types::DataType::Float32 => {
304                Ok(starrocks_data_type.contains("float"))
305            }
306            risingwave_common::types::DataType::Float64 => {
307                Ok(starrocks_data_type.contains("double"))
308            }
309            risingwave_common::types::DataType::Decimal => {
310                Ok(starrocks_data_type.contains("decimal"))
311            }
312            risingwave_common::types::DataType::Date => Ok(starrocks_data_type.contains("date")),
313            risingwave_common::types::DataType::Varchar => {
314                Ok(starrocks_data_type.contains("varchar"))
315            }
316            risingwave_common::types::DataType::Time => Err(SinkError::Starrocks(
317                "TIME is not supported for Starrocks sink. Please convert to VARCHAR or other supported types.".to_owned(),
318            )),
319            risingwave_common::types::DataType::Timestamp => {
320                Ok(starrocks_data_type.contains("datetime"))
321            }
322            risingwave_common::types::DataType::Timestamptz => {
323                // StarRocks JSON encoding writes timestamptz as a timestamp string without a
324                // timezone suffix, so StarRocks can parse it into DATETIME or store it as text.
325                Ok(Self::starrocks_data_type_contains_any(
326                    starrocks_data_type,
327                    &["datetime", "varchar", "char", "string"],
328                ))
329            }
330            risingwave_common::types::DataType::Interval => Err(SinkError::Starrocks(
331                "INTERVAL is not supported for Starrocks sink. Please convert to VARCHAR or other supported types.".to_owned(),
332            )),
333            risingwave_common::types::DataType::Struct(_) => Err(SinkError::Starrocks(
334                "STRUCT is not supported for Starrocks sink.".to_owned(),
335            )),
336            risingwave_common::types::DataType::List(list) => {
337                // For compatibility with older versions starrocks
338                if starrocks_data_type.contains("unknown") {
339                    return Ok(true);
340                }
341                let check_result = Self::check_and_correct_column_type(list.elem(), starrocks_data_type)?;
342                Ok(check_result && starrocks_data_type.contains("array"))
343            }
344            risingwave_common::types::DataType::Bytea => Err(SinkError::Starrocks(
345                "BYTEA is not supported for Starrocks sink. Please convert to VARCHAR or other supported types.".to_owned(),
346            )),
347            risingwave_common::types::DataType::Jsonb => Ok(starrocks_data_type.contains("json")),
348            risingwave_common::types::DataType::Variant => Err(SinkError::Starrocks(
349                "VARIANT is not supported for Starrocks sink.".to_owned(),
350            )),
351            risingwave_common::types::DataType::Serial => {
352                Ok(starrocks_data_type.contains("bigint"))
353            }
354            risingwave_common::types::DataType::Int256 => Err(SinkError::Starrocks(
355                "INT256 is not supported for Starrocks sink.".to_owned(),
356            )),
357            risingwave_common::types::DataType::Map(_) => Err(SinkError::Starrocks(
358                "MAP is not supported for Starrocks sink.".to_owned(),
359            )),
360            DataType::Vector(_) => Err(SinkError::Starrocks(
361                "VECTOR is not supported for Starrocks sink.".to_owned(),
362            )),
363        }
364    }
365}
366
367impl Sink for StarrocksSink {
368    type LogSinker = DecoupleCheckpointLogSinkerOf<StarrocksSinkWriter>;
369
370    const SINK_NAME: &'static str = STARROCKS_SINK;
371
372    crate::impl_validate_sink_unknown_fields!();
373
374    async fn validate(&self) -> Result<()> {
375        if !self.is_append_only && self.pk_indices.is_empty() {
376            return Err(SinkError::Config(anyhow!(
377                "Primary key not defined for upsert starrocks sink (please define in `primary_key` field)"
378            )));
379        }
380        // check reachability
381        let mut client = StarrocksSchemaClient::new(
382            self.config.common.host.clone(),
383            self.config.common.mysql_port.clone(),
384            self.config.common.table.clone(),
385            self.config.common.database.clone(),
386            self.config.common.user.clone(),
387            self.config.common.password.clone(),
388        )
389        .await?;
390        let (read_model, pks) = client.get_pk_from_starrocks().await?;
391
392        if !self.is_append_only && read_model.ne("PRIMARY_KEYS") {
393            return Err(SinkError::Config(anyhow!(
394                "If you want to use upsert, please set the keysType of starrocks to PRIMARY_KEY"
395            )));
396        }
397
398        for (index, filed) in self.schema.fields().iter().enumerate() {
399            if self.pk_indices.contains(&index) && !pks.contains(&filed.name) {
400                return Err(SinkError::Starrocks(format!(
401                    "Can't find pk {:?} in starrocks",
402                    filed.name
403                )));
404            }
405        }
406
407        let starrocks_columns_desc = client.get_columns_from_starrocks().await?;
408
409        self.check_column_name_and_type(starrocks_columns_desc)?;
410        Ok(())
411    }
412
413    fn validate_alter_config(config: &BTreeMap<String, String>) -> Result<()> {
414        StarrocksConfig::from_btreemap(config.clone())?;
415        Ok(())
416    }
417
418    async fn new_log_sinker(&self, writer_param: SinkWriterParam) -> Result<Self::LogSinker> {
419        let commit_checkpoint_interval =
420            NonZeroU64::new(self.config.commit_checkpoint_interval).expect(
421                "commit_checkpoint_interval should be greater than 0, and it should be checked in config validation",
422            );
423
424        let writer = StarrocksSinkWriter::new(
425            self.config.clone(),
426            self.schema.clone(),
427            self.pk_indices.clone(),
428            self.is_append_only,
429            writer_param.time_zone,
430        )?;
431
432        let metrics = SinkWriterMetrics::new(&writer_param);
433
434        Ok(DecoupleCheckpointLogSinkerOf::new(
435            writer,
436            metrics,
437            commit_checkpoint_interval,
438        ))
439    }
440}
441
442pub struct StarrocksSinkWriter {
443    #[expect(dead_code)]
444    schema: Schema,
445    #[expect(dead_code)]
446    pk_indices: Vec<usize>,
447    is_append_only: bool,
448    client: Option<StarrocksClient>,
449    txn_client: Arc<StarrocksTxnClient>,
450    row_encoder: JsonEncoder,
451    curr_txn_label: Option<String>,
452    max_batch_size_bytes: Option<u64>,
453    current_batch_size_bytes: u64,
454}
455
456impl TryFrom<SinkParam> for StarrocksSink {
457    type Error = SinkError;
458
459    fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
460        let schema = param.schema();
461        let config = StarrocksConfig::from_btreemap(param.properties.clone())?;
462        StarrocksSink::new(param, config, schema)
463    }
464}
465
466impl StarrocksSinkWriter {
467    pub fn new(
468        config: StarrocksConfig,
469        schema: Schema,
470        pk_indices: Vec<usize>,
471        is_append_only: bool,
472        time_zone: Tz,
473    ) -> Result<Self> {
474        let mut field_names = schema.names_str();
475        if !is_append_only {
476            field_names.push(STARROCKS_DELETE_SIGN);
477        };
478        // we should quote field names in `MySQL` style to prevent `StarRocks` from rejecting the request due to
479        // a field name being a reserved word. For example, `order`, 'from`, etc.
480        let field_names = field_names
481            .into_iter()
482            .map(|name| format!("`{}`", name))
483            .collect::<Vec<String>>();
484        let field_names_str = field_names
485            .iter()
486            .map(|name| name.as_str())
487            .collect::<Vec<&str>>();
488
489        let header = HeaderBuilder::new()
490            .add_common_header()
491            .set_user_password(config.common.user.clone(), config.common.password.clone())
492            .add_json_format()
493            .set_partial_update(config.partial_update.clone())
494            .set_columns_name(field_names_str)
495            .set_db(config.common.database.clone())
496            .set_table(config.common.table.clone())
497            .build();
498
499        let url = if config.common.use_https {
500            format!("https://{}:{}", config.common.host, config.common.http_port)
501        } else {
502            format!("http://{}:{}", config.common.host, config.common.http_port)
503        };
504        let txn_request_builder =
505            StarrocksTxnRequestBuilder::new(url, header, config.stream_load_http_timeout_ms)?;
506
507        Ok(Self {
508            schema: schema.clone(),
509            pk_indices,
510            is_append_only,
511            client: None,
512            txn_client: Arc::new(StarrocksTxnClient::new(txn_request_builder)),
513            row_encoder: JsonEncoder::new_with_starrocks(schema, None, time_zone),
514            curr_txn_label: None,
515            max_batch_size_bytes: config.max_batch_size_bytes,
516            current_batch_size_bytes: 0,
517        })
518    }
519
520    async fn finish_load_request(&mut self) -> Result<()> {
521        if let Some(client) = self.client.take() {
522            client.finish().await?;
523            self.current_batch_size_bytes = 0;
524        }
525        Ok(())
526    }
527
528    async fn ensure_load_request(&mut self) -> Result<()> {
529        if self.client.is_none() {
530            let txn_label = self.curr_txn_label.clone().ok_or_else(|| {
531                SinkError::Starrocks("Can't find current starrocks transaction label".to_owned())
532            })?;
533            self.client = Some(StarrocksClient::new(self.txn_client.load(txn_label).await?));
534            self.current_batch_size_bytes = 0;
535        }
536        Ok(())
537    }
538
539    async fn write_row_json(&mut self, row_json_string: String) -> Result<()> {
540        let row_size = row_json_string.len() as u64;
541        let size_decision = self
542            .max_batch_size_bytes
543            .map(|max_batch_size_bytes| {
544                decide_load_request_size(
545                    self.current_batch_size_bytes,
546                    row_size,
547                    max_batch_size_bytes,
548                )
549            })
550            .transpose()?;
551        if size_decision
552            .as_ref()
553            .is_some_and(|decision| decision.finish_current_load)
554        {
555            self.finish_load_request().await?;
556        }
557        self.ensure_load_request().await?;
558        self.client
559            .as_mut()
560            .ok_or_else(|| SinkError::Starrocks("Can't find starrocks sink insert".to_owned()))?
561            .write(row_json_string.into())
562            .await?;
563        if let Some(size_decision) = size_decision {
564            self.current_batch_size_bytes = size_decision.next_batch_size_bytes;
565        }
566        Ok(())
567    }
568
569    async fn append_only(&mut self, chunk: StreamChunk) -> Result<()> {
570        for (op, row) in chunk.rows() {
571            if op != Op::Insert {
572                continue;
573            }
574            let row_json_string = Value::Object(self.row_encoder.encode(row)?).to_string();
575            self.write_row_json(row_json_string).await?;
576        }
577        Ok(())
578    }
579
580    async fn upsert(&mut self, chunk: StreamChunk) -> Result<()> {
581        for (op, row) in chunk.rows() {
582            match op {
583                Op::Insert => {
584                    let mut row_json_value = self.row_encoder.encode(row)?;
585                    row_json_value.insert(
586                        STARROCKS_DELETE_SIGN.to_owned(),
587                        Value::String("0".to_owned()),
588                    );
589                    let row_json_string = serde_json::to_string(&row_json_value).map_err(|e| {
590                        SinkError::Starrocks(format!("Json serialize error: {}", e.as_report()))
591                    })?;
592                    self.write_row_json(row_json_string).await?;
593                }
594                Op::Delete => {
595                    let mut row_json_value = self.row_encoder.encode(row)?;
596                    row_json_value.insert(
597                        STARROCKS_DELETE_SIGN.to_owned(),
598                        Value::String("1".to_owned()),
599                    );
600                    let row_json_string = serde_json::to_string(&row_json_value).map_err(|e| {
601                        SinkError::Starrocks(format!("Json serialize error: {}", e.as_report()))
602                    })?;
603                    self.write_row_json(row_json_string).await?;
604                }
605                Op::UpdateDelete => {}
606                Op::UpdateInsert => {
607                    let mut row_json_value = self.row_encoder.encode(row)?;
608                    row_json_value.insert(
609                        STARROCKS_DELETE_SIGN.to_owned(),
610                        Value::String("0".to_owned()),
611                    );
612                    let row_json_string = serde_json::to_string(&row_json_value).map_err(|e| {
613                        SinkError::Starrocks(format!("Json serialize error: {}", e.as_report()))
614                    })?;
615                    self.write_row_json(row_json_string).await?;
616                }
617            }
618        }
619        Ok(())
620    }
621
622    /// Generating a new transaction label, should be unique across all `SinkWriters` even under rewinding.
623    #[inline(always)]
624    fn new_txn_label(&self) -> String {
625        format!(
626            "rw-txn-{}-{}",
627            Uuid::new_v4(),
628            chrono::Utc::now().timestamp_micros()
629        )
630    }
631
632    async fn prepare_and_commit(&self, txn_label: String) -> Result<()> {
633        tracing::debug!(?txn_label, "prepare transaction");
634        let txn_label_res = self.txn_client.prepare(txn_label.clone()).await?;
635        if txn_label != txn_label_res {
636            return Err(SinkError::Starrocks(format!(
637                "label {} returned from prepare transaction {} differs from the current one",
638                txn_label, txn_label_res
639            )));
640        }
641        tracing::debug!(?txn_label, "commit transaction");
642        let txn_label_res = self.txn_client.commit(txn_label.clone()).await?;
643        if txn_label != txn_label_res {
644            return Err(SinkError::Starrocks(format!(
645                "label {} returned from commit transaction {} differs from the current one",
646                txn_label, txn_label_res
647            )));
648        }
649        Ok(())
650    }
651}
652
653impl Drop for StarrocksSinkWriter {
654    fn drop(&mut self) {
655        if let Some(txn_label) = self.curr_txn_label.take() {
656            let txn_client = self.txn_client.clone();
657            tokio::spawn(async move {
658                if let Err(e) = txn_client.rollback(txn_label.clone()).await {
659                    tracing::error!(
660                        "starrocks rollback transaction error: {:?}, txn label: {}",
661                        e.as_report(),
662                        txn_label
663                    );
664                }
665            });
666        }
667    }
668}
669
670#[async_trait]
671impl SinkWriter for StarrocksSinkWriter {
672    async fn begin_epoch(&mut self, _epoch: u64) -> Result<()> {
673        Ok(())
674    }
675
676    async fn write_batch(&mut self, chunk: StreamChunk) -> Result<()> {
677        // We check whether start a new transaction in `write_batch`. Therefore, if no data has been written
678        // within the `commit_checkpoint_interval` period, no meta requests will be made. Otherwise if we request
679        // `prepare` against an empty transaction, the `StarRocks` will report a `hasn't send any data yet` error.
680        if self.curr_txn_label.is_none() {
681            let txn_label = self.new_txn_label();
682            tracing::debug!(?txn_label, "begin transaction");
683            let txn_label_res = self.txn_client.begin(txn_label.clone()).await?;
684            if txn_label != txn_label_res {
685                return Err(SinkError::Starrocks(format!(
686                    "label {} returned from StarRocks {} differs from generated one",
687                    txn_label, txn_label_res
688                )));
689            }
690            self.curr_txn_label = Some(txn_label);
691        }
692        if self.is_append_only {
693            self.append_only(chunk).await
694        } else {
695            self.upsert(chunk).await
696        }
697    }
698
699    async fn barrier(&mut self, is_checkpoint: bool) -> Result<()> {
700        // Here we finish the `/api/transaction/load` request when a barrier is received. Therefore,
701        // one or more load requests should be made within one commit_checkpoint_interval period.
702        // StarRocks will take care of merging those splits into a larger one during prepare transaction.
703        // Thus, only one version will be produced when the transaction is committed. See Stream Load
704        // transaction interface for more information.
705        self.finish_load_request().await?;
706
707        if is_checkpoint
708            && let Some(txn_label) = self.curr_txn_label.take()
709            && let Err(err) = self.prepare_and_commit(txn_label.clone()).await
710        {
711            match self.txn_client.rollback(txn_label.clone()).await {
712                Ok(_) => tracing::warn!(
713                    ?txn_label,
714                    "transaction is successfully rolled back due to commit failure"
715                ),
716                Err(err) => {
717                    tracing::warn!(?txn_label, error = ?err.as_report(), "Couldn't roll back transaction after commit failed")
718                }
719            }
720
721            return Err(err);
722        }
723        Ok(())
724    }
725
726    async fn abort(&mut self) -> Result<()> {
727        if let Some(txn_label) = self.curr_txn_label.take() {
728            tracing::debug!(?txn_label, "rollback transaction");
729            self.txn_client.rollback(txn_label).await?;
730        }
731        Ok(())
732    }
733}
734
735pub struct StarrocksSchemaClient {
736    table: String,
737    db: String,
738    conn: mysql_async::Conn,
739}
740
741impl StarrocksSchemaClient {
742    pub async fn new(
743        host: String,
744        port: String,
745        table: String,
746        db: String,
747        user: String,
748        password: String,
749    ) -> Result<Self> {
750        // username & password may contain special chars, so we need to do URL encoding on them.
751        // Otherwise, Opts::from_url may report a `Parse error`
752        let user = form_urlencoded::byte_serialize(user.as_bytes()).collect::<String>();
753        let password = form_urlencoded::byte_serialize(password.as_bytes()).collect::<String>();
754
755        let conn_uri = format!(
756            "mysql://{}:{}@{}:{}/{}?prefer_socket={}&max_allowed_packet={}&wait_timeout={}",
757            user,
758            password,
759            host,
760            port,
761            db,
762            STARROCK_MYSQL_PREFER_SOCKET,
763            STARROCK_MYSQL_MAX_ALLOWED_PACKET,
764            STARROCK_MYSQL_WAIT_TIMEOUT
765        );
766        let pool = mysql_async::Pool::new(
767            Opts::from_url(&conn_uri)
768                .map_err(|err| SinkError::DorisStarrocksConnect(anyhow!(err)))?,
769        );
770        let conn = pool
771            .get_conn()
772            .await
773            .map_err(|err| SinkError::DorisStarrocksConnect(anyhow!(err)))?;
774
775        Ok(Self { table, db, conn })
776    }
777
778    pub async fn get_columns_from_starrocks(&mut self) -> Result<HashMap<String, String>> {
779        let query = format!(
780            "select column_name, column_type from information_schema.columns where table_name = {:?} and table_schema = {:?};",
781            self.table, self.db
782        );
783        let mut query_map: HashMap<String, String> = HashMap::default();
784        self.conn
785            .query_map(query, |(column_name, column_type)| {
786                query_map.insert(column_name, column_type)
787            })
788            .await
789            .map_err(|err| SinkError::DorisStarrocksConnect(anyhow!(err)))?;
790        Ok(query_map)
791    }
792
793    pub async fn get_pk_from_starrocks(&mut self) -> Result<(String, String)> {
794        let query = format!(
795            "select table_model, primary_key, sort_key from information_schema.tables_config where table_name = {:?} and table_schema = {:?};",
796            self.table, self.db
797        );
798        let table_mode_pk: (String, String) = self
799            .conn
800            .query_map(
801                query,
802                |(table_model, primary_key, sort_key): (String, String, String)| match table_model
803                    .as_str()
804                {
805                    // Get primary key of aggregate table from the sort_key field
806                    // https://docs.starrocks.io/docs/table_design/table_types/table_capabilities/
807                    // https://docs.starrocks.io/docs/sql-reference/information_schema/tables_config/
808                    "AGG_KEYS" => (table_model, sort_key),
809                    _ => (table_model, primary_key),
810                },
811            )
812            .await
813            .map_err(|err| SinkError::DorisStarrocksConnect(anyhow!(err)))?
814            .first()
815            .ok_or_else(|| {
816                SinkError::Starrocks(format!(
817                    "Can't find schema for StarRocks table {:?} in database {:?}. Please check that the table exists and the StarRocks user has write permission on the target table.",
818                    self.table, self.db
819                ))
820            })?
821            .clone();
822        Ok(table_mode_pk)
823    }
824}
825
826#[derive(Debug, Serialize, Deserialize)]
827pub struct StarrocksInsertResultResponse {
828    #[serde(rename = "TxnId")]
829    pub txn_id: Option<i64>,
830    #[serde(rename = "Seq")]
831    pub seq: Option<i64>,
832    #[serde(rename = "Label")]
833    pub label: Option<String>,
834    #[serde(rename = "Status")]
835    pub status: String,
836    #[serde(rename = "Message")]
837    pub message: String,
838    #[serde(rename = "NumberTotalRows")]
839    pub number_total_rows: Option<i64>,
840    #[serde(rename = "NumberLoadedRows")]
841    pub number_loaded_rows: Option<i64>,
842    #[serde(rename = "NumberFilteredRows")]
843    pub number_filtered_rows: Option<i32>,
844    #[serde(rename = "NumberUnselectedRows")]
845    pub number_unselected_rows: Option<i32>,
846    #[serde(rename = "LoadBytes")]
847    pub load_bytes: Option<i64>,
848    #[serde(rename = "LoadTimeMs")]
849    pub load_time_ms: Option<i32>,
850    #[serde(rename = "BeginTxnTimeMs")]
851    pub begin_txn_time_ms: Option<i32>,
852    #[serde(rename = "ReadDataTimeMs")]
853    pub read_data_time_ms: Option<i32>,
854    #[serde(rename = "WriteDataTimeMs")]
855    pub write_data_time_ms: Option<i32>,
856    #[serde(rename = "CommitAndPublishTimeMs")]
857    pub commit_and_publish_time_ms: Option<i32>,
858    #[serde(rename = "StreamLoadPlanTimeMs")]
859    pub stream_load_plan_time_ms: Option<i32>,
860    #[serde(rename = "ExistingJobStatus")]
861    pub existing_job_status: Option<String>,
862    #[serde(rename = "ErrorURL")]
863    pub error_url: Option<String>,
864}
865
866pub struct StarrocksClient {
867    insert: InserterInner,
868}
869impl StarrocksClient {
870    pub fn new(insert: InserterInner) -> Self {
871        Self { insert }
872    }
873
874    pub async fn write(&mut self, data: Bytes) -> Result<()> {
875        self.insert.write(data).await?;
876        Ok(())
877    }
878
879    pub async fn finish(self) -> Result<StarrocksInsertResultResponse> {
880        let raw = self.insert.finish().await?;
881        let res: StarrocksInsertResultResponse = serde_json::from_slice(&raw)
882            .map_err(|err| SinkError::DorisStarrocksConnect(anyhow!(err)))?;
883
884        if !STARROCKS_SUCCESS_STATUS.contains(&res.status.as_str()) {
885            return Err(SinkError::DorisStarrocksConnect(anyhow::anyhow!(
886                "Insert error: {}, {}, {:?}",
887                res.status,
888                res.message,
889                res.error_url,
890            )));
891        };
892        Ok(res)
893    }
894}
895
896pub struct StarrocksTxnClient {
897    request_builder: StarrocksTxnRequestBuilder,
898}
899
900impl StarrocksTxnClient {
901    pub fn new(request_builder: StarrocksTxnRequestBuilder) -> Self {
902        Self { request_builder }
903    }
904
905    fn check_response_and_extract_label(&self, res: Bytes) -> Result<String> {
906        let res: StarrocksInsertResultResponse = serde_json::from_slice(&res)
907            .map_err(|err| SinkError::DorisStarrocksConnect(anyhow!(err)))?;
908        if !STARROCKS_SUCCESS_STATUS.contains(&res.status.as_str()) {
909            return Err(SinkError::DorisStarrocksConnect(anyhow::anyhow!(
910                "transaction error: {}, {}, {:?}",
911                res.status,
912                res.message,
913                res.error_url,
914            )));
915        }
916        res.label.ok_or_else(|| {
917            SinkError::DorisStarrocksConnect(anyhow::anyhow!("Can't get label from response"))
918        })
919    }
920
921    pub async fn begin(&self, label: String) -> Result<String> {
922        let res = self
923            .request_builder
924            .build_begin_request_sender(label)?
925            .send()
926            .await?;
927        self.check_response_and_extract_label(res)
928    }
929
930    pub async fn prepare(&self, label: String) -> Result<String> {
931        let res = self
932            .request_builder
933            .build_prepare_request_sender(label)?
934            .send()
935            .await?;
936        self.check_response_and_extract_label(res)
937    }
938
939    pub async fn commit(&self, label: String) -> Result<String> {
940        let res = self
941            .request_builder
942            .build_commit_request_sender(label)?
943            .send()
944            .await?;
945        self.check_response_and_extract_label(res)
946    }
947
948    pub async fn rollback(&self, label: String) -> Result<String> {
949        let res = self
950            .request_builder
951            .build_rollback_request_sender(label)?
952            .send()
953            .await?;
954        self.check_response_and_extract_label(res)
955    }
956
957    pub async fn load(&self, label: String) -> Result<InserterInner> {
958        self.request_builder.build_txn_inserter(label).await
959    }
960}
961
962#[cfg(test)]
963mod tests {
964    use risingwave_common::types::DataType;
965
966    use super::*;
967
968    fn is_compatible(rw_data_type: DataType, starrocks_data_type: &str) -> bool {
969        StarrocksSink::check_and_correct_column_type(&rw_data_type, starrocks_data_type).unwrap()
970    }
971
972    #[test]
973    fn test_timestamptz_compatible_starrocks_types() {
974        for starrocks_data_type in [
975            "datetime",
976            "datetime(6)",
977            "varchar(64)",
978            "char(32)",
979            "string",
980        ] {
981            assert!(
982                is_compatible(DataType::Timestamptz, starrocks_data_type),
983                "{starrocks_data_type} should be compatible with timestamptz"
984            );
985        }
986    }
987
988    #[test]
989    fn test_timestamptz_incompatible_starrocks_types() {
990        for starrocks_data_type in ["date", "int", "bigint", "json", "boolean"] {
991            assert!(
992                !is_compatible(DataType::Timestamptz, starrocks_data_type),
993                "{starrocks_data_type} should not be compatible with timestamptz"
994            );
995        }
996    }
997
998    fn base_properties() -> BTreeMap<String, String> {
999        BTreeMap::from([
1000            ("starrocks.host".to_owned(), "127.0.0.1".to_owned()),
1001            ("starrocks.mysqlport".to_owned(), "9030".to_owned()),
1002            ("starrocks.httpport".to_owned(), "8030".to_owned()),
1003            ("starrocks.user".to_owned(), "root".to_owned()),
1004            ("starrocks.password".to_owned(), "".to_owned()),
1005            ("starrocks.database".to_owned(), "demo".to_owned()),
1006            ("starrocks.table".to_owned(), "sink_table".to_owned()),
1007            ("type".to_owned(), SINK_TYPE_APPEND_ONLY.to_owned()),
1008        ])
1009    }
1010
1011    #[test]
1012    fn starrocks_max_batch_size_bytes_defaults_to_none() {
1013        let config = StarrocksConfig::from_btreemap(base_properties()).unwrap();
1014
1015        assert_eq!(config.max_batch_size_bytes, None);
1016    }
1017
1018    #[test]
1019    fn starrocks_max_batch_size_bytes_parses() {
1020        let mut properties = base_properties();
1021        properties.insert("starrocks.max_batch_size_bytes".to_owned(), "2".to_owned());
1022
1023        let config = StarrocksConfig::from_btreemap(properties).unwrap();
1024
1025        assert_eq!(config.max_batch_size_bytes, Some(2));
1026    }
1027
1028    #[test]
1029    fn starrocks_max_batch_size_bytes_rejects_zero() {
1030        let mut properties = base_properties();
1031        properties.insert("starrocks.max_batch_size_bytes".to_owned(), "0".to_owned());
1032
1033        let err = StarrocksConfig::from_btreemap(properties).unwrap_err();
1034
1035        assert!(
1036            err.to_string()
1037                .contains("`starrocks.max_batch_size_bytes` must be greater than 0")
1038        );
1039    }
1040
1041    #[test]
1042    fn starrocks_batch_size_allows_exact_limit() {
1043        assert_eq!(
1044            decide_load_request_size(3, 2, 5).unwrap(),
1045            LoadRequestSizeDecision {
1046                finish_current_load: false,
1047                next_batch_size_bytes: 5,
1048            }
1049        );
1050    }
1051
1052    #[test]
1053    fn starrocks_batch_size_rolls_over_before_exceeding_limit() {
1054        assert_eq!(
1055            decide_load_request_size(4, 2, 5).unwrap(),
1056            LoadRequestSizeDecision {
1057                finish_current_load: true,
1058                next_batch_size_bytes: 2,
1059            }
1060        );
1061    }
1062
1063    #[test]
1064    fn starrocks_batch_size_rejects_single_oversized_row() {
1065        let err = decide_load_request_size(0, 6, 5).unwrap_err();
1066
1067        assert!(err.to_string().contains(
1068            "single row payload size 6 bytes exceeds `starrocks.max_batch_size_bytes` limit 5 bytes"
1069        ));
1070    }
1071
1072    #[test]
1073    fn starrocks_batch_size_rolls_over_on_u64_overflow() {
1074        assert_eq!(
1075            decide_load_request_size(u64::MAX, 1, u64::MAX).unwrap(),
1076            LoadRequestSizeDecision {
1077                finish_current_load: true,
1078                next_batch_size_bytes: 1,
1079            }
1080        );
1081    }
1082}
risingwave_connector/sink/starrocks.rs

risingwave_connector/sink/
starrocks.rs
