mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
feat: 扩展 Rust gateway 全功能模块,新增 billing/crypto/wallet crate 及完整数据层
- 新增 aether-billing、aether-crypto、aether-wallet 独立 crate - aether-data 扩展 repository 层:announcements、auth_modules、billing、 candidate_selection、gemini_file_mappings、global_models、management_tokens、 oauth_providers、proxy_nodes、quota、users、wallet 等模块 - aether-gateway 新增 api/auth/billing/control/middleware/scheduler/usage/ video_tasks/hooks/maintenance/model_fetch/provider_transport 等功能模块 - 重构 executor decision 和 gateway state 为模块目录结构 - 新增 gateway router、frontdoor 路由层及对应测试 - Python 侧 API 路由重构,新增 compat/support 模块 - 前端 Logo 组件更新及 Provider 管理页面调整
This commit is contained in:
342
crates/aether-data/src/repository/management_tokens/memory.rs
Normal file
342
crates/aether-data/src/repository/management_tokens/memory.rs
Normal file
@@ -0,0 +1,342 @@
|
||||
use std::sync::RwLock;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
CreateManagementTokenRecord, ManagementTokenListQuery, ManagementTokenReadRepository,
|
||||
ManagementTokenWriteRepository, RegenerateManagementTokenSecret, StoredManagementToken,
|
||||
StoredManagementTokenListPage, StoredManagementTokenWithUser, UpdateManagementTokenRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryManagementTokenRepository {
|
||||
items: RwLock<Vec<StoredManagementTokenWithUser>>,
|
||||
}
|
||||
|
||||
impl InMemoryManagementTokenRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredManagementTokenWithUser>,
|
||||
{
|
||||
Self {
|
||||
items: RwLock::new(items.into_iter().collect()),
|
||||
}
|
||||
}
|
||||
|
||||
fn now_unix_secs() -> Option<u64> {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ManagementTokenReadRepository for InMemoryManagementTokenRepository {
|
||||
async fn list_management_tokens(
|
||||
&self,
|
||||
query: &ManagementTokenListQuery,
|
||||
) -> Result<StoredManagementTokenListPage, DataLayerError> {
|
||||
let items = self.items.read().expect("management token repository lock");
|
||||
let mut filtered = items
|
||||
.iter()
|
||||
.filter(|item| match query.user_id.as_deref() {
|
||||
Some(user_id) => item.token.user_id == user_id,
|
||||
None => true,
|
||||
})
|
||||
.filter(|item| match query.is_active {
|
||||
Some(is_active) => item.token.is_active == is_active,
|
||||
None => true,
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
filtered.sort_by(|left, right| {
|
||||
right
|
||||
.token
|
||||
.created_at_unix_secs
|
||||
.cmp(&left.token.created_at_unix_secs)
|
||||
.then_with(|| right.token.id.cmp(&left.token.id))
|
||||
});
|
||||
|
||||
let total = filtered.len();
|
||||
let items = filtered
|
||||
.into_iter()
|
||||
.skip(query.offset)
|
||||
.take(query.limit)
|
||||
.collect();
|
||||
Ok(StoredManagementTokenListPage { items, total })
|
||||
}
|
||||
|
||||
async fn get_management_token_with_user(
|
||||
&self,
|
||||
token_id: &str,
|
||||
) -> Result<Option<StoredManagementTokenWithUser>, DataLayerError> {
|
||||
let items = self.items.read().expect("management token repository lock");
|
||||
Ok(items.iter().find(|item| item.token.id == token_id).cloned())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ManagementTokenWriteRepository for InMemoryManagementTokenRepository {
|
||||
async fn create_management_token(
|
||||
&self,
|
||||
record: &CreateManagementTokenRecord,
|
||||
) -> Result<StoredManagementToken, DataLayerError> {
|
||||
record.validate()?;
|
||||
|
||||
let mut items = self
|
||||
.items
|
||||
.write()
|
||||
.expect("management token repository lock");
|
||||
if items
|
||||
.iter()
|
||||
.any(|item| item.token.user_id == record.user_id && item.token.name == record.name)
|
||||
{
|
||||
return Err(DataLayerError::InvalidInput(format!(
|
||||
"已存在名为 '{}' 的 Token",
|
||||
record.name
|
||||
)));
|
||||
}
|
||||
|
||||
let now = Self::now_unix_secs();
|
||||
let token = StoredManagementToken::new(
|
||||
record.id.clone(),
|
||||
record.user_id.clone(),
|
||||
record.name.clone(),
|
||||
)?
|
||||
.with_display_fields(
|
||||
record.description.clone(),
|
||||
record.token_prefix.clone(),
|
||||
record.allowed_ips.clone(),
|
||||
)
|
||||
.with_runtime_fields(record.expires_at_unix_secs, None, None, 0, record.is_active)
|
||||
.with_timestamps(now, now);
|
||||
items.push(StoredManagementTokenWithUser::new(
|
||||
token.clone(),
|
||||
record.user.clone(),
|
||||
));
|
||||
Ok(token)
|
||||
}
|
||||
|
||||
async fn update_management_token(
|
||||
&self,
|
||||
record: &UpdateManagementTokenRecord,
|
||||
) -> Result<Option<StoredManagementToken>, DataLayerError> {
|
||||
record.validate()?;
|
||||
|
||||
let mut items = self
|
||||
.items
|
||||
.write()
|
||||
.expect("management token repository lock");
|
||||
let Some(index) = items
|
||||
.iter()
|
||||
.position(|item| item.token.id == record.token_id)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if let Some(name) = &record.name {
|
||||
if items.iter().enumerate().any(|(position, item)| {
|
||||
position != index
|
||||
&& item.token.user_id == items[index].token.user_id
|
||||
&& item.token.name == *name
|
||||
}) {
|
||||
return Err(DataLayerError::InvalidInput(format!(
|
||||
"已存在名为 '{}' 的 Token",
|
||||
name
|
||||
)));
|
||||
}
|
||||
items[index].token.name = name.clone();
|
||||
}
|
||||
|
||||
if record.clear_description {
|
||||
items[index].token.description = None;
|
||||
} else if let Some(description) = &record.description {
|
||||
items[index].token.description = Some(description.clone());
|
||||
}
|
||||
|
||||
if record.clear_allowed_ips {
|
||||
items[index].token.allowed_ips = None;
|
||||
} else if let Some(allowed_ips) = &record.allowed_ips {
|
||||
items[index].token.allowed_ips = Some(allowed_ips.clone());
|
||||
}
|
||||
|
||||
if record.clear_expires_at {
|
||||
items[index].token.expires_at_unix_secs = None;
|
||||
} else if let Some(expires_at_unix_secs) = record.expires_at_unix_secs {
|
||||
items[index].token.expires_at_unix_secs = Some(expires_at_unix_secs);
|
||||
}
|
||||
|
||||
if let Some(is_active) = record.is_active {
|
||||
items[index].token.is_active = is_active;
|
||||
}
|
||||
|
||||
items[index].token.updated_at_unix_secs = Self::now_unix_secs();
|
||||
Ok(Some(items[index].token.clone()))
|
||||
}
|
||||
|
||||
async fn delete_management_token(&self, token_id: &str) -> Result<bool, DataLayerError> {
|
||||
let mut items = self
|
||||
.items
|
||||
.write()
|
||||
.expect("management token repository lock");
|
||||
let original_len = items.len();
|
||||
items.retain(|item| item.token.id != token_id);
|
||||
Ok(items.len() != original_len)
|
||||
}
|
||||
|
||||
async fn set_management_token_active(
|
||||
&self,
|
||||
token_id: &str,
|
||||
is_active: bool,
|
||||
) -> Result<Option<StoredManagementToken>, DataLayerError> {
|
||||
let mut items = self
|
||||
.items
|
||||
.write()
|
||||
.expect("management token repository lock");
|
||||
let Some(item) = items.iter_mut().find(|item| item.token.id == token_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
item.token.is_active = is_active;
|
||||
item.token.updated_at_unix_secs = Self::now_unix_secs();
|
||||
Ok(Some(item.token.clone()))
|
||||
}
|
||||
|
||||
async fn regenerate_management_token_secret(
|
||||
&self,
|
||||
mutation: &RegenerateManagementTokenSecret,
|
||||
) -> Result<Option<StoredManagementToken>, DataLayerError> {
|
||||
mutation.validate()?;
|
||||
|
||||
let mut items = self
|
||||
.items
|
||||
.write()
|
||||
.expect("management token repository lock");
|
||||
let Some(item) = items
|
||||
.iter_mut()
|
||||
.find(|item| item.token.id == mutation.token_id)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
item.token.token_prefix = mutation.token_prefix.clone();
|
||||
item.token.updated_at_unix_secs = Self::now_unix_secs();
|
||||
Ok(Some(item.token.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryManagementTokenRepository;
|
||||
use crate::repository::management_tokens::{
|
||||
CreateManagementTokenRecord, ManagementTokenListQuery, ManagementTokenReadRepository,
|
||||
ManagementTokenWriteRepository, RegenerateManagementTokenSecret, StoredManagementToken,
|
||||
StoredManagementTokenUserSummary, StoredManagementTokenWithUser,
|
||||
UpdateManagementTokenRecord,
|
||||
};
|
||||
|
||||
fn sample_token(id: &str, user_id: &str, is_active: bool) -> StoredManagementTokenWithUser {
|
||||
let token = StoredManagementToken::new(id.to_string(), user_id.to_string(), id.to_string())
|
||||
.expect("token should build")
|
||||
.with_runtime_fields(None, None, None, 2, is_active)
|
||||
.with_timestamps(Some(1_700_000_000), Some(1_700_000_100));
|
||||
let user = StoredManagementTokenUserSummary::new(
|
||||
user_id.to_string(),
|
||||
Some(format!("{user_id}@example.com")),
|
||||
format!("{user_id}-name"),
|
||||
"admin".to_string(),
|
||||
)
|
||||
.expect("user should build");
|
||||
StoredManagementTokenWithUser::new(token, user)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lists_filters_and_mutates_management_tokens() {
|
||||
let repository = InMemoryManagementTokenRepository::seed(vec![
|
||||
sample_token("token-1", "user-1", true),
|
||||
sample_token("token-2", "user-2", false),
|
||||
]);
|
||||
|
||||
let page = repository
|
||||
.list_management_tokens(&ManagementTokenListQuery {
|
||||
user_id: None,
|
||||
is_active: Some(true),
|
||||
offset: 0,
|
||||
limit: 10,
|
||||
})
|
||||
.await
|
||||
.expect("list should succeed");
|
||||
assert_eq!(page.total, 1);
|
||||
assert_eq!(page.items[0].token.id, "token-1");
|
||||
|
||||
let toggled = repository
|
||||
.set_management_token_active("token-2", true)
|
||||
.await
|
||||
.expect("toggle should succeed")
|
||||
.expect("token should exist");
|
||||
assert!(toggled.is_active);
|
||||
|
||||
let created = repository
|
||||
.create_management_token(&CreateManagementTokenRecord {
|
||||
id: "token-3".to_string(),
|
||||
user_id: "user-1".to_string(),
|
||||
user: StoredManagementTokenUserSummary::new(
|
||||
"user-1".to_string(),
|
||||
Some("user-1@example.com".to_string()),
|
||||
"user-1-name".to_string(),
|
||||
"user".to_string(),
|
||||
)
|
||||
.expect("user should build"),
|
||||
token_hash: "hash-3".to_string(),
|
||||
token_prefix: Some("ae_1234".to_string()),
|
||||
name: "created".to_string(),
|
||||
description: Some("created token".to_string()),
|
||||
allowed_ips: Some(serde_json::json!(["127.0.0.1"])),
|
||||
expires_at_unix_secs: Some(1_800_000_000),
|
||||
is_active: true,
|
||||
})
|
||||
.await
|
||||
.expect("create should succeed");
|
||||
assert_eq!(created.name, "created");
|
||||
|
||||
let updated = repository
|
||||
.update_management_token(&UpdateManagementTokenRecord {
|
||||
token_id: "token-3".to_string(),
|
||||
name: Some("renamed".to_string()),
|
||||
description: None,
|
||||
clear_description: true,
|
||||
allowed_ips: Some(serde_json::json!(["10.0.0.1"])),
|
||||
clear_allowed_ips: false,
|
||||
expires_at_unix_secs: None,
|
||||
clear_expires_at: true,
|
||||
is_active: Some(false),
|
||||
})
|
||||
.await
|
||||
.expect("update should succeed")
|
||||
.expect("token should exist");
|
||||
assert_eq!(updated.name, "renamed");
|
||||
assert_eq!(updated.description, None);
|
||||
assert_eq!(updated.allowed_ips, Some(serde_json::json!(["10.0.0.1"])));
|
||||
assert_eq!(updated.expires_at_unix_secs, None);
|
||||
assert!(!updated.is_active);
|
||||
|
||||
let regenerated = repository
|
||||
.regenerate_management_token_secret(&RegenerateManagementTokenSecret {
|
||||
token_id: "token-3".to_string(),
|
||||
token_hash: "hash-3b".to_string(),
|
||||
token_prefix: Some("ae_5678".to_string()),
|
||||
})
|
||||
.await
|
||||
.expect("regenerate should succeed")
|
||||
.expect("token should exist");
|
||||
assert_eq!(regenerated.token_prefix.as_deref(), Some("ae_5678"));
|
||||
|
||||
let deleted = repository
|
||||
.delete_management_token("token-1")
|
||||
.await
|
||||
.expect("delete should succeed");
|
||||
assert!(deleted);
|
||||
}
|
||||
}
|
||||
12
crates/aether-data/src/repository/management_tokens/mod.rs
Normal file
12
crates/aether-data/src/repository/management_tokens/mod.rs
Normal file
@@ -0,0 +1,12 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryManagementTokenRepository;
|
||||
pub use sql::SqlxManagementTokenRepository;
|
||||
pub use types::{
|
||||
CreateManagementTokenRecord, ManagementTokenListQuery, ManagementTokenReadRepository,
|
||||
ManagementTokenWriteRepository, RegenerateManagementTokenSecret, StoredManagementToken,
|
||||
StoredManagementTokenListPage, StoredManagementTokenUserSummary, StoredManagementTokenWithUser,
|
||||
UpdateManagementTokenRecord,
|
||||
};
|
||||
432
crates/aether-data/src/repository/management_tokens/sql.rs
Normal file
432
crates/aether-data/src/repository/management_tokens/sql.rs
Normal file
@@ -0,0 +1,432 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
CreateManagementTokenRecord, ManagementTokenListQuery, ManagementTokenReadRepository,
|
||||
ManagementTokenWriteRepository, RegenerateManagementTokenSecret, StoredManagementToken,
|
||||
StoredManagementTokenListPage, StoredManagementTokenUserSummary, StoredManagementTokenWithUser,
|
||||
UpdateManagementTokenRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const LIST_MANAGEMENT_TOKENS_SQL: &str = r#"
|
||||
SELECT
|
||||
mt.id,
|
||||
mt.user_id,
|
||||
mt.name,
|
||||
mt.description,
|
||||
mt.token_prefix,
|
||||
mt.allowed_ips,
|
||||
EXTRACT(EPOCH FROM mt.expires_at)::bigint AS expires_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM mt.last_used_at)::bigint AS last_used_at_unix_secs,
|
||||
mt.last_used_ip,
|
||||
COALESCE(mt.usage_count, 0) AS usage_count,
|
||||
mt.is_active,
|
||||
EXTRACT(EPOCH FROM mt.created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM mt.updated_at)::bigint AS updated_at_unix_secs,
|
||||
u.id AS user_row_id,
|
||||
u.email AS user_email,
|
||||
u.username AS user_username,
|
||||
u.role::text AS user_role
|
||||
FROM management_tokens mt
|
||||
JOIN users u ON u.id = mt.user_id
|
||||
WHERE ($1::text IS NULL OR mt.user_id = $1)
|
||||
AND ($2::boolean IS NULL OR mt.is_active = $2)
|
||||
ORDER BY mt.created_at DESC, mt.id DESC
|
||||
OFFSET $3
|
||||
LIMIT $4
|
||||
"#;
|
||||
|
||||
const COUNT_MANAGEMENT_TOKENS_SQL: &str = r#"
|
||||
SELECT COUNT(mt.id) AS total
|
||||
FROM management_tokens mt
|
||||
WHERE ($1::text IS NULL OR mt.user_id = $1)
|
||||
AND ($2::boolean IS NULL OR mt.is_active = $2)
|
||||
"#;
|
||||
|
||||
const GET_MANAGEMENT_TOKEN_WITH_USER_SQL: &str = r#"
|
||||
SELECT
|
||||
mt.id,
|
||||
mt.user_id,
|
||||
mt.name,
|
||||
mt.description,
|
||||
mt.token_prefix,
|
||||
mt.allowed_ips,
|
||||
EXTRACT(EPOCH FROM mt.expires_at)::bigint AS expires_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM mt.last_used_at)::bigint AS last_used_at_unix_secs,
|
||||
mt.last_used_ip,
|
||||
COALESCE(mt.usage_count, 0) AS usage_count,
|
||||
mt.is_active,
|
||||
EXTRACT(EPOCH FROM mt.created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM mt.updated_at)::bigint AS updated_at_unix_secs,
|
||||
u.id AS user_row_id,
|
||||
u.email AS user_email,
|
||||
u.username AS user_username,
|
||||
u.role::text AS user_role
|
||||
FROM management_tokens mt
|
||||
JOIN users u ON u.id = mt.user_id
|
||||
WHERE mt.id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const DELETE_MANAGEMENT_TOKEN_SQL: &str = r#"
|
||||
DELETE FROM management_tokens
|
||||
WHERE id = $1
|
||||
"#;
|
||||
|
||||
const CREATE_MANAGEMENT_TOKEN_SQL: &str = r#"
|
||||
INSERT INTO management_tokens (
|
||||
id,
|
||||
user_id,
|
||||
token_hash,
|
||||
token_prefix,
|
||||
name,
|
||||
description,
|
||||
allowed_ips,
|
||||
expires_at,
|
||||
is_active
|
||||
)
|
||||
VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4,
|
||||
$5,
|
||||
$6,
|
||||
$7,
|
||||
CASE
|
||||
WHEN $8::bigint IS NULL THEN NULL
|
||||
ELSE to_timestamp($8::double precision)
|
||||
END,
|
||||
$9
|
||||
)
|
||||
RETURNING
|
||||
id,
|
||||
user_id,
|
||||
name,
|
||||
description,
|
||||
token_prefix,
|
||||
allowed_ips,
|
||||
EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM last_used_at)::bigint AS last_used_at_unix_secs,
|
||||
last_used_ip,
|
||||
COALESCE(usage_count, 0) AS usage_count,
|
||||
is_active,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const UPDATE_MANAGEMENT_TOKEN_SQL: &str = r#"
|
||||
UPDATE management_tokens
|
||||
SET name = COALESCE($2, name),
|
||||
description = CASE
|
||||
WHEN $3 THEN NULL
|
||||
WHEN $4::text IS NULL THEN description
|
||||
ELSE $4
|
||||
END,
|
||||
allowed_ips = CASE
|
||||
WHEN $5 THEN NULL
|
||||
WHEN $6::json IS NULL THEN allowed_ips
|
||||
ELSE $6
|
||||
END,
|
||||
expires_at = CASE
|
||||
WHEN $7 THEN NULL
|
||||
WHEN $8::bigint IS NULL THEN expires_at
|
||||
ELSE to_timestamp($8::double precision)
|
||||
END,
|
||||
is_active = COALESCE($9, is_active),
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
RETURNING
|
||||
id,
|
||||
user_id,
|
||||
name,
|
||||
description,
|
||||
token_prefix,
|
||||
allowed_ips,
|
||||
EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM last_used_at)::bigint AS last_used_at_unix_secs,
|
||||
last_used_ip,
|
||||
COALESCE(usage_count, 0) AS usage_count,
|
||||
is_active,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const SET_MANAGEMENT_TOKEN_ACTIVE_SQL: &str = r#"
|
||||
UPDATE management_tokens
|
||||
SET is_active = $2,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
RETURNING
|
||||
id,
|
||||
user_id,
|
||||
name,
|
||||
description,
|
||||
token_prefix,
|
||||
allowed_ips,
|
||||
EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM last_used_at)::bigint AS last_used_at_unix_secs,
|
||||
last_used_ip,
|
||||
COALESCE(usage_count, 0) AS usage_count,
|
||||
is_active,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const REGENERATE_MANAGEMENT_TOKEN_SECRET_SQL: &str = r#"
|
||||
UPDATE management_tokens
|
||||
SET token_hash = $2,
|
||||
token_prefix = $3,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
RETURNING
|
||||
id,
|
||||
user_id,
|
||||
name,
|
||||
description,
|
||||
token_prefix,
|
||||
allowed_ips,
|
||||
EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM last_used_at)::bigint AS last_used_at_unix_secs,
|
||||
last_used_ip,
|
||||
COALESCE(usage_count, 0) AS usage_count,
|
||||
is_active,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxManagementTokenRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxManagementTokenRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ManagementTokenReadRepository for SqlxManagementTokenRepository {
|
||||
async fn list_management_tokens(
|
||||
&self,
|
||||
query: &ManagementTokenListQuery,
|
||||
) -> Result<StoredManagementTokenListPage, DataLayerError> {
|
||||
let count_row = sqlx::query(COUNT_MANAGEMENT_TOKENS_SQL)
|
||||
.bind(query.user_id.as_deref())
|
||||
.bind(query.is_active)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
let total = count_row.try_get::<i64, _>("total")?;
|
||||
|
||||
let rows = sqlx::query(LIST_MANAGEMENT_TOKENS_SQL)
|
||||
.bind(query.user_id.as_deref())
|
||||
.bind(query.is_active)
|
||||
.bind(i64::try_from(query.offset).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(query.limit).unwrap_or(i64::MAX))
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(StoredManagementTokenListPage {
|
||||
items: rows
|
||||
.iter()
|
||||
.map(map_token_with_user_row)
|
||||
.collect::<Result<Vec<_>, _>>()?,
|
||||
total: usize::try_from(total.max(0)).unwrap_or(usize::MAX),
|
||||
})
|
||||
}
|
||||
|
||||
async fn get_management_token_with_user(
|
||||
&self,
|
||||
token_id: &str,
|
||||
) -> Result<Option<StoredManagementTokenWithUser>, DataLayerError> {
|
||||
let row = sqlx::query(GET_MANAGEMENT_TOKEN_WITH_USER_SQL)
|
||||
.bind(token_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_token_with_user_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ManagementTokenWriteRepository for SqlxManagementTokenRepository {
|
||||
async fn create_management_token(
|
||||
&self,
|
||||
record: &CreateManagementTokenRecord,
|
||||
) -> Result<StoredManagementToken, DataLayerError> {
|
||||
record.validate()?;
|
||||
let row = sqlx::query(CREATE_MANAGEMENT_TOKEN_SQL)
|
||||
.bind(&record.id)
|
||||
.bind(&record.user_id)
|
||||
.bind(&record.token_hash)
|
||||
.bind(record.token_prefix.as_deref())
|
||||
.bind(&record.name)
|
||||
.bind(record.description.as_deref())
|
||||
.bind(record.allowed_ips.as_ref())
|
||||
.bind(
|
||||
record
|
||||
.expires_at_unix_secs
|
||||
.and_then(|value| i64::try_from(value).ok()),
|
||||
)
|
||||
.bind(record.is_active)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_err(|err| map_management_token_write_error(err, Some(record.name.as_str())))?;
|
||||
map_token_row(&row)
|
||||
}
|
||||
|
||||
async fn update_management_token(
|
||||
&self,
|
||||
record: &UpdateManagementTokenRecord,
|
||||
) -> Result<Option<StoredManagementToken>, DataLayerError> {
|
||||
record.validate()?;
|
||||
let row = sqlx::query(UPDATE_MANAGEMENT_TOKEN_SQL)
|
||||
.bind(&record.token_id)
|
||||
.bind(record.name.as_deref())
|
||||
.bind(record.clear_description)
|
||||
.bind(record.description.as_deref())
|
||||
.bind(record.clear_allowed_ips)
|
||||
.bind(record.allowed_ips.as_ref())
|
||||
.bind(record.clear_expires_at)
|
||||
.bind(
|
||||
record
|
||||
.expires_at_unix_secs
|
||||
.and_then(|value| i64::try_from(value).ok()),
|
||||
)
|
||||
.bind(record.is_active)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_err(|err| map_management_token_write_error(err, record.name.as_deref()))?;
|
||||
row.as_ref().map(map_token_row).transpose()
|
||||
}
|
||||
|
||||
async fn delete_management_token(&self, token_id: &str) -> Result<bool, DataLayerError> {
|
||||
let result = sqlx::query(DELETE_MANAGEMENT_TOKEN_SQL)
|
||||
.bind(token_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
async fn set_management_token_active(
|
||||
&self,
|
||||
token_id: &str,
|
||||
is_active: bool,
|
||||
) -> Result<Option<StoredManagementToken>, DataLayerError> {
|
||||
let row = sqlx::query(SET_MANAGEMENT_TOKEN_ACTIVE_SQL)
|
||||
.bind(token_id)
|
||||
.bind(is_active)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_token_row).transpose()
|
||||
}
|
||||
|
||||
async fn regenerate_management_token_secret(
|
||||
&self,
|
||||
mutation: &RegenerateManagementTokenSecret,
|
||||
) -> Result<Option<StoredManagementToken>, DataLayerError> {
|
||||
mutation.validate()?;
|
||||
let row = sqlx::query(REGENERATE_MANAGEMENT_TOKEN_SECRET_SQL)
|
||||
.bind(&mutation.token_id)
|
||||
.bind(&mutation.token_hash)
|
||||
.bind(mutation.token_prefix.as_deref())
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_token_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_unix_secs(value: Option<i64>) -> Option<u64> {
|
||||
value.and_then(|value| u64::try_from(value).ok())
|
||||
}
|
||||
|
||||
fn map_management_token_write_error(
|
||||
err: sqlx::Error,
|
||||
requested_name: Option<&str>,
|
||||
) -> DataLayerError {
|
||||
let conflict = err.as_database_error().and_then(|db_err| {
|
||||
let code = db_err.code().map(|value| value.as_ref().to_string());
|
||||
let constraint = db_err.constraint().map(|value| value.to_string());
|
||||
match (code.as_deref(), constraint.as_deref()) {
|
||||
(Some("23505"), Some("uq_management_tokens_user_name")) => Some(
|
||||
requested_name
|
||||
.map(|name| format!("已存在名为 '{}' 的 Token", name))
|
||||
.unwrap_or_else(|| "Management Token 名称已存在".to_string()),
|
||||
),
|
||||
(Some("23514"), Some("check_allowed_ips_not_empty")) => {
|
||||
Some("IP 白名单不能为空,如需取消限制请不提供此字段".to_string())
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
});
|
||||
|
||||
match conflict {
|
||||
Some(detail) => DataLayerError::InvalidInput(detail),
|
||||
None => DataLayerError::Postgres(err),
|
||||
}
|
||||
}
|
||||
|
||||
fn map_token_row(row: &PgRow) -> Result<StoredManagementToken, DataLayerError> {
|
||||
Ok(StoredManagementToken::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("user_id")?,
|
||||
row.try_get("name")?,
|
||||
)?
|
||||
.with_display_fields(
|
||||
row.try_get("description")?,
|
||||
row.try_get("token_prefix")?,
|
||||
row.try_get("allowed_ips")?,
|
||||
)
|
||||
.with_runtime_fields(
|
||||
optional_unix_secs(row.try_get("expires_at_unix_secs")?),
|
||||
optional_unix_secs(row.try_get("last_used_at_unix_secs")?),
|
||||
row.try_get("last_used_ip")?,
|
||||
u64::try_from(row.try_get::<i32, _>("usage_count")?).unwrap_or(0),
|
||||
row.try_get("is_active")?,
|
||||
)
|
||||
.with_timestamps(
|
||||
optional_unix_secs(row.try_get("created_at_unix_secs")?),
|
||||
optional_unix_secs(row.try_get("updated_at_unix_secs")?),
|
||||
))
|
||||
}
|
||||
|
||||
fn map_user_summary_row(row: &PgRow) -> Result<StoredManagementTokenUserSummary, DataLayerError> {
|
||||
StoredManagementTokenUserSummary::new(
|
||||
row.try_get("user_row_id")?,
|
||||
row.try_get("user_email")?,
|
||||
row.try_get("user_username")?,
|
||||
row.try_get("user_role")?,
|
||||
)
|
||||
}
|
||||
|
||||
fn map_token_with_user_row(row: &PgRow) -> Result<StoredManagementTokenWithUser, DataLayerError> {
|
||||
Ok(StoredManagementTokenWithUser::new(
|
||||
map_token_row(row)?,
|
||||
map_user_summary_row(row)?,
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxManagementTokenRepository;
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_constructs_from_lazy_pool() {
|
||||
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("factory should build");
|
||||
|
||||
let pool = factory.connect_lazy().expect("pool should build");
|
||||
let _repository = SqlxManagementTokenRepository::new(pool);
|
||||
}
|
||||
}
|
||||
335
crates/aether-data/src/repository/management_tokens/types.rs
Normal file
335
crates/aether-data/src/repository/management_tokens/types.rs
Normal file
@@ -0,0 +1,335 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredManagementTokenUserSummary {
|
||||
pub id: String,
|
||||
pub email: Option<String>,
|
||||
pub username: String,
|
||||
pub role: String,
|
||||
}
|
||||
|
||||
impl StoredManagementTokenUserSummary {
|
||||
pub fn new(
|
||||
id: String,
|
||||
email: Option<String>,
|
||||
username: String,
|
||||
role: String,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"users.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if username.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"users.username is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if role.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"users.role is empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(Self {
|
||||
id,
|
||||
email,
|
||||
username,
|
||||
role,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredManagementToken {
|
||||
pub id: String,
|
||||
pub user_id: String,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub token_prefix: Option<String>,
|
||||
pub allowed_ips: Option<serde_json::Value>,
|
||||
pub expires_at_unix_secs: Option<u64>,
|
||||
pub last_used_at_unix_secs: Option<u64>,
|
||||
pub last_used_ip: Option<String>,
|
||||
pub usage_count: u64,
|
||||
pub is_active: bool,
|
||||
pub created_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl StoredManagementToken {
|
||||
pub fn new(id: String, user_id: String, name: String) -> Result<Self, crate::DataLayerError> {
|
||||
if id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"management_tokens.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if user_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"management_tokens.user_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"management_tokens.name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(Self {
|
||||
id,
|
||||
user_id,
|
||||
name,
|
||||
description: None,
|
||||
token_prefix: None,
|
||||
allowed_ips: None,
|
||||
expires_at_unix_secs: None,
|
||||
last_used_at_unix_secs: None,
|
||||
last_used_ip: None,
|
||||
usage_count: 0,
|
||||
is_active: true,
|
||||
created_at_unix_secs: None,
|
||||
updated_at_unix_secs: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn with_display_fields(
|
||||
mut self,
|
||||
description: Option<String>,
|
||||
token_prefix: Option<String>,
|
||||
allowed_ips: Option<serde_json::Value>,
|
||||
) -> Self {
|
||||
self.description = description;
|
||||
self.token_prefix = token_prefix;
|
||||
self.allowed_ips = allowed_ips;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_runtime_fields(
|
||||
mut self,
|
||||
expires_at_unix_secs: Option<u64>,
|
||||
last_used_at_unix_secs: Option<u64>,
|
||||
last_used_ip: Option<String>,
|
||||
usage_count: u64,
|
||||
is_active: bool,
|
||||
) -> Self {
|
||||
self.expires_at_unix_secs = expires_at_unix_secs;
|
||||
self.last_used_at_unix_secs = last_used_at_unix_secs;
|
||||
self.last_used_ip = last_used_ip;
|
||||
self.usage_count = usage_count;
|
||||
self.is_active = is_active;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_timestamps(
|
||||
mut self,
|
||||
created_at_unix_secs: Option<u64>,
|
||||
updated_at_unix_secs: Option<u64>,
|
||||
) -> Self {
|
||||
self.created_at_unix_secs = created_at_unix_secs;
|
||||
self.updated_at_unix_secs = updated_at_unix_secs;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn token_display(&self) -> String {
|
||||
self.token_prefix
|
||||
.as_deref()
|
||||
.map(|prefix| format!("{prefix}...****"))
|
||||
.unwrap_or_else(|| "ae_****".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredManagementTokenWithUser {
|
||||
pub token: StoredManagementToken,
|
||||
pub user: StoredManagementTokenUserSummary,
|
||||
}
|
||||
|
||||
impl StoredManagementTokenWithUser {
|
||||
pub fn new(token: StoredManagementToken, user: StoredManagementTokenUserSummary) -> Self {
|
||||
Self { token, user }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct ManagementTokenListQuery {
|
||||
pub user_id: Option<String>,
|
||||
pub is_active: Option<bool>,
|
||||
pub offset: usize,
|
||||
pub limit: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct CreateManagementTokenRecord {
|
||||
pub id: String,
|
||||
pub user_id: String,
|
||||
pub user: StoredManagementTokenUserSummary,
|
||||
pub token_hash: String,
|
||||
pub token_prefix: Option<String>,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub allowed_ips: Option<serde_json::Value>,
|
||||
pub expires_at_unix_secs: Option<u64>,
|
||||
pub is_active: bool,
|
||||
}
|
||||
|
||||
impl CreateManagementTokenRecord {
|
||||
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||
if self.id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"token_id is required".to_string(),
|
||||
));
|
||||
}
|
||||
if self.user_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"user_id is required".to_string(),
|
||||
));
|
||||
}
|
||||
if self.user.id != self.user_id {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"management token user summary does not match user_id".to_string(),
|
||||
));
|
||||
}
|
||||
if self.token_hash.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"token_hash is required".to_string(),
|
||||
));
|
||||
}
|
||||
if self.name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"name is required".to_string(),
|
||||
));
|
||||
}
|
||||
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(),
|
||||
));
|
||||
};
|
||||
if items.is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"allowed_ips must not be empty".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(),
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct UpdateManagementTokenRecord {
|
||||
pub token_id: String,
|
||||
pub name: Option<String>,
|
||||
pub description: Option<String>,
|
||||
pub clear_description: bool,
|
||||
pub allowed_ips: Option<serde_json::Value>,
|
||||
pub clear_allowed_ips: bool,
|
||||
pub expires_at_unix_secs: Option<u64>,
|
||||
pub clear_expires_at: bool,
|
||||
pub is_active: Option<bool>,
|
||||
}
|
||||
|
||||
impl UpdateManagementTokenRecord {
|
||||
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||
if self.token_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"token_id is required".to_string(),
|
||||
));
|
||||
}
|
||||
if let Some(name) = &self.name {
|
||||
if name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"name must not be empty".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
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(),
|
||||
));
|
||||
};
|
||||
if items.is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"allowed_ips must not be empty".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(),
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct RegenerateManagementTokenSecret {
|
||||
pub token_id: String,
|
||||
pub token_hash: String,
|
||||
pub token_prefix: Option<String>,
|
||||
}
|
||||
|
||||
impl RegenerateManagementTokenSecret {
|
||||
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||
if self.token_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"token_id is required".to_string(),
|
||||
));
|
||||
}
|
||||
if self.token_hash.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"token_hash is required".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredManagementTokenListPage {
|
||||
pub items: Vec<StoredManagementTokenWithUser>,
|
||||
pub total: usize,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ManagementTokenReadRepository: Send + Sync {
|
||||
async fn list_management_tokens(
|
||||
&self,
|
||||
query: &ManagementTokenListQuery,
|
||||
) -> Result<StoredManagementTokenListPage, crate::DataLayerError>;
|
||||
|
||||
async fn get_management_token_with_user(
|
||||
&self,
|
||||
token_id: &str,
|
||||
) -> Result<Option<StoredManagementTokenWithUser>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ManagementTokenWriteRepository: Send + Sync {
|
||||
async fn create_management_token(
|
||||
&self,
|
||||
record: &CreateManagementTokenRecord,
|
||||
) -> Result<StoredManagementToken, crate::DataLayerError>;
|
||||
|
||||
async fn update_management_token(
|
||||
&self,
|
||||
record: &UpdateManagementTokenRecord,
|
||||
) -> Result<Option<StoredManagementToken>, crate::DataLayerError>;
|
||||
|
||||
async fn delete_management_token(&self, token_id: &str) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn set_management_token_active(
|
||||
&self,
|
||||
token_id: &str,
|
||||
is_active: bool,
|
||||
) -> Result<Option<StoredManagementToken>, crate::DataLayerError>;
|
||||
|
||||
async fn regenerate_management_token_secret(
|
||||
&self,
|
||||
mutation: &RegenerateManagementTokenSecret,
|
||||
) -> Result<Option<StoredManagementToken>, crate::DataLayerError>;
|
||||
}
|
||||
Reference in New Issue
Block a user