Use cte based approach to create managed Databricks table (#120)

Author: sgrebnov

src/args/mod.rs

@@ -101,15 +101,15 @@ pub struct CommonArgs {
     pub(crate) system_adapter_env: Vec<(String, String)>,
 
     /// S3 bucket name for the ETL source and target
-    #[arg(long)]
+    #[arg(long, default_value = "spiceai-public-datasets")]
     pub(crate) etl_bucket: String,
 
     /// S3 key prefix (the `{prefix}` portion of `{prefix}/{scenario}/{version}/`)
-    #[arg(long, default_value = "")]
+    #[arg(long, default_value = "data-gen")]
     pub(crate) etl_prefix: String,
 
     /// Version identifier for the data generation to read from.
-    #[arg(long)]
+    #[arg(long, default_value = "1")]
     pub(crate) etl_version: String,
 
     /// Base S3 key prefix for the ETL target (rehydrated) data.
@@ -118,7 +118,7 @@ pub struct CommonArgs {
     pub(crate) etl_target_base_prefix: String,
 
     /// AWS region for the ETL S3 bucket
-    #[arg(long)]
+    #[arg(long, default_value = "us-east-1")]
     pub(crate) etl_region: Option<String>,
 
     /// S3 endpoint URL for the ETL bucket (for MinIO/LocalStack)

src/main.rs

@@ -221,14 +221,27 @@ async fn main() -> anyhow::Result<()> {
 
     let cli = Cli::parse();
 
-    // --- Connect to S3 and read version metadata ---
+    let scenario_name = cli.common.scenario.to_string();
+    let version_prefix = build_version_prefix(
+        &cli.common.etl_prefix,
+        &scenario_name,
+        &cli.common.etl_version,
+    );
+    tracing::info!(
+        etl_source = %format!("s3://{}/{}/tables/", cli.common.etl_bucket, version_prefix),
+        etl_bucket = %cli.common.etl_bucket,
+        etl_prefix = %cli.common.etl_prefix,
+        etl_version = %cli.common.etl_version,
+        etl_region = ?cli.common.etl_region,
+        table_format = %cli.common.table_format,
+        scenario = %scenario_name,
+        concurrency = cli.common.concurrency,
+        "ETL configuration"
+    );
+
     let source_config = TargetConfig {
         bucket: cli.common.etl_bucket.clone(),
-        prefix: build_version_prefix(
-            &cli.common.etl_prefix,
-            &cli.common.scenario.to_string(),
-            &cli.common.etl_version,
-        ),
+        prefix: version_prefix,
         region: cli.common.etl_region.clone(),
         endpoint: cli.common.etl_endpoint.clone(),
     };

system-adapters/databricks/src/main.rs

@@ -162,6 +162,7 @@ enum TableFormat {
 }
 
 impl TableFormat {
+    #[allow(dead_code)]
     fn as_sql_using(self) -> &'static str {
         match self {
             Self::Parquet => "PARQUET",
@@ -182,6 +183,7 @@ impl TableFormat {
 
 #[derive(Debug, Clone)]
 struct RunState {
+    #[allow(dead_code)]
     table_format: TableFormat,
     created_tables: Vec<String>,
     cluster_id: Option<String>,
@@ -384,6 +386,7 @@ impl DatabricksAdapter {
         )
     }
 
+    #[allow(dead_code)]
     fn sql_type_for_arrow(data_type: &DataType) -> Result<String> {
         match data_type {
             DataType::Boolean => Ok("BOOLEAN".to_string()),
@@ -408,6 +411,7 @@ impl DatabricksAdapter {
         }
     }
 
+    #[allow(dead_code)]
     fn create_table_ddl(
         &self,
         table_name: &str,
@@ -436,6 +440,23 @@ impl DatabricksAdapter {
         ))
     }
 
+    /// Build a CTAS statement that creates the table by reading parquet files
+    /// from S3.
+    ///
+    /// ```sql
+    /// CREATE OR REPLACE TABLE catalog.schema.table
+    ///   AS SELECT * FROM parquet.`s3://bucket/prefix/scenario/version/tables/table/`
+    /// ```
+    fn create_table_ctas(&self, table_name: &str) -> String {
+        let source_uri = format!(
+            "s3://spiceai-public-datasets/data-gen/tpch/4/tables/{table_name}/"
+        );
+        format!(
+            "CREATE OR REPLACE TABLE {} AS SELECT * FROM parquet.`{source_uri}`",
+            self.lakebase_table_full_name(table_name),
+        )
+    }
+
     fn table_format_from_setup_metadata(
         &self,
         metadata: &HashMap<String, Value>,
@@ -1084,19 +1105,11 @@ impl Handler for DatabricksAdapter {
             },
         );
 
-        let table_format = {
-            let state = self
-                .runs
-                .get(&run_id)
-                .ok_or_else(|| format!("Unknown run_id: {run_id}"))?;
-            state.table_format
-        };
-
         let mut created_tables = Vec::with_capacity(datasets.len());
 
         match self.config.variant {
             DatabricksVariant::Databricks | DatabricksVariant::Lakebase => {
-                for (table_name, dataset_cfg) in datasets {
+                for (table_name, _dataset_cfg) in datasets {
                     let drop_sql = format!(
                         "DROP TABLE IF EXISTS {}",
                         self.lakebase_table_full_name(&table_name)
@@ -1107,11 +1120,7 @@ impl Handler for DatabricksAdapter {
                         )
                     })?;
 
-                    let create_sql = self
-                        .create_table_ddl(&table_name, &dataset_cfg, table_format)
-                        .map_err(|e| {
-                            format!("Failed to build Lakebase table DDL for '{table_name}': {e}")
-                        })?;
+                    let create_sql = self.create_table_ctas(&table_name);
 
                     eprintln!("[databricks-adapter] create_table '{table_name}': {create_sql}");