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

Merge upstream main

Browse Source
This commit is contained in:
ZheFox
2026-05-20 19:28:49 +08:00
parent cfa32a2b4d d0981c2fd5
commit a31e237cc3
123 changed files with 5970 additions and 595 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/aether-admin/src/observability/stats.rs
View File

@@ -2121,6 +2121,7 @@ mod tests {
api_key_allowed_providers: None,
api_key_allowed_api_formats: None,
api_key_allowed_models: None,
api_key_ip_rules: None,
}
}

531
crates/aether-admin/src/provider/ops/actions.rs
View File

@@ -1,3 +1,4 @@
use super::architectures::normalize_architecture_id;
use super::verify::admin_provider_ops_value_as_f64;
use serde_json::{json, Map, Value};
@@ -13,7 +14,7 @@ pub fn parse_query_balance_payload(
action_config: &Map<String, Value>,
response_json: &Value,
) -> Result<Value, String> {
match architecture_id {
match normalize_architecture_id(architecture_id) {
"generic_api" | "new_api" | "anyrouter" | "done_hub" => {
parse_new_api_balance_payload(action_config, response_json)
}
@@ -113,6 +114,73 @@ pub fn parse_sub2api_balance_payload(
))
}
pub fn parse_sub2api_api_key_usage_payload(
action_config: &Map<String, Value>,
response_json: &Value,
) -> Result<Value, String> {
let usage_data = sub2api_usage_response_object(response_json)?;
let is_valid = bool_value_from_candidates(
response_json,
usage_data,
&["is_active", "data.is_active", "isValid", "data.isValid"],
)
.unwrap_or(true);
if !is_valid {
return Err(response_json
.get("invalidMessage")
.or_else(|| response_json.get("message"))
.or_else(|| usage_data.get("invalidMessage"))
.or_else(|| usage_data.get("message"))
.and_then(Value::as_str)
.unwrap_or("API Key 已禁用或无效")
.to_string());
}
let remaining = balance_value_from_candidates(
response_json,
usage_data,
action_config,
&["remaining_path", "available_path", "balance_path"],
&[
"remaining",
"data.remaining",
"quota.remaining",
"data.quota.remaining",
"balance",
"data.balance",
],
)
.ok_or_else(|| "响应格式无效".to_string())?;
let currency = string_value_from_candidates(
response_json,
usage_data,
&["unit", "data.unit", "quota.unit", "data.quota.unit"],
)
.or_else(|| {
action_config
.get("currency")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
.unwrap_or_else(|| "USD".to_string());
let mut extra = Map::new();
extra.insert("is_active".to_string(), json!(is_valid));
if let Some(quota) = usage_data.get("quota").filter(|value| value.is_object()) {
extra.insert("quota".to_string(), quota.clone());
}
Ok(build_balance_data(
None,
None,
Some(remaining),
&currency,
extra,
))
}
pub fn attach_balance_checkin_outcome(
action_payload: &mut Value,
outcome: &ProviderOpsCheckinOutcome,
@@ -171,39 +239,316 @@ fn parse_new_api_balance_payload(
action_config: &Map<String, Value>,
response_json: &Value,
) -> Result<Value, String> {
let user_data = if response_json.get("success").and_then(Value::as_bool) == Some(true)
&& response_json.get("data").is_some_and(Value::is_object)
{
response_json.get("data")
} else if response_json.get("success").and_then(Value::as_bool) == Some(false) {
return Err(response_json
.get("message")
.and_then(Value::as_str)
.unwrap_or("业务状态码表示失败")
.to_string());
} else {
Some(response_json)
};
let Some(user_data) = user_data.and_then(Value::as_object) else {
return Err("响应格式无效".to_string());
};
let quota_divisor = quota_divisor(action_config);
let total_available =
admin_provider_ops_value_as_f64(user_data.get("quota")).map(|value| value / quota_divisor);
let total_used = admin_provider_ops_value_as_f64(user_data.get("used_quota"))
.map(|value| value / quota_divisor);
let user_data = balance_response_object(response_json)?;
let quota_divisor = balance_divisor(action_config);
let total_available_raw = balance_value_from_candidates(
response_json,
user_data,
action_config,
&[
"balance_path",
"available_path",
"total_available_path",
"quota_path",
],
&[
"total_available",
"data.total_available",
"balance",
"data.balance",
"available",
"data.available",
"remaining",
"data.remaining",
"quota",
"data.quota",
"balance_infos.0.total_balance",
"balance_infos[0].total_balance",
"data.balance_infos.0.total_balance",
"data.balance_infos[0].total_balance",
"balance_total",
"data.balance_total",
"total_balance",
"data.total_balance",
],
);
let total_used_raw = balance_value_from_candidates(
response_json,
user_data,
action_config,
&[
"used_path",
"used_quota_path",
"spent_path",
"usage_path",
"total_used_path",
],
&[
"used_quota",
"data.used_quota",
"used",
"data.used",
"total_used",
"data.total_used",
"spent",
"data.spent",
"usage",
"data.usage",
],
);
let total_granted_raw = balance_value_from_candidates(
response_json,
user_data,
action_config,
&[
"granted_path",
"total_granted_path",
"limit_path",
"total_quota_path",
],
&[
"total_granted",
"data.total_granted",
"total_quota",
"data.total_quota",
"granted",
"data.granted",
"limit",
"data.limit",
"balance_total",
"data.balance_total",
"total_balance",
"data.total_balance",
],
);
let total_available_raw = total_available_raw.or(match (total_granted_raw, total_used_raw) {
(Some(granted), Some(used)) => Some((granted - used).max(0.0)),
_ => None,
});
let total_used_raw = total_used_raw.or(match (total_granted_raw, total_available_raw) {
(Some(granted), Some(available)) => Some((granted - available).max(0.0)),
_ => None,
});
let total_granted_raw = total_granted_raw.or(match (total_available_raw, total_used_raw) {
(Some(available), Some(used)) => Some(available + used),
_ => None,
});
let mut extra = Map::new();
if let Some(plan_name) = new_api_plan_name(user_data) {
extra.insert("plan_name".to_string(), json!(plan_name));
}
Ok(build_balance_data(
None,
total_used,
total_available,
total_granted_raw.map(|value| value / quota_divisor),
total_used_raw.map(|value| value / quota_divisor),
total_available_raw.map(|value| value / quota_divisor),
action_config
.get("currency")
.and_then(Value::as_str)
.unwrap_or("USD"),
Map::new(),
extra,
))
}
fn new_api_plan_name(user_data: &Value) -> Option<String> {
for key in ["group", "plan_name", "planName", "plan", "package"] {
let value = user_data.get(key).and_then(Value::as_str)?.trim();
if !value.is_empty() {
return Some(value.to_string());
}
}
None
}
fn balance_response_object(response_json: &Value) -> Result<&Value, String> {
if let Some(success) = response_json.get("success").and_then(Value::as_bool) {
if !success {
return Err(response_json
.get("message")
.and_then(Value::as_str)
.unwrap_or("业务状态码表示失败")
.to_string());
}
if let Some(data) = response_json.get("data").filter(|value| value.is_object()) {
return Ok(data);
}
}
if let Some(code) = response_json.get("code").and_then(Value::as_i64) {
if code != 0 {
return Err(response_json
.get("message")
.and_then(Value::as_str)
.unwrap_or("查询余额失败")
.to_string());
}
if let Some(data) = response_json.get("data").filter(|value| value.is_object()) {
return Ok(data);
}
}
response_json
.as_object()
.map(|_| response_json)
.ok_or_else(|| "响应格式无效".to_string())
}
fn balance_value_from_candidates(
full_response: &Value,
response_data: &Value,
action_config: &Map<String, Value>,
config_keys: &[&str],
candidate_paths: &[&str],
) -> Option<f64> {
for key in config_keys {
if let Some(path) = action_config.get(*key).and_then(Value::as_str) {
let path = path.trim();
if path.is_empty() {
continue;
}
if let Some(value) = balance_value_at_path(full_response, path)
.or_else(|| balance_value_at_path(response_data, path))
{
return Some(value);
}
}
}
for path in candidate_paths {
if let Some(value) = balance_value_at_path(full_response, path)
.or_else(|| balance_value_at_path(response_data, path))
{
return Some(value);
}
}
None
}
fn balance_value_at_path(value: &Value, path: &str) -> Option<f64> {
let mut current = value;
for segment in path
.split('.')
.map(str::trim)
.filter(|segment| !segment.is_empty())
{
current = value_at_path_segment(current, segment)?;
}
admin_provider_ops_value_as_f64(Some(current))
}
fn value_at_path<'a>(value: &'a Value, path: &str) -> Option<&'a Value> {
let mut current = value;
for segment in path
.split('.')
.map(str::trim)
.filter(|segment| !segment.is_empty())
{
current = value_at_path_segment(current, segment)?;
}
Some(current)
}
fn value_at_path_segment<'a>(mut current: &'a Value, segment: &str) -> Option<&'a Value> {
if !segment.contains('[') {
return if current.is_array() {
segment
.parse::<usize>()
.ok()
.and_then(|index| current.get(index))
} else {
current.get(segment)
};
}
let mut rest = segment;
if let Some(open) = rest.find('[') {
let head = rest[..open].trim();
if !head.is_empty() {
current = current.get(head)?;
}
rest = &rest[open..];
}
while let Some(stripped) = rest.strip_prefix('[') {
let close = stripped.find(']')?;
let index = stripped[..close].trim().parse::<usize>().ok()?;
current = current.get(index)?;
rest = stripped[close + 1..].trim();
if rest.is_empty() {
return Some(current);
}
}
current.get(rest)
}
fn bool_value_from_candidates(
full_response: &Value,
response_data: &Value,
candidate_paths: &[&str],
) -> Option<bool> {
for path in candidate_paths {
if let Some(value) = value_at_path(full_response, path)
.or_else(|| value_at_path(response_data, path))
.and_then(Value::as_bool)
{
return Some(value);
}
}
None
}
fn string_value_from_candidates(
full_response: &Value,
response_data: &Value,
candidate_paths: &[&str],
) -> Option<String> {
for path in candidate_paths {
if let Some(value) = value_at_path(full_response, path)
.or_else(|| value_at_path(response_data, path))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
return Some(value.to_string());
}
}
None
}
fn sub2api_usage_response_object(response_json: &Value) -> Result<&Value, String> {
if let Some(success) = response_json.get("success").and_then(Value::as_bool) {
if !success {
return Err(response_json
.get("message")
.and_then(Value::as_str)
.unwrap_or("查询失败")
.to_string());
}
if let Some(data) = response_json.get("data").filter(|value| value.is_object()) {
return Ok(data);
}
}
if let Some(code) = response_json.get("code").and_then(Value::as_i64) {
if code != 0 {
return Err(response_json
.get("message")
.and_then(Value::as_str)
.unwrap_or("查询失败")
.to_string());
}
if let Some(data) = response_json.get("data").filter(|value| value.is_object()) {
return Ok(data);
}
}
response_json
.as_object()
.map(|_| response_json)
.ok_or_else(|| "响应格式无效".to_string())
}
fn parse_cubence_balance_payload(
action_config: &Map<String, Value>,
response_json: &Value,
@@ -414,6 +759,12 @@ fn quota_divisor(action_config: &Map<String, Value>) -> f64 {
.unwrap_or(500000.0)
}
fn balance_divisor(action_config: &Map<String, Value>) -> f64 {
admin_provider_ops_value_as_f64(action_config.get("balance_divisor"))
.filter(|value| *value > 0.0)
.unwrap_or_else(|| quota_divisor(action_config))
}
fn parse_rfc3339_unix_secs(value: Option<&Value>) -> Option<i64> {
let raw = value?.as_str()?.trim();
if raw.is_empty() {
@@ -466,7 +817,8 @@ fn parse_subscription(value: &Value) -> Option<Value> {
#[cfg(test)]
mod tests {
use super::{
attach_balance_checkin_outcome, parse_query_balance_payload, parse_sub2api_balance_payload,
attach_balance_checkin_outcome, parse_query_balance_payload,
parse_sub2api_api_key_usage_payload, parse_sub2api_balance_payload,
ProviderOpsCheckinOutcome,
};
use serde_json::json;
@@ -490,6 +842,109 @@ mod tests {
assert_eq!(payload["total_used"], json!(1.0));
}
#[test]
fn new_api_alias_parser_supports_balance_field_fallbacks() {
let payload = parse_query_balance_payload(
"oneapi",
&json!({ "quota_divisor": 1, "currency": "CNY" })
.as_object()
.cloned()
.expect("config"),
&json!({
"success": true,
"data": {
"balance": 12.5,
"used": 2.25
}
}),
)
.expect("payload should parse");
assert_eq!(payload["currency"], json!("CNY"));
assert_eq!(payload["total_available"], json!(12.5));
assert_eq!(payload["total_used"], json!(2.25));
}
#[test]
fn new_api_parser_supports_total_available_and_total_used_fields() {
let payload = parse_query_balance_payload(
"new_api",
&json!({ "quota_divisor": 1 })
.as_object()
.cloned()
.expect("config"),
&json!({
"data": {
"total_available": 9.75,
"total_used": 1.25
}
}),
)
.expect("payload should parse");
assert_eq!(payload["total_available"], json!(9.75));
assert_eq!(payload["total_used"], json!(1.25));
assert_eq!(payload["total_granted"], json!(11.0));
}
#[test]
fn new_api_parser_matches_cc_switch_usage_script_shape() {
let payload = parse_query_balance_payload(
"new_api",
&json!({ "quota_divisor": 500000, "currency": "USD" })
.as_object()
.cloned()
.expect("config"),
&json!({
"success": true,
"data": {
"group": "默认套餐",
"quota": 2_500_000,
"used_quota": 500_000
}
}),
)
.expect("payload should parse");
assert_eq!(payload["total_available"], json!(5.0));
assert_eq!(payload["total_used"], json!(1.0));
assert_eq!(payload["total_granted"], json!(6.0));
assert_eq!(payload["currency"], json!("USD"));
assert_eq!(payload["extra"]["plan_name"], json!("默认套餐"));
}
#[test]
fn generic_api_parser_supports_deepseek_balance_shape() {
let payload = parse_query_balance_payload(
"generic_api",
&json!({
"quota_divisor": 1,
"currency": "CNY",
"balance_path": "balance_infos[0].total_balance"
})
.as_object()
.cloned()
.expect("config"),
&json!({
"is_available": true,
"balance_infos": [
{
"currency": "CNY",
"total_balance": "128.50",
"granted_balance": "8.50",
"topped_up_balance": "120.00"
}
]
}),
)
.expect("payload should parse");
assert_eq!(payload["currency"], json!("CNY"));
assert_eq!(payload["total_available"], json!(128.5));
assert_eq!(payload["total_used"], json!(null));
assert_eq!(payload["total_granted"], json!(null));
}
#[test]
fn done_hub_single_request_parser_reads_wrapped_quota() {
let payload = parse_query_balance_payload(
@@ -540,6 +995,28 @@ mod tests {
assert_eq!(payload["extra"]["active_subscriptions"], json!(2));
}
#[test]
fn sub2api_api_key_usage_parser_matches_cc_switch_shape() {
let payload = parse_sub2api_api_key_usage_payload(
&json!({ "currency": "USD" })
.as_object()
.cloned()
.expect("config"),
&json!({
"is_active": true,
"quota": {
"remaining": "12.5",
"unit": "USD"
}
}),
)
.expect("payload should parse");
assert_eq!(payload["total_available"], json!(12.5));
assert_eq!(payload["currency"], json!("USD"));
assert_eq!(payload["extra"]["is_active"], json!(true));
}
#[test]
fn cubence_parser_reads_wrapped_dashboard_overview() {
let payload = parse_query_balance_payload(

19
crates/aether-admin/src/provider/ops/architectures/mod.rs
View File

@@ -119,12 +119,19 @@ pub fn get_architecture(architecture_id: &str) -> Option<ProviderOpsArchitecture
}
pub fn normalize_architecture_id(architecture_id: &str) -> &'static str {
match architecture_id.trim() {
let compact = architecture_id
.trim()
.chars()
.filter(|ch| ch.is_ascii_alphanumeric())
.collect::<String>()
.to_ascii_lowercase();
match compact.as_str() {
"" => "generic_api",
"generic_api" => "generic_api",
"new_api" => "new_api",
"genericapi" => "generic_api",
"newapi" | "oneapi" => "new_api",
"cubence" => "cubence",
"done_hub" => "done_hub",
"donehub" => "done_hub",
"yescode" => "yescode",
"nekocode" => "nekocode",
"anyrouter" => "anyrouter",
@@ -136,7 +143,7 @@ pub fn normalize_architecture_id(architecture_id: &str) -> &'static str {
pub fn admin_provider_ops_is_supported_auth_type(auth_type: &str) -> bool {
matches!(
auth_type,
"api_key" | "session_login" | "oauth" | "cookie" | "none"
"api_key" | "refresh_token" | "session_login" | "oauth" | "cookie" | "none"
)
}
@@ -220,6 +227,8 @@ mod tests {
assert_eq!(normalize_architecture_id(""), "generic_api");
assert_eq!(normalize_architecture_id("done_hub"), "done_hub");
assert_eq!(normalize_architecture_id("new_api"), "new_api");
assert_eq!(normalize_architecture_id("newapi"), "new_api");
assert_eq!(normalize_architecture_id("one-api"), "new_api");
assert_eq!(normalize_architecture_id("unknown"), "generic_api");
}

13
crates/aether-admin/src/provider/ops/architectures/new_api.rs
View File

@@ -10,8 +10,8 @@ pub(super) fn spec() -> ProviderOpsArchitectureSpec {
"properties": {
"api_key": {
"type": "string",
"title": "访问令牌 (API Key)",
"description": "New API 的访问令牌，与 Cookie 二选一",
"title": "访问令牌",
"description": "在 New API 个人安全设置中获取的访问令牌，与 Cookie 二选一",
"x-sensitive": true,
"x-input-type": "password"
},
@@ -30,7 +30,7 @@ pub(super) fn spec() -> ProviderOpsArchitectureSpec {
"user_id": {
"type": "string",
"title": "用户 ID",
"description": "使用访问令牌时必填，使用 Cookie 时可选"
"description": "可选；使用 Cookie 时可自动解析"
}
},
"required": [],
@@ -64,13 +64,6 @@ pub(super) fn spec() -> ProviderOpsArchitectureSpec {
"type": "any_required",
"fields": ["api_key", "cookie"],
"message": "访问令牌和 Cookie 至少需要填写一个"
},
{
"type": "conditional_required",
"if": "api_key",
"then": ["user_id"],
"unless": "cookie",
"message": "使用访问令牌时，用户 ID 不能为空"
}
]
});

48
crates/aether-admin/src/provider/ops/architectures/sub2api.rs
View File

@@ -5,6 +5,38 @@ use super::{
use serde_json::{json, Map, Value};
pub(super) fn spec() -> ProviderOpsArchitectureSpec {
let api_key_usage_schema = json!({
"type": "object",
"properties": {
"base_url": {
"type": "string",
"title": "站点地址",
"description": "API 基础地址"
},
"api_key": {
"type": "string",
"title": "API Key",
"description": "用于访问 Sub2API /v1/usage 的模型 API Key",
"x-sensitive": true,
"x-input-type": "password",
"x-help": "请求 GET /v1/usage，并通过 Authorization: Bearer <API Key> 查询余量"
}
},
"required": ["api_key"],
"x-auth-method": "bearer",
"x-auth-type": "api_key",
"x-field-groups": [
{ "fields": ["base_url"] },
{ "fields": ["api_key"] }
],
"x-validation": [
{
"type": "required",
"fields": ["api_key"],
"message": "请填写 API Key"
}
]
});
let session_login_schema = json!({
"type": "object",
"properties": {
@@ -61,7 +93,7 @@ pub(super) fn spec() -> ProviderOpsArchitectureSpec {
},
"required": ["refresh_token"],
"x-auth-method": "bearer",
"x-auth-type": "api_key",
"x-auth-type": "refresh_token",
"x-field-groups": [
{ "fields": ["base_url"] },
{ "fields": ["refresh_token"] }
@@ -80,20 +112,25 @@ pub(super) fn spec() -> ProviderOpsArchitectureSpec {
display_name: "Sub2API",
description: "Sub2API 风格中转站的预设配置",
hidden: false,
credentials_schema: session_login_schema.clone(),
credentials_schema: api_key_usage_schema.clone(),
verify_endpoint: "/api/v1/auth/me?timezone=Asia/Shanghai",
verify_mode: ProviderOpsVerifyMode::Sub2ApiExchange,
balance_mode: ProviderOpsBalanceMode::Sub2ApiDualRequest,
checkin_mode: ProviderOpsCheckinMode::None,
query_balance_cookie_auth_errors: false,
supported_auth_types: vec![
ProviderOpsAuthSpec {
auth_type: "api_key",
display_name: "API Key 用量接口",
credentials_schema: api_key_usage_schema,
},
ProviderOpsAuthSpec {
auth_type: "session_login",
display_name: "账号密码",
credentials_schema: session_login_schema,
},
ProviderOpsAuthSpec {
auth_type: "api_key",
auth_type: "refresh_token",
display_name: "Refresh Token",
credentials_schema: refresh_token_schema,
},
@@ -101,7 +138,7 @@ pub(super) fn spec() -> ProviderOpsArchitectureSpec {
supported_actions: vec![ProviderOpsActionSpec {
action_type: "query_balance",
display_name: "查询余额",
description: "查询 Sub2API 账户余额和订阅信息",
description: "查询 Sub2API API Key 用量或账户余额和订阅信息",
config_schema: json!({
"type": "object",
"properties": {
@@ -114,7 +151,7 @@ pub(super) fn spec() -> ProviderOpsArchitectureSpec {
"required": []
}),
}],
default_connector: Some("session_login"),
default_connector: Some("api_key"),
}
}
@@ -123,6 +160,7 @@ pub(super) fn default_action_config(action_type: &str) -> Option<Map<String, Val
"query_balance" => Some(json_object(json!({
"endpoint": "/api/v1/auth/me?timezone=Asia/Shanghai",
"subscription_endpoint": "/api/v1/subscriptions/summary",
"api_key_usage_endpoint": "/v1/usage",
"method": "GET",
"currency": "USD"
}))),

3
crates/aether-admin/src/provider/ops/mod.rs
View File

@@ -4,7 +4,8 @@ pub mod config;
pub mod verify;
pub use self::actions::{
attach_balance_checkin_outcome, parse_query_balance_payload, parse_sub2api_balance_payload,
attach_balance_checkin_outcome, parse_query_balance_payload,
parse_sub2api_api_key_usage_payload, parse_sub2api_balance_payload,
parse_yescode_combined_balance_payload, ProviderOpsCheckinOutcome,
};
pub use self::architectures::{

13
crates/aether-admin/src/provider/ops/verify.rs
View File

@@ -379,18 +379,9 @@ pub fn admin_provider_ops_verify_headers(
}
"new_api" => {
for (name, value) in [
(
"User-Agent",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.7339.249 Electron/38.7.0 Safari/537.36",
),
("User-Agent", "cc-switch/1.0"),
("Content-Type", "application/json"),
("Accept", "application/json"),
("Accept-Language", "zh-CN"),
("sec-ch-ua", "\"Not=A?Brand\";v=\"24\", \"Chromium\";v=\"140\""),
("sec-ch-ua-mobile", "?0"),
("sec-ch-ua-platform", "\"macOS\""),
("Sec-Fetch-Site", "cross-site"),
("Sec-Fetch-Mode", "cors"),
("Sec-Fetch-Dest", "empty"),
] {
insert_header(&mut headers, name, value)?;
}

2
crates/aether-data/migrations/mysql/20260520000000_add_api_key_ip_rules.sql Normal file
View File

@@ -0,0 +1,2 @@
ALTER TABLE api_keys
ADD COLUMN ip_rules TEXT NULL AFTER allowed_models;

5
crates/aether-data/migrations/postgres/20260520000000_add_api_key_ip_rules.sql Normal file
View File

@@ -0,0 +1,5 @@
ALTER TABLE api_keys
ADD COLUMN IF NOT EXISTS ip_rules jsonb NULL;
ALTER TABLE api_keys
ALTER COLUMN ip_rules TYPE jsonb USING ip_rules::jsonb;

1
crates/aether-data/migrations/sqlite/20260520000000_add_api_key_ip_rules.sql Normal file
View File

@@ -0,0 +1 @@
ALTER TABLE api_keys ADD COLUMN ip_rules TEXT;

1
crates/aether-data/schema/bootstrap/postgres/001_types_and_tables.sql
View File

@@ -176,6 +176,7 @@ CREATE TABLE IF NOT EXISTS public.api_keys (
allowed_providers json,
allowed_api_formats json,
allowed_models json,
ip_rules jsonb,
rate_limit integer DEFAULT 100,
concurrent_limit integer,
force_capabilities json,

1
crates/aether-data/schema/generated/mysql/baseline/001_identity.sql
View File

@@ -75,6 +75,7 @@ CREATE TABLE IF NOT EXISTS api_keys (
`allowed_models` JSON,
`allowed_providers` JSON,
`allowed_api_formats` JSON,
`ip_rules` JSON,
`rate_limit` INT DEFAULT 100,
`concurrent_limit` INT,
`force_capabilities` JSON,

1
crates/aether-data/schema/generated/postgres/baseline/001_identity.sql
View File

@@ -78,6 +78,7 @@ CREATE TABLE IF NOT EXISTS public.api_keys (
allowed_models jsonb,
allowed_providers jsonb,
allowed_api_formats jsonb,
ip_rules jsonb,
rate_limit integer DEFAULT 100,
concurrent_limit integer,
force_capabilities jsonb,

1
crates/aether-data/schema/generated/sqlite/baseline/001_identity.sql
View File

@@ -73,6 +73,7 @@ CREATE TABLE IF NOT EXISTS api_keys (
allowed_models TEXT,
allowed_providers TEXT,
allowed_api_formats TEXT,
ip_rules TEXT,
rate_limit INTEGER DEFAULT 100,
concurrent_limit INTEGER,
force_capabilities TEXT,

5
crates/aether-data/schema/logical/001_identity.toml
View File

@@ -333,6 +333,11 @@ name = "allowed_api_formats"
type = "json"
nullable = true
[[table.api_keys.columns]]
name = "ip_rules"
type = "json"
nullable = true
[[table.api_keys.columns]]
name = "rate_limit"
type = "int32"

2
crates/aether-data/src/lifecycle/bootstrap/postgres.rs
View File

@@ -7,7 +7,7 @@ use tracing::info;
// Generated by build.rs from schema/bootstrap/postgres.
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260519130000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260520000000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT

28
crates/aether-data/src/lifecycle/migrate/tests.rs
View File

@@ -310,6 +310,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260519000000,
20260519120000,
20260519130000,
20260520000000,
]
);
}
@@ -469,6 +470,30 @@ fn management_tokens_json_columns_are_normalized_to_jsonb_in_postgres_schema_pat
assert!(generated_identity.contains("permissions jsonb,"));
}
#[test]
fn api_key_ip_rules_is_jsonb_in_postgres_schema_paths() {
let api_key_ip_rules_migration = POSTGRES_MIGRATOR
.iter()
.find(|migration| migration.version == 20260520000000)
.expect("api key IP rules migration should be embedded");
assert!(api_key_ip_rules_migration
.sql
.contains("ADD COLUMN IF NOT EXISTS ip_rules jsonb NULL"));
assert!(api_key_ip_rules_migration
.sql
.contains("ALTER COLUMN ip_rules TYPE jsonb USING ip_rules::jsonb"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("ip_rules jsonb,"));
let bootstrap_schema =
include_str!("../../../schema/bootstrap/postgres/001_types_and_tables.sql");
assert!(bootstrap_schema.contains("ip_rules jsonb,"));
let generated_identity =
include_str!("../../../schema/generated/postgres/baseline/001_identity.sql");
assert!(generated_identity.contains("ip_rules jsonb,"));
}
#[test]
fn provider_api_keys_api_key_is_nullable() {
let baseline_migration = POSTGRES_MIGRATOR
@@ -605,6 +630,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260519000000,
20260519120000,
20260519130000,
20260520000000,
]
);
assert_eq!(
@@ -626,6 +652,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260519000000,
20260519120000,
20260519130000,
20260520000000,
]
);
}
@@ -1147,6 +1174,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260519000000,
20260519120000,
20260519130000,
20260520000000,
]
);
}

52
crates/aether-data/src/repository/auth/memory.rs
View File

@@ -78,6 +78,14 @@ impl InMemoryAuthApiKeySnapshotRepository {
0.0,
snapshot.api_key_is_standalone,
)
.and_then(|record| {
record.with_ip_rules(
snapshot
.api_key_ip_rules
.as_ref()
.map(|value| serde_json::json!(value)),
)
})
.expect("derived auth api key export record should build"),
);
if let Some(key_hash) = key_hash {
@@ -496,6 +504,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
api_key_allowed_providers: record.allowed_providers.clone(),
api_key_allowed_api_formats: record.allowed_api_formats.clone(),
api_key_allowed_models: record.allowed_models.clone(),
api_key_ip_rules: record.ip_rules.clone(),
..template
}
} else {
@@ -534,6 +543,12 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
.as_ref()
.map(|value| serde_json::json!(value)),
)?
.with_api_key_ip_rules(
record
.ip_rules
.as_ref()
.map(|value| serde_json::json!(value)),
)?
};
let now_unix_secs = current_unix_secs() as i64;
@@ -566,6 +581,12 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
record.total_cost_usd,
false,
)?
.with_ip_rules(
record
.ip_rules
.as_ref()
.map(|value| serde_json::json!(value)),
)?
.with_activity_timestamps(None, Some(now_unix_secs), Some(now_unix_secs))?;
index
@@ -619,6 +640,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
api_key_allowed_providers: record.allowed_providers.clone(),
api_key_allowed_api_formats: record.allowed_api_formats.clone(),
api_key_allowed_models: record.allowed_models.clone(),
api_key_ip_rules: record.ip_rules.clone(),
..template
}
} else {
@@ -657,6 +679,12 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
.as_ref()
.map(|value| serde_json::json!(value)),
)?
.with_api_key_ip_rules(
record
.ip_rules
.as_ref()
.map(|value| serde_json::json!(value)),
)?
};
let now_unix_secs = current_unix_secs() as i64;
@@ -689,6 +717,12 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
record.total_cost_usd,
true,
)?
.with_ip_rules(
record
.ip_rules
.as_ref()
.map(|value| serde_json::json!(value)),
)?
.with_activity_timestamps(None, Some(now_unix_secs), Some(now_unix_secs))?;
index
@@ -741,6 +775,14 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
export.concurrent_limit = Some(concurrent_limit);
}
}
if let Some(ip_rules) = record.ip_rules {
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
snapshot.api_key_ip_rules = ip_rules.clone();
}
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
export.ip_rules = ip_rules;
}
}
Ok(index.export_by_api_key_id.get(&record.api_key_id).cloned())
}
@@ -806,6 +848,14 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
export.allowed_models = allowed_models;
}
}
if let Some(ip_rules) = record.ip_rules {
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
snapshot.api_key_ip_rules = ip_rules.clone();
}
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
export.ip_rules = ip_rules;
}
}
if record.expires_at_present {
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
snapshot.api_key_expires_at_unix_secs = record.expires_at_unix_secs;
@@ -1239,6 +1289,7 @@ mod tests {
name: None,
rate_limit: None,
concurrent_limit: Some(11),
ip_rules: None,
})
.await
.expect("update should succeed")
@@ -1273,6 +1324,7 @@ mod tests {
allowed_providers: None,
allowed_api_formats: None,
allowed_models: None,
ip_rules: None,
expires_at_present: false,
expires_at_unix_secs: None,
auto_delete_on_expiry_present: false,

39
crates/aether-data/src/repository/auth/mysql.rs
View File

@@ -34,7 +34,8 @@ SELECT
api_keys.expires_at AS api_key_expires_at_unix_secs,
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models
api_keys.allowed_models AS api_key_allowed_models,
api_keys.ip_rules AS api_key_ip_rules
FROM api_keys
JOIN users ON users.id = api_keys.user_id
"#;
@@ -49,6 +50,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -112,12 +114,12 @@ impl MysqlAuthApiKeyReadRepository {
r#"
INSERT INTO api_keys (
id, user_id, key_hash, key_encrypted, name, allowed_providers,
allowed_api_formats, allowed_models, rate_limit, concurrent_limit,
allowed_api_formats, allowed_models, ip_rules, rate_limit, concurrent_limit,
force_capabilities, feature_settings, is_active, expires_at, auto_delete_on_expiry,
total_requests, total_tokens, total_cost_usd, is_standalone,
created_at, updated_at
)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&record.api_key_id)
@@ -137,6 +139,10 @@ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?)
record.allowed_models.as_ref(),
"api_keys.allowed_models",
)?)
.bind(json_string_from_string_list(
record.ip_rules.as_ref(),
"api_keys.ip_rules",
)?)
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.bind(optional_json_to_string(
@@ -175,6 +181,7 @@ struct CreateApiKeyInsertRecord {
allowed_providers: Option<Vec<String>>,
allowed_api_formats: Option<Vec<String>>,
allowed_models: Option<Vec<String>>,
ip_rules: Option<Vec<String>>,
rate_limit: Option<i32>,
concurrent_limit: Option<i32>,
force_capabilities: Option<serde_json::Value>,
@@ -417,6 +424,7 @@ WHERE id = ?
allowed_providers: record.allowed_providers,
allowed_api_formats: record.allowed_api_formats,
allowed_models: record.allowed_models,
ip_rules: record.ip_rules,
rate_limit: Some(record.rate_limit),
concurrent_limit: record.concurrent_limit,
force_capabilities: record.force_capabilities,
@@ -444,6 +452,7 @@ WHERE id = ?
allowed_providers: record.allowed_providers,
allowed_api_formats: record.allowed_api_formats,
allowed_models: record.allowed_models,
ip_rules: record.ip_rules,
rate_limit: record.rate_limit,
concurrent_limit: record.concurrent_limit,
force_capabilities: record.force_capabilities,
@@ -469,6 +478,7 @@ UPDATE api_keys
SET name = COALESCE(?, name),
rate_limit = COALESCE(?, rate_limit),
concurrent_limit = COALESCE(?, concurrent_limit),
ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END,
updated_at = ?
WHERE id = ?
AND user_id = ?
@@ -478,6 +488,11 @@ WHERE id = ?
.bind(record.name.as_deref())
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.bind(record.ip_rules.is_some())
.bind(json_string_from_nested_string_list(
&record.ip_rules,
"api_keys.ip_rules",
)?)
.bind(now)
.bind(&record.api_key_id)
.bind(&record.user_id)
@@ -501,6 +516,7 @@ SET name = COALESCE(?, name),
allowed_providers = CASE WHEN ? THEN ? ELSE allowed_providers END,
allowed_api_formats = CASE WHEN ? THEN ? ELSE allowed_api_formats END,
allowed_models = CASE WHEN ? THEN ? ELSE allowed_models END,
ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END,
expires_at = CASE WHEN ? THEN ? ELSE expires_at END,
auto_delete_on_expiry = CASE WHEN ? THEN ? ELSE auto_delete_on_expiry END,
updated_at = ?
@@ -528,6 +544,11 @@ WHERE id = ?
&record.allowed_models,
"api_keys.allowed_models",
)?)
.bind(record.ip_rules.is_some())
.bind(json_string_from_nested_string_list(
&record.ip_rules,
"api_keys.ip_rules",
)?)
.bind(record.expires_at_present)
.bind(optional_i64_from_u64(
record.expires_at_unix_secs,
@@ -922,7 +943,11 @@ fn map_auth_api_key_snapshot_row(
row.try_get("api_key_allowed_models").map_sql_err()?,
"api_keys.allowed_models",
)?,
)?;
)?
.with_api_key_ip_rules(optional_json_from_string(
row.try_get("api_key_ip_rules").map_sql_err()?,
"api_keys.ip_rules",
)?)?;
Ok(snapshot.with_user_rate_limit(row.try_get("user_rate_limit").map_sql_err()?))
}
@@ -965,6 +990,12 @@ fn map_auth_api_key_export_row(
row.try_get("total_cost_usd").map_sql_err()?,
row.try_get("is_standalone").map_sql_err()?,
)
.and_then(|record| {
record.with_ip_rules(optional_json_from_string(
row.try_get("ip_rules").map_sql_err()?,
"api_keys.ip_rules",
)?)
})
.map(|record| record.with_feature_settings(feature_settings))
.and_then(|record| {
record.with_activity_timestamps(

96
crates/aether-data/src/repository/auth/postgres.rs
View File

@@ -36,7 +36,8 @@ SELECT
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models
api_keys.allowed_models AS api_key_allowed_models,
api_keys.ip_rules AS api_key_ip_rules
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.key_hash = $1
@@ -66,7 +67,8 @@ SELECT
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models
api_keys.allowed_models AS api_key_allowed_models,
api_keys.ip_rules AS api_key_ip_rules
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.id = $1
@@ -96,7 +98,8 @@ SELECT
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models
api_keys.allowed_models AS api_key_allowed_models,
api_keys.ip_rules AS api_key_ip_rules
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.id = $1 AND users.id = $2
@@ -126,7 +129,8 @@ SELECT
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models
api_keys.allowed_models AS api_key_allowed_models,
api_keys.ip_rules AS api_key_ip_rules
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.id = ANY($1::TEXT[])
@@ -143,6 +147,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -173,6 +178,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -202,6 +208,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -231,6 +238,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -260,6 +268,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -333,6 +342,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -369,6 +379,7 @@ INSERT INTO api_keys (
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -396,15 +407,16 @@ VALUES (
$9,
$10,
$11,
NULL,
$12,
NULL,
$13,
$14,
FALSE,
FALSE,
$15,
FALSE,
FALSE,
$16,
$17,
$18,
NOW(),
NOW()
)
@@ -417,6 +429,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -443,6 +456,7 @@ INSERT INTO api_keys (
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -470,15 +484,16 @@ VALUES (
$9,
$10,
$11,
NULL,
$12,
NULL,
$13,
$14,
$15,
FALSE,
TRUE,
$15,
$16,
$17,
$18,
NOW(),
NOW()
)
@@ -491,6 +506,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -513,6 +529,7 @@ SET
name = COALESCE($3, name),
rate_limit = COALESCE($4, rate_limit),
concurrent_limit = COALESCE($5, concurrent_limit),
ip_rules = CASE WHEN $6 THEN $7::jsonb ELSE ip_rules END,
updated_at = NOW()
WHERE user_id = $1
AND id = $2
@@ -526,6 +543,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -551,8 +569,9 @@ SET
allowed_providers = CASE WHEN $7 THEN $8::json ELSE allowed_providers END,
allowed_api_formats = CASE WHEN $9 THEN $10::json ELSE allowed_api_formats END,
allowed_models = CASE WHEN $11 THEN $12::json ELSE allowed_models END,
expires_at = CASE WHEN $13 THEN $14::timestamptz ELSE expires_at END,
auto_delete_on_expiry = CASE WHEN $15 THEN $16 ELSE auto_delete_on_expiry END,
ip_rules = CASE WHEN $13 THEN $14::jsonb ELSE ip_rules END,
expires_at = CASE WHEN $15 THEN $16::timestamptz ELSE expires_at END,
auto_delete_on_expiry = CASE WHEN $17 THEN $18 ELSE auto_delete_on_expiry END,
updated_at = NOW()
WHERE id = $1
AND is_standalone = TRUE
@@ -565,6 +584,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -598,6 +618,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -630,6 +651,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -673,6 +695,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -1122,6 +1145,11 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let ip_rules = record
.ip_rules
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let expires_at = record
.expires_at_unix_secs
.map(|value| {
@@ -1139,6 +1167,7 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(allowed_providers)
.bind(allowed_api_formats)
.bind(allowed_models)
.bind(ip_rules)
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.bind(record.force_capabilities)
@@ -1173,6 +1202,11 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let ip_rules = record
.ip_rules
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let expires_at = record
.expires_at_unix_secs
.map(|value| {
@@ -1190,6 +1224,7 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(allowed_providers)
.bind(allowed_api_formats)
.bind(allowed_models)
.bind(ip_rules)
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.bind(record.force_capabilities)
@@ -1209,12 +1244,21 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
&self,
record: UpdateUserApiKeyBasicRecord,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
let ip_rules = record
.ip_rules
.clone()
.flatten()
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let row = sqlx::query(UPDATE_USER_API_KEY_BASIC_SQL)
.bind(record.user_id)
.bind(record.api_key_id)
.bind(record.name)
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.bind(record.ip_rules.is_some())
.bind(ip_rules)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
@@ -1246,6 +1290,13 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let ip_rules = record
.ip_rules
.clone()
.flatten()
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let expires_at = record
.expires_at_unix_secs
.map(|value| {
@@ -1267,6 +1318,8 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(allowed_api_formats)
.bind(record.allowed_models.is_some())
.bind(allowed_models)
.bind(record.ip_rules.is_some())
.bind(ip_rules)
.bind(record.expires_at_present)
.bind(expires_at)
.bind(record.auto_delete_on_expiry_present)
@@ -1495,7 +1548,8 @@ fn map_auth_api_key_snapshot_row(
row_get(row, "api_key_allowed_providers")?,
row_get(row, "api_key_allowed_api_formats")?,
row_get(row, "api_key_allowed_models")?,
)?;
)?
.with_api_key_ip_rules(row_get(row, "api_key_ip_rules")?)?;
Ok(snapshot.with_user_rate_limit(row_get(row, "user_rate_limit")?))
}
@@ -1523,6 +1577,7 @@ fn map_auth_api_key_export_row(
row_get(row, "total_cost_usd")?,
row_get(row, "is_standalone")?,
)
.and_then(|record| record.with_ip_rules(row_get(row, "ip_rules")?))
.map(|record| record.with_feature_settings(feature_settings))
.and_then(|record| {
record.with_activity_timestamps(
@@ -1538,6 +1593,7 @@ mod tests {
use super::{
SqlxAuthApiKeySnapshotReadRepository, CREATE_STANDALONE_API_KEY_SQL,
CREATE_USER_API_KEY_SQL, UPDATE_STANDALONE_API_KEY_BASIC_SQL,
UPDATE_USER_API_KEY_BASIC_SQL,
};
use crate::driver::postgres::{PostgresPoolConfig, PostgresPoolFactory};
@@ -1546,12 +1602,12 @@ mod tests {
assert!(CREATE_USER_API_KEY_SQL
.contains("expires_at,\n auto_delete_on_expiry,\n is_locked,\n is_standalone,"));
assert!(
CREATE_USER_API_KEY_SQL.contains("$12,\n $13,\n $14,\n FALSE,\n FALSE,\n $15,")
CREATE_USER_API_KEY_SQL.contains("$13,\n $14,\n $15,\n FALSE,\n FALSE,\n $16,")
);
assert!(CREATE_STANDALONE_API_KEY_SQL
.contains("expires_at,\n auto_delete_on_expiry,\n is_locked,\n is_standalone,"));
assert!(CREATE_STANDALONE_API_KEY_SQL
.contains("$12,\n $13,\n $14,\n FALSE,\n TRUE,\n $15,"));
.contains("$13,\n $14,\n $15,\n FALSE,\n TRUE,\n $16,"));
}
#[test]
@@ -1565,15 +1621,23 @@ mod tests {
));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("allowed_models = CASE WHEN $11 THEN $12::json ELSE allowed_models END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("ip_rules = CASE WHEN $13 THEN $14::jsonb ELSE ip_rules END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("rate_limit = CASE WHEN $3 THEN $4 ELSE rate_limit END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("expires_at = CASE WHEN $13 THEN $14::timestamptz ELSE expires_at END"));
.contains("expires_at = CASE WHEN $15 THEN $16::timestamptz ELSE expires_at END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL.contains(
"auto_delete_on_expiry = CASE WHEN $15 THEN $16 ELSE auto_delete_on_expiry END"
"auto_delete_on_expiry = CASE WHEN $17 THEN $18 ELSE auto_delete_on_expiry END"
));
}
#[test]
fn update_user_api_key_basic_sql_casts_ip_rules_as_jsonb() {
assert!(UPDATE_USER_API_KEY_BASIC_SQL
.contains("ip_rules = CASE WHEN $6 THEN $7::jsonb ELSE ip_rules END"));
}
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {

43
crates/aether-data/src/repository/auth/sqlite.rs
View File

@@ -34,7 +34,8 @@ SELECT
api_keys.expires_at AS api_key_expires_at_unix_secs,
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models
api_keys.allowed_models AS api_key_allowed_models,
api_keys.ip_rules AS api_key_ip_rules
FROM api_keys
JOIN users ON users.id = api_keys.user_id
"#;
@@ -49,6 +50,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -112,12 +114,12 @@ impl SqliteAuthApiKeyReadRepository {
r#"
INSERT INTO api_keys (
id, user_id, key_hash, key_encrypted, name, allowed_providers,
allowed_api_formats, allowed_models, rate_limit, concurrent_limit,
allowed_api_formats, allowed_models, ip_rules, rate_limit, concurrent_limit,
force_capabilities, feature_settings, is_active, expires_at, auto_delete_on_expiry,
total_requests, total_tokens, total_cost_usd, is_standalone,
created_at, updated_at
)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&record.api_key_id)
@@ -137,6 +139,10 @@ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?)
record.allowed_models.as_ref(),
"api_keys.allowed_models",
)?)
.bind(json_string_from_string_list(
record.ip_rules.as_ref(),
"api_keys.ip_rules",
)?)
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.bind(optional_json_to_string(
@@ -175,6 +181,7 @@ struct CreateApiKeyInsertRecord {
allowed_providers: Option<Vec<String>>,
allowed_api_formats: Option<Vec<String>>,
allowed_models: Option<Vec<String>>,
ip_rules: Option<Vec<String>>,
rate_limit: Option<i32>,
concurrent_limit: Option<i32>,
force_capabilities: Option<serde_json::Value>,
@@ -417,6 +424,7 @@ WHERE id = ?
allowed_providers: record.allowed_providers,
allowed_api_formats: record.allowed_api_formats,
allowed_models: record.allowed_models,
ip_rules: record.ip_rules,
rate_limit: Some(record.rate_limit),
concurrent_limit: record.concurrent_limit,
force_capabilities: record.force_capabilities,
@@ -444,6 +452,7 @@ WHERE id = ?
allowed_providers: record.allowed_providers,
allowed_api_formats: record.allowed_api_formats,
allowed_models: record.allowed_models,
ip_rules: record.ip_rules,
rate_limit: record.rate_limit,
concurrent_limit: record.concurrent_limit,
force_capabilities: record.force_capabilities,
@@ -469,6 +478,7 @@ UPDATE api_keys
SET name = COALESCE(?, name),
rate_limit = COALESCE(?, rate_limit),
concurrent_limit = COALESCE(?, concurrent_limit),
ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END,
updated_at = ?
WHERE id = ?
AND user_id = ?
@@ -478,6 +488,11 @@ WHERE id = ?
.bind(record.name.as_deref())
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.bind(record.ip_rules.is_some())
.bind(json_string_from_nested_string_list(
&record.ip_rules,
"api_keys.ip_rules",
)?)
.bind(now)
.bind(&record.api_key_id)
.bind(&record.user_id)
@@ -501,6 +516,7 @@ SET name = COALESCE(?, name),
allowed_providers = CASE WHEN ? THEN ? ELSE allowed_providers END,
allowed_api_formats = CASE WHEN ? THEN ? ELSE allowed_api_formats END,
allowed_models = CASE WHEN ? THEN ? ELSE allowed_models END,
ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END,
expires_at = CASE WHEN ? THEN ? ELSE expires_at END,
auto_delete_on_expiry = CASE WHEN ? THEN ? ELSE auto_delete_on_expiry END,
updated_at = ?
@@ -528,6 +544,11 @@ WHERE id = ?
&record.allowed_models,
"api_keys.allowed_models",
)?)
.bind(record.ip_rules.is_some())
.bind(json_string_from_nested_string_list(
&record.ip_rules,
"api_keys.ip_rules",
)?)
.bind(record.expires_at_present)
.bind(optional_i64_from_u64(
record.expires_at_unix_secs,
@@ -922,7 +943,11 @@ fn map_auth_api_key_snapshot_row(
row.try_get("api_key_allowed_models").map_sql_err()?,
"api_keys.allowed_models",
)?,
)?;
)?
.with_api_key_ip_rules(optional_json_from_string(
row.try_get("api_key_ip_rules").map_sql_err()?,
"api_keys.ip_rules",
)?)?;
Ok(snapshot.with_user_rate_limit(row.try_get("user_rate_limit").map_sql_err()?))
}
@@ -965,6 +990,12 @@ fn map_auth_api_key_export_row(
sqlite_real(row, "total_cost_usd")?,
row.try_get("is_standalone").map_sql_err()?,
)
.and_then(|record| {
record.with_ip_rules(optional_json_from_string(
row.try_get("ip_rules").map_sql_err()?,
"api_keys.ip_rules",
)?)
})
.map(|record| record.with_feature_settings(feature_settings))
.and_then(|record| {
record.with_activity_timestamps(
@@ -1081,6 +1112,7 @@ mod tests {
allowed_providers: Some(vec!["openai".to_string()]),
allowed_api_formats: Some(vec!["openai:chat".to_string()]),
allowed_models: Some(vec!["gpt-4.1".to_string()]),
ip_rules: Some(vec!["203.0.113.10".to_string()]),
rate_limit: 100,
concurrent_limit: Some(5),
force_capabilities: Some(json!({"cache": true})),
@@ -1104,6 +1136,7 @@ mod tests {
name: Some("Updated User".to_string()),
rate_limit: Some(150),
concurrent_limit: Some(6),
ip_rules: Some(Some(vec!["10.0.0.0/24".to_string()])),
})
.await
.expect("user key should update")
@@ -1179,6 +1212,7 @@ mod tests {
allowed_providers: Some(vec!["openai".to_string()]),
allowed_api_formats: None,
allowed_models: None,
ip_rules: None,
rate_limit: None,
concurrent_limit: Some(2),
force_capabilities: None,
@@ -1205,6 +1239,7 @@ mod tests {
allowed_providers: Some(None),
allowed_api_formats: Some(Some(vec!["openai:responses".to_string()])),
allowed_models: Some(Some(vec!["gpt-4.1-mini".to_string()])),
ip_rules: None,
expires_at_present: true,
expires_at_unix_secs: Some(2_100_000_000),
auto_delete_on_expiry_present: true,

26
crates/aether-data/src/repository/auth/types.rs
View File

@@ -24,6 +24,7 @@ pub struct StoredAuthApiKeySnapshot {
pub api_key_allowed_providers: Option<Vec<String>>,
pub api_key_allowed_api_formats: Option<Vec<String>>,
pub api_key_allowed_models: Option<Vec<String>>,
pub api_key_ip_rules: Option<Vec<String>>,
}
impl StoredAuthApiKeySnapshot {
@@ -97,9 +98,18 @@ impl StoredAuthApiKeySnapshot {
api_key_allowed_models,
"api_keys.allowed_models",
)?,
api_key_ip_rules: None,
})
}
pub fn with_api_key_ip_rules(
mut self,
api_key_ip_rules: Option<serde_json::Value>,
) -> Result<Self, crate::DataLayerError> {
self.api_key_ip_rules = parse_string_list(api_key_ip_rules, "api_keys.ip_rules")?;
Ok(self)
}
pub fn is_currently_usable(&self, now_unix_secs: u64) -> bool {
if !self.user_is_active || self.user_is_deleted {
return false;
@@ -148,6 +158,7 @@ pub struct ResolvedAuthApiKeySnapshot {
pub api_key_allowed_providers: Option<Vec<String>>,
pub api_key_allowed_api_formats: Option<Vec<String>>,
pub api_key_allowed_models: Option<Vec<String>>,
pub api_key_ip_rules: Option<Vec<String>>,
pub currently_usable: bool,
}
@@ -177,6 +188,7 @@ impl ResolvedAuthApiKeySnapshot {
api_key_allowed_providers: snapshot.api_key_allowed_providers,
api_key_allowed_api_formats: snapshot.api_key_allowed_api_formats,
api_key_allowed_models: snapshot.api_key_allowed_models,
api_key_ip_rules: snapshot.api_key_ip_rules,
currently_usable,
};
resolved.constrain_non_standalone_api_key_policy_to_user_policy();
@@ -332,6 +344,7 @@ pub struct StoredAuthApiKeyExportRecord {
pub allowed_providers: Option<Vec<String>>,
pub allowed_api_formats: Option<Vec<String>>,
pub allowed_models: Option<Vec<String>>,
pub ip_rules: Option<Vec<String>>,
pub rate_limit: Option<i32>,
pub concurrent_limit: Option<i32>,
pub force_capabilities: Option<serde_json::Value>,
@@ -403,6 +416,7 @@ impl StoredAuthApiKeyExportRecord {
"api_keys.allowed_api_formats",
)?,
allowed_models: parse_string_list(allowed_models, "api_keys.allowed_models")?,
ip_rules: None,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -444,6 +458,14 @@ impl StoredAuthApiKeyExportRecord {
self.feature_settings = normalize_optional_json(feature_settings);
self
}
pub fn with_ip_rules(
mut self,
ip_rules: Option<serde_json::Value>,
) -> Result<Self, crate::DataLayerError> {
self.ip_rules = parse_string_list(ip_rules, "api_keys.ip_rules")?;
Ok(self)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
@@ -469,6 +491,7 @@ pub struct CreateUserApiKeyRecord {
pub allowed_providers: Option<Vec<String>>,
pub allowed_api_formats: Option<Vec<String>>,
pub allowed_models: Option<Vec<String>>,
pub ip_rules: Option<Vec<String>>,
pub rate_limit: i32,
pub concurrent_limit: Option<i32>,
pub force_capabilities: Option<serde_json::Value>,
@@ -487,6 +510,7 @@ pub struct UpdateUserApiKeyBasicRecord {
pub name: Option<String>,
pub rate_limit: Option<i32>,
pub concurrent_limit: Option<i32>,
pub ip_rules: Option<Option<Vec<String>>>,
}
#[derive(Debug, Clone, PartialEq)]
@@ -499,6 +523,7 @@ pub struct CreateStandaloneApiKeyRecord {
pub allowed_providers: Option<Vec<String>>,
pub allowed_api_formats: Option<Vec<String>>,
pub allowed_models: Option<Vec<String>>,
pub ip_rules: Option<Vec<String>>,
pub rate_limit: Option<i32>,
pub concurrent_limit: Option<i32>,
pub force_capabilities: Option<serde_json::Value>,
@@ -521,6 +546,7 @@ pub struct UpdateStandaloneApiKeyBasicRecord {
pub allowed_providers: Option<Option<Vec<String>>>,
pub allowed_api_formats: Option<Option<Vec<String>>>,
pub allowed_models: Option<Option<Vec<String>>>,
pub ip_rules: Option<Option<Vec<String>>>,
pub expires_at_present: bool,
pub expires_at_unix_secs: Option<u64>,
pub auto_delete_on_expiry_present: bool,

2
crates/aether-data/src/repository/management_tokens/postgres.rs
View File

@@ -561,7 +561,7 @@ fn map_management_token_write_error(
.unwrap_or_else(|| "Management Token 名称已存在".to_string()),
),
(Some("23514"), Some("check_allowed_ips_not_empty")) => {
Some("IP 白名单不能为空，如需取消限制请不提供此字段".to_string())
Some("IP 限制规则不能为空，如需取消限制请不提供此字段".to_string())
}
_ => None,
}

12
crates/aether-data/src/repository/management_tokens/types.rs
View File

@@ -208,17 +208,17 @@ impl CreateManagementTokenRecord {
if let Some(allowed_ips) = &self.allowed_ips {
let Some(items) = allowed_ips.as_array() else {
return Err(crate::DataLayerError::InvalidInput(
"allowed_ips must be an array".to_string(),
"IP 限制规则必须是数组".to_string(),
));
};
if items.is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"allowed_ips must not be empty".to_string(),
"IP 限制规则不能为空".to_string(),
));
}
if items.iter().any(|value| value.as_str().is_none()) {
return Err(crate::DataLayerError::InvalidInput(
"allowed_ips must contain only strings".to_string(),
"IP 限制规则只能包含字符串".to_string(),
));
}
}
@@ -258,17 +258,17 @@ impl UpdateManagementTokenRecord {
if let Some(allowed_ips) = &self.allowed_ips {
let Some(items) = allowed_ips.as_array() else {
return Err(crate::DataLayerError::InvalidInput(
"allowed_ips must be an array".to_string(),
"IP 限制规则必须是数组".to_string(),
));
};
if items.is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"allowed_ips must not be empty".to_string(),
"IP 限制规则不能为空".to_string(),
));
}
if items.iter().any(|value| value.as_str().is_none()) {
return Err(crate::DataLayerError::InvalidInput(
"allowed_ips must contain only strings".to_string(),
"IP 限制规则只能包含字符串".to_string(),
));
}
}

104
crates/aether-data/src/repository/users/memory.rs
View File

@@ -360,6 +360,58 @@ fn memory_group_members(
.collect()
}
fn filter_memory_export_rows(
repository: &InMemoryUserReadRepository,
query: &UserExportListQuery,
) -> Vec<StoredUserExportRow> {
let mut rows = repository
.export_rows
.read()
.expect("user repository lock")
.clone();
if let Some(role) = query.role.as_deref() {
rows.retain(|row| row.role.eq_ignore_ascii_case(role));
}
if let Some(is_active) = query.is_active {
rows.retain(|row| row.is_active == is_active);
}
if let Some(group_id) = query
.group_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
let member_ids = repository
.group_members
.read()
.expect("user repository lock")
.keys()
.filter(|(candidate_group_id, _)| candidate_group_id == group_id)
.map(|(_, user_id)| user_id.clone())
.collect::<std::collections::BTreeSet<_>>();
rows.retain(|row| member_ids.contains(&row.id));
}
if let Some(search) = query
.search
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
let search = search.to_ascii_lowercase();
rows.retain(|row| {
row.id.to_ascii_lowercase().contains(&search)
|| row.username.to_ascii_lowercase().contains(&search)
|| row
.email
.as_deref()
.unwrap_or_default()
.to_ascii_lowercase()
.contains(&search)
});
}
rows
}
fn memory_export_row_from_auth_user(
repository: &InMemoryUserReadRepository,
user: &StoredUserAuthRecord,
@@ -477,59 +529,17 @@ impl UserReadRepository for InMemoryUserReadRepository {
&self,
query: &UserExportListQuery,
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
let mut rows = self
.export_rows
.read()
.expect("user repository lock")
.clone();
if let Some(role) = query.role.as_deref() {
rows.retain(|row| row.role.eq_ignore_ascii_case(role));
}
if let Some(is_active) = query.is_active {
rows.retain(|row| row.is_active == is_active);
}
if let Some(group_id) = query
.group_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
let member_ids = self
.group_members
.read()
.expect("user repository lock")
.keys()
.filter_map(|(candidate_group_id, user_id)| {
(candidate_group_id == group_id).then(|| user_id.clone())
})
.collect::<std::collections::BTreeSet<_>>();
rows.retain(|row| member_ids.contains(&row.id));
}
if let Some(search) = query
.search
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
let search = search.to_ascii_lowercase();
rows.retain(|row| {
row.id.to_ascii_lowercase().contains(&search)
|| row.username.to_ascii_lowercase().contains(&search)
|| row
.email
.as_deref()
.unwrap_or_default()
.to_ascii_lowercase()
.contains(&search)
});
}
Ok(rows
Ok(filter_memory_export_rows(self, query)
.into_iter()
.skip(query.skip)
.take(query.limit)
.collect())
}
async fn count_export_users(&self, query: &UserExportListQuery) -> Result<u64, DataLayerError> {
Ok(filter_memory_export_rows(self, query).len() as u64)
}
async fn summarize_export_users(&self) -> Result<UserExportSummary, DataLayerError> {
let rows = self.export_rows.read().expect("user repository lock");
Ok(UserExportSummary {

42
crates/aether-data/src/repository/users/mysql.rs
View File

@@ -325,6 +325,48 @@ impl UserReadRepository for MysqlUserReadRepository {
self.fetch_export_rows(builder).await
}
async fn count_export_users(&self, query: &UserExportListQuery) -> Result<u64, DataLayerError> {
let mut builder = QueryBuilder::<MySql>::new("SELECT COUNT(*) AS total FROM users");
builder.push(" WHERE is_deleted = 0");
if let Some(role) = query.role.as_deref() {
builder
.push(" AND LOWER(role) = ")
.push_bind(role.trim().to_ascii_lowercase());
}
if let Some(is_active) = query.is_active {
builder.push(" AND is_active = ").push_bind(is_active);
}
if let Some(group_id) = query
.group_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
builder.push(" AND id IN (SELECT user_id FROM user_group_members WHERE group_id = ");
builder.push_bind(group_id);
builder.push(")");
}
if let Some(search) = query
.search
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
let pattern = format!("%{}%", search.to_ascii_lowercase());
builder
.push(" AND (LOWER(id) LIKE ")
.push_bind(pattern.clone())
.push(" OR LOWER(username) LIKE ")
.push_bind(pattern.clone())
.push(" OR LOWER(COALESCE(email, '')) LIKE ")
.push_bind(pattern)
.push(")");
}
let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?;
Ok(row.try_get::<i64, _>("total").map_sql_err()?.max(0) as u64)
}
async fn summarize_export_users(&self) -> Result<UserExportSummary, DataLayerError> {
let row = sqlx::query(
r#"

55
crates/aether-data/src/repository/users/postgres.rs
View File

@@ -1015,6 +1015,57 @@ WHERE user_group_members.user_id IN (
collect_query_rows(query.fetch(&self.pool), map_user_export_row).await
}
pub async fn count_export_users(
&self,
query: &UserExportListQuery,
) -> Result<u64, DataLayerError> {
let mut builder =
QueryBuilder::<Postgres>::new("SELECT COUNT(*)::BIGINT AS total FROM users");
builder.push(" WHERE is_deleted IS FALSE");
if let Some(role) = query.role.as_deref() {
builder
.push(" AND LOWER(role::text) = ")
.push_bind(role.trim().to_ascii_lowercase());
}
if let Some(is_active) = query.is_active {
builder.push(" AND is_active = ").push_bind(is_active);
}
if let Some(group_id) = query
.group_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
builder.push(" AND id IN (SELECT user_id FROM user_group_members WHERE group_id = ");
builder.push_bind(group_id);
builder.push(")");
}
if let Some(search) = query
.search
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
let pattern = format!("%{}%", search.to_ascii_lowercase());
builder
.push(" AND (LOWER(id) LIKE ")
.push_bind(pattern.clone())
.push(" OR LOWER(username) LIKE ")
.push_bind(pattern.clone())
.push(" OR LOWER(COALESCE(email, '')) LIKE ")
.push_bind(pattern)
.push(")");
}
let row = builder
.build()
.fetch_one(&self.pool)
.await
.map_postgres_err()?;
Ok(row.try_get::<i64, _>("total").map_postgres_err()?.max(0) as u64)
}
pub async fn summarize_export_users(&self) -> Result<UserExportSummary, DataLayerError> {
let row = sqlx::query(SUMMARIZE_EXPORT_USERS_SQL)
.fetch_one(&self.pool)
@@ -2306,6 +2357,10 @@ impl UserReadRepository for SqlxUserReadRepository {
self.list_export_users_page(query).await
}
async fn count_export_users(&self, query: &UserExportListQuery) -> Result<u64, DataLayerError> {
self.count_export_users(query).await
}
async fn summarize_export_users(&self) -> Result<UserExportSummary, DataLayerError> {
self.summarize_export_users().await
}

42
crates/aether-data/src/repository/users/sqlite.rs
View File

@@ -325,6 +325,48 @@ impl UserReadRepository for SqliteUserReadRepository {
self.fetch_export_rows(builder).await
}
async fn count_export_users(&self, query: &UserExportListQuery) -> Result<u64, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new("SELECT COUNT(*) AS total FROM users");
builder.push(" WHERE is_deleted = 0");
if let Some(role) = query.role.as_deref() {
builder
.push(" AND LOWER(role) = ")
.push_bind(role.trim().to_ascii_lowercase());
}
if let Some(is_active) = query.is_active {
builder.push(" AND is_active = ").push_bind(is_active);
}
if let Some(group_id) = query
.group_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
builder.push(" AND id IN (SELECT user_id FROM user_group_members WHERE group_id = ");
builder.push_bind(group_id);
builder.push(")");
}
if let Some(search) = query
.search
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
let pattern = format!("%{}%", search.to_ascii_lowercase());
builder
.push(" AND (LOWER(id) LIKE ")
.push_bind(pattern.clone())
.push(" OR LOWER(username) LIKE ")
.push_bind(pattern.clone())
.push(" OR LOWER(COALESCE(email, '')) LIKE ")
.push_bind(pattern)
.push(")");
}
let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?;
Ok(row.try_get::<i64, _>("total").map_sql_err()?.max(0) as u64)
}
async fn summarize_export_users(&self) -> Result<UserExportSummary, DataLayerError> {
let row = sqlx::query(
r#"

5
crates/aether-data/src/repository/users/types.rs
View File

@@ -669,6 +669,11 @@ pub trait UserReadRepository: Send + Sync {
query: &UserExportListQuery,
) -> Result<Vec<StoredUserExportRow>, crate::DataLayerError>;
async fn count_export_users(
&self,
query: &UserExportListQuery,
) -> Result<u64, crate::DataLayerError>;
async fn summarize_export_users(&self) -> Result<UserExportSummary, crate::DataLayerError>;
async fn find_export_user_by_id(