risingwave_connector/sink/
redis.rs

1// Copyright 2022 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};
16
17use anyhow::anyhow;
18use async_trait::async_trait;
19use phf::phf_set;
20use redis::aio::MultiplexedConnection;
21use redis::cluster::{ClusterClient, ClusterConnection, ClusterPipeline};
22use redis::{Client as RedisClient, Pipeline};
23use risingwave_common::array::StreamChunk;
24use risingwave_common::catalog::Schema;
25use risingwave_common::types::DataType;
26use serde::Deserialize;
27use serde_json::Value;
28use serde_with::serde_as;
29use with_options::WithOptions;
30
31use super::catalog::SinkFormatDesc;
32use super::encoder::template::{RedisSinkPayloadWriterInput, TemplateStringEncoder};
33use super::formatter::SinkFormatterImpl;
34use super::writer::FormattedSink;
35use super::{SinkError, SinkParam};
36use crate::dispatch_sink_formatter_str_key_impl;
37use crate::enforce_secret::EnforceSecret;
38use crate::error::ConnectorResult;
39use crate::sink::log_store::DeliveryFutureManagerAddFuture;
40use crate::sink::writer::{
41    AsyncTruncateLogSinkerOf, AsyncTruncateSinkWriter, AsyncTruncateSinkWriterExt,
42};
43use crate::sink::{Result, Sink, SinkWriterParam};
44
45pub const REDIS_SINK: &str = "redis";
46pub const KEY_FORMAT: &str = "key_format";
47pub const VALUE_FORMAT: &str = "value_format";
48pub const REDIS_VALUE_TYPE: &str = "redis_value_type";
49pub const REDIS_VALUE_TYPE_STRING: &str = "string";
50pub const REDIS_VALUE_TYPE_GEO: &str = "geospatial";
51pub const REDIS_VALUE_TYPE_PUBSUB: &str = "pubsub";
52pub const REDIS_VALUE_TYPE_STREAM: &str = "stream";
53pub const LON_NAME: &str = "longitude";
54pub const LAT_NAME: &str = "latitude";
55pub const MEMBER_NAME: &str = "member";
56pub const CHANNEL: &str = "channel";
57pub const CHANNEL_COLUMN: &str = "channel_column";
58pub const STREAM: &str = "stream";
59pub const STREAM_COLUMN: &str = "stream_column";
60
61#[derive(Deserialize, Debug, Clone, WithOptions)]
62pub struct RedisCommon {
63    #[serde(rename = "redis.url")]
64    pub url: String,
65}
66
67impl EnforceSecret for RedisCommon {
68    const ENFORCE_SECRET_PROPERTIES: phf::Set<&'static str> = phf_set! {
69        "redis.url"
70    };
71}
72
73pub enum RedisPipe {
74    Cluster(ClusterPipeline),
75    Single(Pipeline),
76}
77impl RedisPipe {
78    pub async fn query<T: redis::FromRedisValue>(
79        &self,
80        conn: &mut RedisConn,
81    ) -> ConnectorResult<T> {
82        match (self, conn) {
83            (RedisPipe::Cluster(pipe), RedisConn::Cluster(conn)) => Ok(pipe.query(conn)?),
84            (RedisPipe::Single(pipe), RedisConn::Single(conn)) => {
85                Ok(pipe.query_async(conn).await?)
86            }
87            _ => Err(SinkError::Redis("RedisPipe and RedisConn not match".to_owned()).into()),
88        }
89    }
90
91    pub fn clear(&mut self) {
92        match self {
93            RedisPipe::Cluster(pipe) => pipe.clear(),
94            RedisPipe::Single(pipe) => pipe.clear(),
95        }
96    }
97
98    pub fn set(
99        &mut self,
100        k: RedisSinkPayloadWriterInput,
101        v: RedisSinkPayloadWriterInput,
102    ) -> Result<()> {
103        match self {
104            RedisPipe::Cluster(pipe) => match (k, v) {
105                (
106                    RedisSinkPayloadWriterInput::String(k),
107                    RedisSinkPayloadWriterInput::String(v),
108                ) => {
109                    pipe.set(k, v);
110                }
111                (
112                    RedisSinkPayloadWriterInput::RedisGeoKey((key, member)),
113                    RedisSinkPayloadWriterInput::RedisGeoValue((lat, lon)),
114                ) => {
115                    pipe.geo_add(key, (lon, lat, member));
116                }
117                (
118                    RedisSinkPayloadWriterInput::RedisPubSubStreamKey(key),
119                    RedisSinkPayloadWriterInput::String(v),
120                ) => {
121                    pipe.publish(key, v);
122                }
123                (
124                    RedisSinkPayloadWriterInput::RedisPubSubStreamKey(key),
125                    RedisSinkPayloadWriterInput::RedisStreamValue((field, value)),
126                ) => {
127                    pipe.xadd(key, "*", &[(&field, &value)]);
128                }
129                _ => return Err(SinkError::Redis("RedisPipe set not match".to_owned())),
130            },
131            RedisPipe::Single(pipe) => match (k, v) {
132                (
133                    RedisSinkPayloadWriterInput::String(k),
134                    RedisSinkPayloadWriterInput::String(v),
135                ) => {
136                    pipe.set(k, v);
137                }
138                (
139                    RedisSinkPayloadWriterInput::RedisGeoKey((key, member)),
140                    RedisSinkPayloadWriterInput::RedisGeoValue((lat, lon)),
141                ) => {
142                    pipe.geo_add(key, (lon, lat, member));
143                }
144                (
145                    RedisSinkPayloadWriterInput::RedisPubSubStreamKey(key),
146                    RedisSinkPayloadWriterInput::String(v),
147                ) => {
148                    pipe.publish(key, v);
149                }
150                (
151                    RedisSinkPayloadWriterInput::RedisPubSubStreamKey(key),
152                    RedisSinkPayloadWriterInput::RedisStreamValue((field, value)),
153                ) => {
154                    pipe.xadd(key, "*", &[(&field, &value)]);
155                }
156                _ => return Err(SinkError::Redis("RedisPipe set not match".to_owned())),
157            },
158        };
159        Ok(())
160    }
161
162    pub fn del(&mut self, k: RedisSinkPayloadWriterInput) -> Result<()> {
163        match self {
164            RedisPipe::Cluster(pipe) => match k {
165                RedisSinkPayloadWriterInput::String(k) => {
166                    pipe.del(k);
167                }
168                RedisSinkPayloadWriterInput::RedisGeoKey((key, member)) => {
169                    pipe.zrem(key, member);
170                }
171                _ => return Err(SinkError::Redis("RedisPipe del not match".to_owned())),
172            },
173            RedisPipe::Single(pipe) => match k {
174                RedisSinkPayloadWriterInput::String(k) => {
175                    pipe.del(k);
176                }
177                RedisSinkPayloadWriterInput::RedisGeoKey((key, member)) => {
178                    pipe.zrem(key, member);
179                }
180                _ => return Err(SinkError::Redis("RedisPipe del not match".to_owned())),
181            },
182        };
183        Ok(())
184    }
185}
186pub enum RedisConn {
187    // Redis deployed as a cluster, clusters with only one node should also use this conn
188    Cluster(ClusterConnection),
189    // Redis is not deployed as a cluster
190    Single(MultiplexedConnection),
191}
192
193impl RedisCommon {
194    pub async fn build_conn_and_pipe(&self) -> ConnectorResult<(RedisConn, RedisPipe)> {
195        match serde_json::from_str(&self.url).map_err(|e| SinkError::Config(anyhow!(e))) {
196            Ok(v) => {
197                if let Value::Array(list) = v {
198                    let list = list
199                        .into_iter()
200                        .map(|s| {
201                            if let Value::String(s) = s {
202                                Ok(s)
203                            } else {
204                                Err(SinkError::Redis(
205                                    "redis.url must be array of string".to_owned(),
206                                )
207                                .into())
208                            }
209                        })
210                        .collect::<ConnectorResult<Vec<String>>>()?;
211
212                    let client = ClusterClient::new(list)?;
213                    Ok((
214                        RedisConn::Cluster(client.get_connection()?),
215                        RedisPipe::Cluster(redis::cluster::cluster_pipe()),
216                    ))
217                } else {
218                    Err(SinkError::Redis("redis.url must be array or string".to_owned()).into())
219                }
220            }
221            Err(_) => {
222                let client = RedisClient::open(self.url.clone())?;
223                Ok((
224                    RedisConn::Single(client.get_multiplexed_async_connection().await?),
225                    RedisPipe::Single(redis::pipe()),
226                ))
227            }
228        }
229    }
230}
231
232#[serde_as]
233#[derive(Clone, Debug, Deserialize, WithOptions)]
234pub struct RedisConfig {
235    #[serde(flatten)]
236    pub common: RedisCommon,
237
238    #[serde(flatten)]
239    pub unknown_fields: std::collections::HashMap<String, String>,
240}
241
242crate::impl_sink_unknown_fields!(RedisConfig);
243
244impl EnforceSecret for RedisConfig {
245    fn enforce_secret<'a>(prop_iter: impl Iterator<Item = &'a str>) -> ConnectorResult<()> {
246        for prop in prop_iter {
247            RedisCommon::enforce_one(prop)?;
248        }
249        Ok(())
250    }
251}
252
253impl RedisConfig {
254    pub fn from_btreemap(properties: BTreeMap<String, String>) -> Result<Self> {
255        let config =
256            serde_json::from_value::<RedisConfig>(serde_json::to_value(properties).unwrap())
257                .map_err(|e| SinkError::Config(anyhow!(e)))?;
258        Ok(config)
259    }
260}
261
262#[derive(Debug)]
263pub struct RedisSink {
264    config: RedisConfig,
265    schema: Schema,
266    pk_indices: Vec<usize>,
267    format_desc: SinkFormatDesc,
268    db_name: String,
269    sink_from_name: String,
270}
271
272impl EnforceSecret for RedisSink {
273    fn enforce_secret<'a>(prop_iter: impl Iterator<Item = &'a str>) -> ConnectorResult<()> {
274        for prop in prop_iter {
275            RedisConfig::enforce_one(prop)?;
276        }
277        Ok(())
278    }
279}
280
281#[async_trait]
282impl TryFrom<SinkParam> for RedisSink {
283    type Error = SinkError;
284
285    fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
286        let Some(pk_indices) = param.downstream_pk.clone() else {
287            return Err(SinkError::Config(anyhow!(
288                "Redis Sink Primary Key must be specified."
289            )));
290        };
291        let config = RedisConfig::from_btreemap(param.properties.clone())?;
292        Ok(Self {
293            config,
294            schema: param.schema(),
295            pk_indices,
296            format_desc: param
297                .format_desc
298                .ok_or_else(|| SinkError::Config(anyhow!("missing FORMAT ... ENCODE ...")))?,
299            db_name: param.db_name,
300            sink_from_name: param.sink_from_name,
301        })
302    }
303}
304
305impl Sink for RedisSink {
306    type LogSinker = AsyncTruncateLogSinkerOf<RedisSinkWriter>;
307
308    const SINK_NAME: &'static str = "redis";
309
310    crate::impl_validate_sink_unknown_fields!();
311
312    async fn new_log_sinker(&self, _writer_param: SinkWriterParam) -> Result<Self::LogSinker> {
313        Ok(RedisSinkWriter::new(
314            self.config.clone(),
315            self.schema.clone(),
316            self.pk_indices.clone(),
317            &self.format_desc,
318            self.db_name.clone(),
319            self.sink_from_name.clone(),
320        )
321        .await?
322        .into_log_sinker(usize::MAX))
323    }
324
325    async fn validate(&self) -> Result<()> {
326        let all_map: HashMap<String, DataType> = self
327            .schema
328            .fields()
329            .iter()
330            .map(|f| (f.name.clone(), f.data_type.clone()))
331            .collect();
332        let pk_map: HashMap<String, DataType> = self
333            .schema
334            .fields()
335            .iter()
336            .enumerate()
337            .filter(|(k, _)| self.pk_indices.contains(k))
338            .map(|(_, v)| (v.name.clone(), v.data_type.clone()))
339            .collect();
340        if matches!(
341            self.format_desc.encode,
342            super::catalog::SinkEncode::Template
343        ) {
344            match self
345                .format_desc
346                .options
347                .get(REDIS_VALUE_TYPE)
348                .map(|s| s.as_str())
349            {
350                // if not set, default to string
351                Some(REDIS_VALUE_TYPE_STRING) | None => {
352                    let key_format = self.format_desc.options.get(KEY_FORMAT).ok_or_else(|| {
353                        SinkError::Config(anyhow!(
354                            "Cannot find '{KEY_FORMAT}', please set it or use JSON"
355                        ))
356                    })?;
357                    TemplateStringEncoder::check_string_format(key_format, &pk_map)?;
358                    let value_format =
359                        self.format_desc.options.get(VALUE_FORMAT).ok_or_else(|| {
360                            SinkError::Config(anyhow!(
361                                "Cannot find `{VALUE_FORMAT}`, please set it or use JSON"
362                            ))
363                        })?;
364                    TemplateStringEncoder::check_string_format(value_format, &all_map)?;
365                }
366                Some(REDIS_VALUE_TYPE_GEO) => {
367                    let key_format = self.format_desc.options.get(KEY_FORMAT).ok_or_else(|| {
368                        SinkError::Config(anyhow!(
369                            "Cannot find '{KEY_FORMAT}', please set it or use JSON"
370                        ))
371                    })?;
372                    TemplateStringEncoder::check_string_format(key_format, &pk_map)?;
373
374                    let lon_name = self.format_desc.options.get(LON_NAME).ok_or_else(|| {
375                        SinkError::Config(anyhow!(
376                            "Cannot find `{LON_NAME}`, please set it or use JSON or set `{REDIS_VALUE_TYPE}` to `{REDIS_VALUE_TYPE_STRING}`"
377                        ))
378                    })?;
379                    let lat_name = self.format_desc.options.get(LAT_NAME).ok_or_else(|| {
380                        SinkError::Config(anyhow!(
381                            "Cannot find `{LAT_NAME}`, please set it or use JSON or set `{REDIS_VALUE_TYPE}` to `{REDIS_VALUE_TYPE_STRING}`"
382                        ))
383                    })?;
384                    let member_name = self.format_desc.options.get(MEMBER_NAME).ok_or_else(|| {
385                        SinkError::Config(anyhow!(
386                            "Cannot find `{MEMBER_NAME}`, please set it or use JSON or set `{REDIS_VALUE_TYPE}` to `{REDIS_VALUE_TYPE_STRING}`"
387                        ))
388                    })?;
389                    if let Some(lon_type) = all_map.get(lon_name)
390                        && (lon_type == &DataType::Float64
391                            || lon_type == &DataType::Float32
392                            || lon_type == &DataType::Varchar)
393                    {
394                        // do nothing
395                    } else {
396                        return Err(SinkError::Config(anyhow!(
397                            "`{LON_NAME}` must be set to `float64` or `float32` or `varchar`"
398                        )));
399                    }
400                    if let Some(lat_type) = all_map.get(lat_name)
401                        && (lat_type == &DataType::Float64
402                            || lat_type == &DataType::Float32
403                            || lat_type == &DataType::Varchar)
404                    {
405                        // do nothing
406                    } else {
407                        return Err(SinkError::Config(anyhow!(
408                            "`{LAT_NAME}` must be set to `float64` or `float32` or `varchar`"
409                        )));
410                    }
411                    if let Some(member_type) = pk_map.get(member_name)
412                        && member_type == &DataType::Varchar
413                    {
414                        // do nothing
415                    } else {
416                        return Err(SinkError::Config(anyhow!(
417                            "`{MEMBER_NAME}` must be set to `varchar` and `primary_key`"
418                        )));
419                    }
420                }
421                Some(REDIS_VALUE_TYPE_PUBSUB) => {
422                    let channel = self.format_desc.options.get(CHANNEL);
423                    let channel_column = self.format_desc.options.get(CHANNEL_COLUMN);
424                    if (channel.is_none() && channel_column.is_none())
425                        || (channel.is_some() && channel_column.is_some())
426                    {
427                        return Err(SinkError::Config(anyhow!(
428                            "`{CHANNEL}` and `{CHANNEL_COLUMN}` only one can be set"
429                        )));
430                    }
431
432                    if let Some(channel_column) = channel_column
433                        && let Some(channel_column_type) = all_map.get(channel_column)
434                        && (channel_column_type != &DataType::Varchar)
435                    {
436                        return Err(SinkError::Config(anyhow!(
437                            "`{CHANNEL_COLUMN}` must be set to `varchar`"
438                        )));
439                    }
440
441                    let value_format =
442                        self.format_desc.options.get(VALUE_FORMAT).ok_or_else(|| {
443                            SinkError::Config(anyhow!("Cannot find `{VALUE_FORMAT}`"))
444                        })?;
445                    TemplateStringEncoder::check_string_format(value_format, &all_map)?;
446                }
447                Some(REDIS_VALUE_TYPE_STREAM) => {
448                    tracing::error!("test:for bug");
449                    risingwave_common::license::Feature::RedisSinkStream
450                        .check_available()
451                        .map_err(|e| anyhow::anyhow!(e))?;
452                    let stream = self.format_desc.options.get(STREAM);
453                    let stream_column = self.format_desc.options.get(STREAM_COLUMN);
454                    if (stream.is_none() && stream_column.is_none())
455                        || (stream.is_some() && stream_column.is_some())
456                    {
457                        return Err(SinkError::Config(anyhow!(
458                            "Please specific either `{STREAM}` or `{STREAM_COLUMN}`. They are mutually exclusive options."
459                        )));
460                    }
461
462                    if let Some(stream_column) = stream_column
463                        && let Some(stream_column_type) = all_map.get(stream_column)
464                        && (stream_column_type != &DataType::Varchar)
465                    {
466                        return Err(SinkError::Config(anyhow!(
467                            "`{STREAM_COLUMN}` must be set to `varchar`"
468                        )));
469                    }
470
471                    let value_format =
472                        self.format_desc.options.get(VALUE_FORMAT).ok_or_else(|| {
473                            SinkError::Config(anyhow!("Cannot find `{VALUE_FORMAT}`"))
474                        })?;
475                    let key_format = self.format_desc.options.get(KEY_FORMAT).ok_or_else(|| {
476                        SinkError::Config(anyhow!(
477                            "Cannot find '{KEY_FORMAT}', please set it or use JSON"
478                        ))
479                    })?;
480                    TemplateStringEncoder::check_string_format(key_format, &pk_map)?;
481                    TemplateStringEncoder::check_string_format(value_format, &all_map)?;
482                }
483                _ => {
484                    return Err(SinkError::Config(anyhow!(
485                        "`{REDIS_VALUE_TYPE}` must be set to `{REDIS_VALUE_TYPE_STRING}` or `{REDIS_VALUE_TYPE_GEO}` or `{REDIS_VALUE_TYPE_PUBSUB}` or `{REDIS_VALUE_TYPE_STREAM}`"
486                    )));
487                }
488            }
489        }
490        self.config.common.build_conn_and_pipe().await?;
491        Ok(())
492    }
493}
494
495pub struct RedisSinkWriter {
496    #[expect(dead_code)]
497    epoch: u64,
498    #[expect(dead_code)]
499    schema: Schema,
500    #[expect(dead_code)]
501    pk_indices: Vec<usize>,
502    formatter: SinkFormatterImpl,
503    payload_writer: RedisSinkPayloadWriter,
504}
505
506struct RedisSinkPayloadWriter {
507    // connection to redis, one per executor
508    conn: Option<RedisConn>,
509    // the command pipeline for write-commit
510    pipe: RedisPipe,
511}
512
513impl RedisSinkPayloadWriter {
514    pub async fn new(config: RedisConfig) -> Result<Self> {
515        let (conn, pipe) = config.common.build_conn_and_pipe().await?;
516        let conn = Some(conn);
517
518        Ok(Self { conn, pipe })
519    }
520
521    #[cfg(test)]
522    pub fn mock() -> Self {
523        let conn = None;
524        let pipe = RedisPipe::Single(redis::pipe());
525        Self { conn, pipe }
526    }
527
528    pub async fn commit(&mut self) -> Result<()> {
529        #[cfg(test)]
530        {
531            if self.conn.is_none() {
532                return Ok(());
533            }
534        }
535        self.pipe.query::<()>(self.conn.as_mut().unwrap()).await?;
536        self.pipe.clear();
537        Ok(())
538    }
539}
540
541impl FormattedSink for RedisSinkPayloadWriter {
542    type K = RedisSinkPayloadWriterInput;
543    type V = RedisSinkPayloadWriterInput;
544
545    async fn write_one(&mut self, k: Option<Self::K>, v: Option<Self::V>) -> Result<()> {
546        let k = k.ok_or_else(|| SinkError::Redis("The redis key cannot be null".to_owned()))?;
547        match v {
548            Some(v) => self.pipe.set(k, v)?,
549            None => self.pipe.del(k)?,
550        };
551        Ok(())
552    }
553}
554
555impl RedisSinkWriter {
556    pub async fn new(
557        config: RedisConfig,
558        schema: Schema,
559        pk_indices: Vec<usize>,
560        format_desc: &SinkFormatDesc,
561        db_name: String,
562        sink_from_name: String,
563    ) -> Result<Self> {
564        let payload_writer = RedisSinkPayloadWriter::new(config.clone()).await?;
565        let formatter = SinkFormatterImpl::new(
566            format_desc,
567            schema.clone(),
568            pk_indices.clone(),
569            db_name,
570            sink_from_name,
571            "NO_TOPIC",
572        )
573        .await?;
574
575        Ok(Self {
576            schema,
577            pk_indices,
578            epoch: 0,
579            formatter,
580            payload_writer,
581        })
582    }
583
584    #[cfg(test)]
585    pub async fn mock(
586        schema: Schema,
587        pk_indices: Vec<usize>,
588        format_desc: &SinkFormatDesc,
589    ) -> Result<Self> {
590        let formatter = SinkFormatterImpl::new(
591            format_desc,
592            schema.clone(),
593            pk_indices.clone(),
594            "d1".to_owned(),
595            "t1".to_owned(),
596            "NO_TOPIC",
597        )
598        .await?;
599        Ok(Self {
600            schema,
601            pk_indices,
602            epoch: 0,
603            formatter,
604            payload_writer: RedisSinkPayloadWriter::mock(),
605        })
606    }
607}
608
609impl AsyncTruncateSinkWriter for RedisSinkWriter {
610    async fn write_chunk<'a>(
611        &'a mut self,
612        chunk: StreamChunk,
613        _add_future: DeliveryFutureManagerAddFuture<'a, Self::DeliveryFuture>,
614    ) -> Result<()> {
615        dispatch_sink_formatter_str_key_impl!(&self.formatter, formatter, {
616            self.payload_writer.write_chunk(chunk, formatter).await?;
617            self.payload_writer.commit().await
618        })
619    }
620}
621
622#[cfg(test)]
623mod test {
624    use core::panic;
625
626    use rdkafka::message::FromBytes;
627    use risingwave_common::array::{Array, I32Array, Op, Utf8Array};
628    use risingwave_common::catalog::Field;
629    use risingwave_common::types::DataType;
630    use risingwave_common::util::iter_util::ZipEqDebug;
631
632    use super::*;
633    use crate::sink::catalog::{SinkEncode, SinkFormat};
634    use crate::sink::log_store::DeliveryFutureManager;
635
636    #[tokio::test]
637    async fn test_write() {
638        let schema = Schema::new(vec![
639            Field {
640                data_type: DataType::Int32,
641                name: "id".to_owned(),
642            },
643            Field {
644                data_type: DataType::Varchar,
645                name: "name".to_owned(),
646            },
647        ]);
648
649        let format_desc = SinkFormatDesc {
650            format: SinkFormat::AppendOnly,
651            encode: SinkEncode::Json,
652            options: BTreeMap::default(),
653            secret_refs: BTreeMap::default(),
654            key_encode: None,
655            connection_id: None,
656        };
657
658        let mut redis_sink_writer = RedisSinkWriter::mock(schema, vec![0], &format_desc)
659            .await
660            .unwrap();
661
662        let chunk_a = StreamChunk::new(
663            vec![Op::Insert, Op::Insert, Op::Insert],
664            vec![
665                I32Array::from_iter(vec![1, 2, 3]).into_ref(),
666                Utf8Array::from_iter(vec!["Alice", "Bob", "Clare"]).into_ref(),
667            ],
668        );
669
670        let mut manager = DeliveryFutureManager::new(0);
671
672        redis_sink_writer
673            .write_chunk(chunk_a, manager.start_write_chunk(0, 0))
674            .await
675            .expect("failed to write batch");
676        let expected_a = vec![
677            (
678                0,
679                "*3\r\n$3\r\nSET\r\n$8\r\n{\"id\":1}\r\n$23\r\n{\"id\":1,\"name\":\"Alice\"}\r\n",
680            ),
681            (
682                1,
683                "*3\r\n$3\r\nSET\r\n$8\r\n{\"id\":2}\r\n$21\r\n{\"id\":2,\"name\":\"Bob\"}\r\n",
684            ),
685            (
686                2,
687                "*3\r\n$3\r\nSET\r\n$8\r\n{\"id\":3}\r\n$23\r\n{\"id\":3,\"name\":\"Clare\"}\r\n",
688            ),
689        ];
690
691        if let RedisPipe::Single(pipe) = &redis_sink_writer.payload_writer.pipe {
692            pipe.cmd_iter()
693                .enumerate()
694                .zip_eq_debug(expected_a.clone())
695                .for_each(|((i, cmd), (exp_i, exp_cmd))| {
696                    if exp_i == i {
697                        assert_eq!(exp_cmd, str::from_bytes(&cmd.get_packed_command()).unwrap())
698                    }
699                });
700        } else {
701            panic!("pipe type not match")
702        }
703    }
704
705    #[tokio::test]
706    async fn test_format_write() {
707        let schema = Schema::new(vec![
708            Field {
709                data_type: DataType::Int32,
710                name: "id".to_owned(),
711            },
712            Field {
713                data_type: DataType::Varchar,
714                name: "name".to_owned(),
715            },
716        ]);
717
718        let mut btree_map = BTreeMap::default();
719        btree_map.insert(KEY_FORMAT.to_owned(), "key-{id}".to_owned());
720        btree_map.insert(
721            VALUE_FORMAT.to_owned(),
722            "values:\\{id:{id},name:{name}\\}".to_owned(),
723        );
724        let format_desc = SinkFormatDesc {
725            format: SinkFormat::AppendOnly,
726            encode: SinkEncode::Template,
727            options: btree_map,
728            secret_refs: Default::default(),
729            key_encode: None,
730            connection_id: None,
731        };
732
733        let mut redis_sink_writer = RedisSinkWriter::mock(schema, vec![0], &format_desc)
734            .await
735            .unwrap();
736
737        let mut future_manager = DeliveryFutureManager::new(0);
738
739        let chunk_a = StreamChunk::new(
740            vec![Op::Insert, Op::Insert, Op::Insert],
741            vec![
742                I32Array::from_iter(vec![1, 2, 3]).into_ref(),
743                Utf8Array::from_iter(vec!["Alice", "Bob", "Clare"]).into_ref(),
744            ],
745        );
746
747        redis_sink_writer
748            .write_chunk(chunk_a, future_manager.start_write_chunk(0, 0))
749            .await
750            .expect("failed to write batch");
751        let expected_a = vec![
752            (
753                0,
754                "*3\r\n$3\r\nSET\r\n$5\r\nkey-1\r\n$24\r\nvalues:{id:1,name:Alice}\r\n",
755            ),
756            (
757                1,
758                "*3\r\n$3\r\nSET\r\n$5\r\nkey-2\r\n$22\r\nvalues:{id:2,name:Bob}\r\n",
759            ),
760            (
761                2,
762                "*3\r\n$3\r\nSET\r\n$5\r\nkey-3\r\n$24\r\nvalues:{id:3,name:Clare}\r\n",
763            ),
764        ];
765
766        if let RedisPipe::Single(pipe) = &redis_sink_writer.payload_writer.pipe {
767            pipe.cmd_iter()
768                .enumerate()
769                .zip_eq_debug(expected_a.clone())
770                .for_each(|((i, cmd), (exp_i, exp_cmd))| {
771                    if exp_i == i {
772                        assert_eq!(exp_cmd, str::from_bytes(&cmd.get_packed_command()).unwrap())
773                    }
774                });
775        } else {
776            panic!("pipe type not match")
777        };
778    }
779}
risingwave_connector/sink/redis.rs

risingwave_connector/sink/
redis.rs
