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!(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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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); }
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 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 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 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 vec![records]
796 } else {
797 records.chunks(self.config.batch_size).collect()
798 };
799
800 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 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 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 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 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 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 #[test]
1050 fn qualified_table_ref_unqualified_is_bare_quoted_table() {
1051 assert_eq!(qualified_table_ref(None, "events"), "\"events\"");
1053 }
1054
1055 #[test]
1056 fn qualified_table_ref_with_schema_is_schema_dot_table() {
1057 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 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 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 assert_eq!(
1102 pg_bind_text(Some(&json!("Bob")), "text").as_deref(),
1103 Some("Bob")
1104 );
1105 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 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 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 assert_eq!(
1142 pg_bind_text(Some(&json!({"a": 1})), "int4").as_deref(),
1143 Some("{\"a\":1}")
1144 );
1145 }
1146}