risingwave_connector/parser/
scalar_adapter.rs

1// Copyright 2024 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::str::FromStr;
16
17use anyhow::anyhow;
18use bytes::{Buf, BufMut, BytesMut};
19use pg_bigdecimal::PgNumeric;
20use risingwave_common::types::{
21    DataType, Decimal, Int256, ListValue, ScalarImpl, ScalarRefImpl, StructValue,
22};
23use thiserror_ext::AsReport;
24use tokio_postgres::types::{FromSql, IsNull, Kind, ToSql, Type, to_sql_checked};
25
26use crate::connector_common::postgres::postgres_point_type;
27use crate::error::ConnectorResult;
28
29#[derive(Clone, Debug)]
30pub struct EnumString(pub String);
31
32impl<'a> FromSql<'a> for EnumString {
33    fn from_sql(
34        _ty: &Type,
35        raw: &'a [u8],
36    ) -> Result<Self, Box<dyn std::error::Error + 'static + Sync + Send>> {
37        Ok(EnumString(String::from_utf8_lossy(raw).into_owned()))
38    }
39
40    fn accepts(ty: &Type) -> bool {
41        matches!(ty.kind(), Kind::Enum(_))
42    }
43}
44
45impl ToSql for EnumString {
46    to_sql_checked!();
47
48    fn to_sql(
49        &self,
50        ty: &Type,
51        out: &mut BytesMut,
52    ) -> Result<IsNull, Box<dyn std::error::Error + Sync + Send>>
53    where
54        Self: Sized,
55    {
56        match ty.kind() {
57            Kind::Enum(e) => {
58                if e.contains(&self.0) {
59                    out.extend_from_slice(self.0.as_bytes());
60                    Ok(IsNull::No)
61                } else {
62                    Err(format!(
63                        "EnumString value {} is not in the enum type {:?}",
64                        self.0, e
65                    )
66                    .into())
67                }
68            }
69            _ => Err("EnumString can only be used with ENUM types".into()),
70        }
71    }
72
73    fn accepts(ty: &Type) -> bool {
74        matches!(ty.kind(), Kind::Enum(_))
75    }
76}
77
78#[derive(Debug)]
79pub(crate) struct PgPoint {
80    x: f64,
81    y: f64,
82}
83
84impl<'a> FromSql<'a> for PgPoint {
85    fn from_sql(
86        _ty: &Type,
87        raw: &'a [u8],
88    ) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
89        if raw.len() != 2 * std::mem::size_of::<f64>() {
90            return Err("invalid point binary payload length".into());
91        }
92
93        let mut buf = raw;
94
95        Ok(Self {
96            x: buf.get_f64(),
97            y: buf.get_f64(),
98        })
99    }
100
101    fn accepts(ty: &Type) -> bool {
102        *ty == Type::POINT
103    }
104}
105
106impl ToSql for PgPoint {
107    to_sql_checked!();
108
109    fn to_sql(
110        &self,
111        _ty: &Type,
112        out: &mut BytesMut,
113    ) -> Result<IsNull, Box<dyn std::error::Error + Sync + Send>> {
114        out.put_f64(self.x);
115        out.put_f64(self.y);
116
117        Ok(IsNull::No)
118    }
119
120    fn accepts(ty: &Type) -> bool {
121        *ty == Type::POINT
122    }
123}
124
125/// Adapter for `ScalarImpl` to Postgres data type,
126/// which can be used to encode/decode to/from Postgres value.
127#[derive(Debug)]
128pub(crate) enum ScalarAdapter {
129    Builtin(ScalarImpl),
130    Uuid(uuid::Uuid),
131    Point(PgPoint),
132    // Currently in order to handle the decimal beyond RustDecimal,
133    // we use the PgNumeric type to convert the decimal to a string/decimal/rw_int256.
134    Numeric(PgNumeric),
135    Enum(EnumString),
136    NumericList(Vec<Option<PgNumeric>>),
137    EnumList(Vec<Option<EnumString>>),
138    // UuidList is covered by List, while NumericList and EnumList are special cases.
139    // Note: The IntervalList is not supported.
140    List(Vec<Option<ScalarAdapter>>),
141}
142
143impl ToSql for ScalarAdapter {
144    to_sql_checked!();
145
146    fn to_sql(
147        &self,
148        ty: &Type,
149        out: &mut bytes::BytesMut,
150    ) -> Result<IsNull, Box<dyn std::error::Error + Sync + Send>> {
151        match self {
152            ScalarAdapter::Builtin(v) => v.to_sql(ty, out),
153            ScalarAdapter::Uuid(v) => v.to_sql(ty, out),
154            ScalarAdapter::Point(v) => v.to_sql(ty, out),
155            ScalarAdapter::Numeric(v) => v.to_sql(ty, out),
156            ScalarAdapter::Enum(v) => v.to_sql(ty, out),
157            ScalarAdapter::NumericList(v) => v.to_sql(ty, out),
158            ScalarAdapter::EnumList(v) => v.to_sql(ty, out),
159            ScalarAdapter::List(v) => v.to_sql(ty, out),
160        }
161    }
162
163    fn accepts(_ty: &Type) -> bool {
164        true
165    }
166}
167
168/// convert from Postgres uuid, numeric and enum to `ScalarAdapter`
169impl<'a> FromSql<'a> for ScalarAdapter {
170    fn from_sql(
171        ty: &Type,
172        raw: &'a [u8],
173    ) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
174        match ty.kind() {
175            Kind::Simple => match *ty {
176                Type::UUID => Ok(ScalarAdapter::Uuid(uuid::Uuid::from_sql(ty, raw)?)),
177                Type::POINT => Ok(ScalarAdapter::Point(PgPoint::from_sql(ty, raw)?)),
178                // In order to cover the decimal beyond RustDecimal(only 28 digits are supported),
179                // we use the PgNumeric to handle decimal from postgres.
180                Type::NUMERIC => Ok(ScalarAdapter::Numeric(PgNumeric::from_sql(ty, raw)?)),
181                _ => Ok(ScalarAdapter::Builtin(ScalarImpl::from_sql(ty, raw)?)),
182            },
183            Kind::Enum(_) => Ok(ScalarAdapter::Enum(EnumString::from_sql(ty, raw)?)),
184            Kind::Array(Type::NUMERIC) => {
185                Ok(ScalarAdapter::NumericList(FromSql::from_sql(ty, raw)?))
186            }
187            Kind::Array(inner_type) if let Kind::Enum(_) = inner_type.kind() => {
188                Ok(ScalarAdapter::EnumList(FromSql::from_sql(ty, raw)?))
189            }
190            Kind::Array(_) => Ok(ScalarAdapter::List(FromSql::from_sql(ty, raw)?)),
191            _ => Err(anyhow!("failed to convert type {:?} to ScalarAdapter", ty).into()),
192        }
193    }
194
195    fn accepts(ty: &Type) -> bool {
196        match ty.kind() {
197            Kind::Simple => {
198                matches!(ty, &Type::UUID | &Type::NUMERIC | &Type::POINT)
199                    || <ScalarImpl as FromSql>::accepts(ty)
200            }
201            Kind::Enum(_) => true,
202            Kind::Array(inner_type) => <ScalarAdapter as FromSql>::accepts(inner_type),
203            _ => false,
204        }
205    }
206}
207
208impl ScalarAdapter {
209    pub fn name(&self) -> &'static str {
210        match self {
211            ScalarAdapter::Builtin(_) => "Builtin",
212            ScalarAdapter::Uuid(_) => "Uuid",
213            ScalarAdapter::Point(_) => "Point",
214            ScalarAdapter::Numeric(_) => "Numeric",
215            ScalarAdapter::Enum(_) => "Enum",
216            ScalarAdapter::EnumList(_) => "EnumList",
217            ScalarAdapter::NumericList(_) => "NumericList",
218            ScalarAdapter::List(_) => "List",
219        }
220    }
221
222    /// convert `ScalarRefImpl` to `ScalarAdapter` so that we can correctly encode to postgres value
223    pub(crate) fn from_scalar(
224        scalar: ScalarRefImpl<'_>,
225        ty: &Type,
226    ) -> ConnectorResult<ScalarAdapter> {
227        Ok(match (scalar, ty, ty.kind()) {
228            (ScalarRefImpl::Utf8(s), &Type::UUID, _) => ScalarAdapter::Uuid(s.parse()?),
229            (ScalarRefImpl::Struct(point), &Type::POINT, _) => {
230                let mut fields = point.iter_fields_ref();
231
232                let (
233                    Some(Some(ScalarRefImpl::Float64(x))),
234                    Some(Some(ScalarRefImpl::Float64(y))),
235                    None,
236                ) = (fields.next(), fields.next(), fields.next())
237                else {
238                    return Err(anyhow!(
239                        "failed to convert struct to PostgreSQL point: expected exactly two non-null float64 fields"
240                    )
241                    .into());
242                };
243
244                ScalarAdapter::Point(PgPoint {
245                    x: x.into_inner(),
246                    y: y.into_inner(),
247                })
248            }
249            (ScalarRefImpl::Utf8(s), &Type::NUMERIC, _) => {
250                ScalarAdapter::Numeric(string_to_pg_numeric(s))
251            }
252            (ScalarRefImpl::Int256(s), &Type::NUMERIC, _) => {
253                ScalarAdapter::Numeric(string_to_pg_numeric(&s.to_string()))
254            }
255            (ScalarRefImpl::Utf8(s), _, Kind::Enum(_)) => {
256                ScalarAdapter::Enum(EnumString(s.to_owned()))
257            }
258            (ScalarRefImpl::List(list), &Type::NUMERIC_ARRAY, _) => {
259                let mut vec = vec![];
260                for datum in list.iter() {
261                    vec.push(match datum {
262                        Some(ScalarRefImpl::Int256(s)) => Some(string_to_pg_numeric(&s.to_string())),
263                        Some(ScalarRefImpl::Decimal(s)) => Some(rw_numeric_to_pg_numeric(s)),
264                        Some(ScalarRefImpl::Utf8(s)) => Some(string_to_pg_numeric(s)),
265                        None => None,
266                        _ => {
267                            unreachable!("Only rw-numeric[], rw_int256[] and varchar[] are supported to convert to pg-numeric[]");
268                        }
269                    })
270                }
271                ScalarAdapter::NumericList(vec)
272            }
273            (ScalarRefImpl::List(list), _, Kind::Array(inner_type)) => match inner_type.kind() {
274                Kind::Enum(_) => {
275                    let mut vec = vec![];
276                    for datum in list.iter() {
277                        vec.push(match datum {
278                            Some(ScalarRefImpl::Utf8(s)) => Some(EnumString(s.to_owned())),
279                            _ => unreachable!(
280                                "Only non-null varchar[] is supported to convert to enum[]"
281                            ),
282                        })
283                    }
284                    ScalarAdapter::EnumList(vec)
285                }
286                _ => {
287                    let mut vec = vec![];
288                    for datum in list.iter() {
289                        vec.push(
290                            datum
291                                .map(|s| ScalarAdapter::from_scalar(s, inner_type))
292                                .transpose()?,
293                        );
294                    }
295                    ScalarAdapter::List(vec)
296                }
297            },
298            _ => ScalarAdapter::Builtin(scalar.into_scalar_impl()),
299        })
300    }
301
302    pub fn into_scalar(self, ty: &DataType) -> Option<ScalarImpl> {
303        match (self, &ty) {
304            (ScalarAdapter::Builtin(scalar), _) => Some(scalar),
305            (ScalarAdapter::Uuid(uuid), &DataType::Varchar) => {
306                Some(ScalarImpl::from(uuid.to_string()))
307            }
308            (ScalarAdapter::Point(PgPoint { x, y }), &DataType::Struct(_)) => {
309                assert_eq!(
310                    ty,
311                    &postgres_point_type(),
312                    "PostgreSQL point must map to struct<x float64, y float64>"
313                );
314
315                Some(StructValue::new(vec![Some(x.into()), Some(y.into())]).into())
316            }
317            (ScalarAdapter::Numeric(numeric), &DataType::Varchar) => {
318                Some(ScalarImpl::from(pg_numeric_to_string(&numeric)))
319            }
320            (ScalarAdapter::Numeric(numeric), &DataType::Int256) => {
321                pg_numeric_to_rw_int256(&numeric)
322            }
323            (ScalarAdapter::Numeric(numeric), &DataType::Decimal) => {
324                pg_numeric_to_rw_numeric(&numeric)
325            }
326            (ScalarAdapter::Enum(EnumString(s)), &DataType::Varchar) => Some(ScalarImpl::from(s)),
327            (ScalarAdapter::NumericList(vec), &DataType::List(list)) => {
328                let elem = list.elem();
329                let mut builder = elem.create_array_builder(0);
330                for val in vec {
331                    let scalar = match (val, &elem) {
332                        // A numeric array contains special values like NaN, Inf, -Inf, which are not supported in Debezium,
333                        // when we encounter these special values, we fallback the array to NULL, returning None directly.
334                        (Some(numeric), DataType::Varchar) => {
335                            if pg_numeric_is_special(&numeric) {
336                                return None;
337                            } else {
338                                ScalarAdapter::Numeric(numeric).into_scalar(elem)
339                            }
340                        }
341                        (Some(numeric), DataType::Int256 | DataType::Decimal) => {
342                            if pg_numeric_is_special(&numeric) {
343                                return None;
344                            } else {
345                                // A PgNumeric can sometimes exceeds the range of Int256 and RwNumeric.
346                                // In our json parsing, we fallback the array to NULL in this case.
347                                // Here we keep the behavior consistent and return None directly.
348                                match ScalarAdapter::Numeric(numeric).into_scalar(elem) {
349                                    Some(scalar) => Some(scalar),
350                                    None => {
351                                        return None;
352                                    }
353                                }
354                            }
355                        }
356                        (Some(_), _) => unreachable!(
357                            "Only rw-numeric[], rw_int256[] and varchar[] are supported to convert to pg-numeric[]"
358                        ),
359                        // This item is NULL, continue to handle next item.
360                        (None, _) => None,
361                    };
362                    builder.append(scalar);
363                }
364                Some(ScalarImpl::from(ListValue::new(builder.finish())))
365            }
366            (ScalarAdapter::EnumList(vec), &DataType::List(list)) => {
367                let mut builder = list.elem().create_array_builder(0);
368                for val in vec {
369                    match val {
370                        Some(EnumString(s)) => {
371                            builder.append(Some(ScalarImpl::from(s)));
372                        }
373                        None => {
374                            return None;
375                        }
376                    }
377                }
378                Some(ScalarImpl::from(ListValue::new(builder.finish())))
379            }
380            (ScalarAdapter::List(vec), &DataType::List(list)) => {
381                let elem = list.elem();
382                // Due to https://github.com/risingwavelabs/risingwave/issues/16882, INTERVAL_ARRAY is not supported in Debezium, so we keep backfilling and CDC consistent.
383                if matches!(elem, DataType::Interval) {
384                    return None;
385                }
386                let mut builder = elem.create_array_builder(0);
387                for val in vec {
388                    builder.append(val.and_then(|v| v.into_scalar(elem)));
389                }
390                Some(ScalarImpl::from(ListValue::new(builder.finish())))
391            }
392            (scaler, ty) => {
393                tracing::error!(
394                    adapter = scaler.name(),
395                    rw_type = ty.pg_name(),
396                    "failed to convert from ScalarAdapter: invalid conversion"
397                );
398                None
399            }
400        }
401    }
402}
403
404pub fn validate_pg_type_to_rw_type(pg_type: &DataType, rw_type: &DataType) -> bool {
405    if pg_type == rw_type {
406        return true;
407    }
408    match rw_type {
409        DataType::Varchar => matches!(pg_type, DataType::Decimal | DataType::Int256),
410        DataType::List(list) if list.elem() == &DataType::Varchar => {
411            matches!(
412                pg_type,
413                DataType::List(list) if matches!(list.elem(), DataType::Decimal | DataType::Int256)
414            )
415        }
416        _ => false,
417    }
418}
419
420fn pg_numeric_is_special(val: &PgNumeric) -> bool {
421    matches!(
422        val,
423        PgNumeric::NegativeInf | PgNumeric::PositiveInf | PgNumeric::NaN
424    )
425}
426
427fn pg_numeric_to_rw_int256(val: &PgNumeric) -> Option<ScalarImpl> {
428    match Int256::from_str(pg_numeric_to_string(val).as_str()) {
429        Ok(num) => Some(ScalarImpl::from(num)),
430        Err(err) => {
431            tracing::error!(error = %err.as_report(), "failed to convert PgNumeric to Int256");
432            None
433        }
434    }
435}
436
437fn pg_numeric_to_rw_numeric(val: &PgNumeric) -> Option<ScalarImpl> {
438    match val {
439        PgNumeric::NegativeInf => Some(ScalarImpl::from(Decimal::NegativeInf)),
440        PgNumeric::Normalized(big_decimal) => {
441            match Decimal::from_str(big_decimal.to_string().as_str()) {
442                Ok(num) => Some(ScalarImpl::from(num)),
443                Err(err) => {
444                    tracing::error!(error = %err.as_report(), "parse pg-numeric as rw-numeric failed (likely out-of-range");
445                    None
446                }
447            }
448        }
449        PgNumeric::PositiveInf => Some(ScalarImpl::from(Decimal::PositiveInf)),
450        PgNumeric::NaN => Some(ScalarImpl::from(Decimal::NaN)),
451    }
452}
453
454fn pg_numeric_to_string(val: &PgNumeric) -> String {
455    // TODO(kexiang): NEGATIVE_INFINITY -> -Infinity, POSITIVE_INFINITY -> Infinity, NAN -> NaN
456    // The current implementation is to ensure consistency with the behavior of cdc event parsor.
457    match val {
458        PgNumeric::NegativeInf => String::from("NEGATIVE_INFINITY"),
459        PgNumeric::Normalized(big_decimal) => big_decimal.to_string(),
460        PgNumeric::PositiveInf => String::from("POSITIVE_INFINITY"),
461        PgNumeric::NaN => String::from("NAN"),
462    }
463}
464
465fn string_to_pg_numeric(s: &str) -> PgNumeric {
466    match s {
467        "NEGATIVE_INFINITY" => PgNumeric::NegativeInf,
468        "POSITIVE_INFINITY" => PgNumeric::PositiveInf,
469        "NAN" => PgNumeric::NaN,
470        _ => PgNumeric::Normalized(s.parse().unwrap()),
471    }
472}
473
474fn rw_numeric_to_pg_numeric(val: Decimal) -> PgNumeric {
475    match val {
476        Decimal::NegativeInf => PgNumeric::NegativeInf,
477        Decimal::Normalized(inner) => PgNumeric::Normalized(inner.to_string().parse().unwrap()),
478        Decimal::PositiveInf => PgNumeric::PositiveInf,
479        Decimal::NaN => PgNumeric::NaN,
480    }
481}
482
483#[cfg(test)]
484mod tests {
485    use risingwave_common::types::{ScalarImpl, StructValue};
486    use tokio_postgres::types::{FromSql, Type};
487
488    use super::ScalarAdapter;
489    use crate::connector_common::postgres::postgres_point_type;
490
491    #[test]
492    fn test_postgres_point_into_scalar() {
493        let mut raw = vec![];
494        raw.extend_from_slice(&16777217.25f64.to_be_bytes());
495        raw.extend_from_slice(&(-0.987654321098765f64).to_be_bytes());
496
497        let adapter = ScalarAdapter::from_sql(&Type::POINT, &raw).unwrap();
498        assert_eq!(
499            adapter.into_scalar(&postgres_point_type()),
500            Some(ScalarImpl::Struct(StructValue::new(vec![
501                Some(ScalarImpl::Float64(16777217.25.into())),
502                Some(ScalarImpl::Float64((-0.987654321098765).into())),
503            ])))
504        );
505    }
506}
risingwave_connector/parser/scalar_adapter.rs

risingwave_connector/parser/
scalar_adapter.rs
