1use 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
11fn 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 other => Some(other.to_string()),
42 }
43 }
44 }
45 }
46}
47
48fn 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
65fn 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
89pub struct PostgresSink {
91 config: PostgresSinkConfig,
92 pool: PgPool,
93}
94
95impl PostgresSink {
96 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 vec![records]
569 } else {
570 records.chunks(self.config.batch_size).collect()
571 };
572
573 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 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 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 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 #[test]
753 fn qualified_table_ref_unqualified_is_bare_quoted_table() {
754 assert_eq!(qualified_table_ref(None, "events"), "\"events\"");
756 }
757
758 #[test]
759 fn qualified_table_ref_with_schema_is_schema_dot_table() {
760 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 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 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 assert_eq!(
805 pg_bind_text(Some(&json!("Bob")), "text").as_deref(),
806 Some("Bob")
807 );
808 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 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 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 assert_eq!(
845 pg_bind_text(Some(&json!({"a": 1})), "int4").as_deref(),
846 Some("{\"a\":1}")
847 );
848 }
849}