risingwave_connector/connector_common/
postgres.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::collections::{BTreeMap, HashMap};
16use std::fmt;
17
18use anyhow::{Context, anyhow};
19use openssl::ssl::{SslConnector, SslMethod, SslVerifyMode};
20use postgres_openssl::MakeTlsConnector;
21use risingwave_common::bail;
22use risingwave_common::catalog::{ColumnDesc, ColumnId};
23use risingwave_common::types::{DataType, ScalarImpl, StructType};
24use sea_schema::postgres::def::{ColumnType as SeaType, TableDef, TableInfo};
25use sea_schema::postgres::discovery::SchemaDiscovery;
26use sea_schema::sea_query::{Alias, IntoIden};
27use serde::Deserialize;
28use serde_with::{DisplayFromStr, serde_as};
29use sqlx::postgres::{PgConnectOptions, PgSslMode};
30use sqlx::{PgPool, Row};
31use thiserror_ext::AsReport;
32use tokio_postgres::types::Kind as PgKind;
33use tokio_postgres::{Client as PgClient, NoTls};
34
35#[cfg(not(madsim))]
36use super::maybe_tls_connector::MaybeMakeTlsConnector;
37use crate::error::ConnectorResult;
38
39/// SQL query to discover primary key columns directly from PostgreSQL system tables.
40/// This bypasses querying `information_schema.table_constraints` to avoid permission issues.
41/// Match `pg_class` and `pg_namespace` by exact catalog names instead of casting a
42/// constructed string to `regclass`, as unquoted `regclass` input folds mixed-case
43/// table names to lower case.
44const DISCOVER_PRIMARY_KEY_QUERY: &str = r#"
45    SELECT a.attname as column_name
46    FROM pg_index i
47    JOIN pg_class c ON c.oid = i.indrelid
48    JOIN pg_namespace n ON n.oid = c.relnamespace
49    JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
50    WHERE n.nspname = $1
51      AND c.relname = $2
52      AND i.indisprimary = true
53    ORDER BY array_position(i.indkey, a.attnum)
54"#;
55
56/// Discover pgvector columns with both `atttypmod` (dimension) and `format_type` text.
57/// `vector(n)` is stored as `atttypmod = n`, while dimension-less `vector` uses `-1`.
58/// We rely on this to keep user-defined type modifiers that are not preserved by sea-schema.
59const DISCOVER_PGVECTOR_COLUMNS_QUERY: &str = r#"
60    SELECT
61      a.attname as column_name,
62      a.atttypmod as atttypmod,
63      format_type(a.atttypid, a.atttypmod) as formatted_type
64    FROM pg_attribute a
65    JOIN pg_class c ON c.oid = a.attrelid
66    JOIN pg_namespace n ON n.oid = c.relnamespace
67    JOIN pg_type t ON t.oid = a.atttypid
68    WHERE n.nspname = $1
69      AND c.relname = $2
70      AND t.typname = 'vector'
71      AND a.attnum > 0
72      AND NOT a.attisdropped
73    ORDER BY a.attnum
74"#;
75
76/// Canonical Postgres connection parameters shared across sink, source CDC, batch
77/// executor, and frontend `postgres_query` table function. Each caller constructs
78/// this from its own user-facing config struct before invoking the shared helpers
79/// like `create_pg_client` or `PostgresExternalTable::connect`.
80#[derive(Debug, Clone)]
81pub struct PgConnectionConfig {
82    pub host: String,
83    pub port: u16,
84    pub user: String,
85    pub password: String,
86    pub database: String,
87    pub ssl_mode: SslMode,
88    pub ssl_root_cert: Option<String>,
89}
90
91impl PgConnectionConfig {
92    fn to_sqlx_connect_options(&self) -> PgConnectOptions {
93        let mut options = PgConnectOptions::new()
94            .username(&self.user)
95            .password(&self.password)
96            .host(&self.host)
97            .port(self.port)
98            .database(&self.database)
99            .ssl_mode(match self.ssl_mode {
100                SslMode::Disabled => PgSslMode::Disable,
101                SslMode::Preferred => PgSslMode::Prefer,
102                SslMode::Required => PgSslMode::Require,
103                SslMode::VerifyCa => PgSslMode::VerifyCa,
104                SslMode::VerifyFull => PgSslMode::VerifyFull,
105            });
106
107        if matches!(self.ssl_mode, SslMode::VerifyCa | SslMode::VerifyFull)
108            && let Some(root_cert) = &self.ssl_root_cert
109        {
110            options = options.ssl_root_cert(root_cert.as_str());
111        }
112
113        options
114    }
115}
116
117/// TCP keepalive knobs for the long-lived Postgres client used by the sink.
118/// Lives in `connector_common` so both the sink config and the shared
119/// `create_pg_client` helper reference the same definition.
120#[serde_as]
121#[derive(Debug, Clone, Deserialize)]
122pub struct TcpKeepaliveConfig {
123    #[serde(rename = "tcp.keepalive.idle")]
124    #[serde_as(as = "DisplayFromStr")]
125    pub tcp_keepalive_idle: u32,
126    #[serde(rename = "tcp.keepalive.interval")]
127    #[serde_as(as = "DisplayFromStr")]
128    pub tcp_keepalive_interval: u32,
129    #[serde(rename = "tcp.keepalive.count")]
130    #[serde_as(as = "DisplayFromStr")]
131    pub tcp_keepalive_count: u32,
132}
133
134impl Default for TcpKeepaliveConfig {
135    fn default() -> Self {
136        Self {
137            tcp_keepalive_idle: 10 * 60,
138            tcp_keepalive_interval: 10,
139            tcp_keepalive_count: 3,
140        }
141    }
142}
143
144pub fn pg_connection_config_from_properties(
145    props: &BTreeMap<String, String>,
146) -> ConnectorResult<PgConnectionConfig> {
147    Ok(PgConnectionConfig {
148        host: props
149            .get("hostname")
150            .context("missing `hostname` in postgres-cdc properties")?
151            .clone(),
152        port: {
153            let raw = props
154                .get("port")
155                .context("missing `port` in postgres-cdc properties")?;
156            raw.parse::<u16>()
157                .with_context(|| format!("invalid postgres port `{}`", raw))?
158        },
159        user: props
160            .get("username")
161            .context("missing `username` in postgres-cdc properties")?
162            .clone(),
163        password: props.get("password").cloned().unwrap_or_default(),
164        database: props
165            .get("database.name")
166            .context("missing `database.name` in postgres-cdc properties")?
167            .clone(),
168        ssl_mode: props
169            .get("ssl.mode")
170            .and_then(|v| v.parse::<SslMode>().ok())
171            .unwrap_or_default(),
172        ssl_root_cert: props.get("ssl.root.cert").cloned(),
173    })
174}
175
176pub async fn create_pg_client_from_properties(
177    props: &BTreeMap<String, String>,
178    tcp_keepalive: Option<TcpKeepaliveConfig>,
179) -> ConnectorResult<PgClient> {
180    let config = pg_connection_config_from_properties(props)?;
181    create_pg_client(&config, tcp_keepalive)
182        .await
183        .map_err(Into::into)
184}
185
186pub async fn discover_pgvector_dimensions(
187    client: &PgClient,
188    schema: &str,
189    table: &str,
190) -> ConnectorResult<HashMap<String, usize>> {
191    let rows = client
192        .query(DISCOVER_PGVECTOR_COLUMNS_QUERY, &[&schema, &table])
193        .await?;
194
195    let mut dims = HashMap::new();
196    for row in rows {
197        let col_name: String = row.get("column_name");
198        let atttypmod: i32 = row.get("atttypmod");
199        if atttypmod > 0
200            && let Ok(dim) = usize::try_from(atttypmod)
201        {
202            dims.insert(col_name, dim);
203        }
204    }
205    Ok(dims)
206}
207
208#[derive(Debug, Clone, PartialEq, Deserialize, Default)]
209#[serde(rename_all = "lowercase")]
210pub enum SslMode {
211    #[serde(alias = "disable")]
212    Disabled,
213    #[serde(alias = "prefer")]
214    #[default]
215    Preferred,
216    #[serde(alias = "require")]
217    Required,
218    /// verify that the server is trustworthy by checking the certificate chain
219    /// up to the root certificate stored on the client.
220    #[serde(alias = "verify-ca")]
221    VerifyCa,
222    /// Besides verify the certificate, will also verify that the serverhost name
223    /// matches the name stored in the server certificate.
224    #[serde(alias = "verify-full")]
225    VerifyFull,
226}
227
228pub struct PostgresExternalTable {
229    column_descs: Vec<ColumnDesc>,
230    pk_names: Vec<String>,
231}
232
233impl PostgresExternalTable {
234    /// Discover primary key columns directly from PostgreSQL system tables.
235    /// This bypasses querying `information_schema.table_constraints` to avoid requiring table owner permissions.
236    async fn discover_primary_key(
237        connection: &PgPool,
238        schema_name: &str,
239        table_name: &str,
240    ) -> ConnectorResult<Vec<String>> {
241        let rows = sqlx::query(DISCOVER_PRIMARY_KEY_QUERY)
242            .bind(schema_name)
243            .bind(table_name)
244            .fetch_all(connection)
245            .await
246            .context("Failed to discover primary key columns")?;
247
248        let pk_columns = rows
249            .into_iter()
250            .map(|row| row.get::<String, _>("column_name"))
251            .collect();
252
253        Ok(pk_columns)
254    }
255
256    /// Discover schema with workaround for primary key discovery
257    /// This method uses direct PostgreSQL system table queries for primary keys
258    /// to avoid permission issues when querying `information_schema.table_constraints`
259    async fn discover_pk_and_full_columns(
260        config: &PgConnectionConfig,
261        schema: &str,
262        table: &str,
263    ) -> ConnectorResult<(Vec<sea_schema::postgres::def::ColumnInfo>, Vec<String>)> {
264        let options = config.to_sqlx_connect_options();
265        let connection = PgPool::connect_with(options).await?;
266
267        // Use sea-schema only for column discovery (no permission issues)
268        let schema_discovery = SchemaDiscovery::new(connection.clone(), schema);
269        let empty_map: HashMap<String, Vec<String>> = HashMap::new();
270        let columns = schema_discovery
271            .discover_columns(
272                Alias::new(schema).into_iden(),
273                Alias::new(table).into_iden(),
274                &empty_map,
275            )
276            .await?;
277
278        let pgvector_columns = sqlx::query(DISCOVER_PGVECTOR_COLUMNS_QUERY)
279            .bind(schema)
280            .bind(table)
281            .fetch_all(&connection)
282            .await
283            .context("Failed to discover PostgreSQL pgvector columns")?;
284        let formatted_type_by_column: HashMap<String, String> = pgvector_columns
285            .into_iter()
286            .map(|row| {
287                (
288                    row.get::<String, _>("column_name"),
289                    row.get::<String, _>("formatted_type"),
290                )
291            })
292            .collect();
293
294        // sea-schema reports pgvector as `Unknown("vector")` and drops the dimension.
295        // Patch it with PostgreSQL's formatted type text so we can derive vector(n).
296        let mut columns = columns;
297        for col in &mut columns {
298            if let SeaType::Unknown(name) = &col.col_type
299                && name.eq_ignore_ascii_case("vector")
300                && let Some(formatted_type) = formatted_type_by_column.get(&col.name)
301            {
302                col.col_type = SeaType::Unknown(formatted_type.clone());
303            }
304        }
305
306        // Use direct system table query for primary key discovery
307        let pk_columns = Self::discover_primary_key(&connection, schema, table).await?;
308
309        Ok((columns, pk_columns))
310    }
311
312    async fn discover_schema(
313        config: &PgConnectionConfig,
314        schema: &str,
315        table: &str,
316    ) -> ConnectorResult<TableDef> {
317        let options = config.to_sqlx_connect_options();
318        let connection = PgPool::connect_with(options).await?;
319        let schema_discovery = SchemaDiscovery::new(connection, schema);
320        // fetch column schema and primary key
321        let empty_map = HashMap::new();
322        let table_schema = schema_discovery
323            .discover_table(
324                TableInfo {
325                    name: table.to_owned(),
326                    of_type: None,
327                },
328                &empty_map,
329            )
330            .await?;
331        Ok(table_schema)
332    }
333
334    pub async fn connect(
335        config: &PgConnectionConfig,
336        schema: &str,
337        table: &str,
338        is_append_only: bool,
339    ) -> ConnectorResult<Self> {
340        tracing::debug!("connect to postgres external table");
341
342        let (columns, pk_names) = Self::discover_pk_and_full_columns(config, schema, table).await?;
343
344        let mut column_descs = vec![];
345        for col in &columns {
346            let rw_data_type = sea_type_to_rw_type(&col.col_type)?;
347            let column_desc = if let Some(ref default_expr) = col.default {
348                // parse the value of "column_default" field in information_schema.columns,
349                // non number data type will be stored as "'value'::type"
350                let val_text = default_expr
351                    .0
352                    .split("::")
353                    .map(|s| s.trim_matches('\''))
354                    .next()
355                    .expect("default value expression");
356
357                match ScalarImpl::from_text(val_text, &rw_data_type) {
358                    Ok(scalar) => ColumnDesc::named_with_default_value(
359                        col.name.clone(),
360                        ColumnId::placeholder(),
361                        rw_data_type.clone(),
362                        Some(scalar),
363                    ),
364                    Err(err) => {
365                        tracing::warn!(error=%err.as_report(), "failed to parse postgres default value expression, only constant is supported");
366                        ColumnDesc::named(col.name.clone(), ColumnId::placeholder(), rw_data_type)
367                    }
368                }
369            } else {
370                ColumnDesc::named(col.name.clone(), ColumnId::placeholder(), rw_data_type)
371            };
372            column_descs.push(column_desc);
373        }
374
375        // Check primary key existence using the directly discovered pk_names
376        if !is_append_only && pk_names.is_empty() {
377            return Err(anyhow!(
378                "Postgres table should define the primary key for non-append-only tables"
379            )
380            .into());
381        }
382
383        Ok(Self {
384            column_descs,
385            pk_names,
386        })
387    }
388
389    // return the mapping from column name to pg type, the pg type is used for writing data to postgres
390    pub async fn type_mapping(
391        config: &PgConnectionConfig,
392        schema: &str,
393        table: &str,
394        is_append_only: bool,
395    ) -> ConnectorResult<HashMap<String, tokio_postgres::types::Type>> {
396        tracing::debug!("connect to postgres external table to get type mapping");
397        let table_schema = Self::discover_schema(config, schema, table).await?;
398        let mut column_name_to_pg_type = HashMap::new();
399        for col in &table_schema.columns {
400            let pg_type = sea_type_to_pg_type(&col.col_type)?;
401            column_name_to_pg_type.insert(col.name.clone(), pg_type);
402        }
403        if !is_append_only && table_schema.primary_key_constraints.is_empty() {
404            return Err(anyhow!(
405                "Postgres table should define the primary key for non-append-only tables"
406            )
407            .into());
408        }
409        Ok(column_name_to_pg_type)
410    }
411
412    pub fn column_descs(&self) -> &Vec<ColumnDesc> {
413        &self.column_descs
414    }
415
416    pub fn pk_names(&self) -> &Vec<String> {
417        &self.pk_names
418    }
419}
420
421impl fmt::Display for SslMode {
422    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
423        f.write_str(match self {
424            SslMode::Disabled => "disabled",
425            SslMode::Preferred => "preferred",
426            SslMode::Required => "required",
427            SslMode::VerifyCa => "verify-ca",
428            SslMode::VerifyFull => "verify-full",
429        })
430    }
431}
432
433impl std::str::FromStr for SslMode {
434    type Err = serde_json::Error;
435
436    fn from_str(s: &str) -> Result<Self, Self::Err> {
437        serde_json::from_value(serde_json::Value::String(s.to_owned()))
438    }
439}
440
441pub async fn create_pg_client(
442    config: &PgConnectionConfig,
443    tcp_keepalive: Option<TcpKeepaliveConfig>,
444) -> anyhow::Result<PgClient> {
445    let mut pg_config = tokio_postgres::Config::new();
446    pg_config
447        .user(&config.user)
448        .password(&config.password)
449        .host(&config.host)
450        .port(config.port)
451        .dbname(&config.database);
452
453    // Configure TCP keepalive if provided
454    if let Some(keepalive) = tcp_keepalive {
455        pg_config.keepalives(true);
456        pg_config.keepalives_idle(std::time::Duration::from_secs(
457            keepalive.tcp_keepalive_idle as u64,
458        ));
459        #[cfg(not(target_os = "windows"))]
460        {
461            pg_config.keepalives_interval(std::time::Duration::from_secs(
462                keepalive.tcp_keepalive_interval as u64,
463            ));
464            pg_config.keepalives_retries(keepalive.tcp_keepalive_count);
465        }
466        tracing::info!(
467            "TCP keepalive enabled: idle={}s, interval={}s, retries={}",
468            keepalive.tcp_keepalive_idle,
469            keepalive.tcp_keepalive_interval,
470            keepalive.tcp_keepalive_count
471        );
472    }
473
474    let verify_hostname = matches!(config.ssl_mode, SslMode::VerifyFull);
475
476    #[cfg(not(madsim))]
477    let connector = match config.ssl_mode {
478        SslMode::Disabled => {
479            pg_config.ssl_mode(tokio_postgres::config::SslMode::Disable);
480            MaybeMakeTlsConnector::NoTls(NoTls)
481        }
482        SslMode::Preferred => {
483            pg_config.ssl_mode(tokio_postgres::config::SslMode::Prefer);
484            match SslConnector::builder(SslMethod::tls()) {
485                Ok(mut builder) => {
486                    // disable certificate verification for `prefer`
487                    builder.set_verify(SslVerifyMode::NONE);
488                    MaybeMakeTlsConnector::Tls(MakeTlsConnector::new(builder.build()))
489                }
490                Err(e) => {
491                    tracing::warn!(error = %e.as_report(), "SSL connector error");
492                    MaybeMakeTlsConnector::NoTls(NoTls)
493                }
494            }
495        }
496        SslMode::Required => {
497            pg_config.ssl_mode(tokio_postgres::config::SslMode::Require);
498            let mut builder = SslConnector::builder(SslMethod::tls())?;
499            // disable certificate verification for `require`
500            builder.set_verify(SslVerifyMode::NONE);
501            MaybeMakeTlsConnector::Tls(MakeTlsConnector::new(builder.build()))
502        }
503
504        SslMode::VerifyCa | SslMode::VerifyFull => {
505            pg_config.ssl_mode(tokio_postgres::config::SslMode::Require);
506            let mut builder = SslConnector::builder(SslMethod::tls())?;
507            if let Some(ssl_root_cert) = &config.ssl_root_cert {
508                builder.set_ca_file(ssl_root_cert).map_err(|e| {
509                    anyhow!(format!("bad ssl root cert error: {}", e.to_report_string()))
510                })?;
511            }
512            let mut connector = MakeTlsConnector::new(builder.build());
513            if !verify_hostname {
514                connector.set_callback(|c, _| {
515                    c.set_verify_hostname(false);
516                    Ok(())
517                });
518            }
519            MaybeMakeTlsConnector::Tls(connector)
520        }
521    };
522    #[cfg(madsim)]
523    let connector = NoTls;
524
525    let (client, connection) = pg_config.connect(connector).await?;
526
527    tokio::spawn(async move {
528        if let Err(e) = connection.await {
529            tracing::error!(error = %e.as_report(), "postgres connection error");
530        }
531    });
532
533    Ok(client)
534}
535
536// Used for both source and sink connector
537pub fn sea_type_to_rw_type(col_type: &SeaType) -> ConnectorResult<DataType> {
538    let dtype = match col_type {
539        SeaType::SmallInt | SeaType::SmallSerial => DataType::Int16,
540        SeaType::Integer | SeaType::Serial => DataType::Int32,
541        SeaType::BigInt | SeaType::BigSerial => DataType::Int64,
542        SeaType::Money | SeaType::Decimal(_) | SeaType::Numeric(_) => DataType::Decimal,
543        SeaType::Real => DataType::Float32,
544        SeaType::DoublePrecision => DataType::Float64,
545        SeaType::Varchar(_) | SeaType::Char(_) | SeaType::Text => DataType::Varchar,
546        SeaType::Bytea => DataType::Bytea,
547        SeaType::Timestamp(_) => DataType::Timestamp,
548        SeaType::TimestampWithTimeZone(_) => DataType::Timestamptz,
549        SeaType::Date => DataType::Date,
550        SeaType::Time(_) | SeaType::TimeWithTimeZone(_) => DataType::Time,
551        SeaType::Interval(_) => DataType::Interval,
552        SeaType::Boolean => DataType::Boolean,
553        SeaType::Point => DataType::Struct(StructType::new(vec![
554            ("x", DataType::Float32),
555            ("y", DataType::Float32),
556        ])),
557        SeaType::Uuid => DataType::Varchar,
558        SeaType::Xml => DataType::Varchar,
559        SeaType::Json => DataType::Jsonb,
560        SeaType::JsonBinary => DataType::Jsonb,
561        SeaType::Array(def) => {
562            let item_type = match def.col_type.as_ref() {
563                Some(ty) => sea_type_to_rw_type(ty.as_ref())?,
564                None => {
565                    return Err(anyhow!("ARRAY type missing element type").into());
566                }
567            };
568
569            DataType::list(item_type)
570        }
571        SeaType::PgLsn => DataType::Int64,
572        SeaType::Cidr
573        | SeaType::Inet
574        | SeaType::MacAddr
575        | SeaType::MacAddr8
576        | SeaType::Int4Range
577        | SeaType::Int8Range
578        | SeaType::NumRange
579        | SeaType::TsRange
580        | SeaType::TsTzRange
581        | SeaType::DateRange
582        | SeaType::Enum(_) => DataType::Varchar,
583        SeaType::Line
584        | SeaType::Lseg
585        | SeaType::Box
586        | SeaType::Path
587        | SeaType::Polygon
588        | SeaType::Circle
589        | SeaType::Bit(_)
590        | SeaType::VarBit(_)
591        | SeaType::TsVector
592        | SeaType::TsQuery => {
593            bail!("{:?} type not supported", col_type);
594        }
595        SeaType::Unknown(name) => {
596            if let Some(dim) = parse_pgvector_dimension(name)? {
597                DataType::Vector(dim)
598            } else {
599                // NOTES: user-defined enum type is classified as `Unknown`
600                tracing::warn!("Unknown Postgres data type: {name}, map to varchar");
601                DataType::Varchar
602            }
603        }
604    };
605
606    Ok(dtype)
607}
608
609fn parse_pgvector_dimension(type_name: &str) -> ConnectorResult<Option<usize>> {
610    let normalized = type_name.trim().to_ascii_lowercase();
611    if normalized == "vector" {
612        bail!("pgvector type `vector` is missing dimension, expected `vector(n)`")
613    }
614    if !normalized.starts_with("vector(") || !normalized.ends_with(')') {
615        return Ok(None);
616    }
617
618    let dim_text = normalized
619        .trim_start_matches("vector(")
620        .trim_end_matches(')')
621        .trim();
622    let dim = dim_text
623        .parse::<usize>()
624        .map_err(|_| anyhow!("invalid pgvector dimension in type `{type_name}`"))?;
625
626    if !(1..=DataType::VEC_MAX_SIZE).contains(&dim) {
627        bail!(
628            "pgvector dimension out of range in type `{}`: expect 1..={}",
629            type_name,
630            DataType::VEC_MAX_SIZE
631        );
632    }
633
634    Ok(Some(dim))
635}
636
637// Used for sink connector
638// We use `sea-schema` for table schema discovery.
639// So we have to map `sea-schema` pg types
640// to `tokio-postgres` pg types (which we use for query binding).
641fn sea_type_to_pg_type(sea_type: &SeaType) -> ConnectorResult<tokio_postgres::types::Type> {
642    use tokio_postgres::types::Type as PgType;
643    match sea_type {
644        SeaType::SmallInt => Ok(PgType::INT2),
645        SeaType::Integer => Ok(PgType::INT4),
646        SeaType::BigInt => Ok(PgType::INT8),
647        SeaType::Decimal(_) => Ok(PgType::NUMERIC),
648        SeaType::Numeric(_) => Ok(PgType::NUMERIC),
649        SeaType::Real => Ok(PgType::FLOAT4),
650        SeaType::DoublePrecision => Ok(PgType::FLOAT8),
651        SeaType::Varchar(_) => Ok(PgType::VARCHAR),
652        SeaType::Char(_) => Ok(PgType::CHAR),
653        SeaType::Text => Ok(PgType::TEXT),
654        SeaType::Bytea => Ok(PgType::BYTEA),
655        SeaType::Timestamp(_) => Ok(PgType::TIMESTAMP),
656        SeaType::TimestampWithTimeZone(_) => Ok(PgType::TIMESTAMPTZ),
657        SeaType::Date => Ok(PgType::DATE),
658        SeaType::Time(_) => Ok(PgType::TIME),
659        SeaType::TimeWithTimeZone(_) => Ok(PgType::TIMETZ),
660        SeaType::Interval(_) => Ok(PgType::INTERVAL),
661        SeaType::Boolean => Ok(PgType::BOOL),
662        SeaType::Point => Ok(PgType::POINT),
663        SeaType::Uuid => Ok(PgType::UUID),
664        SeaType::Json => Ok(PgType::JSON),
665        SeaType::JsonBinary => Ok(PgType::JSONB),
666        SeaType::Array(t) => {
667            let Some(t) = t.col_type.as_ref() else {
668                bail!("missing array type")
669            };
670            match t.as_ref() {
671                // RW only supports 1 level of nesting.
672                SeaType::SmallInt => Ok(PgType::INT2_ARRAY),
673                SeaType::Integer => Ok(PgType::INT4_ARRAY),
674                SeaType::BigInt => Ok(PgType::INT8_ARRAY),
675                SeaType::Decimal(_) => Ok(PgType::NUMERIC_ARRAY),
676                SeaType::Numeric(_) => Ok(PgType::NUMERIC_ARRAY),
677                SeaType::Real => Ok(PgType::FLOAT4_ARRAY),
678                SeaType::DoublePrecision => Ok(PgType::FLOAT8_ARRAY),
679                SeaType::Varchar(_) => Ok(PgType::VARCHAR_ARRAY),
680                SeaType::Char(_) => Ok(PgType::CHAR_ARRAY),
681                SeaType::Text => Ok(PgType::TEXT_ARRAY),
682                SeaType::Bytea => Ok(PgType::BYTEA_ARRAY),
683                SeaType::Timestamp(_) => Ok(PgType::TIMESTAMP_ARRAY),
684                SeaType::TimestampWithTimeZone(_) => Ok(PgType::TIMESTAMPTZ_ARRAY),
685                SeaType::Date => Ok(PgType::DATE_ARRAY),
686                SeaType::Time(_) => Ok(PgType::TIME_ARRAY),
687                SeaType::TimeWithTimeZone(_) => Ok(PgType::TIMETZ_ARRAY),
688                SeaType::Interval(_) => Ok(PgType::INTERVAL_ARRAY),
689                SeaType::Boolean => Ok(PgType::BOOL_ARRAY),
690                SeaType::Point => Ok(PgType::POINT_ARRAY),
691                SeaType::Uuid => Ok(PgType::UUID_ARRAY),
692                SeaType::Json => Ok(PgType::JSON_ARRAY),
693                SeaType::JsonBinary => Ok(PgType::JSONB_ARRAY),
694                SeaType::Array(_) => bail!("nested array type is not supported"),
695                SeaType::Unknown(name) => {
696                    // Treat as enum type
697                    Ok(PgType::new(
698                        name.clone(),
699                        0,
700                        PgKind::Array(PgType::new(
701                            name.clone(),
702                            0,
703                            PgKind::Enum(vec![]),
704                            "".into(),
705                        )),
706                        "".into(),
707                    ))
708                }
709                _ => bail!("unsupported array type: {:?}", t),
710            }
711        }
712        SeaType::Unknown(name) => {
713            // Treat as enum type
714            Ok(PgType::new(
715                name.clone(),
716                0,
717                PgKind::Enum(vec![]),
718                "".into(),
719            ))
720        }
721        _ => bail!("unsupported type: {:?}", sea_type),
722    }
723}
724
725#[cfg(test)]
726mod tests {
727    use super::parse_pgvector_dimension;
728
729    #[test]
730    fn test_parse_pgvector_dimension() {
731        assert_eq!(parse_pgvector_dimension("vector(3)").unwrap(), Some(3));
732        assert_eq!(parse_pgvector_dimension("VECTOR(768)").unwrap(), Some(768));
733        assert_eq!(parse_pgvector_dimension("varchar").unwrap(), None);
734    }
735
736    #[test]
737    fn test_parse_pgvector_dimension_requires_size() {
738        let err = parse_pgvector_dimension("vector").unwrap_err();
739        assert!(err.to_string().contains("missing dimension"));
740    }
741}
risingwave_connector/connector_common/postgres.rs

risingwave_connector/connector_common/
postgres.rs
