use crate::error::FaucetError;
use crate::traits::{RowOutcome, Sink};
use chrono::Utc;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use std::sync::atomic::{AtomicU64, Ordering};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum MetadataColumn {
ExtractedAt,
LoadedAt,
RunId,
Source,
Sequence,
}
impl MetadataColumn {
pub fn suffix(self) -> &'static str {
match self {
MetadataColumn::ExtractedAt => "extracted_at",
MetadataColumn::LoadedAt => "loaded_at",
MetadataColumn::RunId => "run_id",
MetadataColumn::Source => "source",
MetadataColumn::Sequence => "sequence",
}
}
}
fn default_prefix() -> String {
"_faucet".to_owned()
}
fn default_true() -> bool {
true
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct MetadataColumnsSpec {
#[serde(default = "default_true")]
pub enabled: bool,
#[serde(default = "default_prefix")]
pub prefix: String,
#[serde(default)]
pub columns: Vec<MetadataColumn>,
}
impl Default for MetadataColumnsSpec {
fn default() -> Self {
Self {
enabled: true,
prefix: default_prefix(),
columns: Vec::new(),
}
}
}
const DEFAULT_COLUMNS: &[MetadataColumn] = &[
MetadataColumn::ExtractedAt,
MetadataColumn::LoadedAt,
MetadataColumn::RunId,
MetadataColumn::Source,
];
#[derive(Debug, Clone)]
pub struct CompiledMetadata {
prefix: String,
columns: Vec<MetadataColumn>,
}
impl CompiledMetadata {
pub fn compile(spec: &MetadataColumnsSpec) -> Result<Option<Self>, FaucetError> {
if !spec.enabled {
return Ok(None);
}
if spec.prefix.trim().is_empty() {
return Err(FaucetError::Config(
"metadata_columns: `prefix` must not be empty".into(),
));
}
let columns = if spec.columns.is_empty() {
DEFAULT_COLUMNS.to_vec()
} else {
spec.columns.clone()
};
Ok(Some(Self {
prefix: spec.prefix.clone(),
columns,
}))
}
fn column_name(&self, col: MetadataColumn) -> String {
format!("{}_{}", self.prefix, col.suffix())
}
}
#[derive(Debug, Clone)]
pub struct MetadataContext {
pub run_id: String,
pub source: String,
}
pub struct MetadataSink {
inner: Box<dyn Sink>,
meta: CompiledMetadata,
ctx: MetadataContext,
seq: AtomicU64,
collided: std::sync::Mutex<std::collections::BTreeSet<String>>,
columns_ready: tokio::sync::OnceCell<()>,
}
impl std::fmt::Debug for MetadataSink {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MetadataSink")
.field("prefix", &self.meta.prefix)
.field("columns", &self.meta.columns)
.finish()
}
}
impl MetadataSink {
pub fn new(inner: Box<dyn Sink>, meta: CompiledMetadata, ctx: MetadataContext) -> Self {
Self {
inner,
meta,
ctx,
seq: AtomicU64::new(0),
collided: std::sync::Mutex::new(std::collections::BTreeSet::new()),
columns_ready: tokio::sync::OnceCell::new(),
}
}
fn stamp(&self, records: &[Value]) -> Vec<Value> {
self.stamp_with(records, true)
}
fn preview(&self, records: &[Value]) -> Vec<Value> {
self.stamp_with(records, false)
}
fn stamp_with(&self, records: &[Value], advance: bool) -> Vec<Value> {
let now = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
let mut next = self.seq.load(Ordering::Relaxed);
records
.iter()
.map(|rec| match rec {
Value::Object(map) => {
let mut map = map.clone();
let seq = if advance {
self.seq.fetch_add(1, Ordering::Relaxed)
} else {
next += 1;
next - 1
};
self.inject(&mut map, &now, seq, advance);
Value::Object(map)
}
other => other.clone(),
})
.collect()
}
fn inject(&self, map: &mut Map<String, Value>, now: &str, seq: u64, note: bool) {
for &col in &self.meta.columns {
let name = self.meta.column_name(col);
let value = match col {
MetadataColumn::ExtractedAt | MetadataColumn::LoadedAt => {
Value::String(now.to_owned())
}
MetadataColumn::RunId => Value::String(self.ctx.run_id.clone()),
MetadataColumn::Source => Value::String(self.ctx.source.clone()),
MetadataColumn::Sequence => Value::from(seq),
};
if map.insert(name.clone(), value).is_some() && note {
self.note_collision(name);
}
}
}
async fn ensure_columns(&self) -> Result<(), FaucetError> {
self.columns_ready
.get_or_try_init(|| async {
let Some(schema) = self.inner.current_schema().await? else {
return Ok(());
};
let props = schema.get("properties").and_then(Value::as_object);
let additions: Vec<crate::drift::ColumnChange> = self
.meta
.columns
.iter()
.map(|&c| (self.meta.column_name(c), c))
.filter(|(name, _)| props.is_none_or(|p| !p.contains_key(name)))
.map(|(name, c)| crate::drift::ColumnChange {
name,
from: None,
to: serde_json::json!({ "type": [
if c == MetadataColumn::Sequence { "integer" } else { "string" },
"null"
] }),
})
.collect();
if additions.is_empty() {
return Ok(());
}
let names: Vec<String> = additions.iter().map(|a| a.name.clone()).collect();
if !self.inner.supports_schema_evolution() {
return Err(FaucetError::Config(format!(
"metadata_columns: the destination has no column {} and sink '{}' \
cannot add columns; add them to the table or disable metadata_columns",
names.join(", "),
self.inner.connector_name()
)));
}
tracing::info!(columns = ?names, "adding missing metadata columns to the destination");
self.inner
.evolve_schema(&crate::drift::SchemaEvolution {
additions,
..Default::default()
})
.await
})
.await
.map(|_| ())
}
fn note_collision(&self, name: String) {
let Ok(mut seen) = self.collided.lock() else {
return;
};
if seen.insert(name.clone()) {
tracing::warn!(
column = %name,
"a source field has the same name as a metadata column; the metadata \
value replaces it (choose another `metadata_columns.prefix` to keep it)"
);
}
}
#[cfg(test)]
fn collisions(&self) -> Vec<String> {
self.collided.lock().unwrap().iter().cloned().collect()
}
}
#[async_trait::async_trait]
impl Sink for MetadataSink {
crate::forward_sink!(inner; except
admit_page, write_batch, write_batch_reported, write_governed, write_batch_partial,
write_batch_idempotent, write_batch_idempotent_reported,
begin_overwrite,
supports_columnar, write_batch_columnar, native_load_capabilities, load_native);
async fn admit_page(&self, records: &[Value]) -> Result<(), FaucetError> {
self.inner.admit_page(&self.preview(records)).await
}
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
self.ensure_columns().await?;
self.inner.write_batch(&self.stamp(records)).await
}
async fn write_batch_partial(&self, records: &[Value]) -> Result<Vec<RowOutcome>, FaucetError> {
self.ensure_columns().await?;
self.inner.write_batch_partial(&self.stamp(records)).await
}
async fn write_batch_reported(
&self,
records: &[Value],
) -> Result<crate::accounting::WriteReport, FaucetError> {
self.ensure_columns().await?;
self.inner.write_batch_reported(&self.stamp(records)).await
}
async fn write_governed(
&self,
page: &crate::governance::GovernedPage<'_>,
) -> Result<crate::accounting::WriteReport, FaucetError> {
self.ensure_columns().await?;
self.inner
.write_governed(&page.restamped(self.stamp(page.records())))
.await
}
#[allow(deprecated)] async fn write_batch_idempotent(
&self,
records: &[Value],
scope: &str,
token: &str,
) -> Result<usize, FaucetError> {
self.ensure_columns().await?;
self.inner
.write_batch_idempotent(&self.stamp(records), scope, token)
.await
}
#[allow(deprecated)] async fn write_batch_idempotent_reported(
&self,
records: &[Value],
scope: &str,
token: &str,
) -> Result<crate::accounting::WriteReport, FaucetError> {
self.ensure_columns().await?;
self.inner
.write_batch_idempotent_reported(&self.stamp(records), scope, token)
.await
}
#[allow(deprecated)] async fn begin_overwrite(&self) -> Result<(), FaucetError> {
self.ensure_columns().await?;
self.inner.begin_overwrite().await
}
}
#[allow(deprecated)] #[cfg(test)]
mod tests {
use super::*;
struct Mapped {
evolutions: std::sync::Mutex<Vec<Vec<String>>>,
evolves: bool,
}
#[async_trait::async_trait]
impl Sink for Mapped {
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
Ok(records.len())
}
async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
Ok(Some(json!({"type": "object", "properties": {
"id": {"type": "integer"}, "_faucet_source": {"type": "string"}
}})))
}
fn supports_schema_evolution(&self) -> bool {
self.evolves
}
async fn evolve_schema(
&self,
evolution: &crate::drift::SchemaEvolution,
) -> Result<(), FaucetError> {
self.evolutions.lock().unwrap().push(
evolution
.additions
.iter()
.map(|a| format!("{}:{}", a.name, a.to["type"][0]))
.collect(),
);
Ok(())
}
}
#[tokio::test]
async fn missing_metadata_columns_are_added_once_before_the_first_write() {
let meta = CompiledMetadata::compile(&spec(&[
MetadataColumn::RunId,
MetadataColumn::Source,
MetadataColumn::Sequence,
]))
.unwrap()
.unwrap();
let inner = std::sync::Arc::new(Mapped {
evolutions: std::sync::Mutex::new(Vec::new()),
evolves: true,
});
struct Shared(std::sync::Arc<Mapped>);
#[async_trait::async_trait]
impl Sink for Shared {
async fn write_batch(&self, r: &[Value]) -> Result<usize, FaucetError> {
self.0.write_batch(r).await
}
async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
self.0.current_schema().await
}
fn supports_schema_evolution(&self) -> bool {
self.0.supports_schema_evolution()
}
async fn evolve_schema(
&self,
e: &crate::drift::SchemaEvolution,
) -> Result<(), FaucetError> {
self.0.evolve_schema(e).await
}
}
let sink = MetadataSink::new(Box::new(Shared(inner.clone())), meta, ctx());
sink.write_batch(&[json!({"id": 1})]).await.unwrap();
sink.write_batch(&[json!({"id": 2})]).await.unwrap();
assert_eq!(
*inner.evolutions.lock().unwrap(),
vec![vec![
"_faucet_run_id:\"string\"".to_string(),
"_faucet_sequence:\"integer\"".to_string()
]]
);
let meta = CompiledMetadata::compile(&spec(&[MetadataColumn::RunId]))
.unwrap()
.unwrap();
let rigid = MetadataSink::new(
Box::new(Mapped {
evolutions: std::sync::Mutex::new(Vec::new()),
evolves: false,
}),
meta,
ctx(),
);
let err = rigid.write_batch(&[json!({"id": 1})]).await.unwrap_err();
assert!(err.to_string().contains("_faucet_run_id"), "{err}");
}
#[test]
fn a_colliding_source_field_is_reported() {
let meta = CompiledMetadata::compile(&spec(&[MetadataColumn::RunId]))
.unwrap()
.unwrap();
let sink = MetadataSink::new(Box::new(CapturingSink::default()), meta, ctx());
let out = sink.stamp(&[
json!({"_faucet_run_id": "upstream", "a": 1}),
json!({"a": 2}),
]);
assert_eq!(out[0]["_faucet_run_id"], json!("run-42"));
assert_eq!(sink.collisions(), vec!["_faucet_run_id".to_string()]);
sink.stamp(&[json!({"_faucet_run_id": "again"})]);
assert_eq!(sink.collisions().len(), 1);
}
#[tokio::test]
async fn metadata_sink_forwards_every_hook_but_the_unstamped_paths() {
let probe = crate::forwarding::HookSink::default();
let meta = CompiledMetadata::compile(&spec(&[MetadataColumn::RunId]))
.unwrap()
.unwrap();
let sink = MetadataSink::new(Box::new(probe.clone()), meta, ctx());
crate::forwarding::assert_sink_forwards_every_hook(
&sink,
&probe,
&[
"supports_columnar",
"write_batch_columnar",
"native_load_capabilities",
"load_native",
],
)
.await;
}
use serde_json::json;
use std::sync::Mutex;
#[derive(Debug, Default)]
struct CapturingSink {
rows: Mutex<Vec<Value>>,
}
#[async_trait::async_trait]
impl Sink for CapturingSink {
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
self.rows.lock().unwrap().extend_from_slice(records);
Ok(records.len())
}
fn config_schema(&self) -> Value {
json!({})
}
}
fn spec(cols: &[MetadataColumn]) -> MetadataColumnsSpec {
MetadataColumnsSpec {
enabled: true,
prefix: "_faucet".into(),
columns: cols.to_vec(),
}
}
fn ctx() -> MetadataContext {
MetadataContext {
run_id: "run-42".into(),
source: "rest".into(),
}
}
#[test]
fn compile_resolves_defaults_and_respects_disable() {
let c = CompiledMetadata::compile(&spec(&[])).unwrap().unwrap();
assert_eq!(c.columns, DEFAULT_COLUMNS.to_vec());
let disabled = MetadataColumnsSpec {
enabled: false,
..Default::default()
};
assert!(CompiledMetadata::compile(&disabled).unwrap().is_none());
let bad = MetadataColumnsSpec {
prefix: " ".into(),
..Default::default()
};
assert!(CompiledMetadata::compile(&bad).is_err());
}
#[tokio::test]
async fn stamps_all_column_kinds_with_prefix() {
let meta = CompiledMetadata::compile(&spec(&[
MetadataColumn::ExtractedAt,
MetadataColumn::LoadedAt,
MetadataColumn::RunId,
MetadataColumn::Source,
MetadataColumn::Sequence,
]))
.unwrap()
.unwrap();
let inner = Box::new(CapturingSink::default());
let sink = MetadataSink::new(inner, meta, ctx());
assert_eq!(
sink.batch_atomicity(),
crate::dlq::BatchAtomicity::BestEffort
);
let n = sink
.write_batch(&[json!({"id": 1}), json!({"id": 2})])
.await
.unwrap();
assert_eq!(n, 2);
let stamped = sink.stamp(&[json!({"id": 1}), json!({"id": 2})]);
let r0 = &stamped[0];
assert_eq!(r0["id"], 1);
assert_eq!(r0["_faucet_run_id"], "run-42");
assert_eq!(r0["_faucet_source"], "rest");
assert!(r0["_faucet_extracted_at"].is_string());
assert!(r0["_faucet_loaded_at"].is_string());
assert_eq!(
stamped[0]["_faucet_sequence"].as_u64().unwrap() + 1,
stamped[1]["_faucet_sequence"].as_u64().unwrap()
);
}
#[tokio::test]
async fn admit_page_forwards_the_stamped_page() {
#[derive(Default)]
struct Admits(std::sync::Mutex<Vec<Value>>);
#[async_trait::async_trait]
impl Sink for Admits {
async fn admit_page(&self, records: &[Value]) -> Result<(), FaucetError> {
self.0.lock().unwrap().extend_from_slice(records);
Ok(())
}
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
Ok(records.len())
}
}
let meta =
CompiledMetadata::compile(&spec(&[MetadataColumn::RunId, MetadataColumn::Sequence]))
.unwrap()
.unwrap();
let inner = std::sync::Arc::new(Admits::default());
struct Shared(std::sync::Arc<Admits>);
#[async_trait::async_trait]
impl Sink for Shared {
async fn admit_page(&self, records: &[Value]) -> Result<(), FaucetError> {
self.0.admit_page(records).await
}
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
self.0.write_batch(records).await
}
}
let sink = MetadataSink::new(Box::new(Shared(std::sync::Arc::clone(&inner))), meta, ctx());
sink.admit_page(&[json!({"id": 1}), json!("x")])
.await
.unwrap();
let admitted = inner.0.lock().unwrap().clone();
assert_eq!(admitted[0]["_faucet_run_id"], "run-42");
assert_eq!(admitted[0]["_faucet_sequence"], 0);
assert_eq!(admitted[1], json!("x"));
let stamped = sink.stamp(&[json!({"id": 1})]);
assert_eq!(
stamped[0]["_faucet_sequence"], 0,
"admission consumed no sequence"
);
}
#[test]
fn non_object_records_pass_through() {
let meta = CompiledMetadata::compile(&spec(&[MetadataColumn::RunId]))
.unwrap()
.unwrap();
let sink = MetadataSink::new(Box::new(CapturingSink::default()), meta, ctx());
let out = sink.stamp(&[json!("scalar"), json!(42)]);
assert_eq!(out, vec![json!("scalar"), json!(42)]);
}
#[tokio::test]
async fn delegates_every_sink_method_to_inner() {
let meta = CompiledMetadata::compile(&spec(&[MetadataColumn::RunId]))
.unwrap()
.unwrap();
let sink = MetadataSink::new(Box::new(CapturingSink::default()), meta, ctx());
assert_eq!(sink.write_batch(&[json!({"a": 1})]).await.unwrap(), 1);
assert_eq!(
sink.write_batch_partial(&[json!({"a": 1})])
.await
.unwrap()
.len(),
1
);
let _ = sink
.write_batch_idempotent(&[json!({"a": 1})], "scope", "tok")
.await;
sink.flush().await.unwrap();
let _ = sink.check(&crate::check::CheckContext::default()).await;
assert!(!sink.supports_cleanup());
let _ = sink
.cleanup_scope(
&std::collections::BTreeMap::new(),
&crate::cleanup::SeenKeys::new(),
)
.await;
assert!(!sink.supports_idempotent_writes());
assert!(sink.last_committed_token("s").await.unwrap().is_none());
let _ = sink.supported_write_modes();
let _ = sink.dedups_by_key();
let _ = sink.sink_guarantee();
let _ = sink.current_schema().await.unwrap();
let _ = sink.supports_schema_evolution();
let _ = sink
.evolve_schema(&crate::drift::SchemaEvolution::default())
.await;
let _ = sink.config_schema();
let _ = sink.connector_name();
let _ = sink.dataset_uri();
assert!(!sink.is_overwrite());
let _ = sink.begin_overwrite().await;
let _ = sink.commit_overwrite().await;
sink.abort_overwrite().await.unwrap();
sink.complete_run().await.unwrap();
assert!(format!("{sink:?}").contains("MetadataSink"));
}
#[test]
fn custom_prefix_and_subset() {
let meta = CompiledMetadata::compile(&MetadataColumnsSpec {
enabled: true,
prefix: "_dt".into(),
columns: vec![MetadataColumn::RunId],
})
.unwrap()
.unwrap();
let sink = MetadataSink::new(Box::new(CapturingSink::default()), meta, ctx());
let out = sink.stamp(&[json!({"a": 1})]);
assert_eq!(out[0]["_dt_run_id"], "run-42");
assert!(out[0].get("_dt_loaded_at").is_none());
}
}