1use serde::{Deserialize, Serialize};
56use std::collections::{HashMap, HashSet, VecDeque};
57
58#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
64pub struct NodeId(u64);
65
66impl NodeId {
67 pub fn value(self) -> u64 {
69 self.0
70 }
71}
72
73impl std::fmt::Display for NodeId {
74 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75 write!(f, "node_{}", self.0)
76 }
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
85pub enum DataSource {
86 File(String),
88 Database(String),
90 InMemory,
92 Generated,
94}
95
96impl DataSource {
97 pub fn label(&self) -> String {
99 match self {
100 DataSource::File(p) => format!("file:{}", p),
101 DataSource::Database(s) => format!("db:{}", s),
102 DataSource::InMemory => "in_memory".to_string(),
103 DataSource::Generated => "generated".to_string(),
104 }
105 }
106}
107
108#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
114pub enum ColumnType {
115 Integer,
117 Float,
119 Boolean,
121 Text,
123 Other(String),
125}
126
127impl ColumnType {
128 fn label(&self) -> &str {
129 match self {
130 ColumnType::Integer => "integer",
131 ColumnType::Float => "float",
132 ColumnType::Boolean => "boolean",
133 ColumnType::Text => "text",
134 ColumnType::Other(s) => s.as_str(),
135 }
136 }
137}
138
139#[derive(Debug, Clone, Serialize, Deserialize)]
145pub struct DataNode {
146 pub id: NodeId,
148 pub name: String,
150 pub source: DataSource,
152 pub schema: Vec<(String, ColumnType)>,
154 pub created_at: String,
156 pub tags: HashMap<String, String>,
158}
159
160impl DataNode {
161 pub fn new(name: String, source: DataSource, schema: Vec<(String, ColumnType)>) -> Self {
164 DataNode {
165 id: NodeId(0),
166 name,
167 source,
168 schema,
169 created_at: chrono_now(),
170 tags: HashMap::new(),
171 }
172 }
173
174 pub fn with_tag(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
176 self.tags.insert(key.into(), value.into());
177 self
178 }
179}
180
181#[derive(Debug, Clone, Serialize, Deserialize)]
188pub struct Transformation {
189 pub input_nodes: Vec<NodeId>,
191 pub output_node: NodeId,
193 pub op_name: String,
195 pub params: HashMap<String, String>,
197 pub created_at: String,
199}
200
201impl Transformation {
202 pub fn new(
204 input_nodes: Vec<NodeId>,
205 output_node: NodeId,
206 op_name: impl Into<String>,
207 params: HashMap<String, String>,
208 ) -> Self {
209 Transformation {
210 input_nodes,
211 output_node,
212 op_name: op_name.into(),
213 params,
214 created_at: chrono_now(),
215 }
216 }
217}
218
219#[derive(Debug, Default, Serialize, Deserialize)]
230pub struct DataLineage {
231 nodes: HashMap<NodeId, DataNode>,
232 transformations: Vec<Transformation>,
233 next_id: u64,
234}
235
236impl DataLineage {
237 pub fn new() -> Self {
239 DataLineage {
240 nodes: HashMap::new(),
241 transformations: Vec::new(),
242 next_id: 1,
243 }
244 }
245
246 pub fn add_node(&mut self, mut node: DataNode) -> NodeId {
248 let id = NodeId(self.next_id);
249 self.next_id += 1;
250 node.id = id;
251 self.nodes.insert(id, node);
252 id
253 }
254
255 pub fn get_node(&self, id: NodeId) -> Option<&DataNode> {
257 self.nodes.get(&id)
258 }
259
260 pub fn get_node_mut(&mut self, id: NodeId) -> Option<&mut DataNode> {
262 self.nodes.get_mut(&id)
263 }
264
265 pub fn nodes(&self) -> impl Iterator<Item = &DataNode> {
267 self.nodes.values()
268 }
269
270 pub fn transformations(&self) -> &[Transformation] {
272 &self.transformations
273 }
274
275 pub fn add_transformation(&mut self, t: Transformation) -> bool {
281 let all_known = t
283 .input_nodes
284 .iter()
285 .chain(std::iter::once(&t.output_node))
286 .all(|id| self.nodes.contains_key(id));
287 if all_known {
288 self.transformations.push(t);
289 true
290 } else {
291 false
292 }
293 }
294
295 fn producing_transformations(&self, node_id: NodeId) -> Vec<&Transformation> {
297 self.transformations
298 .iter()
299 .filter(|t| t.output_node == node_id)
300 .collect()
301 }
302}
303
304pub fn record_transformation(
312 lineage: &mut DataLineage,
313 inputs: Vec<NodeId>,
314 output: NodeId,
315 op_name: impl Into<String>,
316 params: HashMap<String, String>,
317) -> bool {
318 let t = Transformation::new(inputs, output, op_name, params);
319 lineage.add_transformation(t)
320}
321
322pub fn get_provenance(lineage: &DataLineage, node_id: NodeId) -> Vec<DataNode> {
328 let mut visited: HashSet<NodeId> = HashSet::new();
329 let mut queue: VecDeque<NodeId> = VecDeque::new();
330 let mut result: Vec<DataNode> = Vec::new();
331
332 for t in lineage.producing_transformations(node_id) {
334 for &inp in &t.input_nodes {
335 if visited.insert(inp) {
336 queue.push_back(inp);
337 }
338 }
339 }
340
341 while let Some(id) = queue.pop_front() {
342 if let Some(node) = lineage.get_node(id) {
343 result.push(node.clone());
344 }
345 for t in lineage.producing_transformations(id) {
347 for &inp in &t.input_nodes {
348 if visited.insert(inp) {
349 queue.push_back(inp);
350 }
351 }
352 }
353 }
354
355 result
356}
357
358pub fn export_lineage_dot(lineage: &DataLineage) -> String {
363 let mut dot = String::from("digraph lineage {\n rankdir=LR;\n node [shape=box];\n");
364
365 let mut node_ids: Vec<NodeId> = lineage.nodes.keys().copied().collect();
367 node_ids.sort_by_key(|n| n.0);
368 for nid in &node_ids {
369 if let Some(node) = lineage.get_node(*nid) {
370 let schema_str: String = node
372 .schema
373 .iter()
374 .map(|(c, t)| format!("{}:{}", c, t.label()))
375 .collect::<Vec<_>>()
376 .join(", ");
377 let label = format!(
378 "{}\\n[{}]\\n({} col(s))",
379 escape_dot(&node.name),
380 escape_dot(&node.source.label()),
381 node.schema.len(),
382 );
383 dot.push_str(&format!(
385 " {} [label=\"{}\" tooltip=\"{}\"];\n",
386 nid,
387 label,
388 escape_dot(&schema_str),
389 ));
390 }
391 }
392
393 for (tidx, t) in lineage.transformations.iter().enumerate() {
395 let edge_label = format!("{}\\n(step {})", escape_dot(&t.op_name), tidx + 1);
396 for &inp in &t.input_nodes {
397 dot.push_str(&format!(
398 " {} -> {} [label=\"{}\"];\n",
399 inp, t.output_node, edge_label
400 ));
401 }
402 }
403
404 dot.push_str("}\n");
405 dot
406}
407
408pub fn lineage_to_json(lineage: &DataLineage) -> String {
414 let mut out = String::from("{\n");
415
416 out.push_str(" \"nodes\": [\n");
418 let mut node_ids: Vec<NodeId> = lineage.nodes.keys().copied().collect();
419 node_ids.sort_by_key(|n| n.0);
420 for (i, nid) in node_ids.iter().enumerate() {
421 if let Some(node) = lineage.get_node(*nid) {
422 out.push_str(" {\n");
423 out.push_str(&format!(" \"id\": {},\n", nid.0));
424 out.push_str(&format!(
425 " \"name\": \"{}\",\n",
426 json_escape(&node.name)
427 ));
428 out.push_str(&format!(
429 " \"source\": \"{}\",\n",
430 json_escape(&node.source.label())
431 ));
432 out.push_str(&format!(
433 " \"created_at\": \"{}\",\n",
434 json_escape(&node.created_at)
435 ));
436 out.push_str(" \"schema\": [\n");
438 for (si, (col, typ)) in node.schema.iter().enumerate() {
439 out.push_str(&format!(
440 " {{\"column\": \"{}\", \"type\": \"{}\"}}{}",
441 json_escape(col),
442 json_escape(typ.label()),
443 if si + 1 < node.schema.len() {
444 ",\n"
445 } else {
446 "\n"
447 }
448 ));
449 }
450 out.push_str(" ],\n");
451 out.push_str(" \"tags\": {");
453 let tag_pairs: Vec<_> = node.tags.iter().collect();
454 for (ti, (k, v)) in tag_pairs.iter().enumerate() {
455 out.push_str(&format!(
456 "\"{}\": \"{}\"{}",
457 json_escape(k),
458 json_escape(v),
459 if ti + 1 < tag_pairs.len() { ", " } else { "" }
460 ));
461 }
462 out.push_str("}\n");
463 out.push_str(if i + 1 < node_ids.len() {
464 " },\n"
465 } else {
466 " }\n"
467 });
468 }
469 }
470 out.push_str(" ],\n");
471
472 out.push_str(" \"transformations\": [\n");
474 for (i, t) in lineage.transformations.iter().enumerate() {
475 out.push_str(" {\n");
476 let inputs_str: String = t
478 .input_nodes
479 .iter()
480 .map(|n| n.0.to_string())
481 .collect::<Vec<_>>()
482 .join(", ");
483 out.push_str(&format!(" \"input_nodes\": [{}],\n", inputs_str));
484 out.push_str(&format!(" \"output_node\": {},\n", t.output_node.0));
485 out.push_str(&format!(
486 " \"op_name\": \"{}\",\n",
487 json_escape(&t.op_name)
488 ));
489 out.push_str(&format!(
490 " \"created_at\": \"{}\",\n",
491 json_escape(&t.created_at)
492 ));
493 out.push_str(" \"params\": {");
495 let param_pairs: Vec<_> = t.params.iter().collect();
496 for (pi, (k, v)) in param_pairs.iter().enumerate() {
497 out.push_str(&format!(
498 "\"{}\": \"{}\"{}",
499 json_escape(k),
500 json_escape(v),
501 if pi + 1 < param_pairs.len() { ", " } else { "" }
502 ));
503 }
504 out.push_str("}\n");
505 out.push_str(if i + 1 < lineage.transformations.len() {
506 " },\n"
507 } else {
508 " }\n"
509 });
510 }
511 out.push_str(" ]\n");
512
513 out.push_str("}\n");
514 out
515}
516
517fn escape_dot(s: &str) -> String {
524 s.replace('\\', "\\\\").replace('"', "\\\"")
525}
526
527fn json_escape(s: &str) -> String {
529 let mut out = String::with_capacity(s.len());
530 for ch in s.chars() {
531 match ch {
532 '"' => out.push_str("\\\""),
533 '\\' => out.push_str("\\\\"),
534 '\n' => out.push_str("\\n"),
535 '\r' => out.push_str("\\r"),
536 '\t' => out.push_str("\\t"),
537 c => out.push(c),
538 }
539 }
540 out
541}
542
543fn chrono_now() -> String {
546 use std::time::{SystemTime, UNIX_EPOCH};
548 match SystemTime::now().duration_since(UNIX_EPOCH) {
549 Ok(d) => {
550 let secs = d.as_secs();
551 let (days_since_epoch, secs_of_day) = (secs / 86400, secs % 86400);
553 let (h, m, s) = (
554 secs_of_day / 3600,
555 (secs_of_day % 3600) / 60,
556 secs_of_day % 60,
557 );
558 let jdn = days_since_epoch + 2440588; let (y, mo, da) = jdn_to_ymd(jdn);
561 format!("{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z", y, mo, da, h, m, s)
562 }
563 Err(_) => "1970-01-01T00:00:00Z".to_string(),
564 }
565}
566
567fn jdn_to_ymd(jdn: u64) -> (u64, u64, u64) {
570 let jdn = jdn as i64;
571 let f = jdn + 1401 + (((4 * jdn + 274277) / 146097) * 3) / 4 - 38;
572 let e = 4 * f + 3;
573 let g = (e % 1461) / 4;
574 let h = 5 * g + 2;
575 let day = (h % 153) / 5 + 1;
576 let month = (h / 153 + 2) % 12 + 1;
577 let year = e / 1461 - 4716 + (14 - month) / 12;
578 (year as u64, month as u64, day as u64)
579}
580
581#[cfg(test)]
586mod tests {
587 use super::*;
588
589 fn build_simple_lineage() -> (DataLineage, NodeId, NodeId, NodeId) {
590 let mut lineage = DataLineage::new();
591
592 let raw = lineage.add_node(DataNode::new(
593 "raw".to_string(),
594 DataSource::File("/data/raw.csv".to_string()),
595 vec![
596 ("id".to_string(), ColumnType::Integer),
597 ("value".to_string(), ColumnType::Float),
598 ],
599 ));
600
601 let clean = lineage.add_node(DataNode::new(
602 "cleaned".to_string(),
603 DataSource::InMemory,
604 vec![
605 ("id".to_string(), ColumnType::Integer),
606 ("value".to_string(), ColumnType::Float),
607 ],
608 ));
609
610 let agg = lineage.add_node(DataNode::new(
611 "aggregated".to_string(),
612 DataSource::Generated,
613 vec![("mean_value".to_string(), ColumnType::Float)],
614 ));
615
616 let mut params1 = HashMap::new();
617 params1.insert("drop_nulls".to_string(), "true".to_string());
618 record_transformation(&mut lineage, vec![raw], clean, "filter", params1);
619
620 let mut params2 = HashMap::new();
621 params2.insert("fn".to_string(), "mean".to_string());
622 record_transformation(&mut lineage, vec![clean], agg, "aggregate", params2);
623
624 (lineage, raw, clean, agg)
625 }
626
627 #[test]
628 fn test_add_and_get_node() {
629 let mut lineage = DataLineage::new();
630 let id = lineage.add_node(DataNode::new(
631 "test".to_string(),
632 DataSource::InMemory,
633 vec![],
634 ));
635 let node = lineage.get_node(id).expect("node must exist");
636 assert_eq!(node.name, "test");
637 assert_eq!(node.id, id);
638 }
639
640 #[test]
641 fn test_record_transformation_valid() {
642 let (lineage, raw, clean, _agg) = build_simple_lineage();
643 assert_eq!(lineage.transformations().len(), 2);
644 assert_eq!(lineage.transformations()[0].input_nodes, vec![raw]);
645 assert_eq!(lineage.transformations()[0].output_node, clean);
646 }
647
648 #[test]
649 fn test_record_transformation_invalid_node() {
650 let mut lineage = DataLineage::new();
651 let fake_id = NodeId(999);
652 let out_id = lineage.add_node(DataNode::new(
653 "out".to_string(),
654 DataSource::InMemory,
655 vec![],
656 ));
657 let ok = record_transformation(&mut lineage, vec![fake_id], out_id, "op", HashMap::new());
658 assert!(!ok, "should reject unknown input node");
659 }
660
661 #[test]
662 fn test_get_provenance_depth_one() {
663 let (lineage, raw, clean, _agg) = build_simple_lineage();
664 let prov = get_provenance(&lineage, clean);
665 assert_eq!(prov.len(), 1);
666 assert_eq!(prov[0].id, raw);
667 }
668
669 #[test]
670 fn test_get_provenance_depth_two() {
671 let (lineage, raw, clean, agg) = build_simple_lineage();
672 let prov = get_provenance(&lineage, agg);
673 let ids: HashSet<NodeId> = prov.iter().map(|n| n.id).collect();
675 assert!(ids.contains(&raw));
676 assert!(ids.contains(&clean));
677 assert_eq!(prov.len(), 2);
678 }
679
680 #[test]
681 fn test_get_provenance_root_node() {
682 let (lineage, raw, _clean, _agg) = build_simple_lineage();
683 let prov = get_provenance(&lineage, raw);
684 assert!(prov.is_empty(), "root node has no ancestors");
685 }
686
687 #[test]
688 fn test_export_lineage_dot_structure() {
689 let (lineage, _raw, _clean, _agg) = build_simple_lineage();
690 let dot = export_lineage_dot(&lineage);
691
692 assert!(dot.starts_with("digraph lineage {"));
693 assert!(dot.ends_with("}\n"));
694 assert!(dot.contains("raw"));
695 assert!(dot.contains("cleaned"));
696 assert!(dot.contains("aggregated"));
697 assert!(dot.contains("->"));
698 assert!(dot.contains("filter"));
699 assert!(dot.contains("aggregate"));
700 }
701
702 #[test]
703 fn test_lineage_to_json_structure() {
704 let (lineage, _raw, _clean, _agg) = build_simple_lineage();
705 let json = lineage_to_json(&lineage);
706
707 assert!(json.contains("\"nodes\""));
708 assert!(json.contains("\"transformations\""));
709 assert!(json.contains("\"raw\""));
710 assert!(json.contains("\"cleaned\""));
711 assert!(json.contains("\"aggregated\""));
712 assert!(json.contains("\"filter\""));
713 assert!(json.contains("\"aggregate\""));
714 let _: serde_json::Value = serde_json::from_str(&json).expect("must be valid JSON");
716 }
717
718 #[test]
719 fn test_data_source_labels() {
720 assert_eq!(DataSource::File("/a/b.csv".into()).label(), "file:/a/b.csv");
721 assert_eq!(
722 DataSource::Database("pg://localhost/mydb".into()).label(),
723 "db:pg://localhost/mydb"
724 );
725 assert_eq!(DataSource::InMemory.label(), "in_memory");
726 assert_eq!(DataSource::Generated.label(), "generated");
727 }
728
729 #[test]
730 fn test_node_tags() {
731 let mut lineage = DataLineage::new();
732 let id = lineage.add_node(
733 DataNode::new("tagged".to_string(), DataSource::InMemory, vec![])
734 .with_tag("owner", "alice")
735 .with_tag("version", "1"),
736 );
737 let node = lineage.get_node(id).expect("exists");
738 assert_eq!(node.tags.get("owner").map(|s| s.as_str()), Some("alice"));
739 assert_eq!(node.tags.get("version").map(|s| s.as_str()), Some("1"));
740 }
741
742 #[test]
743 fn test_multi_input_transformation() {
744 let mut lineage = DataLineage::new();
745 let a = lineage.add_node(DataNode::new("A".into(), DataSource::InMemory, vec![]));
746 let b = lineage.add_node(DataNode::new("B".into(), DataSource::InMemory, vec![]));
747 let c = lineage.add_node(DataNode::new("C".into(), DataSource::InMemory, vec![]));
748
749 let ok = record_transformation(&mut lineage, vec![a, b], c, "join", HashMap::new());
750 assert!(ok);
751 let prov = get_provenance(&lineage, c);
752 let prov_ids: HashSet<NodeId> = prov.iter().map(|n| n.id).collect();
753 assert!(prov_ids.contains(&a));
754 assert!(prov_ids.contains(&b));
755 }
756
757 #[test]
758 fn test_json_escaping() {
759 let s = r#"he said "hello\nworld""#;
760 let escaped = super::json_escape(s);
761 let json = format!("{{\"v\": \"{}\"}}", escaped);
763 let _: serde_json::Value = serde_json::from_str(&json).expect("valid JSON");
764 }
765
766 #[test]
767 fn test_column_type_labels() {
768 assert_eq!(ColumnType::Integer.label(), "integer");
769 assert_eq!(ColumnType::Float.label(), "float");
770 assert_eq!(ColumnType::Boolean.label(), "boolean");
771 assert_eq!(ColumnType::Text.label(), "text");
772 assert_eq!(ColumnType::Other("blob".into()).label(), "blob");
773 }
774}