From ce38ad80c1cafa0b2d4755e4c42c59ed9434aab1 Mon Sep 17 00:00:00 2001
From: Viktor Yershov <viktor@spice.ai>
Date: Sat, 16 May 2026 12:37:55 -0700
Subject: [PATCH] Fix Lakebase

---
 src/commands/run.rs                   |  5 ++++-
 system-adapters/databricks/Cargo.lock | 13 +++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/src/commands/run.rs b/src/commands/run.rs
index c2b073e8..739b33e8 100644
--- a/src/commands/run.rs
+++ b/src/commands/run.rs
@@ -130,10 +130,13 @@ async fn run_benchmark(
     // primary keys. MERGE INTO handles matching via the ON clause and uses
     // delete+insert execution, which conflicts with Cayenne's automatic
     // on_conflict: Upsert behavior on primary-key tables.
+    // Note: preserve PKs for system adapter setup (e.g. Lakebase synced tables
+    // require primary_key_columns) and only strip them for the ETL pipeline.
     let uses_staging_table = std::env::var("SPICEBENCH_ADBC_UPDATE_STRATEGY")
         .ok()
         .map(|v| v.eq_ignore_ascii_case("staging_table"))
         .unwrap_or(false);
+    let setup_datasets = datasets.clone();
     if uses_staging_table {
         for config in datasets.values_mut() {
             config.primary_key_columns.clear();
@@ -146,7 +149,7 @@ async fn run_benchmark(
             .setup(
                 run_id,
                 setup_metadata.clone(),
-                datasets.clone(),
+                setup_datasets,
                 Some(etl_sink_type),
             )
             .await
diff --git a/system-adapters/databricks/Cargo.lock b/system-adapters/databricks/Cargo.lock
index 669299cb..c560d486 100644
--- a/system-adapters/databricks/Cargo.lock
+++ b/system-adapters/databricks/Cargo.lock
@@ -1470,6 +1470,7 @@ dependencies = [
  "serde",
  "serde_json",
  "tokio",
+ "tracing",
  "uuid",
 ]
 
@@ -1682,9 +1683,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100"
 dependencies = [
  "pin-project-lite",
+ "tracing-attributes",
  "tracing-core",
 ]
 
+[[package]]
+name = "tracing-attributes"
+version = "0.1.31"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "tracing-core"
 version = "0.1.36"