i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix postgres export smoke import compatibility

Browse Source
This commit is contained in:
fawney19
2026-05-06 09:23:57 +08:00
parent 887a58d639
commit beee7a76d2
1 changed files with 169 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

219
crates/aether-data/src/lifecycle/export.rs
View File

@@ -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(|| {
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 object = filter_import_payload("sqlite", table_name, domain, row, target_columns)?;
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(|| {
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 object = filter_import_payload("mysql", table_name, domain, row, target_columns)?;
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)
}
}