1- use datafusion:: execution:: SessionStateBuilder ;
1+ use async_trait:: async_trait;
2+ use datafusion:: error:: DataFusionError ;
3+ use datafusion:: execution:: { SessionState , SessionStateBuilder } ;
4+ use datafusion:: prelude:: SessionContext ;
25
36/// Trait called by the Arrow Flight endpoint that handles distributed parts of a DataFusion
47/// plan for building a DataFusion's [datafusion::prelude::SessionContext].
8+ #[ async_trait]
59pub trait SessionBuilder {
610 /// Takes a [SessionStateBuilder] and adds whatever is necessary for it to work, like
711/// custom extension codecs, custom physical optimization rules, UDFs, UDAFs, config
@@ -10,8 +14,9 @@ pub trait SessionBuilder {
1014/// Example: adding some custom extension plan codecs
1115///
1216/// ```rust
13- ///
1417/// # use std::sync::Arc;
18+ /// # use async_trait::async_trait;
19+ /// # use datafusion::error::DataFusionError;
1520/// # use datafusion::execution::runtime_env::RuntimeEnv;
1621/// # use datafusion::execution::{FunctionRegistry, SessionStateBuilder};
1722/// # use datafusion::physical_plan::ExecutionPlan;
@@ -33,22 +38,81 @@ pub trait SessionBuilder {
3338///
3439/// #[derive(Clone)]
3540/// struct CustomSessionBuilder;
41+ ///
42+ /// #[async_trait]
3643/// impl SessionBuilder for CustomSessionBuilder {
37- /// fn on_new_session (&self, mut builder: SessionStateBuilder) -> SessionStateBuilder {
44+ /// fn session_state_builder (&self, mut builder: SessionStateBuilder) -> Result< SessionStateBuilder, DataFusionError> {
3845/// // Add your UDFs, optimization rules, etc...
39- /// with_user_codec(builder, CustomExecCodec)
46+ /// Ok(with_user_codec(builder, CustomExecCodec))
47+ /// }
48+ /// }
49+ /// ```
50+ fn session_state_builder (
51+ & self ,
52+ builder : SessionStateBuilder ,
53+ ) -> Result < SessionStateBuilder , DataFusionError > {
54+ Ok ( builder)
55+ }
56+
57+ /// Modifies the [SessionState] and returns it. Same as [SessionBuilder::session_state_builder]
58+ /// but operating on an already built [SessionState].
59+ ///
60+ /// Example:
61+ ///
62+ /// ```rust
63+ /// # use async_trait::async_trait;
64+ /// # use datafusion::common::DataFusionError;
65+ /// # use datafusion::execution::SessionState;
66+ /// # use datafusion_distributed::SessionBuilder;
67+ ///
68+ /// #[derive(Clone)]
69+ /// struct CustomSessionBuilder;
70+ ///
71+ /// #[async_trait]
72+ /// impl SessionBuilder for CustomSessionBuilder {
73+ /// async fn session_state(&self, state: SessionState) -> Result<SessionState, DataFusionError> {
74+ /// // mutate the state adding any custom logic
75+ /// Ok(state)
4076/// }
4177/// }
4278/// ```
43- fn on_new_session ( & self , builder : SessionStateBuilder ) -> SessionStateBuilder ;
79+ async fn session_state ( & self , state : SessionState ) -> Result < SessionState , DataFusionError > {
80+ Ok ( state)
81+ }
82+
83+ /// Modifies the [SessionContext] and returns it. Same as [SessionBuilder::session_state_builder]
84+ /// or [SessionBuilder::session_state] but operation on an already built [SessionContext].
85+ ///
86+ /// Example:
87+ ///
88+ /// ```rust
89+ /// # use async_trait::async_trait;
90+ /// # use datafusion::common::DataFusionError;
91+ /// # use datafusion::prelude::SessionContext;
92+ /// # use datafusion_distributed::SessionBuilder;
93+ ///
94+ /// #[derive(Clone)]
95+ /// struct CustomSessionBuilder;
96+ ///
97+ /// #[async_trait]
98+ /// impl SessionBuilder for CustomSessionBuilder {
99+ /// async fn session_context(&self, ctx: SessionContext) -> Result<SessionContext, DataFusionError> {
100+ /// // mutate the context adding any custom logic
101+ /// Ok(ctx)
102+ /// }
103+ /// }
104+ /// ```
105+ async fn session_context (
106+ & self ,
107+ ctx : SessionContext ,
108+ ) -> Result < SessionContext , DataFusionError > {
109+ Ok ( ctx)
110+ }
44111}
45112
46113/// Noop implementation of the [SessionBuilder]. Used by default if no [SessionBuilder] is provided
47114/// while building the Arrow Flight endpoint.
115+ #[ derive( Debug , Clone ) ]
48116pub struct NoopSessionBuilder ;
49117
50- impl SessionBuilder for NoopSessionBuilder {
51- fn on_new_session ( & self , builder : SessionStateBuilder ) -> SessionStateBuilder {
52- builder
53- }
54- }
118+ impl SessionBuilder for NoopSessionBuilder { }
0 commit comments