feat: Added a new trait to expose SchemaProvider

Author: parmesant

src/lib.rs

@@ -58,7 +58,12 @@ pub mod validator;
 use std::time::Duration;
 
 // Public re-exports of crates being used in enterprise
+pub use arrow_array;
+pub use arrow_flight;
+pub use arrow_ipc;
+pub use catalog as parseable_catalog;
 pub use datafusion;
+pub use datafusion_proto;
 pub use handlers::http::modal::{
     ParseableServer, ingest_server::IngestServer, query_server::QueryServer, server::Server,
 };
@@ -68,6 +73,7 @@ use parseable::PARSEABLE;
 use reqwest::{Client, ClientBuilder};
 pub use {opentelemetry, opentelemetry_otlp, opentelemetry_proto, opentelemetry_sdk};
 pub use {tracing_actix_web, tracing_opentelemetry, tracing_subscriber};
+pub use utils as parseable_utils;
 
 // It is very unlikely that panic will occur when dealing with locks.
 pub const LOCK_EXPECT: &str = "Thread shouldn't panic while holding a lock";

src/query/mod.rs

@@ -25,6 +25,7 @@ use arrow_schema::SchemaRef;
 use chrono::NaiveDateTime;
 use chrono::{DateTime, Duration, Utc};
 use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::catalog::SchemaProvider;
 use datafusion::common::tree_node::Transformed;
 use datafusion::execution::disk_manager::DiskManager;
 use datafusion::execution::{
@@ -45,7 +46,7 @@ use datafusion::sql::sqlparser::dialect::PostgreSqlDialect;
 use futures::Stream;
 use futures::stream::select_all;
 use itertools::Itertools;
-use once_cell::sync::Lazy;
+use once_cell::sync::{Lazy, OnceCell};
 use serde::{Deserialize, Serialize};
 use serde_json::{Value, json};
 use std::ops::Bound;
@@ -57,7 +58,6 @@ use sysinfo::System;
 use tokio::runtime::Runtime;
 
 use self::error::ExecuteError;
-use self::stream_schema_provider::GlobalSchemaProvider;
 pub use self::stream_schema_provider::PartialTimeFilter;
 use crate::alerts::alert_structs::Conditions;
 use crate::alerts::alerts_utils::get_filter_string;
@@ -70,7 +70,8 @@ use crate::handlers::http::query::QueryError;
 use crate::metrics::increment_bytes_scanned_in_query_by_date;
 use crate::option::Mode;
 use crate::parseable::{DEFAULT_TENANT, PARSEABLE};
-use crate::storage::{ObjectStorageProvider, ObjectStoreFormat};
+use crate::query::stream_schema_provider::GlobalSchemaProvider;
+use crate::storage::{ObjectStorage, ObjectStorageProvider, ObjectStoreFormat};
 use crate::utils::time::TimeRange;
 
 /// Boxed record-batch stream used as the streaming half of query results.
@@ -95,6 +96,7 @@ type QueryResult = Result<(Either<Vec<RecordBatch>, BoxedBatchStream>, Vec<Strin
 
 // pub static QUERY_SESSION: Lazy<SessionContext> =
 //     Lazy::new(|| Query::create_session_context(PARSEABLE.storage()));
+pub static SCHEMA_PROVIDER: OnceCell<Box<dyn ParseableSchemaProvider>> = OnceCell::new();
 
 pub static QUERY_SESSION_STATE: Lazy<SessionState> =
     Lazy::new(|| Query::create_session_state(PARSEABLE.storage()));
@@ -110,6 +112,15 @@ pub static QUERY_SESSION: Lazy<InMemorySessionContext> = Lazy::new(|| {
     }
 });
 
+/// Trait to enable implementation of SchemaProvider
+pub trait ParseableSchemaProvider: Send + Sync {
+    fn new_provider(
+        &self,
+        storage: Option<Arc<dyn ObjectStorage>>,
+        tenant_id: &Option<String>,
+    ) -> Box<dyn SchemaProvider>;
+}
+
 pub struct InMemorySessionContext {
     session_context: Arc<RwLock<SessionContext>>,
 }
@@ -124,18 +135,23 @@ impl InMemorySessionContext {
     }
 
     pub fn add_schema(&self, tenant_id: &str) {
+        let schema_provider = if let Some(provider) = SCHEMA_PROVIDER.get() {
+            provider.new_provider(
+                Some(PARSEABLE.storage().get_object_store()),
+                &Some(tenant_id.to_owned()),
+            )
+        } else {
+            Box::new(GlobalSchemaProvider {
+                storage: PARSEABLE.storage().get_object_store(),
+                tenant_id: Some(tenant_id.to_owned()),
+            })
+        };
         self.session_context
             .write()
             .expect("SessionContext should be writeable")
             .catalog("datafusion")
             .expect("Default catalog should be available")
-            .register_schema(
-                tenant_id,
-                Arc::new(GlobalSchemaProvider {
-                    storage: PARSEABLE.storage().get_object_store(),
-                    tenant_id: Some(tenant_id.to_owned()),
-                }),
-            )
+            .register_schema(tenant_id, schema_provider.into())
             .expect("Should be able to register new schema");
     }
 
@@ -184,29 +200,41 @@ impl Query {
             // register multiple schemas
             if let Some(tenants) = PARSEABLE.list_tenants() {
                 for t in tenants.iter() {
-                    let schema_provider = Arc::new(GlobalSchemaProvider {
-                        storage: storage.get_object_store(),
-                        tenant_id: Some(t.clone()),
-                    });
-                    let _ = catalog.register_schema(t, schema_provider);
+                    let schema_provider = if let Some(provider) = SCHEMA_PROVIDER.get() {
+                        provider.new_provider(
+                            Some(PARSEABLE.storage().get_object_store()),
+                            &Some(t.to_owned()),
+                        )
+                    } else {
+                        Box::new(GlobalSchemaProvider {
+                            storage: PARSEABLE.storage().get_object_store(),
+                            tenant_id: Some(t.to_owned()),
+                        })
+                    };
+                    let _ = catalog.register_schema(t, schema_provider.into());
                 }
             }
         } else {
             // register just one schema
-            let schema_provider = Arc::new(GlobalSchemaProvider {
-                storage: storage.get_object_store(),
-                tenant_id: None,
-            });
+            let schema_provider = if let Some(provider) = SCHEMA_PROVIDER.get() {
+                provider.new_provider(Some(PARSEABLE.storage().get_object_store()), &None)
+            } else {
+                Box::new(GlobalSchemaProvider {
+                    storage: PARSEABLE.storage().get_object_store(),
+                    tenant_id: None,
+                })
+            };
+
             let _ = catalog.register_schema(
                 &state.config_options().catalog.default_schema,
-                schema_provider,
+                schema_provider.into(),
             );
         }
 
         SessionContext::new_with_state(state)
     }
 
-    fn create_session_state(storage: Arc<dyn ObjectStorageProvider>) -> SessionState {
+    pub fn create_session_state(storage: Arc<dyn ObjectStorageProvider>) -> SessionState {
         let runtime_config = storage
             .get_datafusion_runtime()
             .with_disk_manager_builder(DiskManager::builder());
@@ -288,14 +316,10 @@ impl Query {
             return Ok((Either::Left(vec![]), fields));
         }
 
-        let plan = QUERY_SESSION
-            .get_ctx()
-            .state()
-            .create_physical_plan(df.logical_plan())
-            .await?;
+        let plan = ctx.state().create_physical_plan(df.logical_plan()).await?;
 
         let results = if !is_streaming {
-            let task_ctx = QUERY_SESSION.get_ctx().task_ctx();
+            let task_ctx = ctx.task_ctx();
 
             let batches = collect_partitioned(plan.clone(), task_ctx.clone())
                 .await?
@@ -311,7 +335,7 @@ impl Query {
 
             Either::Left(batches)
         } else {
-            let task_ctx = QUERY_SESSION.get_ctx().task_ctx();
+            let task_ctx = ctx.task_ctx();
 
             let output_partitions = plan.output_partitioning().partition_count();
 

src/query/stream_schema_provider.rs

@@ -702,7 +702,7 @@ pub enum PartialTimeFilter {
 }
 
 impl PartialTimeFilter {
-    fn try_from_expr(expr: &Expr, time_partition: &Option<String>) -> Option<Self> {
+    pub fn try_from_expr(expr: &Expr, time_partition: &Option<String>) -> Option<Self> {
         let Expr::BinaryExpr(binexpr) = expr else {
             return None;
         };