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

laminar_connectors/postgres/sink/
statements.rs
