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}