Add multi-database data layer

Introduce aether-data-schema and driver-specific schema generation for Postgres, MySQL, and SQLite.

Split data backends, lifecycle, repositories, and gateway runtime integration across database drivers.

Verified with cargo fmt --all --check, cargo clippy --workspace --all-targets -- -D warnings, and cargo test --workspace.
This commit is contained in:
fawney19
2026-05-05 18:27:36 +08:00
parent 099653f732
commit fce7e959e5
372 changed files with 86217 additions and 21160 deletions

View File

@@ -1,5 +1,5 @@
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::{query_param_value, unix_secs_to_rfc3339};
use crate::handlers::admin::shared::query_param_value;
use crate::GatewayError;
use axum::{
body::{Body, Bytes},
@@ -9,7 +9,6 @@ use axum::{
};
use regex::Regex;
use serde_json::json;
use sqlx::Row;
const ADMIN_BILLING_DATA_UNAVAILABLE_DETAIL: &str = "Admin billing data unavailable";
@@ -185,20 +184,6 @@ fn admin_billing_validate_safe_expression(expression: &str) -> Result<(), String
Ok(())
}
fn admin_billing_optional_epoch_value(
row: &sqlx::postgres::PgRow,
field: &str,
) -> Result<Option<String>, GatewayError> {
let value = row
.try_get::<Option<i64>, _>(field)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
match value {
None => Ok(None),
Some(value) if value < 0 => Ok(None),
Some(value) => Ok(unix_secs_to_rfc3339(value as u64)),
}
}
pub(crate) async fn maybe_build_local_admin_billing_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,

View File

@@ -1,7 +1,6 @@
use super::{
build_admin_payment_callback_payload, build_admin_payment_callback_payload_from_record,
build_admin_payments_bad_request_response, parse_admin_payments_limit,
parse_admin_payments_offset,
build_admin_payment_callback_payload_from_record, build_admin_payments_bad_request_response,
parse_admin_payments_limit, parse_admin_payments_offset,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::query_param_value;

View File

@@ -12,13 +12,13 @@ mod shared;
use self::shared::{
admin_payment_operator_id, admin_payment_order_id_from_detail_path,
admin_payment_order_id_from_suffix_path, build_admin_payment_callback_payload,
build_admin_payment_callback_payload_from_record, build_admin_payment_order_not_found_response,
build_admin_payment_order_payload, build_admin_payment_orders_page_response,
build_admin_payments_backend_unavailable_response, build_admin_payments_bad_request_response,
build_admin_payments_data_unavailable_response, normalize_admin_payment_currency,
normalize_admin_payment_optional_string, normalize_admin_payment_positive_number,
parse_admin_payments_limit, parse_admin_payments_offset, AdminPaymentOrderCreditRequest,
admin_payment_order_id_from_suffix_path, build_admin_payment_callback_payload_from_record,
build_admin_payment_order_not_found_response, build_admin_payment_order_payload,
build_admin_payment_orders_page_response, build_admin_payments_backend_unavailable_response,
build_admin_payments_bad_request_response, build_admin_payments_data_unavailable_response,
normalize_admin_payment_currency, normalize_admin_payment_optional_string,
normalize_admin_payment_positive_number, parse_admin_payments_limit,
parse_admin_payments_offset, AdminPaymentOrderCreditRequest,
};
pub(crate) async fn maybe_build_local_admin_payments_response(

View File

@@ -1,6 +1,6 @@
use crate::handlers::admin::request::AdminRequestContext;
use crate::handlers::admin::shared::{query_param_value, unix_secs_to_rfc3339};
use crate::{GatewayAdminPaymentCallbackView, GatewayError};
use crate::GatewayAdminPaymentCallbackView;
use axum::{
body::Body,
http,
@@ -8,7 +8,6 @@ use axum::{
Json,
};
use serde_json::json;
use sqlx::Row;
const ADMIN_PAYMENTS_DATA_UNAVAILABLE_DETAIL: &str = "Admin payments data unavailable";
@@ -220,34 +219,6 @@ pub(super) fn build_admin_payment_order_payload(
})
}
pub(super) fn build_admin_payment_callback_payload(
row: &sqlx::postgres::PgRow,
) -> Result<serde_json::Value, GatewayError> {
Ok(json!({
"id": row.try_get::<String, _>("id").map_err(|err| GatewayError::Internal(err.to_string()))?,
"payment_order_id": row.try_get::<Option<String>, _>("payment_order_id").map_err(|err| GatewayError::Internal(err.to_string()))?,
"payment_method": row.try_get::<String, _>("payment_method").map_err(|err| GatewayError::Internal(err.to_string()))?,
"callback_key": row.try_get::<String, _>("callback_key").map_err(|err| GatewayError::Internal(err.to_string()))?,
"order_no": row.try_get::<Option<String>, _>("order_no").map_err(|err| GatewayError::Internal(err.to_string()))?,
"gateway_order_id": row.try_get::<Option<String>, _>("gateway_order_id").map_err(|err| GatewayError::Internal(err.to_string()))?,
"payload_hash": row.try_get::<Option<String>, _>("payload_hash").map_err(|err| GatewayError::Internal(err.to_string()))?,
"signature_valid": row.try_get::<bool, _>("signature_valid").map_err(|err| GatewayError::Internal(err.to_string()))?,
"status": row.try_get::<String, _>("status").map_err(|err| GatewayError::Internal(err.to_string()))?,
"payload": row.try_get::<Option<serde_json::Value>, _>("payload").map_err(|err| GatewayError::Internal(err.to_string()))?,
"error_message": row.try_get::<Option<String>, _>("error_message").map_err(|err| GatewayError::Internal(err.to_string()))?,
"created_at": row
.try_get::<Option<i64>, _>("created_at_unix_ms")
.map_err(|err| GatewayError::Internal(err.to_string()))?
.and_then(|value| u64::try_from(value).ok())
.and_then(unix_secs_to_rfc3339),
"processed_at": row
.try_get::<Option<i64>, _>("processed_at_unix_secs")
.map_err(|err| GatewayError::Internal(err.to_string()))?
.and_then(|value| u64::try_from(value).ok())
.and_then(unix_secs_to_rfc3339),
}))
}
pub(super) fn build_admin_payment_callback_payload_from_record(
record: &GatewayAdminPaymentCallbackView,
) -> serde_json::Value {

View File

@@ -61,7 +61,7 @@ pub(in super::super) async fn build_admin_wallet_adjust_response(
));
}
let operator_id = admin_wallet_operator_id(request_context);
let has_postgres = state.has_postgres_pool();
let has_wallet_writer = state.has_wallet_data_writer();
let Some((wallet, transaction)) = state
.admin_adjust_wallet_balance(
&wallet_id,
@@ -72,7 +72,7 @@ pub(in super::super) async fn build_admin_wallet_adjust_response(
)
.await?
else {
return if has_postgres {
return if has_wallet_writer {
Ok(build_admin_wallet_not_found_response())
} else {
Ok(build_admin_wallets_data_unavailable_response())

View File

@@ -61,7 +61,7 @@ pub(in super::super) async fn build_admin_wallet_recharge_response(
));
}
let operator_id = admin_wallet_operator_id(request_context);
let has_postgres = state.has_postgres_pool();
let has_wallet_writer = state.has_wallet_data_writer();
let Some((wallet, payment_order)) = state
.admin_create_manual_wallet_recharge(
&wallet_id,
@@ -72,7 +72,7 @@ pub(in super::super) async fn build_admin_wallet_recharge_response(
)
.await?
else {
return if has_postgres {
return if has_wallet_writer {
Ok(build_admin_wallet_not_found_response())
} else {
Ok(build_admin_wallets_data_unavailable_response())

View File

@@ -1,8 +1,6 @@
use super::requests::ADMIN_WALLETS_API_KEY_GIFT_ADJUST_DETAIL;
use crate::handlers::admin::request::AdminRequestContext;
use crate::handlers::admin::shared::query_param_value;
use crate::GatewayError;
use sqlx::Row;
pub(in super::super) fn admin_wallet_operator_id(
request_context: &AdminRequestContext<'_>,
@@ -198,17 +196,6 @@ pub(in super::super) fn parse_admin_wallets_owner_type_filter(
}
}
pub(in super::super) fn optional_epoch_value(
row: &sqlx::postgres::PgRow,
key: &str,
) -> Result<Option<String>, GatewayError> {
Ok(row
.try_get::<Option<i64>, _>(key)
.map_err(|err| GatewayError::Internal(err.to_string()))?
.and_then(|value| u64::try_from(value).ok())
.and_then(crate::handlers::admin::shared::unix_secs_to_rfc3339))
}
pub(in super::super) fn admin_wallet_build_order_no(now: chrono::DateTime<chrono::Utc>) -> String {
format!(
"po_{}_{}",

View File

@@ -6,7 +6,6 @@ use super::route_filters::{
};
use crate::constants::INTERNAL_GATEWAY_PATH_PREFIXES;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::query::monitoring as monitoring_query;
use crate::GatewayError;
use aether_admin::observability::monitoring::{
admin_monitoring_bad_request_response, admin_monitoring_user_behavior_user_id_from_path,
@@ -41,33 +40,21 @@ pub(super) async fn build_admin_monitoring_audit_logs_response(
Err(detail) => return Ok(admin_monitoring_bad_request_response(detail)),
};
let Some(pool) = state.postgres_pool() else {
return Ok(build_admin_monitoring_audit_logs_payload_response(
Vec::new(),
0,
limit,
offset,
username,
event_type,
days,
));
};
let cutoff_time = chrono::Utc::now() - chrono::Duration::days(days);
let username_pattern = username
.as_deref()
.map(admin_monitoring_escape_like_pattern)
.map(|value| format!("%{value}%"));
let (items, total) = monitoring_query::list_admin_audit_logs(
&pool,
cutoff_time,
username_pattern.as_deref(),
event_type.as_deref(),
limit,
offset,
)
.await?;
let (items, total) = state
.list_admin_audit_logs(
cutoff_time,
username_pattern.as_deref(),
event_type.as_deref(),
limit,
offset,
)
.await?;
Ok(build_admin_monitoring_audit_logs_payload_response(
items, total, limit, offset, username, event_type, days,
@@ -85,14 +72,8 @@ pub(super) async fn build_admin_monitoring_suspicious_activities_response(
Err(detail) => return Ok(admin_monitoring_bad_request_response(detail)),
};
let Some(pool) = state.postgres_pool() else {
return Ok(
build_admin_monitoring_suspicious_activities_payload_response(Vec::new(), hours),
);
};
let cutoff_time = chrono::Utc::now() - chrono::Duration::hours(hours);
let activities = monitoring_query::list_admin_suspicious_activities(&pool, cutoff_time).await?;
let activities = state.list_admin_suspicious_activities(cutoff_time).await?;
Ok(build_admin_monitoring_suspicious_activities_payload_response(activities, hours))
}
@@ -112,22 +93,11 @@ pub(super) async fn build_admin_monitoring_user_behavior_response(
Err(detail) => return Ok(admin_monitoring_bad_request_response(detail)),
};
let Some(pool) = state.postgres_pool() else {
return Ok(build_admin_monitoring_user_behavior_payload_response(
user_id,
days,
std::collections::BTreeMap::new(),
0,
0,
0,
));
};
let cutoff_time = chrono::Utc::now() - chrono::Duration::days(days);
let event_counts =
monitoring_query::read_admin_user_behavior_event_counts(&pool, &user_id, cutoff_time)
.await?;
let event_counts = state
.read_admin_user_behavior_event_counts(&user_id, cutoff_time)
.await?;
let failed_requests = event_counts
.get("request_failed")

View File

@@ -23,7 +23,7 @@ async fn count_admin_monitoring_cache_affinity_entries(state: &AdminAppState<'_>
}
async fn scan_admin_monitoring_namespaced_keys(
runner: &aether_data::redis::RedisKvRunner,
runner: &aether_data::driver::redis::RedisKvRunner,
pattern: &str,
) -> Result<Vec<String>, GatewayError> {
let mut connection = runner

View File

@@ -1,7 +1,4 @@
use crate::handlers::admin::request::AdminAppState;
use crate::query::usage_heatmap::{
list_usage_heatmap_aggregate_rows, read_stats_daily_cutoff_date,
};
use crate::GatewayError;
use aether_admin::observability::stats::round_to;
use aether_admin::observability::usage::{
@@ -89,49 +86,15 @@ pub(super) async fn build_admin_usage_heatmap_response(
async fn build_admin_heatmap_summaries(
state: &AdminAppState<'_>,
created_from_unix_secs: u64,
start_date: chrono::NaiveDate,
today: chrono::NaiveDate,
_start_date: chrono::NaiveDate,
_today: chrono::NaiveDate,
) -> Result<Vec<StoredUsageDailySummary>, GatewayError> {
let query = UsageDailyHeatmapQuery {
created_from_unix_secs,
user_id: None,
admin_mode: true,
};
let Some(pool) = state.app().postgres_pool() else {
return state.summarize_usage_daily_heatmap(&query).await;
};
let Some(cutoff_date) = read_stats_daily_cutoff_date(&pool).await? else {
return state.summarize_usage_daily_heatmap(&query).await;
};
let cutoff_day = cutoff_date.date_naive().min(today);
let mut summaries =
list_usage_heatmap_aggregate_rows(&pool, start_date, cutoff_day, None).await?;
let raw_start_date = start_date.max(cutoff_day);
if raw_start_date <= today {
let raw_start_of_day = raw_start_date
.and_hms_opt(0, 0, 0)
.expect("heatmap day start should be valid");
let raw_created_from_unix_secs = u64::try_from(
chrono::DateTime::<chrono::Utc>::from_naive_utc_and_offset(
raw_start_of_day,
chrono::Utc,
)
.timestamp(),
)
.unwrap_or_default();
summaries.extend(
state
.summarize_usage_daily_heatmap(&UsageDailyHeatmapQuery {
created_from_unix_secs: raw_created_from_unix_secs,
user_id: None,
admin_mode: true,
})
.await?,
);
}
let mut summaries = state.summarize_usage_daily_heatmap(&query).await?;
summaries.sort_by(|left, right| left.date.cmp(&right.date));
Ok(summaries)
}

View File

@@ -3,7 +3,9 @@ use crate::handlers::admin::provider::shared::support::{
ADMIN_PROVIDER_MAPPING_PREVIEW_MAX_KEYS, ADMIN_PROVIDER_MAPPING_PREVIEW_MAX_MODELS,
};
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::shared::{decrypt_catalog_secret_with_fallbacks, json_string_list};
use crate::handlers::admin::shared::{
decrypt_catalog_secret_with_fallbacks, json_string_list, take_secret_prefix, take_secret_suffix,
};
use crate::handlers::public::matches_model_mapping_for_models;
use crate::{GatewayError, LocalProviderDeleteTaskState};
use aether_data_contracts::repository::global_models::{
@@ -175,14 +177,15 @@ pub(crate) fn mapping_preview_masked_catalog_api_key(
decrypt_catalog_secret_with_fallbacks(state.encryption_key(), ciphertext)
.map(|value| {
if value.len() > 8 {
let char_count = value.chars().count();
if char_count > 8 {
format!(
"{}***{}",
&value[..4],
&value[value.len().saturating_sub(4)..]
take_secret_prefix(&value, 4),
take_secret_suffix(&value, 4)
)
} else if value.len() >= 2 {
format!("{}***", &value[..2])
} else if char_count >= 2 {
format!("{}***", take_secret_prefix(&value, 2))
} else {
"***".to_string()
}

View File

@@ -1,4 +1,4 @@
use aether_data::redis::RedisKeyspace;
use aether_data::driver::redis::RedisKeyspace;
pub(super) fn pool_sticky_pattern(keyspace: &RedisKeyspace, provider_id: &str) -> String {
keyspace.key(&format!("ap:{provider_id}:sticky:*"))

View File

@@ -7,7 +7,7 @@ use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, ADMIN_PROVIDER_POOL_SCAN_BATCH,
};
use crate::GatewayError;
use aether_data::redis::RedisKvRunner;
use aether_data::driver::redis::RedisKvRunner;
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;

View File

@@ -5,7 +5,7 @@ use super::keys::{
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolUnschedulableRule,
};
use aether_data::redis::RedisKvRunner;
use aether_data::driver::redis::RedisKvRunner;
use regex::Regex;
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};

View File

@@ -116,7 +116,7 @@ impl<'a> AdminAppState<'a> {
self.app.mark_provider_key_rpm_reset(key_id, now_unix_secs)
}
pub(crate) fn redis_kv_runner(&self) -> Option<aether_data::redis::RedisKvRunner> {
pub(crate) fn redis_kv_runner(&self) -> Option<aether_data::driver::redis::RedisKvRunner> {
self.app.redis_kv_runner()
}
@@ -128,8 +128,8 @@ impl<'a> AdminAppState<'a> {
self.app.provider_key_rpm_reset_at(key_id, now_unix_secs)
}
pub(crate) fn has_postgres_pool(&self) -> bool {
self.app.postgres_pool().is_some()
pub(crate) fn has_wallet_data_writer(&self) -> bool {
self.app.has_wallet_data_writer()
}
pub(crate) fn mark_admin_monitoring_error_stats_reset(&self, now_unix_secs: u64) {

View File

@@ -12,6 +12,6 @@ pub(crate) use crate::handlers::shared::{
masked_catalog_api_key, normalize_json_array, normalize_json_object, normalize_string_list,
parse_catalog_auth_config_json, provider_catalog_key_supports_format,
provider_key_health_summary, provider_key_status_snapshot_payload, query_param_bool,
query_param_optional_bool, query_param_value, unix_secs_to_rfc3339,
OFFICIAL_EXTERNAL_MODEL_PROVIDERS,
query_param_optional_bool, query_param_value, take_secret_prefix, take_secret_suffix,
unix_secs_to_rfc3339, OFFICIAL_EXTERNAL_MODEL_PROVIDERS,
};

View File

@@ -4,7 +4,6 @@ use super::support::{
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::{query_param_optional_bool, query_param_value};
use crate::query::user_rollups::list_user_usage_totals_from_stats_summary;
use crate::GatewayError;
use axum::{
body::Body,
@@ -43,19 +42,10 @@ pub(in super::super) async fn build_admin_list_users_response(
.iter()
.map(|row| row.id.clone())
.collect::<Vec<_>>();
let usage_totals_future = async {
let Some(pool) = state.app().postgres_pool() else {
return state.summarize_usage_totals_by_user_ids(&user_ids).await;
};
match list_user_usage_totals_from_stats_summary(&pool, &user_ids).await? {
Some(items) => Ok(items),
None => state.summarize_usage_totals_by_user_ids(&user_ids).await,
}
};
let (auth_rows_result, wallet_rows_result, usage_totals_result) = tokio::join!(
state.list_user_auth_by_ids(&user_ids),
state.list_wallet_snapshots_by_user_ids(&user_ids),
usage_totals_future,
state.summarize_usage_totals_by_user_ids(&user_ids),
);
let auth_by_user_id = auth_rows_result?
.into_iter()