datafusion_postgres/hooks/
mod.rs

1pub mod permissions;
2pub mod set_show;
3pub mod transactions;
4
5use async_trait::async_trait;
6
7use datafusion::common::ParamValues;
8use datafusion::logical_expr::LogicalPlan;
9use datafusion::prelude::SessionContext;
10use datafusion::sql::sqlparser::ast::Statement;
11use futures::Sink;
12use pgwire::api::results::Response;
13use pgwire::api::ClientInfo;
14use pgwire::error::{PgWireError, PgWireResult};
15use pgwire::messages::PgWireBackendMessage;
16
17#[async_trait]
18pub trait HookClient: ClientInfo + Send + Sync {
19    async fn send_message(&mut self, item: PgWireBackendMessage) -> PgWireResult<()>;
20}
21
22#[async_trait]
23impl<S> HookClient for S
24where
25    S: ClientInfo + Sink<PgWireBackendMessage> + Send + Sync + Unpin,
26    PgWireError: From<<S as Sink<PgWireBackendMessage>>::Error>,
27{
28    async fn send_message(&mut self, item: PgWireBackendMessage) -> PgWireResult<()> {
29        use futures::SinkExt;
30        self.send(item).await.map_err(PgWireError::from)
31    }
32}
33
34#[async_trait]
35pub trait QueryHook: Send + Sync {
36    /// called in simple query handler to return response directly
37    async fn handle_simple_query(
38        &self,
39        statement: &Statement,
40        session_context: &SessionContext,
41        client: &mut dyn HookClient,
42    ) -> Option<PgWireResult<Response>>;
43
44    /// called at extended query parse phase, for generating `LogicalPlan`from statement
45    async fn handle_extended_parse_query(
46        &self,
47        sql: &Statement,
48        session_context: &SessionContext,
49        client: &(dyn ClientInfo + Send + Sync),
50    ) -> Option<PgWireResult<LogicalPlan>>;
51
52    /// called at extended query execute phase, for query execution
53    async fn handle_extended_query(
54        &self,
55        statement: &Statement,
56        logical_plan: &LogicalPlan,
57        params: &ParamValues,
58        session_context: &SessionContext,
59        client: &mut dyn HookClient,
60    ) -> Option<PgWireResult<Response>>;
61}
datafusion_postgres/hooks/mod.rs

datafusion_postgres/hooks/
mod.rs
