Skip to main content

datafusion_odata/
context.rs

1use std::sync::Arc;
2
3use chrono::{DateTime, Utc};
4use datafusion::{
5    arrow::{datatypes::SchemaRef, record_batch::RecordBatch},
6    dataframe::DataFrame,
7};
8
9use crate::{
10    collection::{CollectionAddr, QueryParams},
11    error::{KeyColumnNotAssigned, ODataError},
12};
13
14///////////////////////////////////////////////////////////////////////////////
15
16pub const DEFAULT_NAMESPACE: &str = "default";
17
18///////////////////////////////////////////////////////////////////////////////
19
20#[async_trait::async_trait]
21pub trait ServiceContext: Send + Sync {
22    fn service_base_url(&self) -> String;
23
24    async fn list_collections(&self) -> Result<Vec<Arc<dyn CollectionContext>>, ODataError>;
25
26    fn on_unsupported_feature(&self) -> OnUnsupported;
27}
28
29///////////////////////////////////////////////////////////////////////////////
30
31#[async_trait::async_trait]
32pub trait CollectionContext: Send + Sync {
33    fn addr(&self) -> Result<&CollectionAddr, ODataError>;
34
35    fn service_base_url(&self) -> Result<String, ODataError>;
36
37    fn collection_base_url(&self) -> Result<String, ODataError>;
38
39    fn collection_namespace(&self) -> Result<String, ODataError> {
40        Ok(DEFAULT_NAMESPACE.to_string())
41    }
42
43    fn collection_name(&self) -> Result<String, ODataError>;
44
45    // Synthetic column name that will be used to propagate entity IDs
46    fn key_column_alias(&self) -> String {
47        "__id__".to_string()
48    }
49
50    fn key_column(&self) -> Result<String, ODataError> {
51        Err(KeyColumnNotAssigned)?
52    }
53
54    async fn last_updated_time(&self) -> DateTime<Utc>;
55
56    async fn schema(&self) -> Result<SchemaRef, ODataError>;
57
58    async fn query(&self, query: QueryParams) -> Result<DataFrame, ODataError>;
59
60    fn on_unsupported_feature(&self) -> OnUnsupported;
61
62    /// Validates the record batches that retunred from datafusion before encode them to xml
63    async fn validate(&self, _record_batches: &[RecordBatch]) -> Result<(), ODataError> {
64        Ok(())
65    }
66}
67
68///////////////////////////////////////////////////////////////////////////////
69
70pub enum OnUnsupported {
71    /// Return an error or crash
72    Error,
73    /// Log error and recover as gracefully as possible
74    Warn,
75}