svilen-mihaylov-db commented on code in PR #12067:
URL: https://github.com/apache/datafusion/pull/12067#discussion_r1722222310
##########
datafusion/core/tests/custom_sources_cases/statistics.rs:
##########
@@ -299,3 +300,240 @@ async fn sql_window() -> Result<()> {
Ok(())
}
+
+struct WrappingTableProvider {
+ underlying: Arc<dyn TableProvider>,
+}
+
+impl WrappingTableProvider {
+ fn new(underlying: Arc<dyn TableProvider>) -> Self {
+ println!("Creating new custom table provider.");
+ Self {
+ underlying: underlying,
+ }
+ }
+
+ fn deregister(&self) {
+ println!("Deregistering custom table provider.");
+ }
+}
+
+#[async_trait]
+impl TableProvider for WrappingTableProvider {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn schema(&self) -> SchemaRef {
+ self.underlying.schema()
+ }
+
+ fn table_type(&self) -> TableType {
+ self.underlying.table_type()
+ }
+
+ async fn scan(
+ &self,
+ state: &dyn Session,
+ projection: Option<&Vec<usize>>,
+ filters: &[Expr],
+ limit: Option<usize>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.underlying
+ .scan(state, projection, filters, limit)
+ .await
+ }
+}
+
+struct TestSchemaProvider {
+ underlying: MemorySchemaProvider,
+}
+
+impl TestSchemaProvider {
+ pub fn new() -> Self {
+ Self {
+ underlying: MemorySchemaProvider::new(),
+ }
+ }
+}
+
+impl Default for TestSchemaProvider {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+#[async_trait]
+impl SchemaProvider for TestSchemaProvider {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn table_names(&self) -> Vec<String> {
+ self.underlying.table_names()
+ }
+
+ async fn table(
+ &self,
+ name: &str,
+ ) -> Result<Option<Arc<dyn TableProvider>>, DataFusionError> {
+ self.underlying.table(name).await
+ }
+
+ fn register_table(
+ &self,
+ name: String,
+ table: Arc<dyn TableProvider>,
+ ) -> Result<Option<Arc<dyn TableProvider>>> {
+ println!("Registering table {}", name);
+
+ let schema = table.schema();
+ for field in &schema.fields {
+ println!("Field {} options", field.name());
+ for (k, v) in field.metadata() {
+ println!("{}={}", k, v)
+ }
+ }
+
+ println!("Metadata");
+ for (k, v) in &schema.metadata {
+ println!("{}={}", k, v)
+ }
+
+ // Here we can register our own TableProvider instance instead of a
memory table 'table'.
+ self.underlying
+ .register_table(name.clone(),
Arc::new(WrappingTableProvider::new(table)))
+ }
+
+ fn deregister_table(&self, name: &str) -> Result<Option<Arc<dyn
TableProvider>>> {
+ println!("Deregistering table {}", name);
+ let result = self.underlying.deregister_table(name);
+ if let Ok(result) = result {
+ if let Some(result) = result {
+ let provider: &WrappingTableProvider =
+ match
result.as_any().downcast_ref::<WrappingTableProvider>() {
+ Some(p) => p,
+ None => panic!("Unexpected table provider"),
+ };
+ provider.deregister();
+ return Ok(Some(result));
+ }
+ return Ok(result);
+ }
+ result
+ }
+
+ fn table_exist(&self, name: &str) -> bool {
+ self.underlying.table_exist(name)
+ }
+}
+
+struct TestCatalogProvider {
+ provider: Arc<TestSchemaProvider>,
+}
+
+impl TestCatalogProvider {
+ pub fn new() -> Self {
+ Self {
+ provider: Arc::new(TestSchemaProvider::new()),
+ }
+ }
+}
+
+impl Default for TestCatalogProvider {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl CatalogProvider for TestCatalogProvider {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn schema_names(&self) -> Vec<String> {
+ vec!["public".to_owned()]
+ }
+
+ fn schema(&self, name: &str) -> Option<Arc<dyn SchemaProvider>> {
+ assert_eq!("public", name);
+ Some(self.provider.clone())
+ }
+}
+
+struct TestCatalogProviderList {
+ underlying: MemoryCatalogProviderList,
+}
+
+impl TestCatalogProviderList {
+ pub fn new() -> Self {
+ Self {
+ underlying: MemoryCatalogProviderList::new(),
+ }
+ }
+}
+
+impl Default for TestCatalogProviderList {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl CatalogProviderList for TestCatalogProviderList {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn register_catalog(
+ &self,
+ name: String,
+ catalog: Arc<dyn CatalogProvider>,
+ ) -> Option<Arc<dyn CatalogProvider>> {
+ self.underlying.register_catalog(name, catalog)
+ }
+
+ fn catalog_names(&self) -> Vec<String> {
+ self.underlying.catalog_names()
+ }
+
+ fn catalog(&self, name: &str) -> Option<Arc<dyn CatalogProvider>> {
+ self.underlying.catalog(name)
+ }
+}
+
+#[tokio::test]
+async fn sql_basic1() -> Result<()> {
Review Comment:
Test output is
```
Registering table t
Field a options
dimension=1
lower_bound='0'
Field b options
Metadata
v1=4
v=2
engine_option='etc etc'
Creating new custom table provider.
Deregistering table t
Deregistering custom table provider.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
