From e25f5e7485ffcd810f96c7be096b04b3cacf30b3 Mon Sep 17 00:00:00 2001
From: Leonardo Yvens <leoyvens@gmail.com>
Date: Fri, 15 Nov 2024 01:39:46 +0000
Subject: [PATCH] impl table_type for DefaultTableSource (#13416)

---
 .../src/datasource/default_table_source.rs    | 45 ++++++++++++++++++-
 1 file changed, 44 insertions(+), 1 deletion(-)

diff --git a/datafusion/core/src/datasource/default_table_source.rs b/datafusion/core/src/datasource/default_table_source.rs
index c37c3b97f4fe..5efabd000d68 100644
--- a/datafusion/core/src/datasource/default_table_source.rs
+++ b/datafusion/core/src/datasource/default_table_source.rs
@@ -24,7 +24,7 @@ use crate::datasource::TableProvider;
 
 use arrow::datatypes::SchemaRef;
 use datafusion_common::{internal_err, Constraints};
-use datafusion_expr::{Expr, TableProviderFilterPushDown, TableSource};
+use datafusion_expr::{Expr, TableProviderFilterPushDown, TableSource, TableType};
 
 /// DataFusion default table source, wrapping TableProvider.
 ///
@@ -61,6 +61,11 @@ impl TableSource for DefaultTableSource {
         self.table_provider.constraints()
     }
 
+    /// Get the type of this table for metadata/catalog purposes.
+    fn table_type(&self) -> TableType {
+        self.table_provider.table_type()
+    }
+
     /// Tests whether the table provider can make use of any or all filter expressions
     /// to optimise data retrieval.
     fn supports_filters_pushdown(
@@ -100,3 +105,41 @@ pub fn source_as_provider(
         _ => internal_err!("TableSource was not DefaultTableSource"),
     }
 }
+
+#[test]
+fn preserves_table_type() {
+    use async_trait::async_trait;
+    use datafusion_common::DataFusionError;
+
+    #[derive(Debug)]
+    struct TestTempTable;
+
+    #[async_trait]
+    impl TableProvider for TestTempTable {
+        fn as_any(&self) -> &dyn Any {
+            self
+        }
+
+        fn table_type(&self) -> TableType {
+            TableType::Temporary
+        }
+
+        fn schema(&self) -> SchemaRef {
+            unimplemented!()
+        }
+
+        async fn scan(
+            &self,
+            _: &dyn datafusion_catalog::Session,
+            _: Option<&Vec<usize>>,
+            _: &[Expr],
+            _: Option<usize>,
+        ) -> Result<Arc<dyn datafusion_physical_plan::ExecutionPlan>, DataFusionError>
+        {
+            unimplemented!()
+        }
+    }
+
+    let table_source = DefaultTableSource::new(Arc::new(TestTempTable));
+    assert_eq!(table_source.table_type(), TableType::Temporary);
+}