refactor: table options

crates/catalog/src/session_catalog.rs

@@ -18,11 +18,7 @@ use protogen::metastore::types::catalog::{
     TableEntry,
     TunnelEntry,
 };
-use protogen::metastore::types::options::{
-    InternalColumnDefinition,
-    TableOptions,
-    TableOptionsInternal,
-};
+use protogen::metastore::types::options::{InternalColumnDefinition, TableOptionsInternal};
 use tracing::debug;
 
 use super::client::MetastoreClientHandle;
@@ -461,12 +457,10 @@ impl TempCatalog {
                     external: false,
                     is_temp: true,
                 },
-                options: TableOptions::Internal(TableOptionsInternal {
-                    columns: columns.to_owned(),
-                }),
+                options: TableOptionsInternal { columns }.into(),
                 tunnel_id: None,
                 access_mode: SourceAccessMode::ReadWrite,
-                columns: Some(columns.to_owned()),
+                schema: Some(schema.as_ref().clone()),
             }
         })
     }
@@ -505,12 +499,13 @@ impl TempCatalog {
                     external: false,
                     is_temp: true,
                 },
-                options: TableOptions::Internal(TableOptionsInternal {
+                options: TableOptionsInternal {
                     columns: Vec::new(),
-                }),
+                }
+                .into(),
                 tunnel_id: None,
                 access_mode: SourceAccessMode::ReadWrite,
-                columns: Some(Vec::new()),
+                schema: None,
             });
         }
 

crates/datasources/src/debug/errors.rs

@@ -8,4 +8,5 @@ pub enum DebugError {
     InvalidTunnel(String),
 }
 
+
 pub type Result<T, E = DebugError> = std::result::Result<T, E>;

crates/datasources/src/debug/mod.rs

@@ -1,23 +1,15 @@
 //! A collection of debug datasources.
 pub mod errors;
+pub mod options;
 
 use std::any::Any;
 use std::fmt;
 use std::pin::Pin;
-use std::str::FromStr;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
 use async_trait::async_trait;
-use datafusion::arrow::array::Int32Array;
-use datafusion::arrow::datatypes::{
-    DataType,
-    Field,
-    Fields,
-    Schema as ArrowSchema,
-    SchemaRef as ArrowSchemaRef,
-};
-use datafusion::arrow::error::Result as ArrowResult;
+use datafusion::arrow::datatypes::{DataType, Field, Fields, SchemaRef as ArrowSchemaRef};
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::datasource::TableProvider;
 use datafusion::error::{DataFusionError, Result as DatafusionResult};
@@ -37,52 +29,17 @@ use datafusion_ext::errors::ExtensionError;
 use datafusion_ext::functions::VirtualLister;
 use errors::DebugError;
 use futures::Stream;
-use parser::errors::ParserError;
-use parser::options::{OptionValue, ParseOptionValue};
-use protogen::metastore::types::options::TunnelOptions;
-use serde::{Deserialize, Serialize};
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub enum DebugTableType {
-    /// A table that will always return an error on the record batch stream.
-    ErrorDuringExecution,
-    /// A table that never stops sending record batches.
-    NeverEnding,
-}
-
-impl ParseOptionValue<DebugTableType> for OptionValue {
-    fn parse_opt(self) -> Result<DebugTableType, ParserError> {
-        let opt = match self {
-            Self::QuotedLiteral(s) | Self::UnquotedLiteral(s) => s
-                .parse()
-                .map_err(|e: DebugError| ParserError::ParserError(e.to_string()))?,
-            o => {
-                return Err(ParserError::ParserError(format!(
-                    "Expected a string, got: {}",
-                    o
-                )))
-            }
-        };
-        Ok(opt)
-    }
-}
-impl fmt::Display for DebugTableType {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}", self.as_str())
-    }
-}
+use parser::options::StatementOptions;
+use protogen::metastore::types::options::{
+    CredentialsOptions,
+    DatabaseOptions,
+    TableOptions,
+    TunnelOptions,
+};
 
-impl FromStr for DebugTableType {
-    type Err = DebugError;
+pub use self::options::{DebugTableType, TableOptionsDebug};
+use crate::DatasourceError;
 
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        Ok(match s {
-            "error_during_execution" => DebugTableType::ErrorDuringExecution,
-            "never_ending" => DebugTableType::NeverEnding,
-            other => return Err(DebugError::UnknownDebugTableType(other.to_string())),
-        })
-    }
-}
 
 /// Validates if the tunnel is supported and returns whether a tunnel is going
 /// to be used or not for  the connection.
@@ -96,82 +53,6 @@ pub fn validate_tunnel_connections(
     }
 }
 
-impl DebugTableType {
-    /// Get the arrow schema for the debug table type.
-    pub fn arrow_schema(&self) -> ArrowSchema {
-        match self {
-            DebugTableType::ErrorDuringExecution => {
-                ArrowSchema::new(vec![Field::new("a", DataType::Int32, false)])
-            }
-            DebugTableType::NeverEnding => ArrowSchema::new(vec![
-                Field::new("a", DataType::Int32, false),
-                Field::new("b", DataType::Int32, false),
-                Field::new("c", DataType::Int32, false),
-            ]),
-        }
-    }
-
-    /// Get the projected arrow schema.
-    pub fn projected_arrow_schema(
-        &self,
-        projection: Option<&Vec<usize>>,
-    ) -> ArrowResult<ArrowSchema> {
-        match projection {
-            Some(proj) => self.arrow_schema().project(proj),
-            None => Ok(self.arrow_schema()),
-        }
-    }
-
-    /// Produces a record batch that matches this debug table's schema.
-    pub fn record_batch(&self, tunnel: bool) -> RecordBatch {
-        let base = if tunnel { 10_i32 } else { 1_i32 };
-        match self {
-            DebugTableType::ErrorDuringExecution => RecordBatch::try_new(
-                Arc::new(self.arrow_schema()),
-                vec![Arc::new(Int32Array::from_value(base, 30))],
-            )
-            .unwrap(),
-            DebugTableType::NeverEnding => RecordBatch::try_new(
-                Arc::new(self.arrow_schema()),
-                vec![
-                    Arc::new(Int32Array::from_value(base, 30)),
-                    Arc::new(Int32Array::from_value(base * 2, 30)),
-                    Arc::new(Int32Array::from_value(base * 3, 30)),
-                ],
-            )
-            .unwrap(),
-        }
-    }
-
-    /// Get a projected record batch for this debug table type.
-    pub fn projected_record_batch(
-        &self,
-        tunnel: bool,
-        projection: Option<&Vec<usize>>,
-    ) -> ArrowResult<RecordBatch> {
-        match projection {
-            Some(proj) => self.record_batch(tunnel).project(proj),
-            None => Ok(self.record_batch(tunnel)),
-        }
-    }
-
-    pub fn as_str(&self) -> &'static str {
-        match self {
-            DebugTableType::ErrorDuringExecution => "error_during_execution",
-            DebugTableType::NeverEnding => "never_ending",
-        }
-    }
-
-    pub fn into_table_provider(
-        self,
-        tunnel_opts: Option<&TunnelOptions>,
-    ) -> Arc<dyn TableProvider> {
-        let tunnel = validate_tunnel_connections(tunnel_opts)
-            .expect("datasources should be validated with tunnels before dispatch");
-        Arc::new(DebugTableProvider { typ: self, tunnel })
-    }
-}
-
 pub struct DebugVirtualLister;
 
 #[async_trait]
@@ -202,8 +83,8 @@ impl VirtualLister for DebugVirtualLister {
 }
 
 pub struct DebugTableProvider {
-    typ: DebugTableType,
-    tunnel: bool,
+    pub typ: DebugTableType,
+    pub tunnel: bool,
 }
 
 #[async_trait]
@@ -369,3 +250,78 @@ impl RecordBatchStream for NeverEndingStream {
         self.batch.schema()
     }
 }
+pub struct DebugDatasource;
+
+
+#[async_trait]
+impl crate::Datasource for DebugDatasource {
+    fn name(&self) -> &'static str {
+        "debug"
+    }
+
+    fn validate_tunnel_connections(
+        &self,
+        tunnel_opts: Option<&TunnelOptions>,
+    ) -> Result<(), DatasourceError> {
+        match tunnel_opts {
+            None => Ok(()),
+            Some(TunnelOptions::Debug(_)) => Ok(()),
+            Some(other) => Err(DebugError::InvalidTunnel(other.to_string()).into()),
+        }
+    }
+
+    fn validate_credentials(
+        &self,
+        _creds: Option<CredentialsOptions>,
+    ) -> Result<(), DatasourceError> {
+        Ok(())
+    }
+
+    fn table_options_from_stmt(
+        &self,
+        opts: &mut StatementOptions,
+        creds: Option<CredentialsOptions>,
+        tunnel_opts: Option<TunnelOptions>,
+    ) -> Result<TableOptions, DatasourceError> {
+        validate_tunnel_connections(tunnel_opts.as_ref()).unwrap();
+
+        let typ: Option<DebugTableType> = match creds {
+            Some(CredentialsOptions::Debug(c)) => c.table_type.parse().ok(),
+            Some(other) => unreachable!("invalid credentials {other} for debug datasource"),
+            None => None,
+        };
+        let typ: DebugTableType = opts
+            .remove_required_or("table_type", typ)
+            .map_err(|e| DebugError::UnknownDebugTableType(e.to_string()))?;
+
+        Ok(TableOptionsDebug { table_type: typ }.into())
+    }
+
+    async fn create_table_provider(
+        &self,
+        options: &TableOptions,
+        tunnel_opts: Option<&TunnelOptions>,
+    ) -> Result<Arc<dyn TableProvider>, DatasourceError> {
+        let options = TableOptionsDebug::try_from(options).unwrap();
+
+        Ok(Arc::new(DebugTableProvider {
+            typ: options.table_type,
+            tunnel: tunnel_opts.is_some(),
+        }))
+    }
+
+    async fn table_provider_from_db_options(
+        &self,
+        _schema: &str,
+        name: &str,
+        _options: &DatabaseOptions,
+        tunnel_opts: Option<&TunnelOptions>,
+    ) -> Result<Option<Arc<dyn TableProvider>>, DatasourceError> {
+        let tbl_type = name.parse()?;
+        self.validate_tunnel_connections(tunnel_opts)?;
+        Ok(Some(Arc::new(DebugTableProvider {
+            typ: tbl_type,
+            tunnel: tunnel_opts.is_some(),
+        })))
+    }
+}