1use 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
12pub(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 other => Some(other.to_string()),
43 }
44 }
45 }
46 }
47}
48
49fn 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
66fn 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
90fn 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
105fn 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
115fn 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
124fn 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
133fn 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
152pub struct PostgresSink {
154 config: PostgresSinkConfig,
155 pool: PgPool,
156}
157
158impl PostgresSink {
159 pub async fn new(config: PostgresSinkConfig) -> Result<Self, FaucetError> {
161 config.write.validate()?;
162 if matches!(
163 config.write.write_mode,
164 faucet_core::WriteMode::Upsert | faucet_core::WriteMode::Delete
165 ) && !matches!(config.column_mapping, PostgresColumnMapping::AutoMap)
166 {
167 return Err(FaucetError::Config(
168 "postgres sink: write_mode upsert/delete requires column_mapping: auto_map \
169 (key columns must be real columns, not inside a JSONB blob)"
170 .into(),
171 ));
172 }
173 if matches!(config.write_method, PostgresWriteMethod::Copy)
177 && matches!(
178 config.write.write_mode,
179 faucet_core::WriteMode::Upsert | faucet_core::WriteMode::Delete
180 )
181 {
182 return Err(FaucetError::Config(format!(
183 "postgres sink: write_method: copy is append-only (COPY has no ON CONFLICT); \
184 it cannot be combined with write_mode: {} — use write_method: insert",
185 config.write.write_mode.as_str()
186 )));
187 }
188
189 let pool = PgPoolOptions::new()
190 .max_connections(config.max_connections)
191 .connect(&config.connection_url)
192 .await
193 .map_err(|e| FaucetError::Sink(format!("PostgreSQL connection failed: {e}")))?;
194
195 Ok(Self { config, pool })
196 }
197
198 fn staging_table_name(&self) -> String {
201 format!("{}__faucet_ovw", self.config.table_name)
202 }
203
204 fn effective_table_name(&self) -> String {
209 if self.config.write.is_overwrite() {
210 self.staging_table_name()
211 } else {
212 self.config.table_name.clone()
213 }
214 }
215
216 async fn discover_columns(
222 &self,
223 conn: &mut sqlx::PgConnection,
224 table_ref: &str,
225 ) -> Result<Vec<(String, String)>, FaucetError> {
226 let columns: Vec<(String, String)> = sqlx::query(
227 "SELECT a.attname::text AS column_name, t.typname::text AS udt_name \
228 FROM pg_catalog.pg_attribute a \
229 JOIN pg_catalog.pg_type t ON t.oid = a.atttypid \
230 WHERE a.attrelid = to_regclass($1)::oid \
231 AND a.attnum > 0 AND NOT a.attisdropped \
232 ORDER BY a.attnum",
233 )
234 .bind(table_ref)
235 .fetch_all(&mut *conn)
236 .await
237 .map_err(|e| FaucetError::Sink(format!("failed to query table columns: {e}")))?
238 .iter()
239 .map(|row| {
240 (
241 row.get::<String, _>("column_name"),
242 row.get::<String, _>("udt_name"),
243 )
244 })
245 .collect();
246
247 if columns.is_empty() {
248 return Err(FaucetError::Sink(format!(
249 "table {table_ref} has no columns or does not exist"
250 )));
251 }
252 Ok(columns)
253 }
254
255 async fn copy_batch(
261 &self,
262 conn: &mut sqlx::PgConnection,
263 records: &[Value],
264 ) -> Result<usize, FaucetError> {
265 if records.is_empty() {
266 return Ok(0);
267 }
268 let table_ref =
269 qualified_table_ref(self.config.schema.as_deref(), &self.effective_table_name());
270
271 let (statement, payload) = match &self.config.column_mapping {
272 PostgresColumnMapping::Jsonb { column } => {
273 let payload = build_jsonb_payload(records);
274 (
275 copy_statement(&table_ref, std::slice::from_ref(column)),
276 payload,
277 )
278 }
279 PostgresColumnMapping::AutoMap => {
280 let columns = self.discover_columns(&mut *conn, &table_ref).await?;
281 let Some(payload) =
282 build_auto_map_payload(records, &columns).map_err(FaucetError::Sink)?
283 else {
284 return Ok(0);
285 };
286 (copy_statement(&table_ref, &payload.columns), payload)
287 }
288 };
289
290 let mut copy_in = conn
291 .copy_in_raw(&statement)
292 .await
293 .map_err(|e| FaucetError::Sink(format!("PostgreSQL COPY start failed: {e}")))?;
294 const SEND_CHUNK: usize = 1 << 20;
297 for chunk in payload.data.as_bytes().chunks(SEND_CHUNK) {
298 if let Err(e) = copy_in.send(chunk).await {
299 return Err(FaucetError::Sink(format!(
302 "PostgreSQL COPY send failed: {e}"
303 )));
304 }
305 }
306 copy_in
307 .finish()
308 .await
309 .map_err(|e| FaucetError::Sink(format!("PostgreSQL COPY failed: {e}")))?;
310 Ok(payload.rows)
311 }
312
313 async fn insert_jsonb(
320 &self,
321 conn: &mut sqlx::PgConnection,
322 records: &[Value],
323 column: &str,
324 ) -> Result<usize, FaucetError> {
325 if records.is_empty() {
326 return Ok(0);
327 }
328
329 let json_values: Vec<serde_json::Value> = records.to_vec();
331 let query = format!(
332 "INSERT INTO {} ({}) SELECT * FROM unnest($1::jsonb[])",
333 qualified_table_ref(self.config.schema.as_deref(), &self.effective_table_name()),
334 quote_ident(column)
335 );
336
337 sqlx::query(&query)
338 .bind(json_values)
339 .execute(&mut *conn)
340 .await
341 .map_err(|e| FaucetError::Sink(format!("PostgreSQL insert failed: {e}")))?;
342
343 Ok(records.len())
344 }
345
346 async fn insert_auto_map_with_conflict(
369 &self,
370 conn: &mut sqlx::PgConnection,
371 records: &[Value],
372 conflict_key: Option<&[String]>,
373 ) -> Result<usize, FaucetError> {
374 if records.is_empty() {
375 return Ok(0);
376 }
377
378 let table_ref =
392 qualified_table_ref(self.config.schema.as_deref(), &self.effective_table_name());
393 let columns = self.discover_columns(&mut *conn, &table_ref).await?;
394
395 let mut matched_rows: Vec<Vec<(&String, &String, &Value)>> =
402 Vec::with_capacity(records.len());
403 let mut used: std::collections::HashSet<&str> = std::collections::HashSet::new();
404
405 for record in records {
406 let obj = record
407 .as_object()
408 .ok_or_else(|| FaucetError::Sink("AutoMap requires JSON object records".into()))?;
409
410 let matching: Vec<(&String, &String, &Value)> = columns
411 .iter()
412 .filter_map(|(col, udt)| obj.get(col).map(|v| (col, udt, v)))
413 .collect();
414
415 if matching.is_empty() {
416 tracing::warn!(
417 record_keys = ?obj.keys().collect::<Vec<_>>(),
418 table_columns = ?columns,
419 "record has no keys matching table columns, skipping"
420 );
421 continue;
422 }
423
424 for (c, _, _) in &matching {
425 used.insert(c.as_str());
426 }
427 matched_rows.push(matching);
428 }
429
430 if matched_rows.is_empty() {
431 return Ok(0);
432 }
433
434 let insert_columns: Vec<(String, String)> = columns
437 .iter()
438 .filter(|(c, _)| used.contains(c.as_str()))
439 .cloned()
440 .collect();
441
442 let num_cols = insert_columns.len();
443 let num_rows = matched_rows.len();
444 let col_names: Vec<String> = insert_columns.iter().map(|(c, _)| quote_ident(c)).collect();
445
446 const MAX_PG_PARAMS: usize = 65535;
451 let max_rows_per_insert = (MAX_PG_PARAMS / num_cols).max(1);
452
453 for sub in matched_rows.chunks(max_rows_per_insert) {
454 let mut value_tuples: Vec<String> = Vec::with_capacity(sub.len());
458 for row_idx in 0..sub.len() {
459 let start = row_idx * num_cols + 1;
460 let placeholders: Vec<String> = (0..num_cols)
461 .map(|c| format!("${}::{}", start + c, insert_columns[c].1))
462 .collect();
463 value_tuples.push(format!("({})", placeholders.join(", ")));
464 }
465
466 let query = format!(
467 "INSERT INTO {} ({}) VALUES {}",
468 table_ref,
469 col_names.join(", "),
470 value_tuples.join(", ")
471 );
472 let query = match conflict_key {
473 Some(key) => format!(
474 "{query} {}",
475 on_conflict_clause(
476 key,
477 &insert_columns
478 .iter()
479 .map(|(c, _)| c.clone())
480 .collect::<Vec<_>>()
481 )
482 ),
483 None => query,
484 };
485
486 let mut q = sqlx::query(&query);
487 for matched in sub {
488 for (col, udt) in &insert_columns {
492 let val = matched
493 .iter()
494 .find(|(c, _, _)| *c == col)
495 .map(|(_, _, v)| *v);
496 q = q.bind(pg_bind_text(val, udt));
497 }
498 }
499
500 q.execute(&mut *conn)
501 .await
502 .map_err(|e| FaucetError::Sink(format!("PostgreSQL insert failed: {e}")))?;
503 }
504
505 Ok(num_rows)
506 }
507
508 async fn insert_auto_map(
515 &self,
516 conn: &mut sqlx::PgConnection,
517 records: &[Value],
518 ) -> Result<usize, FaucetError> {
519 self.insert_auto_map_with_conflict(conn, records, None)
520 .await
521 }
522
523 async fn delete_by_keys(
527 &self,
528 conn: &mut sqlx::PgConnection,
529 deletes: &[faucet_core::KeyTuple],
530 ) -> Result<usize, FaucetError> {
531 if deletes.is_empty() {
532 return Ok(0);
533 }
534 let key = &self.config.write.key;
535 let table_ref = qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name);
536
537 let udts: std::collections::HashMap<String, String> = self
540 .discover_columns(&mut *conn, &table_ref)
541 .await?
542 .into_iter()
543 .collect();
544 let key_udts: Vec<String> = key
545 .iter()
546 .map(|k| udts.get(k).cloned().unwrap_or_else(|| "text".to_string()))
547 .collect();
548 let col_list = key
549 .iter()
550 .map(|k| quote_ident(k))
551 .collect::<Vec<_>>()
552 .join(", ");
553
554 const MAX_PG_PARAMS: usize = 65535;
555 let per = (MAX_PG_PARAMS / key.len().max(1)).max(1);
556 let mut total = 0usize;
557 for chunk in deletes.chunks(per) {
558 let mut ph = 1usize;
559 let tuples: Vec<String> = chunk
560 .iter()
561 .map(|_| {
562 let group = key_udts
563 .iter()
564 .map(|udt| {
565 let p = format!("${ph}::{udt}");
566 ph += 1;
567 p
568 })
569 .collect::<Vec<_>>()
570 .join(", ");
571 format!("({group})")
572 })
573 .collect();
574 let sql = format!(
575 "DELETE FROM {table_ref} WHERE ({col_list}) IN ({})",
576 tuples.join(", ")
577 );
578 let mut q = sqlx::query(&sql);
579 for kt in chunk {
580 for ((_, v), udt) in kt.0.iter().zip(key_udts.iter()) {
581 q = q.bind(pg_bind_text(Some(v), udt));
582 }
583 }
584 let res = q
585 .execute(&mut *conn)
586 .await
587 .map_err(|e| FaucetError::Sink(format!("PostgreSQL delete failed: {e}")))?;
588 total += res.rows_affected() as usize;
589 }
590 Ok(total)
591 }
592
593 async fn cleanup_scope_impl(
605 &self,
606 scope: &std::collections::BTreeMap<String, Value>,
607 seen: &faucet_core::SeenKeys,
608 ) -> Result<u64, FaucetError> {
609 let key = &self.config.write.key;
610 if key.is_empty() {
611 return Err(FaucetError::Sink(
612 "cleanup requires a non-empty `key`".to_string(),
613 ));
614 }
615 let table_ref = qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name);
616
617 let mut tx = self
618 .pool
619 .begin()
620 .await
621 .map_err(|e| FaucetError::Sink(format!("PostgreSQL begin failed: {e}")))?;
622
623 let udts: std::collections::HashMap<String, String> = self
624 .discover_columns(&mut tx, &table_ref)
625 .await?
626 .into_iter()
627 .collect();
628
629 for col in scope.keys().chain(key.iter()) {
633 if !udts.contains_key(col) {
634 return Err(FaucetError::Sink(format!(
635 "cleanup: column '{col}' does not exist on {table_ref} — the \
636 completeness claim and `key` are in destination column terms"
637 )));
638 }
639 }
640 let udt_of = |c: &str| udts.get(c).cloned().unwrap_or_else(|| "text".to_string());
641
642 let temp_cols = key
646 .iter()
647 .map(|k| format!("{} {}", quote_ident(k), udt_of(k)))
648 .collect::<Vec<_>>()
649 .join(", ");
650 sqlx::query(&format!(
651 "CREATE TEMP TABLE faucet_cleanup_keys ({temp_cols}) ON COMMIT DROP"
652 ))
653 .execute(&mut *tx)
654 .await
655 .map_err(|e| FaucetError::Sink(format!("cleanup: temp table creation failed: {e}")))?;
656
657 const MAX_PG_PARAMS: usize = 65535;
659 let per = (MAX_PG_PARAMS / key.len()).max(1);
660 let key_udts: Vec<String> = key.iter().map(|k| udt_of(k)).collect();
661 let col_list = key
662 .iter()
663 .map(|k| quote_ident(k))
664 .collect::<Vec<_>>()
665 .join(", ");
666 for chunk in seen.keys().chunks(per) {
667 let mut ph = 1usize;
668 let tuples: Vec<String> = chunk
669 .iter()
670 .map(|_| {
671 let group = key_udts
672 .iter()
673 .map(|udt| {
674 let s = format!("${ph}::{udt}");
675 ph += 1;
676 s
677 })
678 .collect::<Vec<_>>()
679 .join(", ");
680 format!("({group})")
681 })
682 .collect();
683 let sql = format!(
684 "INSERT INTO faucet_cleanup_keys ({col_list}) VALUES {}",
685 tuples.join(", ")
686 );
687 let mut q = sqlx::query(&sql);
688 for kt in chunk {
689 for ((_, v), udt) in kt.0.iter().zip(key_udts.iter()) {
690 q = q.bind(pg_bind_text(Some(v), udt));
691 }
692 }
693 q.execute(&mut *tx)
694 .await
695 .map_err(|e| FaucetError::Sink(format!("cleanup: loading keys failed: {e}")))?;
696 }
697
698 let mut ph = 1usize;
700 let scope_pred = scope
701 .keys()
702 .map(|c| {
703 let s = format!("t.{} = ${}::{}", quote_ident(c), ph, udt_of(c));
704 ph += 1;
705 s
706 })
707 .collect::<Vec<_>>()
708 .join(" AND ");
709 let join_pred = key
710 .iter()
711 .map(|k| {
712 let q = quote_ident(k);
713 format!("c.{q} = t.{q}")
714 })
715 .collect::<Vec<_>>()
716 .join(" AND ");
717 let sql = format!(
718 "DELETE FROM {table_ref} t WHERE {scope_pred} \
719 AND NOT EXISTS (SELECT 1 FROM faucet_cleanup_keys c WHERE {join_pred})"
720 );
721 let mut q = sqlx::query(&sql);
722 for (col, v) in scope {
723 q = q.bind(pg_bind_text(Some(v), &udt_of(col)));
724 }
725 let res = q
726 .execute(&mut *tx)
727 .await
728 .map_err(|e| FaucetError::Sink(format!("cleanup: delete failed: {e}")))?;
729
730 tx.commit()
731 .await
732 .map_err(|e| FaucetError::Sink(format!("cleanup: commit failed: {e}")))?;
733 Ok(res.rows_affected())
734 }
735
736 async fn apply_plan(
738 &self,
739 conn: &mut sqlx::PgConnection,
740 plan: &faucet_core::WritePlan,
741 ) -> Result<usize, FaucetError> {
742 let mut affected = 0usize;
743 if !plan.upserts.is_empty() {
744 affected += self
745 .insert_auto_map_with_conflict(conn, &plan.upserts, Some(&self.config.write.key))
746 .await?;
747 }
748 if !plan.deletes.is_empty() {
749 affected += self.delete_by_keys(conn, &plan.deletes).await?;
750 }
751 Ok(affected)
752 }
753
754 async fn ensure_commit_table(&self) -> Result<(), FaucetError> {
756 let sql = format!(
757 "CREATE TABLE IF NOT EXISTS {t} ({s} TEXT PRIMARY KEY, {k} TEXT NOT NULL, updated_at TIMESTAMPTZ DEFAULT now())",
758 t = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TABLE),
759 s = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_SCOPE_COL),
760 k = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TOKEN_COL),
761 );
762 sqlx::query(&sql).execute(&self.pool).await.map_err(|e| {
763 FaucetError::Sink(format!("PostgreSQL commit-table create failed: {e}"))
764 })?;
765 Ok(())
766 }
767}
768
769#[async_trait]
770impl faucet_core::Sink for PostgresSink {
771 fn connector_name(&self) -> &'static str {
772 "postgres"
773 }
774
775 fn config_schema(&self) -> serde_json::Value {
776 serde_json::to_value(faucet_core::schema_for!(PostgresSinkConfig))
777 .expect("schema serialization")
778 }
779
780 fn supports_cleanup(&self) -> bool {
781 matches!(self.config.column_mapping, PostgresColumnMapping::AutoMap)
784 }
785
786 async fn cleanup_scope(
787 &self,
788 scope: &std::collections::BTreeMap<String, Value>,
789 seen: &faucet_core::SeenKeys,
790 ) -> Result<u64, FaucetError> {
791 self.cleanup_scope_impl(scope, seen).await
792 }
793
794 fn supported_write_modes(&self) -> &'static [faucet_core::WriteMode] {
795 &[
796 faucet_core::WriteMode::Append,
797 faucet_core::WriteMode::Upsert,
798 faucet_core::WriteMode::Delete,
799 faucet_core::WriteMode::Overwrite,
800 ]
801 }
802
803 fn is_overwrite(&self) -> bool {
804 self.config.write.is_overwrite()
805 }
806
807 async fn begin_overwrite(&self) -> Result<(), FaucetError> {
813 let staging =
814 qualified_table_ref(self.config.schema.as_deref(), &self.staging_table_name());
815 let target = qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name);
816 let mut conn = self
817 .pool
818 .acquire()
819 .await
820 .map_err(|e| FaucetError::Sink(format!("PostgreSQL pool acquire failed: {e}")))?;
821 sqlx::query(&format!("DROP TABLE IF EXISTS {staging}"))
822 .execute(&mut *conn)
823 .await
824 .map_err(|e| {
825 FaucetError::Sink(format!("postgres overwrite: drop stale staging: {e}"))
826 })?;
827 sqlx::query(&format!(
828 "CREATE TABLE {staging} (LIKE {target} INCLUDING DEFAULTS)"
829 ))
830 .execute(&mut *conn)
831 .await
832 .map_err(|e| {
833 FaucetError::Sink(format!(
834 "postgres overwrite: create staging from '{}' (does the table exist?): {e}",
835 self.config.table_name
836 ))
837 })?;
838 Ok(())
839 }
840
841 async fn commit_overwrite(&self) -> Result<(), FaucetError> {
848 let staging =
849 qualified_table_ref(self.config.schema.as_deref(), &self.staging_table_name());
850 let target = qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name);
851 let clear = match &self.config.scope {
853 Some(scope) => {
854 let col = quote_ident(scope.column());
855 format!(
856 "DELETE FROM {target} WHERE {}",
857 scope.render_where_literal(&col)
858 )
859 }
860 None => format!("TRUNCATE TABLE {target}"),
861 };
862 let mut tx = self
863 .pool
864 .begin()
865 .await
866 .map_err(|e| FaucetError::Sink(format!("postgres overwrite: begin swap: {e}")))?;
867 for stmt in [
868 clear,
869 format!("INSERT INTO {target} SELECT * FROM {staging}"),
870 format!("DROP TABLE {staging}"),
871 ] {
872 sqlx::query(&stmt)
873 .execute(&mut *tx)
874 .await
875 .map_err(|e| FaucetError::Sink(format!("postgres overwrite swap failed: {e}")))?;
876 }
877 tx.commit()
878 .await
879 .map_err(|e| FaucetError::Sink(format!("postgres overwrite: commit swap: {e}")))?;
880 Ok(())
881 }
882
883 async fn abort_overwrite(&self) -> Result<(), FaucetError> {
886 let staging =
887 qualified_table_ref(self.config.schema.as_deref(), &self.staging_table_name());
888 let mut conn = self
889 .pool
890 .acquire()
891 .await
892 .map_err(|e| FaucetError::Sink(format!("PostgreSQL pool acquire failed: {e}")))?;
893 sqlx::query(&format!("DROP TABLE IF EXISTS {staging}"))
894 .execute(&mut *conn)
895 .await
896 .map_err(|e| FaucetError::Sink(format!("postgres overwrite: drop staging: {e}")))?;
897 Ok(())
898 }
899
900 fn dedups_by_key(&self) -> bool {
901 self.config.write.dedups_by_key()
902 }
903
904 fn supports_schema_evolution(&self) -> bool {
905 true
906 }
907
908 async fn current_schema(&self) -> Result<Option<serde_json::Value>, FaucetError> {
916 let table_ref = qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name);
917 let rows: Vec<(String, String, bool)> = sqlx::query(
918 "SELECT a.attname::text AS column_name, t.typname::text AS udt_name, a.attnotnull \
919 FROM pg_catalog.pg_attribute a \
920 JOIN pg_catalog.pg_type t ON t.oid = a.atttypid \
921 WHERE a.attrelid = to_regclass($1)::oid \
922 AND a.attnum > 0 AND NOT a.attisdropped \
923 ORDER BY a.attnum",
924 )
925 .bind(&table_ref)
926 .fetch_all(&self.pool)
927 .await
928 .map_err(|e| FaucetError::Sink(format!("postgres current_schema query failed: {e}")))?
929 .iter()
930 .map(|row| {
931 (
932 row.get::<String, _>("column_name"),
933 row.get::<String, _>("udt_name"),
934 row.get::<bool, _>("attnotnull"),
935 )
936 })
937 .collect();
938
939 if rows.is_empty() {
940 return Ok(None); }
942
943 let mut props = serde_json::Map::new();
944 for (name, udt, notnull) in rows {
945 props.insert(name, pg_udt_to_json_schema(&udt, !notnull));
946 }
947 Ok(Some(
948 serde_json::json!({ "type": "object", "properties": props }),
949 ))
950 }
951
952 async fn evolve_schema(
957 &self,
958 evolution: &faucet_core::SchemaEvolution,
959 ) -> Result<(), FaucetError> {
960 let table_ref = qualified_table_ref(self.config.schema.as_deref(), &self.config.table_name);
961 let mut conn = self
962 .pool
963 .acquire()
964 .await
965 .map_err(|e| FaucetError::Sink(format!("postgres evolve acquire failed: {e}")))?;
966
967 for c in &evolution.additions {
968 let t =
969 faucet_core::json_schema_base_type(&c.to).unwrap_or(faucet_core::SqlBaseType::Text);
970 sqlx::query(&build_add_column_sql(&table_ref, &c.name, t))
971 .execute(&mut *conn)
972 .await
973 .map_err(|e| {
974 FaucetError::Sink(format!("postgres ADD COLUMN {} failed: {e}", c.name))
975 })?;
976 }
977 for c in &evolution.widenings {
978 let t =
979 faucet_core::json_schema_base_type(&c.to).unwrap_or(faucet_core::SqlBaseType::Text);
980 sqlx::query(&build_alter_type_sql(&table_ref, &c.name, t))
981 .execute(&mut *conn)
982 .await
983 .map_err(|e| {
984 FaucetError::Sink(format!("postgres ALTER TYPE {} failed: {e}", c.name))
985 })?;
986 }
987 for col in &evolution.relax_nullability {
988 sqlx::query(&build_drop_not_null_sql(&table_ref, col))
989 .execute(&mut *conn)
990 .await
991 .map_err(|e| {
992 FaucetError::Sink(format!("postgres DROP NOT NULL {col} failed: {e}"))
993 })?;
994 }
995 Ok(())
996 }
997
998 fn dataset_uri(&self) -> String {
999 let table = match &self.config.schema {
1000 Some(s) => format!("{}.{}", s, self.config.table_name),
1001 None => self.config.table_name.clone(),
1002 };
1003 format!(
1004 "{}?table={}",
1005 faucet_core::redact_uri_credentials(&self.config.connection_url),
1006 table
1007 )
1008 }
1009
1010 async fn check(
1016 &self,
1017 ctx: &faucet_core::check::CheckContext,
1018 ) -> Result<faucet_core::check::CheckReport, FaucetError> {
1019 use faucet_core::check::{CheckReport, Probe};
1020
1021 let started = std::time::Instant::now();
1022 let probe =
1023 match tokio::time::timeout(ctx.timeout, sqlx::query("SELECT 1").execute(&self.pool))
1024 .await
1025 {
1026 Ok(Ok(_)) => Probe::pass("auth", started.elapsed()),
1027 Ok(Err(e)) => Probe::fail_hint(
1028 "auth",
1029 started.elapsed(),
1030 e.to_string(),
1031 "check connection_url / credentials / that the database is reachable",
1032 ),
1033 Err(_) => Probe::fail_hint(
1034 "auth",
1035 started.elapsed(),
1036 "timed out",
1037 "check connection_url / credentials / that the database is reachable",
1038 ),
1039 };
1040 Ok(CheckReport::single(probe))
1041 }
1042
1043 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
1059 if records.is_empty() {
1060 return Ok(0);
1061 }
1062
1063 if matches!(
1064 self.config.write.write_mode,
1065 faucet_core::WriteMode::Upsert | faucet_core::WriteMode::Delete
1066 ) {
1067 let plan = faucet_core::plan_writes(records, &self.config.write);
1068 if let Some((idx, msg)) = plan.failed.first() {
1069 return Err(FaucetError::Sink(format!(
1070 "postgres {}: row {idx}: {msg}",
1071 self.config.write.write_mode.as_str()
1072 )));
1073 }
1074 let mut conn =
1075 self.pool.acquire().await.map_err(|e| {
1076 FaucetError::Sink(format!("PostgreSQL pool acquire failed: {e}"))
1077 })?;
1078 return self.apply_plan(&mut conn, &plan).await;
1079 }
1080 let chunks: Vec<&[Value]> = if self.config.batch_size == 0 {
1084 vec![records]
1088 } else {
1089 records.chunks(self.config.batch_size).collect()
1090 };
1091
1092 let mut conn = self
1095 .pool
1096 .acquire()
1097 .await
1098 .map_err(|e| FaucetError::Sink(format!("PostgreSQL pool acquire failed: {e}")))?;
1099
1100 let mut total = 0;
1101 for chunk in chunks {
1102 total += match self.config.write_method {
1103 PostgresWriteMethod::Copy => self.copy_batch(&mut conn, chunk).await?,
1106 PostgresWriteMethod::Insert => match &self.config.column_mapping {
1107 PostgresColumnMapping::Jsonb { column } => {
1108 self.insert_jsonb(&mut conn, chunk, column).await?
1109 }
1110 PostgresColumnMapping::AutoMap => {
1111 self.insert_auto_map(&mut conn, chunk).await?
1112 }
1113 },
1114 };
1115 }
1116
1117 tracing::info!(
1118 table = %self.config.table_name,
1119 rows = total,
1120 "PostgreSQL write complete"
1121 );
1122 Ok(total)
1123 }
1124
1125 async fn write_batch_partial(
1134 &self,
1135 records: &[Value],
1136 ) -> Result<Vec<faucet_core::RowOutcome>, FaucetError> {
1137 if !matches!(
1138 self.config.write.write_mode,
1139 faucet_core::WriteMode::Upsert | faucet_core::WriteMode::Delete
1140 ) {
1141 self.write_batch(records).await?;
1143 return Ok(records.iter().map(|_| Ok(())).collect());
1144 }
1145
1146 let plan = faucet_core::plan_writes(records, &self.config.write);
1147 let mut conn = self
1148 .pool
1149 .acquire()
1150 .await
1151 .map_err(|e| FaucetError::Sink(format!("PostgreSQL pool acquire failed: {e}")))?;
1152 self.apply_plan(&mut conn, &plan).await?;
1153
1154 let mut outcomes: Vec<faucet_core::RowOutcome> = records.iter().map(|_| Ok(())).collect();
1155 for (idx, msg) in &plan.failed {
1156 outcomes[*idx] = Err(FaucetError::Sink(format!(
1157 "postgres {}: {msg}",
1158 self.config.write.write_mode.as_str()
1159 )));
1160 }
1161 Ok(outcomes)
1162 }
1163
1164 fn supports_idempotent_writes(&self) -> bool {
1165 true
1166 }
1167
1168 async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
1169 self.ensure_commit_table().await?;
1170 let sql = format!(
1171 "SELECT {k} FROM {t} WHERE {s} = $1",
1172 t = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TABLE),
1173 k = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TOKEN_COL),
1174 s = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_SCOPE_COL),
1175 );
1176 let row = sqlx::query(&sql)
1177 .bind(scope)
1178 .fetch_optional(&self.pool)
1179 .await
1180 .map_err(|e| FaucetError::Sink(format!("PostgreSQL token read failed: {e}")))?;
1181 Ok(row.map(|r| r.get::<String, _>(0)))
1182 }
1183
1184 async fn write_batch_idempotent(
1185 &self,
1186 records: &[Value],
1187 scope: &str,
1188 token: &str,
1189 ) -> Result<usize, FaucetError> {
1190 self.ensure_commit_table().await?;
1191
1192 let plan = if matches!(self.config.write.write_mode, faucet_core::WriteMode::Append) {
1195 None
1196 } else {
1197 let plan = faucet_core::plan_writes(records, &self.config.write);
1198 if let Some((idx, msg)) = plan.failed.first() {
1199 return Err(FaucetError::Sink(format!(
1200 "postgres {}: row {idx}: {msg}",
1201 self.config.write.write_mode.as_str()
1202 )));
1203 }
1204 Some(plan)
1205 };
1206
1207 let mut tx =
1208 self.pool.begin().await.map_err(|e| {
1209 FaucetError::Sink(format!("PostgreSQL transaction begin failed: {e}"))
1210 })?;
1211
1212 let written = match &plan {
1218 Some(plan) => self.apply_plan(&mut tx, plan).await?,
1219 None => match &self.config.column_mapping {
1220 PostgresColumnMapping::Jsonb { column } => {
1221 self.insert_jsonb(&mut tx, records, column).await?
1222 }
1223 PostgresColumnMapping::AutoMap => self.insert_auto_map(&mut tx, records).await?,
1224 },
1225 };
1226
1227 let upsert = format!(
1228 "INSERT INTO {t} ({s}, {k}) VALUES ($1, $2) ON CONFLICT ({s}) DO UPDATE SET {k} = EXCLUDED.{k}, updated_at = now()",
1229 t = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TABLE),
1230 s = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_SCOPE_COL),
1231 k = quote_ident(faucet_core::idempotency::COMMIT_TOKEN_TOKEN_COL),
1232 );
1233 sqlx::query(&upsert)
1234 .bind(scope)
1235 .bind(token)
1236 .execute(&mut *tx)
1237 .await
1238 .map_err(|e| FaucetError::Sink(format!("PostgreSQL token upsert failed: {e}")))?;
1239
1240 tx.commit()
1241 .await
1242 .map_err(|e| FaucetError::Sink(format!("PostgreSQL transaction commit failed: {e}")))?;
1243 Ok(written)
1244 }
1245}
1246
1247#[cfg(test)]
1248mod tests {
1249 use super::{
1250 build_add_column_sql, build_alter_type_sql, build_drop_not_null_sql, on_conflict_clause,
1251 pg_bind_text, pg_udt_to_json_schema, qualified_table_ref,
1252 };
1253 use serde_json::json;
1254
1255 #[test]
1256 fn pg_add_column_ddl() {
1257 let sql = build_add_column_sql("\"public\".\"t\"", "email", faucet_core::SqlBaseType::Text);
1258 assert_eq!(
1259 sql,
1260 "ALTER TABLE \"public\".\"t\" ADD COLUMN IF NOT EXISTS \"email\" text"
1261 );
1262 }
1263
1264 #[test]
1265 fn pg_widen_column_ddl() {
1266 let sql = build_alter_type_sql(
1267 "\"public\".\"t\"",
1268 "score",
1269 faucet_core::SqlBaseType::Double,
1270 );
1271 assert_eq!(
1272 sql,
1273 "ALTER TABLE \"public\".\"t\" ALTER COLUMN \"score\" TYPE double precision USING \"score\"::double precision"
1274 );
1275 }
1276
1277 #[test]
1278 fn pg_drop_not_null_ddl() {
1279 let sql = build_drop_not_null_sql("\"t\"", "created_at");
1280 assert_eq!(
1281 sql,
1282 "ALTER TABLE \"t\" ALTER COLUMN \"created_at\" DROP NOT NULL"
1283 );
1284 }
1285
1286 #[test]
1287 fn pg_udt_round_trips_to_json_schema() {
1288 assert_eq!(
1289 pg_udt_to_json_schema("int8", false),
1290 json!({"type":"integer"})
1291 );
1292 assert_eq!(
1293 pg_udt_to_json_schema("float8", false),
1294 json!({"type":"number"})
1295 );
1296 assert_eq!(
1297 pg_udt_to_json_schema("bool", false),
1298 json!({"type":"boolean"})
1299 );
1300 assert_eq!(
1301 pg_udt_to_json_schema("jsonb", false),
1302 json!({"type":"object"})
1303 );
1304 assert_eq!(
1305 pg_udt_to_json_schema("text", false),
1306 json!({"type":"string"})
1307 );
1308 assert_eq!(
1310 pg_udt_to_json_schema("timestamptz", true),
1311 json!({"type":["string","null"]})
1312 );
1313 }
1314
1315 #[test]
1316 fn upsert_on_conflict_clause_for_keys() {
1317 let clause = on_conflict_clause(
1318 &["id".to_string()],
1319 &["id".to_string(), "name".to_string(), "email".to_string()],
1320 );
1321 assert_eq!(
1322 clause,
1323 r#"ON CONFLICT ("id") DO UPDATE SET "name" = EXCLUDED."name", "email" = EXCLUDED."email""#
1324 );
1325 }
1326
1327 #[test]
1328 fn upsert_on_conflict_all_columns_are_key_does_nothing() {
1329 let clause = on_conflict_clause(&["id".to_string()], &["id".to_string()]);
1330 assert_eq!(clause, r#"ON CONFLICT ("id") DO NOTHING"#);
1331 }
1332
1333 #[test]
1334 fn commit_token_table_is_the_shared_constant() {
1335 assert_eq!(
1336 faucet_core::idempotency::COMMIT_TOKEN_TABLE,
1337 "_faucet_commit_token"
1338 );
1339 }
1340
1341 #[test]
1346 fn qualified_table_ref_unqualified_is_bare_quoted_table() {
1347 assert_eq!(qualified_table_ref(None, "events"), "\"events\"");
1349 }
1350
1351 #[test]
1352 fn qualified_table_ref_with_schema_is_schema_dot_table() {
1353 assert_eq!(
1356 qualified_table_ref(Some("analytics"), "events"),
1357 "\"analytics\".\"events\""
1358 );
1359 }
1360
1361 #[test]
1362 fn qualified_table_ref_escapes_embedded_quotes() {
1363 assert_eq!(
1365 qualified_table_ref(Some("we\"ird"), "ta\"ble"),
1366 "\"we\"\"ird\".\"ta\"\"ble\""
1367 );
1368 }
1369
1370 #[test]
1371 fn null_and_absent_bind_sql_null() {
1372 assert_eq!(pg_bind_text(None, "text"), None);
1373 assert_eq!(pg_bind_text(Some(&json!(null)), "int4"), None);
1374 assert_eq!(pg_bind_text(Some(&json!(null)), "jsonb"), None);
1375 }
1376
1377 #[test]
1378 fn scalars_bind_plain_text_for_typed_columns() {
1379 assert_eq!(
1381 pg_bind_text(Some(&json!(42)), "int4").as_deref(),
1382 Some("42")
1383 );
1384 assert_eq!(
1385 pg_bind_text(Some(&json!(1.5)), "numeric").as_deref(),
1386 Some("1.5")
1387 );
1388 assert_eq!(
1389 pg_bind_text(Some(&json!(true)), "bool").as_deref(),
1390 Some("true")
1391 );
1392 assert_eq!(
1393 pg_bind_text(Some(&json!("2025-01-01T00:00:00Z")), "timestamptz").as_deref(),
1394 Some("2025-01-01T00:00:00Z")
1395 );
1396 assert_eq!(
1398 pg_bind_text(Some(&json!("Bob")), "text").as_deref(),
1399 Some("Bob")
1400 );
1401 assert_eq!(
1403 pg_bind_text(Some(&json!(18446744073709551615u64)), "numeric").as_deref(),
1404 Some("18446744073709551615")
1405 );
1406 }
1407
1408 #[test]
1409 fn json_columns_get_json_text_with_quotes_preserved() {
1410 assert_eq!(
1414 pg_bind_text(Some(&json!("Bob")), "jsonb").as_deref(),
1415 Some("\"Bob\"")
1416 );
1417 assert_eq!(
1418 pg_bind_text(Some(&json!({"a": 1})), "jsonb").as_deref(),
1419 Some("{\"a\":1}")
1420 );
1421 assert_eq!(
1422 pg_bind_text(Some(&json!([1, 2])), "json").as_deref(),
1423 Some("[1,2]")
1424 );
1425 assert_eq!(pg_bind_text(Some(&json!(5)), "jsonb").as_deref(), Some("5"));
1426 assert_eq!(
1428 pg_bind_text(Some(&json!("x")), "JSONB").as_deref(),
1429 Some("\"x\"")
1430 );
1431 }
1432
1433 #[test]
1434 fn objects_into_non_json_columns_emit_json_text_so_the_cast_fails_loudly() {
1435 assert_eq!(
1438 pg_bind_text(Some(&json!({"a": 1})), "int4").as_deref(),
1439 Some("{\"a\":1}")
1440 );
1441 }
1442}