datafusion_flight_sql_server/
session.rs

1use async_trait::async_trait;
2use datafusion::execution::context::SessionState;
3use tonic::{Request, Status};
4
5type Result<T, E = Status> = std::result::Result<T, E>;
6
7// SessionStateProvider is a trait used to provide a SessionState for a given
8// request.
9#[async_trait]
10pub trait SessionStateProvider: Sync + Send {
11    async fn new_context(&self, request: &Request<()>) -> Result<SessionState>;
12}
13
14// StaticSessionStateProvider is a simple implementation of SessionStateProvider that
15// uses a static SessionState.
16pub(crate) struct StaticSessionStateProvider {
17    state: SessionState,
18}
19
20impl StaticSessionStateProvider {
21    pub fn new(state: SessionState) -> Self {
22        Self { state }
23    }
24}
25
26#[async_trait]
27impl SessionStateProvider for StaticSessionStateProvider {
28    async fn new_context(&self, _request: &Request<()>) -> Result<SessionState> {
29        Ok(self.state.clone())
30    }
31}