Support for Lakebase - Autoscaling (#225)

* Fix

* Fix

Author: krinart

.github/workflows/run_spicebench.yml

@@ -214,7 +214,7 @@ jobs:
           SCENARIO: ${{ github.event.inputs.scenario || 'tpch' }}
           SYSTEM_UNDER_TEST: ${{ github.event.inputs.system_under_test || 'spice_cloud' }}
           SYSTEM_ADAPTER: ${{ github.event.inputs.system_under_test || 'spice_cloud' }}
-          NUM_QUERY_CLIENTS: '8'
+          NUM_QUERY_CLIENTS: '2'
           ETL_BUCKET: 'spicebench'
           ETL_PREFIX: ${{ github.event.inputs.etl_type == 'changes' && 'data-gen-mutable' || 'data-gen' }}
           SCALE_FACTOR: ${{ github.event.inputs.scale_factor || '1' }}
@@ -232,9 +232,9 @@ jobs:
           S3_AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
           LAKEBASE_PG_HOST: ${{ secrets.LAKEBASE_PG_HOST }}
           LAKEBASE_PG_USER: ${{ secrets.LAKEBASE_PG_USER }}
-          LAKEBASE_PG_DB_NAME: ${{ secrets.LAKEBASE_PG_DB_NAME }}
           LAKEBASE_PG_SCHEMA: ${{ secrets.LAKEBASE_PG_SCHEMA }}
-          LAKEBASE_DATABASE_INSTANCE: ${{ secrets.LAKEBASE_DATABASE_INSTANCE }}
+          LAKEBASE_PROJECT: ${{ secrets.LAKEBASE_PROJECT }}
+          LAKEBASE_BRANCH: ${{ secrets.LAKEBASE_BRANCH }}
           SPIDAPTER_ICEBERG_REGION: us-west-1
           SPIDAPTER_ICEBERG_CATALOG_FROM: iceberg:https://glue.us-west-1.amazonaws.com/iceberg/v1/catalogs/211125479522/namespaces
         run: |
@@ -309,11 +309,11 @@ jobs:
 
           if [ "${SYSTEM_UNDER_TEST}" = "databricks-lakebase" ]; then
             ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env DATABRICKS_COMPUTE_MODE=lakebase"
-            ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PG_HOST=${LAKEBASE_PG_HOST}"
+            ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PG_HOST=ep-sweet-morning-d1z2teqy.database.us-west-2.cloud.databricks.com"
             ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PG_USER=${LAKEBASE_PG_USER}"
-            ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PG_DB_NAME=${LAKEBASE_PG_DB_NAME}"
             ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PG_SCHEMA=${LAKEBASE_PG_SCHEMA}"
-            ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_DATABASE_INSTANCE=${LAKEBASE_DATABASE_INSTANCE}"
+            ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PROJECT=${LAKEBASE_PROJECT}"
+            ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_BRANCH=${LAKEBASE_BRANCH}"
           fi
 
           ~/.spice/bin/spicebench run \

crates/data-generation/src/storage/file.rs

@@ -157,15 +157,6 @@ impl DataStorage for FileStorage {
         .await
         .map_err(|e| anyhow::anyhow!("spawn_blocking panicked reading parquet: {e}"))??;
 
-        if table_name == "lineitem"
-            && let Some(result) = result.as_ref()
-        {
-            eprintln!(
-                "[etl-read] table={table_name} batch_id={batch_id} rows={}",
-                result.rows_read,
-            );
-        }
-
         Ok(result)
     }
 

src/commands/load/mod.rs

@@ -532,7 +532,8 @@ async fn run_checkpoint_validation(
 ) -> CheckpointValidationResult {
     let deadline = tokio::time::Instant::now() + max_wait;
 
-    let probe_query = &queries[0];
+    // Use query2 as the probe query.  Query1 is too slow.
+    let probe_query = &queries[1];
     let Some(probe_expected) = expected_results.get(&probe_query.name) else {
         eprintln!(
             "Checkpoint {checkpoint_idx}: no expected results for probe query '{}', skipping validation",

src/commands/run.rs

@@ -125,7 +125,6 @@ async fn run_benchmark(
         &mutations,
         target_config.clone(),
     )?;
-
     let (setup_response, mut pipeline) = {
         let setup_response = system_adapter_client
             .lock()

system-adapters/databricks/src/main.rs

@@ -449,7 +449,7 @@ impl AdapterConfig {
 struct DatabaseCredentialResponse {
     token: String,
     #[serde(alias = "expiration_time", alias = "expire_time")]
-    expiration_time: Option<String>,
+    _expiration_time: Option<String>,
 }
 
 impl DatabricksAdapter {
@@ -1759,12 +1759,12 @@ impl DatabricksAdapter {
                 });
                 (url, payload)
             }
-            LakebaseSyncTarget::Project { name, branch } => {
+            LakebaseSyncTarget::Project { name: _, branch: _ } => {
                 // Autoscaling uses the postgres API path and endpoint-based credential generation
                 let endpoint_path =
-                    format!("projects/{}/branches/{}/endpoints/default", name, branch);
+                    format!("projects/spicebench/branches/production/endpoints/primary");
                 let url = format!(
-                    "https://{}/api/2.0/postgres/generate-database-credential",
+                    "https://{}/api/2.0/postgres/credentials",
                     self.config.endpoint
                 );
                 let payload = json!({
@@ -1775,8 +1775,6 @@ impl DatabricksAdapter {
             }
         };
 
-        eprintln!("[databricks-adapter] generating fresh Lakebase PG OAuth token");
-
         let response = self
             .client
             .post(&url)
@@ -1794,10 +1792,6 @@ impl DatabricksAdapter {
         }
 
         let cred: DatabaseCredentialResponse = response.json().await?;
-        eprintln!(
-            "[databricks-adapter] Lakebase PG token generated, expires: {}",
-            cred.expiration_time.as_deref().unwrap_or("unknown")
-        );
 
         Ok(cred.token)
     }
@@ -2030,9 +2024,6 @@ impl Handler for DatabricksAdapter {
         match variant {
             DatabricksVariant::Databricks => {}
             DatabricksVariant::Lakebase => {
-                eprintln!("[databricks-adapter] Waiting 2 minutes for schema to initialize");
-                std::thread::sleep(Duration::from_secs(120));
-
                 let lakebase_config = match &self.config.compute_target {
                     ComputeTarget::Lakebase(cfg) => cfg,
                     _ => {
@@ -2387,7 +2378,10 @@ impl Handler for DatabricksAdapter {
 
 #[tokio::main]
 async fn main() -> Result<()> {
-    let cli = Cli::parse();
+    let cli = Cli::try_parse().map_err(|e| {
+        eprintln!("[databricks-adapter] CLI parse error: {e}");
+        e.exit();
+    })?;
 
     match cli.command {
         Commands::Stdio(args) => {