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
62        let (table_schema, remote_schema) = if let Some(table_schema) = table_schema {
63            let remote_schema = if transform.is_some() {
64                // Infer remote schema
65                let conn = pool.get().await?;
66                match conn.infer_schema(&sql).await {
67                    Ok((remote_schema, _)) => Some(remote_schema),
68                    Err(_) => None,
69                }
70            } else {
71                None
72            };
73            (table_schema, remote_schema)
74        } else {
75            // Infer table schema
76            let conn = pool.get().await?;
77            match conn.infer_schema(&sql).await {
78                Ok((remote_schema, inferred_table_schema)) => {
79                    (inferred_table_schema, Some(remote_schema))
80                }
81                Err(e) => {
82                    return Err(DataFusionError::Execution(format!(
83                        "Failed to infer schema: {e}"
84                    )));
85                }
86            }
87        };
88
89        let transformed_table_schema = transform_schema(
90            table_schema.clone(),
91            transform.as_ref(),
92            remote_schema.as_ref(),
93        )?;
94
95        Ok(RemoteTable {
96            conn_options,
97            sql,
98            table_schema,
99            transformed_table_schema,
100            remote_schema,
101            transform,
102            pool,
103        })
104    }
105
106    pub fn remote_schema(&self) -> Option<RemoteSchemaRef> {
107        self.remote_schema.clone()
108    }
109}
110
111#[async_trait::async_trait]
112impl TableProvider for RemoteTable {
113    fn as_any(&self) -> &dyn Any {
114        self
115    }
116
117    fn schema(&self) -> SchemaRef {
118        self.transformed_table_schema.clone()
119    }
120
121    fn table_type(&self) -> TableType {
122        TableType::View
123    }
124
125    async fn scan(
126        &self,
127        _state: &dyn Session,
128        projection: Option<&Vec<usize>>,
129        filters: &[Expr],
130        limit: Option<usize>,
131    ) -> DFResult<Arc<dyn ExecutionPlan>> {
132        Ok(Arc::new(RemoteTableExec::try_new(
133            self.conn_options.clone(),
134            self.sql.clone(),
135            self.table_schema.clone(),
136            self.remote_schema.clone(),
137            projection.cloned(),
138            filters.to_vec(),
139            limit,
140            self.transform.clone(),
141            self.pool.get().await?,
142        )?))
143    }
144
145    fn supports_filters_pushdown(
146        &self,
147        filters: &[&Expr],
148    ) -> DFResult<Vec<TableProviderFilterPushDown>> {
149        Ok(filters
150            .iter()
151            .map(|f| support_filter_pushdown(self.conn_options.db_type(), &self.sql, f))
152            .collect())
153    }
154}
155
156pub(crate) fn support_filter_pushdown(
157    db_type: RemoteDbType,
158    sql: &str,
159    filter: &Expr,
160) -> TableProviderFilterPushDown {
161    if !db_type.support_rewrite_with_filters_limit(sql) {
162        return TableProviderFilterPushDown::Unsupported;
163    }
164    let unparser = match db_type {
165        RemoteDbType::Mysql => Unparser::new(&MySqlDialect {}),
166        RemoteDbType::Postgres => Unparser::new(&PostgreSqlDialect {}),
167        RemoteDbType::Sqlite => Unparser::new(&SqliteDialect {}),
168        RemoteDbType::Oracle => return TableProviderFilterPushDown::Unsupported,
169    };
170    if unparser.expr_to_sql(filter).is_err() {
171        return TableProviderFilterPushDown::Unsupported;
172    }
173
174    let mut pushdown = TableProviderFilterPushDown::Exact;
175    filter
176        .apply(|e| {
177            if matches!(e, Expr::ScalarFunction(_)) {
178                pushdown = TableProviderFilterPushDown::Unsupported;
179            }
180            Ok(TreeNodeRecursion::Continue)
181        })
182        .expect("won't fail");
183
184    pushdown
185}
datafusion_remote_table/table.rs

datafusion_remote_table/
table.rs
