Skip to main content

faucet_sink_postgres/
sink.rs

1//! PostgreSQL sink implementation.
2
3use crate::config::{PostgresColumnMapping, PostgresSinkConfig};
4use async_trait::async_trait;
5use faucet_core::FaucetError;
6use faucet_core::util::quote_ident;
7use serde_json::Value;
8use sqlx::postgres::PgPoolOptions;
9use sqlx::{PgPool, Row};
10
11/// Render a JSON value as the text to bind for a PostgreSQL column whose
12/// underlying type is `udt` (`information_schema.columns.udt_name`), or `None`
13/// for SQL `NULL`.
14///
15/// The accompanying placeholder is emitted as `$N::<udt>`, so PostgreSQL runs
16/// the destination column type's input function over this text. That makes
17/// `string → timestamptz/uuid/date`, `number → int4/numeric/float8`,
18/// `bool → bool`, and `json → jsonb` all work — instead of binding every value
19/// as `serde_json::Value` (which sqlx encodes as `jsonb`, so an insert into any
20/// non-`jsonb` column fails at runtime with *"column is of type … but
21/// expression is of type jsonb"*; this was the C1 bug in audit #146).
22///
23/// For `json`/`jsonb` columns the value is bound as its JSON text (so a string
24/// keeps its quotes and objects/arrays round-trip); the `::jsonb` cast then
25/// parses it. For every other type the scalar's plain text form is bound and
26/// the column's input function parses it via the cast.
27fn pg_bind_text(value: Option<&Value>, udt: &str) -> Option<String> {
28    match value {
29        None | Some(Value::Null) => None,
30        Some(v) => {
31            if udt.eq_ignore_ascii_case("json") || udt.eq_ignore_ascii_case("jsonb") {
32                Some(v.to_string())
33            } else {
34                match v {
35                    Value::Bool(b) => Some(b.to_string()),
36                    Value::Number(n) => Some(n.to_string()),
37                    Value::String(s) => Some(s.clone()),
38                    // Arrays/objects have no scalar text form for a non-JSON
39                    // column; bind their JSON text so the `::<type>` cast fails
40                    // loudly rather than silently coercing.
41                    other => Some(other.to_string()),
42                }
43            }
44        }
45    }
46}
47
48/// Build the SQL relation reference for the configured table, optionally
49/// schema-qualified.
50///
51/// Both the AutoMap column-discovery probe and the `INSERT` statements use this
52/// single helper, so column discovery is always scoped to the *exact* relation
53/// the `INSERT` targets (#146 M13). With no schema the bare quoted table name
54/// resolves against the connection's `search_path`; with a schema it becomes
55/// `"schema"."table"`, pinning both discovery and insert to that namespace —
56/// otherwise a table of the same name in another schema pollutes the
57/// AutoMap column set (duplicate / wrong columns).
58fn qualified_table_ref(schema: Option<&str>, table: &str) -> String {
59    match schema {
60        Some(s) => format!("{}.{}", quote_ident(s), quote_ident(table)),
61        None => quote_ident(table),
62    }
63}
64
65/// Build the `ON CONFLICT (key) DO UPDATE …` tail for an upsert INSERT.
66/// Non-key columns are SET from EXCLUDED. If every column is a key column,
67/// emit `DO NOTHING`.
68fn on_conflict_clause(key: &[String], all_cols: &[String]) -> String {
69    let key_list = key
70        .iter()
71        .map(|k| quote_ident(k))
72        .collect::<Vec<_>>()
73        .join(", ");
74    let updates: Vec<String> = all_cols
75        .iter()
76        .filter(|c| !key.iter().any(|k| k == *c))
77        .map(|c| format!("{q} = EXCLUDED.{q}", q = quote_ident(c)))
78        .collect();
79    if updates.is_empty() {
80        format!("ON CONFLICT ({key_list}) DO NOTHING")
81    } else {
82        format!(
83            "ON CONFLICT ({key_list}) DO UPDATE SET {}",
84            updates.join(", ")
85        )
86    }
87}
88
89/// A sink that writes JSON records to a PostgreSQL table.
90pub struct PostgresSink {
91    config: PostgresSinkConfig,
92    pool: PgPool,
93}
94
95impl PostgresSink {
96    /// Create a new PostgreSQL sink. Establishes a connection pool.
97    pub async fn new(config: PostgresSinkConfig) -> Result<Self, FaucetError> {
98        config.write.validate()?;
99        if !matches!(config.write.write_mode, faucet_core::WriteMode::Append)
100            && !matches!(config.column_mapping, PostgresColumnMapping::AutoMap)
101        {
102            return Err(FaucetError::Config(
103                "postgres sink: write_mode upsert/delete requires column_mapping: auto_map \
104                 (key columns must be real columns, not inside a JSONB blob)"
105                    .into(),
106            ));
107        }
108
109        let pool = PgPoolOptions::new()
110            .max_connections(config.max_connections)
111            .connect(&config.connection_url)
112            .await
113            .map_err(|e| FaucetError::Sink(format!("PostgreSQL connection failed: {e}")))?;
114
115        Ok(Self { config, pool })
116    }
117
118    /// Insert a batch of records using JSONB column mode, on the given connection.
119    ///
120    /// Accepts `&mut sqlx::PgConnection` so the same logic runs both standalone
121    /// (via a pool-acquired connection, autocommit) and inside the idempotent
122    /// transaction (where `&mut *tx` is passed — `Transaction<'_, Postgres>`
123    /// derefs to `PgConnection`).
124    async fn insert_jsonb(
125        &self,
126        conn: &mut sqlx::PgConnection,
127        records: &[Value],
128        column: &str,
129    ) -> Result<usize, FaucetError> {
130        if records.is_empty() {
131            return Ok(0);
132        }
133
134        // Use a single INSERT with unnest for efficiency.
135        let json_values: Vec<serde_json::Value> = records.to_vec();
136        let query = format!(
137            "INSERT INTO {} ({}) SELECT * FROM unnest($1::jsonb[])",
138            qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name),
139            quote_ident(column)
140        );
141
142        sqlx::query(&query)
143            .bind(json_values)
144            .execute(&mut *conn)
145            .await
146            .map_err(|e| FaucetError::Sink(format!("PostgreSQL insert failed: {e}")))?;
147
148        Ok(records.len())
149    }
150
151    /// Insert a batch of records using auto-mapped columns, on the given connection.
152    ///
153    /// Accepts `&mut sqlx::PgConnection` so the same logic runs both standalone
154    /// (via a pool-acquired connection, autocommit) and inside the idempotent
155    /// transaction (where `&mut *tx` is passed — `Transaction<'_, Postgres>`
156    /// derefs to `PgConnection`). Running the column-discovery query on the same
157    /// connection is harmless and avoids any cross-connection visibility surprise.
158    ///
159    /// Discovers each column's name *and* underlying type (`udt_name`) from the
160    /// table schema and maps top-level JSON fields to columns. Each placeholder
161    /// is emitted as `$N::<udt>` and the value is bound as text (see
162    /// [`pg_bind_text`]), so the destination column's input function parses it —
163    /// numbers, booleans, timestamps, uuids, and JSON all land in their native
164    /// column types. (Previously every value was bound as `serde_json::Value`,
165    /// which sqlx encodes as `jsonb`, so an insert into any non-`jsonb` column
166    /// failed at runtime — audit #146 C1.) Uses a single multi-row INSERT
167    /// (sub-chunked at the 65535-parameter cap) for efficiency.
168    ///
169    /// When `conflict_key` is `Some(key)`, each sub-chunk's INSERT is given an
170    /// `ON CONFLICT (key) DO UPDATE …` tail so it upserts by the key columns
171    /// (last-write-wins within the batch is already handled by the planner's
172    /// dedup, so a single sub-chunk never double-hits the same conflict target).
173    async fn insert_auto_map_with_conflict(
174        &self,
175        conn: &mut sqlx::PgConnection,
176        records: &[Value],
177        conflict_key: Option<&[String]>,
178    ) -> Result<usize, FaucetError> {
179        if records.is_empty() {
180            return Ok(0);
181        }
182
183        // Get column names AND their underlying types for the *exact* relation
184        // the INSERT will target. Scoping by `to_regclass(<qualified ref>)`
185        // resolves the relation the same way the INSERT does — by the configured
186        // schema if set, otherwise by the connection's `search_path` — so a
187        // table of the same name in another schema can no longer pollute the
188        // column set with duplicate/wrong columns (#146 M13). The previous query
189        // filtered `information_schema.columns` by `table_name` alone (no schema
190        // predicate), merging every same-named table across all schemas.
191        //
192        // `pg_type.typname` is the concrete type (`int4`, `timestamptz`,
193        // `numeric`, `jsonb`, `uuid`, `text`, …) — identical to the old
194        // `information_schema.columns.udt_name` — used as the per-placeholder
195        // cast target below. `::text` casts the `name`-typed catalog columns so
196        // sqlx decodes them as `String`.
197        let table_ref = qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name);
198        let columns: Vec<(String, String)> = sqlx::query(
199            "SELECT a.attname::text AS column_name, t.typname::text AS udt_name \
200             FROM pg_catalog.pg_attribute a \
201             JOIN pg_catalog.pg_type t ON t.oid = a.atttypid \
202             WHERE a.attrelid = to_regclass($1)::oid \
203               AND a.attnum > 0 AND NOT a.attisdropped \
204             ORDER BY a.attnum",
205        )
206        .bind(&table_ref)
207        .fetch_all(&mut *conn)
208        .await
209        .map_err(|e| FaucetError::Sink(format!("failed to query table columns: {e}")))?
210        .iter()
211        .map(|row| {
212            (
213                row.get::<String, _>("column_name"),
214                row.get::<String, _>("udt_name"),
215            )
216        })
217        .collect();
218
219        if columns.is_empty() {
220            return Err(FaucetError::Sink(format!(
221                "table {table_ref} has no columns or does not exist"
222            )));
223        }
224
225        // Pre-validate all records and collect matched (column, udt, value)
226        // triples per record. The INSERT column set is the UNION of table
227        // columns present in ANY record (in declared table order), not just the
228        // first record's keys — otherwise a field present only in a later record
229        // of the batch would be silently dropped (audit #146 H1). A row missing
230        // a unioned column binds SQL NULL.
231        let mut matched_rows: Vec<Vec<(&String, &String, &Value)>> =
232            Vec::with_capacity(records.len());
233        let mut used: std::collections::HashSet<&str> = std::collections::HashSet::new();
234
235        for record in records {
236            let obj = record
237                .as_object()
238                .ok_or_else(|| FaucetError::Sink("AutoMap requires JSON object records".into()))?;
239
240            let matching: Vec<(&String, &String, &Value)> = columns
241                .iter()
242                .filter_map(|(col, udt)| obj.get(col).map(|v| (col, udt, v)))
243                .collect();
244
245            if matching.is_empty() {
246                tracing::warn!(
247                    record_keys = ?obj.keys().collect::<Vec<_>>(),
248                    table_columns = ?columns,
249                    "record has no keys matching table columns, skipping"
250                );
251                continue;
252            }
253
254            for (c, _, _) in &matching {
255                used.insert(c.as_str());
256            }
257            matched_rows.push(matching);
258        }
259
260        if matched_rows.is_empty() {
261            return Ok(0);
262        }
263
264        // Table columns (in declared order, with their udt) present in at least
265        // one record.
266        let insert_columns: Vec<(String, String)> = columns
267            .iter()
268            .filter(|(c, _)| used.contains(c.as_str()))
269            .cloned()
270            .collect();
271
272        let num_cols = insert_columns.len();
273        let num_rows = matched_rows.len();
274        let col_names: Vec<String> = insert_columns.iter().map(|(c, _)| quote_ident(c)).collect();
275
276        // PostgreSQL caps bind parameters per statement at 65535. A multi-row
277        // INSERT binds `rows × num_cols` parameters, so a wide table at a large
278        // batch_size can exceed it and fail at runtime (#78/#21). Split into
279        // sub-INSERTs of at most floor(MAX_PARAMS / num_cols) rows.
280        const MAX_PG_PARAMS: usize = 65535;
281        let max_rows_per_insert = (MAX_PG_PARAMS / num_cols).max(1);
282
283        for sub in matched_rows.chunks(max_rows_per_insert) {
284            // Build multi-row VALUES clause with per-column casts so the column
285            // type's input function parses the bound text:
286            //   ($1::int4, $2::timestamptz), ($3::int4, $4::timestamptz), ...
287            let mut value_tuples: Vec<String> = Vec::with_capacity(sub.len());
288            for row_idx in 0..sub.len() {
289                let start = row_idx * num_cols + 1;
290                let placeholders: Vec<String> = (0..num_cols)
291                    .map(|c| format!("${}::{}", start + c, insert_columns[c].1))
292                    .collect();
293                value_tuples.push(format!("({})", placeholders.join(", ")));
294            }
295
296            let query = format!(
297                "INSERT INTO {} ({}) VALUES {}",
298                table_ref,
299                col_names.join(", "),
300                value_tuples.join(", ")
301            );
302            let query = match conflict_key {
303                Some(key) => format!(
304                    "{query} {}",
305                    on_conflict_clause(
306                        key,
307                        &insert_columns
308                            .iter()
309                            .map(|(c, _)| c.clone())
310                            .collect::<Vec<_>>()
311                    )
312                ),
313                None => query,
314            };
315
316            let mut q = sqlx::query(&query);
317            for matched in sub {
318                // Bind values in the fixed column order, as text matching each
319                // column's type. A record missing a column that appeared in the
320                // first record binds SQL NULL.
321                for (col, udt) in &insert_columns {
322                    let val = matched
323                        .iter()
324                        .find(|(c, _, _)| *c == col)
325                        .map(|(_, _, v)| *v);
326                    q = q.bind(pg_bind_text(val, udt));
327                }
328            }
329
330            q.execute(&mut *conn)
331                .await
332                .map_err(|e| FaucetError::Sink(format!("PostgreSQL insert failed: {e}")))?;
333        }
334
335        Ok(num_rows)
336    }
337
338    /// Insert a batch of records using auto-mapped columns, on the given
339    /// connection, with plain append semantics (no `ON CONFLICT` tail).
340    ///
341    /// Thin wrapper over [`insert_auto_map_with_conflict`](Self::insert_auto_map_with_conflict)
342    /// so the append path and the idempotent-write path keep their original
343    /// signature.
344    async fn insert_auto_map(
345        &self,
346        conn: &mut sqlx::PgConnection,
347        records: &[Value],
348    ) -> Result<usize, FaucetError> {
349        self.insert_auto_map_with_conflict(conn, records, None)
350            .await
351    }
352
353    /// Delete rows whose key columns match any of `deletes`, using
354    /// `DELETE FROM t WHERE (k1, …) IN ((v1, …), …)` with per-column `::udt`
355    /// casts (the key columns' underlying types), chunked at the param cap.
356    async fn delete_by_keys(
357        &self,
358        conn: &mut sqlx::PgConnection,
359        deletes: &[faucet_core::KeyTuple],
360    ) -> Result<usize, FaucetError> {
361        if deletes.is_empty() {
362            return Ok(0);
363        }
364        let key = &self.config.write.key;
365        let table_ref = qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name);
366
367        // Underlying types for the key columns → drives the ::udt casts, same
368        // source the insert path uses.
369        let udts: std::collections::HashMap<String, String> = sqlx::query(
370            "SELECT a.attname::text AS column_name, t.typname::text AS udt_name \
371             FROM pg_catalog.pg_attribute a \
372             JOIN pg_catalog.pg_type t ON t.oid = a.atttypid \
373             WHERE a.attrelid = to_regclass($1)::oid AND a.attnum > 0 AND NOT a.attisdropped",
374        )
375        .bind(&table_ref)
376        .fetch_all(&mut *conn)
377        .await
378        .map_err(|e| FaucetError::Sink(format!("failed to query table columns: {e}")))?
379        .iter()
380        .map(|row| {
381            (
382                row.get::<String, _>("column_name"),
383                row.get::<String, _>("udt_name"),
384            )
385        })
386        .collect();
387        let key_udts: Vec<String> = key
388            .iter()
389            .map(|k| udts.get(k).cloned().unwrap_or_else(|| "text".to_string()))
390            .collect();
391        let col_list = key
392            .iter()
393            .map(|k| quote_ident(k))
394            .collect::<Vec<_>>()
395            .join(", ");
396
397        const MAX_PG_PARAMS: usize = 65535;
398        let per = (MAX_PG_PARAMS / key.len().max(1)).max(1);
399        let mut total = 0usize;
400        for chunk in deletes.chunks(per) {
401            let mut ph = 1usize;
402            let tuples: Vec<String> = chunk
403                .iter()
404                .map(|_| {
405                    let group = key_udts
406                        .iter()
407                        .map(|udt| {
408                            let p = format!("${ph}::{udt}");
409                            ph += 1;
410                            p
411                        })
412                        .collect::<Vec<_>>()
413                        .join(", ");
414                    format!("({group})")
415                })
416                .collect();
417            let sql = format!(
418                "DELETE FROM {table_ref} WHERE ({col_list}) IN ({})",
419                tuples.join(", ")
420            );
421            let mut q = sqlx::query(&sql);
422            for kt in chunk {
423                for ((_, v), udt) in kt.0.iter().zip(key_udts.iter()) {
424                    q = q.bind(pg_bind_text(Some(v), udt));
425                }
426            }
427            let res = q
428                .execute(&mut *conn)
429                .await
430                .map_err(|e| FaucetError::Sink(format!("PostgreSQL delete failed: {e}")))?;
431            total += res.rows_affected() as usize;
432        }
433        Ok(total)
434    }
435
436    /// Apply a planned upsert/delete batch on one connection.
437    async fn apply_plan(
438        &self,
439        conn: &mut sqlx::PgConnection,
440        plan: &faucet_core::WritePlan,
441    ) -> Result<usize, FaucetError> {
442        let mut affected = 0usize;
443        if !plan.upserts.is_empty() {
444            affected += self
445                .insert_auto_map_with_conflict(conn, &plan.upserts, Some(&self.config.write.key))
446                .await?;
447        }
448        if !plan.deletes.is_empty() {
449            affected += self.delete_by_keys(conn, &plan.deletes).await?;
450        }
451        Ok(affected)
452    }
453
454    /// Ensure the commit-token watermark table exists.
455    async fn ensure_commit_table(&self) -> Result<(), FaucetError> {
456        let sql = format!(
457            "CREATE TABLE IF NOT EXISTS {t} ({s} TEXT PRIMARY KEY, {k} TEXT NOT NULL, updated_at TIMESTAMPTZ DEFAULT now())",
458            t = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TABLE),
459            s = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_SCOPE_COL),
460            k = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TOKEN_COL),
461        );
462        sqlx::query(&sql).execute(&self.pool).await.map_err(|e| {
463            FaucetError::Sink(format!("PostgreSQL commit-table create failed: {e}"))
464        })?;
465        Ok(())
466    }
467}
468
469#[async_trait]
470impl faucet_core::Sink for PostgresSink {
471    fn config_schema(&self) -> serde_json::Value {
472        serde_json::to_value(faucet_core::schema_for!(PostgresSinkConfig))
473            .expect("schema serialization")
474    }
475
476    fn supported_write_modes(&self) -> &'static [faucet_core::WriteMode] {
477        &[
478            faucet_core::WriteMode::Append,
479            faucet_core::WriteMode::Upsert,
480            faucet_core::WriteMode::Delete,
481        ]
482    }
483
484    fn dataset_uri(&self) -> String {
485        let table = match &self.config.schema {
486            Some(s) => format!("{}.{}", s, self.config.table_name),
487            None => self.config.table_name.clone(),
488        };
489        format!(
490            "{}?table={}",
491            faucet_core::redact_uri_credentials(&self.config.connection_url),
492            table
493        )
494    }
495
496    /// Preflight connectivity probe (`faucet doctor`).
497    ///
498    /// Acquires a connection from the existing pool and runs `SELECT 1`. This
499    /// is non-mutating and idempotent — it validates that the database is
500    /// reachable and the credentials are accepted without writing anything.
501    async fn check(
502        &self,
503        ctx: &faucet_core::check::CheckContext,
504    ) -> Result<faucet_core::check::CheckReport, FaucetError> {
505        use faucet_core::check::{CheckReport, Probe};
506
507        let started = std::time::Instant::now();
508        let probe =
509            match tokio::time::timeout(ctx.timeout, sqlx::query("SELECT 1").execute(&self.pool))
510                .await
511            {
512                Ok(Ok(_)) => Probe::pass("auth", started.elapsed()),
513                Ok(Err(e)) => Probe::fail_hint(
514                    "auth",
515                    started.elapsed(),
516                    e.to_string(),
517                    "check connection_url / credentials / that the database is reachable",
518                ),
519                Err(_) => Probe::fail_hint(
520                    "auth",
521                    started.elapsed(),
522                    "timed out",
523                    "check connection_url / credentials / that the database is reachable",
524                ),
525            };
526        Ok(CheckReport::single(probe))
527    }
528
529    /// Write records to PostgreSQL.
530    ///
531    /// When `config.batch_size > 0` and the input slice is larger than
532    /// `batch_size`, the slice is split into chunks of `batch_size` rows and
533    /// each chunk is sent as a separate multi-row `INSERT`. When
534    /// `config.batch_size == 0`, the entire slice is sent in a single
535    /// `INSERT` — useful when upstream `StreamPage`s are already sized for
536    /// Postgres' per-statement bind-parameter limit (~65 535 / num_columns
537    /// in AutoMap mode).
538    ///
539    /// Acquires one connection from the pool and routes all chunks through it.
540    /// Each INSERT executes as its own autocommit statement — identical
541    /// observable behaviour to executing directly on the pool, while keeping the
542    /// same connection for the entire call (avoids repeated pool-checkout
543    /// overhead on large batches).
544    async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
545        if records.is_empty() {
546            return Ok(0);
547        }
548
549        if !matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
550            let plan = faucet_core::plan_writes(records, &self.config.write);
551            if let Some((idx, msg)) = plan.failed.first() {
552                return Err(FaucetError::Sink(format!(
553                    "postgres {}: row {idx}: {msg}",
554                    self.config.write.write_mode.as_str()
555                )));
556            }
557            let mut conn =
558                self.pool.acquire().await.map_err(|e| {
559                    FaucetError::Sink(format!("PostgreSQL pool acquire failed: {e}"))
560                })?;
561            return self.apply_plan(&mut conn, &plan).await;
562        }
563
564        let chunks: Vec<&[Value]> = if self.config.batch_size == 0 {
565            // Sentinel: pass the entire upstream page through in a single
566            // INSERT statement. Subject to Postgres' 65 535 bind-parameter
567            // limit in AutoMap mode; JSONB mode binds a single array.
568            vec![records]
569        } else {
570            records.chunks(self.config.batch_size).collect()
571        };
572
573        // Acquire once; reuse for all chunks (each statement autocommits —
574        // no BEGIN is issued, so behaviour is identical to using the pool).
575        let mut conn = self
576            .pool
577            .acquire()
578            .await
579            .map_err(|e| FaucetError::Sink(format!("PostgreSQL pool acquire failed: {e}")))?;
580
581        let mut total = 0;
582        for chunk in chunks {
583            total += match &self.config.column_mapping {
584                PostgresColumnMapping::Jsonb { column } => {
585                    self.insert_jsonb(&mut conn, chunk, column).await?
586                }
587                PostgresColumnMapping::AutoMap => self.insert_auto_map(&mut conn, chunk).await?,
588            };
589        }
590
591        tracing::info!(
592            table = %self.config.table_name,
593            rows = total,
594            "PostgreSQL write complete"
595        );
596        Ok(total)
597    }
598
599    /// Write a batch and report per-row outcomes.
600    ///
601    /// In append mode this delegates to [`write_batch`](faucet_core::Sink::write_batch) and
602    /// maps a single success onto an all-`Ok(())` vector (the trait default).
603    /// In upsert/delete mode the good rows are applied (upserts + deletes), and
604    /// only the rows whose key could not be extracted (missing / null key) are
605    /// reported as `Err` so the pipeline routes them to the DLQ per-row instead
606    /// of sending the whole page.
607    async fn write_batch_partial(
608        &self,
609        records: &[Value],
610    ) -> Result<Vec<faucet_core::RowOutcome>, FaucetError> {
611        if matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
612            self.write_batch(records).await?;
613            return Ok(records.iter().map(|_| Ok(())).collect());
614        }
615
616        let plan = faucet_core::plan_writes(records, &self.config.write);
617        let mut conn = self
618            .pool
619            .acquire()
620            .await
621            .map_err(|e| FaucetError::Sink(format!("PostgreSQL pool acquire failed: {e}")))?;
622        self.apply_plan(&mut conn, &plan).await?;
623
624        let mut outcomes: Vec<faucet_core::RowOutcome> = records.iter().map(|_| Ok(())).collect();
625        for (idx, msg) in &plan.failed {
626            outcomes[*idx] = Err(FaucetError::Sink(format!(
627                "postgres {}: {msg}",
628                self.config.write.write_mode.as_str()
629            )));
630        }
631        Ok(outcomes)
632    }
633
634    fn supports_idempotent_writes(&self) -> bool {
635        true
636    }
637
638    async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
639        self.ensure_commit_table().await?;
640        let sql = format!(
641            "SELECT {k} FROM {t} WHERE {s} = $1",
642            t = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TABLE),
643            k = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TOKEN_COL),
644            s = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_SCOPE_COL),
645        );
646        let row = sqlx::query(&sql)
647            .bind(scope)
648            .fetch_optional(&self.pool)
649            .await
650            .map_err(|e| FaucetError::Sink(format!("PostgreSQL token read failed: {e}")))?;
651        Ok(row.map(|r| r.get::<String, _>(0)))
652    }
653
654    async fn write_batch_idempotent(
655        &self,
656        records: &[Value],
657        scope: &str,
658        token: &str,
659    ) -> Result<usize, FaucetError> {
660        self.ensure_commit_table().await?;
661
662        // For upsert/delete modes, plan the page before opening the transaction
663        // so a key-extraction failure aborts without leaving an open tx.
664        let plan = if matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
665            None
666        } else {
667            let plan = faucet_core::plan_writes(records, &self.config.write);
668            if let Some((idx, msg)) = plan.failed.first() {
669                return Err(FaucetError::Sink(format!(
670                    "postgres {}: row {idx}: {msg}",
671                    self.config.write.write_mode.as_str()
672                )));
673            }
674            Some(plan)
675        };
676
677        let mut tx =
678            self.pool.begin().await.map_err(|e| {
679                FaucetError::Sink(format!("PostgreSQL transaction begin failed: {e}"))
680            })?;
681
682        // Data write(s) and the commit-token upsert share ONE transaction so
683        // the page is committed atomically with its watermark: on crash either
684        // both land or neither does, which is what makes a replay skip-on-resume
685        // produce zero duplicates. For upsert/delete this means the planned
686        // upserts/deletes commit together with the watermark in the same tx.
687        let written = match &plan {
688            Some(plan) => self.apply_plan(&mut tx, plan).await?,
689            None => match &self.config.column_mapping {
690                PostgresColumnMapping::Jsonb { column } => {
691                    self.insert_jsonb(&mut tx, records, column).await?
692                }
693                PostgresColumnMapping::AutoMap => self.insert_auto_map(&mut tx, records).await?,
694            },
695        };
696
697        let upsert = format!(
698            "INSERT INTO {t} ({s}, {k}) VALUES ($1, $2) ON CONFLICT ({s}) DO UPDATE SET {k} = EXCLUDED.{k}, updated_at = now()",
699            t = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TABLE),
700            s = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_SCOPE_COL),
701            k = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TOKEN_COL),
702        );
703        sqlx::query(&upsert)
704            .bind(scope)
705            .bind(token)
706            .execute(&mut *tx)
707            .await
708            .map_err(|e| FaucetError::Sink(format!("PostgreSQL token upsert failed: {e}")))?;
709
710        tx.commit()
711            .await
712            .map_err(|e| FaucetError::Sink(format!("PostgreSQL transaction commit failed: {e}")))?;
713        Ok(written)
714    }
715}
716
717#[cfg(test)]
718mod tests {
719    use super::{on_conflict_clause, pg_bind_text, qualified_table_ref};
720    use serde_json::json;
721
722    #[test]
723    fn upsert_on_conflict_clause_for_keys() {
724        let clause = on_conflict_clause(
725            &["id".to_string()],
726            &["id".to_string(), "name".to_string(), "email".to_string()],
727        );
728        assert_eq!(
729            clause,
730            r#"ON CONFLICT ("id") DO UPDATE SET "name" = EXCLUDED."name", "email" = EXCLUDED."email""#
731        );
732    }
733
734    #[test]
735    fn upsert_on_conflict_all_columns_are_key_does_nothing() {
736        let clause = on_conflict_clause(&["id".to_string()], &["id".to_string()]);
737        assert_eq!(clause, r#"ON CONFLICT ("id") DO NOTHING"#);
738    }
739
740    #[test]
741    fn commit_token_table_is_the_shared_constant() {
742        assert_eq!(
743            faucet_core::idempotency::COMMIT_TOKEN_TABLE,
744            "_faucet_commit_token"
745        );
746    }
747
748    // dataset_uri test is skipped: PostgresSink::new() requires a live pool
749    // (connects to PostgreSQL in new()), and no offline constructor exists.
750    // The URI format is covered by unit tests in faucet-core's redact tests.
751
752    #[test]
753    fn qualified_table_ref_unqualified_is_bare_quoted_table() {
754        // No schema → bare quoted table, resolved against the search_path.
755        assert_eq!(qualified_table_ref(None, "events"), "\"events\"");
756    }
757
758    #[test]
759    fn qualified_table_ref_with_schema_is_schema_dot_table() {
760        // With a schema → "schema"."table", so discovery and INSERT both
761        // target the same explicit relation (#146 M13).
762        assert_eq!(
763            qualified_table_ref(Some("analytics"), "events"),
764            "\"analytics\".\"events\""
765        );
766    }
767
768    #[test]
769    fn qualified_table_ref_escapes_embedded_quotes() {
770        // SQL-injection safety: embedded double-quotes are doubled.
771        assert_eq!(
772            qualified_table_ref(Some("we\"ird"), "ta\"ble"),
773            "\"we\"\"ird\".\"ta\"\"ble\""
774        );
775    }
776
777    #[test]
778    fn null_and_absent_bind_sql_null() {
779        assert_eq!(pg_bind_text(None, "text"), None);
780        assert_eq!(pg_bind_text(Some(&json!(null)), "int4"), None);
781        assert_eq!(pg_bind_text(Some(&json!(null)), "jsonb"), None);
782    }
783
784    #[test]
785    fn scalars_bind_plain_text_for_typed_columns() {
786        // The `$N::<udt>` cast parses these via the column's input function.
787        assert_eq!(
788            pg_bind_text(Some(&json!(42)), "int4").as_deref(),
789            Some("42")
790        );
791        assert_eq!(
792            pg_bind_text(Some(&json!(1.5)), "numeric").as_deref(),
793            Some("1.5")
794        );
795        assert_eq!(
796            pg_bind_text(Some(&json!(true)), "bool").as_deref(),
797            Some("true")
798        );
799        assert_eq!(
800            pg_bind_text(Some(&json!("2025-01-01T00:00:00Z")), "timestamptz").as_deref(),
801            Some("2025-01-01T00:00:00Z")
802        );
803        // A plain string into TEXT keeps NO JSON quotes (the bug bound `"Bob"`).
804        assert_eq!(
805            pg_bind_text(Some(&json!("Bob")), "text").as_deref(),
806            Some("Bob")
807        );
808        // Large u64 beyond i64 keeps exact text (no f64 precision loss).
809        assert_eq!(
810            pg_bind_text(Some(&json!(18446744073709551615u64)), "numeric").as_deref(),
811            Some("18446744073709551615")
812        );
813    }
814
815    #[test]
816    fn json_columns_get_json_text_with_quotes_preserved() {
817        // For jsonb/json columns the value is bound as JSON text so the
818        // `::jsonb` cast parses it: a string keeps its quotes, objects/arrays
819        // round-trip.
820        assert_eq!(
821            pg_bind_text(Some(&json!("Bob")), "jsonb").as_deref(),
822            Some("\"Bob\"")
823        );
824        assert_eq!(
825            pg_bind_text(Some(&json!({"a": 1})), "jsonb").as_deref(),
826            Some("{\"a\":1}")
827        );
828        assert_eq!(
829            pg_bind_text(Some(&json!([1, 2])), "json").as_deref(),
830            Some("[1,2]")
831        );
832        assert_eq!(pg_bind_text(Some(&json!(5)), "jsonb").as_deref(), Some("5"));
833        // udt match is case-insensitive.
834        assert_eq!(
835            pg_bind_text(Some(&json!("x")), "JSONB").as_deref(),
836            Some("\"x\"")
837        );
838    }
839
840    #[test]
841    fn objects_into_non_json_columns_emit_json_text_so_the_cast_fails_loudly() {
842        // No scalar text form for an object targeting e.g. an int column; the
843        // `::int4` cast will reject this rather than silently coercing.
844        assert_eq!(
845            pg_bind_text(Some(&json!({"a": 1})), "int4").as_deref(),
846            Some("{\"a\":1}")
847        );
848    }
849}