Fix postgres export smoke import compatibility

crates/aether-data/src/lifecycle/export.rs

@@ -135,6 +135,7 @@ struct PostgresImportColumn {
 }
 
 type PostgresImportColumns = BTreeMap<String, PostgresImportColumn>;
+type ImportColumnNames = BTreeSet<String>;
 
 pub fn encode_jsonl(records: &[DataExportRecord]) -> Result<String, DataLayerError> {
     validate_export_records(records)?;
@@ -399,24 +400,27 @@ pub async fn import_sqlite_plan(
     plan: &DataImportPlan,
 ) -> Result<usize, DataLayerError> {
     let mut imported = 0usize;
+    let mut column_cache = BTreeMap::<String, ImportColumnNames>::new();
     for domain in &plan.manifest.domains {
         if *domain == ExportDomain::Billing {
             for row in plan.rows(*domain) {
-                import_sqlite_billing_row(pool, row).await?;
+                import_sqlite_billing_row(pool, row, &mut column_cache).await?;
                 imported = imported.saturating_add(1);
             }
             continue;
         }
         if *domain == ExportDomain::Wallets {
             for row in plan.rows(*domain) {
-                import_sqlite_wallet_row(pool, row).await?;
+                import_sqlite_wallet_row(pool, row, &mut column_cache).await?;
                 imported = imported.saturating_add(1);
             }
             continue;
         }
         let (table_name, _id_column) = sqlite_domain_table(*domain)?;
+        let target_columns =
+            sqlite_import_columns_cached(pool, &mut column_cache, table_name).await?;
         for row in plan.rows(*domain) {
-            import_sqlite_row(pool, table_name, *domain, row).await?;
+            import_sqlite_row(pool, table_name, *domain, row, &target_columns).await?;
             imported = imported.saturating_add(1);
         }
     }
@@ -572,24 +576,27 @@ pub async fn import_mysql_plan(
     plan: &DataImportPlan,
 ) -> Result<usize, DataLayerError> {
     let mut imported = 0usize;
+    let mut column_cache = BTreeMap::<String, ImportColumnNames>::new();
     for domain in &plan.manifest.domains {
         if *domain == ExportDomain::Billing {
             for row in plan.rows(*domain) {
-                import_mysql_billing_row(pool, row).await?;
+                import_mysql_billing_row(pool, row, &mut column_cache).await?;
                 imported = imported.saturating_add(1);
             }
             continue;
         }
         if *domain == ExportDomain::Wallets {
             for row in plan.rows(*domain) {
-                import_mysql_wallet_row(pool, row).await?;
+                import_mysql_wallet_row(pool, row, &mut column_cache).await?;
                 imported = imported.saturating_add(1);
             }
             continue;
         }
         let (table_name, _id_column) = mysql_domain_table(*domain)?;
+        let target_columns =
+            mysql_import_columns_cached(pool, &mut column_cache, table_name).await?;
         for row in plan.rows(*domain) {
-            import_mysql_row(pool, table_name, *domain, row).await?;
+            import_mysql_row(pool, table_name, *domain, row, &target_columns).await?;
             imported = imported.saturating_add(1);
         }
     }
@@ -689,21 +696,9 @@ async fn import_sqlite_row(
     table_name: &str,
     domain: ExportDomain,
     row: &ExportRow,
+    target_columns: &ImportColumnNames,
 ) -> Result<(), DataLayerError> {
-    let object = row.payload.as_object().ok_or_else(|| {
+    let object = filter_import_payload("sqlite", table_name, domain, row, target_columns)?;
-        DataLayerError::InvalidInput(format!(
-            "{} export row '{}' payload must be a JSON object",
-            domain.as_str(),
-            row.id
-        ))
-    })?;
-    if object.is_empty() {
-        return Err(DataLayerError::InvalidInput(format!(
-            "{} export row '{}' payload cannot be empty",
-            domain.as_str(),
-            row.id
-        )));
-    }
 
     let columns = object.keys().map(String::as_str).collect::<Vec<_>>();
     let column_sql = columns
@@ -728,16 +723,20 @@ async fn import_sqlite_row(
 async fn import_sqlite_billing_row(
     pool: &crate::driver::sqlite::SqlitePool,
     row: &ExportRow,
+    column_cache: &mut BTreeMap<String, ImportColumnNames>,
 ) -> Result<(), DataLayerError> {
     let (table_name, payload) = billing_payload_table(row)?;
+    let table_name = sqlite_billing_table_name(&table_name)?;
+    let target_columns = sqlite_import_columns_cached(pool, column_cache, table_name).await?;
     import_sqlite_row(
         pool,
-        sqlite_billing_table_name(&table_name)?,
+        table_name,
         ExportDomain::Billing,
         &ExportRow {
             id: row.id.clone(),
             payload,
         },
+        &target_columns,
     )
     .await
 }
@@ -756,16 +755,20 @@ fn sqlite_billing_table_name(table_name: &str) -> Result<&'static str, DataLayer
 async fn import_sqlite_wallet_row(
     pool: &crate::driver::sqlite::SqlitePool,
     row: &ExportRow,
+    column_cache: &mut BTreeMap<String, ImportColumnNames>,
 ) -> Result<(), DataLayerError> {
     let (table_name, payload) = domain_payload_table(row, "wallet", Some("wallets"))?;
+    let table_name = sqlite_wallet_table_name(&table_name)?;
+    let target_columns = sqlite_import_columns_cached(pool, column_cache, table_name).await?;
     import_sqlite_row(
         pool,
-        sqlite_wallet_table_name(&table_name)?,
+        table_name,
         ExportDomain::Wallets,
         &ExportRow {
             id: row.id.clone(),
             payload,
         },
+        &target_columns,
     )
     .await
 }
@@ -795,6 +798,40 @@ fn sqlite_wallet_table_name(table_name: &str) -> Result<&'static str, DataLayerE
         })
 }
 
+async fn sqlite_import_columns_cached(
+    pool: &crate::driver::sqlite::SqlitePool,
+    cache: &mut BTreeMap<String, ImportColumnNames>,
+    table_name: &str,
+) -> Result<ImportColumnNames, DataLayerError> {
+    if let Some(columns) = cache.get(table_name) {
+        return Ok(columns.clone());
+    }
+
+    let columns = load_sqlite_import_columns(pool, table_name).await?;
+    cache.insert(table_name.to_string(), columns.clone());
+    Ok(columns)
+}
+
+async fn load_sqlite_import_columns(
+    pool: &crate::driver::sqlite::SqlitePool,
+    table_name: &str,
+) -> Result<ImportColumnNames, DataLayerError> {
+    let sql = format!("PRAGMA table_info({table_name})");
+    let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?;
+    let mut columns = ImportColumnNames::new();
+    for row in rows {
+        columns.insert(row.try_get::<String, _>("name").map_sql_err()?);
+    }
+
+    if columns.is_empty() {
+        return Err(DataLayerError::UnexpectedValue(format!(
+            "sqlite import target table '{table_name}' has no visible columns"
+        )));
+    }
+
+    Ok(columns)
+}
+
 fn sqlite_quote_identifier(identifier: &str) -> Result<String, DataLayerError> {
     if identifier.trim().is_empty() {
         return Err(DataLayerError::InvalidInput(
@@ -1418,21 +1455,9 @@ async fn import_mysql_row(
     table_name: &str,
     domain: ExportDomain,
     row: &ExportRow,
+    target_columns: &ImportColumnNames,
 ) -> Result<(), DataLayerError> {
-    let object = row.payload.as_object().ok_or_else(|| {
+    let object = filter_import_payload("mysql", table_name, domain, row, target_columns)?;
-        DataLayerError::InvalidInput(format!(
-            "{} export row '{}' payload must be a JSON object",
-            domain.as_str(),
-            row.id
-        ))
-    })?;
-    if object.is_empty() {
-        return Err(DataLayerError::InvalidInput(format!(
-            "{} export row '{}' payload cannot be empty",
-            domain.as_str(),
-            row.id
-        )));
-    }
 
     let columns = object.keys().map(String::as_str).collect::<Vec<_>>();
     let column_sql = columns
@@ -1466,16 +1491,20 @@ async fn import_mysql_row(
 async fn import_mysql_billing_row(
     pool: &crate::driver::mysql::MysqlPool,
     row: &ExportRow,
+    column_cache: &mut BTreeMap<String, ImportColumnNames>,
 ) -> Result<(), DataLayerError> {
     let (table_name, payload) = billing_payload_table(row)?;
+    let table_name = mysql_billing_table_name(&table_name)?;
+    let target_columns = mysql_import_columns_cached(pool, column_cache, table_name).await?;
     import_mysql_row(
         pool,
-        mysql_billing_table_name(&table_name)?,
+        table_name,
         ExportDomain::Billing,
         &ExportRow {
             id: row.id.clone(),
             payload,
         },
+        &target_columns,
     )
     .await
 }
@@ -1494,16 +1523,20 @@ fn mysql_billing_table_name(table_name: &str) -> Result<&'static str, DataLayerE
 async fn import_mysql_wallet_row(
     pool: &crate::driver::mysql::MysqlPool,
     row: &ExportRow,
+    column_cache: &mut BTreeMap<String, ImportColumnNames>,
 ) -> Result<(), DataLayerError> {
     let (table_name, payload) = domain_payload_table(row, "wallet", Some("wallets"))?;
+    let table_name = mysql_wallet_table_name(&table_name)?;
+    let target_columns = mysql_import_columns_cached(pool, column_cache, table_name).await?;
     import_mysql_row(
         pool,
-        mysql_wallet_table_name(&table_name)?,
+        table_name,
         ExportDomain::Wallets,
         &ExportRow {
             id: row.id.clone(),
             payload,
         },
+        &target_columns,
     )
     .await
 }
@@ -1533,6 +1566,52 @@ fn mysql_wallet_table_name(table_name: &str) -> Result<&'static str, DataLayerEr
         })
 }
 
+async fn mysql_import_columns_cached(
+    pool: &crate::driver::mysql::MysqlPool,
+    cache: &mut BTreeMap<String, ImportColumnNames>,
+    table_name: &str,
+) -> Result<ImportColumnNames, DataLayerError> {
+    if let Some(columns) = cache.get(table_name) {
+        return Ok(columns.clone());
+    }
+
+    let columns = load_mysql_import_columns(pool, table_name).await?;
+    cache.insert(table_name.to_string(), columns.clone());
+    Ok(columns)
+}
+
+async fn load_mysql_import_columns(
+    pool: &crate::driver::mysql::MysqlPool,
+    table_name: &str,
+) -> Result<ImportColumnNames, DataLayerError> {
+    let relation_name = table_name.trim_matches('`');
+    let rows = sqlx::query(
+        r#"
+SELECT COLUMN_NAME AS column_name
+FROM information_schema.columns
+WHERE table_schema = DATABASE()
+  AND table_name = ?
+"#,
+    )
+    .bind(relation_name)
+    .fetch_all(pool)
+    .await
+    .map_sql_err()?;
+
+    let mut columns = ImportColumnNames::new();
+    for row in rows {
+        columns.insert(row.try_get::<String, _>("column_name").map_sql_err()?);
+    }
+
+    if columns.is_empty() {
+        return Err(DataLayerError::UnexpectedValue(format!(
+            "mysql import target table '{table_name}' has no visible columns"
+        )));
+    }
+
+    Ok(columns)
+}
+
 fn mysql_quote_identifier(identifier: &str) -> Result<String, DataLayerError> {
     if identifier.trim().is_empty() {
         return Err(DataLayerError::InvalidInput(
@@ -1584,6 +1663,46 @@ fn bind_mysql_json_value<'q>(
     })
 }
 
+fn filter_import_payload(
+    driver_name: &str,
+    table_name: &str,
+    domain: ExportDomain,
+    row: &ExportRow,
+    target_columns: &ImportColumnNames,
+) -> Result<serde_json::Map<String, Value>, DataLayerError> {
+    let object = row.payload.as_object().ok_or_else(|| {
+        DataLayerError::InvalidInput(format!(
+            "{} export row '{}' payload must be a JSON object",
+            domain.as_str(),
+            row.id
+        ))
+    })?;
+    if object.is_empty() {
+        return Err(DataLayerError::InvalidInput(format!(
+            "{} export row '{}' payload cannot be empty",
+            domain.as_str(),
+            row.id
+        )));
+    }
+
+    let mut filtered = serde_json::Map::new();
+    for (column_name, value) in object {
+        if target_columns.contains(column_name) {
+            filtered.insert(column_name.clone(), value.clone());
+        }
+    }
+
+    if filtered.is_empty() {
+        return Err(DataLayerError::InvalidInput(format!(
+            "{} export row '{}' has no columns supported by {driver_name} table '{table_name}'",
+            domain.as_str(),
+            row.id
+        )));
+    }
+
+    Ok(filtered)
+}
+
 fn payload_with_table(payload: Value, table_name: &str) -> Result<Value, DataLayerError> {
     let mut object = payload.as_object().cloned().ok_or_else(|| {
         DataLayerError::UnexpectedValue("export row payload must be a JSON object".to_string())
@@ -2152,7 +2271,7 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet
         let request_id = format!("export-request-{suffix}");
 
         sqlx::query(
-            "INSERT INTO users (id, email, username, auth_source, created_at, updated_at) VALUES ($1, $2, $3, 'local', 1, 2)",
+            "INSERT INTO users (id, email, username, auth_source, email_verified, created_at, updated_at) VALUES ($1, $2, $3, 'local', TRUE, to_timestamp(1), to_timestamp(2))",
         )
         .bind(&user_id)
         .bind(format!("{user_id}@example.com"))
@@ -2161,7 +2280,7 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet
         .await
         .expect("user should seed");
         sqlx::query(
-            "INSERT INTO api_keys (id, user_id, key_hash, key_encrypted, name, created_at, updated_at) VALUES ($1, $2, $3, 'ciphertext-1', 'Default', 1, 2)",
+            "INSERT INTO api_keys (id, user_id, key_hash, key_encrypted, name, created_at, updated_at) VALUES ($1, $2, $3, 'ciphertext-1', 'Default', to_timestamp(1), to_timestamp(2))",
         )
         .bind(&api_key_id)
         .bind(&user_id)
@@ -2170,7 +2289,7 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet
         .await
         .expect("api key should seed");
         sqlx::query(
-            "INSERT INTO providers (id, name, provider_type, created_at, updated_at) VALUES ($1, $2, 'openai', 1, 2)",
+            "INSERT INTO providers (id, name, provider_type, created_at, updated_at) VALUES ($1, $2, 'openai', to_timestamp(1), to_timestamp(2))",
         )
         .bind(&provider_id)
         .bind(format!("Provider {suffix}"))
@@ -2178,7 +2297,7 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet
         .await
         .expect("provider should seed");
         sqlx::query(
-            "INSERT INTO provider_api_keys (id, provider_id, name, encrypted_key, created_at, updated_at) VALUES ($1, $2, 'Provider Key', 'ciphertext-provider', 1, 2)",
+            "INSERT INTO provider_api_keys (id, provider_id, name, encrypted_key, total_tokens, total_cost_usd, created_at, updated_at) VALUES ($1, $2, 'Provider Key', 'ciphertext-provider', 0, 0, to_timestamp(1), to_timestamp(2))",
         )
         .bind(&provider_key_id)
         .bind(&provider_id)
@@ -2186,7 +2305,7 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet
         .await
         .expect("provider key should seed");
         sqlx::query(
-            "INSERT INTO provider_endpoints (id, provider_id, name, base_url, created_at, updated_at) VALUES ($1, $2, 'Primary', 'https://example.test', 1, 2)",
+            "INSERT INTO provider_endpoints (id, provider_id, name, base_url, created_at, updated_at) VALUES ($1, $2, 'Primary', 'https://example.test', to_timestamp(1), to_timestamp(2))",
         )
         .bind(&endpoint_id)
         .bind(&provider_id)
@@ -2194,7 +2313,7 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet
         .await
         .expect("endpoint should seed");
         sqlx::query(
-            "INSERT INTO global_models (id, name, created_at, updated_at) VALUES ($1, $2, 1, 2)",
+            "INSERT INTO global_models (id, name, created_at, updated_at) VALUES ($1, $2, to_timestamp(1), to_timestamp(2))",
         )
         .bind(&global_model_id)
         .bind(format!("global-model-{suffix}"))
@@ -2202,7 +2321,7 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet
         .await
         .expect("global model should seed");
         sqlx::query(
-            "INSERT INTO models (id, provider_id, global_model_id, provider_model_name, created_at, updated_at) VALUES ($1, $2, $3, 'provider-model', 1, 2)",
+            "INSERT INTO models (id, provider_id, global_model_id, provider_model_name, created_at, updated_at) VALUES ($1, $2, $3, 'provider-model', to_timestamp(1), to_timestamp(2))",
         )
         .bind(&model_id)
         .bind(&provider_id)
@@ -2211,7 +2330,7 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet
         .await
         .expect("model should seed");
         sqlx::query(
-            "INSERT INTO billing_rules (id, global_model_id, name, task_type, expression, variables, dimension_mappings, is_enabled, created_at, updated_at) VALUES ($1, $2, 'Rule One', 'chat', 'input_tokens * 0.01', '{}', '{\"input\":\"input_tokens\"}', TRUE, 1, 2)",
+            "INSERT INTO billing_rules (id, global_model_id, name, task_type, expression, variables, dimension_mappings, is_enabled, created_at, updated_at) VALUES ($1, $2, 'Rule One', 'chat', 'input_tokens * 0.01', '{}', '{\"input\":\"input_tokens\"}', TRUE, to_timestamp(1), to_timestamp(2))",
         )
         .bind("billing-rule-1")
         .bind(&global_model_id)
@@ -2219,21 +2338,21 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet
         .await
         .expect("billing rule should seed");
         sqlx::query(
-            "INSERT INTO dimension_collectors (id, api_format, task_type, dimension_name, source_type, value_type, transform_expression, priority, is_enabled, created_at, updated_at) VALUES ($1, 'openai', 'chat', 'input_tokens', 'computed', 'float', 'usage.input_tokens', 10, TRUE, 1, 2)",
+            "INSERT INTO dimension_collectors (id, api_format, task_type, dimension_name, source_type, value_type, transform_expression, priority, is_enabled, created_at, updated_at) VALUES ($1, 'openai', 'chat', 'input_tokens', 'computed', 'float', 'usage.input_tokens', 10, TRUE, to_timestamp(1), to_timestamp(2))",
         )
         .bind("collector-1")
         .execute(&pool)
         .await
         .expect("dimension collector should seed");
         sqlx::query(
-            "INSERT INTO system_configs (id, key, value, created_at, updated_at) VALUES ($1, 'billing.enabled', 'true', 1, 2)",
+            "INSERT INTO system_configs (id, key, value, created_at, updated_at) VALUES ($1, 'billing.enabled', 'true', to_timestamp(1), to_timestamp(2))",
         )
         .bind(&config_id)
         .execute(&pool)
         .await
         .expect("system config should seed");
         sqlx::query(
-            "INSERT INTO wallets (id, user_id, created_at, updated_at) VALUES ($1, $2, 1, 2)",
+            "INSERT INTO wallets (id, user_id, created_at, updated_at) VALUES ($1, $2, to_timestamp(1), to_timestamp(2))",
         )
         .bind(&wallet_id)
         .bind(&user_id)
@@ -2487,6 +2606,6 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet
             .duration_since(std::time::UNIX_EPOCH)
             .unwrap_or_default()
             .as_nanos();
-        format!("{}-{nanos}", std::process::id())
+        format!("{:013x}", nanos & 0x1fff_ffff_fffff)
     }
 }