laminar_connectors/postgres/sink/
statements.rs

1//! COPY, UPSERT, DELETE, and table-creation statement construction.
2
3use std::sync::Arc;
4
5use arrow_schema::{Field, SchemaRef};
6
7use crate::error::ConnectorError;
8
9use super::super::sink_config::{quote_sql_identifier, PostgresSinkConfig};
10use super::super::types::{arrow_to_pg_ddl_type, arrow_type_to_pg_array_cast};
11use super::input::{quoted_user_columns, user_fields, validate_sink_schema};
12use super::PostgresSink;
13
14impl PostgresSink {
15    /// Builds the COPY BINARY SQL statement.
16    ///
17    /// ```sql
18    /// COPY "public"."events" ("id", "value", "ts") FROM STDIN BINARY
19    /// ```
20    ///
21    /// # Errors
22    ///
23    /// Returns an error when the schema or sink configuration cannot produce a valid COPY.
24    pub fn build_copy_sql(
25        schema: &SchemaRef,
26        config: &PostgresSinkConfig,
27    ) -> Result<String, ConnectorError> {
28        validate_sink_schema(schema, config)?;
29        let columns = quoted_user_columns(schema);
30        Ok(format!(
31            "COPY {} ({}) FROM STDIN BINARY",
32            config.qualified_table_name(),
33            columns.join(", "),
34        ))
35    }
36
37    /// Builds the UNNEST-based upsert SQL statement.
38    ///
39    /// ```sql
40    /// INSERT INTO "public"."target" ("id", "value", "updated_at")
41    /// SELECT * FROM UNNEST($1::int8[], $2::text[], $3::timestamptz[])
42    /// ON CONFLICT (id) DO UPDATE SET
43    ///     value = EXCLUDED.value,
44    ///     updated_at = EXCLUDED.updated_at
45    /// ```
46    ///
47    /// # Errors
48    ///
49    /// Returns an error when the schema or sink configuration cannot produce a valid upsert.
50    pub fn build_upsert_sql(
51        schema: &SchemaRef,
52        config: &PostgresSinkConfig,
53    ) -> Result<String, ConnectorError> {
54        validate_sink_schema(schema, config)?;
55        let fields = user_fields(schema);
56
57        let columns: Vec<String> = fields
58            .iter()
59            .map(|field| quote_sql_identifier(field.name()))
60            .collect();
61
62        let unnest_params: Vec<String> = fields
63            .iter()
64            .enumerate()
65            .map(|(i, field)| arrow_type_to_pg_array_cast(field.data_type(), i + 1))
66            .collect::<Result<_, _>>()?;
67
68        let non_key_columns: Vec<&Arc<Field>> = fields
69            .iter()
70            .copied()
71            .filter(|field| {
72                !config
73                    .primary_key_columns
74                    .iter()
75                    .any(|primary_key| primary_key == field.name())
76            })
77            .collect();
78
79        let update_clause: Vec<String> = non_key_columns
80            .iter()
81            .map(|field| {
82                let column = quote_sql_identifier(field.name());
83                format!("{column} = EXCLUDED.{column}")
84            })
85            .collect();
86
87        let pk_list = config
88            .primary_key_columns
89            .iter()
90            .map(|column| quote_sql_identifier(column))
91            .collect::<Vec<_>>()
92            .join(", ");
93
94        if update_clause.is_empty() {
95            // Key-only table: use DO NOTHING
96            Ok(format!(
97                "INSERT INTO {} ({}) \
98                 SELECT * FROM UNNEST({}) \
99                 ON CONFLICT ({}) DO NOTHING",
100                config.qualified_table_name(),
101                columns.join(", "),
102                unnest_params.join(", "),
103                pk_list,
104            ))
105        } else {
106            Ok(format!(
107                "INSERT INTO {} ({}) \
108                 SELECT * FROM UNNEST({}) \
109                 ON CONFLICT ({}) DO UPDATE SET {}",
110                config.qualified_table_name(),
111                columns.join(", "),
112                unnest_params.join(", "),
113                pk_list,
114                update_clause.join(", "),
115            ))
116        }
117    }
118
119    /// Builds the DELETE SQL for changelog deletes. One array parameter is bound per primary-key
120    /// column (`$1`, `$2`, …), each holding that column's values for the batch's deleted keys.
121    ///
122    /// ```sql
123    /// -- single PK
124    /// DELETE FROM "public"."events" WHERE "id" = ANY($1::int8[])
125    /// -- composite PK: match tuple-wise via UNNEST, not the cross-product
126    /// DELETE FROM "public"."events" AS "target"
127    ///   USING UNNEST($1::int8[], $2::text[]) AS "keys"("id", "name")
128    /// ```
129    ///
130    /// # Errors
131    ///
132    /// Returns an error when the configured primary key cannot produce a valid delete.
133    pub fn build_delete_sql(
134        schema: &SchemaRef,
135        config: &PostgresSinkConfig,
136    ) -> Result<String, ConnectorError> {
137        validate_sink_schema(schema, config)?;
138        let key_parameter = |column: &str, parameter: usize| -> Result<String, ConnectorError> {
139            let field = schema.field_with_name(column).map_err(|_| {
140                ConnectorError::ConfigurationError(format!(
141                    "primary key column '{column}' is not present in PostgreSQL sink schema"
142                ))
143            })?;
144            arrow_type_to_pg_array_cast(field.data_type(), parameter)
145        };
146        let pk = &config.primary_key_columns;
147
148        // A single PK column can use a plain ANY(); a composite PK must match keys tuple-wise, or
149        // `col1 = ANY($1) AND col2 = ANY($2)` deletes the cross-product — e.g. deleting (1,'a') and
150        // (2,'b') would also delete (1,'b') and (2,'a') (CN-2). UNNEST zips the arrays positionally.
151        if pk.len() <= 1 {
152            let column = pk.first().ok_or_else(|| {
153                ConnectorError::ConfigurationError(
154                    "PostgreSQL changelog delete requires a primary key".into(),
155                )
156            })?;
157            let quoted_column = quote_sql_identifier(column);
158            Ok(format!(
159                "DELETE FROM {} WHERE {quoted_column} = ANY({})",
160                config.qualified_table_name(),
161                key_parameter(column, 1)?,
162            ))
163        } else {
164            let unnest_args: Vec<String> = pk
165                .iter()
166                .enumerate()
167                .map(|(i, column)| key_parameter(column, i + 1))
168                .collect::<Result<_, ConnectorError>>()?;
169            let quoted_keys: Vec<String> = pk
170                .iter()
171                .map(|column| quote_sql_identifier(column))
172                .collect();
173            let target_alias = quote_sql_identifier("target");
174            let key_alias = quote_sql_identifier("keys");
175            let match_conditions: Vec<String> = quoted_keys
176                .iter()
177                .map(|column| format!("{target_alias}.{column} = {key_alias}.{column}"))
178                .collect();
179            Ok(format!(
180                "DELETE FROM {} AS {target_alias} USING UNNEST({}) AS {key_alias}({}) WHERE {}",
181                config.qualified_table_name(),
182                unnest_args.join(", "),
183                quoted_keys.join(", "),
184                match_conditions.join(" AND "),
185            ))
186        }
187    }
188
189    /// Builds CREATE TABLE DDL from the Arrow schema.
190    ///
191    /// ```sql
192    /// CREATE TABLE IF NOT EXISTS "public"."events" (
193    ///     "id" BIGINT NOT NULL,
194    ///     "value" TEXT,
195    ///     "ts" TIMESTAMPTZ,
196    ///     PRIMARY KEY ("id")
197    /// )
198    /// ```
199    ///
200    /// # Errors
201    ///
202    /// Returns an error when an Arrow field or identifier cannot be represented in `PostgreSQL`.
203    pub fn build_create_table_sql(
204        schema: &SchemaRef,
205        config: &PostgresSinkConfig,
206    ) -> Result<String, ConnectorError> {
207        validate_sink_schema(schema, config)?;
208        let fields = user_fields(schema);
209
210        let column_defs: Vec<String> = fields
211            .iter()
212            .map(|field| {
213                let pg_type = arrow_to_pg_ddl_type(field.data_type())?;
214                let nullable = if field.is_nullable() { "" } else { " NOT NULL" };
215                Ok(format!(
216                    "    {} {}{}",
217                    quote_sql_identifier(field.name()),
218                    pg_type,
219                    nullable
220                ))
221            })
222            .collect::<Result<_, ConnectorError>>()?;
223
224        let mut ddl = format!(
225            "CREATE TABLE IF NOT EXISTS {} (\n{}\n",
226            config.qualified_table_name(),
227            column_defs.join(",\n"),
228        );
229
230        if !config.primary_key_columns.is_empty() {
231            use std::fmt::Write;
232            let primary_keys = config
233                .primary_key_columns
234                .iter()
235                .map(|column| quote_sql_identifier(column))
236                .collect::<Vec<_>>()
237                .join(", ");
238            let _ = write!(ddl, ",\n    PRIMARY KEY ({primary_keys})\n");
239        }
240
241        ddl.push(')');
242        Ok(ddl)
243    }
244}
laminar_connectors/postgres/sink/statements.rs

laminar_connectors/postgres/sink/
statements.rs
