Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,236 @@
|
||||
//! Lineage tracker implementation.
|
||||
|
||||
use chrono::Utc;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
use tracing::{debug, info, instrument};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::dag::Task;
|
||||
use crate::engine::ExecutionStatus;
|
||||
use crate::{DataValue, Result, state::StateManager};
|
||||
|
||||
use super::cache::{LineageCache, LineagePath};
|
||||
use super::config::LineageConfig;
|
||||
use super::config::LineageExportFormat;
|
||||
use super::graph::LineageGraph;
|
||||
use super::impact::{ImpactAnalysis, ImpactAnalyzer};
|
||||
use super::node::{LineageEdge, LineageNode};
|
||||
use super::schema::SchemaEvolutionTracker;
|
||||
use super::types::{DataSinkType, DataSourceType, LineageEdgeType, LineageNodeType};
|
||||
|
||||
/// Data lineage tracker
|
||||
pub struct LineageTracker {
|
||||
/// Configuration
|
||||
config: LineageConfig,
|
||||
/// State manager for persistence
|
||||
state_manager: Arc<StateManager>,
|
||||
/// Lineage graph
|
||||
lineage_graph: Arc<RwLock<LineageGraph>>,
|
||||
/// Schema evolution tracker
|
||||
schema_tracker: Arc<RwLock<SchemaEvolutionTracker>>,
|
||||
/// Impact analyzer
|
||||
impact_analyzer: Arc<ImpactAnalyzer>,
|
||||
/// Lineage cache for performance
|
||||
lineage_cache: Arc<RwLock<LineageCache>>,
|
||||
}
|
||||
|
||||
impl LineageTracker {
|
||||
/// Create a new lineage tracker
|
||||
pub fn new(config: LineageConfig, state_manager: Arc<StateManager>) -> Result<Self> {
|
||||
let lineage_graph = Arc::new(RwLock::new(LineageGraph::new()));
|
||||
let schema_tracker = Arc::new(RwLock::new(SchemaEvolutionTracker::new()));
|
||||
let impact_analyzer = Arc::new(ImpactAnalyzer::new(config.impact_analysis.clone()));
|
||||
let lineage_cache = Arc::new(RwLock::new(LineageCache::new()));
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
state_manager,
|
||||
lineage_graph,
|
||||
schema_tracker,
|
||||
impact_analyzer,
|
||||
lineage_cache,
|
||||
})
|
||||
}
|
||||
|
||||
/// Track the start of a task
|
||||
#[instrument(skip(self, task))]
|
||||
pub async fn track_task_start(&self, task_id: &str, task: &Task) -> Result<()> {
|
||||
if !self.config.auto_capture {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
debug!("Tracking task start: {}", task_id);
|
||||
|
||||
let mut graph = self.lineage_graph.write().await;
|
||||
|
||||
// Create task node
|
||||
let task_node = LineageNode {
|
||||
node_id: task_id.to_string(),
|
||||
node_type: LineageNodeType::Task {
|
||||
task_id: task_id.to_string(),
|
||||
task_type: format!("{:?}", task.task_type),
|
||||
},
|
||||
name: task.name.clone(),
|
||||
description: task.description.clone(),
|
||||
schema: None,
|
||||
created_at: Utc::now(),
|
||||
modified_at: Utc::now(),
|
||||
properties: HashMap::new(),
|
||||
tags: vec!["task".to_string()],
|
||||
};
|
||||
|
||||
graph.add_node(task_node)?;
|
||||
|
||||
// Track source relationships
|
||||
if let Some(ref source) = task.source {
|
||||
let source_node_id = format!("source_{}", source.id());
|
||||
if !graph.has_node(&source_node_id) {
|
||||
let source_node = LineageNode {
|
||||
node_id: source_node_id.clone(),
|
||||
node_type: LineageNodeType::DataSource {
|
||||
source_type: DataSourceType::Other("generic".to_string()),
|
||||
connection_info: HashMap::new(),
|
||||
},
|
||||
name: source.id().to_string(),
|
||||
description: None,
|
||||
schema: None,
|
||||
created_at: Utc::now(),
|
||||
modified_at: Utc::now(),
|
||||
properties: HashMap::new(),
|
||||
tags: vec!["source".to_string()],
|
||||
};
|
||||
graph.add_node(source_node)?;
|
||||
}
|
||||
|
||||
// Add edge from source to task
|
||||
graph.add_edge(
|
||||
&source_node_id,
|
||||
task_id,
|
||||
LineageEdge {
|
||||
edge_id: Uuid::new_v4().to_string(),
|
||||
edge_type: LineageEdgeType::DataFlow,
|
||||
strength: 1.0,
|
||||
created_at: Utc::now(),
|
||||
properties: HashMap::new(),
|
||||
},
|
||||
)?;
|
||||
}
|
||||
|
||||
// Track sink relationships
|
||||
if let Some(ref sink) = task.sink {
|
||||
let sink_node_id = format!("sink_{}", sink.id());
|
||||
if !graph.has_node(&sink_node_id) {
|
||||
let sink_node = LineageNode {
|
||||
node_id: sink_node_id.clone(),
|
||||
node_type: LineageNodeType::DataSink {
|
||||
sink_type: DataSinkType::Other("generic".to_string()),
|
||||
connection_info: HashMap::new(),
|
||||
},
|
||||
name: sink.id().to_string(),
|
||||
description: None,
|
||||
schema: None,
|
||||
created_at: Utc::now(),
|
||||
modified_at: Utc::now(),
|
||||
properties: HashMap::new(),
|
||||
tags: vec!["sink".to_string()],
|
||||
};
|
||||
graph.add_node(sink_node)?;
|
||||
}
|
||||
|
||||
// Add edge from task to sink
|
||||
graph.add_edge(
|
||||
task_id,
|
||||
&sink_node_id,
|
||||
LineageEdge {
|
||||
edge_id: Uuid::new_v4().to_string(),
|
||||
edge_type: LineageEdgeType::DataFlow,
|
||||
strength: 1.0,
|
||||
created_at: Utc::now(),
|
||||
properties: HashMap::new(),
|
||||
},
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Track the completion of a task
|
||||
#[instrument(skip(self))]
|
||||
pub async fn track_task_completion(
|
||||
&self,
|
||||
task_id: &str,
|
||||
status: &ExecutionStatus,
|
||||
records_processed: u64,
|
||||
) -> Result<()> {
|
||||
if !self.config.auto_capture {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
debug!(
|
||||
"Tracking task completion: {} (status: {:?})",
|
||||
task_id, status
|
||||
);
|
||||
|
||||
let mut graph = self.lineage_graph.write().await;
|
||||
|
||||
// Update task node with completion information
|
||||
if let Some(node) = graph.get_node_mut(task_id) {
|
||||
node.modified_at = Utc::now();
|
||||
node.properties.insert(
|
||||
"execution_status".to_string(),
|
||||
DataValue::String(format!("{status:?}")),
|
||||
);
|
||||
node.properties.insert(
|
||||
"records_processed".to_string(),
|
||||
DataValue::Int(records_processed as i64),
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Analyze downstream impact of changes to a node
|
||||
#[instrument(skip(self))]
|
||||
pub async fn analyze_downstream_impact(&self, node_id: &str) -> Result<ImpactAnalysis> {
|
||||
info!("Analyzing downstream impact for node: {}", node_id);
|
||||
|
||||
self.impact_analyzer
|
||||
.analyze_downstream_impact(node_id, &*self.lineage_graph.read().await)
|
||||
}
|
||||
|
||||
/// Analyze upstream dependencies of a node
|
||||
#[instrument(skip(self))]
|
||||
pub async fn analyze_upstream_dependencies(&self, node_id: &str) -> Result<ImpactAnalysis> {
|
||||
info!("Analyzing upstream dependencies for node: {}", node_id);
|
||||
|
||||
self.impact_analyzer
|
||||
.analyze_upstream_dependencies(node_id, &*self.lineage_graph.read().await)
|
||||
}
|
||||
|
||||
/// Get lineage path between two nodes
|
||||
pub async fn get_lineage_path(&self, source: &str, target: &str) -> Result<Vec<LineagePath>> {
|
||||
let graph = self.lineage_graph.read().await;
|
||||
let mut cache = self.lineage_cache.write().await;
|
||||
|
||||
// Check cache first
|
||||
if let Some(paths) = cache.get_path(source, target) {
|
||||
return Ok(paths.clone());
|
||||
}
|
||||
|
||||
// Calculate paths
|
||||
let paths = graph.find_paths(source, target)?;
|
||||
|
||||
// Cache results
|
||||
cache.cache_path((source.to_string(), target.to_string()), paths.clone());
|
||||
|
||||
Ok(paths)
|
||||
}
|
||||
|
||||
/// Export lineage graph for visualization
|
||||
pub async fn export_graph(&self, format: LineageExportFormat) -> Result<String> {
|
||||
let graph = self.lineage_graph.read().await;
|
||||
graph.export(format)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user