1use crate::error::FaucetError;
16use crate::traits::{RowOutcome, Sink};
17use chrono::Utc;
18use schemars::JsonSchema;
19use serde::{Deserialize, Serialize};
20use serde_json::{Map, Value};
21use std::sync::atomic::{AtomicU64, Ordering};
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
25#[serde(rename_all = "snake_case")]
26pub enum MetadataColumn {
27 ExtractedAt,
30 LoadedAt,
32 RunId,
34 Source,
36 Sequence,
38}
39
40impl MetadataColumn {
41 pub fn suffix(self) -> &'static str {
43 match self {
44 MetadataColumn::ExtractedAt => "extracted_at",
45 MetadataColumn::LoadedAt => "loaded_at",
46 MetadataColumn::RunId => "run_id",
47 MetadataColumn::Source => "source",
48 MetadataColumn::Sequence => "sequence",
49 }
50 }
51}
52
53fn default_prefix() -> String {
54 "_faucet".to_owned()
55}
56fn default_true() -> bool {
57 true
58}
59
60#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
62#[serde(deny_unknown_fields)]
63pub struct MetadataColumnsSpec {
64 #[serde(default = "default_true")]
66 pub enabled: bool,
67 #[serde(default = "default_prefix")]
69 pub prefix: String,
70 #[serde(default)]
73 pub columns: Vec<MetadataColumn>,
74}
75
76impl Default for MetadataColumnsSpec {
77 fn default() -> Self {
78 Self {
79 enabled: true,
80 prefix: default_prefix(),
81 columns: Vec::new(),
82 }
83 }
84}
85
86const DEFAULT_COLUMNS: &[MetadataColumn] = &[
87 MetadataColumn::ExtractedAt,
88 MetadataColumn::LoadedAt,
89 MetadataColumn::RunId,
90 MetadataColumn::Source,
91];
92
93#[derive(Debug, Clone)]
95pub struct CompiledMetadata {
96 prefix: String,
97 columns: Vec<MetadataColumn>,
98}
99
100impl CompiledMetadata {
101 pub fn compile(spec: &MetadataColumnsSpec) -> Result<Option<Self>, FaucetError> {
104 if !spec.enabled {
105 return Ok(None);
106 }
107 if spec.prefix.trim().is_empty() {
108 return Err(FaucetError::Config(
109 "metadata_columns: `prefix` must not be empty".into(),
110 ));
111 }
112 let columns = if spec.columns.is_empty() {
113 DEFAULT_COLUMNS.to_vec()
114 } else {
115 spec.columns.clone()
116 };
117 Ok(Some(Self {
118 prefix: spec.prefix.clone(),
119 columns,
120 }))
121 }
122
123 fn column_name(&self, col: MetadataColumn) -> String {
124 format!("{}_{}", self.prefix, col.suffix())
125 }
126}
127
128#[derive(Debug, Clone)]
130pub struct MetadataContext {
131 pub run_id: String,
133 pub source: String,
135}
136
137pub struct MetadataSink {
140 inner: Box<dyn Sink>,
141 meta: CompiledMetadata,
142 ctx: MetadataContext,
143 seq: AtomicU64,
144 collided: std::sync::Mutex<std::collections::BTreeSet<String>>,
146 columns_ready: tokio::sync::OnceCell<()>,
149}
150
151impl std::fmt::Debug for MetadataSink {
152 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
153 f.debug_struct("MetadataSink")
154 .field("prefix", &self.meta.prefix)
155 .field("columns", &self.meta.columns)
156 .finish()
157 }
158}
159
160impl MetadataSink {
161 pub fn new(inner: Box<dyn Sink>, meta: CompiledMetadata, ctx: MetadataContext) -> Self {
163 Self {
164 inner,
165 meta,
166 ctx,
167 seq: AtomicU64::new(0),
168 collided: std::sync::Mutex::new(std::collections::BTreeSet::new()),
169 columns_ready: tokio::sync::OnceCell::new(),
170 }
171 }
172
173 fn stamp(&self, records: &[Value]) -> Vec<Value> {
176 self.stamp_with(records, true)
177 }
178
179 fn preview(&self, records: &[Value]) -> Vec<Value> {
182 self.stamp_with(records, false)
183 }
184
185 fn stamp_with(&self, records: &[Value], advance: bool) -> Vec<Value> {
186 let now = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
187 let mut next = self.seq.load(Ordering::Relaxed);
188 records
189 .iter()
190 .map(|rec| match rec {
191 Value::Object(map) => {
192 let mut map = map.clone();
193 let seq = if advance {
194 self.seq.fetch_add(1, Ordering::Relaxed)
195 } else {
196 next += 1;
197 next - 1
198 };
199 self.inject(&mut map, &now, seq, advance);
200 Value::Object(map)
201 }
202 other => other.clone(),
203 })
204 .collect()
205 }
206
207 fn inject(&self, map: &mut Map<String, Value>, now: &str, seq: u64, note: bool) {
208 for &col in &self.meta.columns {
209 let name = self.meta.column_name(col);
210 let value = match col {
211 MetadataColumn::ExtractedAt | MetadataColumn::LoadedAt => {
212 Value::String(now.to_owned())
213 }
214 MetadataColumn::RunId => Value::String(self.ctx.run_id.clone()),
215 MetadataColumn::Source => Value::String(self.ctx.source.clone()),
216 MetadataColumn::Sequence => Value::from(seq),
217 };
218 if map.insert(name.clone(), value).is_some() && note {
219 self.note_collision(name);
220 }
221 }
222 }
223
224 async fn ensure_columns(&self) -> Result<(), FaucetError> {
232 self.columns_ready
233 .get_or_try_init(|| async {
234 let Some(schema) = self.inner.current_schema().await? else {
235 return Ok(());
236 };
237 let props = schema.get("properties").and_then(Value::as_object);
238 let additions: Vec<crate::drift::ColumnChange> = self
239 .meta
240 .columns
241 .iter()
242 .map(|&c| (self.meta.column_name(c), c))
243 .filter(|(name, _)| props.is_none_or(|p| !p.contains_key(name)))
244 .map(|(name, c)| crate::drift::ColumnChange {
245 name,
246 from: None,
247 to: serde_json::json!({ "type": [
248 if c == MetadataColumn::Sequence { "integer" } else { "string" },
249 "null"
250 ] }),
251 })
252 .collect();
253 if additions.is_empty() {
254 return Ok(());
255 }
256 let names: Vec<String> = additions.iter().map(|a| a.name.clone()).collect();
257 if !self.inner.supports_schema_evolution() {
258 return Err(FaucetError::Config(format!(
259 "metadata_columns: the destination has no column {} and sink '{}' \
260 cannot add columns; add them to the table or disable metadata_columns",
261 names.join(", "),
262 self.inner.connector_name()
263 )));
264 }
265 tracing::info!(columns = ?names, "adding missing metadata columns to the destination");
266 self.inner
267 .evolve_schema(&crate::drift::SchemaEvolution {
268 additions,
269 ..Default::default()
270 })
271 .await
272 })
273 .await
274 .map(|_| ())
275 }
276
277 fn note_collision(&self, name: String) {
280 let Ok(mut seen) = self.collided.lock() else {
281 return;
282 };
283 if seen.insert(name.clone()) {
284 tracing::warn!(
285 column = %name,
286 "a source field has the same name as a metadata column; the metadata \
287 value replaces it (choose another `metadata_columns.prefix` to keep it)"
288 );
289 }
290 }
291
292 #[cfg(test)]
293 fn collisions(&self) -> Vec<String> {
294 self.collided.lock().unwrap().iter().cloned().collect()
295 }
296}
297
298#[async_trait::async_trait]
299impl Sink for MetadataSink {
300 crate::forward_sink!(inner; except
303 admit_page, write_batch, write_batch_reported, write_governed, write_batch_partial,
304 write_batch_idempotent, write_batch_idempotent_reported,
305 begin_overwrite,
306 supports_columnar, write_batch_columnar, native_load_capabilities, load_native);
307
308 async fn admit_page(&self, records: &[Value]) -> Result<(), FaucetError> {
309 self.inner.admit_page(&self.preview(records)).await
310 }
311 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
312 self.ensure_columns().await?;
313 self.inner.write_batch(&self.stamp(records)).await
314 }
315
316 async fn write_batch_partial(&self, records: &[Value]) -> Result<Vec<RowOutcome>, FaucetError> {
317 self.ensure_columns().await?;
318 self.inner.write_batch_partial(&self.stamp(records)).await
319 }
320
321 async fn write_batch_reported(
322 &self,
323 records: &[Value],
324 ) -> Result<crate::accounting::WriteReport, FaucetError> {
325 self.ensure_columns().await?;
326 self.inner.write_batch_reported(&self.stamp(records)).await
327 }
328
329 async fn write_governed(
330 &self,
331 page: &crate::governance::GovernedPage<'_>,
332 ) -> Result<crate::accounting::WriteReport, FaucetError> {
333 self.ensure_columns().await?;
334 self.inner
335 .write_governed(&page.restamped(self.stamp(page.records())))
336 .await
337 }
338
339 #[allow(deprecated)] async fn write_batch_idempotent(
341 &self,
342 records: &[Value],
343 scope: &str,
344 token: &str,
345 ) -> Result<usize, FaucetError> {
346 self.ensure_columns().await?;
347 self.inner
348 .write_batch_idempotent(&self.stamp(records), scope, token)
349 .await
350 }
351
352 #[allow(deprecated)] async fn write_batch_idempotent_reported(
354 &self,
355 records: &[Value],
356 scope: &str,
357 token: &str,
358 ) -> Result<crate::accounting::WriteReport, FaucetError> {
359 self.ensure_columns().await?;
360 self.inner
361 .write_batch_idempotent_reported(&self.stamp(records), scope, token)
362 .await
363 }
364
365 #[allow(deprecated)] async fn begin_overwrite(&self) -> Result<(), FaucetError> {
367 self.ensure_columns().await?;
369 self.inner.begin_overwrite().await
370 }
371}
372
373#[allow(deprecated)] #[cfg(test)]
375mod tests {
376 use super::*;
377
378 struct Mapped {
379 evolutions: std::sync::Mutex<Vec<Vec<String>>>,
380 evolves: bool,
381 }
382 #[async_trait::async_trait]
383 impl Sink for Mapped {
384 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
385 Ok(records.len())
386 }
387 async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
388 Ok(Some(json!({"type": "object", "properties": {
389 "id": {"type": "integer"}, "_faucet_source": {"type": "string"}
390 }})))
391 }
392 fn supports_schema_evolution(&self) -> bool {
393 self.evolves
394 }
395 async fn evolve_schema(
396 &self,
397 evolution: &crate::drift::SchemaEvolution,
398 ) -> Result<(), FaucetError> {
399 self.evolutions.lock().unwrap().push(
400 evolution
401 .additions
402 .iter()
403 .map(|a| format!("{}:{}", a.name, a.to["type"][0]))
404 .collect(),
405 );
406 Ok(())
407 }
408 }
409
410 #[tokio::test]
411 async fn missing_metadata_columns_are_added_once_before_the_first_write() {
412 let meta = CompiledMetadata::compile(&spec(&[
413 MetadataColumn::RunId,
414 MetadataColumn::Source,
415 MetadataColumn::Sequence,
416 ]))
417 .unwrap()
418 .unwrap();
419 let inner = std::sync::Arc::new(Mapped {
420 evolutions: std::sync::Mutex::new(Vec::new()),
421 evolves: true,
422 });
423 struct Shared(std::sync::Arc<Mapped>);
424 #[async_trait::async_trait]
425 impl Sink for Shared {
426 async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
427 self.0.write_batch(r).await
428 }
429 async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
430 self.0.current_schema().await
431 }
432 fn supports_schema_evolution(&self) -> bool {
433 self.0.supports_schema_evolution()
434 }
435 async fn evolve_schema(
436 &self,
437 e: &crate::drift::SchemaEvolution,
438 ) -> Result<(), FaucetError> {
439 self.0.evolve_schema(e).await
440 }
441 }
442 let sink = MetadataSink::new(Box::new(Shared(inner.clone())), meta, ctx());
443 sink.write_batch(&[json!({"id": 1})]).await.unwrap();
444 sink.write_batch(&[json!({"id": 2})]).await.unwrap();
445 assert_eq!(
446 *inner.evolutions.lock().unwrap(),
447 vec![vec![
448 "_faucet_run_id:\"string\"".to_string(),
449 "_faucet_sequence:\"integer\"".to_string()
450 ]]
451 );
452
453 let meta = CompiledMetadata::compile(&spec(&[MetadataColumn::RunId]))
454 .unwrap()
455 .unwrap();
456 let rigid = MetadataSink::new(
457 Box::new(Mapped {
458 evolutions: std::sync::Mutex::new(Vec::new()),
459 evolves: false,
460 }),
461 meta,
462 ctx(),
463 );
464 let err = rigid.write_batch(&[json!({"id": 1})]).await.unwrap_err();
465 assert!(err.to_string().contains("_faucet_run_id"), "{err}");
466 }
467
468 #[test]
469 fn a_colliding_source_field_is_reported() {
470 let meta = CompiledMetadata::compile(&spec(&[MetadataColumn::RunId]))
471 .unwrap()
472 .unwrap();
473 let sink = MetadataSink::new(Box::new(CapturingSink::default()), meta, ctx());
474 let out = sink.stamp(&[
475 json!({"_faucet_run_id": "upstream", "a": 1}),
476 json!({"a": 2}),
477 ]);
478 assert_eq!(out[0]["_faucet_run_id"], json!("run-42"));
479 assert_eq!(sink.collisions(), vec!["_faucet_run_id".to_string()]);
480 sink.stamp(&[json!({"_faucet_run_id": "again"})]);
481 assert_eq!(sink.collisions().len(), 1);
482 }
483
484 #[tokio::test]
485 async fn metadata_sink_forwards_every_hook_but_the_unstamped_paths() {
486 let probe = crate::forwarding::HookSink::default();
487 let meta = CompiledMetadata::compile(&spec(&[MetadataColumn::RunId]))
488 .unwrap()
489 .unwrap();
490 let sink = MetadataSink::new(Box::new(probe.clone()), meta, ctx());
491 crate::forwarding::assert_sink_forwards_every_hook(
492 &sink,
493 &probe,
494 &[
495 "supports_columnar",
496 "write_batch_columnar",
497 "native_load_capabilities",
498 "load_native",
499 ],
500 )
501 .await;
502 }
503 use serde_json::json;
504 use std::sync::Mutex;
505
506 #[derive(Debug, Default)]
507 struct CapturingSink {
508 rows: Mutex<Vec<Value>>,
509 }
510 #[async_trait::async_trait]
511 impl Sink for CapturingSink {
512 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
513 self.rows.lock().unwrap().extend_from_slice(records);
514 Ok(records.len())
515 }
516 fn config_schema(&self) -> Value {
517 json!({})
518 }
519 }
520
521 fn spec(cols: &[MetadataColumn]) -> MetadataColumnsSpec {
522 MetadataColumnsSpec {
523 enabled: true,
524 prefix: "_faucet".into(),
525 columns: cols.to_vec(),
526 }
527 }
528
529 fn ctx() -> MetadataContext {
530 MetadataContext {
531 run_id: "run-42".into(),
532 source: "rest".into(),
533 }
534 }
535
536 #[test]
537 fn compile_resolves_defaults_and_respects_disable() {
538 let c = CompiledMetadata::compile(&spec(&[])).unwrap().unwrap();
539 assert_eq!(c.columns, DEFAULT_COLUMNS.to_vec());
540 let disabled = MetadataColumnsSpec {
542 enabled: false,
543 ..Default::default()
544 };
545 assert!(CompiledMetadata::compile(&disabled).unwrap().is_none());
546 let bad = MetadataColumnsSpec {
548 prefix: " ".into(),
549 ..Default::default()
550 };
551 assert!(CompiledMetadata::compile(&bad).is_err());
552 }
553
554 #[tokio::test]
555 async fn stamps_all_column_kinds_with_prefix() {
556 let meta = CompiledMetadata::compile(&spec(&[
557 MetadataColumn::ExtractedAt,
558 MetadataColumn::LoadedAt,
559 MetadataColumn::RunId,
560 MetadataColumn::Source,
561 MetadataColumn::Sequence,
562 ]))
563 .unwrap()
564 .unwrap();
565 let inner = Box::new(CapturingSink::default());
566 let sink = MetadataSink::new(inner, meta, ctx());
567 assert_eq!(
568 sink.batch_atomicity(),
569 crate::dlq::BatchAtomicity::BestEffort
570 );
571 let n = sink
572 .write_batch(&[json!({"id": 1}), json!({"id": 2})])
573 .await
574 .unwrap();
575 assert_eq!(n, 2);
576 let stamped = sink.stamp(&[json!({"id": 1}), json!({"id": 2})]);
579 let r0 = &stamped[0];
580 assert_eq!(r0["id"], 1);
581 assert_eq!(r0["_faucet_run_id"], "run-42");
582 assert_eq!(r0["_faucet_source"], "rest");
583 assert!(r0["_faucet_extracted_at"].is_string());
584 assert!(r0["_faucet_loaded_at"].is_string());
585 assert_eq!(
587 stamped[0]["_faucet_sequence"].as_u64().unwrap() + 1,
588 stamped[1]["_faucet_sequence"].as_u64().unwrap()
589 );
590 }
591
592 #[tokio::test]
596 async fn admit_page_forwards_the_stamped_page() {
597 #[derive(Default)]
598 struct Admits(std::sync::Mutex<Vec<Value>>);
599 #[async_trait::async_trait]
600 impl Sink for Admits {
601 async fn admit_page(&self, records: &[Value]) -> Result<(), FaucetError> {
602 self.0.lock().unwrap().extend_from_slice(records);
603 Ok(())
604 }
605 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
606 Ok(records.len())
607 }
608 }
609 let meta =
610 CompiledMetadata::compile(&spec(&[MetadataColumn::RunId, MetadataColumn::Sequence]))
611 .unwrap()
612 .unwrap();
613 let inner = std::sync::Arc::new(Admits::default());
614 struct Shared(std::sync::Arc<Admits>);
615 #[async_trait::async_trait]
616 impl Sink for Shared {
617 async fn admit_page(&self, records: &[Value]) -> Result<(), FaucetError> {
618 self.0.admit_page(records).await
619 }
620 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
621 self.0.write_batch(records).await
622 }
623 }
624 let sink = MetadataSink::new(Box::new(Shared(std::sync::Arc::clone(&inner))), meta, ctx());
625 sink.admit_page(&[json!({"id": 1}), json!("x")])
626 .await
627 .unwrap();
628 let admitted = inner.0.lock().unwrap().clone();
629 assert_eq!(admitted[0]["_faucet_run_id"], "run-42");
630 assert_eq!(admitted[0]["_faucet_sequence"], 0);
631 assert_eq!(admitted[1], json!("x"));
632 let stamped = sink.stamp(&[json!({"id": 1})]);
633 assert_eq!(
634 stamped[0]["_faucet_sequence"], 0,
635 "admission consumed no sequence"
636 );
637 }
638
639 #[test]
640 fn non_object_records_pass_through() {
641 let meta = CompiledMetadata::compile(&spec(&[MetadataColumn::RunId]))
642 .unwrap()
643 .unwrap();
644 let sink = MetadataSink::new(Box::new(CapturingSink::default()), meta, ctx());
645 let out = sink.stamp(&[json!("scalar"), json!(42)]);
646 assert_eq!(out, vec![json!("scalar"), json!(42)]);
647 }
648
649 #[tokio::test]
650 async fn delegates_every_sink_method_to_inner() {
651 let meta = CompiledMetadata::compile(&spec(&[MetadataColumn::RunId]))
652 .unwrap()
653 .unwrap();
654 let sink = MetadataSink::new(Box::new(CapturingSink::default()), meta, ctx());
655
656 assert_eq!(sink.write_batch(&[json!({"a": 1})]).await.unwrap(), 1);
658 assert_eq!(
659 sink.write_batch_partial(&[json!({"a": 1})])
660 .await
661 .unwrap()
662 .len(),
663 1
664 );
665 let _ = sink
666 .write_batch_idempotent(&[json!({"a": 1})], "scope", "tok")
667 .await;
668 sink.flush().await.unwrap();
669
670 let _ = sink.check(&crate::check::CheckContext::default()).await;
672 assert!(!sink.supports_cleanup());
673 let _ = sink
674 .cleanup_scope(
675 &std::collections::BTreeMap::new(),
676 &crate::cleanup::SeenKeys::new(),
677 )
678 .await;
679 assert!(!sink.supports_idempotent_writes());
680 assert!(sink.last_committed_token("s").await.unwrap().is_none());
681 let _ = sink.supported_write_modes();
682 let _ = sink.dedups_by_key();
683 let _ = sink.sink_guarantee();
684 let _ = sink.current_schema().await.unwrap();
685 let _ = sink.supports_schema_evolution();
686 let _ = sink
687 .evolve_schema(&crate::drift::SchemaEvolution::default())
688 .await;
689 let _ = sink.config_schema();
690 let _ = sink.connector_name();
691 let _ = sink.dataset_uri();
692 assert!(!sink.is_overwrite());
693 let _ = sink.begin_overwrite().await;
694 let _ = sink.commit_overwrite().await;
695 sink.abort_overwrite().await.unwrap();
696 sink.complete_run().await.unwrap();
697
698 assert!(format!("{sink:?}").contains("MetadataSink"));
699 }
700
701 #[test]
702 fn custom_prefix_and_subset() {
703 let meta = CompiledMetadata::compile(&MetadataColumnsSpec {
704 enabled: true,
705 prefix: "_dt".into(),
706 columns: vec![MetadataColumn::RunId],
707 })
708 .unwrap()
709 .unwrap();
710 let sink = MetadataSink::new(Box::new(CapturingSink::default()), meta, ctx());
711 let out = sink.stamp(&[json!({"a": 1})]);
712 assert_eq!(out[0]["_dt_run_id"], "run-42");
713 assert!(out[0].get("_dt_loaded_at").is_none());
714 }
715}