datafusion_openlineage/
session.rs

1//! The customer-facing entry point: [`OpenLineage::builder`].
2//!
3//! One builder assembles the three pieces an integration needs — a
4//! [`Transport`]/[`OpenLineageClient`], a [`LineageContextProvider`], and an
5//! [`OpenLineageConfig`] — and installs them on a `SessionState`. The common
6//! case is [`OpenLineageBuilder::from_env`] followed by
7//! [`OpenLineageBuilder::instrument`].
8
9use std::sync::Arc;
10
11use datafusion::execution::SessionStateBuilder;
12use datafusion::execution::context::SessionState;
13
14use openlineage_client::{ClientError, LineageContext, OpenLineageClient, Transport};
15
16use crate::config::{DataFusionConfig, OpenLineageConfig};
17use crate::context::{LineageContextProvider, StaticContextProvider};
18use crate::rule::OpenLineageQueryPlanner;
19
20/// Entry point for instrumenting a DataFusion session with OpenLineage.
21///
22/// Construct a [`OpenLineageBuilder`] with [`OpenLineage::builder`].
23#[derive(Debug)]
24pub struct OpenLineage;
25
26impl OpenLineage {
27    /// Start configuring OpenLineage instrumentation.
28    pub fn builder() -> OpenLineageBuilder {
29        OpenLineageBuilder::default()
30    }
31}
32
33/// Builds and installs OpenLineage instrumentation on a [`SessionState`].
34///
35/// Set a transport (or client) and, optionally, a context provider and config;
36/// then call [`Self::instrument`]. [`Self::from_env`] is the one-call path that
37/// derives everything from the standard OpenLineage environment.
38#[derive(Default)]
39pub struct OpenLineageBuilder {
40    client: Option<OpenLineageClient>,
41    transport: Option<Arc<dyn Transport>>,
42    context: Option<Arc<dyn LineageContextProvider>>,
43    config: Option<OpenLineageConfig>,
44}
45
46impl std::fmt::Debug for OpenLineageBuilder {
47    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        f.debug_struct("OpenLineageBuilder")
49            .field("has_client", &self.client.is_some())
50            .field("has_transport", &self.transport.is_some())
51            .field("has_context", &self.context.is_some())
52            .field("config", &self.config)
53            .finish()
54    }
55}
56
57impl OpenLineageBuilder {
58    /// Derive client, context, and config from the standard OpenLineage
59    /// environment in one call.
60    ///
61    /// Builds the client via [`OpenLineageClient::from_env`] (reads
62    /// `OPENLINEAGE_URL` / `OPENLINEAGE_ENDPOINT` / `OPENLINEAGE_API_KEY`), the
63    /// config via [`OpenLineageConfig::for_datafusion_from_env`] (reads
64    /// `OPENLINEAGE_NAMESPACE` / `OPENLINEAGE_TIMEOUT_MS`), and the context via
65    /// [`LineageContext::from_env`] (reads the parent-run variables). Anything
66    /// set explicitly on the builder beforehand is preserved.
67    ///
68    /// # Errors
69    /// Returns [`ClientError`] if the client cannot be built — e.g. called
70    /// outside a Tokio runtime, or `OPENLINEAGE_URL` is malformed.
71    pub fn from_env(mut self) -> Result<Self, ClientError> {
72        let config = self
73            .config
74            .take()
75            .unwrap_or_else(OpenLineageConfig::for_datafusion_from_env);
76        if self.client.is_none() && self.transport.is_none() {
77            self.client = Some(OpenLineageClient::from_env()?);
78        }
79        if self.context.is_none() {
80            let ctx = LineageContext::from_env(&config);
81            self.context = Some(Arc::new(StaticContextProvider(ctx)));
82        }
83        self.config = Some(config);
84        Ok(self)
85    }
86
87    /// Use a pre-built [`OpenLineageClient`]. Takes precedence over
88    /// [`Self::transport`].
89    pub fn client(mut self, client: OpenLineageClient) -> Self {
90        self.client = Some(client);
91        self
92    }
93
94    /// Drain events into `transport`. A client is constructed around it at
95    /// [`Self::instrument`] time. Ignored if [`Self::client`] is also set.
96    pub fn transport(mut self, transport: Arc<dyn Transport>) -> Self {
97        self.transport = Some(transport);
98        self
99    }
100
101    /// Set the per-query context provider. Defaults to an empty
102    /// [`StaticContextProvider`].
103    pub fn context(mut self, context: Arc<dyn LineageContextProvider>) -> Self {
104        self.context = Some(context);
105        self
106    }
107
108    /// Set the static config. Defaults to [`OpenLineageConfig::for_datafusion`].
109    pub fn config(mut self, config: OpenLineageConfig) -> Self {
110        self.config = Some(config);
111        self
112    }
113
114    /// Install the instrumentation on `state`, returning the wired
115    /// `SessionState`.
116    ///
117    /// Sets the session's query planner to an [`OpenLineageQueryPlanner`]. A
118    /// pre-existing custom `QueryPlanner` is replaced; standard physical planning
119    /// (including extension nodes) is preserved because the new planner delegates
120    /// to a `DefaultPhysicalPlanner`.
121    ///
122    /// # Panics
123    /// Panics if no client or transport was set and a default client must be
124    /// built outside a Tokio runtime. Prefer [`Self::from_env`] (which surfaces a
125    /// [`ClientError`]) or set a client built with
126    /// [`OpenLineageClient::try_new`].
127    pub fn instrument(self, state: SessionState) -> SessionState {
128        let client = self.client.unwrap_or_else(|| match self.transport {
129            Some(transport) => OpenLineageClient::new(transport),
130            None => OpenLineageClient::noop(),
131        });
132        let context = self
133            .context
134            .unwrap_or_else(|| Arc::new(StaticContextProvider::default()));
135        let config = self
136            .config
137            .unwrap_or_else(OpenLineageConfig::for_datafusion);
138
139        let planner = Arc::new(OpenLineageQueryPlanner::new(
140            client,
141            context,
142            config,
143            Vec::new(),
144        ));
145        SessionStateBuilder::from(state)
146            .with_query_planner(planner)
147            .build()
148    }
149}
150
151/// Instrument `state` with an explicit client, context provider, and config.
152///
153/// A thin wrapper over [`OpenLineage::builder`] for callers that already hold
154/// the three pieces. Prefer the builder ([`OpenLineage::builder`]) for new code.
155pub fn instrument_session_state(
156    state: SessionState,
157    client: OpenLineageClient,
158    context: Arc<dyn LineageContextProvider>,
159    config: OpenLineageConfig,
160) -> SessionState {
161    OpenLineage::builder()
162        .client(client)
163        .context(context)
164        .config(config)
165        .instrument(state)
166}
167
168/// Instrument `state` with an empty [`StaticContextProvider`].
169///
170/// A thin wrapper over [`OpenLineage::builder`]. Prefer the builder for new code.
171pub fn instrument_session_state_simple(
172    state: SessionState,
173    client: OpenLineageClient,
174    config: OpenLineageConfig,
175) -> SessionState {
176    OpenLineage::builder()
177        .client(client)
178        .config(config)
179        .instrument(state)
180}
datafusion_openlineage/session.rs

datafusion_openlineage/
session.rs
