Merge pull request #502 from splitgraph/http-health

Add a http health endpoint

Author: gruuya

.github/workflows/ci.yml

@@ -11,17 +11,6 @@ jobs:
       # set debuginfo to 1 (line tables only, for panic tracebacks)
       RUSTFLAGS: "-C debuginfo=1"
 
-    services:
-      postgres:
-        image: postgres:14
-        env:
-          POSTGRES_PASSWORD: postgres
-          POSTGRES_DB: db_test
-        ports:
-          - 5432/tcp
-        options: >-
-          --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5
-
     steps:
       - uses: actions/checkout@v3
 
@@ -68,19 +57,15 @@ jobs:
         run: |
           cargo build
 
-      - name: Spin up the test object stores
-        run: docker compose up -d
+      - name: Spin up the test object stores and postgres catalog
+        run: docker compose up -d --wait || true
 
       # TODO split tests into unit and integration (one requires postgres?)
       - name: Run tests
         run: |
           cargo test --workspace
         env:
-          # Database URL for end-to-end + postgres repository tests; the host is `localhost` since we don't specify
-          # a child container to run the job in (and instead run everything in the root specified by `runs-on`).
-          # The port has been randomly assigned so we need to fetch it.
-          DATABASE_URL:
-            "postgres://postgres:postgres@localhost:${{ job.services.postgres.ports[5432] }}/db_test"
+          DATABASE_URL: "postgres://postgres:test@localhost:5432/db_test"
 
       # TODO recompiles the whole thing with different flags (busting the cache?,
       #      also codecov needs a token for private repos; also this breaks in Docker

Cargo.lock

@@ -116,8 +116,8 @@ hex = ">=0.4.0"
 itertools = { workspace = true }
 lazy_static = ">=1.4.0"
 # Pick up hashbrown >0.13.1 to resolve conflicts with wasm crates
-metrics = { git = "https://github.com/metrics-rs/metrics", rev = "bc29c4f984605ccadb0daede451b28523c1aff86", optional = true }
-metrics-exporter-prometheus = { git = "https://github.com/metrics-rs/metrics", rev = "bc29c4f984605ccadb0daede451b28523c1aff86", optional = true }
+metrics = { version = "0.22.1", optional = true }
+metrics-exporter-prometheus = { version = "0.13.1", optional = true }
 moka = { version = "0.12.2", default_features = false, features = ["future", "atomic64", "quanta"] }
 object_store = "0.9"
 parking_lot = "0.12.1"

Cargo.toml

@@ -1,5 +1,14 @@
 version: "3.9"
 services:
+  postgres:
+    image: postgres:16
+    environment:
+      POSTGRES_USER: postgres
+      POSTGRES_PASSWORD: test
+    ports:
+      - "5432:5432"
+    restart: unless-stopped
+
   minio:
     image: minio/minio:latest
     ports:

docker-compose.yml

@@ -189,7 +189,7 @@ pub fn setup_metrics(metrics: &schema::Metrics) {
     describe_counter!(GRPC_REQUESTS, "Counter tracking gRPC request statistics");
 }
 
-pub async fn build_context(cfg: &schema::SeafowlConfig) -> Result<SeafowlContext> {
+pub async fn build_context(cfg: schema::SeafowlConfig) -> Result<SeafowlContext> {
     let mut runtime_config = RuntimeConfig::new();
     if let Some(max_memory) = cfg.runtime.max_memory {
         runtime_config = runtime_config
@@ -217,7 +217,7 @@ pub async fn build_context(cfg: &schema::SeafowlConfig) -> Result<SeafowlContext
     // Register the HTTP object store for external tables
     add_http_object_store(&context, &cfg.misc.ssl_cert_file);
 
-    let metastore = build_metastore(cfg, internal_object_store.clone()).await;
+    let metastore = build_metastore(&cfg, internal_object_store.clone()).await;
 
     // Create default DB/collection
     if let Err(CatalogError::CatalogDoesNotExist { .. }) =
@@ -244,12 +244,12 @@ pub async fn build_context(cfg: &schema::SeafowlConfig) -> Result<SeafowlContext
     // (it will reload its schema before running the query)
 
     Ok(SeafowlContext {
+        config: cfg,
         inner: context,
         metastore: Arc::new(metastore),
         internal_object_store,
         default_catalog: DEFAULT_DB.to_string(),
         default_schema: DEFAULT_SCHEMA.to_string(),
-        max_partition_size: cfg.misc.max_partition_size,
     })
 }
 
@@ -297,7 +297,7 @@ mod tests {
             },
         };
 
-        let context = build_context(&config).await.unwrap();
+        let context = build_context(config).await.unwrap();
 
         // Run a query against the context to test it works
         let results = context

src/config/context.rs

@@ -199,8 +199,8 @@ impl ObjectCacheProperties {
         Arc::new(CachingObjectStore::new(
             inner,
             &path,
-            self.capacity,
             self.min_fetch_size,
+            self.capacity,
             Duration::from_secs(self.ttl_s),
         ))
     }
@@ -473,7 +473,9 @@ If Seafowl is running on GCP a token should be fetched using the GCP metadata en
 
 pub fn load_config(path: &Path) -> Result<SeafowlConfig, ConfigError> {
     let config = Config::builder()
-        .add_source(File::with_name(path.to_str().expect("Error parsing path")))
+        .add_source(
+            File::with_name(path.to_str().expect("Error parsing path")).required(false),
+        )
         .add_source(Environment::with_prefix(ENV_PREFIX).separator(ENV_SEPARATOR));
 
     config.build()?.try_deserialize().and_then(validate_config)
@@ -486,7 +488,7 @@ pub fn load_config_from_string(
     env_override: Option<Map<String, String>>,
 ) -> Result<SeafowlConfig, ConfigError> {
     let config = Config::builder()
-        .add_source(File::from_str(config_str, FileFormat::Toml))
+        .add_source(File::from_str(config_str, FileFormat::Toml).required(false))
         .add_source(
             Environment::with_prefix(ENV_PREFIX)
                 .separator(ENV_SEPARATOR)

src/config/schema.rs

@@ -381,7 +381,7 @@ impl SeafowlContext {
             plan,
             self.internal_object_store.clone(),
             table_prefix,
-            self.max_partition_size,
+            self.config.misc.max_partition_size,
         )
         .await?;
 

src/context/delta.rs

@@ -9,6 +9,7 @@ use crate::object_store::wrapped::InternalObjectStore;
 use crate::wasm_udf::data_types::{get_volatility, CreateFunctionDetails};
 use crate::wasm_udf::wasm::create_udf_from_wasm;
 
+use crate::config::schema::SeafowlConfig;
 use base64::{engine::general_purpose::STANDARD, Engine};
 pub use datafusion::error::{DataFusionError as Error, Result};
 use datafusion::{error::DataFusionError, prelude::SessionContext, sql::TableReference};
@@ -21,12 +22,12 @@ use uuid::Uuid;
 // The core Seafowl object, responsible for parsing, logical and physical planning, as well as
 // interacting with the catalog and object store.
 pub struct SeafowlContext {
+    pub config: SeafowlConfig,
     pub inner: SessionContext,
     pub metastore: Arc<Metastore>,
     pub internal_object_store: Arc<InternalObjectStore>,
     pub default_catalog: String,
     pub default_schema: String,
-    pub max_partition_size: u32,
 }
 
 impl SeafowlContext {
@@ -42,12 +43,12 @@ impl SeafowlContext {
             build_state_with_table_factories(session_config, self.inner().runtime_env());
 
         Arc::from(SeafowlContext {
+            config: self.config.clone(),
             inner: SessionContext::new_with_state(state),
             metastore: self.metastore.clone(),
             internal_object_store: self.internal_object_store.clone(),
             default_catalog: catalog,
             default_schema: schema,
-            max_partition_size: self.max_partition_size,
         })
     }
 
@@ -226,7 +227,7 @@ pub mod test_utils {
             runtime: Default::default(),
             misc: Default::default(),
         };
-        build_context(&config).await.unwrap()
+        build_context(config).await.unwrap()
     }
 
     pub async fn in_memory_context_with_test_db() -> Arc<SeafowlContext> {

src/context/mod.rs

@@ -347,7 +347,7 @@ impl SeafowlContext {
                         &update_plan,
                         self.internal_object_store.clone(),
                         table_prefix,
-                        self.max_partition_size,
+                        self.config.misc.max_partition_size,
                     )
                     .await?;
 
@@ -460,7 +460,7 @@ impl SeafowlContext {
                                 &filter_plan,
                                 self.internal_object_store.clone(),
                                 table_prefix,
-                                self.max_partition_size,
+                                self.config.misc.max_partition_size,
                             )
                             .await?;
 

src/context/physical.rs

@@ -9,6 +9,8 @@ use warp::{hyper, Rejection};
 
 use arrow::json::writer::record_batches_to_json_rows;
 use arrow::record_batch::RecordBatch;
+#[cfg(feature = "frontend-arrow-flight")]
+use arrow_flight::flight_service_client::FlightServiceClient;
 
 use arrow_integration_test::{schema_from_json, schema_to_json};
 use arrow_schema::SchemaRef;
@@ -42,9 +44,9 @@ use crate::auth::{token_to_principal, AccessPolicy, Action, UserContext};
 use crate::catalog::DEFAULT_DB;
 #[cfg(feature = "metrics")]
 use crate::config::context::HTTP_REQUESTS;
-use crate::config::schema::{AccessSettings, MEBIBYTES};
+use crate::config::schema::{AccessSettings, HttpFrontend, MEBIBYTES};
 use crate::{
-    config::schema::{str_to_hex_hash, HttpFrontend},
+    config::schema::str_to_hex_hash,
     context::logical::{is_read_only, is_statement_read_only},
     context::SeafowlContext,
 };
@@ -482,8 +484,29 @@ async fn load_part(mut part: Part) -> Result<Vec<u8>, ApiError> {
     Ok(bytes)
 }
 
-// We need the allow to silence the compiler: it asks us to add warp::generic::Tuple to the first
-// parameter of the return type, but that struct is not exportable (generic is private).
+/// GET /healthz or /readyz
+pub async fn health_endpoint(context: Arc<SeafowlContext>) -> Result<Response, ApiError> {
+    #[cfg(feature = "frontend-arrow-flight")]
+    if let Some(flight) = &context.config.frontend.flight {
+        // TODO: run SELECT 1 or something similar?
+        if let Err(err) = FlightServiceClient::connect(format!(
+            "http://{}:{}",
+            flight.bind_host, flight.bind_port
+        ))
+        .await
+        {
+            warn!(%err, "Arrow Flight client can't connect, health check failed");
+            return Ok(warp::reply::with_status(
+                "not_ready",
+                StatusCode::SERVICE_UNAVAILABLE,
+            )
+            .into_response());
+        };
+    };
+
+    Ok(warp::reply::with_status("ready", StatusCode::OK).into_response())
+}
+
 pub fn filters(
     context: Arc<SeafowlContext>,
     config: HttpFrontend,
@@ -501,22 +524,34 @@ pub fn filters(
         .max_age(CORS_MAXAGE);
 
     let log = warp::log::custom(|info: Info<'_>| {
+        let path = info.path();
+
         #[cfg(feature = "metrics")]
-        counter!(
-            HTTP_REQUESTS,
-            "method" => info.method().as_str().to_string(),
-            // Omit a potential db prefix or url-encoded query from the path
-            "route" => if info.path().contains("/upload/") { "/upload" } else { "/q" },
-            "status" => info.status().as_u16().to_string(),
-        )
-        .increment(1);
+        {
+            let route = if path.contains("/upload/") {
+                "/upload".to_string()
+            } else if path.contains("/q") {
+                "/q".to_string()
+            } else {
+                path.to_string()
+            };
+
+            counter!(
+                HTTP_REQUESTS,
+                "method" => info.method().as_str().to_string(),
+                // Omit a potential db prefix or url-encoded query from the path
+                "route" => route,
+                "status" => info.status().as_u16().to_string(),
+            )
+            .increment(1);
+        }
 
         info!(
             target: module_path!(),
             "{} \"{} {} {:?}\" {} \"{}\" \"{}\" {:?}",
             info.remote_addr().map(|addr| addr.to_string()).unwrap_or("-".to_string()),
             info.method(),
-            info.path(),
+            path,
             info.version(),
             info.status().as_u16(),
             info.referer().unwrap_or("-"),
@@ -581,7 +616,7 @@ pub fn filters(
         .map(into_response);
 
     // Upload endpoint
-    let ctx = context;
+    let ctx = context.clone();
     let upload_route = warp::path!(String / "upload" / String / String)
         .or(warp::any()
             .map(move || DEFAULT_DB.to_string())
@@ -596,9 +631,21 @@ pub fn filters(
         .then(upload)
         .map(into_response);
 
+    // Health-check/readiness probe
+    let ctx = context;
+    let health_route = warp::path!("healthz")
+        .or(warp::path!("readyz"))
+        .and(warp::path::end())
+        .and(warp::get())
+        .unify()
+        .and(warp::any().map(move || ctx.clone()))
+        .then(health_endpoint)
+        .map(into_response);
+
     cached_read_query_route
         .or(uncached_read_write_query_route)
         .or(upload_route)
+        .or(health_route)
         .with(cors)
         .with(log)
         .map(|r| with_header(r, header::VARY, VARY))