feat: System adapter SetupResponse supports specifying catalog/namespace path (#137)

* feat: System adapter SetupResponse supports specifying catalog/namespace path

* chore: auto-fix cargo fmt + clippy

* Update lib.rs

* feat: Add ETL sink type to system adapter setup request

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: peasee

crates/system-adapter-protocol/src/client.rs

@@ -181,11 +181,13 @@ impl Client {
         run_id: uuid::Uuid,
         metadata: std::collections::HashMap<String, serde_json::Value>,
         datasets: std::collections::HashMap<String, crate::DatasetConfig>,
+        etl_sink_type: Option<crate::EtlSinkType>,
     ) -> Result<crate::SetupResponse> {
         let request = crate::SetupRequest {
             run_id,
             metadata,
             datasets,
+            etl_sink_type,
         };
         let rpc_request = JsonRpcRequest::new(1, crate::methods::SETUP, request);
         let response = self.call_typed(rpc_request).await?;

crates/system-adapter-protocol/src/lib.rs

@@ -41,7 +41,9 @@ limitations under the License.
 //!
 //! // Setup a benchmark run
 //! let run_id = Uuid::new_v4();
-//! let setup_response = client.setup(run_id, HashMap::new(), HashMap::new()).await?;
+//! let setup_response = client
+//!     .setup(run_id, HashMap::new(), HashMap::new(), None)
+//!     .await?;
 //!
 //! println!("Driver: {:?}", setup_response.driver);
 //!
@@ -57,7 +59,7 @@ limitations under the License.
 //! # #[cfg(feature = "server")]
 //! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
 //! use system_adapter_protocol::{
-//!     AdbcDriver, DatasetConfig, Handler, Server, SetupResponse, TeardownResponse,
+//!     AdbcDriver, DatasetConfig, EtlSinkType, Handler, Server, SetupResponse, TeardownResponse,
 //! };
 //! use async_trait::async_trait;
 //! use std::collections::HashMap;
@@ -72,11 +74,13 @@ limitations under the License.
 //!         run_id: Uuid,
 //!         metadata: HashMap<String, serde_json::Value>,
 //!         datasets: HashMap<String, DatasetConfig>,
+//!         etl_sink_type: Option<EtlSinkType>,
 //!     ) -> Result<SetupResponse, String> {
-//!         let _ = (metadata, datasets);
+//!         let _ = (metadata, datasets, etl_sink_type);
 //!         Ok(SetupResponse {
 //!             driver: AdbcDriver::Flightsql,
 //!             db_kwargs: HashMap::new(),
+//!             catalog_namespace: None,
 //!         })
 //!     }
 //!
@@ -121,6 +125,17 @@ pub enum AdbcDriver {
     Postgresql,
 }
 
+/// ETL sink type used by spicebench for this run.
+///
+/// This is provided to adapters in [`SetupRequest`] so they can optionally
+/// adjust setup behavior based on how data is loaded.
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "snake_case")]
+pub enum EtlSinkType {
+    Hive,
+    Adbc,
+}
+
 impl std::fmt::Display for AdbcDriver {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -160,6 +175,9 @@ pub struct SetupRequest {
     pub metadata: HashMap<String, serde_json::Value>,
     /// Map of dataset name to dataset definition
     pub datasets: HashMap<String, DatasetConfig>,
+    /// Optional ETL sink type selected by spicebench.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub etl_sink_type: Option<EtlSinkType>,
 }
 
 /// Response from setup request containing ADBC connection information
@@ -169,6 +187,10 @@ pub struct SetupResponse {
     pub driver: AdbcDriver,
     /// Driver-specific connection parameters
     pub db_kwargs: HashMap<String, serde_json::Value>,
+    /// Optional catalog/namespace path where benchmark tables were created
+    /// (e.g. "catalog.schema").
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub catalog_namespace: Option<String>,
 }
 /// Request to teardown a benchmark run
 ///

crates/system-adapter-protocol/src/server.rs

@@ -17,8 +17,8 @@ limitations under the License.
 //! Server implementations for system adapter JSON-RPC protocol.
 
 use crate::{
-    DatasetConfig, JsonRpcError, JsonRpcResponse, MetricsRequest, MetricsResponse, SetupRequest,
-    SetupResponse, TeardownRequest, TeardownResponse, error_codes, methods,
+    DatasetConfig, EtlSinkType, JsonRpcError, JsonRpcResponse, MetricsRequest, MetricsResponse,
+    SetupRequest, SetupResponse, TeardownRequest, TeardownResponse, error_codes, methods,
 };
 use async_trait::async_trait;
 use serde::de::DeserializeOwned;
@@ -76,6 +76,7 @@ pub trait Handler: Send + Sync {
         run_id: Uuid,
         metadata: HashMap<String, serde_json::Value>,
         datasets: HashMap<String, DatasetConfig>,
+        etl_sink_type: Option<EtlSinkType>,
     ) -> std::result::Result<SetupResponse, String>;
 
     /// Teardown a benchmark run
@@ -235,7 +236,7 @@ impl<H: Handler> Server<H> {
         };
         Self::handler_response(
             self.handler
-                .setup(req.run_id, req.metadata, req.datasets)
+                .setup(req.run_id, req.metadata, req.datasets, req.etl_sink_type)
                 .await,
             id,
         )
@@ -285,10 +286,12 @@ mod tests {
             _run_id: Uuid,
             _metadata: HashMap<String, serde_json::Value>,
             _datasets: HashMap<String, DatasetConfig>,
+            _etl_sink_type: Option<EtlSinkType>,
         ) -> std::result::Result<SetupResponse, String> {
             Ok(SetupResponse {
                 driver: crate::AdbcDriver::Flightsql,
                 db_kwargs: HashMap::new(),
+                catalog_namespace: None,
             })
         }
 

crates/test-framework/src/queries/mod.rs

@@ -225,6 +225,10 @@ impl Query {
     /// becomes:
     ///   `SELECT * FROM arrow.customer WHERE c_custkey = 1`
     ///
+    /// Multi-part namespaces are also supported. For example, if `reference_schema`
+    /// is `catalog.schema`, the query becomes:
+    ///   `SELECT * FROM catalog.schema.customer WHERE c_custkey = 1`
+    ///
     /// Uses `DataFusion`'s SQL parser to parse the query, rewrite all table references,
     /// and unparse back to SQL. This works with any valid SQL query.
     ///
@@ -237,6 +241,12 @@ impl Query {
         use sqlparser::parser::Parser;
         use std::ops::ControlFlow;
 
+        let namespace_parts: Vec<ObjectNamePart> = reference_schema
+            .split('.')
+            .filter(|part| !part.is_empty())
+            .map(|part| ObjectNamePart::Identifier(Ident::new(part)))
+            .collect();
+
         // Parse the SQL query using sqlparser
         let dialect = sqlparser::dialect::PostgreSqlDialect {};
         let mut statements = Parser::parse_sql(&dialect, &self.sql).map_err(|e| {
@@ -262,10 +272,9 @@ impl Query {
         let _ = visit_relations_mut(statement, |table_name| {
             // Only rewrite if the table doesn't already have a schema prefix (single-part name)
             if table_name.0.len() == 1 {
-                // Prepend the reference schema to the table name
-                table_name
-                    .0
-                    .insert(0, ObjectNamePart::Identifier(Ident::new(reference_schema)));
+                for namespace_part in namespace_parts.iter().rev() {
+                    table_name.0.insert(0, namespace_part.clone());
+                }
             }
             ControlFlow::<()>::Continue(())
         });
@@ -992,6 +1001,23 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_rewrite_with_multi_part_namespace() {
+        let query = Query::new(
+            "test_multi_part".into(),
+            "SELECT * FROM customer JOIN orders ON customer.c_custkey = orders.o_custkey".into(),
+            false,
+        );
+
+        let rewritten = query
+            .rewrite_with_reference_schema("main.catalog")
+            .expect("Failed to rewrite query with multi-part namespace");
+
+        let sql = rewritten.sql.as_ref();
+        assert!(sql.contains("main.catalog.customer"));
+        assert!(sql.contains("main.catalog.orders"));
+    }
+
     #[test]
     fn test_rewrite_with_existing_schema_prefix() {
         let query = Query::new(

src/commands/load/mod.rs

@@ -331,6 +331,7 @@ pub(crate) async fn run(
     etl_pipeline: &mut ETLPipeline,
     checkpoint_steps: Option<usize>,
     checkpoint_dir: Option<&Path>,
+    query_catalog_namespace: Option<String>,
 ) -> anyhow::Result<()> {
     let metric_attributes = run_metric_attributes(common_args);
 
@@ -419,8 +420,12 @@ pub(crate) async fn run(
     let has_checkpoint_validation =
         common_args.validate_results && checkpoint_steps.is_some() && checkpoint_dir.is_some();
 
-    let (query_set, test_builder) =
-        super::build_test_with_validation(scenario, test_builder).await?;
+    let (query_set, test_builder) = super::build_test_with_validation(
+        scenario,
+        test_builder,
+        query_catalog_namespace.as_deref(),
+    )
+    .await?;
 
     // Build ordered query names for mapping checkpoint query_idx → query name.
     let queries = query_set.get_queries(None, None, None).await?;

src/commands/mod.rs

@@ -50,6 +50,22 @@ pub(crate) fn create_telemetry_with_resource(common: &CommonArgs, resource: Reso
     Telemetry::new_with_resource(&resource, "SPICEAI_BENCHMARK_METRICS_KEY")
 }
 
+fn rewrite_queries_with_catalog_namespace(
+    queries: Vec<test_framework::queries::Query>,
+    query_catalog_namespace: Option<&str>,
+) -> anyhow::Result<Vec<test_framework::queries::Query>> {
+    if let Some(catalog_namespace) = query_catalog_namespace
+        && !catalog_namespace.trim().is_empty()
+    {
+        return queries
+            .into_iter()
+            .map(|query| query.rewrite_with_reference_schema(catalog_namespace))
+            .collect::<anyhow::Result<Vec<_>>>();
+    }
+
+    Ok(queries)
+}
+
 /// Build a test configuration with validation data if applicable
 ///
 /// This is a common helper for bench, throughput, and load tests that:
@@ -64,9 +80,13 @@ pub(crate) fn create_telemetry_with_resource(common: &CommonArgs, resource: Reso
 pub(crate) async fn build_test_with_validation(
     scenario: &Scenario,
     test_builder: NotStarted,
+    query_catalog_namespace: Option<&str>,
 ) -> anyhow::Result<(QuerySet, NotStarted)> {
     let query_set = scenario.load_query_set()?;
-    let queries = query_set.get_queries(None, None, None).await?;
+    let queries = rewrite_queries_with_catalog_namespace(
+        query_set.get_queries(None, None, None).await?,
+        query_catalog_namespace,
+    )?;
 
     let test_builder = test_builder
         .with_query_set(queries)
@@ -117,3 +137,26 @@ macro_rules! wait_test_and_memory {
         }
     };
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_rewrite_queries_with_catalog_namespace() {
+        let queries = vec![test_framework::queries::Query::new(
+            "q1".into(),
+            "SELECT * FROM customer".into(),
+            false,
+        )];
+
+        let rewritten = rewrite_queries_with_catalog_namespace(queries, Some("catalog.schema"))
+            .expect("rewrite should succeed");
+
+        assert_eq!(rewritten.len(), 1);
+        assert_eq!(
+            rewritten[0].sql.as_ref(),
+            "SELECT * FROM catalog.schema.customer"
+        );
+    }
+}

src/main.rs

@@ -127,6 +127,10 @@ async fn run_benchmark(
     let data_source: Arc<dyn DataStorage> = source.clone();
 
     let mut setup_response_for_run: Option<system_adapter_protocol::SetupResponse> = None;
+    let etl_sink_type = match common.etl_sink {
+        EtlSink::Hive => system_adapter_protocol::EtlSinkType::Hive,
+        EtlSink::Adbc => system_adapter_protocol::EtlSinkType::Adbc,
+    };
 
     let (target, target_config, target_kind, adbc_sink): (
         Arc<dyn Sink>,
@@ -177,6 +181,7 @@ async fn run_benchmark(
                     run_id,
                     setup_metadata.clone(),
                     setup_pipeline.create_tables_request_datasets(),
+                    Some(etl_sink_type),
                 )
                 .await
                 .map_err(|e| anyhow::anyhow!("Failed to setup system adapter: {e}"))?;
@@ -245,12 +250,14 @@ async fn run_benchmark(
                 run_id,
                 setup_metadata,
                 pipeline.create_tables_request_datasets(),
+                Some(etl_sink_type),
             )
             .await
             .map_err(|e| anyhow::anyhow!("Failed to setup system adapter: {e}"))?
     };
 
     let driver_name = setup_response.driver.to_string();
+    let query_catalog_namespace = setup_response.catalog_namespace.clone();
     let db_kwargs = setup_response.db_kwargs;
 
     let load_conn = match AdbcConnection::create(&driver_name, db_kwargs) {
@@ -273,6 +280,7 @@ async fn run_benchmark(
         &mut pipeline,
         checkpoint_steps,
         Some(checkpoint_dir.path()),
+        query_catalog_namespace,
     )
     .await?;
 

system-adapters/databricks/src/main.rs

@@ -25,7 +25,7 @@ use reqwest::StatusCode;
 use serde::{Deserialize, Serialize};
 use serde_json::{Value, json};
 use system_adapter_protocol::{
-    AdbcDriver, DatasetConfig, Handler, Server, SetupResponse, TeardownResponse,
+    AdbcDriver, DatasetConfig, EtlSinkType, Handler, Server, SetupResponse, TeardownResponse,
 };
 use uuid::Uuid;
 
@@ -1707,7 +1707,9 @@ impl Handler for DatabricksAdapter {
         run_id: Uuid,
         metadata: HashMap<String, Value>,
         datasets: HashMap<String, DatasetConfig>,
+        etl_sink_type: Option<EtlSinkType>,
     ) -> std::result::Result<SetupResponse, String> {
+        let _ = etl_sink_type;
         eprintln!("[databricks-adapter] setup: run_id={run_id}");
         eprintln!("[databricks-adapter] endpoint={}", self.config.endpoint);
 
@@ -1838,6 +1840,7 @@ impl Handler for DatabricksAdapter {
                 Ok(SetupResponse {
                     driver: AdbcDriver::Postgresql,
                     db_kwargs: HashMap::from([("uri".to_string(), Value::String(pg_uri))]),
+                    catalog_namespace: Some(format!("{}.{}", self.config.catalog, self.config.schema)),
                 })
             }
             _ => Ok(SetupResponse {
@@ -1846,6 +1849,7 @@ impl Handler for DatabricksAdapter {
                     "uri".to_string(),
                     Value::String(self.databricks_uri()),
                 )]),
+                catalog_namespace: Some(format!("{}.{}", self.config.catalog, self.config.schema)),
             }),
         }
     }