Skip to main content

datafusion_openlineage/
rule.rs

1//! Plan-carried lineage marker and its lowering into the terminal node.
2//!
3//! OpenLineage instrumentation has three concerns with different needs (see ADR
4//! 0005): lineage *extraction* needs the optimized `LogicalPlan`; the START event
5//! and orchestration context need `&SessionState` and are async; the terminal
6//! COMPLETE/FAIL node needs to sit at the physical root and observe execution.
7//! Only the [`QueryPlanner`] seam has `&SessionState`, so the planning-time work
8//! lives there — but the terminal node is installed the composable, DataFusion-
9//! idiomatic way: a registered [`ExtensionPlanner`] lowers a plan-carried marker
10//! into [`OpenLineageExec`], rather than the planner hand-wrapping the physical
11//! root.
12//!
13//! Flow, all under one `run_id`. First, [`OpenLineageQueryPlanner`] (the
14//! [`QueryPlanner`]) extracts lineage from the optimized logical plan, resolves
15//! the async [`LineageContextProvider`], emits START, builds the COMPLETE
16//! template, and wraps the *logical* plan in a [`LineageMarker`] carrying that
17//! template. Then physical planning lowers the marker via
18//! [`LineageExtensionPlanner`] into an [`OpenLineageExec`] at the root, which
19//! emits COMPLETE/FAIL at end of execution.
20//!
21//! A `LogicalPlan::Extension` requires a registered `ExtensionPlanner` (the
22//! default physical planner errors on unknown extension nodes), so the planner
23//! delegates physical planning to a [`DefaultPhysicalPlanner`] configured with
24//! [`LineageExtensionPlanner`].
25
26use std::cmp::Ordering;
27use std::fmt;
28use std::hash::{Hash, Hasher};
29use std::sync::Arc;
30
31use async_trait::async_trait;
32use datafusion::common::{DFSchemaRef, Result};
33use datafusion::dataframe::DataFrame;
34use datafusion::execution::context::{QueryPlanner, SessionContext, SessionState};
35use datafusion::logical_expr::{
36    Expr, Extension, InvariantLevel, LogicalPlan, UserDefinedLogicalNode,
37    UserDefinedLogicalNodeCore,
38};
39use datafusion::physical_plan::ExecutionPlan;
40use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner};
41use uuid::Uuid;
42
43use crate::builder::{complete_event, fail_event, start_event};
44use crate::client::OpenLineageClient;
45use crate::config::OpenLineageConfig;
46use crate::context::{LineageContext, LineageContextProvider};
47use crate::event::RunEvent;
48use crate::exec::OpenLineageExec;
49use crate::extract::{QueryLineage, extract};
50
51tokio::task_local! {
52    /// Set while [`OpenLineageQueryPlanner::execute_ddl_with_lineage`] runs the DDL,
53    /// so the *nested* `create_physical_plan` that DataFusion triggers for the body
54    /// (e.g. `create_memory_table` collecting the CTAS SELECT) does not emit a
55    /// second, spurious run. Task-local rather than a shared flag so it is scoped to
56    /// this one execution's task tree and never suppresses a concurrent query on
57    /// another task. See [`suppressing_nested_lineage`].
58    static SUPPRESS_NESTED_LINEAGE: ();
59}
60
61/// True when the current task is inside an `execute_ddl_with_lineage` call, i.e.
62/// any `create_physical_plan` here is the nested body of a DDL run already being
63/// reported and must not emit its own events.
64fn nested_lineage_suppressed() -> bool {
65    SUPPRESS_NESTED_LINEAGE.try_with(|()| ()).is_ok()
66}
67
68// ---------------------------------------------------------------------------
69// The plan-carried marker.
70// ---------------------------------------------------------------------------
71
72/// A logical no-op wrapping the real plan, carrying the per-query lineage payload
73/// from [`OpenLineageQueryPlanner`] (which has `&SessionState`) to
74/// [`LineageExtensionPlanner`] (which installs the terminal node). Schema-
75/// transparent: it reports its input's schema so optimization and physical
76/// planning treat it as a pass-through.
77#[derive(Clone)]
78pub struct LineageMarker {
79    input: LogicalPlan,
80    /// COMPLETE event template, built at plan time; cloned into the terminal
81    /// [`OpenLineageExec`] at lowering and mutated into FAIL there on error.
82    complete: RunEvent,
83    client: OpenLineageClient,
84    producer: String,
85}
86
87impl LineageMarker {
88    /// Wrap `input`, carrying the COMPLETE template the terminal
89    /// [`OpenLineageExec`] emits at end of execution. Usually built for you by
90    /// [`LineageHandle::into_marker`]; public so a host can construct the marker
91    /// when composing lineage into its own planner.
92    pub fn new(
93        input: LogicalPlan,
94        complete: RunEvent,
95        client: OpenLineageClient,
96        producer: String,
97    ) -> Self {
98        Self {
99            input,
100            complete,
101            client,
102            producer,
103        }
104    }
105}
106
107impl fmt::Debug for LineageMarker {
108    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109        f.debug_struct("LineageMarker").finish_non_exhaustive()
110    }
111}
112
113// Identity is the run id plus the wrapped plan: enough to distinguish markers,
114// and the payload (client/template) is behavioral rather than structural.
115impl PartialEq for LineageMarker {
116    fn eq(&self, other: &Self) -> bool {
117        self.complete.run.run_id == other.complete.run.run_id && self.input == other.input
118    }
119}
120impl Eq for LineageMarker {}
121impl PartialOrd for LineageMarker {
122    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
123        self.complete
124            .run
125            .run_id
126            .partial_cmp(&other.complete.run.run_id)
127    }
128}
129impl Hash for LineageMarker {
130    fn hash<H: Hasher>(&self, state: &mut H) {
131        self.complete.run.run_id.hash(state);
132    }
133}
134
135impl UserDefinedLogicalNodeCore for LineageMarker {
136    fn name(&self) -> &str {
137        "LineageMarker"
138    }
139
140    fn inputs(&self) -> Vec<&LogicalPlan> {
141        vec![&self.input]
142    }
143
144    fn schema(&self) -> &DFSchemaRef {
145        self.input.schema()
146    }
147
148    fn check_invariants(&self, _check: InvariantLevel) -> Result<()> {
149        Ok(())
150    }
151
152    fn expressions(&self) -> Vec<Expr> {
153        vec![]
154    }
155
156    fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
157        write!(f, "LineageMarker")
158    }
159
160    fn with_exprs_and_inputs(
161        &self,
162        _exprs: Vec<Expr>,
163        mut inputs: Vec<LogicalPlan>,
164    ) -> Result<Self> {
165        Ok(Self {
166            input: inputs.pop().expect("LineageMarker has one input"),
167            complete: self.complete.clone(),
168            client: self.client.clone(),
169            producer: self.producer.clone(),
170        })
171    }
172}
173
174// ---------------------------------------------------------------------------
175// Lowering: marker -> OpenLineageExec.
176// ---------------------------------------------------------------------------
177
178/// Lowers a [`LineageMarker`] into an [`OpenLineageExec`] at physical-planning
179/// time. Register it on the session's physical planner (see
180/// [`crate::session::instrument_session_state`]).
181#[derive(Debug, Default)]
182pub struct LineageExtensionPlanner;
183
184#[async_trait]
185impl ExtensionPlanner for LineageExtensionPlanner {
186    async fn plan_extension(
187        &self,
188        _planner: &dyn PhysicalPlanner,
189        node: &dyn UserDefinedLogicalNode,
190        _logical_inputs: &[&LogicalPlan],
191        physical_inputs: &[Arc<dyn ExecutionPlan>],
192        _session_state: &SessionState,
193    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
194        // Not our node: let another extension planner handle it.
195        let Some(marker) = node.as_any().downcast_ref::<LineageMarker>() else {
196            return Ok(None);
197        };
198        let inner = physical_inputs
199            .first()
200            .expect("LineageMarker has one physical input")
201            .clone();
202        Ok(Some(OpenLineageExec::new(
203            inner,
204            marker.client.clone(),
205            marker.complete.clone(),
206            marker.producer.clone(),
207        )))
208    }
209}
210
211// ---------------------------------------------------------------------------
212// The reusable planning-time lineage step.
213// ---------------------------------------------------------------------------
214
215/// The per-query lineage payload produced by [`begin_lineage`], carried under one
216/// `run_id` from the planning-time step to the terminal event.
217///
218/// A host composing lineage with another concern (e.g. running lineage's START
219/// step *after* a policy gate, from its own single [`QueryPlanner`]) holds this
220/// between [`begin_lineage`] and [`Self::into_marker`] / the `emit_*` methods.
221pub struct LineageHandle {
222    run_id: Uuid,
223    lineage: QueryLineage,
224    context: LineageContext,
225}
226
227impl LineageHandle {
228    /// The run id START was emitted under; COMPLETE/FAIL must share it.
229    pub fn run_id(&self) -> Uuid {
230        self.run_id
231    }
232
233    /// Wrap `plan` in a [`LineageMarker`] carrying the COMPLETE template, so a
234    /// registered [`LineageExtensionPlanner`] lowers it into the terminal
235    /// [`OpenLineageExec`] at the physical root (the composable path).
236    ///
237    /// Borrows so the handle stays available for [`Self::emit_fail`] if the
238    /// subsequent physical planning errors.
239    pub fn to_marker(
240        &self,
241        plan: LogicalPlan,
242        client: OpenLineageClient,
243        config: &OpenLineageConfig,
244    ) -> LogicalPlan {
245        let complete = complete_event(self.run_id, &self.lineage, &self.context, config);
246        LogicalPlan::Extension(Extension {
247            node: Arc::new(LineageMarker::new(
248                plan,
249                complete,
250                client,
251                config.producer.clone(),
252            )),
253        })
254    }
255
256    /// Emit FAIL for a planning error that occurs *after* START, before any
257    /// [`OpenLineageExec`] exists to observe execution. Under the same `run_id`.
258    pub fn emit_fail(&self, client: &OpenLineageClient, config: &OpenLineageConfig, err: &str) {
259        client.emit(fail_event(
260            self.run_id,
261            &self.lineage,
262            &self.context,
263            config,
264            err,
265        ));
266    }
267
268    /// Emit COMPLETE directly, under the same `run_id`, with `eventTime` refreshed
269    /// to now (so run duration is meaningful). For paths that run the query
270    /// *outside* a physical plan — there is no [`OpenLineageExec`] to emit the
271    /// terminal event, so the caller emits it (e.g. the CTAS / `CREATE VIEW` DDL
272    /// path, which materializes internally and hands back an empty result).
273    pub fn emit_complete(&self, client: &OpenLineageClient, config: &OpenLineageConfig) {
274        let mut event = complete_event(self.run_id, &self.lineage, &self.context, config);
275        event.event_time = chrono::Utc::now().to_rfc3339();
276        client.emit(event);
277    }
278
279    /// Fold SQL text into the lineage when the context provider didn't supply it
280    /// (a path that has the raw statement in hand — e.g. the DDL path).
281    pub fn set_sql_if_absent(&mut self, sql: &str) {
282        if self.lineage.sql.is_none() {
283            self.lineage.sql = Some(sql.to_string());
284        }
285    }
286}
287
288/// Planning-time lineage work, decoupled from the [`QueryPlanner`] trait: extract
289/// lineage from `plan`, resolve the async [`LineageContextProvider`], and — unless
290/// the query touches no datasets — mint a `run_id` and emit START.
291///
292/// Returns a [`LineageHandle`] to carry the run under one id to the terminal event
293/// (via [`LineageHandle::into_marker`]), or `None` when lineage is suppressed (no
294/// inputs and no outputs — `information_schema` introspection, `SET`/`SHOW`,
295/// metadata probes; or a nested DDL body), in which case no START fired and the
296/// caller must emit nothing.
297///
298/// This is the reusable step for hosts that sequence lineage into their own
299/// single `QueryPlanner` (e.g. after a policy gate) rather than installing
300/// [`OpenLineageQueryPlanner`] as the session planner.
301pub async fn begin_lineage(
302    client: &OpenLineageClient,
303    context: &dyn LineageContextProvider,
304    config: &OpenLineageConfig,
305    plan: &LogicalPlan,
306    session_state: &SessionState,
307) -> Option<LineageHandle> {
308    // A `create_physical_plan` nested inside `execute_ddl_with_lineage` is the
309    // DDL body (e.g. the CTAS SELECT that `create_memory_table` collects); the
310    // enclosing DDL run already reports it, so emit nothing here.
311    if nested_lineage_suppressed() {
312        return None;
313    }
314
315    let mut lineage = extract(plan, config);
316    let cx = context.context(session_state).await;
317    // The SQL text isn't recoverable from the plan; take it from the
318    // host-supplied context (absent on non-SQL paths, e.g. ingest).
319    lineage.sql = cx.sql.clone();
320
321    // Suppress lineage for queries that touch no datasets — information_schema
322    // introspection, `SET`/`SHOW`, metadata-RPC probes. They carry no input or
323    // output, so a START/COMPLETE pair only adds a dangling job node to the graph.
324    if lineage.inputs.is_empty() && lineage.outputs.is_empty() {
325        return None;
326    }
327
328    let run_id = cx.run_id.unwrap_or_else(Uuid::now_v7);
329    client.emit(start_event(run_id, &lineage, &cx, config));
330    Some(LineageHandle {
331        run_id,
332        lineage,
333        context: cx,
334    })
335}
336
337// ---------------------------------------------------------------------------
338// The query planner: extract + START + inject the marker.
339// ---------------------------------------------------------------------------
340
341/// A [`QueryPlanner`] that emits OpenLineage events around a query.
342///
343/// It does the `&SessionState`-bound, async work — extract lineage, resolve
344/// context, emit START, mint the `run_id`, emit FAIL on a planning error — then
345/// hands off to physical planning by wrapping the logical plan in a
346/// [`LineageMarker`]. The registered [`LineageExtensionPlanner`] lowers that
347/// marker into the terminal [`OpenLineageExec`]. Built by
348/// [`crate::session::instrument_session_state`].
349pub struct OpenLineageQueryPlanner {
350    client: OpenLineageClient,
351    context: Arc<dyn LineageContextProvider>,
352    config: OpenLineageConfig,
353    /// Physical planner that knows how to lower [`LineageMarker`]; composes any
354    /// extension planners the host already had.
355    physical: Arc<DefaultPhysicalPlanner>,
356}
357
358impl OpenLineageQueryPlanner {
359    /// Build a planner whose physical planning lowers our marker plus
360    /// `extra_extension_planners` (any the host session already registered).
361    pub fn new(
362        client: OpenLineageClient,
363        context: Arc<dyn LineageContextProvider>,
364        config: OpenLineageConfig,
365        extra_extension_planners: Vec<Arc<dyn ExtensionPlanner + Send + Sync>>,
366    ) -> Self {
367        let mut planners: Vec<Arc<dyn ExtensionPlanner + Send + Sync>> =
368            vec![Arc::new(LineageExtensionPlanner)];
369        planners.extend(extra_extension_planners);
370        Self {
371            client,
372            context,
373            config,
374            physical: Arc::new(DefaultPhysicalPlanner::with_extension_planners(planners)),
375        }
376    }
377
378    /// Planning-time lineage work shared by the `QueryPlanner` path and the
379    /// `SessionContext`-level DDL path (see [`crate::session::OpenLineageSqlExt`]):
380    /// see the free [`begin_lineage`] function this delegates to.
381    async fn begin_lineage(
382        &self,
383        plan: &LogicalPlan,
384        session_state: &SessionState,
385    ) -> Option<LineageHandle> {
386        begin_lineage(
387            &self.client,
388            self.context.as_ref(),
389            &self.config,
390            plan,
391            session_state,
392        )
393        .await
394    }
395
396    /// Execute a DDL-with-input statement (CTAS / CREATE VIEW) that DataFusion
397    /// runs *outside* the `QueryPlanner` hook, emitting lineage around it.
398    ///
399    /// `SessionContext::execute_logical_plan` dispatches these DDL variants to its
400    /// own `create_memory_table` / `create_view` before any `QueryPlanner` sees the
401    /// wrapper (it only ever plans the stripped SELECT body), so the planner path
402    /// captures the inputs but never the created table as an output. This runs
403    /// [`extract`] on the *full* DDL plan (so the output dataset, its schema, and
404    /// column lineage are captured), emits START, delegates the actual creation to
405    /// `execute_logical_plan` — reusing DataFusion's registration logic, including
406    /// every `if_not_exists` / `or_replace` branch — then emits COMPLETE on success
407    /// or FAIL on error, under the same `run_id`.
408    ///
409    /// `raw_sql` is folded into the lineage when the context provider didn't supply
410    /// SQL text, since this path *does* have the original statement in hand.
411    ///
412    /// COMPLETE here does not carry an `outputStatistics.rowCount`: DataFusion
413    /// materializes the CTAS body internally and hands back an empty result, so
414    /// there is no stream for an `OpenLineageExec` to count. The output edge,
415    /// schema, lifecycle, and column lineage are all present; runtime row stats for
416    /// this path are a documented follow-up.
417    pub(crate) async fn execute_ddl_with_lineage(
418        &self,
419        ctx: &SessionContext,
420        plan: LogicalPlan,
421        raw_sql: &str,
422    ) -> Result<DataFrame> {
423        let Some(mut handle) = self.begin_lineage(&plan, &ctx.state()).await else {
424            // No datasets touched — nothing to report; just run it.
425            return ctx.execute_logical_plan(plan).await;
426        };
427        // This path has the SQL in hand even when the context provider omitted it.
428        handle.set_sql_if_absent(raw_sql);
429
430        // Run the DDL with nested-lineage suppression set: `execute_logical_plan`
431        // dispatches CTAS/CREATE VIEW to `create_memory_table`/`create_view`, which
432        // collect the SELECT body back through *this* planner; without the guard
433        // that body would emit its own (input-only) run alongside this DDL run.
434        let result = SUPPRESS_NESTED_LINEAGE
435            .scope((), ctx.execute_logical_plan(plan))
436            .await;
437        match result {
438            Ok(df) => {
439                // No `OpenLineageExec` on this path (DataFusion materializes the CTAS
440                // body internally), so emit COMPLETE directly under the run id.
441                handle.emit_complete(&self.client, &self.config);
442                Ok(df)
443            }
444            Err(err) => {
445                handle.emit_fail(&self.client, &self.config, &err.to_string());
446                Err(err)
447            }
448        }
449    }
450}
451
452impl fmt::Debug for OpenLineageQueryPlanner {
453    // `DefaultPhysicalPlanner` is not `Debug`, so don't try to print it.
454    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
455        f.debug_struct("OpenLineageQueryPlanner")
456            .finish_non_exhaustive()
457    }
458}
459
460#[async_trait]
461impl QueryPlanner for OpenLineageQueryPlanner {
462    async fn create_physical_plan(
463        &self,
464        logical_plan: &LogicalPlan,
465        session_state: &SessionState,
466    ) -> Result<Arc<dyn ExecutionPlan>> {
467        // Extract lineage, resolve context, and emit START — or, when the query
468        // touches no datasets, plan straight through without a marker so no events
469        // fire.
470        let Some(handle) = self.begin_lineage(logical_plan, session_state).await else {
471            return self
472                .physical
473                .create_physical_plan(logical_plan, session_state)
474                .await;
475        };
476
477        // Carry the COMPLETE template into the physical phase via the plan itself;
478        // the extension planner lowers it into an OpenLineageExec at the root that
479        // emits COMPLETE/FAIL at end of execution, under this same run id.
480        let wrapped = handle.to_marker(logical_plan.clone(), self.client.clone(), &self.config);
481
482        match self
483            .physical
484            .create_physical_plan(&wrapped, session_state)
485            .await
486        {
487            Ok(plan) => Ok(plan),
488            Err(err) => {
489                // Planning failed outright — no execution to observe, emit FAIL now.
490                handle.emit_fail(&self.client, &self.config, &err.to_string());
491                Err(err)
492            }
493        }
494    }
495}
496
497#[cfg(test)]
498mod tests {
499    use std::collections::hash_map::DefaultHasher;
500
501    use datafusion::logical_expr::LogicalPlanBuilder;
502
503    use super::*;
504    use crate::QueryLineage;
505    use crate::context::LineageContext;
506    use crate::transport::NoopTransport;
507
508    // `name`/`inputs`/`schema`/`expressions`/`with_exprs_and_inputs` exist on both
509    // `UserDefinedLogicalNodeCore` and the object-safe `UserDefinedLogicalNode`
510    // (blanket impl), so calls go through this alias to disambiguate.
511    use datafusion::logical_expr::UserDefinedLogicalNodeCore as NodeCore;
512
513    /// A `LineageMarker` over a trivial empty-relation plan, tagged with `run_id`.
514    /// The marker fields are private, so these tests live inline.
515    fn marker(run_id: Uuid) -> LineageMarker {
516        let input = LogicalPlanBuilder::empty(false).build().unwrap();
517        let config = OpenLineageConfig::default();
518        let complete = complete_event(
519            run_id,
520            &QueryLineage::default(),
521            &LineageContext::default(),
522            &config,
523        );
524        LineageMarker {
525            input,
526            complete,
527            client: OpenLineageClient::new(Arc::new(NoopTransport)),
528            producer: config.producer,
529        }
530    }
531
532    fn hash_of(m: &LineageMarker) -> u64 {
533        let mut h = DefaultHasher::new();
534        m.hash(&mut h);
535        h.finish()
536    }
537
538    // `OpenLineageClient::new` spawns a background drain, so these need a runtime.
539    #[tokio::test]
540    async fn node_core_is_schema_transparent_and_expr_free() {
541        let m = marker(Uuid::now_v7());
542        assert_eq!(NodeCore::name(&m), "LineageMarker");
543        assert_eq!(NodeCore::inputs(&m).len(), 1);
544        // Schema-transparent: it reports its single input's schema verbatim.
545        assert_eq!(NodeCore::schema(&m), NodeCore::inputs(&m)[0].schema());
546        assert!(NodeCore::expressions(&m).is_empty());
547        assert!(NodeCore::check_invariants(&m, InvariantLevel::Always).is_ok());
548        assert_eq!(format!("{m:?}"), "LineageMarker { .. }");
549    }
550
551    #[tokio::test]
552    async fn with_exprs_and_inputs_rebuilds_preserving_payload() {
553        let run_id = Uuid::now_v7();
554        let m = marker(run_id);
555        let new_input = LogicalPlanBuilder::empty(true).build().unwrap();
556        let rebuilt = NodeCore::with_exprs_and_inputs(&m, vec![], vec![new_input.clone()]).unwrap();
557        // The wrapped input swaps; the run-id payload is carried through.
558        assert_eq!(NodeCore::inputs(&rebuilt)[0], &new_input);
559        assert_eq!(rebuilt.complete.run.run_id, run_id);
560    }
561
562    #[tokio::test]
563    async fn identity_keys_on_run_id_and_input() {
564        let run_id = Uuid::now_v7();
565        // Same run id + same (empty) plan → equal, same hash, equal ordering.
566        let a = marker(run_id);
567        let b = marker(run_id);
568        assert_eq!(a, b);
569        assert_eq!(hash_of(&a), hash_of(&b));
570        assert_eq!(a.partial_cmp(&b), Some(Ordering::Equal));
571
572        // Different run id → not equal, ordering follows the run id.
573        let c = marker(Uuid::now_v7());
574        assert_ne!(a, c);
575        assert_eq!(
576            a.partial_cmp(&c),
577            a.complete.run.run_id.partial_cmp(&c.complete.run.run_id)
578        );
579    }
580
581    // --- the reusable `begin_lineage` step (for hosts sequencing it themselves) ---
582
583    use std::sync::Mutex;
584
585    use datafusion::execution::context::SessionContext;
586
587    use crate::context::StaticContextProvider;
588    use crate::event::{RunEvent, RunEventType};
589    use crate::transport::{Transport, TransportError};
590
591    /// Records emitted events for assertions.
592    #[derive(Clone, Default, Debug)]
593    struct Recording {
594        events: std::sync::Arc<Mutex<Vec<RunEvent>>>,
595    }
596    #[async_trait]
597    impl Transport for Recording {
598        async fn emit(&self, event: &RunEvent) -> std::result::Result<(), TransportError> {
599            self.events.lock().unwrap().push(event.clone());
600            Ok(())
601        }
602    }
603
604    async fn table_select() -> (SessionState, LogicalPlan) {
605        let ctx = SessionContext::new();
606        ctx.sql("CREATE TABLE t AS VALUES (1)").await.unwrap();
607        let plan = ctx
608            .state()
609            .create_logical_plan("SELECT * FROM t")
610            .await
611            .unwrap();
612        (ctx.state(), plan)
613    }
614
615    #[tokio::test]
616    async fn begin_lineage_emits_start_and_yields_a_marker() {
617        let (state, plan) = table_select().await;
618        let rec = Recording::default();
619        let client = OpenLineageClient::new(std::sync::Arc::new(rec.clone()));
620        let config = OpenLineageConfig::default();
621        let ctx = StaticContextProvider::default();
622
623        let handle = begin_lineage(&client, &ctx, &config, &plan, &state)
624            .await
625            .expect("query touches a dataset -> a run begins");
626
627        // START was emitted under the handle's run id.
628        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
629        let events = rec.events.lock().unwrap();
630        let start = events
631            .iter()
632            .find(|e| e.event_type == RunEventType::Start)
633            .expect("START emitted");
634        assert_eq!(start.run.run_id, handle.run_id());
635
636        // The handle wraps a plan in a LineageMarker for the extension planner.
637        let wrapped = handle.to_marker(plan.clone(), client.clone(), &config);
638        assert!(matches!(&wrapped, LogicalPlan::Extension(e) if e.node.name() == "LineageMarker"));
639    }
640
641    #[tokio::test]
642    async fn begin_lineage_returns_none_for_no_dataset_query() {
643        // A metadata-only statement touches no datasets -> no run, no START.
644        let ctx_df = SessionContext::new();
645        let plan = ctx_df
646            .state()
647            .create_logical_plan("SET a = 1")
648            .await
649            .unwrap();
650        let rec = Recording::default();
651        let client = OpenLineageClient::new(std::sync::Arc::new(rec.clone()));
652        let config = OpenLineageConfig::default();
653        let ctx = StaticContextProvider::default();
654
655        let handle = begin_lineage(&client, &ctx, &config, &plan, &ctx_df.state()).await;
656        assert!(handle.is_none(), "no datasets -> no run begins");
657
658        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
659        assert!(
660            rec.events.lock().unwrap().is_empty(),
661            "no events for a no-dataset query"
662        );
663    }
664}