risingwave_connector/parser/
postgres.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::sync::LazyLock;
16
17use anyhow::{Context, anyhow, bail};
18use bytes::Buf;
19use risingwave_common::array::Finite32;
20use risingwave_common::catalog::Schema;
21use risingwave_common::log::LogSuppressor;
22use risingwave_common::row::OwnedRow;
23use risingwave_common::types::{DataType, Datum, Decimal, ScalarImpl, VectorVal};
24use thiserror_ext::AsReport;
25use tokio_postgres::types::{FromSql, Type};
26
27use crate::parser::scalar_adapter::ScalarAdapter;
28use crate::parser::utils::log_error;
29
30static LOG_SUPPRESSOR: LazyLock<LogSuppressor> = LazyLock::new(LogSuppressor::default);
31
32/// Adapter for PostgreSQL `vector` type in CDC snapshot reads.
33/// It parses pgvector binary format.
34struct PgVectorAdapter(Vec<f32>);
35
36impl<'a> FromSql<'a> for PgVectorAdapter {
37    fn accepts(ty: &Type) -> bool {
38        ty.name() == "vector"
39    }
40
41    fn from_sql(
42        _ty: &Type,
43        raw: &'a [u8],
44    ) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
45        Self::parse_binary(raw)
46    }
47}
48
49impl PgVectorAdapter {
50    fn parse_binary(raw: &[u8]) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
51        // Binary format from pgvector extension:
52        // int16 dimension, int16 unused, repeated float4 values.
53        if raw.len() < 4 {
54            return Err("invalid vector binary payload".into());
55        }
56        let mut buf = raw;
57        let dim = buf.get_u16() as usize;
58        let _unused = buf.get_u16();
59        if buf.remaining() != dim * std::mem::size_of::<f32>() {
60            return Err("invalid vector binary payload length".into());
61        }
62        let mut elems = Vec::with_capacity(dim);
63        for _ in 0..dim {
64            elems.push(buf.get_f32());
65        }
66        Ok(Self(elems))
67    }
68}
69
70macro_rules! try_handle_data_type {
71    ($row:expr, $i:expr, $name:expr, $type:ty) => {{
72        $row.try_get::<_, Option<$type>>($i)
73            .map(|value| value.map(ScalarImpl::from))
74            .with_context(|| {
75                format!(
76                    "failed to decode PostgreSQL snapshot column `{}` as {}",
77                    $name,
78                    stringify!($type)
79                )
80            })
81    }};
82}
83
84pub fn postgres_row_to_owned_row(row: tokio_postgres::Row, schema: &Schema) -> OwnedRow {
85    let mut datums = vec![];
86    for i in 0..schema.fields.len() {
87        let rw_field = &schema.fields[i];
88        let name = rw_field.name.as_str();
89        let datum = postgres_cell_to_scalar_impl(&row, &rw_field.data_type, i, name);
90        datums.push(datum);
91    }
92    OwnedRow::new(datums)
93}
94
95/// Decode primary-key columns strictly while preserving the legacy lenient behavior for all
96/// other columns in a PostgreSQL CDC snapshot row.
97pub fn postgres_row_to_owned_row_with_strict_pk(
98    row: tokio_postgres::Row,
99    schema: &Schema,
100    pk_indices: &[usize],
101) -> anyhow::Result<OwnedRow> {
102    super::decode_row_with_strict_pk(
103        "PostgreSQL",
104        schema,
105        pk_indices,
106        |index, field| {
107            postgres_cell_to_scalar_impl_strict(&row, &field.data_type, index, &field.name)
108        },
109        |name, err| log_error!(name, err, "parse column failed"),
110    )
111}
112
113pub fn postgres_cell_to_scalar_impl(
114    row: &tokio_postgres::Row,
115    data_type: &DataType,
116    i: usize,
117    name: &str,
118) -> Option<ScalarImpl> {
119    match postgres_cell_to_scalar_impl_strict(row, data_type, i, name) {
120        Ok(datum) => datum,
121        Err(err) => {
122            log_error!(name, err, "parse column failed");
123            None
124        }
125    }
126}
127
128pub fn postgres_cell_to_scalar_impl_strict(
129    row: &tokio_postgres::Row,
130    data_type: &DataType,
131    i: usize,
132    name: &str,
133) -> anyhow::Result<Datum> {
134    // We observe several incompatibility issue in Debezium's Postgres connector. We summarize them here:
135    // Issue #1. The null of enum list is not supported in Debezium. An enum list contains `NULL` will fallback to `NULL`.
136    // Issue #2. In our parser, when there's inf, -inf, nan or invalid item in a list, the whole list will fallback null.
137    match data_type {
138        DataType::Boolean
139        | DataType::Int16
140        | DataType::Int32
141        | DataType::Int64
142        | DataType::Float32
143        | DataType::Float64
144        | DataType::Date
145        | DataType::Time
146        | DataType::Timestamp
147        | DataType::Timestamptz
148        | DataType::Jsonb
149        | DataType::Interval
150        | DataType::Bytea => {
151            // ScalarAdapter is also fine. But ScalarImpl is more efficient
152            row.try_get::<_, Option<ScalarImpl>>(i)
153                .with_context(|| format!("failed to decode PostgreSQL snapshot column `{name}`"))
154        }
155        DataType::Decimal => {
156            // Decimal is more efficient than PgNumeric in ScalarAdapter
157            try_handle_data_type!(row, i, name, Decimal)
158        }
159        DataType::Varchar | DataType::Int256 => {
160            match row
161                .try_get::<_, Option<ScalarAdapter>>(i)
162                .with_context(|| format!("failed to decode PostgreSQL snapshot column `{name}`"))?
163            {
164                Some(value) => value.into_scalar(data_type).map(Some).ok_or_else(|| {
165                    anyhow!("failed to convert PostgreSQL snapshot column `{name}` to {data_type}")
166                }),
167                None => Ok(None),
168            }
169        }
170        DataType::Vector(expected_size) => {
171            match row
172                .try_get::<_, Option<PgVectorAdapter>>(i)
173                .with_context(|| format!("failed to decode PostgreSQL snapshot column `{name}`"))?
174            {
175                Some(PgVectorAdapter(v)) => {
176                    if v.len() != *expected_size {
177                        bail!(
178                            "PostgreSQL snapshot column `{name}` vector dimension mismatch: \
179                             expected {}, got {}",
180                            expected_size,
181                            v.len()
182                        );
183                    }
184                    let finite = v
185                        .into_iter()
186                        .map(Finite32::try_from)
187                        .collect::<Result<Vec<_>, _>>()
188                        .map_err(anyhow::Error::msg)
189                        .with_context(|| {
190                            format!(
191                                "PostgreSQL snapshot column `{name}` contains a non-finite vector \
192                                 element"
193                            )
194                        })?;
195                    Ok(Some(ScalarImpl::Vector(VectorVal::from(finite))))
196                }
197                None => Ok(None),
198            }
199        }
200        DataType::List(list) => match list.elem() {
201            // TODO(Kexiang): allow DataType::List(_)
202            elem @ (DataType::Struct(_) | DataType::List(_) | DataType::Serial) => {
203                bail!("unsupported PostgreSQL snapshot list element type {elem}")
204            }
205            _ => {
206                match row
207                    .try_get::<_, Option<ScalarAdapter>>(i)
208                    .with_context(|| {
209                        format!("failed to decode PostgreSQL snapshot list column `{name}`")
210                    })? {
211                    Some(value) => value.into_scalar(data_type).map(Some).ok_or_else(|| {
212                        anyhow!(
213                            "failed to convert PostgreSQL snapshot column `{name}` to {data_type}"
214                        )
215                    }),
216                    None => Ok(None),
217                }
218            }
219        },
220        DataType::Struct(_) | DataType::Serial | DataType::Map(_) | DataType::Variant => {
221            bail!("unsupported PostgreSQL snapshot data type {data_type} for column `{name}`")
222        }
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use tokio_postgres::NoTls;
229
230    use crate::parser::postgres::PgVectorAdapter;
231    use crate::parser::scalar_adapter::EnumString;
232    const DB: &str = "postgres";
233    const USER: &str = "kexiang";
234
235    #[test]
236    fn test_pg_vector_adapter_parse_binary() {
237        let mut raw = vec![];
238        // dim = 3
239        raw.extend_from_slice(&(3u16.to_be_bytes()));
240        // unused
241        raw.extend_from_slice(&(0u16.to_be_bytes()));
242        raw.extend_from_slice(&1.5f32.to_be_bytes());
243        raw.extend_from_slice(&(-2.25f32).to_be_bytes());
244        raw.extend_from_slice(&3.0f32.to_be_bytes());
245
246        let v = PgVectorAdapter::parse_binary(&raw).unwrap();
247        assert_eq!(v.0, vec![1.5, -2.25, 3.0]);
248    }
249
250    #[ignore]
251    #[tokio::test]
252    async fn enum_string_integration_test() {
253        let connect = format!(
254            "host=localhost port=5432 user={} password={} dbname={}",
255            USER, DB, DB
256        );
257        let (client, connection) = tokio_postgres::connect(connect.as_str(), NoTls)
258            .await
259            .unwrap();
260
261        // The connection object performs the actual communication with the database,
262        // so spawn it off to run on its own.
263        tokio::spawn(async move {
264            if let Err(e) = connection.await {
265                eprintln!("connection error: {}", e);
266            }
267        });
268
269        // allow type existed
270        let _ = client
271            .execute("CREATE TYPE mood AS ENUM ('sad', 'ok', 'happy')", &[])
272            .await;
273        client
274            .execute(
275                "CREATE TABLE IF NOT EXISTS person(id int PRIMARY KEY, current_mood mood)",
276                &[],
277            )
278            .await
279            .unwrap();
280        client.execute("DELETE FROM person;", &[]).await.unwrap();
281        client
282            .execute("INSERT INTO person VALUES (1, 'happy')", &[])
283            .await
284            .unwrap();
285
286        // test from_sql
287        let got: EnumString = client
288            .query_one("SELECT * FROM person", &[])
289            .await
290            .unwrap()
291            .get::<usize, Option<EnumString>>(1)
292            .unwrap();
293        assert_eq!("happy", got.0.as_str());
294
295        client.execute("DELETE FROM person", &[]).await.unwrap();
296
297        // test to_sql
298        client
299            .execute("INSERT INTO person VALUES (2, $1)", &[&got])
300            .await
301            .unwrap();
302
303        let got_new: EnumString = client
304            .query_one("SELECT * FROM person", &[])
305            .await
306            .unwrap()
307            .get::<usize, Option<EnumString>>(1)
308            .unwrap();
309        assert_eq!("happy", got_new.0.as_str());
310        client.execute("DROP TABLE person", &[]).await.unwrap();
311        client.execute("DROP TYPE mood", &[]).await.unwrap();
312    }
313}
risingwave_connector/parser/postgres.rs

risingwave_connector/parser/
postgres.rs
