Add Observability (#300)

Author: matthewmturner

.github/workflows/test.yml

@@ -221,6 +221,31 @@ jobs:
       - name: Run UDFs WASM tests
         run: |
           cargo test --manifest-path crates/datafusion-functions-parquet/Cargo.toml
+  test-crate-datafusion-app:
+    name: Crate / DataFusion-App
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        arch: [amd64]
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          submodules: true
+      - name: Cache Cargo
+        uses: actions/cache@v4
+        with:
+          path: /home/runner/.cargo
+          key: cargo-dft-cache-
+      - name: Cache Rust dependencies
+        uses: actions/cache@v4
+        with:
+          path: target
+          key: ${{ runner.os }}-cargo-target-${{ hashFiles('Cargo.lock') }}
+      - name: Setup Rust Toolchain
+        uses: ./.github/actions/setup-rust
+      - name: Run UDFs WASM tests
+        run: |
+          cargo test --all-features --manifest-path crates/datafusion-app/Cargo.toml
   test-crate-udfs-wasm:
     name: Crate / UDFs-WASM
     runs-on: ubuntu-latest

Cargo.lock

@@ -29,6 +29,7 @@ futures = "0.3.30"
 http = "1"
 http-body = "1"
 itertools = "0.13.0"
+jiff = { version = "0.2.5", optional = true }
 lazy_static = "1.4.0"
 log = "0.4.22"
 metrics = { version = "0.24.0", optional = true }
@@ -59,7 +60,7 @@ tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
 tui-logger = { version = "0.12", features = ["tracing-support"] }
 tui-textarea = { version = "0.6.1", features = ["search"] }
 url = { version = "2.5.2", optional = true }
-uuid = { version = "1.10.0" }
+uuid = { version = "1.10.0", optional = true }
 
 [dev-dependencies]
 assert_cmd = "2.0.16"
@@ -85,15 +86,18 @@ flightsql = [
   "dep:metrics-exporter-prometheus",
   "dep:tonic",
   "dep:tower-http",
+  "dep:uuid",
 ]
 functions-json = ["datafusion-app/functions-json"]
 functions-parquet = ["datafusion-app/functions-parquet"]
 http = [
   "axum",
   "datafusion-app/observability",
+  "dep:jiff",
   "dep:metrics",
   "dep:metrics-exporter-prometheus",
   "dep:tower-http",
+  "dep:uuid",
 ]
 hudi = ["datafusion-app/hudi"]
 huggingface = ["datafusion-app/huggingface"]

Cargo.toml

@@ -91,9 +91,9 @@ pub struct ExecutionConfig {
     pub wasm_udf: WasmUdfConfig,
     #[serde(default = "default_catalog")]
     pub catalog: CatalogConfig,
-    // #[cfg(feature = "observability")]
-    // #[serde(default)]
-    // pub observability: ObservabilityConfig,
+    #[cfg(feature = "observability")]
+    #[serde(default)]
+    pub observability: ObservabilityConfig,
 }
 
 impl Default for ExecutionConfig {
@@ -109,8 +109,8 @@ impl Default for ExecutionConfig {
             #[cfg(feature = "udfs-wasm")]
             wasm_udf: default_wasm_udf(),
             catalog: default_catalog(),
-            // #[cfg(feature = "observability")]
-            // observability: default_observability(),
+            #[cfg(feature = "observability")]
+            observability: default_observability(),
         }
     }
 }
@@ -302,20 +302,25 @@ fn default_catalog_name() -> String {
 #[cfg(feature = "observability")]
 #[derive(Clone, Debug, Deserialize)]
 pub struct ObservabilityConfig {
-    #[serde(default = "default_observability_catalog_name")]
-    pub catalog_name: String,
+    #[serde(default = "default_observability_schema_name")]
+    pub schema_name: String,
 }
 
 #[cfg(feature = "observability")]
 impl Default for ObservabilityConfig {
     fn default() -> Self {
         Self {
-            catalog_name: default_observability_catalog_name(),
+            schema_name: default_observability_schema_name(),
         }
     }
 }
 
 #[cfg(feature = "observability")]
-fn default_observability_catalog_name() -> String {
+fn default_observability() -> ObservabilityConfig {
+    ObservabilityConfig::default()
+}
+
+#[cfg(feature = "observability")]
+fn default_observability_schema_name() -> String {
     "observability".to_string()
 }

crates/datafusion-app/src/config.rs

@@ -97,18 +97,8 @@ impl ExecutionContext {
             executor = Some(dedicated_executor)
         }
 
-        #[cfg(any(
-            feature = "udfs-wasm",
-            feature = "observability",
-            feature = "functions-json"
-        ))]
+        #[allow(unused_mut)]
         let mut session_ctx = SessionContext::new_with_state(session_state);
-        #[cfg(all(
-            not(feature = "udfs-wasm"),
-            not(feature = "observability"),
-            not(feature = "functions-json")
-        ))]
-        let session_ctx = SessionContext::new_with_state(session_state);
 
         #[cfg(feature = "functions-json")]
         datafusion_functions_json::register_all(&mut session_ctx)?;
@@ -132,26 +122,42 @@ impl ExecutionContext {
         let catalog = create_app_catalog(config, app_name, app_version)?;
         session_ctx.register_catalog(&config.catalog.name, catalog);
 
-        // #[cfg(feature = "observability")]
-        // {
-        //     let obs = ObservabilityContext::new(config.observability.clone());
-        // }
-
-        #[cfg(feature = "observability")]
-        let ctx = Self {
-            config: config.clone(),
-            session_ctx,
-            ddl_path: config.ddl_path.as_ref().map(PathBuf::from),
-            executor,
-            observability: ObservabilityContext::default(),
-        };
-
-        #[cfg(not(feature = "observability"))]
-        let ctx = Self {
-            config: config.clone(),
-            session_ctx,
-            ddl_path: config.ddl_path.as_ref().map(PathBuf::from),
-            executor,
+        let ctx = {
+            #[cfg(feature = "observability")]
+            {
+                let observability =
+                    ObservabilityContext::try_new(config.observability.clone(), app_name)?;
+                if let Some(cat) = session_ctx.catalog(&config.catalog.name) {
+                    match cat
+                        .register_schema(&config.observability.schema_name, observability.schema())
+                    {
+                        Ok(_) => {
+                            info!("Registered observability schema")
+                        }
+                        Err(e) => {
+                            error!("Error registering observability schema: {}", e.to_string())
+                        }
+                    }
+                } else {
+                    error!("Missing catalog to register observability schema")
+                }
+                Self {
+                    config: config.clone(),
+                    session_ctx,
+                    ddl_path: config.ddl_path.as_ref().map(PathBuf::from),
+                    executor,
+                    observability,
+                }
+            }
+            #[cfg(not(feature = "observability"))]
+            {
+                Self {
+                    config: config.clone(),
+                    session_ctx,
+                    ddl_path: config.ddl_path.as_ref().map(PathBuf::from),
+                    executor,
+                }
+            }
         };
 
         Ok(ctx)