Merge pull request #8 from twintag/twintag-odata5

add validate function to CollectionContext
kamu-data · Sep 23, 2024 · 33ff333 · 33ff333
2 parents 519e01f + 3f1747c
commit 33ff333
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/src/context.rs b/src/context.rs
@@ -1,7 +1,10 @@
 use std::sync::Arc;
 
 use chrono::{DateTime, Utc};
-use datafusion::{arrow::datatypes::SchemaRef, dataframe::DataFrame};
+use datafusion::{
+    arrow::{datatypes::SchemaRef, record_batch::RecordBatch},
+    dataframe::DataFrame,
+};
 
 use crate::{
     collection::{CollectionAddr, QueryParams},
@@ -55,6 +58,11 @@ pub trait CollectionContext: Send + Sync {
     async fn query(&self, query: QueryParams) -> Result<DataFrame, ODataError>;
 
     fn on_unsupported_feature(&self) -> OnUnsupported;
+
+    /// Validates the record batches that retunred from datafusion before encode them to xml
+    async fn validate(&self, _record_batches: &[RecordBatch]) -> Result<(), ODataError> {
+        Ok(())
+    }
 }
 
 ///////////////////////////////////////////////////////////////////////////////

diff --git a/src/handlers.rs b/src/handlers.rs
@@ -150,6 +150,8 @@ pub async fn odata_collection_handler(
     let schema: datafusion::arrow::datatypes::Schema = df.schema().clone().into();
     let record_batches = df.collect().await.map_err(ODataError::internal)?;
 
+    ctx.validate(&record_batches).await?;
+
     let num_rows: usize = record_batches.iter().map(|b| b.num_rows()).sum();
     let raw_bytes: usize = record_batches
         .iter()