datafusion_federation/sql/
executor.rs

1use async_trait::async_trait;
2use core::fmt;
3use datafusion::{
4    arrow::datatypes::SchemaRef, error::Result, physical_plan::SendableRecordBatchStream,
5    sql::sqlparser::ast, sql::unparser::dialect::Dialect,
6};
7use std::sync::Arc;
8
9pub type SQLExecutorRef = Arc<dyn SQLExecutor>;
10pub type AstAnalyzer = Box<dyn Fn(ast::Statement) -> Result<ast::Statement>>;
11
12#[async_trait]
13pub trait SQLExecutor: Sync + Send {
14    /// Executor name
15    fn name(&self) -> &str;
16
17    /// Executor compute context allows differentiating the remote compute context
18    /// such as authorization or active database.
19    ///
20    /// Note: returning None here may cause incorrect federation with other providers of the
21    /// same name that also have a compute_context of None.
22    /// Instead try to return a unique string that will never match any other
23    /// provider's context.
24    fn compute_context(&self) -> Option<String>;
25
26    /// The specific SQL dialect (currently supports 'sqlite', 'postgres', 'flight')
27    fn dialect(&self) -> Arc<dyn Dialect>;
28
29    /// Returns an AST analyzer specific for this engine to modify the AST before execution
30    fn ast_analyzer(&self) -> Option<AstAnalyzer> {
31        None
32    }
33
34    /// Execute a SQL query
35    fn execute(&self, query: &str, schema: SchemaRef) -> Result<SendableRecordBatchStream>;
36
37    /// Returns the tables provided by the remote
38    async fn table_names(&self) -> Result<Vec<String>>;
39
40    /// Returns the schema of table_name within this [`SQLExecutor`]
41    async fn get_table_schema(&self, table_name: &str) -> Result<SchemaRef>;
42}
43
44impl fmt::Debug for dyn SQLExecutor {
45    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
46        write!(f, "{} {:?}", self.name(), self.compute_context())
47    }
48}
49
50impl fmt::Display for dyn SQLExecutor {
51    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
52        write!(f, "{} {:?}", self.name(), self.compute_context())
53    }
54}
datafusion_federation/sql/executor.rs

datafusion_federation/sql/
executor.rs
