Skip to main content

faucet_sink_postgres/
sink.rs

1//! PostgreSQL sink implementation.
2
3use crate::config::{PostgresColumnMapping, PostgresSinkConfig, PostgresWriteMethod};
4use crate::copy::{build_auto_map_payload, build_jsonb_payload, copy_statement};
5use async_trait::async_trait;
6use faucet_core::FaucetError;
7use faucet_core::util::quote_ident;
8use serde_json::Value;
9use sqlx::postgres::PgPoolOptions;
10use sqlx::{PgPool, Row};
11
12/// Render a JSON value as the text to bind for a PostgreSQL column whose
13/// underlying type is `udt` (`information_schema.columns.udt_name`), or `None`
14/// for SQL `NULL`.
15///
16/// The accompanying placeholder is emitted as `$N::<udt>`, so PostgreSQL runs
17/// the destination column type's input function over this text. That makes
18/// `string → timestamptz/uuid/date`, `number → int4/numeric/float8`,
19/// `bool → bool`, and `json → jsonb` all work — instead of binding every value
20/// as `serde_json::Value` (which sqlx encodes as `jsonb`, so an insert into any
21/// non-`jsonb` column fails at runtime with *"column is of type … but
22/// expression is of type jsonb"*; this was the C1 bug in audit #146).
23///
24/// For `json`/`jsonb` columns the value is bound as its JSON text (so a string
25/// keeps its quotes and objects/arrays round-trip); the `::jsonb` cast then
26/// parses it. For every other type the scalar's plain text form is bound and
27/// the column's input function parses it via the cast.
28pub(crate) fn pg_bind_text(value: Option<&Value>, udt: &str) -> Option<String> {
29    match value {
30        None | Some(Value::Null) => None,
31        Some(v) => {
32            if udt.eq_ignore_ascii_case("json") || udt.eq_ignore_ascii_case("jsonb") {
33                Some(v.to_string())
34            } else {
35                match v {
36                    Value::Bool(b) => Some(b.to_string()),
37                    Value::Number(n) => Some(n.to_string()),
38                    Value::String(s) => Some(s.clone()),
39                    // Arrays/objects have no scalar text form for a non-JSON
40                    // column; bind their JSON text so the `::<type>` cast fails
41                    // loudly rather than silently coercing.
42                    other => Some(other.to_string()),
43                }
44            }
45        }
46    }
47}
48
49/// Build the SQL relation reference for the configured table, optionally
50/// schema-qualified.
51///
52/// Both the AutoMap column-discovery probe and the `INSERT` statements use this
53/// single helper, so column discovery is always scoped to the *exact* relation
54/// the `INSERT` targets (#146 M13). With no schema the bare quoted table name
55/// resolves against the connection's `search_path`; with a schema it becomes
56/// `"schema"."table"`, pinning both discovery and insert to that namespace —
57/// otherwise a table of the same name in another schema pollutes the
58/// AutoMap column set (duplicate / wrong columns).
59fn qualified_table_ref(schema: Option<&str>, table: &str) -> String {
60    match schema {
61        Some(s) => format!("{}.{}", quote_ident(s), quote_ident(table)),
62        None => quote_ident(table),
63    }
64}
65
66/// Build the `ON CONFLICT (key) DO UPDATE …` tail for an upsert INSERT.
67/// Non-key columns are SET from EXCLUDED. If every column is a key column,
68/// emit `DO NOTHING`.
69fn on_conflict_clause(key: &[String], all_cols: &[String]) -> String {
70    let key_list = key
71        .iter()
72        .map(|k| quote_ident(k))
73        .collect::<Vec<_>>()
74        .join(", ");
75    let updates: Vec<String> = all_cols
76        .iter()
77        .filter(|c| !key.iter().any(|k| k == *c))
78        .map(|c| format!("{q} = EXCLUDED.{q}", q = quote_ident(c)))
79        .collect();
80    if updates.is_empty() {
81        format!("ON CONFLICT ({key_list}) DO NOTHING")
82    } else {
83        format!(
84            "ON CONFLICT ({key_list}) DO UPDATE SET {}",
85            updates.join(", ")
86        )
87    }
88}
89
90/// Map a [`faucet_core::SqlBaseType`] to the PostgreSQL type keyword used when
91/// adding/widening a column during schema evolution (issue #194). Integers
92/// always widen to `bigint` and floats to `double precision` so a later, wider
93/// value never overflows a narrower column.
94fn pg_keyword(t: faucet_core::SqlBaseType) -> &'static str {
95    use faucet_core::SqlBaseType::*;
96    match t {
97        Integer => "bigint",
98        Double => "double precision",
99        Boolean => "boolean",
100        Text => "text",
101        Json => "jsonb",
102    }
103}
104
105/// `ALTER TABLE <ref> ADD COLUMN IF NOT EXISTS "<col>" <kw>` — idempotent column
106/// addition. `table_ref` is already quoted (`"schema"."table"`).
107fn build_add_column_sql(table_ref: &str, col: &str, t: faucet_core::SqlBaseType) -> String {
108    format!(
109        "ALTER TABLE {table_ref} ADD COLUMN IF NOT EXISTS {} {}",
110        quote_ident(col),
111        pg_keyword(t)
112    )
113}
114
115/// `ALTER TABLE <ref> ALTER COLUMN "<col>" TYPE <kw> USING "<col>"::<kw>` — widen
116/// an existing column's type. Naturally idempotent (re-running the same TYPE
117/// change is a no-op).
118fn build_alter_type_sql(table_ref: &str, col: &str, t: faucet_core::SqlBaseType) -> String {
119    let q = quote_ident(col);
120    let kw = pg_keyword(t);
121    format!("ALTER TABLE {table_ref} ALTER COLUMN {q} TYPE {kw} USING {q}::{kw}")
122}
123
124/// `ALTER TABLE <ref> ALTER COLUMN "<col>" DROP NOT NULL` — relax a NOT NULL
125/// constraint. Naturally idempotent.
126fn build_drop_not_null_sql(table_ref: &str, col: &str) -> String {
127    format!(
128        "ALTER TABLE {table_ref} ALTER COLUMN {} DROP NOT NULL",
129        quote_ident(col)
130    )
131}
132
133/// Map a PostgreSQL type name (`pg_type.typname`, e.g. `int4`, `float8`, `bool`,
134/// `jsonb`) back to a JSON-Schema type fragment so [`PostgresSink::current_schema`]
135/// round-trips with [`faucet_core::diff_schema`]. `nullable` reflects whether the
136/// column allows NULL (`NOT a.attnotnull`).
137fn pg_udt_to_json_schema(udt: &str, nullable: bool) -> serde_json::Value {
138    let base = match udt {
139        "int2" | "int4" | "int8" => "integer",
140        "float4" | "float8" | "numeric" => "number",
141        "bool" => "boolean",
142        "json" | "jsonb" => "object",
143        _ => "string",
144    };
145    if nullable {
146        serde_json::json!({ "type": [base, "null"] })
147    } else {
148        serde_json::json!({ "type": base })
149    }
150}
151
152/// A sink that writes JSON records to a PostgreSQL table.
153pub struct PostgresSink {
154    config: PostgresSinkConfig,
155    pool: PgPool,
156}
157
158impl PostgresSink {
159    /// Create a new PostgreSQL sink. Establishes a connection pool.
160    pub async fn new(config: PostgresSinkConfig) -> Result<Self, FaucetError> {
161        config.write.validate()?;
162        if !matches!(config.write.write_mode, faucet_core::WriteMode::Append)
163            && !matches!(config.column_mapping, PostgresColumnMapping::AutoMap)
164        {
165            return Err(FaucetError::Config(
166                "postgres sink: write_mode upsert/delete requires column_mapping: auto_map \
167                 (key columns must be real columns, not inside a JSONB blob)"
168                    .into(),
169            ));
170        }
171        if matches!(config.write_method, PostgresWriteMethod::Copy)
172            && !matches!(config.write.write_mode, faucet_core::WriteMode::Append)
173        {
174            return Err(FaucetError::Config(format!(
175                "postgres sink: write_method: copy is append-only (COPY has no ON CONFLICT); \
176                 it cannot be combined with write_mode: {} — use write_method: insert",
177                config.write.write_mode.as_str()
178            )));
179        }
180
181        let pool = PgPoolOptions::new()
182            .max_connections(config.max_connections)
183            .connect(&config.connection_url)
184            .await
185            .map_err(|e| FaucetError::Sink(format!("PostgreSQL connection failed: {e}")))?;
186
187        Ok(Self { config, pool })
188    }
189
190    /// Discover the target relation's column names and underlying types
191    /// (`pg_type.typname`), scoped to the *exact* relation the writes target
192    /// via `to_regclass` (#146 M13). Shared by the INSERT and COPY paths so
193    /// both see an identical column set. `::text` casts the `name`-typed
194    /// catalog columns so sqlx decodes them as `String`.
195    async fn discover_columns(
196        &self,
197        conn: &mut sqlx::PgConnection,
198        table_ref: &str,
199    ) -> Result<Vec<(String, String)>, FaucetError> {
200        let columns: Vec<(String, String)> = sqlx::query(
201            "SELECT a.attname::text AS column_name, t.typname::text AS udt_name \
202             FROM pg_catalog.pg_attribute a \
203             JOIN pg_catalog.pg_type t ON t.oid = a.atttypid \
204             WHERE a.attrelid = to_regclass($1)::oid \
205               AND a.attnum > 0 AND NOT a.attisdropped \
206             ORDER BY a.attnum",
207        )
208        .bind(table_ref)
209        .fetch_all(&mut *conn)
210        .await
211        .map_err(|e| FaucetError::Sink(format!("failed to query table columns: {e}")))?
212        .iter()
213        .map(|row| {
214            (
215                row.get::<String, _>("column_name"),
216                row.get::<String, _>("udt_name"),
217            )
218        })
219        .collect();
220
221        if columns.is_empty() {
222            return Err(FaucetError::Sink(format!(
223                "table {table_ref} has no columns or does not exist"
224            )));
225        }
226        Ok(columns)
227    }
228
229    /// Write one chunk via `COPY … FROM STDIN (FORMAT text)` — the bulk-load
230    /// fast-path (issue #308). Append-only (validated at construction); the
231    /// server parses each field with the destination column's input function,
232    /// so type semantics match the `INSERT` path exactly. A bad row fails the
233    /// whole COPY (all-or-nothing, like a failed multi-row `INSERT`).
234    async fn copy_batch(
235        &self,
236        conn: &mut sqlx::PgConnection,
237        records: &[Value],
238    ) -> Result<usize, FaucetError> {
239        if records.is_empty() {
240            return Ok(0);
241        }
242        let table_ref = qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name);
243
244        let (statement, payload) = match &self.config.column_mapping {
245            PostgresColumnMapping::Jsonb { column } => {
246                let payload = build_jsonb_payload(records);
247                (
248                    copy_statement(&table_ref, std::slice::from_ref(column)),
249                    payload,
250                )
251            }
252            PostgresColumnMapping::AutoMap => {
253                let columns = self.discover_columns(&mut *conn, &table_ref).await?;
254                let Some(payload) =
255                    build_auto_map_payload(records, &columns).map_err(FaucetError::Sink)?
256                else {
257                    return Ok(0);
258                };
259                (copy_statement(&table_ref, &payload.columns), payload)
260            }
261        };
262
263        let mut copy_in = conn
264            .copy_in_raw(&statement)
265            .await
266            .map_err(|e| FaucetError::Sink(format!("PostgreSQL COPY start failed: {e}")))?;
267        // Ship in ~1 MiB slices so a huge page never materializes a second
268        // time inside sqlx's write buffer.
269        const SEND_CHUNK: usize = 1 << 20;
270        for chunk in payload.data.as_bytes().chunks(SEND_CHUNK) {
271            if let Err(e) = copy_in.send(chunk).await {
272                // Dropping the handle aborts the COPY server-side; surface
273                // the original error.
274                return Err(FaucetError::Sink(format!(
275                    "PostgreSQL COPY send failed: {e}"
276                )));
277            }
278        }
279        copy_in
280            .finish()
281            .await
282            .map_err(|e| FaucetError::Sink(format!("PostgreSQL COPY failed: {e}")))?;
283        Ok(payload.rows)
284    }
285
286    /// Insert a batch of records using JSONB column mode, on the given connection.
287    ///
288    /// Accepts `&mut sqlx::PgConnection` so the same logic runs both standalone
289    /// (via a pool-acquired connection, autocommit) and inside the idempotent
290    /// transaction (where `&mut *tx` is passed — `Transaction<'_, Postgres>`
291    /// derefs to `PgConnection`).
292    async fn insert_jsonb(
293        &self,
294        conn: &mut sqlx::PgConnection,
295        records: &[Value],
296        column: &str,
297    ) -> Result<usize, FaucetError> {
298        if records.is_empty() {
299            return Ok(0);
300        }
301
302        // Use a single INSERT with unnest for efficiency.
303        let json_values: Vec<serde_json::Value> = records.to_vec();
304        let query = format!(
305            "INSERT INTO {} ({}) SELECT * FROM unnest($1::jsonb[])",
306            qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name),
307            quote_ident(column)
308        );
309
310        sqlx::query(&query)
311            .bind(json_values)
312            .execute(&mut *conn)
313            .await
314            .map_err(|e| FaucetError::Sink(format!("PostgreSQL insert failed: {e}")))?;
315
316        Ok(records.len())
317    }
318
319    /// Insert a batch of records using auto-mapped columns, on the given connection.
320    ///
321    /// Accepts `&mut sqlx::PgConnection` so the same logic runs both standalone
322    /// (via a pool-acquired connection, autocommit) and inside the idempotent
323    /// transaction (where `&mut *tx` is passed — `Transaction<'_, Postgres>`
324    /// derefs to `PgConnection`). Running the column-discovery query on the same
325    /// connection is harmless and avoids any cross-connection visibility surprise.
326    ///
327    /// Discovers each column's name *and* underlying type (`udt_name`) from the
328    /// table schema and maps top-level JSON fields to columns. Each placeholder
329    /// is emitted as `$N::<udt>` and the value is bound as text (see
330    /// [`pg_bind_text`]), so the destination column's input function parses it —
331    /// numbers, booleans, timestamps, uuids, and JSON all land in their native
332    /// column types. (Previously every value was bound as `serde_json::Value`,
333    /// which sqlx encodes as `jsonb`, so an insert into any non-`jsonb` column
334    /// failed at runtime — audit #146 C1.) Uses a single multi-row INSERT
335    /// (sub-chunked at the 65535-parameter cap) for efficiency.
336    ///
337    /// When `conflict_key` is `Some(key)`, each sub-chunk's INSERT is given an
338    /// `ON CONFLICT (key) DO UPDATE …` tail so it upserts by the key columns
339    /// (last-write-wins within the batch is already handled by the planner's
340    /// dedup, so a single sub-chunk never double-hits the same conflict target).
341    async fn insert_auto_map_with_conflict(
342        &self,
343        conn: &mut sqlx::PgConnection,
344        records: &[Value],
345        conflict_key: Option<&[String]>,
346    ) -> Result<usize, FaucetError> {
347        if records.is_empty() {
348            return Ok(0);
349        }
350
351        // Get column names AND their underlying types for the *exact* relation
352        // the INSERT will target. Scoping by `to_regclass(<qualified ref>)`
353        // resolves the relation the same way the INSERT does — by the configured
354        // schema if set, otherwise by the connection's `search_path` — so a
355        // table of the same name in another schema can no longer pollute the
356        // column set with duplicate/wrong columns (#146 M13). The previous query
357        // filtered `information_schema.columns` by `table_name` alone (no schema
358        // predicate), merging every same-named table across all schemas.
359        //
360        // `pg_type.typname` is the concrete type (`int4`, `timestamptz`,
361        // `numeric`, `jsonb`, `uuid`, `text`, …) — identical to the old
362        // `information_schema.columns.udt_name` — used as the per-placeholder
363        // cast target below.
364        let table_ref = qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name);
365        let columns = self.discover_columns(&mut *conn, &table_ref).await?;
366
367        // Pre-validate all records and collect matched (column, udt, value)
368        // triples per record. The INSERT column set is the UNION of table
369        // columns present in ANY record (in declared table order), not just the
370        // first record's keys — otherwise a field present only in a later record
371        // of the batch would be silently dropped (audit #146 H1). A row missing
372        // a unioned column binds SQL NULL.
373        let mut matched_rows: Vec<Vec<(&String, &String, &Value)>> =
374            Vec::with_capacity(records.len());
375        let mut used: std::collections::HashSet<&str> = std::collections::HashSet::new();
376
377        for record in records {
378            let obj = record
379                .as_object()
380                .ok_or_else(|| FaucetError::Sink("AutoMap requires JSON object records".into()))?;
381
382            let matching: Vec<(&String, &String, &Value)> = columns
383                .iter()
384                .filter_map(|(col, udt)| obj.get(col).map(|v| (col, udt, v)))
385                .collect();
386
387            if matching.is_empty() {
388                tracing::warn!(
389                    record_keys = ?obj.keys().collect::<Vec<_>>(),
390                    table_columns = ?columns,
391                    "record has no keys matching table columns, skipping"
392                );
393                continue;
394            }
395
396            for (c, _, _) in &matching {
397                used.insert(c.as_str());
398            }
399            matched_rows.push(matching);
400        }
401
402        if matched_rows.is_empty() {
403            return Ok(0);
404        }
405
406        // Table columns (in declared order, with their udt) present in at least
407        // one record.
408        let insert_columns: Vec<(String, String)> = columns
409            .iter()
410            .filter(|(c, _)| used.contains(c.as_str()))
411            .cloned()
412            .collect();
413
414        let num_cols = insert_columns.len();
415        let num_rows = matched_rows.len();
416        let col_names: Vec<String> = insert_columns.iter().map(|(c, _)| quote_ident(c)).collect();
417
418        // PostgreSQL caps bind parameters per statement at 65535. A multi-row
419        // INSERT binds `rows × num_cols` parameters, so a wide table at a large
420        // batch_size can exceed it and fail at runtime (#78/#21). Split into
421        // sub-INSERTs of at most floor(MAX_PARAMS / num_cols) rows.
422        const MAX_PG_PARAMS: usize = 65535;
423        let max_rows_per_insert = (MAX_PG_PARAMS / num_cols).max(1);
424
425        for sub in matched_rows.chunks(max_rows_per_insert) {
426            // Build multi-row VALUES clause with per-column casts so the column
427            // type's input function parses the bound text:
428            //   ($1::int4, $2::timestamptz), ($3::int4, $4::timestamptz), ...
429            let mut value_tuples: Vec<String> = Vec::with_capacity(sub.len());
430            for row_idx in 0..sub.len() {
431                let start = row_idx * num_cols + 1;
432                let placeholders: Vec<String> = (0..num_cols)
433                    .map(|c| format!("${}::{}", start + c, insert_columns[c].1))
434                    .collect();
435                value_tuples.push(format!("({})", placeholders.join(", ")));
436            }
437
438            let query = format!(
439                "INSERT INTO {} ({}) VALUES {}",
440                table_ref,
441                col_names.join(", "),
442                value_tuples.join(", ")
443            );
444            let query = match conflict_key {
445                Some(key) => format!(
446                    "{query} {}",
447                    on_conflict_clause(
448                        key,
449                        &insert_columns
450                            .iter()
451                            .map(|(c, _)| c.clone())
452                            .collect::<Vec<_>>()
453                    )
454                ),
455                None => query,
456            };
457
458            let mut q = sqlx::query(&query);
459            for matched in sub {
460                // Bind values in the fixed column order, as text matching each
461                // column's type. A record missing a column that appeared in the
462                // first record binds SQL NULL.
463                for (col, udt) in &insert_columns {
464                    let val = matched
465                        .iter()
466                        .find(|(c, _, _)| *c == col)
467                        .map(|(_, _, v)| *v);
468                    q = q.bind(pg_bind_text(val, udt));
469                }
470            }
471
472            q.execute(&mut *conn)
473                .await
474                .map_err(|e| FaucetError::Sink(format!("PostgreSQL insert failed: {e}")))?;
475        }
476
477        Ok(num_rows)
478    }
479
480    /// Insert a batch of records using auto-mapped columns, on the given
481    /// connection, with plain append semantics (no `ON CONFLICT` tail).
482    ///
483    /// Thin wrapper over [`insert_auto_map_with_conflict`](Self::insert_auto_map_with_conflict)
484    /// so the append path and the idempotent-write path keep their original
485    /// signature.
486    async fn insert_auto_map(
487        &self,
488        conn: &mut sqlx::PgConnection,
489        records: &[Value],
490    ) -> Result<usize, FaucetError> {
491        self.insert_auto_map_with_conflict(conn, records, None)
492            .await
493    }
494
495    /// Delete rows whose key columns match any of `deletes`, using
496    /// `DELETE FROM t WHERE (k1, …) IN ((v1, …), …)` with per-column `::udt`
497    /// casts (the key columns' underlying types), chunked at the param cap.
498    async fn delete_by_keys(
499        &self,
500        conn: &mut sqlx::PgConnection,
501        deletes: &[faucet_core::KeyTuple],
502    ) -> Result<usize, FaucetError> {
503        if deletes.is_empty() {
504            return Ok(0);
505        }
506        let key = &self.config.write.key;
507        let table_ref = qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name);
508
509        // Underlying types for the key columns → drives the ::udt casts, same
510        // source the insert path uses.
511        let udts: std::collections::HashMap<String, String> = self
512            .discover_columns(&mut *conn, &table_ref)
513            .await?
514            .into_iter()
515            .collect();
516        let key_udts: Vec<String> = key
517            .iter()
518            .map(|k| udts.get(k).cloned().unwrap_or_else(|| "text".to_string()))
519            .collect();
520        let col_list = key
521            .iter()
522            .map(|k| quote_ident(k))
523            .collect::<Vec<_>>()
524            .join(", ");
525
526        const MAX_PG_PARAMS: usize = 65535;
527        let per = (MAX_PG_PARAMS / key.len().max(1)).max(1);
528        let mut total = 0usize;
529        for chunk in deletes.chunks(per) {
530            let mut ph = 1usize;
531            let tuples: Vec<String> = chunk
532                .iter()
533                .map(|_| {
534                    let group = key_udts
535                        .iter()
536                        .map(|udt| {
537                            let p = format!("${ph}::{udt}");
538                            ph += 1;
539                            p
540                        })
541                        .collect::<Vec<_>>()
542                        .join(", ");
543                    format!("({group})")
544                })
545                .collect();
546            let sql = format!(
547                "DELETE FROM {table_ref} WHERE ({col_list}) IN ({})",
548                tuples.join(", ")
549            );
550            let mut q = sqlx::query(&sql);
551            for kt in chunk {
552                for ((_, v), udt) in kt.0.iter().zip(key_udts.iter()) {
553                    q = q.bind(pg_bind_text(Some(v), udt));
554                }
555            }
556            let res = q
557                .execute(&mut *conn)
558                .await
559                .map_err(|e| FaucetError::Sink(format!("PostgreSQL delete failed: {e}")))?;
560            total += res.rows_affected() as usize;
561        }
562        Ok(total)
563    }
564
565    /// Apply a planned upsert/delete batch on one connection.
566    async fn apply_plan(
567        &self,
568        conn: &mut sqlx::PgConnection,
569        plan: &faucet_core::WritePlan,
570    ) -> Result<usize, FaucetError> {
571        let mut affected = 0usize;
572        if !plan.upserts.is_empty() {
573            affected += self
574                .insert_auto_map_with_conflict(conn, &plan.upserts, Some(&self.config.write.key))
575                .await?;
576        }
577        if !plan.deletes.is_empty() {
578            affected += self.delete_by_keys(conn, &plan.deletes).await?;
579        }
580        Ok(affected)
581    }
582
583    /// Ensure the commit-token watermark table exists.
584    async fn ensure_commit_table(&self) -> Result<(), FaucetError> {
585        let sql = format!(
586            "CREATE TABLE IF NOT EXISTS {t} ({s} TEXT PRIMARY KEY, {k} TEXT NOT NULL, updated_at TIMESTAMPTZ DEFAULT now())",
587            t = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TABLE),
588            s = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_SCOPE_COL),
589            k = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TOKEN_COL),
590        );
591        sqlx::query(&sql).execute(&self.pool).await.map_err(|e| {
592            FaucetError::Sink(format!("PostgreSQL commit-table create failed: {e}"))
593        })?;
594        Ok(())
595    }
596}
597
598#[async_trait]
599impl faucet_core::Sink for PostgresSink {
600    fn config_schema(&self) -> serde_json::Value {
601        serde_json::to_value(faucet_core::schema_for!(PostgresSinkConfig))
602            .expect("schema serialization")
603    }
604
605    fn supported_write_modes(&self) -> &'static [faucet_core::WriteMode] {
606        &[
607            faucet_core::WriteMode::Append,
608            faucet_core::WriteMode::Upsert,
609            faucet_core::WriteMode::Delete,
610        ]
611    }
612
613    fn dedups_by_key(&self) -> bool {
614        self.config.write.dedups_by_key()
615    }
616
617    fn supports_schema_evolution(&self) -> bool {
618        true
619    }
620
621    /// Read the live destination schema from `pg_catalog` as an
622    /// `infer_schema`-shaped object (`{"type":"object","properties":{…}}`), or
623    /// `None` when the target table does not exist yet (issue #194).
624    ///
625    /// Reuses the AutoMap column-discovery query shape (scoped to the exact
626    /// relation via `to_regclass`), additionally reading `a.attnotnull` so
627    /// nullability round-trips through `pg_udt_to_json_schema`.
628    async fn current_schema(&self) -> Result<Option<serde_json::Value>, FaucetError> {
629        let table_ref = qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name);
630        let rows: Vec<(String, String, bool)> = sqlx::query(
631            "SELECT a.attname::text AS column_name, t.typname::text AS udt_name, a.attnotnull \
632             FROM pg_catalog.pg_attribute a \
633             JOIN pg_catalog.pg_type t ON t.oid = a.atttypid \
634             WHERE a.attrelid = to_regclass($1)::oid \
635               AND a.attnum > 0 AND NOT a.attisdropped \
636             ORDER BY a.attnum",
637        )
638        .bind(&table_ref)
639        .fetch_all(&self.pool)
640        .await
641        .map_err(|e| FaucetError::Sink(format!("postgres current_schema query failed: {e}")))?
642        .iter()
643        .map(|row| {
644            (
645                row.get::<String, _>("column_name"),
646                row.get::<String, _>("udt_name"),
647                row.get::<bool, _>("attnotnull"),
648            )
649        })
650        .collect();
651
652        if rows.is_empty() {
653            return Ok(None); // table does not exist yet
654        }
655
656        let mut props = serde_json::Map::new();
657        for (name, udt, notnull) in rows {
658            props.insert(name, pg_udt_to_json_schema(&udt, !notnull));
659        }
660        Ok(Some(
661            serde_json::json!({ "type": "object", "properties": props }),
662        ))
663    }
664
665    /// Apply an additive schema evolution (new columns, lossless widenings,
666    /// nullability relaxations) to the destination table. Idempotent —
667    /// `ADD COLUMN IF NOT EXISTS`, and re-running the same TYPE / DROP NOT NULL
668    /// is a no-op (issue #194).
669    async fn evolve_schema(
670        &self,
671        evolution: &faucet_core::SchemaEvolution,
672    ) -> Result<(), FaucetError> {
673        let table_ref = qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name);
674        let mut conn = self
675            .pool
676            .acquire()
677            .await
678            .map_err(|e| FaucetError::Sink(format!("postgres evolve acquire failed: {e}")))?;
679
680        for c in &evolution.additions {
681            let t =
682                faucet_core::json_schema_base_type(&c.to).unwrap_or(faucet_core::SqlBaseType::Text);
683            sqlx::query(&build_add_column_sql(&table_ref, &c.name, t))
684                .execute(&mut *conn)
685                .await
686                .map_err(|e| {
687                    FaucetError::Sink(format!("postgres ADD COLUMN {} failed: {e}", c.name))
688                })?;
689        }
690        for c in &evolution.widenings {
691            let t =
692                faucet_core::json_schema_base_type(&c.to).unwrap_or(faucet_core::SqlBaseType::Text);
693            sqlx::query(&build_alter_type_sql(&table_ref, &c.name, t))
694                .execute(&mut *conn)
695                .await
696                .map_err(|e| {
697                    FaucetError::Sink(format!("postgres ALTER TYPE {} failed: {e}", c.name))
698                })?;
699        }
700        for col in &evolution.relax_nullability {
701            sqlx::query(&build_drop_not_null_sql(&table_ref, col))
702                .execute(&mut *conn)
703                .await
704                .map_err(|e| {
705                    FaucetError::Sink(format!("postgres DROP NOT NULL {col} failed: {e}"))
706                })?;
707        }
708        Ok(())
709    }
710
711    fn dataset_uri(&self) -> String {
712        let table = match &self.config.schema {
713            Some(s) => format!("{}.{}", s, self.config.table_name),
714            None => self.config.table_name.clone(),
715        };
716        format!(
717            "{}?table={}",
718            faucet_core::redact_uri_credentials(&self.config.connection_url),
719            table
720        )
721    }
722
723    /// Preflight connectivity probe (`faucet doctor`).
724    ///
725    /// Acquires a connection from the existing pool and runs `SELECT 1`. This
726    /// is non-mutating and idempotent — it validates that the database is
727    /// reachable and the credentials are accepted without writing anything.
728    async fn check(
729        &self,
730        ctx: &faucet_core::check::CheckContext,
731    ) -> Result<faucet_core::check::CheckReport, FaucetError> {
732        use faucet_core::check::{CheckReport, Probe};
733
734        let started = std::time::Instant::now();
735        let probe =
736            match tokio::time::timeout(ctx.timeout, sqlx::query("SELECT 1").execute(&self.pool))
737                .await
738            {
739                Ok(Ok(_)) => Probe::pass("auth", started.elapsed()),
740                Ok(Err(e)) => Probe::fail_hint(
741                    "auth",
742                    started.elapsed(),
743                    e.to_string(),
744                    "check connection_url / credentials / that the database is reachable",
745                ),
746                Err(_) => Probe::fail_hint(
747                    "auth",
748                    started.elapsed(),
749                    "timed out",
750                    "check connection_url / credentials / that the database is reachable",
751                ),
752            };
753        Ok(CheckReport::single(probe))
754    }
755
756    /// Write records to PostgreSQL.
757    ///
758    /// When `config.batch_size > 0` and the input slice is larger than
759    /// `batch_size`, the slice is split into chunks of `batch_size` rows and
760    /// each chunk is sent as a separate multi-row `INSERT`. When
761    /// `config.batch_size == 0`, the entire slice is sent in a single
762    /// `INSERT` — useful when upstream `StreamPage`s are already sized for
763    /// Postgres' per-statement bind-parameter limit (~65 535 / num_columns
764    /// in AutoMap mode).
765    ///
766    /// Acquires one connection from the pool and routes all chunks through it.
767    /// Each INSERT executes as its own autocommit statement — identical
768    /// observable behaviour to executing directly on the pool, while keeping the
769    /// same connection for the entire call (avoids repeated pool-checkout
770    /// overhead on large batches).
771    async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
772        if records.is_empty() {
773            return Ok(0);
774        }
775
776        if !matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
777            let plan = faucet_core::plan_writes(records, &self.config.write);
778            if let Some((idx, msg)) = plan.failed.first() {
779                return Err(FaucetError::Sink(format!(
780                    "postgres {}: row {idx}: {msg}",
781                    self.config.write.write_mode.as_str()
782                )));
783            }
784            let mut conn =
785                self.pool.acquire().await.map_err(|e| {
786                    FaucetError::Sink(format!("PostgreSQL pool acquire failed: {e}"))
787                })?;
788            return self.apply_plan(&mut conn, &plan).await;
789        }
790
791        let chunks: Vec<&[Value]> = if self.config.batch_size == 0 {
792            // Sentinel: pass the entire upstream page through in a single
793            // INSERT statement. Subject to Postgres' 65 535 bind-parameter
794            // limit in AutoMap mode; JSONB mode binds a single array.
795            vec![records]
796        } else {
797            records.chunks(self.config.batch_size).collect()
798        };
799
800        // Acquire once; reuse for all chunks (each statement autocommits —
801        // no BEGIN is issued, so behaviour is identical to using the pool).
802        let mut conn = self
803            .pool
804            .acquire()
805            .await
806            .map_err(|e| FaucetError::Sink(format!("PostgreSQL pool acquire failed: {e}")))?;
807
808        let mut total = 0;
809        for chunk in chunks {
810            total += match self.config.write_method {
811                // Bulk-load fast-path: COPY the chunk instead of a multi-row
812                // INSERT (append-only; validated at construction, #308).
813                PostgresWriteMethod::Copy => self.copy_batch(&mut conn, chunk).await?,
814                PostgresWriteMethod::Insert => match &self.config.column_mapping {
815                    PostgresColumnMapping::Jsonb { column } => {
816                        self.insert_jsonb(&mut conn, chunk, column).await?
817                    }
818                    PostgresColumnMapping::AutoMap => {
819                        self.insert_auto_map(&mut conn, chunk).await?
820                    }
821                },
822            };
823        }
824
825        tracing::info!(
826            table = %self.config.table_name,
827            rows = total,
828            "PostgreSQL write complete"
829        );
830        Ok(total)
831    }
832
833    /// Write a batch and report per-row outcomes.
834    ///
835    /// In append mode this delegates to [`write_batch`](faucet_core::Sink::write_batch) and
836    /// maps a single success onto an all-`Ok(())` vector (the trait default).
837    /// In upsert/delete mode the good rows are applied (upserts + deletes), and
838    /// only the rows whose key could not be extracted (missing / null key) are
839    /// reported as `Err` so the pipeline routes them to the DLQ per-row instead
840    /// of sending the whole page.
841    async fn write_batch_partial(
842        &self,
843        records: &[Value],
844    ) -> Result<Vec<faucet_core::RowOutcome>, FaucetError> {
845        if matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
846            self.write_batch(records).await?;
847            return Ok(records.iter().map(|_| Ok(())).collect());
848        }
849
850        let plan = faucet_core::plan_writes(records, &self.config.write);
851        let mut conn = self
852            .pool
853            .acquire()
854            .await
855            .map_err(|e| FaucetError::Sink(format!("PostgreSQL pool acquire failed: {e}")))?;
856        self.apply_plan(&mut conn, &plan).await?;
857
858        let mut outcomes: Vec<faucet_core::RowOutcome> = records.iter().map(|_| Ok(())).collect();
859        for (idx, msg) in &plan.failed {
860            outcomes[*idx] = Err(FaucetError::Sink(format!(
861                "postgres {}: {msg}",
862                self.config.write.write_mode.as_str()
863            )));
864        }
865        Ok(outcomes)
866    }
867
868    fn supports_idempotent_writes(&self) -> bool {
869        true
870    }
871
872    async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
873        self.ensure_commit_table().await?;
874        let sql = format!(
875            "SELECT {k} FROM {t} WHERE {s} = $1",
876            t = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TABLE),
877            k = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TOKEN_COL),
878            s = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_SCOPE_COL),
879        );
880        let row = sqlx::query(&sql)
881            .bind(scope)
882            .fetch_optional(&self.pool)
883            .await
884            .map_err(|e| FaucetError::Sink(format!("PostgreSQL token read failed: {e}")))?;
885        Ok(row.map(|r| r.get::<String, _>(0)))
886    }
887
888    async fn write_batch_idempotent(
889        &self,
890        records: &[Value],
891        scope: &str,
892        token: &str,
893    ) -> Result<usize, FaucetError> {
894        self.ensure_commit_table().await?;
895
896        // For upsert/delete modes, plan the page before opening the transaction
897        // so a key-extraction failure aborts without leaving an open tx.
898        let plan = if matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
899            None
900        } else {
901            let plan = faucet_core::plan_writes(records, &self.config.write);
902            if let Some((idx, msg)) = plan.failed.first() {
903                return Err(FaucetError::Sink(format!(
904                    "postgres {}: row {idx}: {msg}",
905                    self.config.write.write_mode.as_str()
906                )));
907            }
908            Some(plan)
909        };
910
911        let mut tx =
912            self.pool.begin().await.map_err(|e| {
913                FaucetError::Sink(format!("PostgreSQL transaction begin failed: {e}"))
914            })?;
915
916        // Data write(s) and the commit-token upsert share ONE transaction so
917        // the page is committed atomically with its watermark: on crash either
918        // both land or neither does, which is what makes a replay skip-on-resume
919        // produce zero duplicates. For upsert/delete this means the planned
920        // upserts/deletes commit together with the watermark in the same tx.
921        let written = match &plan {
922            Some(plan) => self.apply_plan(&mut tx, plan).await?,
923            None => match &self.config.column_mapping {
924                PostgresColumnMapping::Jsonb { column } => {
925                    self.insert_jsonb(&mut tx, records, column).await?
926                }
927                PostgresColumnMapping::AutoMap => self.insert_auto_map(&mut tx, records).await?,
928            },
929        };
930
931        let upsert = format!(
932            "INSERT INTO {t} ({s}, {k}) VALUES ($1, $2) ON CONFLICT ({s}) DO UPDATE SET {k} = EXCLUDED.{k}, updated_at = now()",
933            t = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TABLE),
934            s = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_SCOPE_COL),
935            k = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TOKEN_COL),
936        );
937        sqlx::query(&upsert)
938            .bind(scope)
939            .bind(token)
940            .execute(&mut *tx)
941            .await
942            .map_err(|e| FaucetError::Sink(format!("PostgreSQL token upsert failed: {e}")))?;
943
944        tx.commit()
945            .await
946            .map_err(|e| FaucetError::Sink(format!("PostgreSQL transaction commit failed: {e}")))?;
947        Ok(written)
948    }
949}
950
951#[cfg(test)]
952mod tests {
953    use super::{
954        build_add_column_sql, build_alter_type_sql, build_drop_not_null_sql, on_conflict_clause,
955        pg_bind_text, pg_udt_to_json_schema, qualified_table_ref,
956    };
957    use serde_json::json;
958
959    #[test]
960    fn pg_add_column_ddl() {
961        let sql = build_add_column_sql("\"public\".\"t\"", "email", faucet_core::SqlBaseType::Text);
962        assert_eq!(
963            sql,
964            "ALTER TABLE \"public\".\"t\" ADD COLUMN IF NOT EXISTS \"email\" text"
965        );
966    }
967
968    #[test]
969    fn pg_widen_column_ddl() {
970        let sql = build_alter_type_sql(
971            "\"public\".\"t\"",
972            "score",
973            faucet_core::SqlBaseType::Double,
974        );
975        assert_eq!(
976            sql,
977            "ALTER TABLE \"public\".\"t\" ALTER COLUMN \"score\" TYPE double precision USING \"score\"::double precision"
978        );
979    }
980
981    #[test]
982    fn pg_drop_not_null_ddl() {
983        let sql = build_drop_not_null_sql("\"t\"", "created_at");
984        assert_eq!(
985            sql,
986            "ALTER TABLE \"t\" ALTER COLUMN \"created_at\" DROP NOT NULL"
987        );
988    }
989
990    #[test]
991    fn pg_udt_round_trips_to_json_schema() {
992        assert_eq!(
993            pg_udt_to_json_schema("int8", false),
994            json!({"type":"integer"})
995        );
996        assert_eq!(
997            pg_udt_to_json_schema("float8", false),
998            json!({"type":"number"})
999        );
1000        assert_eq!(
1001            pg_udt_to_json_schema("bool", false),
1002            json!({"type":"boolean"})
1003        );
1004        assert_eq!(
1005            pg_udt_to_json_schema("jsonb", false),
1006            json!({"type":"object"})
1007        );
1008        assert_eq!(
1009            pg_udt_to_json_schema("text", false),
1010            json!({"type":"string"})
1011        );
1012        // Unknown types fall back to string; nullable widens the type array.
1013        assert_eq!(
1014            pg_udt_to_json_schema("timestamptz", true),
1015            json!({"type":["string","null"]})
1016        );
1017    }
1018
1019    #[test]
1020    fn upsert_on_conflict_clause_for_keys() {
1021        let clause = on_conflict_clause(
1022            &["id".to_string()],
1023            &["id".to_string(), "name".to_string(), "email".to_string()],
1024        );
1025        assert_eq!(
1026            clause,
1027            r#"ON CONFLICT ("id") DO UPDATE SET "name" = EXCLUDED."name", "email" = EXCLUDED."email""#
1028        );
1029    }
1030
1031    #[test]
1032    fn upsert_on_conflict_all_columns_are_key_does_nothing() {
1033        let clause = on_conflict_clause(&["id".to_string()], &["id".to_string()]);
1034        assert_eq!(clause, r#"ON CONFLICT ("id") DO NOTHING"#);
1035    }
1036
1037    #[test]
1038    fn commit_token_table_is_the_shared_constant() {
1039        assert_eq!(
1040            faucet_core::idempotency::COMMIT_TOKEN_TABLE,
1041            "_faucet_commit_token"
1042        );
1043    }
1044
1045    // dataset_uri test is skipped: PostgresSink::new() requires a live pool
1046    // (connects to PostgreSQL in new()), and no offline constructor exists.
1047    // The URI format is covered by unit tests in faucet-core's redact tests.
1048
1049    #[test]
1050    fn qualified_table_ref_unqualified_is_bare_quoted_table() {
1051        // No schema → bare quoted table, resolved against the search_path.
1052        assert_eq!(qualified_table_ref(None, "events"), "\"events\"");
1053    }
1054
1055    #[test]
1056    fn qualified_table_ref_with_schema_is_schema_dot_table() {
1057        // With a schema → "schema"."table", so discovery and INSERT both
1058        // target the same explicit relation (#146 M13).
1059        assert_eq!(
1060            qualified_table_ref(Some("analytics"), "events"),
1061            "\"analytics\".\"events\""
1062        );
1063    }
1064
1065    #[test]
1066    fn qualified_table_ref_escapes_embedded_quotes() {
1067        // SQL-injection safety: embedded double-quotes are doubled.
1068        assert_eq!(
1069            qualified_table_ref(Some("we\"ird"), "ta\"ble"),
1070            "\"we\"\"ird\".\"ta\"\"ble\""
1071        );
1072    }
1073
1074    #[test]
1075    fn null_and_absent_bind_sql_null() {
1076        assert_eq!(pg_bind_text(None, "text"), None);
1077        assert_eq!(pg_bind_text(Some(&json!(null)), "int4"), None);
1078        assert_eq!(pg_bind_text(Some(&json!(null)), "jsonb"), None);
1079    }
1080
1081    #[test]
1082    fn scalars_bind_plain_text_for_typed_columns() {
1083        // The `$N::<udt>` cast parses these via the column's input function.
1084        assert_eq!(
1085            pg_bind_text(Some(&json!(42)), "int4").as_deref(),
1086            Some("42")
1087        );
1088        assert_eq!(
1089            pg_bind_text(Some(&json!(1.5)), "numeric").as_deref(),
1090            Some("1.5")
1091        );
1092        assert_eq!(
1093            pg_bind_text(Some(&json!(true)), "bool").as_deref(),
1094            Some("true")
1095        );
1096        assert_eq!(
1097            pg_bind_text(Some(&json!("2025-01-01T00:00:00Z")), "timestamptz").as_deref(),
1098            Some("2025-01-01T00:00:00Z")
1099        );
1100        // A plain string into TEXT keeps NO JSON quotes (the bug bound `"Bob"`).
1101        assert_eq!(
1102            pg_bind_text(Some(&json!("Bob")), "text").as_deref(),
1103            Some("Bob")
1104        );
1105        // Large u64 beyond i64 keeps exact text (no f64 precision loss).
1106        assert_eq!(
1107            pg_bind_text(Some(&json!(18446744073709551615u64)), "numeric").as_deref(),
1108            Some("18446744073709551615")
1109        );
1110    }
1111
1112    #[test]
1113    fn json_columns_get_json_text_with_quotes_preserved() {
1114        // For jsonb/json columns the value is bound as JSON text so the
1115        // `::jsonb` cast parses it: a string keeps its quotes, objects/arrays
1116        // round-trip.
1117        assert_eq!(
1118            pg_bind_text(Some(&json!("Bob")), "jsonb").as_deref(),
1119            Some("\"Bob\"")
1120        );
1121        assert_eq!(
1122            pg_bind_text(Some(&json!({"a": 1})), "jsonb").as_deref(),
1123            Some("{\"a\":1}")
1124        );
1125        assert_eq!(
1126            pg_bind_text(Some(&json!([1, 2])), "json").as_deref(),
1127            Some("[1,2]")
1128        );
1129        assert_eq!(pg_bind_text(Some(&json!(5)), "jsonb").as_deref(), Some("5"));
1130        // udt match is case-insensitive.
1131        assert_eq!(
1132            pg_bind_text(Some(&json!("x")), "JSONB").as_deref(),
1133            Some("\"x\"")
1134        );
1135    }
1136
1137    #[test]
1138    fn objects_into_non_json_columns_emit_json_text_so_the_cast_fails_loudly() {
1139        // No scalar text form for an object targeting e.g. an int column; the
1140        // `::int4` cast will reject this rather than silently coercing.
1141        assert_eq!(
1142            pg_bind_text(Some(&json!({"a": 1})), "int4").as_deref(),
1143            Some("{\"a\":1}")
1144        );
1145    }
1146}