datafusion_remote_table/
table.rs

1use crate::connection::RemoteDbType;
2use crate::{
3    ConnectionOptions, DFResult, Pool, RemoteSchemaRef, RemoteTableExec, Transform, connect,
4    transform_schema,
5};
6use datafusion::arrow::datatypes::SchemaRef;
7use datafusion::catalog::{Session, TableProvider};
8use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
9use datafusion::datasource::TableType;
10use datafusion::error::DataFusionError;
11use datafusion::logical_expr::{Expr, TableProviderFilterPushDown};
12use datafusion::physical_plan::ExecutionPlan;
13use datafusion::sql::unparser::Unparser;
14use datafusion::sql::unparser::dialect::{MySqlDialect, PostgreSqlDialect, SqliteDialect};
15use std::any::Any;
16use std::sync::Arc;
17
18#[derive(Debug)]
19pub struct RemoteTable {
20    pub(crate) conn_options: ConnectionOptions,
21    pub(crate) sql: String,
22    pub(crate) table_schema: SchemaRef,
23    pub(crate) transformed_table_schema: SchemaRef,
24    pub(crate) remote_schema: Option<RemoteSchemaRef>,
25    pub(crate) transform: Option<Arc<dyn Transform>>,
26    pub(crate) pool: Arc<dyn Pool>,
27}
28
29impl RemoteTable {
30    pub async fn try_new(
31        conn_options: ConnectionOptions,
32        sql: impl Into<String>,
33    ) -> DFResult<Self> {
34        Self::try_new_with_schema_transform(conn_options, sql, None, None).await
35    }
36
37    pub async fn try_new_with_schema(
38        conn_options: ConnectionOptions,
39        sql: impl Into<String>,
40        table_schema: SchemaRef,
41    ) -> DFResult<Self> {
42        Self::try_new_with_schema_transform(conn_options, sql, Some(table_schema), None).await
43    }
44
45    pub async fn try_new_with_transform(
46        conn_options: ConnectionOptions,
47        sql: impl Into<String>,
48        transform: Arc<dyn Transform>,
49    ) -> DFResult<Self> {
50        Self::try_new_with_schema_transform(conn_options, sql, None, Some(transform)).await
51    }
52
53    pub async fn try_new_with_schema_transform(
54        conn_options: ConnectionOptions,
55        sql: impl Into<String>,
56        table_schema: Option<SchemaRef>,
57        transform: Option<Arc<dyn Transform>>,
58    ) -> DFResult<Self> {
59        let sql = sql.into();
60        let pool = connect(&conn_options).await?;
61        let conn = pool.get().await?;
62        let (table_schema, remote_schema) = match conn.infer_schema(&sql).await {
63            Ok((remote_schema, inferred_table_schema)) => (
64                table_schema.unwrap_or(inferred_table_schema),
65                Some(remote_schema),
66            ),
67            Err(e) => {
68                if let Some(table_schema) = table_schema {
69                    (table_schema, None)
70                } else {
71                    return Err(DataFusionError::Execution(format!(
72                        "Failed to infer schema: {e}"
73                    )));
74                }
75            }
76        };
77        let transformed_table_schema = transform_schema(
78            table_schema.clone(),
79            transform.as_ref(),
80            remote_schema.as_ref(),
81        )?;
82        Ok(RemoteTable {
83            conn_options,
84            sql,
85            table_schema,
86            transformed_table_schema,
87            remote_schema,
88            transform,
89            pool,
90        })
91    }
92
93    pub fn remote_schema(&self) -> Option<RemoteSchemaRef> {
94        self.remote_schema.clone()
95    }
96}
97
98#[async_trait::async_trait]
99impl TableProvider for RemoteTable {
100    fn as_any(&self) -> &dyn Any {
101        self
102    }
103
104    fn schema(&self) -> SchemaRef {
105        self.transformed_table_schema.clone()
106    }
107
108    fn table_type(&self) -> TableType {
109        TableType::View
110    }
111
112    async fn scan(
113        &self,
114        _state: &dyn Session,
115        projection: Option<&Vec<usize>>,
116        filters: &[Expr],
117        limit: Option<usize>,
118    ) -> DFResult<Arc<dyn ExecutionPlan>> {
119        Ok(Arc::new(RemoteTableExec::try_new(
120            self.conn_options.clone(),
121            self.sql.clone(),
122            self.table_schema.clone(),
123            self.remote_schema.clone(),
124            projection.cloned(),
125            filters.to_vec(),
126            limit,
127            self.transform.clone(),
128            self.pool.get().await?,
129        )?))
130    }
131
132    fn supports_filters_pushdown(
133        &self,
134        filters: &[&Expr],
135    ) -> DFResult<Vec<TableProviderFilterPushDown>> {
136        Ok(filters
137            .iter()
138            .map(|f| support_filter_pushdown(self.conn_options.db_type(), &self.sql, f))
139            .collect())
140    }
141}
142
143pub(crate) fn support_filter_pushdown(
144    db_type: RemoteDbType,
145    sql: &str,
146    filter: &Expr,
147) -> TableProviderFilterPushDown {
148    if !db_type.support_rewrite_with_filters_limit(sql) {
149        return TableProviderFilterPushDown::Unsupported;
150    }
151    let unparser = match db_type {
152        RemoteDbType::Mysql => Unparser::new(&MySqlDialect {}),
153        RemoteDbType::Postgres => Unparser::new(&PostgreSqlDialect {}),
154        RemoteDbType::Sqlite => Unparser::new(&SqliteDialect {}),
155        RemoteDbType::Oracle => return TableProviderFilterPushDown::Unsupported,
156    };
157    if unparser.expr_to_sql(filter).is_err() {
158        return TableProviderFilterPushDown::Unsupported;
159    }
160
161    let mut pushdown = TableProviderFilterPushDown::Exact;
162    filter
163        .apply(|e| {
164            if matches!(e, Expr::ScalarFunction(_)) {
165                pushdown = TableProviderFilterPushDown::Unsupported;
166            }
167            Ok(TreeNodeRecursion::Continue)
168        })
169        .expect("won't fail");
170
171    pushdown
172}
datafusion_remote_table/table.rs

datafusion_remote_table/
table.rs
