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

risingwave_connector/sink/
redis.rs
