use super::*; pub async fn export_sqlite_core_jsonl( pool: &crate::driver::sqlite::SqlitePool, created_at_unix_secs: u64, ) -> Result { export_sqlite_jsonl(pool, sqlite_core_export_domains(), created_at_unix_secs).await } pub async fn export_sqlite_jsonl( pool: &crate::driver::sqlite::SqlitePool, domains: Vec, created_at_unix_secs: u64, ) -> Result { let manifest = DataExportManifest::new( created_at_unix_secs, Some(DatabaseDriver::Sqlite), domains.clone(), ); let mut records = vec![DataExportRecord::manifest(manifest)]; for domain in domains { if domain == ExportDomain::Billing { export_sqlite_billing_records(pool, &mut records).await?; continue; } if domain == ExportDomain::Wallets { export_sqlite_wallet_records(pool, &mut records).await?; continue; } let (table_name, id_column) = sqlite_domain_table(domain)?; let order_by = export_order_by(domain, id_column); let sql = format!("SELECT * FROM {table_name} ORDER BY {order_by}"); let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?; for row in rows { let id = sqlite_export_row_id(domain, &row, id_column)?; records.push(DataExportRecord::row(domain, id, sqlite_row_payload(&row)?)); } } encode_jsonl(&records) } pub async fn import_sqlite_jsonl( pool: &crate::driver::sqlite::SqlitePool, input: &str, ) -> Result { let plan = build_import_plan(input)?; import_sqlite_plan(pool, &plan).await } pub async fn import_sqlite_plan( pool: &crate::driver::sqlite::SqlitePool, plan: &DataImportPlan, ) -> Result { let mut imported = 0usize; let mut column_cache = BTreeMap::::new(); for domain in &plan.manifest.domains { if *domain == ExportDomain::Billing { for row in plan.rows(*domain) { 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, &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, &target_columns).await?; imported = imported.saturating_add(1); } } Ok(imported) } fn sqlite_domain_table( domain: ExportDomain, ) -> Result<(&'static str, &'static str), DataLayerError> { match domain { ExportDomain::Users => Ok(("users", "id")), ExportDomain::ApiKeys => Ok(("api_keys", "id")), ExportDomain::Providers => Ok(("providers", "id")), ExportDomain::ProviderKeys => Ok(("provider_api_keys", "id")), ExportDomain::Endpoints => Ok(("provider_endpoints", "id")), ExportDomain::Models => Ok(("models", "id")), ExportDomain::GlobalModels => Ok(("global_models", "id")), ExportDomain::AuthModules => Ok(("auth_modules", "id")), ExportDomain::OAuthProviders => Ok(("oauth_providers", "provider_type")), ExportDomain::UserOAuthLinks => Ok(("user_oauth_links", "id")), ExportDomain::UserGroups => Ok(("user_groups", "id")), ExportDomain::UserGroupMembers => Ok(("user_group_members", "group_id")), ExportDomain::ProxyNodes => Ok(("proxy_nodes", "id")), ExportDomain::SystemConfigs => Ok(("system_configs", "id")), ExportDomain::Wallets => Err(DataLayerError::InvalidInput( "sqlite wallet export uses multiple tables and must be handled as a domain".to_string(), )), ExportDomain::Usage => Ok((r#""usage""#, "request_id")), ExportDomain::Billing => Err(DataLayerError::InvalidInput( "sqlite billing export uses multiple tables and must be handled as a domain" .to_string(), )), } } fn sqlite_export_row_id( domain: ExportDomain, row: &sqlx::sqlite::SqliteRow, id_column: &str, ) -> Result { if domain == ExportDomain::UserGroupMembers { let group_id = sqlite_required_export_text(row, "group_id", domain)?; let user_id = sqlite_required_export_text(row, "user_id", domain)?; return Ok(format!("{group_id}:{user_id}")); } sqlite_required_export_text(row, id_column, domain) } fn sqlite_required_export_text( row: &sqlx::sqlite::SqliteRow, column: &str, domain: ExportDomain, ) -> Result { row.try_get::, _>(column) .map_sql_err()? .ok_or_else(|| { DataLayerError::UnexpectedValue(format!( "{} export row has null id column '{}'", domain.as_str(), column )) }) } async fn export_sqlite_billing_records( pool: &crate::driver::sqlite::SqlitePool, records: &mut Vec, ) -> Result<(), DataLayerError> { for table_name in [ "billing_rules", "dimension_collectors", "usage_settlement_snapshots", ] { let id_column = if table_name == "usage_settlement_snapshots" { "request_id" } else { "id" }; let sql = format!("SELECT * FROM {table_name} ORDER BY {id_column} ASC"); let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?; for row in rows { let id = row .try_get::, _>(id_column) .map_sql_err()? .ok_or_else(|| { DataLayerError::UnexpectedValue(format!( "billing export row in table '{table_name}' has null id" )) })?; records.push(DataExportRecord::row( ExportDomain::Billing, format!("{table_name}:{id}"), payload_with_table(sqlite_row_payload(&row)?, table_name)?, )); } } Ok(()) } async fn export_sqlite_wallet_records( pool: &crate::driver::sqlite::SqlitePool, records: &mut Vec, ) -> Result<(), DataLayerError> { for (table_name, id_column) in sqlite_wallet_tables() { let sql = format!("SELECT * FROM {table_name} ORDER BY {id_column} ASC"); let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?; for row in rows { let id = row .try_get::, _>(id_column) .map_sql_err()? .ok_or_else(|| { DataLayerError::UnexpectedValue(format!( "wallet export row in table '{table_name}' has null id" )) })?; records.push(DataExportRecord::row( ExportDomain::Wallets, format!("{table_name}:{id}"), payload_with_table(sqlite_row_payload(&row)?, table_name)?, )); } } Ok(()) } async fn import_sqlite_row( pool: &crate::driver::sqlite::SqlitePool, table_name: &str, domain: ExportDomain, row: &ExportRow, target_columns: &ImportColumnNames, ) -> Result<(), DataLayerError> { let object = filter_import_payload("sqlite", table_name, domain, row, target_columns)?; let columns = object.keys().map(String::as_str).collect::>(); let column_sql = columns .iter() .map(|column| sqlite_quote_identifier(column)) .collect::, _>>()? .join(", "); let placeholder_sql = vec!["?"; columns.len()].join(", "); let sql = format!("INSERT OR REPLACE INTO {table_name} ({column_sql}) VALUES ({placeholder_sql})"); let mut query = sqlx::query(&sql); for column in columns { let value = object .get(column) .expect("column name came from payload object keys"); query = bind_sqlite_json_value(query, value)?; } query.execute(pool).await.map_sql_err()?; Ok(()) } async fn import_sqlite_billing_row( pool: &crate::driver::sqlite::SqlitePool, row: &ExportRow, column_cache: &mut BTreeMap, ) -> 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, table_name, ExportDomain::Billing, &ExportRow { id: row.id.clone(), payload, }, &target_columns, ) .await } fn sqlite_billing_table_name(table_name: &str) -> Result<&'static str, DataLayerError> { match table_name { "billing_rules" => Ok("billing_rules"), "dimension_collectors" => Ok("dimension_collectors"), "usage_settlement_snapshots" => Ok("usage_settlement_snapshots"), other => Err(DataLayerError::InvalidInput(format!( "unsupported sqlite billing export table '{other}'" ))), } } async fn import_sqlite_wallet_row( pool: &crate::driver::sqlite::SqlitePool, row: &ExportRow, column_cache: &mut BTreeMap, ) -> 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, table_name, ExportDomain::Wallets, &ExportRow { id: row.id.clone(), payload, }, &target_columns, ) .await } fn sqlite_wallet_tables() -> &'static [(&'static str, &'static str)] { &[ ("wallets", "id"), ("wallet_transactions", "id"), ("wallet_daily_usage_ledgers", "id"), ("payment_orders", "id"), ("payment_callbacks", "id"), ("refund_requests", "id"), ("redeem_code_batches", "id"), ("redeem_codes", "id"), ] } fn sqlite_wallet_table_name(table_name: &str) -> Result<&'static str, DataLayerError> { sqlite_wallet_tables() .iter() .find(|(candidate, _)| *candidate == table_name) .map(|(table, _)| *table) .ok_or_else(|| { DataLayerError::InvalidInput(format!( "unsupported sqlite wallet export table '{table_name}'" )) }) } async fn sqlite_import_columns_cached( pool: &crate::driver::sqlite::SqlitePool, cache: &mut BTreeMap, table_name: &str, ) -> Result { 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 { 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::("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_row_payload(row: &sqlx::sqlite::SqliteRow) -> Result { let mut object = serde_json::Map::new(); for (index, column) in row.columns().iter().enumerate() { object.insert( column.name().to_string(), sqlite_value_to_json(row, index, column.name())?, ); } Ok(Value::Object(object)) } fn sqlite_value_to_json( row: &sqlx::sqlite::SqliteRow, index: usize, column_name: &str, ) -> Result { let raw = row.try_get_raw(index).map_sql_err()?; if raw.is_null() { return Ok(Value::Null); } match raw.type_info().name().to_ascii_uppercase().as_str() { "INTEGER" => { let value = row.try_get::(index).map_sql_err()?; if sqlite_integer_column_is_boolean(column_name) { match value { 0 => return Ok(Value::Bool(false)), 1 => return Ok(Value::Bool(true)), _ => {} } } Ok(Value::from(value)) } "REAL" | "FLOAT" | "DOUBLE" => { let value = row.try_get::(index).map_sql_err()?; serde_json::Number::from_f64(value) .map(Value::Number) .ok_or_else(|| { DataLayerError::UnexpectedValue(format!( "sqlite export column {} contains non-finite float", index )) }) } "TEXT" => Ok(Value::String( row.try_get::(index).map_sql_err()?, )), "BLOB" => { let bytes = row.try_get::, _>(index).map_sql_err()?; Ok(Value::Array(bytes.into_iter().map(Value::from).collect())) } other => Err(DataLayerError::UnexpectedValue(format!( "unsupported sqlite export column type '{other}' at index {index}" ))), } } fn sqlite_integer_column_is_boolean(column_name: &str) -> bool { column_name.starts_with("is_") || column_name.starts_with("has_") || column_name.starts_with("supports_") || column_name.starts_with("enable_") || column_name.starts_with("use_") || matches!( column_name, "announcement_notifications" | "auto_delete_on_expiry" | "auto_fetch_models" | "email_notifications" | "email_verified" | "format_converted" | "keep_priority_on_conversion" | "signature_valid" | "tunnel_connected" | "tunnel_mode" | "usage_alerts" | "webhook_sent" ) }