floe_core/lineage/

mod.rs

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use serde_json::{json, Value};

use crate::config::{EntityConfig, LineageConfig};
use crate::run::events::{RunEvent, RunObserver};

struct EntityUris {
    source: String,
    accepted: String,
    rejected: Option<String>,
}

pub struct OpenLineageObserver {
    client: reqwest::blocking::Client,
    config: LineageConfig,
    entity_start_ms: Mutex<HashMap<String, u128>>,
    entity_run_ids: Mutex<HashMap<String, String>>,
    run_start_ms: Mutex<Option<u128>>,
    entity_schemas: HashMap<String, Vec<(String, String)>>,
    entity_uris: HashMap<String, EntityUris>,
    consecutive_failures: AtomicUsize,
    circuit_open: AtomicBool,
}

impl OpenLineageObserver {
    pub fn new(config: &LineageConfig, entities: &[EntityConfig]) -> crate::FloeResult<Self> {
        let timeout = Duration::from_secs(config.timeout_secs.unwrap_or(5));
        let client = reqwest::blocking::Client::builder()
            .timeout(timeout)
            .build()
            .map_err(|e| {
                Box::new(crate::errors::ConfigError(format!(
                    "lineage: failed to build HTTP client: {e}"
                ))) as Box<dyn std::error::Error + Send + Sync>
            })?;

        let entity_schemas = entities
            .iter()
            .map(|e| {
                let fields: Vec<(String, String)> = e
                    .schema
                    .columns
                    .iter()
                    .map(|c| (c.name.clone(), c.column_type.clone()))
                    .collect();
                (e.name.clone(), fields)
            })
            .collect();

        let entity_uris = entities
            .iter()
            .map(|e| {
                (
                    e.name.clone(),
                    EntityUris {
                        source: e.source.path.clone(),
                        accepted: e.sink.accepted.path.clone(),
                        rejected: e.sink.rejected.as_ref().map(|r| r.path.clone()),
                    },
                )
            })
            .collect();

        Ok(Self {
            client,
            config: config.clone(),
            entity_start_ms: Mutex::new(HashMap::new()),
            entity_run_ids: Mutex::new(HashMap::new()),
            run_start_ms: Mutex::new(None),
            entity_schemas,
            entity_uris,
            consecutive_failures: AtomicUsize::new(0),
            circuit_open: AtomicBool::new(false),
        })
    }

    fn attempt_post(&self, url: &str, body: &Value) -> Result<(), bool> {
        let mut req = self.client.post(url).json(body);
        if let Some(api_key) = self.config.api_key.as_deref() {
            req = req.bearer_auth(api_key);
        }
        match req.send() {
            Err(_) => Err(true),
            Ok(resp) => {
                let status = resp.status();
                if status.is_success() {
                    Ok(())
                } else if status.as_u16() == 429 || status.is_server_error() {
                    Err(true)
                } else {
                    Err(false)
                }
            }
        }
    }

    fn post_event(&self, body: Value) {
        if self.circuit_open.load(Ordering::Relaxed) {
            return;
        }

        let url = format!("{}/api/v1/lineage", self.config.url.trim_end_matches('/'));
        let max_failures = self.config.max_failures.unwrap_or(3) as usize;
        let retry_delays_ms: &[u64] = &[0, 100, 500];

        let mut succeeded = false;

        'retry: for (attempt, &delay_ms) in retry_delays_ms.iter().enumerate() {
            if delay_ms > 0 {
                std::thread::sleep(Duration::from_millis(delay_ms));
            }
            match self.attempt_post(&url, &body) {
                Ok(()) => {
                    self.consecutive_failures.store(0, Ordering::Relaxed);
                    succeeded = true;
                    break 'retry;
                }
                Err(is_retryable) => {
                    if !is_retryable || attempt == retry_delays_ms.len() - 1 {
                        break 'retry;
                    }
                }
            }
        }

        if succeeded {
            return;
        }

        // Always warn when an event is dropped, whether the failure was retryable or not.
        crate::warnings::emit(
            "",
            None,
            None,
            Some("lineage_http_error"),
            &format!("lineage: POST {url} failed — event dropped"),
        );

        let failures = self.consecutive_failures.fetch_add(1, Ordering::Relaxed) + 1;
        if failures >= max_failures && !self.circuit_open.swap(true, Ordering::Relaxed) {
            crate::warnings::emit(
                "",
                None,
                None,
                Some("lineage_circuit_open"),
                &format!(
                    "lineage: disabled for this run after {failures} consecutive failures — check endpoint {url}"
                ),
            );
        }
    }

    fn producer(&self) -> &str {
        self.config
            .producer
            .as_deref()
            .unwrap_or("https://github.com/malon64/floe")
    }

    fn parent_run_facet(&self) -> Option<Value> {
        if let Ok(run_id) = std::env::var("AIRFLOW_CTX_DAG_RUN_ID") {
            let dag = std::env::var("AIRFLOW_CTX_DAG_ID").unwrap_or_default();
            let task = std::env::var("AIRFLOW_CTX_TASK_ID").unwrap_or_default();
            let job_name = if task.is_empty() {
                dag.clone()
            } else {
                format!("{dag}.{task}")
            };
            return Some(json!({
                "run": { "runId": run_id },
                "job": {
                    "namespace": self.config.namespace,
                    "name": job_name
                },
                "_producer": self.producer(),
                "_schemaURL": "https://openlineage.io/spec/facets/1-0-0/ParentRunFacet.json"
            }));
        }
        if let Ok(run_id) = std::env::var("DAGSTER_RUN_ID") {
            let job = std::env::var("DAGSTER_JOB_NAME").unwrap_or_default();
            return Some(json!({
                "run": { "runId": run_id },
                "job": {
                    "namespace": self.config.namespace,
                    "name": job
                },
                "_producer": self.producer(),
                "_schemaURL": "https://openlineage.io/spec/facets/1-0-0/ParentRunFacet.json"
            }));
        }
        None
    }

    fn emit_entity_run_event(
        &self,
        entity_run_id: &str,
        name: &str,
        event_type: &str,
        ts_ms: u128,
        stats: Option<EntityStats>,
        uris: Option<&EntityUris>,
    ) {
        let event_time = ms_to_iso8601(ts_ms);
        let job_name = format!("{}.{name}", self.config.namespace);

        let mut run_facets = json!({});
        if let Some(parent) = self.parent_run_facet() {
            run_facets["parent"] = parent;
        }

        // Build inputs: source dataset with schema and quality facets on COMPLETE/FAIL.
        let inputs = match (stats.as_ref(), uris) {
            (Some(s), Some(u)) => {
                let rejection_rate = if s.rows > 0 {
                    s.rejected as f64 / s.rows as f64
                } else {
                    0.0
                };
                let schema_facet = json!({
                    "fields": s.schema_fields.iter().map(|(col_name, col_type)| {
                        json!({ "name": col_name, "type": col_type })
                    }).collect::<Vec<_>>(),
                    "_producer": self.producer(),
                    "_schemaURL": "https://openlineage.io/spec/facets/1-1-1/SchemaDatasetFacet.json"
                });
                let dq_facet = json!({
                    "rowCount": s.rows,
                    "validCount": s.accepted,
                    "invalidCount": s.rejected,
                    "_producer": self.producer(),
                    "_schemaURL": "https://openlineage.io/spec/facets/1-0-2/DataQualityMetricsInputDatasetFacet.json"
                });
                let floe_facet = json!({
                    "entity": name,
                    "rejectionRate": rejection_rate,
                    "files": s.files,
                    "rows": s.rows,
                    "accepted": s.accepted,
                    "rejected": s.rejected,
                    "warnings": s.warnings,
                    "errors": s.errors,
                    "_producer": self.producer(),
                    "_schemaURL": "https://github.com/malon64/floe/schemas/FloeQualityRunFacet.json"
                });
                json!([{
                    "namespace": self.config.namespace,
                    "name": u.source,
                    "facets": {
                        "schema": schema_facet,
                        "dataQualityMetrics": dq_facet,
                        "floeQualityRun": floe_facet
                    }
                }])
            }
            _ => json!([]),
        };

        // Build outputs: accepted sink always present; rejected sink when configured.
        let outputs = match uris {
            Some(u) => {
                let mut out = vec![json!({
                    "namespace": self.config.namespace,
                    "name": u.accepted,
                    "facets": {}
                })];
                if let Some(ref rej) = u.rejected {
                    out.push(json!({
                        "namespace": self.config.namespace,
                        "name": rej,
                        "facets": {}
                    }));
                }
                json!(out)
            }
            None => json!([]),
        };

        let body = json!({
            "eventType": event_type,
            "eventTime": event_time,
            "run": {
                "runId": entity_run_id,
                "facets": run_facets
            },
            "job": {
                "namespace": self.config.namespace,
                "name": job_name,
                "facets": {}
            },
            "inputs": inputs,
            "outputs": outputs,
            "producer": self.producer(),
            "schemaURL": "https://openlineage.io/spec/1-0-5/OpenLineage.json#/$defs/RunEvent"
        });
        self.post_event(body);
    }
}

struct EntityStats {
    files: u64,
    rows: u64,
    accepted: u64,
    rejected: u64,
    warnings: u64,
    errors: u64,
    schema_fields: Vec<(String, String)>,
}

fn ms_to_iso8601(ms: u128) -> String {
    let secs = (ms / 1000) as i64;
    let nanos = ((ms % 1000) * 1_000_000) as i64;
    match time::OffsetDateTime::from_unix_timestamp(secs) {
        Ok(dt) => {
            let ns = time::Duration::nanoseconds(nanos);
            let dt = dt.saturating_add(ns);
            dt.format(&time::format_description::well_known::Rfc3339)
                .unwrap_or_else(|_| ms.to_string())
        }
        Err(_) => ms.to_string(),
    }
}

impl RunObserver for OpenLineageObserver {
    fn on_event(&self, event: RunEvent) {
        match event {
            RunEvent::RunStarted { run_id, ts_ms, .. } => {
                // Reset circuit breaker at the start of each run so a recovered endpoint
                // is retried in subsequent runs within the same long-lived process.
                self.consecutive_failures.store(0, Ordering::Relaxed);
                self.circuit_open.store(false, Ordering::Relaxed);
                if let Ok(mut guard) = self.run_start_ms.lock() {
                    *guard = Some(ts_ms);
                }
                let event_time = ms_to_iso8601(ts_ms);
                let mut run_facets = json!({});
                if let Some(parent) = self.parent_run_facet() {
                    run_facets["parent"] = parent;
                }
                let body = json!({
                    "eventType": "START",
                    "eventTime": event_time,
                    "run": {
                        "runId": run_id,
                        "facets": run_facets
                    },
                    "job": {
                        "namespace": self.config.namespace,
                        "name": run_id,
                        "facets": {}
                    },
                    "inputs": [],
                    "outputs": [],
                    "producer": self.producer(),
                    "schemaURL": "https://openlineage.io/spec/1-0-5/OpenLineage.json#/$defs/RunEvent"
                });
                self.post_event(body);
            }
            RunEvent::EntityStarted {
                run_id,
                name,
                ts_ms,
            } => {
                // Use a per-entity run_id (overall_run_id.entity.name) so that
                // the START and COMPLETE events for the same entity share the same run_id.
                let entity_run_id = format!("{run_id}.entity.{name}");
                if let Ok(mut guard) = self.entity_start_ms.lock() {
                    guard.insert(name.clone(), ts_ms);
                }
                if let Ok(mut guard) = self.entity_run_ids.lock() {
                    guard.insert(name.clone(), entity_run_id.clone());
                }
                self.emit_entity_run_event(&entity_run_id, &name, "START", ts_ms, None, None);
            }
            RunEvent::EntityFinished {
                run_id,
                name,
                status,
                files,
                files_skipped: _,
                rows,
                accepted,
                rejected,
                warnings,
                errors,
                ts_ms,
            } => {
                let entity_run_id = self
                    .entity_run_ids
                    .lock()
                    .ok()
                    .and_then(|g| g.get(&name).cloned())
                    .unwrap_or_else(|| format!("{run_id}.entity.{name}"));
                let event_type = if status == "failed" || status == "aborted" {
                    "FAIL"
                } else {
                    "COMPLETE"
                };
                let schema_fields = self.entity_schemas.get(&name).cloned().unwrap_or_default();
                let stats = EntityStats {
                    files,
                    rows,
                    accepted,
                    rejected,
                    warnings,
                    errors,
                    schema_fields,
                };
                let uris = self.entity_uris.get(&name);
                self.emit_entity_run_event(
                    &entity_run_id,
                    &name,
                    event_type,
                    ts_ms,
                    Some(stats),
                    uris,
                );
            }
            RunEvent::RunFinished {
                run_id,
                status,
                ts_ms,
                ..
            } => {
                let event_type = if status == "failed" || status == "aborted" {
                    "FAIL"
                } else {
                    "COMPLETE"
                };
                let event_time = ms_to_iso8601(ts_ms);
                let mut run_facets = json!({});
                if let Some(parent) = self.parent_run_facet() {
                    run_facets["parent"] = parent;
                }
                let body = json!({
                    "eventType": event_type,
                    "eventTime": event_time,
                    "run": {
                        "runId": run_id,
                        "facets": run_facets
                    },
                    "job": {
                        "namespace": self.config.namespace,
                        "name": run_id,
                        "facets": {}
                    },
                    "inputs": [],
                    "outputs": [],
                    "producer": self.producer(),
                    "schemaURL": "https://openlineage.io/spec/1-0-5/OpenLineage.json#/$defs/RunEvent"
                });
                self.post_event(body);
            }
            _ => {}
        }
    }
}

pub fn build_observer(
    config: &LineageConfig,
    entities: &[EntityConfig],
) -> crate::FloeResult<Arc<dyn RunObserver>> {
    let obs = OpenLineageObserver::new(config, entities)?;
    Ok(Arc::new(obs))
}

impl OpenLineageObserver {
    pub fn is_circuit_open(&self) -> bool {
        self.circuit_open.load(Ordering::Relaxed)
    }

    pub fn consecutive_failures(&self) -> usize {
        self.consecutive_failures.load(Ordering::Relaxed)
    }
}