add on insert do update for upsert pg sink

Author: kwannoel

e2e_test/sink/postgres_sink.slt

@@ -269,8 +269,62 @@ select * from postgres_query('$PGHOST', '$PGPORT', '$PGUSER', '${PGPASSWORD:post
 1 Varcharvalue1 Textvalue1 123 456 789 12.34 56.78 90.12 t 2023-05-22 12:34:56 2023-05-22 12:34:56 1 day {"key": "value"}
 2 Varcharvalue4 Textvalue2 234 400 890 NaN 67.89 1.23 f 2023-05-23 23:45:01 2023-05-23 23:45:01 2 days {"key": "value2"}
 
+################### test duplicate inserts should work
+
+statement ok
+CREATE TABLE rw_types_table_dup (
+    id BIGINT PRIMARY KEY,
+    varchar_column VARCHAR,
+    text_column TEXT,
+    smallint_column SMALLINT,
+    integer_column INTEGER,
+    bigint_column BIGINT,
+    decimal_column DECIMAL,
+    real_column REAL,
+    double_column DOUBLE PRECISION,
+    boolean_column BOOLEAN,
+    date_column DATE,
+    time_column TIME,
+    interval_column INTERVAL,
+    jsonb_column JSONB,
+    timestamp_column TIMESTAMP
+);
+
+statement ok
+INSERT INTO rw_types_table (id, varchar_column, text_column, integer_column, smallint_column, bigint_column, decimal_column, real_column, double_column, boolean_column, date_column, time_column, timestamp_column, interval_column, jsonb_column) VALUES
+    (1, 'Varcharvalue4', 'Textvalue1', 123, 456, 789, 12.34, 56.78, 90.12, TRUE, '2023-05-22', '12:34:56', '2023-05-22 12:34:56', '1 day', '{"key": "value"}'),
+    (2, 'Varcharvalue6', 'Textvalue2', 234, 567, 890, 'NAN'::decimal, 67.89, 01.23, FALSE, '2023-05-23', '23:45:01', '2023-05-23 23:45:01', '2 days', '{"key": "value2"}'),
+
+statement ok
+flush;
+
+statement ok
+CREATE SINK postgres_rw_types_sink_dup FROM rw_types_table_dup WITH (
+    connector='postgres',
+    host='$PGHOST',
+    port='$PGPORT',
+    user='$PGUSER',
+    password='$PGPASSWORD',
+    database='sink_test',
+    table='pg_types_table',
+    type='upsert',
+    primary_key='id',
+);
+
+query I
+select * from postgres_query('$PGHOST', '$PGPORT', '$PGUSER', '${PGPASSWORD:postgres}', 'sink_test', 'select * from pg_types_table order by id;');
+----
+1 Varcharvalue4 Textvalue1 123 456 789 12.34 56.78 90.12 t 2023-05-22 12:34:56 2023-05-22 12:34:56 1 day {"key": "value"}
+2 Varcharvalue6 Textvalue2 234 567 890 NaN 67.89 1.23 f 2023-05-23 23:45:01 2023-05-23 23:45:01 2 days {"key": "value2"}
+
 ################### cleanup sink
 
+statement ok
+DROP SINK postgres_rw_types_sink_dup;
+
+statement ok
+DROP TABLE rw_types_table_dup;
+
 statement ok
 DROP SINK postgres_rw_types_sink;
 

src/connector/src/sink/postgres.rs

@@ -424,6 +424,7 @@ impl PostgresSinkWriter {
             &self.pk_indices,
             parameters,
             remaining,
+            true,
         )
         .await?;
         transaction.commit().await?;
@@ -463,6 +464,7 @@ impl PostgresSinkWriter {
             &self.pk_indices,
             delete_parameters,
             delete_remaining_parameter,
+            false,
         )
         .await?;
         let (insert_parameters, insert_remaining_parameter) = insert_parameter_buffer.into_parts();
@@ -474,6 +476,7 @@ impl PostgresSinkWriter {
             &self.pk_indices,
             insert_parameters,
             insert_remaining_parameter,
+            false,
         )
         .await?;
         transaction.commit().await?;
@@ -489,6 +492,7 @@ impl PostgresSinkWriter {
         pk_indices: &[usize],
         parameters: Vec<Vec<Option<ScalarAdapter>>>,
         remaining_parameter: Vec<Option<ScalarAdapter>>,
+        append_only: bool,
     ) -> Result<()> {
         let column_length = match op {
             Op::Insert => schema.len(),
@@ -506,7 +510,13 @@ impl PostgresSinkWriter {
                 schema.fields(),
             );
             let statement = match op {
-                Op::Insert => create_insert_sql(schema, table_name, rows_length),
+                Op::Insert => {
+                    if append_only {
+                        create_insert_sql(schema, table_name, rows_length)
+                    } else {
+                        create_upsert_sql(schema, table_name, pk_indices, rows_length)
+                    }
+                }
                 Op::Delete => create_delete_sql(schema, table_name, pk_indices, rows_length),
                 _ => unreachable!(),
             };
@@ -523,7 +533,13 @@ impl PostgresSinkWriter {
                 "flattened parameters are unaligned"
             );
             let statement = match op {
-                Op::Insert => create_insert_sql(schema, table_name, rows_length),
+                Op::Insert => {
+                    if append_only {
+                        create_insert_sql(schema, table_name, rows_length)
+                    } else {
+                        create_upsert_sql(schema, table_name, pk_indices, rows_length)
+                    }
+                }
                 Op::Delete => create_delete_sql(schema, table_name, pk_indices, rows_length),
                 _ => unreachable!(),
             };
@@ -602,6 +618,46 @@ fn create_delete_sql(
     format!("DELETE FROM {table_name} WHERE {pk} in ({parameters})")
 }
 
+fn create_upsert_sql(
+    schema: &Schema,
+    table_name: &str,
+    pk_indices: &[usize],
+    number_of_rows: usize,
+) -> String {
+    let number_of_columns = schema.len();
+    let columns: String = schema
+        .fields()
+        .iter()
+        .map(|field| field.name.clone())
+        .collect_vec()
+        .join(", ");
+    let parameters: String = (0..number_of_rows)
+        .map(|i| {
+            let row_parameters = (0..number_of_columns)
+                .map(|j| format!("${}", i * number_of_columns + j + 1))
+                .join(", ");
+            format!("({row_parameters})")
+        })
+        .collect_vec()
+        .join(", ");
+    let pk_columns = pk_indices
+        .iter()
+        .map(|i| schema.fields()[*i].name.clone())
+        .collect_vec()
+        .join(", ");
+    let update_parameters: String = (0..number_of_columns)
+        .filter(|i| !pk_indices.contains(i))
+        .map(|i| {
+            let column = schema.fields()[i].name.clone();
+            format!("{column} = EXCLUDED.{column}")
+        })
+        .collect_vec()
+        .join(", ");
+    format!(
+        "INSERT INTO {table_name} ({columns}) VALUES {parameters} on conflict ({pk_columns}) do update set {update_parameters}"
+    )
+}
+
 #[cfg(test)]
 mod tests {
     use std::fmt::Display;
@@ -662,4 +718,26 @@ mod tests {
             expect!["DELETE FROM test_table WHERE (a, b) in (($1, $2), ($3, $4), ($5, $6))"],
         );
     }
+
+    #[test]
+    fn test_create_upsert_sql() {
+        let schema = Schema::new(vec![
+            Field {
+                data_type: DataType::Int32,
+                name: "a".to_owned(),
+            },
+            Field {
+                data_type: DataType::Int32,
+                name: "b".to_owned(),
+            },
+        ]);
+        let table_name = "test_table";
+        let sql = create_upsert_sql(&schema, table_name, &[1], 3);
+        check(
+            sql,
+            expect![
+                "INSERT INTO test_table (a, b) VALUES ($1, $2), ($3, $4), ($5, $6) on conflict (b) do update set a = EXCLUDED.a"
+            ],
+        );
+    }
 }