mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00: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:
12
crates/aether-billing/Cargo.toml
Normal file
12
crates/aether-billing/Cargo.toml
Normal file
@@ -0,0 +1,12 @@
|
||||
[package]
|
||||
name = "aether-billing"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Shared billing domain core for Aether Rust migration"
|
||||
|
||||
[dependencies]
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
216
crates/aether-billing/src/default_rule.rs
Normal file
216
crates/aether-billing/src/default_rule.rs
Normal file
@@ -0,0 +1,216 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::pricing::BillingModelPricingSnapshot;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct VirtualBillingRule {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub task_type: String,
|
||||
pub expression: String,
|
||||
pub variables: BTreeMap<String, Value>,
|
||||
pub dimension_mappings: BTreeMap<String, Value>,
|
||||
pub scope: String,
|
||||
}
|
||||
|
||||
pub struct DefaultBillingRuleGenerator;
|
||||
|
||||
impl DefaultBillingRuleGenerator {
|
||||
pub fn generate_for_pricing(
|
||||
pricing: &BillingModelPricingSnapshot,
|
||||
task_type: &str,
|
||||
) -> Option<VirtualBillingRule> {
|
||||
let tiers = pricing
|
||||
.effective_tiered_pricing()
|
||||
.and_then(|value| value.get("tiers"))
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
if tiers.is_empty() && pricing.effective_price_per_request().is_none() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let first_tier = tiers.first().cloned().unwrap_or_else(|| json!({}));
|
||||
let base_input_price = tier_value(&first_tier, "input_price_per_1m", 0.0);
|
||||
let base_output_price = tier_value(&first_tier, "output_price_per_1m", 0.0);
|
||||
let base_cache_creation_price =
|
||||
tier_value_with_fallback(&first_tier, "cache_creation_price_per_1m", 1.25);
|
||||
let base_cache_read_price =
|
||||
tier_value_with_fallback(&first_tier, "cache_read_price_per_1m", 0.1);
|
||||
let base_request_price = pricing.effective_price_per_request().unwrap_or(0.0);
|
||||
|
||||
let mut variables = BTreeMap::new();
|
||||
variables.insert("input_price_per_1m".to_string(), json!(base_input_price));
|
||||
variables.insert("output_price_per_1m".to_string(), json!(base_output_price));
|
||||
variables.insert(
|
||||
"cache_creation_price_per_1m".to_string(),
|
||||
json!(base_cache_creation_price),
|
||||
);
|
||||
variables.insert(
|
||||
"cache_read_price_per_1m".to_string(),
|
||||
json!(base_cache_read_price),
|
||||
);
|
||||
variables.insert("price_per_request".to_string(), json!(base_request_price));
|
||||
|
||||
let mut dimension_mappings = BTreeMap::new();
|
||||
for (name, key, default) in [
|
||||
("input_tokens", "input_tokens", json!(0)),
|
||||
("output_tokens", "output_tokens", json!(0)),
|
||||
("cache_creation_tokens", "cache_creation_tokens", json!(0)),
|
||||
("cache_read_tokens", "cache_read_tokens", json!(0)),
|
||||
("request_count", "request_count", json!(1)),
|
||||
] {
|
||||
dimension_mappings.insert(
|
||||
name.to_string(),
|
||||
json!({
|
||||
"source": "dimension",
|
||||
"key": key,
|
||||
"required": false,
|
||||
"allow_zero": true,
|
||||
"default": default,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
for (name, expression) in [
|
||||
("input_cost", "input_tokens * input_price_per_1m / 1000000"),
|
||||
(
|
||||
"output_cost",
|
||||
"output_tokens * output_price_per_1m / 1000000",
|
||||
),
|
||||
(
|
||||
"cache_creation_cost",
|
||||
"cache_creation_tokens * cache_creation_price_per_1m / 1000000",
|
||||
),
|
||||
(
|
||||
"cache_read_cost",
|
||||
"cache_read_tokens * cache_read_price_per_1m / 1000000",
|
||||
),
|
||||
("request_cost", "request_count * price_per_request"),
|
||||
] {
|
||||
dimension_mappings.insert(
|
||||
name.to_string(),
|
||||
json!({
|
||||
"source": "computed",
|
||||
"expression": expression,
|
||||
"required": false,
|
||||
"default": 0,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
if !tiers.is_empty() {
|
||||
dimension_mappings.insert(
|
||||
"input_price_per_1m".to_string(),
|
||||
json!({
|
||||
"source": "tiered",
|
||||
"tier_key": "total_input_context",
|
||||
"allow_zero": true,
|
||||
"tiers": build_tier_entries(&tiers, "input_price_per_1m", None, false),
|
||||
"default": base_input_price,
|
||||
}),
|
||||
);
|
||||
dimension_mappings.insert(
|
||||
"output_price_per_1m".to_string(),
|
||||
json!({
|
||||
"source": "tiered",
|
||||
"tier_key": "total_input_context",
|
||||
"allow_zero": true,
|
||||
"tiers": build_tier_entries(&tiers, "output_price_per_1m", None, false),
|
||||
"default": base_output_price,
|
||||
}),
|
||||
);
|
||||
dimension_mappings.insert(
|
||||
"cache_creation_price_per_1m".to_string(),
|
||||
json!({
|
||||
"source": "tiered",
|
||||
"tier_key": "total_input_context",
|
||||
"allow_zero": true,
|
||||
"ttl_key": "cache_ttl_minutes",
|
||||
"ttl_value_key": "cache_creation_price_per_1m",
|
||||
"tiers": build_tier_entries(&tiers, "cache_creation_price_per_1m", Some(1.25), true),
|
||||
"default": base_cache_creation_price,
|
||||
}),
|
||||
);
|
||||
dimension_mappings.insert(
|
||||
"cache_read_price_per_1m".to_string(),
|
||||
json!({
|
||||
"source": "tiered",
|
||||
"tier_key": "total_input_context",
|
||||
"allow_zero": true,
|
||||
"ttl_key": "cache_ttl_minutes",
|
||||
"ttl_value_key": "cache_read_price_per_1m",
|
||||
"tiers": build_tier_entries(&tiers, "cache_read_price_per_1m", Some(0.1), true),
|
||||
"default": base_cache_read_price,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
Some(VirtualBillingRule {
|
||||
id: "__default__".to_string(),
|
||||
name: format!("Default rule for {}", pricing.global_model_name),
|
||||
task_type: normalize_task_type(task_type).to_string(),
|
||||
expression:
|
||||
"input_cost + output_cost + cache_creation_cost + cache_read_cost + request_cost"
|
||||
.to_string(),
|
||||
variables,
|
||||
dimension_mappings,
|
||||
scope: "default".to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn normalize_task_type(task_type: &str) -> &str {
|
||||
if task_type.trim().eq_ignore_ascii_case("cli") {
|
||||
"chat"
|
||||
} else {
|
||||
task_type.trim()
|
||||
}
|
||||
}
|
||||
|
||||
fn tier_value(tier: &Value, key: &str, default: f64) -> f64 {
|
||||
tier.get(key).and_then(Value::as_f64).unwrap_or(default)
|
||||
}
|
||||
|
||||
fn tier_value_with_fallback(tier: &Value, key: &str, default_multiplier: f64) -> f64 {
|
||||
if let Some(value) = tier.get(key).and_then(Value::as_f64) {
|
||||
return value;
|
||||
}
|
||||
tier.get("input_price_per_1m")
|
||||
.and_then(Value::as_f64)
|
||||
.map(|value| value * default_multiplier)
|
||||
.unwrap_or(0.0)
|
||||
}
|
||||
|
||||
fn build_tier_entries(
|
||||
tiers: &[Value],
|
||||
key: &str,
|
||||
default_multiplier: Option<f64>,
|
||||
include_cache_ttl_pricing: bool,
|
||||
) -> Vec<Value> {
|
||||
tiers
|
||||
.iter()
|
||||
.map(|tier| {
|
||||
let mut value = serde_json::Map::new();
|
||||
value.insert(
|
||||
"up_to".to_string(),
|
||||
tier.get("up_to").cloned().unwrap_or(Value::Null),
|
||||
);
|
||||
let resolved = match default_multiplier {
|
||||
Some(multiplier) => Value::from(tier_value_with_fallback(tier, key, multiplier)),
|
||||
None => Value::from(tier_value(tier, key, 0.0)),
|
||||
};
|
||||
value.insert("value".to_string(), resolved);
|
||||
if include_cache_ttl_pricing {
|
||||
if let Some(ttl_pricing) = tier.get("cache_ttl_pricing").cloned() {
|
||||
value.insert("cache_ttl_pricing".to_string(), ttl_pricing);
|
||||
}
|
||||
}
|
||||
Value::Object(value)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
860
crates/aether-billing/src/formula_engine.rs
Normal file
860
crates/aether-billing/src/formula_engine.rs
Normal file
@@ -0,0 +1,860 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::precision::quantize_cost;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum FormulaEvaluationStatus {
|
||||
Complete,
|
||||
Incomplete,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct FormulaEvaluationResult {
|
||||
pub status: FormulaEvaluationStatus,
|
||||
pub cost: f64,
|
||||
pub resolved_dimensions: BTreeMap<String, serde_json::Value>,
|
||||
pub resolved_variables: BTreeMap<String, serde_json::Value>,
|
||||
pub cost_breakdown: BTreeMap<String, f64>,
|
||||
pub tier_index: Option<i64>,
|
||||
pub tier_info: Option<serde_json::Value>,
|
||||
pub missing_required: Vec<String>,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum UnsafeExpressionError {
|
||||
#[error("unsupported expression syntax: {0}")]
|
||||
Unsupported(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ExpressionEvaluationError {
|
||||
#[error("expression evaluation failed: {0}")]
|
||||
Failed(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
#[error("missing required dimensions: {missing_required:?}")]
|
||||
pub struct BillingIncompleteError {
|
||||
pub missing_required: Vec<String>,
|
||||
}
|
||||
|
||||
pub struct FormulaEngine;
|
||||
|
||||
type TierMetadata = Option<(i64, serde_json::Value)>;
|
||||
type MappingResolution = Result<(serde_json::Value, bool, TierMetadata), ExpressionEvaluationError>;
|
||||
|
||||
impl Default for FormulaEngine {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl FormulaEngine {
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
|
||||
pub fn evaluate(
|
||||
&self,
|
||||
expression: &str,
|
||||
variables: Option<&BTreeMap<String, serde_json::Value>>,
|
||||
dimensions: Option<&BTreeMap<String, serde_json::Value>>,
|
||||
dimension_mappings: Option<&BTreeMap<String, serde_json::Value>>,
|
||||
strict_mode: bool,
|
||||
) -> Result<FormulaEvaluationResult, ExpressionEvaluationError> {
|
||||
let dims = dimensions.cloned().unwrap_or_default();
|
||||
let mut resolved = variables.cloned().unwrap_or_default();
|
||||
let mut missing_required = Vec::new();
|
||||
let mut tier_index = None;
|
||||
let mut tier_info = None;
|
||||
let mut computed = BTreeMap::new();
|
||||
|
||||
if let Some(mappings) = dimension_mappings {
|
||||
for (var_name, mapping) in mappings {
|
||||
let source = mapping
|
||||
.get("source")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("constant")
|
||||
.to_ascii_lowercase();
|
||||
if source == "computed" {
|
||||
computed.insert(var_name.clone(), mapping.clone());
|
||||
continue;
|
||||
}
|
||||
if source == "constant" && resolved.contains_key(var_name) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let (value, is_missing, tier_meta) = resolve_mapping(var_name, mapping, &dims)?;
|
||||
if let Some((idx, info)) = tier_meta {
|
||||
if tier_index.is_none() {
|
||||
tier_index = Some(idx);
|
||||
tier_info = Some(info);
|
||||
}
|
||||
}
|
||||
if is_missing {
|
||||
missing_required.push(var_name.clone());
|
||||
continue;
|
||||
}
|
||||
resolved.insert(var_name.clone(), value);
|
||||
}
|
||||
}
|
||||
|
||||
if !computed.is_empty() {
|
||||
let mut unresolved = computed;
|
||||
for _ in 0..std::cmp::max(4, unresolved.len() + 1) {
|
||||
let mut progressed = false;
|
||||
let names: Vec<String> = unresolved.keys().cloned().collect();
|
||||
for var_name in names {
|
||||
let mapping = match unresolved.get(&var_name).cloned() {
|
||||
Some(value) => value,
|
||||
None => continue,
|
||||
};
|
||||
let (status, value) = try_resolve_computed(&mapping, &dims, &resolved)?;
|
||||
match status {
|
||||
ComputedStatus::Pending => {}
|
||||
ComputedStatus::MissingRequired => {
|
||||
missing_required.push(var_name.clone());
|
||||
unresolved.remove(&var_name);
|
||||
}
|
||||
ComputedStatus::Ok | ComputedStatus::Defaulted => {
|
||||
resolved.insert(var_name.clone(), value);
|
||||
unresolved.remove(&var_name);
|
||||
progressed = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
if !progressed {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
for (var_name, mapping) in unresolved {
|
||||
let required = mapping
|
||||
.get("required")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
if required {
|
||||
missing_required.push(var_name);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !missing_required.is_empty() {
|
||||
if strict_mode {
|
||||
return Err(ExpressionEvaluationError::Failed(
|
||||
BillingIncompleteError { missing_required }.to_string(),
|
||||
));
|
||||
}
|
||||
return Ok(FormulaEvaluationResult {
|
||||
status: FormulaEvaluationStatus::Incomplete,
|
||||
cost: 0.0,
|
||||
resolved_dimensions: dims,
|
||||
resolved_variables: resolved,
|
||||
cost_breakdown: BTreeMap::new(),
|
||||
tier_index,
|
||||
tier_info,
|
||||
missing_required,
|
||||
error: None,
|
||||
});
|
||||
}
|
||||
|
||||
let cost = evaluate_expression(expression, &resolved)
|
||||
.map_err(|err| ExpressionEvaluationError::Failed(err.to_string()))?;
|
||||
if cost < 0.0 {
|
||||
return Ok(FormulaEvaluationResult {
|
||||
status: FormulaEvaluationStatus::Incomplete,
|
||||
cost: 0.0,
|
||||
resolved_dimensions: dims,
|
||||
resolved_variables: resolved,
|
||||
cost_breakdown: BTreeMap::new(),
|
||||
tier_index,
|
||||
tier_info,
|
||||
missing_required: Vec::new(),
|
||||
error: Some("negative_cost".to_string()),
|
||||
});
|
||||
}
|
||||
|
||||
let mut breakdown = BTreeMap::new();
|
||||
for (key, value) in &resolved {
|
||||
if key.ends_with("_cost") {
|
||||
if let Some(number) = as_f64(value) {
|
||||
breakdown.insert(key.clone(), quantize_cost(number));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(FormulaEvaluationResult {
|
||||
status: FormulaEvaluationStatus::Complete,
|
||||
cost: quantize_cost(cost),
|
||||
resolved_dimensions: dims,
|
||||
resolved_variables: resolved,
|
||||
cost_breakdown: breakdown,
|
||||
tier_index,
|
||||
tier_info,
|
||||
missing_required: Vec::new(),
|
||||
error: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn extract_variable_names(expression: &str) -> Result<BTreeSet<String>, UnsafeExpressionError> {
|
||||
let tokens = tokenize(expression)?;
|
||||
let mut names = BTreeSet::new();
|
||||
for index in 0..tokens.len() {
|
||||
if let Token::Identifier(name) = &tokens[index] {
|
||||
if matches!(tokens.get(index + 1), Some(Token::LeftParen)) {
|
||||
continue;
|
||||
}
|
||||
names.insert(name.clone());
|
||||
}
|
||||
}
|
||||
Ok(names)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum ComputedStatus {
|
||||
Ok,
|
||||
Pending,
|
||||
MissingRequired,
|
||||
Defaulted,
|
||||
}
|
||||
|
||||
fn try_resolve_computed(
|
||||
mapping: &serde_json::Value,
|
||||
dims: &BTreeMap<String, serde_json::Value>,
|
||||
resolved: &BTreeMap<String, serde_json::Value>,
|
||||
) -> Result<(ComputedStatus, serde_json::Value), ExpressionEvaluationError> {
|
||||
let required = mapping
|
||||
.get("required")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
let default = mapping
|
||||
.get("default")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| serde_json::json!(0));
|
||||
let expr = mapping
|
||||
.get("expression")
|
||||
.or_else(|| mapping.get("transform_expression"))
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or_default();
|
||||
if expr.is_empty() {
|
||||
return Ok((
|
||||
if required {
|
||||
ComputedStatus::MissingRequired
|
||||
} else {
|
||||
ComputedStatus::Defaulted
|
||||
},
|
||||
default,
|
||||
));
|
||||
}
|
||||
|
||||
let needed = extract_variable_names(expr)
|
||||
.map_err(|err| ExpressionEvaluationError::Failed(err.to_string()))?;
|
||||
let env: BTreeMap<String, serde_json::Value> = dims
|
||||
.iter()
|
||||
.map(|(k, v)| (k.clone(), v.clone()))
|
||||
.chain(resolved.iter().map(|(k, v)| (k.clone(), v.clone())))
|
||||
.collect();
|
||||
if needed.iter().any(|name| !env.contains_key(name)) {
|
||||
return Ok((
|
||||
if required {
|
||||
ComputedStatus::Pending
|
||||
} else {
|
||||
ComputedStatus::Defaulted
|
||||
},
|
||||
default,
|
||||
));
|
||||
}
|
||||
|
||||
match evaluate_expression(expr, &env) {
|
||||
Ok(value) => Ok((ComputedStatus::Ok, serde_json::json!(value))),
|
||||
Err(err) if required => Err(ExpressionEvaluationError::Failed(err.to_string())),
|
||||
Err(_) => Ok((ComputedStatus::Defaulted, default)),
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_mapping(
|
||||
var_name: &str,
|
||||
mapping: &serde_json::Value,
|
||||
dims: &BTreeMap<String, serde_json::Value>,
|
||||
) -> MappingResolution {
|
||||
let source = mapping
|
||||
.get("source")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("constant")
|
||||
.to_ascii_lowercase();
|
||||
|
||||
match source.as_str() {
|
||||
"constant" => Ok((
|
||||
mapping
|
||||
.get("default")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| serde_json::json!(0)),
|
||||
false,
|
||||
None,
|
||||
)),
|
||||
"dimension" => resolve_dimension(var_name, mapping, dims),
|
||||
"tiered" => resolve_tiered(var_name, mapping, dims),
|
||||
"matrix" => resolve_matrix(var_name, mapping, dims),
|
||||
_ => Ok((
|
||||
mapping
|
||||
.get("default")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| serde_json::json!(0)),
|
||||
false,
|
||||
None,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_dimension(
|
||||
var_name: &str,
|
||||
mapping: &serde_json::Value,
|
||||
dims: &BTreeMap<String, serde_json::Value>,
|
||||
) -> MappingResolution {
|
||||
let key = mapping
|
||||
.get("key")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or(var_name);
|
||||
let required = mapping
|
||||
.get("required")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
let allow_zero = mapping
|
||||
.get("allow_zero")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(true);
|
||||
let default = mapping
|
||||
.get("default")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| serde_json::json!(0));
|
||||
|
||||
let Some(value) = dims.get(key).cloned() else {
|
||||
return Ok((default, required, None));
|
||||
};
|
||||
if !allow_zero && is_zero_like(&value) {
|
||||
return Ok((default, required, None));
|
||||
}
|
||||
Ok((value, false, None))
|
||||
}
|
||||
|
||||
fn resolve_tiered(
|
||||
_var_name: &str,
|
||||
mapping: &serde_json::Value,
|
||||
dims: &BTreeMap<String, serde_json::Value>,
|
||||
) -> MappingResolution {
|
||||
let tier_key = mapping
|
||||
.get("tier_key")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("total_input_context");
|
||||
let default = mapping
|
||||
.get("default")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| serde_json::json!(0));
|
||||
let Some(tiers) = mapping.get("tiers").and_then(|v| v.as_array()) else {
|
||||
return Ok((default, false, None));
|
||||
};
|
||||
|
||||
let tier_value = dims.get(tier_key).and_then(as_f64).unwrap_or(0.0);
|
||||
let ttl_key = mapping.get("ttl_key").and_then(|v| v.as_str());
|
||||
let ttl_value_key = mapping.get("ttl_value_key").and_then(|v| v.as_str());
|
||||
let cache_ttl_minutes = ttl_key.and_then(|key| dims.get(key)).and_then(as_f64);
|
||||
|
||||
let mut matched_index = None;
|
||||
let mut matched_tier = None;
|
||||
for (index, tier) in tiers.iter().enumerate() {
|
||||
let up_to = tier.get("up_to").and_then(as_f64);
|
||||
if up_to.is_none() || tier_value <= up_to.unwrap_or_default() {
|
||||
matched_index = Some(index as i64);
|
||||
matched_tier = Some(tier.clone());
|
||||
break;
|
||||
}
|
||||
}
|
||||
let Some(tier) = matched_tier.or_else(|| tiers.last().cloned()) else {
|
||||
return Ok((default, false, None));
|
||||
};
|
||||
|
||||
if let (Some(cache_ttl_minutes), Some(ttl_value_key)) = (cache_ttl_minutes, ttl_value_key) {
|
||||
if let Some(ttl_pricing) = tier.get("cache_ttl_pricing").and_then(|v| v.as_array()) {
|
||||
for ttl_entry in ttl_pricing {
|
||||
let ttl_limit = ttl_entry
|
||||
.get("ttl_minutes")
|
||||
.and_then(as_f64)
|
||||
.unwrap_or_default();
|
||||
if cache_ttl_minutes <= ttl_limit {
|
||||
if let Some(value) = ttl_entry.get(ttl_value_key) {
|
||||
return Ok((
|
||||
value.clone(),
|
||||
false,
|
||||
Some((matched_index.unwrap_or(0), tier)),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(value) = ttl_pricing.last().and_then(|v| v.get(ttl_value_key)) {
|
||||
return Ok((
|
||||
value.clone(),
|
||||
false,
|
||||
Some((matched_index.unwrap_or(0), tier)),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok((
|
||||
tier.get("value").cloned().unwrap_or(default),
|
||||
false,
|
||||
Some((matched_index.unwrap_or(0), tier)),
|
||||
))
|
||||
}
|
||||
|
||||
fn resolve_matrix(
|
||||
var_name: &str,
|
||||
mapping: &serde_json::Value,
|
||||
dims: &BTreeMap<String, serde_json::Value>,
|
||||
) -> MappingResolution {
|
||||
let key = mapping
|
||||
.get("key")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or(var_name);
|
||||
let default = mapping
|
||||
.get("default")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| serde_json::json!(0));
|
||||
let Some(value) = dims.get(key) else {
|
||||
let required = mapping
|
||||
.get("required")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
return Ok((default, required, None));
|
||||
};
|
||||
let Some(entries) = mapping.get("entries").and_then(|v| v.as_object()) else {
|
||||
return Ok((default, false, None));
|
||||
};
|
||||
let lookup_key = match value {
|
||||
serde_json::Value::String(text) => text.clone(),
|
||||
_ => value.to_string(),
|
||||
};
|
||||
Ok((
|
||||
entries.get(&lookup_key).cloned().unwrap_or(default),
|
||||
false,
|
||||
None,
|
||||
))
|
||||
}
|
||||
|
||||
fn is_zero_like(value: &serde_json::Value) -> bool {
|
||||
as_f64(value).map(|number| number == 0.0).unwrap_or(false)
|
||||
}
|
||||
|
||||
fn as_f64(value: &serde_json::Value) -> Option<f64> {
|
||||
value.as_f64().or_else(|| value.as_i64().map(|v| v as f64))
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
enum Token {
|
||||
Number(f64),
|
||||
Identifier(String),
|
||||
Plus,
|
||||
Minus,
|
||||
Star,
|
||||
DoubleStar,
|
||||
Slash,
|
||||
DoubleSlash,
|
||||
Percent,
|
||||
LeftParen,
|
||||
RightParen,
|
||||
Comma,
|
||||
}
|
||||
|
||||
fn tokenize(input: &str) -> Result<Vec<Token>, UnsafeExpressionError> {
|
||||
let chars: Vec<char> = input.chars().collect();
|
||||
let mut tokens = Vec::new();
|
||||
let mut index = 0;
|
||||
while index < chars.len() {
|
||||
match chars[index] {
|
||||
ch if ch.is_ascii_whitespace() => index += 1,
|
||||
ch if ch.is_ascii_digit() || ch == '.' => {
|
||||
let start = index;
|
||||
index += 1;
|
||||
while index < chars.len() && (chars[index].is_ascii_digit() || chars[index] == '.')
|
||||
{
|
||||
index += 1;
|
||||
}
|
||||
let text: String = chars[start..index].iter().collect();
|
||||
let number = text.parse::<f64>().map_err(|_| {
|
||||
UnsafeExpressionError::Unsupported(format!("invalid numeric literal: {text}"))
|
||||
})?;
|
||||
tokens.push(Token::Number(number));
|
||||
}
|
||||
ch if ch.is_ascii_alphabetic() || ch == '_' => {
|
||||
let start = index;
|
||||
index += 1;
|
||||
while index < chars.len()
|
||||
&& (chars[index].is_ascii_alphanumeric() || chars[index] == '_')
|
||||
{
|
||||
index += 1;
|
||||
}
|
||||
tokens.push(Token::Identifier(chars[start..index].iter().collect()));
|
||||
}
|
||||
'+' => {
|
||||
tokens.push(Token::Plus);
|
||||
index += 1;
|
||||
}
|
||||
'-' => {
|
||||
tokens.push(Token::Minus);
|
||||
index += 1;
|
||||
}
|
||||
'*' => {
|
||||
if matches!(chars.get(index + 1), Some('*')) {
|
||||
tokens.push(Token::DoubleStar);
|
||||
index += 2;
|
||||
} else {
|
||||
tokens.push(Token::Star);
|
||||
index += 1;
|
||||
}
|
||||
}
|
||||
'/' => {
|
||||
if matches!(chars.get(index + 1), Some('/')) {
|
||||
tokens.push(Token::DoubleSlash);
|
||||
index += 2;
|
||||
} else {
|
||||
tokens.push(Token::Slash);
|
||||
index += 1;
|
||||
}
|
||||
}
|
||||
'%' => {
|
||||
tokens.push(Token::Percent);
|
||||
index += 1;
|
||||
}
|
||||
'(' => {
|
||||
tokens.push(Token::LeftParen);
|
||||
index += 1;
|
||||
}
|
||||
')' => {
|
||||
tokens.push(Token::RightParen);
|
||||
index += 1;
|
||||
}
|
||||
',' => {
|
||||
tokens.push(Token::Comma);
|
||||
index += 1;
|
||||
}
|
||||
other => {
|
||||
return Err(UnsafeExpressionError::Unsupported(format!(
|
||||
"unsupported token: {other}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(tokens)
|
||||
}
|
||||
|
||||
fn evaluate_expression(
|
||||
expression: &str,
|
||||
variables: &BTreeMap<String, serde_json::Value>,
|
||||
) -> Result<f64, UnsafeExpressionError> {
|
||||
let tokens = tokenize(expression)?;
|
||||
let mut parser = Parser {
|
||||
tokens: &tokens,
|
||||
index: 0,
|
||||
variables,
|
||||
};
|
||||
let value = parser.parse_expression()?;
|
||||
if parser.index != tokens.len() {
|
||||
return Err(UnsafeExpressionError::Unsupported(
|
||||
"unexpected trailing tokens".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
struct Parser<'a> {
|
||||
tokens: &'a [Token],
|
||||
index: usize,
|
||||
variables: &'a BTreeMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
impl<'a> Parser<'a> {
|
||||
fn parse_expression(&mut self) -> Result<f64, UnsafeExpressionError> {
|
||||
let mut left = self.parse_term()?;
|
||||
loop {
|
||||
match self.peek() {
|
||||
Some(Token::Plus) => {
|
||||
self.index += 1;
|
||||
left += self.parse_term()?;
|
||||
}
|
||||
Some(Token::Minus) => {
|
||||
self.index += 1;
|
||||
left -= self.parse_term()?;
|
||||
}
|
||||
_ => break,
|
||||
}
|
||||
}
|
||||
Ok(left)
|
||||
}
|
||||
|
||||
fn parse_term(&mut self) -> Result<f64, UnsafeExpressionError> {
|
||||
let mut left = self.parse_power()?;
|
||||
loop {
|
||||
match self.peek() {
|
||||
Some(Token::Star) => {
|
||||
self.index += 1;
|
||||
left *= self.parse_power()?;
|
||||
}
|
||||
Some(Token::Slash) => {
|
||||
self.index += 1;
|
||||
left /= self.parse_power()?;
|
||||
}
|
||||
Some(Token::DoubleSlash) => {
|
||||
self.index += 1;
|
||||
left = (left / self.parse_power()?).floor();
|
||||
}
|
||||
Some(Token::Percent) => {
|
||||
self.index += 1;
|
||||
left %= self.parse_power()?;
|
||||
}
|
||||
_ => break,
|
||||
}
|
||||
}
|
||||
Ok(left)
|
||||
}
|
||||
|
||||
fn parse_power(&mut self) -> Result<f64, UnsafeExpressionError> {
|
||||
let left = self.parse_unary()?;
|
||||
if matches!(self.peek(), Some(Token::DoubleStar)) {
|
||||
self.index += 1;
|
||||
let right = self.parse_power()?;
|
||||
return Ok(left.powf(right));
|
||||
}
|
||||
Ok(left)
|
||||
}
|
||||
|
||||
fn parse_unary(&mut self) -> Result<f64, UnsafeExpressionError> {
|
||||
match self.peek() {
|
||||
Some(Token::Plus) => {
|
||||
self.index += 1;
|
||||
self.parse_unary()
|
||||
}
|
||||
Some(Token::Minus) => {
|
||||
self.index += 1;
|
||||
Ok(-self.parse_unary()?)
|
||||
}
|
||||
_ => self.parse_primary(),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_primary(&mut self) -> Result<f64, UnsafeExpressionError> {
|
||||
match self.next() {
|
||||
Some(Token::Number(value)) => Ok(*value),
|
||||
Some(Token::Identifier(name)) => {
|
||||
if matches!(self.peek(), Some(Token::LeftParen)) {
|
||||
self.index += 1;
|
||||
let args = self.parse_arguments()?;
|
||||
evaluate_function(name, &args)
|
||||
} else {
|
||||
let value = self.variables.get(name).and_then(as_f64).ok_or_else(|| {
|
||||
UnsafeExpressionError::Unsupported(format!("unknown variable: {name}"))
|
||||
})?;
|
||||
Ok(value)
|
||||
}
|
||||
}
|
||||
Some(Token::LeftParen) => {
|
||||
let value = self.parse_expression()?;
|
||||
match self.next() {
|
||||
Some(Token::RightParen) => Ok(value),
|
||||
_ => Err(UnsafeExpressionError::Unsupported(
|
||||
"missing closing ')'".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
other => Err(UnsafeExpressionError::Unsupported(format!(
|
||||
"unexpected token: {other:?}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_arguments(&mut self) -> Result<Vec<f64>, UnsafeExpressionError> {
|
||||
let mut args = Vec::new();
|
||||
if matches!(self.peek(), Some(Token::RightParen)) {
|
||||
self.index += 1;
|
||||
return Ok(args);
|
||||
}
|
||||
loop {
|
||||
args.push(self.parse_expression()?);
|
||||
match self.next() {
|
||||
Some(Token::Comma) => {}
|
||||
Some(Token::RightParen) => break,
|
||||
other => {
|
||||
return Err(UnsafeExpressionError::Unsupported(format!(
|
||||
"unexpected token in function call: {other:?}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(args)
|
||||
}
|
||||
|
||||
fn peek(&self) -> Option<&'a Token> {
|
||||
self.tokens.get(self.index)
|
||||
}
|
||||
|
||||
fn next(&mut self) -> Option<&'a Token> {
|
||||
let token = self.tokens.get(self.index);
|
||||
if token.is_some() {
|
||||
self.index += 1;
|
||||
}
|
||||
token
|
||||
}
|
||||
}
|
||||
|
||||
fn evaluate_function(name: &str, args: &[f64]) -> Result<f64, UnsafeExpressionError> {
|
||||
match name {
|
||||
"min" => args.iter().copied().reduce(f64::min).ok_or_else(|| {
|
||||
UnsafeExpressionError::Unsupported("min() requires arguments".to_string())
|
||||
}),
|
||||
"max" => args.iter().copied().reduce(f64::max).ok_or_else(|| {
|
||||
UnsafeExpressionError::Unsupported("max() requires arguments".to_string())
|
||||
}),
|
||||
"abs" => args.first().copied().map(f64::abs).ok_or_else(|| {
|
||||
UnsafeExpressionError::Unsupported("abs() requires one argument".to_string())
|
||||
}),
|
||||
"round" => {
|
||||
let Some(value) = args.first().copied() else {
|
||||
return Err(UnsafeExpressionError::Unsupported(
|
||||
"round() requires at least one argument".to_string(),
|
||||
));
|
||||
};
|
||||
let digits = args.get(1).copied().unwrap_or(0.0) as i32;
|
||||
let factor = 10_f64.powi(digits);
|
||||
Ok((value * factor).round() / factor)
|
||||
}
|
||||
"int" => args
|
||||
.first()
|
||||
.copied()
|
||||
.map(|value| value.trunc())
|
||||
.ok_or_else(|| {
|
||||
UnsafeExpressionError::Unsupported("int() requires one argument".to_string())
|
||||
}),
|
||||
"float" => args.first().copied().ok_or_else(|| {
|
||||
UnsafeExpressionError::Unsupported("float() requires one argument".to_string())
|
||||
}),
|
||||
_ => Err(UnsafeExpressionError::Unsupported(format!(
|
||||
"function not allowed: {name}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{extract_variable_names, FormulaEngine, FormulaEvaluationStatus};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
#[test]
|
||||
fn extracts_variable_names_without_functions() {
|
||||
let names = extract_variable_names("min(input_cost, output_cost) + tax")
|
||||
.expect("variables should parse");
|
||||
assert!(names.contains("input_cost"));
|
||||
assert!(names.contains("output_cost"));
|
||||
assert!(names.contains("tax"));
|
||||
assert!(!names.contains("min"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evaluates_formula_with_computed_variables() {
|
||||
let engine = FormulaEngine::new();
|
||||
let dimensions = BTreeMap::from([
|
||||
("input_tokens".to_string(), serde_json::json!(1000)),
|
||||
("output_tokens".to_string(), serde_json::json!(500)),
|
||||
]);
|
||||
let variables = BTreeMap::from([
|
||||
("input_price_per_1m".to_string(), serde_json::json!(3.0)),
|
||||
("output_price_per_1m".to_string(), serde_json::json!(15.0)),
|
||||
]);
|
||||
let mappings = BTreeMap::from([
|
||||
(
|
||||
"input_cost".to_string(),
|
||||
serde_json::json!({
|
||||
"source": "computed",
|
||||
"expression": "input_tokens * input_price_per_1m / 1000000"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"output_cost".to_string(),
|
||||
serde_json::json!({
|
||||
"source": "computed",
|
||||
"expression": "output_tokens * output_price_per_1m / 1000000"
|
||||
}),
|
||||
),
|
||||
]);
|
||||
|
||||
let result = engine
|
||||
.evaluate(
|
||||
"input_cost + output_cost",
|
||||
Some(&variables),
|
||||
Some(&dimensions),
|
||||
Some(&mappings),
|
||||
false,
|
||||
)
|
||||
.expect("formula should evaluate");
|
||||
|
||||
assert_eq!(result.status, FormulaEvaluationStatus::Complete);
|
||||
assert_eq!(result.cost, 0.0105);
|
||||
assert_eq!(result.cost_breakdown.get("input_cost"), Some(&0.003_f64));
|
||||
assert_eq!(result.cost_breakdown.get("output_cost"), Some(&0.0075_f64));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_tiered_mapping() {
|
||||
let engine = FormulaEngine::new();
|
||||
let dimensions =
|
||||
BTreeMap::from([("total_input_context".to_string(), serde_json::json!(2000))]);
|
||||
let mappings = BTreeMap::from([(
|
||||
"input_price_per_1m".to_string(),
|
||||
serde_json::json!({
|
||||
"source": "tiered",
|
||||
"tier_key": "total_input_context",
|
||||
"tiers": [
|
||||
{ "up_to": 1000, "value": 1.0 },
|
||||
{ "up_to": 4000, "value": 2.0 }
|
||||
],
|
||||
"default": 0.0
|
||||
}),
|
||||
)]);
|
||||
|
||||
let result = engine
|
||||
.evaluate(
|
||||
"input_price_per_1m",
|
||||
None,
|
||||
Some(&dimensions),
|
||||
Some(&mappings),
|
||||
false,
|
||||
)
|
||||
.expect("tiered mapping should evaluate");
|
||||
|
||||
assert_eq!(result.cost, 2.0);
|
||||
assert_eq!(result.tier_index, Some(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incomplete_when_required_dimension_missing() {
|
||||
let engine = FormulaEngine::new();
|
||||
let mappings = BTreeMap::from([(
|
||||
"input_tokens".to_string(),
|
||||
serde_json::json!({
|
||||
"source": "dimension",
|
||||
"key": "input_tokens",
|
||||
"required": true
|
||||
}),
|
||||
)]);
|
||||
|
||||
let result = engine
|
||||
.evaluate("input_tokens", None, None, Some(&mappings), false)
|
||||
.expect("incomplete result should return ok");
|
||||
|
||||
assert_eq!(result.status, FormulaEvaluationStatus::Incomplete);
|
||||
assert_eq!(result.missing_required, vec!["input_tokens".to_string()]);
|
||||
}
|
||||
}
|
||||
27
crates/aether-billing/src/lib.rs
Normal file
27
crates/aether-billing/src/lib.rs
Normal file
@@ -0,0 +1,27 @@
|
||||
mod default_rule;
|
||||
mod formula_engine;
|
||||
mod models;
|
||||
mod precision;
|
||||
mod pricing;
|
||||
mod schema;
|
||||
mod service;
|
||||
mod token_normalization;
|
||||
mod usage_mapper;
|
||||
|
||||
pub use default_rule::{normalize_task_type, DefaultBillingRuleGenerator, VirtualBillingRule};
|
||||
pub use formula_engine::{
|
||||
extract_variable_names, BillingIncompleteError, ExpressionEvaluationError, FormulaEngine,
|
||||
FormulaEvaluationResult, FormulaEvaluationStatus, UnsafeExpressionError,
|
||||
};
|
||||
pub use models::{BillingDimension, BillingUnit, CostBreakdown, StandardizedUsage};
|
||||
pub use precision::{
|
||||
quantize_cost, quantize_display, quantize_value, BILLING_DISPLAY_PRECISION,
|
||||
BILLING_STORAGE_PRECISION,
|
||||
};
|
||||
pub use pricing::{BillingComputation, BillingModelPricingSnapshot, BillingUsageInput};
|
||||
pub use schema::{
|
||||
BillingSnapshot, BillingSnapshotStatus, CostResult, BILLING_SNAPSHOT_SCHEMA_VERSION,
|
||||
};
|
||||
pub use service::BillingService;
|
||||
pub use token_normalization::normalize_input_tokens_for_billing;
|
||||
pub use usage_mapper::{map_usage, map_usage_from_response, UsageMapper};
|
||||
140
crates/aether-billing/src/models.rs
Normal file
140
crates/aether-billing/src/models.rs
Normal file
@@ -0,0 +1,140 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum BillingUnit {
|
||||
Per1MTokens,
|
||||
Per1MTokensHour,
|
||||
PerRequest,
|
||||
Fixed,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct BillingDimension {
|
||||
pub name: String,
|
||||
pub usage_field: String,
|
||||
pub price_field: String,
|
||||
pub unit: BillingUnit,
|
||||
pub default_price: f64,
|
||||
}
|
||||
|
||||
impl BillingDimension {
|
||||
pub fn calculate(&self, usage_value: f64, price: f64) -> f64 {
|
||||
if usage_value <= 0.0 || price <= 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
match self.unit {
|
||||
BillingUnit::Per1MTokens | BillingUnit::Per1MTokensHour => {
|
||||
(usage_value / 1_000_000.0) * price
|
||||
}
|
||||
BillingUnit::PerRequest => usage_value * price,
|
||||
BillingUnit::Fixed => price,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, Default)]
|
||||
pub struct StandardizedUsage {
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cache_creation_tokens: i64,
|
||||
pub cache_read_tokens: i64,
|
||||
pub reasoning_tokens: i64,
|
||||
pub cache_storage_token_hours: f64,
|
||||
pub request_count: i64,
|
||||
pub dimensions: BTreeMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
impl StandardizedUsage {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
request_count: 1,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get(&self, field_name: &str) -> Option<serde_json::Value> {
|
||||
match field_name {
|
||||
"input_tokens" => Some(serde_json::json!(self.input_tokens)),
|
||||
"output_tokens" => Some(serde_json::json!(self.output_tokens)),
|
||||
"cache_creation_tokens" => Some(serde_json::json!(self.cache_creation_tokens)),
|
||||
"cache_read_tokens" => Some(serde_json::json!(self.cache_read_tokens)),
|
||||
"reasoning_tokens" => Some(serde_json::json!(self.reasoning_tokens)),
|
||||
"cache_storage_token_hours" => Some(serde_json::json!(self.cache_storage_token_hours)),
|
||||
"request_count" => Some(serde_json::json!(self.request_count)),
|
||||
"extra" | "dimensions" => Some(serde_json::json!(self.dimensions)),
|
||||
_ => self.dimensions.get(field_name).cloned(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set(&mut self, field_name: &str, value: impl Into<serde_json::Value>) {
|
||||
let value = value.into();
|
||||
match field_name {
|
||||
"input_tokens" => self.input_tokens = as_i64(&value, 0),
|
||||
"output_tokens" => self.output_tokens = as_i64(&value, 0),
|
||||
"cache_creation_tokens" => self.cache_creation_tokens = as_i64(&value, 0),
|
||||
"cache_read_tokens" => self.cache_read_tokens = as_i64(&value, 0),
|
||||
"reasoning_tokens" => self.reasoning_tokens = as_i64(&value, 0),
|
||||
"cache_storage_token_hours" => self.cache_storage_token_hours = as_f64(&value, 0.0),
|
||||
"request_count" => self.request_count = as_i64(&value, 0),
|
||||
"extra" | "dimensions" => {
|
||||
self.dimensions = match value {
|
||||
serde_json::Value::Object(map) => map.into_iter().collect(),
|
||||
_ => BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
self.dimensions.insert(field_name.to_string(), value);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, Default)]
|
||||
pub struct CostBreakdown {
|
||||
pub costs: BTreeMap<String, f64>,
|
||||
pub total_cost: f64,
|
||||
pub tier_index: Option<i64>,
|
||||
pub effective_prices: BTreeMap<String, f64>,
|
||||
}
|
||||
|
||||
fn as_i64(value: &serde_json::Value, default: i64) -> i64 {
|
||||
value
|
||||
.as_i64()
|
||||
.or_else(|| value.as_u64().and_then(|v| i64::try_from(v).ok()))
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
fn as_f64(value: &serde_json::Value, default: f64) -> f64 {
|
||||
value.as_f64().unwrap_or(default)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{BillingDimension, BillingUnit, StandardizedUsage};
|
||||
|
||||
#[test]
|
||||
fn dimension_calculates_per_million_tokens() {
|
||||
let dimension = BillingDimension {
|
||||
name: "input".to_string(),
|
||||
usage_field: "input_tokens".to_string(),
|
||||
price_field: "input_price_per_1m".to_string(),
|
||||
unit: BillingUnit::Per1MTokens,
|
||||
default_price: 0.0,
|
||||
};
|
||||
assert_eq!(dimension.calculate(500_000.0, 2.0), 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn standardized_usage_reads_and_writes_known_and_extra_fields() {
|
||||
let mut usage = StandardizedUsage::new();
|
||||
usage.set("input_tokens", 10);
|
||||
usage.set("custom_dimension", "value");
|
||||
|
||||
assert_eq!(usage.get("input_tokens"), Some(serde_json::json!(10)));
|
||||
assert_eq!(
|
||||
usage.get("custom_dimension"),
|
||||
Some(serde_json::json!("value"))
|
||||
);
|
||||
}
|
||||
}
|
||||
33
crates/aether-billing/src/precision.rs
Normal file
33
crates/aether-billing/src/precision.rs
Normal file
@@ -0,0 +1,33 @@
|
||||
pub const BILLING_STORAGE_PRECISION: u32 = 8;
|
||||
pub const BILLING_DISPLAY_PRECISION: u32 = 6;
|
||||
|
||||
pub fn quantize_value(value: f64, precision: u32) -> f64 {
|
||||
if !value.is_finite() {
|
||||
return value;
|
||||
}
|
||||
let factor = 10_f64.powi(precision as i32);
|
||||
(value * factor).round() / factor
|
||||
}
|
||||
|
||||
pub fn quantize_cost(value: f64) -> f64 {
|
||||
quantize_value(value, BILLING_STORAGE_PRECISION)
|
||||
}
|
||||
|
||||
pub fn quantize_display(value: f64) -> f64 {
|
||||
quantize_value(value, BILLING_DISPLAY_PRECISION)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{quantize_cost, quantize_display};
|
||||
|
||||
#[test]
|
||||
fn quantizes_cost_to_storage_precision() {
|
||||
assert_eq!(quantize_cost(1.234567891), 1.23456789);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quantizes_display_to_display_precision() {
|
||||
assert_eq!(quantize_display(1.23456789), 1.234568);
|
||||
}
|
||||
}
|
||||
94
crates/aether-billing/src/pricing.rs
Normal file
94
crates/aether-billing/src/pricing.rs
Normal file
@@ -0,0 +1,94 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct BillingModelPricingSnapshot {
|
||||
pub provider_id: String,
|
||||
pub provider_billing_type: Option<String>,
|
||||
pub provider_api_key_id: Option<String>,
|
||||
pub provider_api_key_rate_multipliers: Option<Value>,
|
||||
pub provider_api_key_cache_ttl_minutes: Option<i64>,
|
||||
pub global_model_id: String,
|
||||
pub global_model_name: String,
|
||||
pub global_model_config: Option<Value>,
|
||||
pub default_price_per_request: Option<f64>,
|
||||
pub default_tiered_pricing: Option<Value>,
|
||||
pub model_id: Option<String>,
|
||||
pub model_provider_model_name: Option<String>,
|
||||
pub model_config: Option<Value>,
|
||||
pub model_price_per_request: Option<f64>,
|
||||
pub model_tiered_pricing: Option<Value>,
|
||||
}
|
||||
|
||||
impl BillingModelPricingSnapshot {
|
||||
pub fn effective_tiered_pricing(&self) -> Option<&Value> {
|
||||
self.model_tiered_pricing
|
||||
.as_ref()
|
||||
.or(self.default_tiered_pricing.as_ref())
|
||||
}
|
||||
|
||||
pub fn effective_price_per_request(&self) -> Option<f64> {
|
||||
self.model_price_per_request
|
||||
.or(self.default_price_per_request)
|
||||
}
|
||||
|
||||
pub fn is_free_tier(&self) -> bool {
|
||||
self.provider_billing_type
|
||||
.as_deref()
|
||||
.map(|value| value.eq_ignore_ascii_case("free_tier"))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
pub fn rate_multiplier_for_api_format(&self, api_format: Option<&str>) -> f64 {
|
||||
let Some(api_format) = api_format.map(str::trim).filter(|value| !value.is_empty()) else {
|
||||
return 1.0;
|
||||
};
|
||||
let normalized = api_format.to_ascii_lowercase();
|
||||
let Some(mapping) = self
|
||||
.provider_api_key_rate_multipliers
|
||||
.as_ref()
|
||||
.and_then(Value::as_object)
|
||||
else {
|
||||
return 1.0;
|
||||
};
|
||||
mapping
|
||||
.get(&normalized)
|
||||
.and_then(|value| value.as_f64())
|
||||
.unwrap_or(1.0)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct BillingUsageInput {
|
||||
pub task_type: String,
|
||||
pub api_format: Option<String>,
|
||||
pub request_count: i64,
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cache_creation_tokens: i64,
|
||||
pub cache_read_tokens: i64,
|
||||
pub cache_ttl_minutes: Option<i64>,
|
||||
}
|
||||
|
||||
impl BillingUsageInput {
|
||||
pub fn new(task_type: impl Into<String>) -> Self {
|
||||
Self {
|
||||
task_type: task_type.into(),
|
||||
api_format: None,
|
||||
request_count: 1,
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
cache_creation_tokens: 0,
|
||||
cache_read_tokens: 0,
|
||||
cache_ttl_minutes: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct BillingComputation {
|
||||
pub cost_result: crate::CostResult,
|
||||
pub actual_total_cost: f64,
|
||||
pub rate_multiplier: f64,
|
||||
pub is_free_tier: bool,
|
||||
}
|
||||
110
crates/aether-billing/src/schema.rs
Normal file
110
crates/aether-billing/src/schema.rs
Normal file
@@ -0,0 +1,110 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
pub const BILLING_SNAPSHOT_SCHEMA_VERSION: &str = "2.0";
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum BillingSnapshotStatus {
|
||||
Complete,
|
||||
Incomplete,
|
||||
NoRule,
|
||||
Legacy,
|
||||
}
|
||||
|
||||
impl BillingSnapshotStatus {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Complete => "complete",
|
||||
Self::Incomplete => "incomplete",
|
||||
Self::NoRule => "no_rule",
|
||||
Self::Legacy => "legacy",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for BillingSnapshotStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct BillingSnapshot {
|
||||
pub schema_version: String,
|
||||
pub rule_id: Option<String>,
|
||||
pub rule_name: Option<String>,
|
||||
pub scope: Option<String>,
|
||||
pub expression: Option<String>,
|
||||
pub resolved_dimensions: BTreeMap<String, serde_json::Value>,
|
||||
pub resolved_variables: BTreeMap<String, serde_json::Value>,
|
||||
pub cost_breakdown: BTreeMap<String, f64>,
|
||||
pub total_cost: f64,
|
||||
pub tier_index: Option<i64>,
|
||||
pub tier_info: Option<serde_json::Value>,
|
||||
pub missing_required: Vec<String>,
|
||||
pub status: BillingSnapshotStatus,
|
||||
pub calculated_at: String,
|
||||
pub engine_version: String,
|
||||
}
|
||||
|
||||
impl Default for BillingSnapshot {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
schema_version: BILLING_SNAPSHOT_SCHEMA_VERSION.to_string(),
|
||||
rule_id: None,
|
||||
rule_name: None,
|
||||
scope: None,
|
||||
expression: None,
|
||||
resolved_dimensions: BTreeMap::new(),
|
||||
resolved_variables: BTreeMap::new(),
|
||||
cost_breakdown: BTreeMap::new(),
|
||||
total_cost: 0.0,
|
||||
tier_index: None,
|
||||
tier_info: None,
|
||||
missing_required: Vec::new(),
|
||||
status: BillingSnapshotStatus::NoRule,
|
||||
calculated_at: String::new(),
|
||||
engine_version: "2.0".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl BillingSnapshot {
|
||||
pub fn dimensions_used(&self) -> &BTreeMap<String, serde_json::Value> {
|
||||
&self.resolved_dimensions
|
||||
}
|
||||
|
||||
pub fn cost(&self) -> f64 {
|
||||
self.total_cost
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct CostResult {
|
||||
pub cost: f64,
|
||||
pub status: BillingSnapshotStatus,
|
||||
pub snapshot: BillingSnapshot,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{BillingSnapshot, BillingSnapshotStatus};
|
||||
|
||||
#[test]
|
||||
fn snapshot_aliases_dimensions_and_cost() {
|
||||
let mut snapshot = BillingSnapshot {
|
||||
total_cost: 1.25,
|
||||
status: BillingSnapshotStatus::Complete,
|
||||
..BillingSnapshot::default()
|
||||
};
|
||||
snapshot
|
||||
.resolved_dimensions
|
||||
.insert("input_tokens".to_string(), serde_json::json!(10));
|
||||
|
||||
assert_eq!(snapshot.cost(), 1.25);
|
||||
assert_eq!(
|
||||
snapshot.dimensions_used().get("input_tokens"),
|
||||
Some(&serde_json::json!(10))
|
||||
);
|
||||
}
|
||||
}
|
||||
237
crates/aether-billing/src/service.rs
Normal file
237
crates/aether-billing/src/service.rs
Normal file
@@ -0,0 +1,237 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::default_rule::{normalize_task_type, DefaultBillingRuleGenerator};
|
||||
use crate::precision::quantize_cost;
|
||||
use crate::pricing::{BillingComputation, BillingModelPricingSnapshot, BillingUsageInput};
|
||||
use crate::schema::{
|
||||
BillingSnapshot, BillingSnapshotStatus, CostResult, BILLING_SNAPSHOT_SCHEMA_VERSION,
|
||||
};
|
||||
use crate::{
|
||||
normalize_input_tokens_for_billing, ExpressionEvaluationError, FormulaEngine,
|
||||
FormulaEvaluationStatus,
|
||||
};
|
||||
|
||||
pub struct BillingService {
|
||||
engine: FormulaEngine,
|
||||
}
|
||||
|
||||
impl BillingService {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
engine: FormulaEngine::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn calculate(
|
||||
&self,
|
||||
pricing: &BillingModelPricingSnapshot,
|
||||
input: &BillingUsageInput,
|
||||
) -> Result<BillingComputation, ExpressionEvaluationError> {
|
||||
let Some(rule) =
|
||||
DefaultBillingRuleGenerator::generate_for_pricing(pricing, &input.task_type)
|
||||
else {
|
||||
return Ok(BillingComputation {
|
||||
cost_result: CostResult {
|
||||
cost: 0.0,
|
||||
status: BillingSnapshotStatus::NoRule,
|
||||
snapshot: BillingSnapshot {
|
||||
schema_version: BILLING_SNAPSHOT_SCHEMA_VERSION.to_string(),
|
||||
rule_id: None,
|
||||
rule_name: None,
|
||||
scope: None,
|
||||
expression: None,
|
||||
resolved_dimensions: build_dimensions(input),
|
||||
resolved_variables: BTreeMap::new(),
|
||||
cost_breakdown: BTreeMap::new(),
|
||||
total_cost: 0.0,
|
||||
tier_index: None,
|
||||
tier_info: None,
|
||||
missing_required: Vec::new(),
|
||||
status: BillingSnapshotStatus::NoRule,
|
||||
calculated_at: now_marker(),
|
||||
engine_version: "2.0".to_string(),
|
||||
},
|
||||
},
|
||||
actual_total_cost: 0.0,
|
||||
rate_multiplier: pricing
|
||||
.rate_multiplier_for_api_format(input.api_format.as_deref()),
|
||||
is_free_tier: pricing.is_free_tier(),
|
||||
});
|
||||
};
|
||||
|
||||
let dims = build_dimensions(input);
|
||||
let result = self.engine.evaluate(
|
||||
&rule.expression,
|
||||
Some(&rule.variables),
|
||||
Some(&dims),
|
||||
Some(&rule.dimension_mappings),
|
||||
false,
|
||||
)?;
|
||||
|
||||
let status = match result.status {
|
||||
FormulaEvaluationStatus::Complete => BillingSnapshotStatus::Complete,
|
||||
FormulaEvaluationStatus::Incomplete => BillingSnapshotStatus::Incomplete,
|
||||
};
|
||||
let total_cost = if matches!(status, BillingSnapshotStatus::Complete) {
|
||||
result.cost
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
let rate_multiplier = pricing.rate_multiplier_for_api_format(input.api_format.as_deref());
|
||||
let is_free_tier = pricing.is_free_tier();
|
||||
let actual_total_cost = if is_free_tier {
|
||||
0.0
|
||||
} else {
|
||||
quantize_cost(total_cost * rate_multiplier)
|
||||
};
|
||||
|
||||
Ok(BillingComputation {
|
||||
cost_result: CostResult {
|
||||
cost: total_cost,
|
||||
status,
|
||||
snapshot: BillingSnapshot {
|
||||
schema_version: BILLING_SNAPSHOT_SCHEMA_VERSION.to_string(),
|
||||
rule_id: Some(rule.id),
|
||||
rule_name: Some(rule.name),
|
||||
scope: Some(rule.scope),
|
||||
expression: Some(rule.expression),
|
||||
resolved_dimensions: result.resolved_dimensions,
|
||||
resolved_variables: result.resolved_variables,
|
||||
cost_breakdown: result.cost_breakdown,
|
||||
total_cost,
|
||||
tier_index: result.tier_index,
|
||||
tier_info: result.tier_info,
|
||||
missing_required: result.missing_required,
|
||||
status,
|
||||
calculated_at: now_marker(),
|
||||
engine_version: "2.0".to_string(),
|
||||
},
|
||||
},
|
||||
actual_total_cost,
|
||||
rate_multiplier,
|
||||
is_free_tier,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for BillingService {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
fn build_dimensions(input: &BillingUsageInput) -> BTreeMap<String, Value> {
|
||||
let normalized_input_tokens = normalize_input_tokens_for_billing(
|
||||
input.api_format.as_deref(),
|
||||
input.input_tokens,
|
||||
input.cache_read_tokens,
|
||||
);
|
||||
let total_input_context = input
|
||||
.input_tokens
|
||||
.saturating_add(input.cache_creation_tokens)
|
||||
.saturating_add(input.cache_read_tokens);
|
||||
|
||||
let mut out = BTreeMap::from([
|
||||
("input_tokens".to_string(), json!(normalized_input_tokens)),
|
||||
("output_tokens".to_string(), json!(input.output_tokens)),
|
||||
(
|
||||
"cache_creation_tokens".to_string(),
|
||||
json!(input.cache_creation_tokens),
|
||||
),
|
||||
(
|
||||
"cache_read_tokens".to_string(),
|
||||
json!(input.cache_read_tokens),
|
||||
),
|
||||
(
|
||||
"request_count".to_string(),
|
||||
json!(input.request_count.max(0)),
|
||||
),
|
||||
(
|
||||
"total_input_context".to_string(),
|
||||
json!(total_input_context),
|
||||
),
|
||||
(
|
||||
"effective_task_type".to_string(),
|
||||
json!(normalize_task_type(&input.task_type)),
|
||||
),
|
||||
]);
|
||||
|
||||
if let Some(cache_ttl_minutes) = input.cache_ttl_minutes {
|
||||
out.insert(
|
||||
"cache_ttl_minutes".to_string(),
|
||||
json!(cache_ttl_minutes.max(0)),
|
||||
);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn now_marker() -> String {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::BillingService;
|
||||
use crate::{BillingModelPricingSnapshot, BillingSnapshotStatus, BillingUsageInput};
|
||||
|
||||
fn pricing() -> BillingModelPricingSnapshot {
|
||||
BillingModelPricingSnapshot {
|
||||
provider_id: "provider-1".to_string(),
|
||||
provider_billing_type: Some("pay_as_you_go".to_string()),
|
||||
provider_api_key_id: Some("key-1".to_string()),
|
||||
provider_api_key_rate_multipliers: Some(json!({"openai:chat": 0.5})),
|
||||
provider_api_key_cache_ttl_minutes: Some(60),
|
||||
global_model_id: "global-model-1".to_string(),
|
||||
global_model_name: "gpt-5".to_string(),
|
||||
global_model_config: None,
|
||||
default_price_per_request: Some(0.02),
|
||||
default_tiered_pricing: Some(json!({
|
||||
"tiers": [{
|
||||
"up_to": null,
|
||||
"input_price_per_1m": 3.0,
|
||||
"output_price_per_1m": 15.0,
|
||||
"cache_creation_price_per_1m": 3.75,
|
||||
"cache_read_price_per_1m": 0.30
|
||||
}]
|
||||
})),
|
||||
model_id: Some("model-1".to_string()),
|
||||
model_provider_model_name: Some("gpt-5-upstream".to_string()),
|
||||
model_config: None,
|
||||
model_price_per_request: None,
|
||||
model_tiered_pricing: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn calculates_complete_snapshot_for_usage() {
|
||||
let result = BillingService::new()
|
||||
.calculate(
|
||||
&pricing(),
|
||||
&BillingUsageInput {
|
||||
task_type: "chat".to_string(),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
request_count: 1,
|
||||
input_tokens: 1_000,
|
||||
output_tokens: 500,
|
||||
cache_creation_tokens: 0,
|
||||
cache_read_tokens: 100,
|
||||
cache_ttl_minutes: Some(60),
|
||||
},
|
||||
)
|
||||
.expect("billing should calculate");
|
||||
|
||||
assert_eq!(result.cost_result.status, BillingSnapshotStatus::Complete);
|
||||
assert!(result.cost_result.cost > 0.0);
|
||||
assert!(result.actual_total_cost > 0.0);
|
||||
assert_eq!(result.rate_multiplier, 0.5);
|
||||
}
|
||||
}
|
||||
69
crates/aether-billing/src/token_normalization.rs
Normal file
69
crates/aether-billing/src/token_normalization.rs
Normal file
@@ -0,0 +1,69 @@
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum ApiFamily {
|
||||
OpenAi,
|
||||
Claude,
|
||||
Gemini,
|
||||
Unknown,
|
||||
}
|
||||
|
||||
fn parse_api_family(api_format: Option<&str>) -> ApiFamily {
|
||||
let Some(api_format) = api_format else {
|
||||
return ApiFamily::Unknown;
|
||||
};
|
||||
let family = api_format
|
||||
.split(':')
|
||||
.next()
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
match family.as_str() {
|
||||
"openai" => ApiFamily::OpenAi,
|
||||
"claude" | "anthropic" => ApiFamily::Claude,
|
||||
"gemini" | "google" => ApiFamily::Gemini,
|
||||
_ => ApiFamily::Unknown,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn normalize_input_tokens_for_billing(
|
||||
api_format: Option<&str>,
|
||||
input_tokens: i64,
|
||||
cache_read_tokens: i64,
|
||||
) -> i64 {
|
||||
if input_tokens <= 0 {
|
||||
return input_tokens.max(0);
|
||||
}
|
||||
if cache_read_tokens <= 0 {
|
||||
return input_tokens;
|
||||
}
|
||||
|
||||
match parse_api_family(api_format) {
|
||||
ApiFamily::Claude => input_tokens,
|
||||
ApiFamily::OpenAi | ApiFamily::Gemini => (input_tokens - cache_read_tokens).max(0),
|
||||
ApiFamily::Unknown => input_tokens,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::normalize_input_tokens_for_billing;
|
||||
|
||||
#[test]
|
||||
fn subtracts_cache_tokens_for_openai_and_gemini() {
|
||||
assert_eq!(
|
||||
normalize_input_tokens_for_billing(Some("openai:chat"), 100, 20),
|
||||
80
|
||||
);
|
||||
assert_eq!(
|
||||
normalize_input_tokens_for_billing(Some("gemini:chat"), 100, 20),
|
||||
80
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_input_tokens_for_claude() {
|
||||
assert_eq!(
|
||||
normalize_input_tokens_for_billing(Some("claude:chat"), 100, 20),
|
||||
100
|
||||
);
|
||||
}
|
||||
}
|
||||
190
crates/aether-billing/src/usage_mapper.rs
Normal file
190
crates/aether-billing/src/usage_mapper.rs
Normal file
@@ -0,0 +1,190 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::models::StandardizedUsage;
|
||||
|
||||
pub struct UsageMapper;
|
||||
|
||||
impl UsageMapper {
|
||||
pub fn map(
|
||||
raw_usage: &serde_json::Value,
|
||||
api_format: &str,
|
||||
extra_mapping: Option<&BTreeMap<String, String>>,
|
||||
) -> StandardizedUsage {
|
||||
if !raw_usage.is_object() {
|
||||
return StandardizedUsage::new();
|
||||
}
|
||||
|
||||
let mut usage = StandardizedUsage::new();
|
||||
let mut mapping = base_mapping(api_format);
|
||||
if let Some(extra_mapping) = extra_mapping {
|
||||
mapping.extend(extra_mapping.clone());
|
||||
}
|
||||
|
||||
for (source_path, target_field) in mapping {
|
||||
if let Some(value) = get_nested_value(raw_usage, &source_path) {
|
||||
usage.set(&target_field, value.clone());
|
||||
}
|
||||
}
|
||||
|
||||
usage
|
||||
}
|
||||
|
||||
pub fn map_from_response(response: &serde_json::Value, api_format: &str) -> StandardizedUsage {
|
||||
let family = api_family(api_format);
|
||||
let usage_value = if family == "gemini" {
|
||||
response
|
||||
.get("usageMetadata")
|
||||
.or_else(|| {
|
||||
response
|
||||
.get("candidates")
|
||||
.and_then(|v| v.get(0))
|
||||
.and_then(|v| v.get("usageMetadata"))
|
||||
})
|
||||
.unwrap_or(&serde_json::Value::Null)
|
||||
} else {
|
||||
response.get("usage").unwrap_or(&serde_json::Value::Null)
|
||||
};
|
||||
Self::map(usage_value, api_format, None)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn map_usage(raw_usage: &serde_json::Value, api_format: &str) -> StandardizedUsage {
|
||||
UsageMapper::map(raw_usage, api_format, None)
|
||||
}
|
||||
|
||||
pub fn map_usage_from_response(
|
||||
response: &serde_json::Value,
|
||||
api_format: &str,
|
||||
) -> StandardizedUsage {
|
||||
UsageMapper::map_from_response(response, api_format)
|
||||
}
|
||||
|
||||
fn api_family(api_format: &str) -> String {
|
||||
api_format
|
||||
.split(':')
|
||||
.next()
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn base_mapping(api_format: &str) -> BTreeMap<String, String> {
|
||||
let mut mapping = BTreeMap::new();
|
||||
match api_family(api_format).as_str() {
|
||||
"openai" => {
|
||||
mapping.insert("prompt_tokens".to_string(), "input_tokens".to_string());
|
||||
mapping.insert("completion_tokens".to_string(), "output_tokens".to_string());
|
||||
mapping.insert(
|
||||
"prompt_tokens_details.cached_tokens".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"completion_tokens_details.reasoning_tokens".to_string(),
|
||||
"reasoning_tokens".to_string(),
|
||||
);
|
||||
}
|
||||
"gemini" => {
|
||||
mapping.insert("promptTokenCount".to_string(), "input_tokens".to_string());
|
||||
mapping.insert(
|
||||
"candidatesTokenCount".to_string(),
|
||||
"output_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"cachedContentTokenCount".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"usageMetadata.promptTokenCount".to_string(),
|
||||
"input_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"usageMetadata.candidatesTokenCount".to_string(),
|
||||
"output_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"usageMetadata.cachedContentTokenCount".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
}
|
||||
_ => {
|
||||
mapping.insert("input_tokens".to_string(), "input_tokens".to_string());
|
||||
mapping.insert("output_tokens".to_string(), "output_tokens".to_string());
|
||||
mapping.insert(
|
||||
"cache_creation_input_tokens".to_string(),
|
||||
"cache_creation_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"cache_read_input_tokens".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
mapping
|
||||
}
|
||||
|
||||
fn get_nested_value<'a>(value: &'a serde_json::Value, path: &str) -> Option<&'a serde_json::Value> {
|
||||
let mut current = value;
|
||||
for segment in path.split('.') {
|
||||
current = current.get(segment)?;
|
||||
}
|
||||
Some(current)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{map_usage, map_usage_from_response};
|
||||
|
||||
#[test]
|
||||
fn maps_openai_usage() {
|
||||
let usage = map_usage(
|
||||
&serde_json::json!({
|
||||
"prompt_tokens": 12,
|
||||
"completion_tokens": 8,
|
||||
"prompt_tokens_details": { "cached_tokens": 2 },
|
||||
"completion_tokens_details": { "reasoning_tokens": 3 }
|
||||
}),
|
||||
"openai:chat",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 12);
|
||||
assert_eq!(usage.output_tokens, 8);
|
||||
assert_eq!(usage.cache_read_tokens, 2);
|
||||
assert_eq!(usage.reasoning_tokens, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_claude_usage() {
|
||||
let usage = map_usage(
|
||||
&serde_json::json!({
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 5,
|
||||
"cache_creation_input_tokens": 4,
|
||||
"cache_read_input_tokens": 1
|
||||
}),
|
||||
"claude:chat",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 10);
|
||||
assert_eq!(usage.output_tokens, 5);
|
||||
assert_eq!(usage.cache_creation_tokens, 4);
|
||||
assert_eq!(usage.cache_read_tokens, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_gemini_usage_from_response() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 14,
|
||||
"candidatesTokenCount": 6,
|
||||
"cachedContentTokenCount": 2
|
||||
}
|
||||
}),
|
||||
"gemini:chat",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 14);
|
||||
assert_eq!(usage.output_tokens, 6);
|
||||
assert_eq!(usage.cache_read_tokens, 2);
|
||||
}
|
||||
}
|
||||
@@ -69,6 +69,10 @@ where
|
||||
.map(|entries| entries.len())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
impl<K, V> ExpiringMap<K, V>
|
||||
@@ -81,9 +85,7 @@ where
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(entry) = entries.get(key).cloned() else {
|
||||
return None;
|
||||
};
|
||||
let entry = entries.get(key).cloned()?;
|
||||
|
||||
if entry.inserted_at.elapsed() > ttl {
|
||||
entries.remove(key);
|
||||
|
||||
17
crates/aether-crypto/Cargo.toml
Normal file
17
crates/aether-crypto/Cargo.toml
Normal file
@@ -0,0 +1,17 @@
|
||||
[package]
|
||||
name = "aether-crypto"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Shared crypto compatibility helpers for Rust migration"
|
||||
|
||||
[dependencies]
|
||||
aes.workspace = true
|
||||
base64.workspace = true
|
||||
cbc.workspace = true
|
||||
hmac.workspace = true
|
||||
pbkdf2.workspace = true
|
||||
sha2.workspace = true
|
||||
thiserror.workspace = true
|
||||
uuid.workspace = true
|
||||
7
crates/aether-crypto/src/lib.rs
Normal file
7
crates/aether-crypto/src/lib.rs
Normal file
@@ -0,0 +1,7 @@
|
||||
mod python_fernet;
|
||||
|
||||
pub use python_fernet::{
|
||||
decrypt_python_fernet_ciphertext, derive_python_fernet_key, encrypt_python_fernet_plaintext,
|
||||
looks_like_python_fernet_ciphertext, PythonFernetCompat, PythonFernetError, APP_SALT_HEX,
|
||||
APP_SALT_SEED, DEVELOPMENT_ENCRYPTION_KEY,
|
||||
};
|
||||
309
crates/aether-crypto/src/python_fernet.rs
Normal file
309
crates/aether-crypto/src/python_fernet.rs
Normal file
@@ -0,0 +1,309 @@
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
|
||||
use base64::engine::general_purpose::{URL_SAFE, URL_SAFE_NO_PAD};
|
||||
use base64::Engine as _;
|
||||
use cbc::{Decryptor, Encryptor};
|
||||
use hmac::{Hmac, Mac};
|
||||
use pbkdf2::pbkdf2_hmac;
|
||||
use sha2::{Digest, Sha256};
|
||||
use uuid::Uuid;
|
||||
|
||||
const FERNET_VERSION: u8 = 0x80;
|
||||
const HMAC_SIZE: usize = 32;
|
||||
const IV_SIZE: usize = 16;
|
||||
const SIGNING_KEY_SIZE: usize = 16;
|
||||
const ENCRYPTION_KEY_SIZE: usize = 16;
|
||||
const MIN_TOKEN_SIZE: usize = 1 + 8 + IV_SIZE + HMAC_SIZE;
|
||||
const PBKDF2_ITERATIONS: u32 = 100_000;
|
||||
|
||||
pub const APP_SALT_SEED: &[u8] = b"aether-v1";
|
||||
pub const APP_SALT_HEX: &str = "8797080a7a4b45b4810e934d1af36261";
|
||||
pub const DEVELOPMENT_ENCRYPTION_KEY: &str = "dev-encryption-key-do-not-use-in-production";
|
||||
|
||||
type Aes128CbcDec = Decryptor<aes::Aes128>;
|
||||
type Aes128CbcEnc = Encryptor<aes::Aes128>;
|
||||
type HmacSha256 = Hmac<Sha256>;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum PythonFernetError {
|
||||
#[error("invalid Python Fernet outer base64 payload")]
|
||||
InvalidOuterBase64,
|
||||
#[error("invalid Python Fernet inner base64 payload")]
|
||||
InvalidInnerBase64,
|
||||
#[error("invalid Python Fernet token structure")]
|
||||
InvalidTokenStructure,
|
||||
#[error("unsupported Python Fernet token version: {0:#x}")]
|
||||
UnsupportedTokenVersion(u8),
|
||||
#[error("invalid Python Fernet token signature")]
|
||||
InvalidTokenSignature,
|
||||
#[error("invalid Python Fernet token padding")]
|
||||
InvalidPadding,
|
||||
#[error("invalid Python Fernet plaintext utf-8")]
|
||||
InvalidUtf8(#[from] std::string::FromUtf8Error),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PythonFernetCompat {
|
||||
signing_key: [u8; SIGNING_KEY_SIZE],
|
||||
encryption_key: [u8; ENCRYPTION_KEY_SIZE],
|
||||
}
|
||||
|
||||
impl PythonFernetCompat {
|
||||
pub fn from_secret(secret: &str) -> Self {
|
||||
let raw_key = raw_fernet_key(secret);
|
||||
Self::from_raw_key(raw_key)
|
||||
}
|
||||
|
||||
pub fn decrypt_ciphertext(&self, ciphertext: &str) -> Result<String, PythonFernetError> {
|
||||
if ciphertext.is_empty() {
|
||||
return Ok(String::new());
|
||||
}
|
||||
|
||||
let outer =
|
||||
decode_urlsafe(ciphertext).map_err(|_| PythonFernetError::InvalidOuterBase64)?;
|
||||
let inner =
|
||||
decode_urlsafe_bytes(&outer).map_err(|_| PythonFernetError::InvalidInnerBase64)?;
|
||||
let plaintext = self.decrypt_token_bytes(&inner)?;
|
||||
String::from_utf8(plaintext).map_err(PythonFernetError::InvalidUtf8)
|
||||
}
|
||||
|
||||
pub fn encrypt_plaintext(&self, plaintext: &str) -> Result<String, PythonFernetError> {
|
||||
let timestamp = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
self.encrypt_token(plaintext, timestamp, *Uuid::new_v4().as_bytes())
|
||||
}
|
||||
|
||||
fn from_raw_key(raw_key: [u8; 32]) -> Self {
|
||||
let mut signing_key = [0u8; SIGNING_KEY_SIZE];
|
||||
let mut encryption_key = [0u8; ENCRYPTION_KEY_SIZE];
|
||||
signing_key.copy_from_slice(&raw_key[..SIGNING_KEY_SIZE]);
|
||||
encryption_key.copy_from_slice(&raw_key[SIGNING_KEY_SIZE..]);
|
||||
Self {
|
||||
signing_key,
|
||||
encryption_key,
|
||||
}
|
||||
}
|
||||
|
||||
fn decrypt_token_bytes(&self, token: &[u8]) -> Result<Vec<u8>, PythonFernetError> {
|
||||
if token.len() < MIN_TOKEN_SIZE {
|
||||
return Err(PythonFernetError::InvalidTokenStructure);
|
||||
}
|
||||
if token[0] != FERNET_VERSION {
|
||||
return Err(PythonFernetError::UnsupportedTokenVersion(token[0]));
|
||||
}
|
||||
|
||||
let signed_len = token.len() - HMAC_SIZE;
|
||||
let (signed, signature) = token.split_at(signed_len);
|
||||
|
||||
let mut mac = HmacSha256::new_from_slice(&self.signing_key)
|
||||
.map_err(|_| PythonFernetError::InvalidTokenSignature)?;
|
||||
mac.update(signed);
|
||||
mac.verify_slice(signature)
|
||||
.map_err(|_| PythonFernetError::InvalidTokenSignature)?;
|
||||
|
||||
let iv_offset = 1 + 8;
|
||||
let ciphertext_offset = iv_offset + IV_SIZE;
|
||||
let iv = &token[iv_offset..ciphertext_offset];
|
||||
let mut ciphertext = token[ciphertext_offset..signed_len].to_vec();
|
||||
let plaintext = Aes128CbcDec::new((&self.encryption_key).into(), iv.into())
|
||||
.decrypt_padded_mut::<Pkcs7>(&mut ciphertext)
|
||||
.map_err(|_| PythonFernetError::InvalidPadding)?;
|
||||
Ok(plaintext.to_vec())
|
||||
}
|
||||
|
||||
fn encrypt_token(
|
||||
&self,
|
||||
plaintext: &str,
|
||||
timestamp: u64,
|
||||
iv: [u8; IV_SIZE],
|
||||
) -> Result<String, PythonFernetError> {
|
||||
let plaintext = plaintext.as_bytes();
|
||||
let mut padded = vec![0u8; plaintext.len() + IV_SIZE];
|
||||
padded[..plaintext.len()].copy_from_slice(plaintext);
|
||||
let ciphertext = Aes128CbcEnc::new((&self.encryption_key).into(), (&iv).into())
|
||||
.encrypt_padded_mut::<Pkcs7>(&mut padded, plaintext.len())
|
||||
.map_err(|_| PythonFernetError::InvalidPadding)?
|
||||
.to_vec();
|
||||
|
||||
let mut signed = Vec::with_capacity(1 + 8 + IV_SIZE + ciphertext.len() + HMAC_SIZE);
|
||||
signed.push(FERNET_VERSION);
|
||||
signed.extend_from_slice(×tamp.to_be_bytes());
|
||||
signed.extend_from_slice(&iv);
|
||||
signed.extend_from_slice(&ciphertext);
|
||||
|
||||
let mut mac = HmacSha256::new_from_slice(&self.signing_key)
|
||||
.map_err(|_| PythonFernetError::InvalidTokenSignature)?;
|
||||
mac.update(&signed);
|
||||
let signature = mac.finalize().into_bytes();
|
||||
signed.extend_from_slice(&signature);
|
||||
|
||||
let inner = URL_SAFE.encode(signed);
|
||||
Ok(URL_SAFE.encode(inner.as_bytes()))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn derive_python_fernet_key(secret: &str) -> String {
|
||||
URL_SAFE.encode(raw_fernet_key(secret))
|
||||
}
|
||||
|
||||
pub fn decrypt_python_fernet_ciphertext(
|
||||
secret: &str,
|
||||
ciphertext: &str,
|
||||
) -> Result<String, PythonFernetError> {
|
||||
PythonFernetCompat::from_secret(secret).decrypt_ciphertext(ciphertext)
|
||||
}
|
||||
|
||||
pub fn looks_like_python_fernet_ciphertext(ciphertext: &str) -> bool {
|
||||
let ciphertext = ciphertext.trim();
|
||||
if ciphertext.is_empty() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let Ok(outer) = decode_urlsafe(ciphertext) else {
|
||||
return false;
|
||||
};
|
||||
let Ok(inner) = decode_urlsafe_bytes(&outer) else {
|
||||
return false;
|
||||
};
|
||||
|
||||
inner.len() >= MIN_TOKEN_SIZE && inner.first().copied() == Some(FERNET_VERSION)
|
||||
}
|
||||
|
||||
pub fn encrypt_python_fernet_plaintext(
|
||||
secret: &str,
|
||||
plaintext: &str,
|
||||
) -> Result<String, PythonFernetError> {
|
||||
PythonFernetCompat::from_secret(secret).encrypt_plaintext(plaintext)
|
||||
}
|
||||
|
||||
fn raw_fernet_key(secret: &str) -> [u8; 32] {
|
||||
if let Ok(raw_key) = decode_direct_fernet_key(secret) {
|
||||
return raw_key;
|
||||
}
|
||||
|
||||
let mut salt = [0u8; 16];
|
||||
salt.copy_from_slice(&Sha256::digest(APP_SALT_SEED)[..16]);
|
||||
|
||||
let mut raw_key = [0u8; 32];
|
||||
pbkdf2_hmac::<Sha256>(secret.as_bytes(), &salt, PBKDF2_ITERATIONS, &mut raw_key);
|
||||
raw_key
|
||||
}
|
||||
|
||||
fn decode_direct_fernet_key(secret: &str) -> Result<[u8; 32], PythonFernetError> {
|
||||
let decoded = URL_SAFE
|
||||
.decode(secret)
|
||||
.map_err(|_| PythonFernetError::InvalidInnerBase64)?;
|
||||
let raw_key: [u8; 32] = decoded
|
||||
.as_slice()
|
||||
.try_into()
|
||||
.map_err(|_| PythonFernetError::InvalidTokenStructure)?;
|
||||
Ok(raw_key)
|
||||
}
|
||||
|
||||
fn decode_urlsafe(value: &str) -> Result<Vec<u8>, base64::DecodeError> {
|
||||
URL_SAFE
|
||||
.decode(value)
|
||||
.or_else(|_| URL_SAFE_NO_PAD.decode(value))
|
||||
}
|
||||
|
||||
fn decode_urlsafe_bytes(value: &[u8]) -> Result<Vec<u8>, base64::DecodeError> {
|
||||
URL_SAFE
|
||||
.decode(value)
|
||||
.or_else(|_| URL_SAFE_NO_PAD.decode(value))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
decrypt_python_fernet_ciphertext, derive_python_fernet_key,
|
||||
encrypt_python_fernet_plaintext, looks_like_python_fernet_ciphertext, PythonFernetCompat,
|
||||
PythonFernetError, APP_SALT_HEX, DEVELOPMENT_ENCRYPTION_KEY,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn derives_python_pbkdf2_key_for_development_secret() {
|
||||
assert_eq!(APP_SALT_HEX, "8797080a7a4b45b4810e934d1af36261");
|
||||
assert_eq!(
|
||||
derive_python_fernet_key(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
"qGVbbzTSey8Hi1DRtS6wkb2jL33pRBHXTQW-GO6qne0="
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn passes_through_existing_fernet_key_secret() {
|
||||
let direct_key = "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY=";
|
||||
assert_eq!(derive_python_fernet_key(direct_key), direct_key);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn treats_unpadded_direct_key_like_python_pbkdf2_secret() {
|
||||
let unpadded_direct_key = "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY";
|
||||
assert_eq!(
|
||||
derive_python_fernet_key(unpadded_direct_key),
|
||||
"cI8mUtZz6AfpTnBy9xP48Wcp7k_r9h6jJ8jtUoc30cY="
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decrypts_python_compatible_outer_wrapped_ciphertext() {
|
||||
let crypto = PythonFernetCompat::from_secret(DEVELOPMENT_ENCRYPTION_KEY);
|
||||
let ciphertext = crypto
|
||||
.encrypt_token(
|
||||
"{\"api_key\":\"sk-test\",\"provider\":\"openai\"}",
|
||||
1_710_000_000,
|
||||
*b"fixed-fernet-iv!",
|
||||
)
|
||||
.expect("ciphertext should build");
|
||||
|
||||
let plaintext = decrypt_python_fernet_ciphertext(DEVELOPMENT_ENCRYPTION_KEY, &ciphertext)
|
||||
.expect("ciphertext should decrypt");
|
||||
|
||||
assert_eq!(
|
||||
plaintext,
|
||||
"{\"api_key\":\"sk-test\",\"provider\":\"openai\"}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_python_fernet_ciphertext_shape() {
|
||||
let ciphertext = encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-test")
|
||||
.expect("ciphertext should build");
|
||||
|
||||
assert!(looks_like_python_fernet_ciphertext(&ciphertext));
|
||||
assert!(!looks_like_python_fernet_ciphertext("sk-plaintext-openai"));
|
||||
assert!(!looks_like_python_fernet_ciphertext(
|
||||
r#"{"headers":{"x-account-id":"acc-1"}}"#
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_tampered_signature() {
|
||||
let crypto = PythonFernetCompat::from_secret(DEVELOPMENT_ENCRYPTION_KEY);
|
||||
let mut ciphertext = crypto
|
||||
.encrypt_token("secret", 1_710_000_000, *b"fixed-fernet-iv!")
|
||||
.expect("ciphertext should build");
|
||||
ciphertext.replace_range(ciphertext.len() - 2.., "AA");
|
||||
|
||||
let err = decrypt_python_fernet_ciphertext(DEVELOPMENT_ENCRYPTION_KEY, &ciphertext)
|
||||
.expect_err("tampered ciphertext should fail");
|
||||
assert!(matches!(
|
||||
err,
|
||||
PythonFernetError::InvalidInnerBase64
|
||||
| PythonFernetError::InvalidTokenSignature
|
||||
| PythonFernetError::InvalidPadding
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encrypt_and_decrypt_round_trip() {
|
||||
let ciphertext =
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-live-openai")
|
||||
.expect("ciphertext should build");
|
||||
let plaintext = decrypt_python_fernet_ciphertext(DEVELOPMENT_ENCRYPTION_KEY, &ciphertext)
|
||||
.expect("ciphertext should decrypt");
|
||||
assert_eq!(plaintext, "sk-live-openai");
|
||||
}
|
||||
}
|
||||
@@ -8,7 +8,9 @@ description = "Shared data access contracts and config for Aether Rust services"
|
||||
|
||||
[dependencies]
|
||||
aether-cache.workspace = true
|
||||
aether-wallet.workspace = true
|
||||
async-trait.workspace = true
|
||||
chrono.workspace = true
|
||||
futures-util.workspace = true
|
||||
redis.workspace = true
|
||||
serde.workspace = true
|
||||
|
||||
@@ -129,6 +129,14 @@ mod tests {
|
||||
assert!(backends.leases().postgres().is_none());
|
||||
assert!(backends.locks().redis().is_none());
|
||||
assert!(backends.read().auth_api_keys().is_none());
|
||||
assert!(backends.read().auth_modules().is_none());
|
||||
assert!(backends.read().billing().is_none());
|
||||
assert!(backends.read().gemini_file_mappings().is_none());
|
||||
assert!(backends.read().global_models().is_none());
|
||||
assert!(backends.read().management_tokens().is_none());
|
||||
assert!(backends.read().oauth_providers().is_none());
|
||||
assert!(backends.read().proxy_nodes().is_none());
|
||||
assert!(backends.read().minimal_candidate_selection().is_none());
|
||||
assert!(backends.read().request_candidates().is_none());
|
||||
assert!(backends.read().provider_catalog().is_none());
|
||||
assert!(backends.read().usage().is_none());
|
||||
@@ -137,6 +145,7 @@ mod tests {
|
||||
assert!(backends.transactions().postgres().is_none());
|
||||
assert!(backends.workers().redis().is_none());
|
||||
assert!(backends.write().shadow_results().is_none());
|
||||
assert!(backends.write().usage().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -160,13 +169,31 @@ mod tests {
|
||||
assert!(backends.postgres().is_some());
|
||||
assert!(backends.leases().postgres().is_some());
|
||||
assert!(backends.read().auth_api_keys().is_some());
|
||||
assert!(backends.read().auth_modules().is_some());
|
||||
assert!(backends.read().billing().is_some());
|
||||
assert!(backends.read().gemini_file_mappings().is_some());
|
||||
assert!(backends.read().global_models().is_some());
|
||||
assert!(backends.read().management_tokens().is_some());
|
||||
assert!(backends.read().oauth_providers().is_some());
|
||||
assert!(backends.read().proxy_nodes().is_some());
|
||||
assert!(backends.read().minimal_candidate_selection().is_some());
|
||||
assert!(backends.read().request_candidates().is_some());
|
||||
assert!(backends.read().provider_catalog().is_some());
|
||||
assert!(backends.read().provider_quotas().is_some());
|
||||
assert!(backends.read().usage().is_some());
|
||||
assert!(backends.read().video_tasks().is_some());
|
||||
assert!(backends.read().wallets().is_some());
|
||||
assert!(backends.read().shadow_results().is_some());
|
||||
assert!(backends.transactions().postgres().is_some());
|
||||
assert!(backends.write().shadow_results().is_some());
|
||||
assert!(backends.write().auth_modules().is_some());
|
||||
assert!(backends.write().gemini_file_mappings().is_some());
|
||||
assert!(backends.write().management_tokens().is_some());
|
||||
assert!(backends.write().oauth_providers().is_some());
|
||||
assert!(backends.write().proxy_nodes().is_some());
|
||||
assert!(backends.write().provider_quotas().is_some());
|
||||
assert!(backends.write().usage().is_some());
|
||||
assert!(backends.write().wallets().is_some());
|
||||
assert!(backends.config().postgres.is_some());
|
||||
}
|
||||
|
||||
@@ -188,8 +215,12 @@ mod tests {
|
||||
assert!(backends.locks().redis().is_some());
|
||||
assert!(backends.workers().redis().is_some());
|
||||
assert!(backends.read().auth_api_keys().is_none());
|
||||
assert!(backends.read().auth_modules().is_none());
|
||||
assert!(backends.read().global_models().is_none());
|
||||
assert!(backends.read().oauth_providers().is_none());
|
||||
assert!(backends.transactions().postgres().is_none());
|
||||
assert!(backends.write().shadow_results().is_none());
|
||||
assert!(backends.write().usage().is_none());
|
||||
assert!(backends.config().redis.is_some());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,19 +4,117 @@ use crate::postgres::{
|
||||
PostgresLeaseRunner, PostgresLeaseRunnerConfig, PostgresPool, PostgresPoolConfig,
|
||||
PostgresPoolFactory, PostgresTransactionRunner,
|
||||
};
|
||||
use crate::repository::auth::{AuthApiKeyReadRepository, SqlxAuthApiKeySnapshotReadRepository};
|
||||
use crate::repository::announcements::{
|
||||
AnnouncementReadRepository, AnnouncementWriteRepository, SqlxAnnouncementReadRepository,
|
||||
};
|
||||
use crate::repository::auth::{
|
||||
AuthApiKeyReadRepository, AuthApiKeyWriteRepository, SqlxAuthApiKeySnapshotReadRepository,
|
||||
};
|
||||
use crate::repository::auth_modules::{
|
||||
AuthModuleReadRepository, AuthModuleWriteRepository, SqlxAuthModuleReadRepository,
|
||||
SqlxAuthModuleRepository,
|
||||
};
|
||||
use crate::repository::billing::{BillingReadRepository, SqlxBillingReadRepository};
|
||||
use crate::repository::candidate_selection::{
|
||||
MinimalCandidateSelectionReadRepository, SqlxMinimalCandidateSelectionReadRepository,
|
||||
};
|
||||
use crate::repository::candidates::{
|
||||
RequestCandidateReadRepository, SqlxRequestCandidateReadRepository,
|
||||
RequestCandidateReadRepository, RequestCandidateWriteRepository,
|
||||
SqlxRequestCandidateReadRepository,
|
||||
};
|
||||
use crate::repository::gemini_file_mappings::{
|
||||
GeminiFileMappingReadRepository, GeminiFileMappingWriteRepository,
|
||||
SqlxGeminiFileMappingRepository,
|
||||
};
|
||||
use crate::repository::global_models::{
|
||||
GlobalModelReadRepository, GlobalModelWriteRepository, SqlxGlobalModelReadRepository,
|
||||
};
|
||||
use crate::repository::management_tokens::{
|
||||
ManagementTokenReadRepository, ManagementTokenWriteRepository, SqlxManagementTokenRepository,
|
||||
};
|
||||
use crate::repository::oauth_providers::{
|
||||
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqlxOAuthProviderRepository,
|
||||
};
|
||||
use crate::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, SqlxProviderCatalogReadRepository,
|
||||
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||
SqlxProviderCatalogReadRepository,
|
||||
};
|
||||
use crate::repository::proxy_nodes::{
|
||||
ProxyNodeReadRepository, ProxyNodeWriteRepository, SqlxProxyNodeRepository,
|
||||
};
|
||||
use crate::repository::quota::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, SqlxProviderQuotaRepository,
|
||||
};
|
||||
use crate::repository::shadow_results::{
|
||||
ShadowResultReadRepository, ShadowResultWriteRepository, SqlxShadowResultRepository,
|
||||
};
|
||||
use crate::repository::usage::{SqlxUsageReadRepository, UsageReadRepository};
|
||||
use crate::repository::video_tasks::{SqlxVideoTaskReadRepository, VideoTaskReadRepository};
|
||||
use crate::repository::usage::{
|
||||
SqlxUsageReadRepository, UsageReadRepository, UsageWriteRepository,
|
||||
};
|
||||
use crate::repository::users::{SqlxUserReadRepository, UserReadRepository};
|
||||
use crate::repository::video_tasks::{
|
||||
SqlxVideoTaskReadRepository, SqlxVideoTaskRepository, VideoTaskReadRepository,
|
||||
VideoTaskWriteRepository,
|
||||
};
|
||||
use crate::repository::wallet::{
|
||||
SqlxWalletRepository, WalletReadRepository, WalletWriteRepository,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
use sqlx::Row;
|
||||
|
||||
const FIND_SYSTEM_CONFIG_VALUE_SQL: &str = r#"
|
||||
SELECT value
|
||||
FROM system_configs
|
||||
WHERE key = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const UPSERT_SYSTEM_CONFIG_VALUE_SQL: &str = r#"
|
||||
INSERT INTO system_configs (id, key, value, description, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, NOW(), NOW())
|
||||
ON CONFLICT (key) DO UPDATE
|
||||
SET value = EXCLUDED.value,
|
||||
description = COALESCE(EXCLUDED.description, system_configs.description),
|
||||
updated_at = NOW()
|
||||
RETURNING value
|
||||
"#;
|
||||
|
||||
const LIST_SYSTEM_CONFIG_ENTRIES_SQL: &str = r#"
|
||||
SELECT
|
||||
key,
|
||||
value,
|
||||
description,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
FROM system_configs
|
||||
ORDER BY key ASC
|
||||
"#;
|
||||
|
||||
const UPSERT_SYSTEM_CONFIG_ENTRY_SQL: &str = r#"
|
||||
INSERT INTO system_configs (id, key, value, description, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, NOW(), NOW())
|
||||
ON CONFLICT (key) DO UPDATE
|
||||
SET value = EXCLUDED.value,
|
||||
description = COALESCE(EXCLUDED.description, system_configs.description),
|
||||
updated_at = NOW()
|
||||
RETURNING
|
||||
key,
|
||||
value,
|
||||
description,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const DELETE_SYSTEM_CONFIG_VALUE_SQL: &str = r#"
|
||||
DELETE FROM system_configs
|
||||
WHERE key = $1
|
||||
"#;
|
||||
|
||||
const READ_ADMIN_SYSTEM_STATS_SQL: &str = r#"
|
||||
SELECT
|
||||
(SELECT COUNT(id) FROM users) AS total_users,
|
||||
(SELECT COUNT(id) FROM users WHERE is_active IS TRUE) AS active_users,
|
||||
(SELECT COUNT(id) FROM api_keys) AS total_api_keys,
|
||||
(SELECT COUNT(id) FROM usage) AS total_requests
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostgresBackend {
|
||||
@@ -48,22 +146,128 @@ impl PostgresBackend {
|
||||
Arc::new(SqlxAuthApiKeySnapshotReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn announcement_read_repository(&self) -> Arc<dyn AnnouncementReadRepository> {
|
||||
Arc::new(SqlxAnnouncementReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn announcement_write_repository(&self) -> Arc<dyn AnnouncementWriteRepository> {
|
||||
Arc::new(SqlxAnnouncementReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn auth_api_key_write_repository(&self) -> Arc<dyn AuthApiKeyWriteRepository> {
|
||||
Arc::new(SqlxAuthApiKeySnapshotReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn auth_module_read_repository(&self) -> Arc<dyn AuthModuleReadRepository> {
|
||||
Arc::new(SqlxAuthModuleReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn auth_module_write_repository(&self) -> Arc<dyn AuthModuleWriteRepository> {
|
||||
Arc::new(SqlxAuthModuleRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn billing_read_repository(&self) -> Arc<dyn BillingReadRepository> {
|
||||
Arc::new(SqlxBillingReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn minimal_candidate_selection_read_repository(
|
||||
&self,
|
||||
) -> Arc<dyn MinimalCandidateSelectionReadRepository> {
|
||||
Arc::new(SqlxMinimalCandidateSelectionReadRepository::new(
|
||||
self.pool_clone(),
|
||||
))
|
||||
}
|
||||
|
||||
pub fn request_candidate_read_repository(&self) -> Arc<dyn RequestCandidateReadRepository> {
|
||||
Arc::new(SqlxRequestCandidateReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn request_candidate_write_repository(&self) -> Arc<dyn RequestCandidateWriteRepository> {
|
||||
Arc::new(SqlxRequestCandidateReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn gemini_file_mapping_read_repository(&self) -> Arc<dyn GeminiFileMappingReadRepository> {
|
||||
Arc::new(SqlxGeminiFileMappingRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn gemini_file_mapping_write_repository(
|
||||
&self,
|
||||
) -> Arc<dyn GeminiFileMappingWriteRepository> {
|
||||
Arc::new(SqlxGeminiFileMappingRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn global_model_read_repository(&self) -> Arc<dyn GlobalModelReadRepository> {
|
||||
Arc::new(SqlxGlobalModelReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn global_model_write_repository(&self) -> Arc<dyn GlobalModelWriteRepository> {
|
||||
Arc::new(SqlxGlobalModelReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn management_token_read_repository(&self) -> Arc<dyn ManagementTokenReadRepository> {
|
||||
Arc::new(SqlxManagementTokenRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn management_token_write_repository(&self) -> Arc<dyn ManagementTokenWriteRepository> {
|
||||
Arc::new(SqlxManagementTokenRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn oauth_provider_read_repository(&self) -> Arc<dyn OAuthProviderReadRepository> {
|
||||
Arc::new(SqlxOAuthProviderRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn oauth_provider_write_repository(&self) -> Arc<dyn OAuthProviderWriteRepository> {
|
||||
Arc::new(SqlxOAuthProviderRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
|
||||
Arc::new(SqlxProxyNodeRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn proxy_node_write_repository(&self) -> Arc<dyn ProxyNodeWriteRepository> {
|
||||
Arc::new(SqlxProxyNodeRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_catalog_read_repository(&self) -> Arc<dyn ProviderCatalogReadRepository> {
|
||||
Arc::new(SqlxProviderCatalogReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_catalog_write_repository(&self) -> Arc<dyn ProviderCatalogWriteRepository> {
|
||||
Arc::new(SqlxProviderCatalogReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_quota_read_repository(&self) -> Arc<dyn ProviderQuotaReadRepository> {
|
||||
Arc::new(SqlxProviderQuotaRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn usage_read_repository(&self) -> Arc<dyn UsageReadRepository> {
|
||||
Arc::new(SqlxUsageReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn user_read_repository(&self) -> Arc<dyn UserReadRepository> {
|
||||
Arc::new(SqlxUserReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn usage_write_repository(&self) -> Arc<dyn UsageWriteRepository> {
|
||||
Arc::new(SqlxUsageReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn wallet_read_repository(&self) -> Arc<dyn WalletReadRepository> {
|
||||
Arc::new(SqlxWalletRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn wallet_write_repository(&self) -> Arc<dyn WalletWriteRepository> {
|
||||
Arc::new(SqlxWalletRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn video_task_read_repository(&self) -> Arc<dyn VideoTaskReadRepository> {
|
||||
Arc::new(SqlxVideoTaskReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn video_task_write_repository(&self) -> Arc<dyn VideoTaskWriteRepository> {
|
||||
Arc::new(SqlxVideoTaskRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn transaction_runner(&self) -> PostgresTransactionRunner {
|
||||
PostgresTransactionRunner::new(self.pool_clone())
|
||||
}
|
||||
@@ -79,9 +283,103 @@ impl PostgresBackend {
|
||||
Arc::new(SqlxShadowResultRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_quota_write_repository(&self) -> Arc<dyn ProviderQuotaWriteRepository> {
|
||||
Arc::new(SqlxProviderQuotaRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn shadow_result_read_repository(&self) -> Arc<dyn ShadowResultReadRepository> {
|
||||
Arc::new(SqlxShadowResultRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub async fn find_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
) -> Result<Option<serde_json::Value>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_SYSTEM_CONFIG_VALUE_SQL)
|
||||
.bind(key)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.map(|row| row.try_get("value"))
|
||||
.transpose()
|
||||
.map_err(Into::into)
|
||||
}
|
||||
|
||||
pub async fn upsert_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
description: Option<&str>,
|
||||
) -> Result<serde_json::Value, DataLayerError> {
|
||||
let row = sqlx::query(UPSERT_SYSTEM_CONFIG_VALUE_SQL)
|
||||
.bind(uuid::Uuid::new_v4().to_string())
|
||||
.bind(key)
|
||||
.bind(value)
|
||||
.bind(description)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
row.try_get("value").map_err(Into::into)
|
||||
}
|
||||
|
||||
pub async fn list_system_config_entries(
|
||||
&self,
|
||||
) -> Result<Vec<(String, serde_json::Value, Option<String>, Option<u64>)>, DataLayerError> {
|
||||
let rows = sqlx::query(LIST_SYSTEM_CONFIG_ENTRIES_SQL)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.into_iter()
|
||||
.map(|row| {
|
||||
Ok((
|
||||
row.try_get("key")?,
|
||||
row.try_get("value")?,
|
||||
row.try_get("description")?,
|
||||
row.try_get::<Option<i64>, _>("updated_at_unix_secs")?
|
||||
.map(|value| value.max(0) as u64),
|
||||
))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn upsert_system_config_entry(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
description: Option<&str>,
|
||||
) -> Result<(String, serde_json::Value, Option<String>, Option<u64>), DataLayerError> {
|
||||
let row = sqlx::query(UPSERT_SYSTEM_CONFIG_ENTRY_SQL)
|
||||
.bind(uuid::Uuid::new_v4().to_string())
|
||||
.bind(key)
|
||||
.bind(value)
|
||||
.bind(description)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
Ok((
|
||||
row.try_get("key")?,
|
||||
row.try_get("value")?,
|
||||
row.try_get("description")?,
|
||||
row.try_get::<Option<i64>, _>("updated_at_unix_secs")?
|
||||
.map(|value| value.max(0) as u64),
|
||||
))
|
||||
}
|
||||
|
||||
pub async fn delete_system_config_value(&self, key: &str) -> Result<bool, DataLayerError> {
|
||||
let result = sqlx::query(DELETE_SYSTEM_CONFIG_VALUE_SQL)
|
||||
.bind(key)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
pub async fn read_admin_system_stats(&self) -> Result<(u64, u64, u64, u64), DataLayerError> {
|
||||
let row = sqlx::query(READ_ADMIN_SYSTEM_STATS_SQL)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
Ok((
|
||||
row.try_get::<i64, _>("total_users")?.max(0) as u64,
|
||||
row.try_get::<i64, _>("active_users")?.max(0) as u64,
|
||||
row.try_get::<i64, _>("total_api_keys")?.max(0) as u64,
|
||||
row.try_get::<i64, _>("total_requests")?.max(0) as u64,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -109,15 +407,38 @@ mod tests {
|
||||
let _pool = backend.pool();
|
||||
let _pool_clone = backend.pool_clone();
|
||||
let _auth_api_key_reader = backend.auth_api_key_read_repository();
|
||||
let _auth_api_key_writer = backend.auth_api_key_write_repository();
|
||||
let _auth_module_reader = backend.auth_module_read_repository();
|
||||
let _billing_reader = backend.billing_read_repository();
|
||||
let _gemini_file_mapping_reader = backend.gemini_file_mapping_read_repository();
|
||||
let _global_model_reader = backend.global_model_read_repository();
|
||||
let _global_model_writer = backend.global_model_write_repository();
|
||||
let _management_token_reader = backend.management_token_read_repository();
|
||||
let _management_token_writer = backend.management_token_write_repository();
|
||||
let _oauth_provider_reader = backend.oauth_provider_read_repository();
|
||||
let _oauth_provider_writer = backend.oauth_provider_write_repository();
|
||||
let _proxy_node_reader = backend.proxy_node_read_repository();
|
||||
let _proxy_node_writer = backend.proxy_node_write_repository();
|
||||
let _minimal_candidate_selection_reader =
|
||||
backend.minimal_candidate_selection_read_repository();
|
||||
let _request_candidate_reader = backend.request_candidate_read_repository();
|
||||
let _request_candidate_writer = backend.request_candidate_write_repository();
|
||||
let _gemini_file_mapping_writer = backend.gemini_file_mapping_write_repository();
|
||||
let _provider_catalog_reader = backend.provider_catalog_read_repository();
|
||||
let _provider_catalog_writer = backend.provider_catalog_write_repository();
|
||||
let _provider_quota_reader = backend.provider_quota_read_repository();
|
||||
let _usage_reader = backend.usage_read_repository();
|
||||
let _usage_writer = backend.usage_write_repository();
|
||||
let _wallet_reader = backend.wallet_read_repository();
|
||||
let _wallet_writer = backend.wallet_write_repository();
|
||||
let _video_task_reader = backend.video_task_read_repository();
|
||||
let _video_task_writer = backend.video_task_write_repository();
|
||||
let _transaction_runner = backend.transaction_runner();
|
||||
let _lease_runner = backend
|
||||
.lease_runner(PostgresLeaseRunnerConfig::default())
|
||||
.expect("lease runner should build");
|
||||
let _shadow_result_reader = backend.shadow_result_read_repository();
|
||||
let _shadow_result_writer = backend.shadow_result_write_repository();
|
||||
let _provider_quota_writer = backend.provider_quota_write_repository();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,20 +2,44 @@ use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::PostgresBackend;
|
||||
use crate::repository::announcements::AnnouncementReadRepository;
|
||||
use crate::repository::auth::AuthApiKeyReadRepository;
|
||||
use crate::repository::auth_modules::AuthModuleReadRepository;
|
||||
use crate::repository::billing::BillingReadRepository;
|
||||
use crate::repository::candidate_selection::MinimalCandidateSelectionReadRepository;
|
||||
use crate::repository::candidates::RequestCandidateReadRepository;
|
||||
use crate::repository::gemini_file_mappings::GeminiFileMappingReadRepository;
|
||||
use crate::repository::global_models::GlobalModelReadRepository;
|
||||
use crate::repository::management_tokens::ManagementTokenReadRepository;
|
||||
use crate::repository::oauth_providers::OAuthProviderReadRepository;
|
||||
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
|
||||
use crate::repository::proxy_nodes::ProxyNodeReadRepository;
|
||||
use crate::repository::quota::ProviderQuotaReadRepository;
|
||||
use crate::repository::shadow_results::ShadowResultReadRepository;
|
||||
use crate::repository::usage::UsageReadRepository;
|
||||
use crate::repository::users::UserReadRepository;
|
||||
use crate::repository::video_tasks::VideoTaskReadRepository;
|
||||
use crate::repository::wallet::WalletReadRepository;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct DataReadRepositories {
|
||||
announcements: Option<Arc<dyn AnnouncementReadRepository>>,
|
||||
auth_api_keys: Option<Arc<dyn AuthApiKeyReadRepository>>,
|
||||
auth_modules: Option<Arc<dyn AuthModuleReadRepository>>,
|
||||
billing: Option<Arc<dyn BillingReadRepository>>,
|
||||
gemini_file_mappings: Option<Arc<dyn GeminiFileMappingReadRepository>>,
|
||||
global_models: Option<Arc<dyn GlobalModelReadRepository>>,
|
||||
management_tokens: Option<Arc<dyn ManagementTokenReadRepository>>,
|
||||
oauth_providers: Option<Arc<dyn OAuthProviderReadRepository>>,
|
||||
proxy_nodes: Option<Arc<dyn ProxyNodeReadRepository>>,
|
||||
minimal_candidate_selection: Option<Arc<dyn MinimalCandidateSelectionReadRepository>>,
|
||||
request_candidates: Option<Arc<dyn RequestCandidateReadRepository>>,
|
||||
provider_catalog: Option<Arc<dyn ProviderCatalogReadRepository>>,
|
||||
provider_quotas: Option<Arc<dyn ProviderQuotaReadRepository>>,
|
||||
usage: Option<Arc<dyn UsageReadRepository>>,
|
||||
users: Option<Arc<dyn UserReadRepository>>,
|
||||
video_tasks: Option<Arc<dyn VideoTaskReadRepository>>,
|
||||
wallets: Option<Arc<dyn WalletReadRepository>>,
|
||||
shadow_results: Option<Arc<dyn ShadowResultReadRepository>>,
|
||||
}
|
||||
|
||||
@@ -23,10 +47,28 @@ impl fmt::Debug for DataReadRepositories {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("DataReadRepositories")
|
||||
.field("has_auth_api_keys", &self.auth_api_keys.is_some())
|
||||
.field("has_announcements", &self.announcements.is_some())
|
||||
.field("has_auth_modules", &self.auth_modules.is_some())
|
||||
.field("has_billing", &self.billing.is_some())
|
||||
.field(
|
||||
"has_gemini_file_mappings",
|
||||
&self.gemini_file_mappings.is_some(),
|
||||
)
|
||||
.field("has_global_models", &self.global_models.is_some())
|
||||
.field("has_management_tokens", &self.management_tokens.is_some())
|
||||
.field("has_oauth_providers", &self.oauth_providers.is_some())
|
||||
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
|
||||
.field(
|
||||
"has_minimal_candidate_selection",
|
||||
&self.minimal_candidate_selection.is_some(),
|
||||
)
|
||||
.field("has_request_candidates", &self.request_candidates.is_some())
|
||||
.field("has_provider_catalog", &self.provider_catalog.is_some())
|
||||
.field("has_provider_quotas", &self.provider_quotas.is_some())
|
||||
.field("has_usage", &self.usage.is_some())
|
||||
.field("has_users", &self.users.is_some())
|
||||
.field("has_video_tasks", &self.video_tasks.is_some())
|
||||
.field("has_wallets", &self.wallets.is_some())
|
||||
.field("has_shadow_results", &self.shadow_results.is_some())
|
||||
.finish()
|
||||
}
|
||||
@@ -35,11 +77,25 @@ impl fmt::Debug for DataReadRepositories {
|
||||
impl DataReadRepositories {
|
||||
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
|
||||
Self {
|
||||
announcements: postgres.map(PostgresBackend::announcement_read_repository),
|
||||
auth_api_keys: postgres.map(PostgresBackend::auth_api_key_read_repository),
|
||||
auth_modules: postgres.map(PostgresBackend::auth_module_read_repository),
|
||||
billing: postgres.map(PostgresBackend::billing_read_repository),
|
||||
gemini_file_mappings: postgres
|
||||
.map(PostgresBackend::gemini_file_mapping_read_repository),
|
||||
global_models: postgres.map(PostgresBackend::global_model_read_repository),
|
||||
management_tokens: postgres.map(PostgresBackend::management_token_read_repository),
|
||||
oauth_providers: postgres.map(PostgresBackend::oauth_provider_read_repository),
|
||||
proxy_nodes: postgres.map(PostgresBackend::proxy_node_read_repository),
|
||||
minimal_candidate_selection: postgres
|
||||
.map(PostgresBackend::minimal_candidate_selection_read_repository),
|
||||
request_candidates: postgres.map(PostgresBackend::request_candidate_read_repository),
|
||||
provider_catalog: postgres.map(PostgresBackend::provider_catalog_read_repository),
|
||||
provider_quotas: postgres.map(PostgresBackend::provider_quota_read_repository),
|
||||
usage: postgres.map(PostgresBackend::usage_read_repository),
|
||||
users: postgres.map(PostgresBackend::user_read_repository),
|
||||
video_tasks: postgres.map(PostgresBackend::video_task_read_repository),
|
||||
wallets: postgres.map(PostgresBackend::wallet_read_repository),
|
||||
shadow_results: postgres.map(PostgresBackend::shadow_result_read_repository),
|
||||
}
|
||||
}
|
||||
@@ -48,6 +104,44 @@ impl DataReadRepositories {
|
||||
self.auth_api_keys.clone()
|
||||
}
|
||||
|
||||
pub fn announcements(&self) -> Option<Arc<dyn AnnouncementReadRepository>> {
|
||||
self.announcements.clone()
|
||||
}
|
||||
|
||||
pub fn auth_modules(&self) -> Option<Arc<dyn AuthModuleReadRepository>> {
|
||||
self.auth_modules.clone()
|
||||
}
|
||||
|
||||
pub fn billing(&self) -> Option<Arc<dyn BillingReadRepository>> {
|
||||
self.billing.clone()
|
||||
}
|
||||
|
||||
pub fn gemini_file_mappings(&self) -> Option<Arc<dyn GeminiFileMappingReadRepository>> {
|
||||
self.gemini_file_mappings.clone()
|
||||
}
|
||||
|
||||
pub fn global_models(&self) -> Option<Arc<dyn GlobalModelReadRepository>> {
|
||||
self.global_models.clone()
|
||||
}
|
||||
|
||||
pub fn management_tokens(&self) -> Option<Arc<dyn ManagementTokenReadRepository>> {
|
||||
self.management_tokens.clone()
|
||||
}
|
||||
|
||||
pub fn oauth_providers(&self) -> Option<Arc<dyn OAuthProviderReadRepository>> {
|
||||
self.oauth_providers.clone()
|
||||
}
|
||||
|
||||
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeReadRepository>> {
|
||||
self.proxy_nodes.clone()
|
||||
}
|
||||
|
||||
pub fn minimal_candidate_selection(
|
||||
&self,
|
||||
) -> Option<Arc<dyn MinimalCandidateSelectionReadRepository>> {
|
||||
self.minimal_candidate_selection.clone()
|
||||
}
|
||||
|
||||
pub fn request_candidates(&self) -> Option<Arc<dyn RequestCandidateReadRepository>> {
|
||||
self.request_candidates.clone()
|
||||
}
|
||||
@@ -56,24 +150,48 @@ impl DataReadRepositories {
|
||||
self.provider_catalog.clone()
|
||||
}
|
||||
|
||||
pub fn provider_quotas(&self) -> Option<Arc<dyn ProviderQuotaReadRepository>> {
|
||||
self.provider_quotas.clone()
|
||||
}
|
||||
|
||||
pub fn usage(&self) -> Option<Arc<dyn UsageReadRepository>> {
|
||||
self.usage.clone()
|
||||
}
|
||||
|
||||
pub fn users(&self) -> Option<Arc<dyn UserReadRepository>> {
|
||||
self.users.clone()
|
||||
}
|
||||
|
||||
pub fn video_tasks(&self) -> Option<Arc<dyn VideoTaskReadRepository>> {
|
||||
self.video_tasks.clone()
|
||||
}
|
||||
|
||||
pub fn wallets(&self) -> Option<Arc<dyn WalletReadRepository>> {
|
||||
self.wallets.clone()
|
||||
}
|
||||
|
||||
pub fn shadow_results(&self) -> Option<Arc<dyn ShadowResultReadRepository>> {
|
||||
self.shadow_results.clone()
|
||||
}
|
||||
|
||||
pub fn has_any(&self) -> bool {
|
||||
self.auth_api_keys.is_some()
|
||||
|| self.announcements.is_some()
|
||||
|| self.auth_modules.is_some()
|
||||
|| self.billing.is_some()
|
||||
|| self.gemini_file_mappings.is_some()
|
||||
|| self.global_models.is_some()
|
||||
|| self.management_tokens.is_some()
|
||||
|| self.oauth_providers.is_some()
|
||||
|| self.proxy_nodes.is_some()
|
||||
|| self.minimal_candidate_selection.is_some()
|
||||
|| self.request_candidates.is_some()
|
||||
|| self.provider_catalog.is_some()
|
||||
|| self.provider_quotas.is_some()
|
||||
|| self.usage.is_some()
|
||||
|| self.users.is_some()
|
||||
|| self.video_tasks.is_some()
|
||||
|| self.wallets.is_some()
|
||||
|| self.shadow_results.is_some()
|
||||
}
|
||||
}
|
||||
@@ -101,11 +219,22 @@ mod tests {
|
||||
let read = DataReadRepositories::from_postgres(Some(&backend));
|
||||
|
||||
assert!(read.has_any());
|
||||
assert!(read.announcements().is_some());
|
||||
assert!(read.auth_api_keys().is_some());
|
||||
assert!(read.auth_modules().is_some());
|
||||
assert!(read.billing().is_some());
|
||||
assert!(read.gemini_file_mappings().is_some());
|
||||
assert!(read.global_models().is_some());
|
||||
assert!(read.management_tokens().is_some());
|
||||
assert!(read.oauth_providers().is_some());
|
||||
assert!(read.proxy_nodes().is_some());
|
||||
assert!(read.minimal_candidate_selection().is_some());
|
||||
assert!(read.request_candidates().is_some());
|
||||
assert!(read.provider_catalog().is_some());
|
||||
assert!(read.provider_quotas().is_some());
|
||||
assert!(read.usage().is_some());
|
||||
assert!(read.video_tasks().is_some());
|
||||
assert!(read.wallets().is_some());
|
||||
assert!(read.shadow_results().is_some());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use crate::redis::{
|
||||
RedisClient, RedisClientConfig, RedisClientFactory, RedisKeyspace, RedisLockRunner,
|
||||
RedisLockRunnerConfig, RedisStreamRunner, RedisStreamRunnerConfig,
|
||||
RedisClient, RedisClientConfig, RedisClientFactory, RedisKeyspace, RedisKvRunner,
|
||||
RedisKvRunnerConfig, RedisLockRunner, RedisLockRunnerConfig, RedisStreamRunner,
|
||||
RedisStreamRunnerConfig,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -46,12 +47,18 @@ impl RedisBackend {
|
||||
) -> Result<RedisStreamRunner, DataLayerError> {
|
||||
RedisStreamRunner::new(self.client_clone(), self.keyspace(), config)
|
||||
}
|
||||
|
||||
pub fn kv_runner(&self, config: RedisKvRunnerConfig) -> Result<RedisKvRunner, DataLayerError> {
|
||||
RedisKvRunner::new(self.client_clone(), self.keyspace(), config)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::RedisBackend;
|
||||
use crate::redis::{RedisClientConfig, RedisLockRunnerConfig, RedisStreamRunnerConfig};
|
||||
use crate::redis::{
|
||||
RedisClientConfig, RedisKvRunnerConfig, RedisLockRunnerConfig, RedisStreamRunnerConfig,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn backend_retains_config_client_and_shared_runners() {
|
||||
@@ -72,5 +79,8 @@ mod tests {
|
||||
let _stream_runner = backend
|
||||
.stream_runner(RedisStreamRunnerConfig::default())
|
||||
.expect("stream runner should build");
|
||||
let _kv_runner = backend
|
||||
.kv_runner(RedisKvRunnerConfig::default())
|
||||
.expect("kv runner should build");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,17 +2,62 @@ use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::PostgresBackend;
|
||||
use crate::repository::announcements::AnnouncementWriteRepository;
|
||||
use crate::repository::auth::AuthApiKeyWriteRepository;
|
||||
use crate::repository::auth_modules::AuthModuleWriteRepository;
|
||||
use crate::repository::candidates::RequestCandidateWriteRepository;
|
||||
use crate::repository::gemini_file_mappings::GeminiFileMappingWriteRepository;
|
||||
use crate::repository::global_models::GlobalModelWriteRepository;
|
||||
use crate::repository::management_tokens::ManagementTokenWriteRepository;
|
||||
use crate::repository::oauth_providers::OAuthProviderWriteRepository;
|
||||
use crate::repository::provider_catalog::ProviderCatalogWriteRepository;
|
||||
use crate::repository::proxy_nodes::ProxyNodeWriteRepository;
|
||||
use crate::repository::quota::ProviderQuotaWriteRepository;
|
||||
use crate::repository::shadow_results::ShadowResultWriteRepository;
|
||||
use crate::repository::usage::UsageWriteRepository;
|
||||
use crate::repository::video_tasks::VideoTaskWriteRepository;
|
||||
use crate::repository::wallet::WalletWriteRepository;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct DataWriteRepositories {
|
||||
announcements: Option<Arc<dyn AnnouncementWriteRepository>>,
|
||||
auth_api_keys: Option<Arc<dyn AuthApiKeyWriteRepository>>,
|
||||
auth_modules: Option<Arc<dyn AuthModuleWriteRepository>>,
|
||||
shadow_results: Option<Arc<dyn ShadowResultWriteRepository>>,
|
||||
request_candidates: Option<Arc<dyn RequestCandidateWriteRepository>>,
|
||||
gemini_file_mappings: Option<Arc<dyn GeminiFileMappingWriteRepository>>,
|
||||
global_models: Option<Arc<dyn GlobalModelWriteRepository>>,
|
||||
management_tokens: Option<Arc<dyn ManagementTokenWriteRepository>>,
|
||||
oauth_providers: Option<Arc<dyn OAuthProviderWriteRepository>>,
|
||||
proxy_nodes: Option<Arc<dyn ProxyNodeWriteRepository>>,
|
||||
provider_catalog: Option<Arc<dyn ProviderCatalogWriteRepository>>,
|
||||
provider_quotas: Option<Arc<dyn ProviderQuotaWriteRepository>>,
|
||||
usage: Option<Arc<dyn UsageWriteRepository>>,
|
||||
video_tasks: Option<Arc<dyn VideoTaskWriteRepository>>,
|
||||
wallets: Option<Arc<dyn WalletWriteRepository>>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for DataWriteRepositories {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("DataWriteRepositories")
|
||||
.field("has_announcements", &self.announcements.is_some())
|
||||
.field("has_auth_api_keys", &self.auth_api_keys.is_some())
|
||||
.field("has_auth_modules", &self.auth_modules.is_some())
|
||||
.field("has_shadow_results", &self.shadow_results.is_some())
|
||||
.field("has_request_candidates", &self.request_candidates.is_some())
|
||||
.field(
|
||||
"has_gemini_file_mappings",
|
||||
&self.gemini_file_mappings.is_some(),
|
||||
)
|
||||
.field("has_global_models", &self.global_models.is_some())
|
||||
.field("has_management_tokens", &self.management_tokens.is_some())
|
||||
.field("has_oauth_providers", &self.oauth_providers.is_some())
|
||||
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
|
||||
.field("has_provider_catalog", &self.provider_catalog.is_some())
|
||||
.field("has_provider_quotas", &self.provider_quotas.is_some())
|
||||
.field("has_usage", &self.usage.is_some())
|
||||
.field("has_video_tasks", &self.video_tasks.is_some())
|
||||
.field("has_wallets", &self.wallets.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -20,7 +65,22 @@ impl fmt::Debug for DataWriteRepositories {
|
||||
impl DataWriteRepositories {
|
||||
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
|
||||
Self {
|
||||
announcements: postgres.map(PostgresBackend::announcement_write_repository),
|
||||
auth_api_keys: postgres.map(PostgresBackend::auth_api_key_write_repository),
|
||||
auth_modules: postgres.map(PostgresBackend::auth_module_write_repository),
|
||||
shadow_results: postgres.map(PostgresBackend::shadow_result_write_repository),
|
||||
request_candidates: postgres.map(PostgresBackend::request_candidate_write_repository),
|
||||
gemini_file_mappings: postgres
|
||||
.map(PostgresBackend::gemini_file_mapping_write_repository),
|
||||
global_models: postgres.map(PostgresBackend::global_model_write_repository),
|
||||
management_tokens: postgres.map(PostgresBackend::management_token_write_repository),
|
||||
oauth_providers: postgres.map(PostgresBackend::oauth_provider_write_repository),
|
||||
proxy_nodes: postgres.map(PostgresBackend::proxy_node_write_repository),
|
||||
provider_catalog: postgres.map(PostgresBackend::provider_catalog_write_repository),
|
||||
provider_quotas: postgres.map(PostgresBackend::provider_quota_write_repository),
|
||||
usage: postgres.map(PostgresBackend::usage_write_repository),
|
||||
video_tasks: postgres.map(PostgresBackend::video_task_write_repository),
|
||||
wallets: postgres.map(PostgresBackend::wallet_write_repository),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -28,8 +88,78 @@ impl DataWriteRepositories {
|
||||
self.shadow_results.clone()
|
||||
}
|
||||
|
||||
pub fn announcements(&self) -> Option<Arc<dyn AnnouncementWriteRepository>> {
|
||||
self.announcements.clone()
|
||||
}
|
||||
|
||||
pub fn auth_api_keys(&self) -> Option<Arc<dyn AuthApiKeyWriteRepository>> {
|
||||
self.auth_api_keys.clone()
|
||||
}
|
||||
|
||||
pub fn auth_modules(&self) -> Option<Arc<dyn AuthModuleWriteRepository>> {
|
||||
self.auth_modules.clone()
|
||||
}
|
||||
|
||||
pub fn usage(&self) -> Option<Arc<dyn UsageWriteRepository>> {
|
||||
self.usage.clone()
|
||||
}
|
||||
|
||||
pub fn request_candidates(&self) -> Option<Arc<dyn RequestCandidateWriteRepository>> {
|
||||
self.request_candidates.clone()
|
||||
}
|
||||
|
||||
pub fn gemini_file_mappings(&self) -> Option<Arc<dyn GeminiFileMappingWriteRepository>> {
|
||||
self.gemini_file_mappings.clone()
|
||||
}
|
||||
|
||||
pub fn global_models(&self) -> Option<Arc<dyn GlobalModelWriteRepository>> {
|
||||
self.global_models.clone()
|
||||
}
|
||||
|
||||
pub fn management_tokens(&self) -> Option<Arc<dyn ManagementTokenWriteRepository>> {
|
||||
self.management_tokens.clone()
|
||||
}
|
||||
|
||||
pub fn oauth_providers(&self) -> Option<Arc<dyn OAuthProviderWriteRepository>> {
|
||||
self.oauth_providers.clone()
|
||||
}
|
||||
|
||||
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeWriteRepository>> {
|
||||
self.proxy_nodes.clone()
|
||||
}
|
||||
|
||||
pub fn provider_quotas(&self) -> Option<Arc<dyn ProviderQuotaWriteRepository>> {
|
||||
self.provider_quotas.clone()
|
||||
}
|
||||
|
||||
pub fn provider_catalog(&self) -> Option<Arc<dyn ProviderCatalogWriteRepository>> {
|
||||
self.provider_catalog.clone()
|
||||
}
|
||||
|
||||
pub fn video_tasks(&self) -> Option<Arc<dyn VideoTaskWriteRepository>> {
|
||||
self.video_tasks.clone()
|
||||
}
|
||||
|
||||
pub fn wallets(&self) -> Option<Arc<dyn WalletWriteRepository>> {
|
||||
self.wallets.clone()
|
||||
}
|
||||
|
||||
pub fn has_any(&self) -> bool {
|
||||
self.shadow_results.is_some()
|
||||
self.announcements.is_some()
|
||||
|| self.auth_api_keys.is_some()
|
||||
|| self.auth_modules.is_some()
|
||||
|| self.shadow_results.is_some()
|
||||
|| self.request_candidates.is_some()
|
||||
|| self.gemini_file_mappings.is_some()
|
||||
|| self.global_models.is_some()
|
||||
|| self.management_tokens.is_some()
|
||||
|| self.oauth_providers.is_some()
|
||||
|| self.proxy_nodes.is_some()
|
||||
|| self.provider_catalog.is_some()
|
||||
|| self.provider_quotas.is_some()
|
||||
|| self.usage.is_some()
|
||||
|| self.video_tasks.is_some()
|
||||
|| self.wallets.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -56,6 +186,20 @@ mod tests {
|
||||
let write = DataWriteRepositories::from_postgres(Some(&backend));
|
||||
|
||||
assert!(write.has_any());
|
||||
assert!(write.announcements().is_some());
|
||||
assert!(write.auth_api_keys().is_some());
|
||||
assert!(write.auth_modules().is_some());
|
||||
assert!(write.shadow_results().is_some());
|
||||
assert!(write.request_candidates().is_some());
|
||||
assert!(write.gemini_file_mappings().is_some());
|
||||
assert!(write.global_models().is_some());
|
||||
assert!(write.management_tokens().is_some());
|
||||
assert!(write.oauth_providers().is_some());
|
||||
assert!(write.proxy_nodes().is_some());
|
||||
assert!(write.provider_catalog().is_some());
|
||||
assert!(write.provider_quotas().is_some());
|
||||
assert!(write.usage().is_some());
|
||||
assert!(write.video_tasks().is_some());
|
||||
assert!(write.wallets().is_some());
|
||||
}
|
||||
}
|
||||
|
||||
171
crates/aether-data/src/redis/kv.rs
Normal file
171
crates/aether-data/src/redis/kv.rs
Normal file
@@ -0,0 +1,171 @@
|
||||
use std::future::Future;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::redis::{RedisClient, RedisKeyspace};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct RedisKvRunnerConfig {
|
||||
pub command_timeout_ms: Option<u64>,
|
||||
pub default_ttl_seconds: u64,
|
||||
}
|
||||
|
||||
impl Default for RedisKvRunnerConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
command_timeout_ms: Some(1_000),
|
||||
default_ttl_seconds: 300,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisKvRunnerConfig {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
if let Some(timeout) = self.command_timeout_ms {
|
||||
if timeout == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis kv command_timeout_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
if self.default_ttl_seconds == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis kv default_ttl_seconds must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RedisKvRunner {
|
||||
client: RedisClient,
|
||||
keyspace: RedisKeyspace,
|
||||
config: RedisKvRunnerConfig,
|
||||
}
|
||||
|
||||
impl RedisKvRunner {
|
||||
pub fn new(
|
||||
client: RedisClient,
|
||||
keyspace: RedisKeyspace,
|
||||
config: RedisKvRunnerConfig,
|
||||
) -> Result<Self, DataLayerError> {
|
||||
config.validate()?;
|
||||
Ok(Self {
|
||||
client,
|
||||
keyspace,
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn client(&self) -> &RedisClient {
|
||||
&self.client
|
||||
}
|
||||
|
||||
pub fn keyspace(&self) -> &RedisKeyspace {
|
||||
&self.keyspace
|
||||
}
|
||||
|
||||
pub fn config(&self) -> RedisKvRunnerConfig {
|
||||
self.config
|
||||
}
|
||||
|
||||
pub async fn setex(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &str,
|
||||
ttl_seconds: Option<u64>,
|
||||
) -> Result<String, DataLayerError> {
|
||||
let resolved_ttl = ttl_seconds.unwrap_or(self.config.default_ttl_seconds);
|
||||
let namespaced_key = self.keyspace.key(key);
|
||||
self.run_with_timeout("redis kv setex", async {
|
||||
let mut connection = self.client.get_multiplexed_async_connection().await?;
|
||||
Ok(redis::cmd("SETEX")
|
||||
.arg(&namespaced_key)
|
||||
.arg(resolved_ttl)
|
||||
.arg(value)
|
||||
.query_async(&mut connection)
|
||||
.await?)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn del(&self, key: &str) -> Result<i64, DataLayerError> {
|
||||
let namespaced_key = self.keyspace.key(key);
|
||||
self.run_with_timeout("redis kv del", async {
|
||||
let mut connection = self.client.get_multiplexed_async_connection().await?;
|
||||
Ok(redis::cmd("DEL")
|
||||
.arg(&namespaced_key)
|
||||
.query_async(&mut connection)
|
||||
.await?)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_with_timeout<T, F>(
|
||||
&self,
|
||||
operation: &'static str,
|
||||
future: F,
|
||||
) -> Result<T, DataLayerError>
|
||||
where
|
||||
F: Future<Output = Result<T, DataLayerError>>,
|
||||
{
|
||||
if let Some(timeout_ms) = self.config.command_timeout_ms {
|
||||
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
DataLayerError::TimedOut(format!("{operation} exceeded {timeout_ms}ms timeout"))
|
||||
})?
|
||||
} else {
|
||||
future.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{RedisKvRunner, RedisKvRunnerConfig};
|
||||
use crate::redis::{RedisClientConfig, RedisClientFactory, RedisKeyspace};
|
||||
|
||||
fn build_runner() -> RedisKvRunner {
|
||||
let config = RedisClientConfig {
|
||||
url: "redis://localhost/0".to_string(),
|
||||
key_prefix: Some("aether-test".to_string()),
|
||||
};
|
||||
let factory = RedisClientFactory::new(config).expect("redis factory");
|
||||
let client = factory.connect_lazy().expect("connect");
|
||||
let keyspace = factory.config().keyspace();
|
||||
RedisKvRunner::new(client, keyspace, RedisKvRunnerConfig::default()).expect("runner build")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runner_reuses_client_keyspace_and_config() {
|
||||
let runner = build_runner();
|
||||
assert_eq!(
|
||||
runner.keyspace().key("kv:setex:1"),
|
||||
"aether-test:kv:setex:1"
|
||||
);
|
||||
assert_eq!(runner.config(), RedisKvRunnerConfig::default());
|
||||
let _client = runner.client();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_zero_default_ttl() {
|
||||
let config = RedisKvRunnerConfig {
|
||||
command_timeout_ms: Some(100),
|
||||
default_ttl_seconds: 0,
|
||||
};
|
||||
assert!(RedisKvRunner::new(
|
||||
RedisClientFactory::new(RedisClientConfig {
|
||||
url: "redis://localhost/0".to_string(),
|
||||
key_prefix: Some("aether-test".to_string()),
|
||||
})
|
||||
.expect("redis factory")
|
||||
.connect_lazy()
|
||||
.expect("redis client"),
|
||||
RedisKeyspace::new(Some("aether-test")),
|
||||
config,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
}
|
||||
@@ -1,9 +1,11 @@
|
||||
mod client;
|
||||
mod kv;
|
||||
mod lock;
|
||||
mod namespace;
|
||||
mod stream;
|
||||
|
||||
pub use client::{RedisClient, RedisClientConfig, RedisClientFactory};
|
||||
pub use kv::{RedisKvRunner, RedisKvRunnerConfig};
|
||||
pub use lock::{RedisLockKey, RedisLockLease, RedisLockRunner, RedisLockRunnerConfig};
|
||||
pub use namespace::RedisKeyspace;
|
||||
pub use stream::{
|
||||
|
||||
@@ -167,6 +167,15 @@ impl RedisStreamRunner {
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
fields: &BTreeMap<String, String>,
|
||||
) -> Result<String, DataLayerError> {
|
||||
self.append_fields_with_maxlen(stream, fields, None).await
|
||||
}
|
||||
|
||||
pub async fn append_fields_with_maxlen(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
fields: &BTreeMap<String, String>,
|
||||
maxlen: Option<usize>,
|
||||
) -> Result<String, DataLayerError> {
|
||||
validate_stream_name(stream)?;
|
||||
if fields.is_empty() {
|
||||
@@ -178,7 +187,11 @@ impl RedisStreamRunner {
|
||||
self.run_with_timeout("redis stream append", async {
|
||||
let mut connection = self.client.get_multiplexed_async_connection().await?;
|
||||
let mut command = redis::cmd("XADD");
|
||||
command.arg(&stream.0).arg("*");
|
||||
command.arg(&stream.0);
|
||||
if let Some(maxlen) = maxlen.filter(|value| *value > 0) {
|
||||
command.arg("MAXLEN").arg("~").arg(maxlen);
|
||||
}
|
||||
command.arg("*");
|
||||
for (key, value) in fields {
|
||||
command.arg(key).arg(value);
|
||||
}
|
||||
@@ -284,6 +297,28 @@ impl RedisStreamRunner {
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn delete(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
ids: &[String],
|
||||
) -> Result<usize, DataLayerError> {
|
||||
validate_stream_name(stream)?;
|
||||
if ids.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
self.run_with_timeout("redis stream delete", async {
|
||||
let mut connection = self.client.get_multiplexed_async_connection().await?;
|
||||
let mut command = redis::cmd("XDEL");
|
||||
command.arg(&stream.0);
|
||||
for id in ids {
|
||||
command.arg(id);
|
||||
}
|
||||
Ok(command.query_async::<usize>(&mut connection).await?)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn claim_stale(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
@@ -592,6 +627,7 @@ mod tests {
|
||||
runner.ack(&stream, &group, &[]).await.expect("empty ack"),
|
||||
0
|
||||
);
|
||||
assert_eq!(runner.delete(&stream, &[]).await.expect("empty delete"), 0);
|
||||
assert!(runner
|
||||
.claim_stale(
|
||||
&stream,
|
||||
|
||||
379
crates/aether-data/src/repository/announcements/memory.rs
Normal file
379
crates/aether-data/src/repository/announcements/memory.rs
Normal file
@@ -0,0 +1,379 @@
|
||||
use std::collections::BTreeSet;
|
||||
use std::sync::RwLock;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::types::{
|
||||
AnnouncementListQuery, AnnouncementReadRepository, AnnouncementWriteRepository,
|
||||
CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, UpdateAnnouncementRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryAnnouncementReadRepository {
|
||||
announcements: RwLock<Vec<StoredAnnouncement>>,
|
||||
announcement_reads: RwLock<BTreeSet<(String, String)>>,
|
||||
}
|
||||
|
||||
impl InMemoryAnnouncementReadRepository {
|
||||
pub fn seed<I>(announcements: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredAnnouncement>,
|
||||
{
|
||||
Self::seed_with_reads(announcements, std::iter::empty::<(String, String)>())
|
||||
}
|
||||
|
||||
pub fn seed_with_reads<I, J>(announcements: I, reads: J) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredAnnouncement>,
|
||||
J: IntoIterator<Item = (String, String)>,
|
||||
{
|
||||
Self {
|
||||
announcements: RwLock::new(announcements.into_iter().collect()),
|
||||
announcement_reads: RwLock::new(reads.into_iter().collect()),
|
||||
}
|
||||
}
|
||||
|
||||
fn now_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AnnouncementReadRepository for InMemoryAnnouncementReadRepository {
|
||||
async fn find_by_id(
|
||||
&self,
|
||||
announcement_id: &str,
|
||||
) -> Result<Option<StoredAnnouncement>, DataLayerError> {
|
||||
Ok(self
|
||||
.announcements
|
||||
.read()
|
||||
.expect("announcement repository lock")
|
||||
.iter()
|
||||
.find(|announcement| announcement.id == announcement_id)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn list_announcements(
|
||||
&self,
|
||||
query: &AnnouncementListQuery,
|
||||
) -> Result<StoredAnnouncementPage, DataLayerError> {
|
||||
let now_unix_secs = query.now_unix_secs.unwrap_or_else(Self::now_unix_secs);
|
||||
let announcements = self
|
||||
.announcements
|
||||
.read()
|
||||
.expect("announcement repository lock");
|
||||
|
||||
let mut items: Vec<_> = announcements
|
||||
.iter()
|
||||
.filter(|announcement| {
|
||||
if !query.active_only {
|
||||
return true;
|
||||
}
|
||||
announcement.is_active
|
||||
&& announcement
|
||||
.start_time_unix_secs
|
||||
.is_none_or(|value| value <= now_unix_secs)
|
||||
&& announcement
|
||||
.end_time_unix_secs
|
||||
.is_none_or(|value| value >= now_unix_secs)
|
||||
})
|
||||
.cloned()
|
||||
.collect();
|
||||
|
||||
items.sort_by(|left, right| {
|
||||
right
|
||||
.is_pinned
|
||||
.cmp(&left.is_pinned)
|
||||
.then_with(|| right.priority.cmp(&left.priority))
|
||||
.then_with(|| right.created_at_unix_secs.cmp(&left.created_at_unix_secs))
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
|
||||
let total = items.len() as u64;
|
||||
let items = items
|
||||
.into_iter()
|
||||
.skip(query.offset)
|
||||
.take(query.limit)
|
||||
.collect();
|
||||
|
||||
Ok(StoredAnnouncementPage { items, total })
|
||||
}
|
||||
|
||||
async fn count_unread_active_announcements(
|
||||
&self,
|
||||
user_id: &str,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
let announcements = self
|
||||
.announcements
|
||||
.read()
|
||||
.expect("announcement repository lock");
|
||||
let reads = self
|
||||
.announcement_reads
|
||||
.read()
|
||||
.expect("announcement reads repository lock");
|
||||
|
||||
let total = announcements
|
||||
.iter()
|
||||
.filter(|announcement| {
|
||||
announcement.is_active
|
||||
&& announcement
|
||||
.start_time_unix_secs
|
||||
.is_none_or(|value| value <= now_unix_secs)
|
||||
&& announcement
|
||||
.end_time_unix_secs
|
||||
.is_none_or(|value| value >= now_unix_secs)
|
||||
&& !reads.contains(&(user_id.to_string(), announcement.id.clone()))
|
||||
})
|
||||
.count() as u64;
|
||||
|
||||
Ok(total)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AnnouncementWriteRepository for InMemoryAnnouncementReadRepository {
|
||||
async fn create_announcement(
|
||||
&self,
|
||||
record: CreateAnnouncementRecord,
|
||||
) -> Result<StoredAnnouncement, DataLayerError> {
|
||||
record.validate()?;
|
||||
let now_unix_secs = Self::now_unix_secs();
|
||||
let announcement = StoredAnnouncement::new(
|
||||
Uuid::new_v4().to_string(),
|
||||
record.title,
|
||||
record.content,
|
||||
record.kind,
|
||||
record.priority,
|
||||
true,
|
||||
record.is_pinned,
|
||||
Some(record.author_id),
|
||||
None,
|
||||
record.start_time_unix_secs.map(|value| value as i64),
|
||||
record.end_time_unix_secs.map(|value| value as i64),
|
||||
now_unix_secs as i64,
|
||||
now_unix_secs as i64,
|
||||
)?;
|
||||
self.announcements
|
||||
.write()
|
||||
.expect("announcement repository lock")
|
||||
.push(announcement.clone());
|
||||
Ok(announcement)
|
||||
}
|
||||
|
||||
async fn update_announcement(
|
||||
&self,
|
||||
record: UpdateAnnouncementRecord,
|
||||
) -> Result<Option<StoredAnnouncement>, DataLayerError> {
|
||||
record.validate()?;
|
||||
let mut announcements = self
|
||||
.announcements
|
||||
.write()
|
||||
.expect("announcement repository lock");
|
||||
let Some(announcement) = announcements
|
||||
.iter_mut()
|
||||
.find(|announcement| announcement.id == record.announcement_id)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if let Some(title) = record.title {
|
||||
announcement.title = title;
|
||||
}
|
||||
if let Some(content) = record.content {
|
||||
announcement.content = content;
|
||||
}
|
||||
if let Some(kind) = record.kind {
|
||||
announcement.kind = kind;
|
||||
}
|
||||
if let Some(priority) = record.priority {
|
||||
announcement.priority = priority;
|
||||
}
|
||||
if let Some(is_active) = record.is_active {
|
||||
announcement.is_active = is_active;
|
||||
}
|
||||
if let Some(is_pinned) = record.is_pinned {
|
||||
announcement.is_pinned = is_pinned;
|
||||
}
|
||||
if let Some(start_time_unix_secs) = record.start_time_unix_secs {
|
||||
announcement.start_time_unix_secs = Some(start_time_unix_secs);
|
||||
}
|
||||
if let Some(end_time_unix_secs) = record.end_time_unix_secs {
|
||||
announcement.end_time_unix_secs = Some(end_time_unix_secs);
|
||||
}
|
||||
announcement.updated_at_unix_secs = Self::now_unix_secs();
|
||||
Ok(Some(announcement.clone()))
|
||||
}
|
||||
|
||||
async fn delete_announcement(&self, announcement_id: &str) -> Result<bool, DataLayerError> {
|
||||
let mut announcements = self
|
||||
.announcements
|
||||
.write()
|
||||
.expect("announcement repository lock");
|
||||
let original_len = announcements.len();
|
||||
announcements.retain(|announcement| announcement.id != announcement_id);
|
||||
Ok(announcements.len() != original_len)
|
||||
}
|
||||
|
||||
async fn mark_announcement_as_read(
|
||||
&self,
|
||||
user_id: &str,
|
||||
announcement_id: &str,
|
||||
_read_at_unix_secs: u64,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let inserted = self
|
||||
.announcement_reads
|
||||
.write()
|
||||
.expect("announcement reads repository lock")
|
||||
.insert((user_id.to_string(), announcement_id.to_string()));
|
||||
Ok(inserted)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryAnnouncementReadRepository;
|
||||
use crate::repository::announcements::{
|
||||
AnnouncementReadRepository, AnnouncementWriteRepository, CreateAnnouncementRecord,
|
||||
StoredAnnouncement, UpdateAnnouncementRecord,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_seeded_announcements() {
|
||||
let repository = InMemoryAnnouncementReadRepository::seed(vec![StoredAnnouncement::new(
|
||||
"announcement-1".to_string(),
|
||||
"系统维护".to_string(),
|
||||
"今天晚些时候维护".to_string(),
|
||||
"maintenance".to_string(),
|
||||
10,
|
||||
true,
|
||||
true,
|
||||
Some("admin-1".to_string()),
|
||||
Some("admin".to_string()),
|
||||
None,
|
||||
None,
|
||||
1_711_000_000,
|
||||
1_711_000_100,
|
||||
)
|
||||
.expect("announcement should build")]);
|
||||
|
||||
let announcement = repository
|
||||
.find_by_id("announcement-1")
|
||||
.await
|
||||
.expect("announcement should load")
|
||||
.expect("announcement should exist");
|
||||
|
||||
assert_eq!(announcement.title, "系统维护");
|
||||
assert_eq!(announcement.author_username.as_deref(), Some("admin"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mutates_seeded_announcements() {
|
||||
let repository = InMemoryAnnouncementReadRepository::seed(vec![]);
|
||||
|
||||
let created = repository
|
||||
.create_announcement(CreateAnnouncementRecord {
|
||||
title: "系统维护".to_string(),
|
||||
content: "今天晚些时候维护".to_string(),
|
||||
kind: "maintenance".to_string(),
|
||||
priority: 10,
|
||||
is_pinned: true,
|
||||
author_id: "admin-1".to_string(),
|
||||
start_time_unix_secs: None,
|
||||
end_time_unix_secs: None,
|
||||
})
|
||||
.await
|
||||
.expect("create should succeed");
|
||||
assert_eq!(created.kind, "maintenance");
|
||||
|
||||
let updated = repository
|
||||
.update_announcement(UpdateAnnouncementRecord {
|
||||
announcement_id: created.id.clone(),
|
||||
title: Some("系统升级".to_string()),
|
||||
content: None,
|
||||
kind: Some("important".to_string()),
|
||||
priority: Some(99),
|
||||
is_active: Some(false),
|
||||
is_pinned: Some(false),
|
||||
start_time_unix_secs: None,
|
||||
end_time_unix_secs: None,
|
||||
})
|
||||
.await
|
||||
.expect("update should succeed")
|
||||
.expect("announcement should exist");
|
||||
assert_eq!(updated.title, "系统升级");
|
||||
assert_eq!(updated.kind, "important");
|
||||
assert!(!updated.is_active);
|
||||
|
||||
let deleted = repository
|
||||
.delete_announcement(&created.id)
|
||||
.await
|
||||
.expect("delete should succeed");
|
||||
assert!(deleted);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tracks_user_announcement_read_state() {
|
||||
let repository = InMemoryAnnouncementReadRepository::seed_with_reads(
|
||||
vec![
|
||||
StoredAnnouncement::new(
|
||||
"announcement-1".to_string(),
|
||||
"系统维护".to_string(),
|
||||
"今天晚些时候维护".to_string(),
|
||||
"maintenance".to_string(),
|
||||
10,
|
||||
true,
|
||||
true,
|
||||
Some("admin-1".to_string()),
|
||||
Some("admin".to_string()),
|
||||
None,
|
||||
None,
|
||||
1_711_000_000,
|
||||
1_711_000_100,
|
||||
)
|
||||
.expect("announcement should build"),
|
||||
StoredAnnouncement::new(
|
||||
"announcement-2".to_string(),
|
||||
"系统升级".to_string(),
|
||||
"升级说明".to_string(),
|
||||
"info".to_string(),
|
||||
5,
|
||||
true,
|
||||
false,
|
||||
Some("admin-1".to_string()),
|
||||
Some("admin".to_string()),
|
||||
None,
|
||||
None,
|
||||
1_711_000_000,
|
||||
1_711_000_100,
|
||||
)
|
||||
.expect("announcement should build"),
|
||||
],
|
||||
[("user-1".to_string(), "announcement-1".to_string())],
|
||||
);
|
||||
|
||||
let unread = repository
|
||||
.count_unread_active_announcements("user-1", 1_711_000_200)
|
||||
.await
|
||||
.expect("count should succeed");
|
||||
assert_eq!(unread, 1);
|
||||
|
||||
let inserted = repository
|
||||
.mark_announcement_as_read("user-1", "announcement-2", 1_711_000_300)
|
||||
.await
|
||||
.expect("mark read should succeed");
|
||||
assert!(inserted);
|
||||
|
||||
let unread = repository
|
||||
.count_unread_active_announcements("user-1", 1_711_000_200)
|
||||
.await
|
||||
.expect("count should succeed");
|
||||
assert_eq!(unread, 0);
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/announcements/mod.rs
Normal file
10
crates/aether-data/src/repository/announcements/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryAnnouncementReadRepository;
|
||||
pub use sql::SqlxAnnouncementReadRepository;
|
||||
pub use types::{
|
||||
AnnouncementListQuery, AnnouncementReadRepository, AnnouncementWriteRepository,
|
||||
CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, UpdateAnnouncementRecord,
|
||||
};
|
||||
369
crates/aether-data/src/repository/announcements/sql.rs
Normal file
369
crates/aether-data/src/repository/announcements/sql.rs
Normal file
@@ -0,0 +1,369 @@
|
||||
use async_trait::async_trait;
|
||||
use chrono::{TimeZone, Utc};
|
||||
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
AnnouncementListQuery, AnnouncementReadRepository, AnnouncementWriteRepository,
|
||||
CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, UpdateAnnouncementRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_ANNOUNCEMENT_BY_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
a.id,
|
||||
a.title,
|
||||
a.content,
|
||||
a.type,
|
||||
a.priority,
|
||||
a.is_active,
|
||||
a.is_pinned,
|
||||
a.author_id,
|
||||
u.username AS author_username,
|
||||
EXTRACT(EPOCH FROM a.start_time)::bigint AS start_time_unix_secs,
|
||||
EXTRACT(EPOCH FROM a.end_time)::bigint AS end_time_unix_secs,
|
||||
EXTRACT(EPOCH FROM a.created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM a.updated_at)::bigint AS updated_at_unix_secs
|
||||
FROM announcements a
|
||||
LEFT JOIN users u ON u.id = a.author_id
|
||||
WHERE a.id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const LIST_ANNOUNCEMENTS_SQL: &str = r#"
|
||||
SELECT
|
||||
a.id,
|
||||
a.title,
|
||||
a.content,
|
||||
a.type,
|
||||
a.priority,
|
||||
a.is_active,
|
||||
a.is_pinned,
|
||||
a.author_id,
|
||||
u.username AS author_username,
|
||||
EXTRACT(EPOCH FROM a.start_time)::bigint AS start_time_unix_secs,
|
||||
EXTRACT(EPOCH FROM a.end_time)::bigint AS end_time_unix_secs,
|
||||
EXTRACT(EPOCH FROM a.created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM a.updated_at)::bigint AS updated_at_unix_secs
|
||||
FROM announcements a
|
||||
LEFT JOIN users u ON u.id = a.author_id
|
||||
WHERE (
|
||||
NOT $1 OR (
|
||||
a.is_active = TRUE
|
||||
AND (a.start_time IS NULL OR a.start_time <= TO_TIMESTAMP($2::double precision))
|
||||
AND (a.end_time IS NULL OR a.end_time >= TO_TIMESTAMP($2::double precision))
|
||||
)
|
||||
)
|
||||
ORDER BY a.is_pinned DESC, a.priority DESC, a.created_at DESC, a.id ASC
|
||||
OFFSET $3
|
||||
LIMIT $4
|
||||
"#;
|
||||
|
||||
const COUNT_ANNOUNCEMENTS_SQL: &str = r#"
|
||||
SELECT COUNT(a.id) AS total
|
||||
FROM announcements a
|
||||
WHERE (
|
||||
NOT $1 OR (
|
||||
a.is_active = TRUE
|
||||
AND (a.start_time IS NULL OR a.start_time <= TO_TIMESTAMP($2::double precision))
|
||||
AND (a.end_time IS NULL OR a.end_time >= TO_TIMESTAMP($2::double precision))
|
||||
)
|
||||
)
|
||||
"#;
|
||||
|
||||
const COUNT_UNREAD_ACTIVE_ANNOUNCEMENTS_SQL: &str = r#"
|
||||
SELECT COUNT(a.id) AS total
|
||||
FROM announcements a
|
||||
WHERE a.is_active = TRUE
|
||||
AND (a.start_time IS NULL OR a.start_time <= TO_TIMESTAMP($2::double precision))
|
||||
AND (a.end_time IS NULL OR a.end_time >= TO_TIMESTAMP($2::double precision))
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM announcement_reads r
|
||||
WHERE r.user_id = $1
|
||||
AND r.announcement_id = a.id
|
||||
)
|
||||
"#;
|
||||
|
||||
const CREATE_ANNOUNCEMENT_SQL: &str = r#"
|
||||
INSERT INTO announcements (
|
||||
id,
|
||||
title,
|
||||
content,
|
||||
type,
|
||||
priority,
|
||||
author_id,
|
||||
is_active,
|
||||
is_pinned,
|
||||
start_time,
|
||||
end_time,
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4,
|
||||
$5,
|
||||
$6,
|
||||
TRUE,
|
||||
$7,
|
||||
$8,
|
||||
$9,
|
||||
NOW(),
|
||||
NOW()
|
||||
)
|
||||
RETURNING
|
||||
id,
|
||||
title,
|
||||
content,
|
||||
type,
|
||||
priority,
|
||||
is_active,
|
||||
is_pinned,
|
||||
author_id,
|
||||
(SELECT username FROM users WHERE id = announcements.author_id) AS author_username,
|
||||
EXTRACT(EPOCH FROM start_time)::bigint AS start_time_unix_secs,
|
||||
EXTRACT(EPOCH FROM end_time)::bigint AS end_time_unix_secs,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const UPDATE_ANNOUNCEMENT_SQL: &str = r#"
|
||||
UPDATE announcements
|
||||
SET
|
||||
title = COALESCE($2, title),
|
||||
content = COALESCE($3, content),
|
||||
type = COALESCE($4, type),
|
||||
priority = COALESCE($5, priority),
|
||||
is_active = COALESCE($6, is_active),
|
||||
is_pinned = COALESCE($7, is_pinned),
|
||||
start_time = COALESCE($8, start_time),
|
||||
end_time = COALESCE($9, end_time),
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
RETURNING
|
||||
id,
|
||||
title,
|
||||
content,
|
||||
type,
|
||||
priority,
|
||||
is_active,
|
||||
is_pinned,
|
||||
author_id,
|
||||
(SELECT username FROM users WHERE id = announcements.author_id) AS author_username,
|
||||
EXTRACT(EPOCH FROM start_time)::bigint AS start_time_unix_secs,
|
||||
EXTRACT(EPOCH FROM end_time)::bigint AS end_time_unix_secs,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const DELETE_ANNOUNCEMENT_SQL: &str = r#"
|
||||
DELETE FROM announcements
|
||||
WHERE id = $1
|
||||
"#;
|
||||
|
||||
const MARK_ANNOUNCEMENT_AS_READ_SQL: &str = r#"
|
||||
INSERT INTO announcement_reads (
|
||||
id,
|
||||
user_id,
|
||||
announcement_id,
|
||||
read_at
|
||||
)
|
||||
VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
TO_TIMESTAMP($4::double precision)
|
||||
)
|
||||
ON CONFLICT (user_id, announcement_id) DO NOTHING
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxAnnouncementReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxAnnouncementReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AnnouncementReadRepository for SqlxAnnouncementReadRepository {
|
||||
async fn find_by_id(
|
||||
&self,
|
||||
announcement_id: &str,
|
||||
) -> Result<Option<StoredAnnouncement>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_ANNOUNCEMENT_BY_ID_SQL)
|
||||
.bind(announcement_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_announcement_row).transpose()
|
||||
}
|
||||
|
||||
async fn list_announcements(
|
||||
&self,
|
||||
query: &AnnouncementListQuery,
|
||||
) -> Result<StoredAnnouncementPage, DataLayerError> {
|
||||
let now_unix_secs = query.now_unix_secs.unwrap_or_else(current_unix_secs);
|
||||
let total_row = sqlx::query(COUNT_ANNOUNCEMENTS_SQL)
|
||||
.bind(query.active_only)
|
||||
.bind(now_unix_secs as f64)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
let total = total_row.try_get::<i64, _>("total")?.max(0) as u64;
|
||||
|
||||
let rows = sqlx::query(LIST_ANNOUNCEMENTS_SQL)
|
||||
.bind(query.active_only)
|
||||
.bind(now_unix_secs as f64)
|
||||
.bind(query.offset as i64)
|
||||
.bind(query.limit as i64)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
let items = rows
|
||||
.iter()
|
||||
.map(map_announcement_row)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
|
||||
Ok(StoredAnnouncementPage { items, total })
|
||||
}
|
||||
|
||||
async fn count_unread_active_announcements(
|
||||
&self,
|
||||
user_id: &str,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
let row = sqlx::query(COUNT_UNREAD_ACTIVE_ANNOUNCEMENTS_SQL)
|
||||
.bind(user_id)
|
||||
.bind(now_unix_secs as f64)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
Ok(row.try_get::<i64, _>("total")?.max(0) as u64)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AnnouncementWriteRepository for SqlxAnnouncementReadRepository {
|
||||
async fn create_announcement(
|
||||
&self,
|
||||
record: CreateAnnouncementRecord,
|
||||
) -> Result<StoredAnnouncement, DataLayerError> {
|
||||
record.validate()?;
|
||||
let row = sqlx::query(CREATE_ANNOUNCEMENT_SQL)
|
||||
.bind(uuid::Uuid::new_v4().to_string())
|
||||
.bind(record.title)
|
||||
.bind(record.content)
|
||||
.bind(record.kind)
|
||||
.bind(record.priority)
|
||||
.bind(record.author_id)
|
||||
.bind(record.is_pinned)
|
||||
.bind(optional_datetime(record.start_time_unix_secs))
|
||||
.bind(optional_datetime(record.end_time_unix_secs))
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
map_announcement_row(&row)
|
||||
}
|
||||
|
||||
async fn update_announcement(
|
||||
&self,
|
||||
record: UpdateAnnouncementRecord,
|
||||
) -> Result<Option<StoredAnnouncement>, DataLayerError> {
|
||||
record.validate()?;
|
||||
let row = sqlx::query(UPDATE_ANNOUNCEMENT_SQL)
|
||||
.bind(record.announcement_id)
|
||||
.bind(record.title)
|
||||
.bind(record.content)
|
||||
.bind(record.kind)
|
||||
.bind(record.priority)
|
||||
.bind(record.is_active)
|
||||
.bind(record.is_pinned)
|
||||
.bind(optional_datetime(record.start_time_unix_secs))
|
||||
.bind(optional_datetime(record.end_time_unix_secs))
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_announcement_row).transpose()
|
||||
}
|
||||
|
||||
async fn delete_announcement(&self, announcement_id: &str) -> Result<bool, DataLayerError> {
|
||||
let result = sqlx::query(DELETE_ANNOUNCEMENT_SQL)
|
||||
.bind(announcement_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
async fn mark_announcement_as_read(
|
||||
&self,
|
||||
user_id: &str,
|
||||
announcement_id: &str,
|
||||
read_at_unix_secs: u64,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let result = sqlx::query(MARK_ANNOUNCEMENT_AS_READ_SQL)
|
||||
.bind(uuid::Uuid::new_v4().to_string())
|
||||
.bind(user_id)
|
||||
.bind(announcement_id)
|
||||
.bind(read_at_unix_secs as f64)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_datetime(unix_secs: Option<u64>) -> Option<chrono::DateTime<Utc>> {
|
||||
unix_secs.and_then(|value| {
|
||||
i64::try_from(value)
|
||||
.ok()
|
||||
.and_then(|value| Utc.timestamp_opt(value, 0).single())
|
||||
})
|
||||
}
|
||||
|
||||
fn current_unix_secs() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
fn map_announcement_row(row: &PgRow) -> Result<StoredAnnouncement, DataLayerError> {
|
||||
StoredAnnouncement::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("title")?,
|
||||
row.try_get("content")?,
|
||||
row.try_get("type")?,
|
||||
row.try_get("priority")?,
|
||||
row.try_get("is_active")?,
|
||||
row.try_get("is_pinned")?,
|
||||
row.try_get("author_id")?,
|
||||
row.try_get("author_username")?,
|
||||
row.try_get("start_time_unix_secs")?,
|
||||
row.try_get("end_time_unix_secs")?,
|
||||
row.try_get("created_at_unix_secs")?,
|
||||
row.try_get("updated_at_unix_secs")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxAnnouncementReadRepository;
|
||||
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 = SqlxAnnouncementReadRepository::new(pool);
|
||||
}
|
||||
}
|
||||
237
crates/aether-data/src/repository/announcements/types.rs
Normal file
237
crates/aether-data/src/repository/announcements/types.rs
Normal file
@@ -0,0 +1,237 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredAnnouncement {
|
||||
pub id: String,
|
||||
pub title: String,
|
||||
pub content: String,
|
||||
pub kind: String,
|
||||
pub priority: i32,
|
||||
pub is_active: bool,
|
||||
pub is_pinned: bool,
|
||||
pub author_id: Option<String>,
|
||||
pub author_username: Option<String>,
|
||||
pub start_time_unix_secs: Option<u64>,
|
||||
pub end_time_unix_secs: Option<u64>,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub updated_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
impl StoredAnnouncement {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
title: String,
|
||||
content: String,
|
||||
kind: String,
|
||||
priority: i32,
|
||||
is_active: bool,
|
||||
is_pinned: bool,
|
||||
author_id: Option<String>,
|
||||
author_username: Option<String>,
|
||||
start_time_unix_secs: Option<i64>,
|
||||
end_time_unix_secs: Option<i64>,
|
||||
created_at_unix_secs: i64,
|
||||
updated_at_unix_secs: i64,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"announcements.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if title.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"announcements.title is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if content.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"announcements.content is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if kind.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"announcements.type is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
title,
|
||||
content,
|
||||
kind,
|
||||
priority,
|
||||
is_active,
|
||||
is_pinned,
|
||||
author_id,
|
||||
author_username,
|
||||
start_time_unix_secs: start_time_unix_secs
|
||||
.map(|value| parse_timestamp(value, "announcements.start_time"))
|
||||
.transpose()?,
|
||||
end_time_unix_secs: end_time_unix_secs
|
||||
.map(|value| parse_timestamp(value, "announcements.end_time"))
|
||||
.transpose()?,
|
||||
created_at_unix_secs: parse_timestamp(
|
||||
created_at_unix_secs,
|
||||
"announcements.created_at",
|
||||
)?,
|
||||
updated_at_unix_secs: parse_timestamp(
|
||||
updated_at_unix_secs,
|
||||
"announcements.updated_at",
|
||||
)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct AnnouncementListQuery {
|
||||
pub active_only: bool,
|
||||
pub offset: usize,
|
||||
pub limit: usize,
|
||||
pub now_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredAnnouncementPage {
|
||||
pub items: Vec<StoredAnnouncement>,
|
||||
pub total: u64,
|
||||
}
|
||||
|
||||
fn parse_timestamp(value: i64, field: &str) -> Result<u64, crate::DataLayerError> {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("{field} is negative: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AnnouncementReadRepository: Send + Sync {
|
||||
async fn find_by_id(
|
||||
&self,
|
||||
announcement_id: &str,
|
||||
) -> Result<Option<StoredAnnouncement>, crate::DataLayerError>;
|
||||
|
||||
async fn list_announcements(
|
||||
&self,
|
||||
query: &AnnouncementListQuery,
|
||||
) -> Result<StoredAnnouncementPage, crate::DataLayerError>;
|
||||
|
||||
async fn count_unread_active_announcements(
|
||||
&self,
|
||||
user_id: &str,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<u64, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct CreateAnnouncementRecord {
|
||||
pub title: String,
|
||||
pub content: String,
|
||||
pub kind: String,
|
||||
pub priority: i32,
|
||||
pub is_pinned: bool,
|
||||
pub author_id: String,
|
||||
pub start_time_unix_secs: Option<u64>,
|
||||
pub end_time_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl CreateAnnouncementRecord {
|
||||
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||
if self.title.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"announcement title cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.content.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"announcement content cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.kind.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"announcement type cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.author_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"announcement author_id cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct UpdateAnnouncementRecord {
|
||||
pub announcement_id: String,
|
||||
pub title: Option<String>,
|
||||
pub content: Option<String>,
|
||||
pub kind: Option<String>,
|
||||
pub priority: Option<i32>,
|
||||
pub is_active: Option<bool>,
|
||||
pub is_pinned: Option<bool>,
|
||||
pub start_time_unix_secs: Option<u64>,
|
||||
pub end_time_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl UpdateAnnouncementRecord {
|
||||
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||
if self.announcement_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"announcement_id cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self
|
||||
.title
|
||||
.as_deref()
|
||||
.is_some_and(|value| value.trim().is_empty())
|
||||
{
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"announcement title cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self
|
||||
.content
|
||||
.as_deref()
|
||||
.is_some_and(|value| value.trim().is_empty())
|
||||
{
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"announcement content cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self
|
||||
.kind
|
||||
.as_deref()
|
||||
.is_some_and(|value| value.trim().is_empty())
|
||||
{
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"announcement type cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AnnouncementWriteRepository: Send + Sync {
|
||||
async fn create_announcement(
|
||||
&self,
|
||||
record: CreateAnnouncementRecord,
|
||||
) -> Result<StoredAnnouncement, crate::DataLayerError>;
|
||||
|
||||
async fn update_announcement(
|
||||
&self,
|
||||
record: UpdateAnnouncementRecord,
|
||||
) -> Result<Option<StoredAnnouncement>, crate::DataLayerError>;
|
||||
|
||||
async fn delete_announcement(
|
||||
&self,
|
||||
announcement_id: &str,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn mark_announcement_as_read(
|
||||
&self,
|
||||
user_id: &str,
|
||||
announcement_id: &str,
|
||||
read_at_unix_secs: u64,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
}
|
||||
@@ -3,13 +3,20 @@ use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot};
|
||||
use super::types::{
|
||||
AuthApiKeyExportSummary, AuthApiKeyLookupKey, AuthApiKeyReadRepository,
|
||||
AuthApiKeyWriteRepository, CreateStandaloneApiKeyRecord, CreateUserApiKeyRecord,
|
||||
StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
|
||||
UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct MemoryAuthApiKeyIndex {
|
||||
by_api_key_id: BTreeMap<String, StoredAuthApiKeySnapshot>,
|
||||
export_by_api_key_id: BTreeMap<String, StoredAuthApiKeyExportRecord>,
|
||||
by_key_hash: BTreeMap<String, String>,
|
||||
touch_counts: BTreeMap<String, usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
@@ -23,8 +30,46 @@ impl InMemoryAuthApiKeySnapshotRepository {
|
||||
I: IntoIterator<Item = (Option<String>, StoredAuthApiKeySnapshot)>,
|
||||
{
|
||||
let mut by_api_key_id = BTreeMap::new();
|
||||
let mut export_by_api_key_id = BTreeMap::new();
|
||||
let mut by_key_hash = BTreeMap::new();
|
||||
for (key_hash, snapshot) in items {
|
||||
let derived_key_hash = key_hash
|
||||
.clone()
|
||||
.unwrap_or_else(|| format!("memory-{}", snapshot.api_key_id));
|
||||
export_by_api_key_id.insert(
|
||||
snapshot.api_key_id.clone(),
|
||||
StoredAuthApiKeyExportRecord::new(
|
||||
snapshot.user_id.clone(),
|
||||
snapshot.api_key_id.clone(),
|
||||
derived_key_hash.clone(),
|
||||
None,
|
||||
snapshot.api_key_name.clone(),
|
||||
snapshot
|
||||
.api_key_allowed_providers
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
snapshot
|
||||
.api_key_allowed_api_formats
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
snapshot
|
||||
.api_key_allowed_models
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
snapshot.api_key_rate_limit,
|
||||
snapshot.api_key_concurrent_limit,
|
||||
None,
|
||||
snapshot.api_key_is_active,
|
||||
snapshot
|
||||
.api_key_expires_at_unix_secs
|
||||
.map(|value| value as i64),
|
||||
false,
|
||||
0,
|
||||
0.0,
|
||||
snapshot.api_key_is_standalone,
|
||||
)
|
||||
.expect("derived auth api key export record should build"),
|
||||
);
|
||||
if let Some(key_hash) = key_hash {
|
||||
by_key_hash.insert(key_hash, snapshot.api_key_id.clone());
|
||||
}
|
||||
@@ -33,10 +78,38 @@ impl InMemoryAuthApiKeySnapshotRepository {
|
||||
Self {
|
||||
index: RwLock::new(MemoryAuthApiKeyIndex {
|
||||
by_api_key_id,
|
||||
export_by_api_key_id,
|
||||
by_key_hash,
|
||||
touch_counts: BTreeMap::new(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_export_records<I>(mut self, items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredAuthApiKeyExportRecord>,
|
||||
{
|
||||
let index = self
|
||||
.index
|
||||
.get_mut()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
for item in items {
|
||||
index
|
||||
.export_by_api_key_id
|
||||
.insert(item.api_key_id.clone(), item);
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
pub fn touch_count(&self, api_key_id: &str) -> usize {
|
||||
self.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock")
|
||||
.touch_counts
|
||||
.get(api_key_id)
|
||||
.copied()
|
||||
.unwrap_or(0)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -68,13 +141,698 @@ impl AuthApiKeyReadRepository for InMemoryAuthApiKeySnapshotRepository {
|
||||
.cloned(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_api_key_snapshots_by_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeySnapshot>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(api_key_ids
|
||||
.iter()
|
||||
.filter_map(|api_key_id| index.by_api_key_id.get(api_key_id).cloned())
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_export_api_keys_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(index
|
||||
.export_by_api_key_id
|
||||
.values()
|
||||
.filter(|record| {
|
||||
!record.is_standalone && user_ids.iter().any(|id| id == &record.user_id)
|
||||
})
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_export_api_keys_by_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(api_key_ids
|
||||
.iter()
|
||||
.filter_map(|api_key_id| index.export_by_api_key_id.get(api_key_id).cloned())
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_export_standalone_api_keys_page(
|
||||
&self,
|
||||
query: &StandaloneApiKeyExportListQuery,
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(index
|
||||
.export_by_api_key_id
|
||||
.values()
|
||||
.filter(|record| {
|
||||
record.is_standalone
|
||||
&& query
|
||||
.is_active
|
||||
.is_none_or(|is_active| record.is_active == is_active)
|
||||
})
|
||||
.skip(query.skip)
|
||||
.take(query.limit)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn count_export_standalone_api_keys(
|
||||
&self,
|
||||
is_active: Option<bool>,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(index
|
||||
.export_by_api_key_id
|
||||
.values()
|
||||
.filter(|record| {
|
||||
record.is_standalone
|
||||
&& is_active.is_none_or(|expected| record.is_active == expected)
|
||||
})
|
||||
.count() as u64)
|
||||
}
|
||||
|
||||
async fn summarize_export_api_keys_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let mut summary = AuthApiKeyExportSummary::default();
|
||||
for record in index.export_by_api_key_id.values().filter(|record| {
|
||||
!record.is_standalone && user_ids.iter().any(|id| id == &record.user_id)
|
||||
}) {
|
||||
summary.total = summary.total.saturating_add(1);
|
||||
if record.is_active
|
||||
&& record
|
||||
.expires_at_unix_secs
|
||||
.is_none_or(|expires_at_unix_secs| expires_at_unix_secs >= now_unix_secs)
|
||||
{
|
||||
summary.active = summary.active.saturating_add(1);
|
||||
}
|
||||
}
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
async fn summarize_export_non_standalone_api_keys(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let mut summary = AuthApiKeyExportSummary::default();
|
||||
for record in index
|
||||
.export_by_api_key_id
|
||||
.values()
|
||||
.filter(|record| !record.is_standalone)
|
||||
{
|
||||
summary.total = summary.total.saturating_add(1);
|
||||
if record.is_active
|
||||
&& record
|
||||
.expires_at_unix_secs
|
||||
.is_none_or(|expires_at_unix_secs| expires_at_unix_secs >= now_unix_secs)
|
||||
{
|
||||
summary.active = summary.active.saturating_add(1);
|
||||
}
|
||||
}
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
async fn find_export_standalone_api_key_by_id(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(index
|
||||
.export_by_api_key_id
|
||||
.get(api_key_id)
|
||||
.filter(|record| record.is_standalone)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn summarize_export_standalone_api_keys(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let mut summary = AuthApiKeyExportSummary::default();
|
||||
for record in index
|
||||
.export_by_api_key_id
|
||||
.values()
|
||||
.filter(|record| record.is_standalone)
|
||||
{
|
||||
summary.total = summary.total.saturating_add(1);
|
||||
if record.is_active
|
||||
&& record
|
||||
.expires_at_unix_secs
|
||||
.is_none_or(|expires_at_unix_secs| expires_at_unix_secs >= now_unix_secs)
|
||||
{
|
||||
summary.active = summary.active.saturating_add(1);
|
||||
}
|
||||
}
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
async fn list_export_standalone_api_keys(
|
||||
&self,
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(index
|
||||
.export_by_api_key_id
|
||||
.values()
|
||||
.filter(|record| record.is_standalone)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
|
||||
async fn touch_last_used_at(&self, api_key_id: &str) -> Result<bool, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
if !index.by_api_key_id.contains_key(api_key_id) {
|
||||
return Ok(false);
|
||||
}
|
||||
let counter = index
|
||||
.touch_counts
|
||||
.entry(api_key_id.to_string())
|
||||
.or_insert(0);
|
||||
*counter += 1;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn create_user_api_key(
|
||||
&self,
|
||||
record: CreateUserApiKeyRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
if index.by_api_key_id.contains_key(&record.api_key_id) {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"duplicate api_keys.id: {}",
|
||||
record.api_key_id
|
||||
)));
|
||||
}
|
||||
if index.by_key_hash.contains_key(&record.key_hash) {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"duplicate api_keys.key_hash: {}",
|
||||
record.key_hash
|
||||
)));
|
||||
}
|
||||
|
||||
let template = index
|
||||
.by_api_key_id
|
||||
.values()
|
||||
.find(|snapshot| snapshot.user_id == record.user_id)
|
||||
.cloned();
|
||||
let snapshot = if let Some(template) = template {
|
||||
StoredAuthApiKeySnapshot {
|
||||
api_key_id: record.api_key_id.clone(),
|
||||
api_key_name: record.name.clone(),
|
||||
api_key_is_active: true,
|
||||
api_key_is_locked: false,
|
||||
api_key_is_standalone: false,
|
||||
api_key_rate_limit: Some(record.rate_limit),
|
||||
api_key_concurrent_limit: Some(record.concurrent_limit),
|
||||
api_key_expires_at_unix_secs: None,
|
||||
api_key_allowed_providers: None,
|
||||
api_key_allowed_api_formats: None,
|
||||
api_key_allowed_models: None,
|
||||
..template
|
||||
}
|
||||
} else {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
record.user_id.clone(),
|
||||
format!(
|
||||
"user-{}",
|
||||
&record.user_id.chars().take(8).collect::<String>()
|
||||
),
|
||||
None,
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
record.api_key_id.clone(),
|
||||
record.name.clone(),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(record.rate_limit),
|
||||
Some(record.concurrent_limit),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)?
|
||||
};
|
||||
|
||||
let export = StoredAuthApiKeyExportRecord::new(
|
||||
record.user_id.clone(),
|
||||
record.api_key_id.clone(),
|
||||
record.key_hash.clone(),
|
||||
record.key_encrypted,
|
||||
record.name,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(record.rate_limit),
|
||||
Some(record.concurrent_limit),
|
||||
None,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
0,
|
||||
0.0,
|
||||
false,
|
||||
)?;
|
||||
|
||||
index
|
||||
.by_key_hash
|
||||
.insert(record.key_hash, record.api_key_id.clone());
|
||||
index
|
||||
.by_api_key_id
|
||||
.insert(record.api_key_id.clone(), snapshot);
|
||||
index
|
||||
.export_by_api_key_id
|
||||
.insert(record.api_key_id, export.clone());
|
||||
Ok(Some(export))
|
||||
}
|
||||
|
||||
async fn create_standalone_api_key(
|
||||
&self,
|
||||
record: CreateStandaloneApiKeyRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
if index.by_api_key_id.contains_key(&record.api_key_id) {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"duplicate api_keys.id: {}",
|
||||
record.api_key_id
|
||||
)));
|
||||
}
|
||||
if index.by_key_hash.contains_key(&record.key_hash) {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"duplicate api_keys.key_hash: {}",
|
||||
record.key_hash
|
||||
)));
|
||||
}
|
||||
|
||||
let template = index
|
||||
.by_api_key_id
|
||||
.values()
|
||||
.find(|snapshot| snapshot.user_id == record.user_id)
|
||||
.cloned();
|
||||
let snapshot = if let Some(template) = template {
|
||||
StoredAuthApiKeySnapshot {
|
||||
api_key_id: record.api_key_id.clone(),
|
||||
api_key_name: record.name.clone(),
|
||||
api_key_is_active: true,
|
||||
api_key_is_locked: false,
|
||||
api_key_is_standalone: true,
|
||||
api_key_rate_limit: Some(record.rate_limit),
|
||||
api_key_concurrent_limit: Some(record.concurrent_limit),
|
||||
api_key_expires_at_unix_secs: None,
|
||||
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(),
|
||||
..template
|
||||
}
|
||||
} else {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
record.user_id.clone(),
|
||||
format!(
|
||||
"admin-{}",
|
||||
&record.user_id.chars().take(8).collect::<String>()
|
||||
),
|
||||
None,
|
||||
"admin".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
record.api_key_id.clone(),
|
||||
record.name.clone(),
|
||||
true,
|
||||
false,
|
||||
true,
|
||||
Some(record.rate_limit),
|
||||
Some(record.concurrent_limit),
|
||||
None,
|
||||
record
|
||||
.allowed_providers
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
record
|
||||
.allowed_api_formats
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
record
|
||||
.allowed_models
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
)?
|
||||
};
|
||||
|
||||
let export = StoredAuthApiKeyExportRecord::new(
|
||||
record.user_id.clone(),
|
||||
record.api_key_id.clone(),
|
||||
record.key_hash.clone(),
|
||||
record.key_encrypted,
|
||||
record.name,
|
||||
record
|
||||
.allowed_providers
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
record
|
||||
.allowed_api_formats
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
record
|
||||
.allowed_models
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
Some(record.rate_limit),
|
||||
Some(record.concurrent_limit),
|
||||
None,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
0,
|
||||
0.0,
|
||||
true,
|
||||
)?;
|
||||
|
||||
index
|
||||
.by_key_hash
|
||||
.insert(record.key_hash, record.api_key_id.clone());
|
||||
index
|
||||
.by_api_key_id
|
||||
.insert(record.api_key_id.clone(), snapshot);
|
||||
index
|
||||
.export_by_api_key_id
|
||||
.insert(record.api_key_id, export.clone());
|
||||
Ok(Some(export))
|
||||
}
|
||||
|
||||
async fn update_user_api_key_basic(
|
||||
&self,
|
||||
record: UpdateUserApiKeyBasicRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(&record.api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if snapshot.user_id != record.user_id || snapshot.api_key_is_standalone {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(name) = record.name {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_name = Some(name.clone());
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.name = Some(name);
|
||||
}
|
||||
}
|
||||
if let Some(rate_limit) = record.rate_limit {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_rate_limit = Some(rate_limit);
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.rate_limit = Some(rate_limit);
|
||||
}
|
||||
}
|
||||
Ok(index.export_by_api_key_id.get(&record.api_key_id).cloned())
|
||||
}
|
||||
|
||||
async fn update_standalone_api_key_basic(
|
||||
&self,
|
||||
record: UpdateStandaloneApiKeyBasicRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(&record.api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !snapshot.api_key_is_standalone {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(name) = record.name {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_name = Some(name.clone());
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.name = Some(name);
|
||||
}
|
||||
}
|
||||
if let Some(rate_limit) = record.rate_limit {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_rate_limit = Some(rate_limit);
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.rate_limit = Some(rate_limit);
|
||||
}
|
||||
}
|
||||
if let Some(allowed_providers) = record.allowed_providers {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_allowed_providers = allowed_providers.clone();
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.allowed_providers = allowed_providers;
|
||||
}
|
||||
}
|
||||
if let Some(allowed_api_formats) = record.allowed_api_formats {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_allowed_api_formats = allowed_api_formats.clone();
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.allowed_api_formats = allowed_api_formats;
|
||||
}
|
||||
}
|
||||
if let Some(allowed_models) = record.allowed_models {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_allowed_models = allowed_models.clone();
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.allowed_models = allowed_models;
|
||||
}
|
||||
}
|
||||
Ok(index.export_by_api_key_id.get(&record.api_key_id).cloned())
|
||||
}
|
||||
|
||||
async fn set_user_api_key_active(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
is_active: bool,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if snapshot.user_id != user_id || snapshot.api_key_is_standalone {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(api_key_id) {
|
||||
snapshot.api_key_is_active = is_active;
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(api_key_id) {
|
||||
export.is_active = is_active;
|
||||
}
|
||||
Ok(index.export_by_api_key_id.get(api_key_id).cloned())
|
||||
}
|
||||
|
||||
async fn set_standalone_api_key_active(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
is_active: bool,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !snapshot.api_key_is_standalone {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(api_key_id) {
|
||||
snapshot.api_key_is_active = is_active;
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(api_key_id) {
|
||||
export.is_active = is_active;
|
||||
}
|
||||
Ok(index.export_by_api_key_id.get(api_key_id).cloned())
|
||||
}
|
||||
|
||||
async fn set_user_api_key_locked(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
is_locked: bool,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
if snapshot.user_id != user_id || snapshot.api_key_is_standalone {
|
||||
return Ok(false);
|
||||
}
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(api_key_id) {
|
||||
snapshot.api_key_is_locked = is_locked;
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn set_user_api_key_allowed_providers(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
allowed_providers: Option<Vec<String>>,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if snapshot.user_id != user_id || snapshot.api_key_is_standalone {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(api_key_id) {
|
||||
snapshot.api_key_allowed_providers = allowed_providers.clone();
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(api_key_id) {
|
||||
export.allowed_providers = allowed_providers;
|
||||
}
|
||||
Ok(index.export_by_api_key_id.get(api_key_id).cloned())
|
||||
}
|
||||
|
||||
async fn set_user_api_key_force_capabilities(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
force_capabilities: Option<serde_json::Value>,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if snapshot.user_id != user_id || snapshot.api_key_is_standalone {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(export) = index.export_by_api_key_id.get_mut(api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
export.force_capabilities = force_capabilities;
|
||||
Ok(Some(export.clone()))
|
||||
}
|
||||
|
||||
async fn delete_user_api_key(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
if snapshot.user_id != user_id || snapshot.api_key_is_standalone {
|
||||
return Ok(false);
|
||||
}
|
||||
index.by_api_key_id.remove(api_key_id);
|
||||
index.export_by_api_key_id.remove(api_key_id);
|
||||
index.by_key_hash.retain(|_, value| value != api_key_id);
|
||||
index.touch_counts.remove(api_key_id);
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn delete_standalone_api_key(&self, api_key_id: &str) -> Result<bool, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
if !snapshot.api_key_is_standalone {
|
||||
return Ok(false);
|
||||
}
|
||||
index.by_api_key_id.remove(api_key_id);
|
||||
index.export_by_api_key_id.remove(api_key_id);
|
||||
index.by_key_hash.retain(|_, value| value != api_key_id);
|
||||
index.touch_counts.remove(api_key_id);
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryAuthApiKeySnapshotRepository;
|
||||
use crate::repository::auth::{
|
||||
AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot,
|
||||
AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthApiKeyWriteRepository,
|
||||
StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
|
||||
fn sample_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
@@ -129,5 +887,132 @@ mod tests {
|
||||
.await
|
||||
.expect("find by user/api key ids should succeed")
|
||||
.is_some());
|
||||
let snapshots = repository
|
||||
.list_api_key_snapshots_by_ids(&["key-1".to_string(), "missing".to_string()])
|
||||
.await
|
||||
.expect("batch lookup should succeed");
|
||||
assert_eq!(snapshots.len(), 1);
|
||||
assert_eq!(snapshots[0].api_key_id, "key-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn touches_last_used_for_existing_key() {
|
||||
let repository = InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_snapshot("key-1", "user-1"),
|
||||
)]);
|
||||
|
||||
assert!(repository
|
||||
.touch_last_used_at("key-1")
|
||||
.await
|
||||
.expect("touch should succeed"));
|
||||
assert_eq!(repository.touch_count("key-1"), 1);
|
||||
assert!(!repository
|
||||
.touch_last_used_at("missing")
|
||||
.await
|
||||
.expect("missing touch should succeed"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lists_export_records_for_user_bound_and_standalone_keys() {
|
||||
let repository = InMemoryAuthApiKeySnapshotRepository::seed(vec![
|
||||
(
|
||||
Some("hash-user".to_string()),
|
||||
sample_snapshot("key-user", "user-1"),
|
||||
),
|
||||
(
|
||||
Some("hash-standalone".to_string()),
|
||||
sample_snapshot("key-standalone", "admin-1"),
|
||||
),
|
||||
])
|
||||
.with_export_records(vec![
|
||||
StoredAuthApiKeyExportRecord::new(
|
||||
"user-1".to_string(),
|
||||
"key-user".to_string(),
|
||||
"hash-user".to_string(),
|
||||
Some("enc-user".to_string()),
|
||||
Some("default".to_string()),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
Some(120),
|
||||
Some(7),
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
true,
|
||||
Some(200),
|
||||
false,
|
||||
14,
|
||||
1.5,
|
||||
false,
|
||||
)
|
||||
.expect("user export record should build"),
|
||||
StoredAuthApiKeyExportRecord::new(
|
||||
"admin-1".to_string(),
|
||||
"key-standalone".to_string(),
|
||||
"hash-standalone".to_string(),
|
||||
Some("enc-standalone".to_string()),
|
||||
Some("standalone".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(1),
|
||||
None,
|
||||
true,
|
||||
None,
|
||||
true,
|
||||
2,
|
||||
0.25,
|
||||
true,
|
||||
)
|
||||
.expect("standalone export record should build"),
|
||||
]);
|
||||
|
||||
let user_records = repository
|
||||
.list_export_api_keys_by_user_ids(&["user-1".to_string()])
|
||||
.await
|
||||
.expect("user export lookup should succeed");
|
||||
assert_eq!(user_records.len(), 1);
|
||||
assert_eq!(user_records[0].api_key_id, "key-user");
|
||||
assert_eq!(user_records[0].key_encrypted.as_deref(), Some("enc-user"));
|
||||
assert_eq!(user_records[0].total_requests, 14);
|
||||
|
||||
let standalone_records = repository
|
||||
.list_export_standalone_api_keys()
|
||||
.await
|
||||
.expect("standalone export lookup should succeed");
|
||||
assert_eq!(standalone_records.len(), 1);
|
||||
assert_eq!(standalone_records[0].api_key_id, "key-standalone");
|
||||
assert!(standalone_records[0].is_standalone);
|
||||
|
||||
let selected_records = repository
|
||||
.list_export_api_keys_by_ids(&[
|
||||
"key-standalone".to_string(),
|
||||
"missing".to_string(),
|
||||
"key-user".to_string(),
|
||||
])
|
||||
.await
|
||||
.expect("api key id export lookup should succeed");
|
||||
assert_eq!(selected_records.len(), 2);
|
||||
assert_eq!(selected_records[0].api_key_id, "key-standalone");
|
||||
assert_eq!(selected_records[1].api_key_id, "key-user");
|
||||
|
||||
let paged_records = repository
|
||||
.list_export_standalone_api_keys_page(&StandaloneApiKeyExportListQuery {
|
||||
skip: 0,
|
||||
limit: 10,
|
||||
is_active: Some(true),
|
||||
})
|
||||
.await
|
||||
.expect("standalone export page should succeed");
|
||||
assert_eq!(paged_records.len(), 1);
|
||||
assert_eq!(paged_records[0].api_key_id, "key-standalone");
|
||||
assert_eq!(
|
||||
repository
|
||||
.count_export_standalone_api_keys(Some(true))
|
||||
.await
|
||||
.expect("standalone export count should succeed"),
|
||||
1
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,5 +5,8 @@ mod types;
|
||||
pub use memory::InMemoryAuthApiKeySnapshotRepository;
|
||||
pub use sql::SqlxAuthApiKeySnapshotReadRepository;
|
||||
pub use types::{
|
||||
AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthRepository, StoredAuthApiKeySnapshot,
|
||||
AuthApiKeyExportSummary, AuthApiKeyLookupKey, AuthApiKeyReadRepository,
|
||||
AuthApiKeyWriteRepository, AuthRepository, CreateStandaloneApiKeyRecord,
|
||||
CreateUserApiKeyRecord, StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord,
|
||||
StoredAuthApiKeySnapshot, UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord,
|
||||
};
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -9,6 +9,7 @@ pub struct StoredAuthApiKeySnapshot {
|
||||
pub user_auth_source: String,
|
||||
pub user_is_active: bool,
|
||||
pub user_is_deleted: bool,
|
||||
pub user_rate_limit: Option<i32>,
|
||||
pub user_allowed_providers: Option<Vec<String>>,
|
||||
pub user_allowed_api_formats: Option<Vec<String>>,
|
||||
pub user_allowed_models: Option<Vec<String>>,
|
||||
@@ -58,6 +59,7 @@ impl StoredAuthApiKeySnapshot {
|
||||
user_auth_source,
|
||||
user_is_active,
|
||||
user_is_deleted,
|
||||
user_rate_limit: None,
|
||||
user_allowed_providers: parse_string_list(
|
||||
user_allowed_providers,
|
||||
"users.allowed_providers",
|
||||
@@ -115,6 +117,157 @@ impl StoredAuthApiKeySnapshot {
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
pub fn with_user_rate_limit(mut self, user_rate_limit: Option<i32>) -> Self {
|
||||
self.user_rate_limit = user_rate_limit;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredAuthApiKeyExportRecord {
|
||||
pub user_id: String,
|
||||
pub api_key_id: String,
|
||||
pub key_hash: String,
|
||||
pub key_encrypted: Option<String>,
|
||||
pub name: Option<String>,
|
||||
pub allowed_providers: Option<Vec<String>>,
|
||||
pub allowed_api_formats: Option<Vec<String>>,
|
||||
pub allowed_models: Option<Vec<String>>,
|
||||
pub rate_limit: Option<i32>,
|
||||
pub concurrent_limit: Option<i32>,
|
||||
pub force_capabilities: Option<serde_json::Value>,
|
||||
pub is_active: bool,
|
||||
pub expires_at_unix_secs: Option<u64>,
|
||||
pub auto_delete_on_expiry: bool,
|
||||
pub total_requests: u64,
|
||||
pub total_cost_usd: f64,
|
||||
pub is_standalone: bool,
|
||||
}
|
||||
|
||||
impl StoredAuthApiKeyExportRecord {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
user_id: String,
|
||||
api_key_id: String,
|
||||
key_hash: String,
|
||||
key_encrypted: Option<String>,
|
||||
name: Option<String>,
|
||||
allowed_providers: Option<serde_json::Value>,
|
||||
allowed_api_formats: Option<serde_json::Value>,
|
||||
allowed_models: Option<serde_json::Value>,
|
||||
rate_limit: Option<i32>,
|
||||
concurrent_limit: Option<i32>,
|
||||
force_capabilities: Option<serde_json::Value>,
|
||||
is_active: bool,
|
||||
expires_at_unix_secs: Option<i64>,
|
||||
auto_delete_on_expiry: bool,
|
||||
total_requests: i64,
|
||||
total_cost_usd: f64,
|
||||
is_standalone: bool,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if user_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"api_keys.user_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if api_key_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"api_keys.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if key_hash.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"api_keys.key_hash is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if !total_cost_usd.is_finite() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"api_keys.total_cost_usd is not finite".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
user_id,
|
||||
api_key_id,
|
||||
key_hash,
|
||||
key_encrypted,
|
||||
name,
|
||||
allowed_providers: parse_string_list(allowed_providers, "api_keys.allowed_providers")?,
|
||||
allowed_api_formats: parse_string_list(
|
||||
allowed_api_formats,
|
||||
"api_keys.allowed_api_formats",
|
||||
)?,
|
||||
allowed_models: parse_string_list(allowed_models, "api_keys.allowed_models")?,
|
||||
rate_limit,
|
||||
concurrent_limit,
|
||||
force_capabilities,
|
||||
is_active,
|
||||
expires_at_unix_secs: expires_at_unix_secs
|
||||
.map(|value| parse_u64_i64(value, "api_keys.expires_at_unix_secs"))
|
||||
.transpose()?,
|
||||
auto_delete_on_expiry,
|
||||
total_requests: parse_u64_i64(total_requests, "api_keys.total_requests")?,
|
||||
total_cost_usd,
|
||||
is_standalone,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct AuthApiKeyExportSummary {
|
||||
pub total: u64,
|
||||
pub active: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct StandaloneApiKeyExportListQuery {
|
||||
pub skip: usize,
|
||||
pub limit: usize,
|
||||
pub is_active: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CreateUserApiKeyRecord {
|
||||
pub user_id: String,
|
||||
pub api_key_id: String,
|
||||
pub key_hash: String,
|
||||
pub key_encrypted: Option<String>,
|
||||
pub name: Option<String>,
|
||||
pub rate_limit: i32,
|
||||
pub concurrent_limit: i32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpdateUserApiKeyBasicRecord {
|
||||
pub user_id: String,
|
||||
pub api_key_id: String,
|
||||
pub name: Option<String>,
|
||||
pub rate_limit: Option<i32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CreateStandaloneApiKeyRecord {
|
||||
pub user_id: String,
|
||||
pub api_key_id: String,
|
||||
pub key_hash: String,
|
||||
pub key_encrypted: Option<String>,
|
||||
pub name: Option<String>,
|
||||
pub allowed_providers: Option<Vec<String>>,
|
||||
pub allowed_api_formats: Option<Vec<String>>,
|
||||
pub allowed_models: Option<Vec<String>>,
|
||||
pub rate_limit: i32,
|
||||
pub concurrent_limit: i32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpdateStandaloneApiKeyBasicRecord {
|
||||
pub api_key_id: String,
|
||||
pub name: Option<String>,
|
||||
pub rate_limit: Option<i32>,
|
||||
pub allowed_providers: Option<Option<Vec<String>>>,
|
||||
pub allowed_api_formats: Option<Option<Vec<String>>>,
|
||||
pub allowed_models: Option<Option<Vec<String>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -133,11 +286,137 @@ pub trait AuthApiKeyReadRepository: Send + Sync {
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, crate::DataLayerError>;
|
||||
|
||||
async fn list_api_key_snapshots_by_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeySnapshot>, crate::DataLayerError>;
|
||||
|
||||
async fn list_export_api_keys_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn list_export_api_keys_by_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn list_export_standalone_api_keys_page(
|
||||
&self,
|
||||
query: &StandaloneApiKeyExportListQuery,
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn count_export_standalone_api_keys(
|
||||
&self,
|
||||
is_active: Option<bool>,
|
||||
) -> Result<u64, crate::DataLayerError>;
|
||||
|
||||
async fn summarize_export_api_keys_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, crate::DataLayerError>;
|
||||
|
||||
async fn summarize_export_non_standalone_api_keys(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, crate::DataLayerError>;
|
||||
|
||||
async fn summarize_export_standalone_api_keys(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, crate::DataLayerError>;
|
||||
|
||||
async fn find_export_standalone_api_key_by_id(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn list_export_standalone_api_keys(
|
||||
&self,
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait AuthRepository: AuthApiKeyReadRepository + Send + Sync {}
|
||||
#[async_trait]
|
||||
pub trait AuthApiKeyWriteRepository: Send + Sync {
|
||||
async fn touch_last_used_at(&self, api_key_id: &str) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
impl<T> AuthRepository for T where T: AuthApiKeyReadRepository + Send + Sync {}
|
||||
async fn create_user_api_key(
|
||||
&self,
|
||||
record: CreateUserApiKeyRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn create_standalone_api_key(
|
||||
&self,
|
||||
record: CreateStandaloneApiKeyRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn update_user_api_key_basic(
|
||||
&self,
|
||||
record: UpdateUserApiKeyBasicRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn update_standalone_api_key_basic(
|
||||
&self,
|
||||
record: UpdateStandaloneApiKeyBasicRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn set_user_api_key_active(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
is_active: bool,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn set_standalone_api_key_active(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
is_active: bool,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn set_user_api_key_locked(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
is_locked: bool,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn set_user_api_key_allowed_providers(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
allowed_providers: Option<Vec<String>>,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn set_user_api_key_force_capabilities(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
force_capabilities: Option<serde_json::Value>,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn delete_user_api_key(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn delete_standalone_api_key(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait AuthRepository:
|
||||
AuthApiKeyReadRepository + AuthApiKeyWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> AuthRepository for T where
|
||||
T: AuthApiKeyReadRepository + AuthApiKeyWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
fn parse_string_list(
|
||||
value: Option<serde_json::Value>,
|
||||
@@ -146,9 +425,43 @@ fn parse_string_list(
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
let array = value.as_array().ok_or_else(|| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("{field_name} is not a JSON array"))
|
||||
})?;
|
||||
parse_string_list_value(&value, field_name)
|
||||
}
|
||||
|
||||
fn parse_string_list_value(
|
||||
value: &serde_json::Value,
|
||||
field_name: &str,
|
||||
) -> Result<Option<Vec<String>>, crate::DataLayerError> {
|
||||
match value {
|
||||
serde_json::Value::Null => Ok(None),
|
||||
serde_json::Value::Array(array) => parse_string_list_array(array, field_name).map(Some),
|
||||
serde_json::Value::String(raw) => parse_embedded_string_list(raw, field_name),
|
||||
_ => Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"{field_name} is not a JSON array"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_embedded_string_list(
|
||||
raw: &str,
|
||||
field_name: &str,
|
||||
) -> Result<Option<Vec<String>>, crate::DataLayerError> {
|
||||
let raw = raw.trim();
|
||||
if raw.is_empty() || raw.eq_ignore_ascii_case("null") {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if let Ok(decoded) = serde_json::from_str::<serde_json::Value>(raw) {
|
||||
return parse_string_list_value(&decoded, field_name);
|
||||
}
|
||||
|
||||
Ok(Some(vec![raw.to_string()]))
|
||||
}
|
||||
|
||||
fn parse_string_list_array(
|
||||
array: &[serde_json::Value],
|
||||
field_name: &str,
|
||||
) -> Result<Vec<String>, crate::DataLayerError> {
|
||||
let mut items = Vec::with_capacity(array.len());
|
||||
for item in array {
|
||||
let Some(item) = item.as_str() else {
|
||||
@@ -156,14 +469,23 @@ fn parse_string_list(
|
||||
"{field_name} contains a non-string item"
|
||||
)));
|
||||
};
|
||||
items.push(item.to_string());
|
||||
let item = item.trim();
|
||||
if !item.is_empty() {
|
||||
items.push(item.to_string());
|
||||
}
|
||||
}
|
||||
Ok(Some(items))
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
fn parse_u64_i64(value: i64, field_name: &str) -> Result<u64, crate::DataLayerError> {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::StoredAuthApiKeySnapshot;
|
||||
use super::{StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot};
|
||||
|
||||
#[test]
|
||||
fn rejects_non_array_allowed_providers() {
|
||||
@@ -193,6 +515,72 @@ mod tests {
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_stringified_allowed_provider_array() {
|
||||
let snapshot = StoredAuthApiKeySnapshot::new(
|
||||
"user-1".to_string(),
|
||||
"alice".to_string(),
|
||||
None,
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!("[\"openai\", \" gemini \"]")),
|
||||
None,
|
||||
None,
|
||||
"key-1".to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("snapshot should build");
|
||||
|
||||
assert_eq!(
|
||||
snapshot.user_allowed_providers,
|
||||
Some(vec!["openai".to_string(), "gemini".to_string()])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_single_string_allowed_provider() {
|
||||
let snapshot = StoredAuthApiKeySnapshot::new(
|
||||
"user-1".to_string(),
|
||||
"alice".to_string(),
|
||||
None,
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!("openai")),
|
||||
None,
|
||||
None,
|
||||
"key-1".to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("snapshot should build");
|
||||
|
||||
assert_eq!(
|
||||
snapshot.user_allowed_providers,
|
||||
Some(vec!["openai".to_string()])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expired_non_standalone_key_is_not_usable() {
|
||||
let snapshot = StoredAuthApiKeySnapshot::new(
|
||||
@@ -222,4 +610,59 @@ mod tests {
|
||||
|
||||
assert!(!snapshot.is_currently_usable(101));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn export_record_rejects_negative_totals() {
|
||||
assert!(StoredAuthApiKeyExportRecord::new(
|
||||
"user-1".to_string(),
|
||||
"key-1".to_string(),
|
||||
"hash-1".to_string(),
|
||||
Some("enc".to_string()),
|
||||
Some("default".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(60),
|
||||
Some(5),
|
||||
None,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
-1,
|
||||
0.0,
|
||||
false,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn export_record_accepts_stringified_allowed_models() {
|
||||
let record = StoredAuthApiKeyExportRecord::new(
|
||||
"user-1".to_string(),
|
||||
"key-1".to_string(),
|
||||
"hash-1".to_string(),
|
||||
Some("enc".to_string()),
|
||||
Some("default".to_string()),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!("[\"gpt-5\", \" gpt-4.1 \"]")),
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
true,
|
||||
Some(200),
|
||||
false,
|
||||
12,
|
||||
1.25,
|
||||
false,
|
||||
)
|
||||
.expect("export record should build");
|
||||
|
||||
assert_eq!(
|
||||
record.allowed_models,
|
||||
Some(vec!["gpt-5".to_string(), "gpt-4.1".to_string()])
|
||||
);
|
||||
assert_eq!(record.total_requests, 12);
|
||||
assert_eq!(record.total_cost_usd, 1.25);
|
||||
}
|
||||
}
|
||||
|
||||
114
crates/aether-data/src/repository/auth_modules/memory.rs
Normal file
114
crates/aether-data/src/repository/auth_modules/memory.rs
Normal file
@@ -0,0 +1,114 @@
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
AuthModuleReadRepository, AuthModuleWriteRepository, StoredLdapModuleConfig,
|
||||
StoredOAuthProviderModuleConfig,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryAuthModuleReadRepository {
|
||||
oauth_providers: RwLock<Vec<StoredOAuthProviderModuleConfig>>,
|
||||
ldap_config: RwLock<Option<StoredLdapModuleConfig>>,
|
||||
}
|
||||
|
||||
impl InMemoryAuthModuleReadRepository {
|
||||
pub fn seed<I>(oauth_providers: I, ldap_config: Option<StoredLdapModuleConfig>) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredOAuthProviderModuleConfig>,
|
||||
{
|
||||
Self {
|
||||
oauth_providers: RwLock::new(oauth_providers.into_iter().collect()),
|
||||
ldap_config: RwLock::new(ldap_config),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AuthModuleReadRepository for InMemoryAuthModuleReadRepository {
|
||||
async fn list_enabled_oauth_providers(
|
||||
&self,
|
||||
) -> Result<Vec<StoredOAuthProviderModuleConfig>, DataLayerError> {
|
||||
Ok(self
|
||||
.oauth_providers
|
||||
.read()
|
||||
.expect("auth module oauth provider repository lock")
|
||||
.clone())
|
||||
}
|
||||
|
||||
async fn get_ldap_config(&self) -> Result<Option<StoredLdapModuleConfig>, DataLayerError> {
|
||||
Ok(self
|
||||
.ldap_config
|
||||
.read()
|
||||
.expect("auth module ldap repository lock")
|
||||
.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AuthModuleWriteRepository for InMemoryAuthModuleReadRepository {
|
||||
async fn upsert_ldap_config(
|
||||
&self,
|
||||
config: &StoredLdapModuleConfig,
|
||||
) -> Result<Option<StoredLdapModuleConfig>, DataLayerError> {
|
||||
self.ldap_config
|
||||
.write()
|
||||
.expect("auth module ldap repository lock")
|
||||
.replace(config.clone());
|
||||
Ok(Some(config.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryAuthModuleReadRepository;
|
||||
use crate::repository::auth_modules::{
|
||||
AuthModuleReadRepository, StoredLdapModuleConfig, StoredOAuthProviderModuleConfig,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_seeded_auth_module_configs() {
|
||||
let repository = InMemoryAuthModuleReadRepository::seed(
|
||||
vec![StoredOAuthProviderModuleConfig::new(
|
||||
"linuxdo".to_string(),
|
||||
"Linux DO".to_string(),
|
||||
"client-id".to_string(),
|
||||
Some("encrypted".to_string()),
|
||||
"https://example.com/callback".to_string(),
|
||||
)
|
||||
.expect("oauth provider should build")],
|
||||
Some(StoredLdapModuleConfig {
|
||||
server_url: "ldaps://ldap.example.com".to_string(),
|
||||
bind_dn: "cn=admin,dc=example,dc=com".to_string(),
|
||||
bind_password_encrypted: Some("encrypted-password".to_string()),
|
||||
base_dn: "dc=example,dc=com".to_string(),
|
||||
user_search_filter: Some("(uid={username})".to_string()),
|
||||
username_attr: Some("uid".to_string()),
|
||||
email_attr: Some("mail".to_string()),
|
||||
display_name_attr: Some("displayName".to_string()),
|
||||
is_enabled: true,
|
||||
is_exclusive: false,
|
||||
use_starttls: true,
|
||||
connect_timeout: Some(10),
|
||||
}),
|
||||
);
|
||||
|
||||
let oauth = repository
|
||||
.list_enabled_oauth_providers()
|
||||
.await
|
||||
.expect("oauth providers should load");
|
||||
let ldap = repository
|
||||
.get_ldap_config()
|
||||
.await
|
||||
.expect("ldap config should load");
|
||||
|
||||
assert_eq!(oauth.len(), 1);
|
||||
assert_eq!(oauth[0].provider_type, "linuxdo");
|
||||
assert_eq!(
|
||||
ldap.expect("ldap config should exist").server_url,
|
||||
"ldaps://ldap.example.com"
|
||||
);
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/auth_modules/mod.rs
Normal file
10
crates/aether-data/src/repository/auth_modules/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryAuthModuleReadRepository;
|
||||
pub use sql::{SqlxAuthModuleReadRepository, SqlxAuthModuleRepository};
|
||||
pub use types::{
|
||||
AuthModuleReadRepository, AuthModuleWriteRepository, StoredLdapModuleConfig,
|
||||
StoredOAuthProviderModuleConfig,
|
||||
};
|
||||
297
crates/aether-data/src/repository/auth_modules/sql.rs
Normal file
297
crates/aether-data/src/repository/auth_modules/sql.rs
Normal file
@@ -0,0 +1,297 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
AuthModuleReadRepository, AuthModuleWriteRepository, StoredLdapModuleConfig,
|
||||
StoredOAuthProviderModuleConfig,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const LIST_ENABLED_OAUTH_PROVIDERS_SQL: &str = r#"
|
||||
SELECT
|
||||
provider_type,
|
||||
display_name,
|
||||
client_id,
|
||||
client_secret_encrypted,
|
||||
redirect_uri
|
||||
FROM oauth_providers
|
||||
WHERE is_enabled = TRUE
|
||||
ORDER BY provider_type ASC
|
||||
"#;
|
||||
|
||||
const GET_LDAP_CONFIG_SQL: &str = r#"
|
||||
SELECT
|
||||
server_url,
|
||||
bind_dn,
|
||||
bind_password_encrypted,
|
||||
base_dn,
|
||||
user_search_filter,
|
||||
username_attr,
|
||||
email_attr,
|
||||
display_name_attr,
|
||||
is_enabled,
|
||||
is_exclusive,
|
||||
use_starttls,
|
||||
connect_timeout
|
||||
FROM ldap_configs
|
||||
ORDER BY id ASC
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const UPDATE_LDAP_CONFIG_SQL: &str = r#"
|
||||
UPDATE ldap_configs
|
||||
SET
|
||||
server_url = $1,
|
||||
bind_dn = $2,
|
||||
bind_password_encrypted = $3,
|
||||
base_dn = $4,
|
||||
user_search_filter = $5,
|
||||
username_attr = $6,
|
||||
email_attr = $7,
|
||||
display_name_attr = $8,
|
||||
is_enabled = $9,
|
||||
is_exclusive = $10,
|
||||
use_starttls = $11,
|
||||
connect_timeout = $12,
|
||||
updated_at = NOW()
|
||||
WHERE id = (
|
||||
SELECT id
|
||||
FROM ldap_configs
|
||||
ORDER BY id ASC
|
||||
LIMIT 1
|
||||
)
|
||||
RETURNING
|
||||
server_url,
|
||||
bind_dn,
|
||||
bind_password_encrypted,
|
||||
base_dn,
|
||||
user_search_filter,
|
||||
username_attr,
|
||||
email_attr,
|
||||
display_name_attr,
|
||||
is_enabled,
|
||||
is_exclusive,
|
||||
use_starttls,
|
||||
connect_timeout
|
||||
"#;
|
||||
|
||||
const INSERT_LDAP_CONFIG_SQL: &str = r#"
|
||||
INSERT INTO ldap_configs (
|
||||
server_url,
|
||||
bind_dn,
|
||||
bind_password_encrypted,
|
||||
base_dn,
|
||||
user_search_filter,
|
||||
username_attr,
|
||||
email_attr,
|
||||
display_name_attr,
|
||||
is_enabled,
|
||||
is_exclusive,
|
||||
use_starttls,
|
||||
connect_timeout,
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4,
|
||||
$5,
|
||||
$6,
|
||||
$7,
|
||||
$8,
|
||||
$9,
|
||||
$10,
|
||||
$11,
|
||||
$12,
|
||||
NOW(),
|
||||
NOW()
|
||||
)
|
||||
RETURNING
|
||||
server_url,
|
||||
bind_dn,
|
||||
bind_password_encrypted,
|
||||
base_dn,
|
||||
user_search_filter,
|
||||
username_attr,
|
||||
email_attr,
|
||||
display_name_attr,
|
||||
is_enabled,
|
||||
is_exclusive,
|
||||
use_starttls,
|
||||
connect_timeout
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxAuthModuleReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxAuthModuleReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxAuthModuleRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxAuthModuleRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AuthModuleReadRepository for SqlxAuthModuleReadRepository {
|
||||
async fn list_enabled_oauth_providers(
|
||||
&self,
|
||||
) -> Result<Vec<StoredOAuthProviderModuleConfig>, DataLayerError> {
|
||||
let rows = sqlx::query(LIST_ENABLED_OAUTH_PROVIDERS_SQL)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_oauth_row).collect()
|
||||
}
|
||||
|
||||
async fn get_ldap_config(&self) -> Result<Option<StoredLdapModuleConfig>, DataLayerError> {
|
||||
let row = sqlx::query(GET_LDAP_CONFIG_SQL)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_ldap_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AuthModuleReadRepository for SqlxAuthModuleRepository {
|
||||
async fn list_enabled_oauth_providers(
|
||||
&self,
|
||||
) -> Result<Vec<StoredOAuthProviderModuleConfig>, DataLayerError> {
|
||||
let rows = sqlx::query(LIST_ENABLED_OAUTH_PROVIDERS_SQL)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_oauth_row).collect()
|
||||
}
|
||||
|
||||
async fn get_ldap_config(&self) -> Result<Option<StoredLdapModuleConfig>, DataLayerError> {
|
||||
let row = sqlx::query(GET_LDAP_CONFIG_SQL)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_ldap_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AuthModuleWriteRepository for SqlxAuthModuleRepository {
|
||||
async fn upsert_ldap_config(
|
||||
&self,
|
||||
config: &StoredLdapModuleConfig,
|
||||
) -> Result<Option<StoredLdapModuleConfig>, DataLayerError> {
|
||||
let updated = sqlx::query(UPDATE_LDAP_CONFIG_SQL)
|
||||
.bind(&config.server_url)
|
||||
.bind(&config.bind_dn)
|
||||
.bind(config.bind_password_encrypted.as_deref())
|
||||
.bind(&config.base_dn)
|
||||
.bind(config.user_search_filter.as_deref())
|
||||
.bind(config.username_attr.as_deref())
|
||||
.bind(config.email_attr.as_deref())
|
||||
.bind(config.display_name_attr.as_deref())
|
||||
.bind(config.is_enabled)
|
||||
.bind(config.is_exclusive)
|
||||
.bind(config.use_starttls)
|
||||
.bind(config.connect_timeout)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
if let Some(row) = updated.as_ref() {
|
||||
return map_ldap_row(row).map(Some);
|
||||
}
|
||||
|
||||
let inserted = sqlx::query(INSERT_LDAP_CONFIG_SQL)
|
||||
.bind(&config.server_url)
|
||||
.bind(&config.bind_dn)
|
||||
.bind(config.bind_password_encrypted.as_deref())
|
||||
.bind(&config.base_dn)
|
||||
.bind(config.user_search_filter.as_deref())
|
||||
.bind(config.username_attr.as_deref())
|
||||
.bind(config.email_attr.as_deref())
|
||||
.bind(config.display_name_attr.as_deref())
|
||||
.bind(config.is_enabled)
|
||||
.bind(config.is_exclusive)
|
||||
.bind(config.use_starttls)
|
||||
.bind(config.connect_timeout)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
inserted.as_ref().map(map_ldap_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
fn map_oauth_row(row: &PgRow) -> Result<StoredOAuthProviderModuleConfig, DataLayerError> {
|
||||
StoredOAuthProviderModuleConfig::new(
|
||||
row.try_get("provider_type")?,
|
||||
row.try_get("display_name")?,
|
||||
row.try_get("client_id")?,
|
||||
row.try_get("client_secret_encrypted")?,
|
||||
row.try_get("redirect_uri")?,
|
||||
)
|
||||
}
|
||||
|
||||
fn map_ldap_row(row: &PgRow) -> Result<StoredLdapModuleConfig, DataLayerError> {
|
||||
Ok(StoredLdapModuleConfig {
|
||||
server_url: row.try_get("server_url")?,
|
||||
bind_dn: row.try_get("bind_dn")?,
|
||||
bind_password_encrypted: row.try_get("bind_password_encrypted")?,
|
||||
base_dn: row.try_get("base_dn")?,
|
||||
user_search_filter: row.try_get("user_search_filter")?,
|
||||
username_attr: row.try_get("username_attr")?,
|
||||
email_attr: row.try_get("email_attr")?,
|
||||
display_name_attr: row.try_get("display_name_attr")?,
|
||||
is_enabled: row.try_get("is_enabled")?,
|
||||
is_exclusive: row.try_get("is_exclusive")?,
|
||||
use_starttls: row.try_get("use_starttls")?,
|
||||
connect_timeout: row.try_get("connect_timeout")?,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{SqlxAuthModuleReadRepository, SqlxAuthModuleRepository};
|
||||
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 = SqlxAuthModuleReadRepository::new(pool);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn writable_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 = SqlxAuthModuleRepository::new(pool);
|
||||
}
|
||||
}
|
||||
73
crates/aether-data/src/repository/auth_modules/types.rs
Normal file
73
crates/aether-data/src/repository/auth_modules/types.rs
Normal file
@@ -0,0 +1,73 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredOAuthProviderModuleConfig {
|
||||
pub provider_type: String,
|
||||
pub display_name: String,
|
||||
pub client_id: String,
|
||||
pub client_secret_encrypted: Option<String>,
|
||||
pub redirect_uri: String,
|
||||
}
|
||||
|
||||
impl StoredOAuthProviderModuleConfig {
|
||||
pub fn new(
|
||||
provider_type: String,
|
||||
display_name: String,
|
||||
client_id: String,
|
||||
client_secret_encrypted: Option<String>,
|
||||
redirect_uri: String,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if provider_type.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"oauth_providers.provider_type is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if display_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"oauth_providers.display_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(Self {
|
||||
provider_type,
|
||||
display_name,
|
||||
client_id,
|
||||
client_secret_encrypted,
|
||||
redirect_uri,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredLdapModuleConfig {
|
||||
pub server_url: String,
|
||||
pub bind_dn: String,
|
||||
pub bind_password_encrypted: Option<String>,
|
||||
pub base_dn: String,
|
||||
pub user_search_filter: Option<String>,
|
||||
pub username_attr: Option<String>,
|
||||
pub email_attr: Option<String>,
|
||||
pub display_name_attr: Option<String>,
|
||||
pub is_enabled: bool,
|
||||
pub is_exclusive: bool,
|
||||
pub use_starttls: bool,
|
||||
pub connect_timeout: Option<i32>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AuthModuleReadRepository: Send + Sync {
|
||||
async fn list_enabled_oauth_providers(
|
||||
&self,
|
||||
) -> Result<Vec<StoredOAuthProviderModuleConfig>, crate::DataLayerError>;
|
||||
|
||||
async fn get_ldap_config(
|
||||
&self,
|
||||
) -> Result<Option<StoredLdapModuleConfig>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AuthModuleWriteRepository: Send + Sync {
|
||||
async fn upsert_ldap_config(
|
||||
&self,
|
||||
config: &StoredLdapModuleConfig,
|
||||
) -> Result<Option<StoredLdapModuleConfig>, crate::DataLayerError>;
|
||||
}
|
||||
113
crates/aether-data/src/repository/billing/memory.rs
Normal file
113
crates/aether-data/src/repository/billing/memory.rs
Normal file
@@ -0,0 +1,113 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{BillingReadRepository, StoredBillingModelContext};
|
||||
use crate::DataLayerError;
|
||||
|
||||
type BillingContextKey = (String, String, Option<String>);
|
||||
type BillingContextMap = BTreeMap<BillingContextKey, StoredBillingModelContext>;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryBillingReadRepository {
|
||||
by_key: RwLock<BillingContextMap>,
|
||||
}
|
||||
|
||||
impl InMemoryBillingReadRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredBillingModelContext>,
|
||||
{
|
||||
let mut by_key = BTreeMap::new();
|
||||
for item in items {
|
||||
by_key.insert(
|
||||
(
|
||||
item.provider_id.clone(),
|
||||
item.global_model_name.clone(),
|
||||
item.provider_api_key_id.clone(),
|
||||
),
|
||||
item,
|
||||
);
|
||||
}
|
||||
Self {
|
||||
by_key: RwLock::new(by_key),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BillingReadRepository for InMemoryBillingReadRepository {
|
||||
async fn find_model_context(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
provider_api_key_id: Option<&str>,
|
||||
global_model_name: &str,
|
||||
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||
let key = (
|
||||
provider_id.to_string(),
|
||||
global_model_name.to_string(),
|
||||
provider_api_key_id.map(ToOwned::to_owned),
|
||||
);
|
||||
let by_key = self.by_key.read().expect("billing repository lock");
|
||||
if let Some(value) = by_key.get(&key) {
|
||||
return Ok(Some(value.clone()));
|
||||
}
|
||||
|
||||
if let Some(value) = by_key
|
||||
.get(&(provider_id.to_string(), global_model_name.to_string(), None))
|
||||
.cloned()
|
||||
{
|
||||
return Ok(Some(value));
|
||||
}
|
||||
|
||||
Ok(by_key
|
||||
.iter()
|
||||
.find(|((stored_provider_id, stored_model_name, _), _)| {
|
||||
stored_provider_id == provider_id && stored_model_name == global_model_name
|
||||
})
|
||||
.map(|(_, value)| value.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::InMemoryBillingReadRepository;
|
||||
use crate::repository::billing::{BillingReadRepository, StoredBillingModelContext};
|
||||
|
||||
fn sample_context() -> StoredBillingModelContext {
|
||||
StoredBillingModelContext::new(
|
||||
"provider-1".to_string(),
|
||||
Some("pay_as_you_go".to_string()),
|
||||
Some("key-1".to_string()),
|
||||
Some(json!({"openai:chat": 0.8})),
|
||||
Some(60),
|
||||
"global-model-1".to_string(),
|
||||
"gpt-5".to_string(),
|
||||
Some(json!({"streaming": true})),
|
||||
Some(0.02),
|
||||
Some(json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]})),
|
||||
Some("model-1".to_string()),
|
||||
Some("gpt-5-upstream".to_string()),
|
||||
None,
|
||||
Some(0.01),
|
||||
None,
|
||||
)
|
||||
.expect("billing context should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn falls_back_to_provider_without_key_scope() {
|
||||
let repository = InMemoryBillingReadRepository::seed(vec![sample_context()]);
|
||||
let stored = repository
|
||||
.find_model_context("provider-1", Some("key-2"), "gpt-5")
|
||||
.await
|
||||
.expect("lookup should succeed")
|
||||
.expect("context should exist");
|
||||
|
||||
assert_eq!(stored.provider_id, "provider-1");
|
||||
assert_eq!(stored.global_model_name, "gpt-5");
|
||||
}
|
||||
}
|
||||
7
crates/aether-data/src/repository/billing/mod.rs
Normal file
7
crates/aether-data/src/repository/billing/mod.rs
Normal file
@@ -0,0 +1,7 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryBillingReadRepository;
|
||||
pub use sql::SqlxBillingReadRepository;
|
||||
pub use types::{BillingReadRepository, StoredBillingModelContext};
|
||||
121
crates/aether-data/src/repository/billing/sql.rs
Normal file
121
crates/aether-data/src/repository/billing/sql.rs
Normal file
@@ -0,0 +1,121 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::types::{BillingReadRepository, StoredBillingModelContext};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_MODEL_CONTEXT_SQL: &str = r#"
|
||||
SELECT
|
||||
p.id AS provider_id,
|
||||
CAST(p.billing_type AS TEXT) AS provider_billing_type,
|
||||
pak.id AS provider_api_key_id,
|
||||
pak.rate_multipliers AS provider_api_key_rate_multipliers,
|
||||
pak.cache_ttl_minutes AS provider_api_key_cache_ttl_minutes,
|
||||
gm.id AS global_model_id,
|
||||
gm.name AS global_model_name,
|
||||
gm.config AS global_model_config,
|
||||
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
|
||||
gm.default_tiered_pricing AS default_tiered_pricing,
|
||||
m.id AS model_id,
|
||||
m.provider_model_name AS model_provider_model_name,
|
||||
m.config AS model_config,
|
||||
CAST(m.price_per_request AS DOUBLE PRECISION) AS model_price_per_request,
|
||||
m.tiered_pricing AS model_tiered_pricing
|
||||
FROM providers p
|
||||
INNER JOIN global_models gm
|
||||
ON gm.name = $2
|
||||
AND gm.is_active = TRUE
|
||||
LEFT JOIN models m
|
||||
ON m.global_model_id = gm.id
|
||||
AND m.provider_id = p.id
|
||||
AND m.is_active = TRUE
|
||||
LEFT JOIN provider_api_keys pak
|
||||
ON pak.id = $3
|
||||
AND pak.provider_id = p.id
|
||||
WHERE p.id = $1
|
||||
ORDER BY COALESCE(m.is_available, FALSE) DESC, m.created_at ASC
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxBillingReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxBillingReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub async fn find_model_context(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
provider_api_key_id: Option<&str>,
|
||||
global_model_name: &str,
|
||||
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_MODEL_CONTEXT_SQL)
|
||||
.bind(provider_id)
|
||||
.bind(global_model_name)
|
||||
.bind(provider_api_key_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BillingReadRepository for SqlxBillingReadRepository {
|
||||
async fn find_model_context(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
provider_api_key_id: Option<&str>,
|
||||
global_model_name: &str,
|
||||
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||
Self::find_model_context(self, provider_id, provider_api_key_id, global_model_name).await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_row(row: &sqlx::postgres::PgRow) -> Result<StoredBillingModelContext, DataLayerError> {
|
||||
StoredBillingModelContext::new(
|
||||
row.try_get("provider_id")?,
|
||||
row.try_get("provider_billing_type")?,
|
||||
row.try_get("provider_api_key_id")?,
|
||||
row.try_get("provider_api_key_rate_multipliers")?,
|
||||
row.try_get::<Option<i32>, _>("provider_api_key_cache_ttl_minutes")?
|
||||
.map(i64::from),
|
||||
row.try_get("global_model_id")?,
|
||||
row.try_get("global_model_name")?,
|
||||
row.try_get("global_model_config")?,
|
||||
row.try_get("default_price_per_request")?,
|
||||
row.try_get("default_tiered_pricing")?,
|
||||
row.try_get("model_id")?,
|
||||
row.try_get("model_provider_model_name")?,
|
||||
row.try_get("model_config")?,
|
||||
row.try_get("model_price_per_request")?,
|
||||
row.try_get("model_tiered_pricing")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxBillingReadRepository;
|
||||
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 = SqlxBillingReadRepository::new(pool);
|
||||
}
|
||||
}
|
||||
85
crates/aether-data/src/repository/billing/types.rs
Normal file
85
crates/aether-data/src/repository/billing/types.rs
Normal file
@@ -0,0 +1,85 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredBillingModelContext {
|
||||
pub provider_id: String,
|
||||
pub provider_billing_type: Option<String>,
|
||||
pub provider_api_key_id: Option<String>,
|
||||
pub provider_api_key_rate_multipliers: Option<Value>,
|
||||
pub provider_api_key_cache_ttl_minutes: Option<i64>,
|
||||
pub global_model_id: String,
|
||||
pub global_model_name: String,
|
||||
pub global_model_config: Option<Value>,
|
||||
pub default_price_per_request: Option<f64>,
|
||||
pub default_tiered_pricing: Option<Value>,
|
||||
pub model_id: Option<String>,
|
||||
pub model_provider_model_name: Option<String>,
|
||||
pub model_config: Option<Value>,
|
||||
pub model_price_per_request: Option<f64>,
|
||||
pub model_tiered_pricing: Option<Value>,
|
||||
}
|
||||
|
||||
impl StoredBillingModelContext {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
provider_id: String,
|
||||
provider_billing_type: Option<String>,
|
||||
provider_api_key_id: Option<String>,
|
||||
provider_api_key_rate_multipliers: Option<Value>,
|
||||
provider_api_key_cache_ttl_minutes: Option<i64>,
|
||||
global_model_id: String,
|
||||
global_model_name: String,
|
||||
global_model_config: Option<Value>,
|
||||
default_price_per_request: Option<f64>,
|
||||
default_tiered_pricing: Option<Value>,
|
||||
model_id: Option<String>,
|
||||
model_provider_model_name: Option<String>,
|
||||
model_config: Option<Value>,
|
||||
model_price_per_request: Option<f64>,
|
||||
model_tiered_pricing: Option<Value>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if provider_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"billing.provider_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if global_model_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"billing.global_model_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if global_model_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"billing.global_model_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(Self {
|
||||
provider_id,
|
||||
provider_billing_type,
|
||||
provider_api_key_id,
|
||||
provider_api_key_rate_multipliers,
|
||||
provider_api_key_cache_ttl_minutes,
|
||||
global_model_id,
|
||||
global_model_name,
|
||||
global_model_config,
|
||||
default_price_per_request,
|
||||
default_tiered_pricing,
|
||||
model_id,
|
||||
model_provider_model_name,
|
||||
model_config,
|
||||
model_price_per_request,
|
||||
model_tiered_pricing,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait BillingReadRepository: Send + Sync {
|
||||
async fn find_model_context(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
provider_api_key_id: Option<&str>,
|
||||
global_model_name: &str,
|
||||
) -> Result<Option<StoredBillingModelContext>, crate::DataLayerError>;
|
||||
}
|
||||
154
crates/aether-data/src/repository/candidate_selection/memory.rs
Normal file
154
crates/aether-data/src/repository/candidate_selection/memory.rs
Normal file
@@ -0,0 +1,154 @@
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryMinimalCandidateSelectionReadRepository {
|
||||
rows: RwLock<Vec<StoredMinimalCandidateSelectionRow>>,
|
||||
}
|
||||
|
||||
impl InMemoryMinimalCandidateSelectionReadRepository {
|
||||
pub fn seed<I>(rows: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredMinimalCandidateSelectionRow>,
|
||||
{
|
||||
Self {
|
||||
rows: RwLock::new(rows.into_iter().collect()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl MinimalCandidateSelectionReadRepository for InMemoryMinimalCandidateSelectionReadRepository {
|
||||
async fn list_for_exact_api_format(
|
||||
&self,
|
||||
api_format: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
let api_format = api_format.trim();
|
||||
let mut rows = self
|
||||
.rows
|
||||
.read()
|
||||
.expect("candidate selection repository lock")
|
||||
.iter()
|
||||
.filter(|row| {
|
||||
row.provider_is_active
|
||||
&& row.endpoint_is_active
|
||||
&& row.key_is_active
|
||||
&& row.model_is_active
|
||||
&& row.model_is_available
|
||||
&& row.endpoint_api_format.eq_ignore_ascii_case(api_format)
|
||||
&& row.key_supports_api_format(api_format)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
rows.sort_by(|left, right| {
|
||||
left.provider_priority
|
||||
.cmp(&right.provider_priority)
|
||||
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
|
||||
.then(left.provider_id.cmp(&right.provider_id))
|
||||
.then(left.endpoint_id.cmp(&right.endpoint_id))
|
||||
.then(left.key_id.cmp(&right.key_id))
|
||||
.then(left.model_id.cmp(&right.model_id))
|
||||
});
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
async fn list_for_exact_api_format_and_global_model(
|
||||
&self,
|
||||
api_format: &str,
|
||||
global_model_name: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
let rows = self.list_for_exact_api_format(api_format).await?;
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.filter(|row| row.global_model_name == global_model_name)
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryMinimalCandidateSelectionReadRepository;
|
||||
use crate::repository::candidate_selection::{
|
||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
};
|
||||
|
||||
fn sample_row(
|
||||
provider_id: &str,
|
||||
api_format: &str,
|
||||
global_model_name: &str,
|
||||
provider_priority: i32,
|
||||
) -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: provider_id.to_string(),
|
||||
provider_name: provider_id.to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority,
|
||||
provider_is_active: true,
|
||||
endpoint_id: format!("endpoint-{provider_id}"),
|
||||
endpoint_api_format: api_format.to_string(),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: format!("key-{provider_id}"),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec![api_format.to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 50,
|
||||
key_global_priority_by_format: None,
|
||||
model_id: format!("model-{provider_id}"),
|
||||
global_model_id: "global-model-1".to_string(),
|
||||
global_model_name: global_model_name.to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: global_model_name.to_string(),
|
||||
model_provider_model_mappings: None,
|
||||
model_supports_streaming: None,
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn filters_by_exact_api_format_and_global_model() {
|
||||
let repository = InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_row("provider-2", "openai:chat", "gpt-4.1", 20),
|
||||
sample_row("provider-1", "openai:chat", "gpt-4.1", 10),
|
||||
sample_row("provider-3", "openai:responses", "gpt-4.1", 5),
|
||||
sample_row("provider-4", "openai:chat", "gpt-4.1-mini", 1),
|
||||
]);
|
||||
|
||||
let rows = repository
|
||||
.list_for_exact_api_format_and_global_model("openai:chat", "gpt-4.1")
|
||||
.await
|
||||
.expect("list should succeed");
|
||||
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(rows[0].provider_id, "provider-1");
|
||||
assert_eq!(rows[1].provider_id, "provider-2");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn filters_by_exact_api_format_only() {
|
||||
let repository = InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_row("provider-2", "openai:chat", "gpt-4.1", 20),
|
||||
sample_row("provider-1", "openai:chat", "gpt-4.1-mini", 10),
|
||||
sample_row("provider-3", "openai:responses", "gpt-4.1", 5),
|
||||
]);
|
||||
|
||||
let rows = repository
|
||||
.list_for_exact_api_format("openai:chat")
|
||||
.await
|
||||
.expect("list should succeed");
|
||||
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(rows[0].provider_id, "provider-1");
|
||||
assert_eq!(rows[1].provider_id, "provider-2");
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/candidate_selection/mod.rs
Normal file
10
crates/aether-data/src/repository/candidate_selection/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryMinimalCandidateSelectionReadRepository;
|
||||
pub use sql::SqlxMinimalCandidateSelectionReadRepository;
|
||||
pub use types::{
|
||||
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
|
||||
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
|
||||
};
|
||||
548
crates/aether-data/src/repository/candidate_selection/sql.rs
Normal file
548
crates/aether-data/src/repository/candidate_selection/sql.rs
Normal file
@@ -0,0 +1,548 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredProviderModelMapping,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const LIST_FOR_EXACT_API_FORMAT_SQL: &str = r#"
|
||||
SELECT
|
||||
p.id AS provider_id,
|
||||
p.name AS provider_name,
|
||||
p.provider_type AS provider_type,
|
||||
p.provider_priority AS provider_priority,
|
||||
p.is_active AS provider_is_active,
|
||||
pe.id AS endpoint_id,
|
||||
pe.api_format AS endpoint_api_format,
|
||||
pe.api_family AS endpoint_api_family,
|
||||
pe.endpoint_kind AS endpoint_kind,
|
||||
pe.is_active AS endpoint_is_active,
|
||||
pak.id AS key_id,
|
||||
pak.name AS key_name,
|
||||
pak.auth_type AS key_auth_type,
|
||||
pak.is_active AS key_is_active,
|
||||
pak.api_formats AS key_api_formats,
|
||||
pak.allowed_models AS key_allowed_models,
|
||||
pak.capabilities AS key_capabilities,
|
||||
pak.internal_priority AS key_internal_priority,
|
||||
pak.global_priority_by_format AS key_global_priority_by_format,
|
||||
m.id AS model_id,
|
||||
m.global_model_id AS global_model_id,
|
||||
gm.name AS global_model_name,
|
||||
CASE
|
||||
WHEN gm.config IS NOT NULL THEN gm.config -> 'model_mappings'
|
||||
ELSE NULL
|
||||
END AS global_model_mappings,
|
||||
CASE
|
||||
WHEN gm.config IS NOT NULL AND gm.config ? 'streaming'
|
||||
THEN (gm.config ->> 'streaming')::BOOLEAN
|
||||
ELSE NULL
|
||||
END AS global_model_supports_streaming,
|
||||
m.provider_model_name AS model_provider_model_name,
|
||||
m.provider_model_mappings AS model_provider_model_mappings,
|
||||
m.supports_streaming AS model_supports_streaming,
|
||||
m.is_active AS model_is_active,
|
||||
m.is_available AS model_is_available
|
||||
FROM providers p
|
||||
INNER JOIN provider_endpoints pe
|
||||
ON pe.provider_id = p.id
|
||||
INNER JOIN provider_api_keys pak
|
||||
ON pak.provider_id = p.id
|
||||
INNER JOIN models m
|
||||
ON m.provider_id = p.id
|
||||
INNER JOIN global_models gm
|
||||
ON gm.id = m.global_model_id
|
||||
WHERE p.is_active = TRUE
|
||||
AND pe.is_active = TRUE
|
||||
AND pak.is_active = TRUE
|
||||
AND m.is_active = TRUE
|
||||
AND m.is_available = TRUE
|
||||
AND gm.is_active = TRUE
|
||||
AND LOWER(pe.api_format) = LOWER($1)
|
||||
AND (
|
||||
pak.api_formats IS NULL
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM json_array_elements_text(pak.api_formats) AS fmt(value)
|
||||
WHERE LOWER(fmt.value) = LOWER($1)
|
||||
)
|
||||
)
|
||||
ORDER BY
|
||||
gm.name ASC,
|
||||
p.provider_priority ASC,
|
||||
pak.internal_priority ASC,
|
||||
p.id ASC,
|
||||
pe.id ASC,
|
||||
pak.id ASC,
|
||||
m.id ASC
|
||||
"#;
|
||||
|
||||
const LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL: &str = r#"
|
||||
SELECT
|
||||
p.id AS provider_id,
|
||||
p.name AS provider_name,
|
||||
p.provider_type AS provider_type,
|
||||
p.provider_priority AS provider_priority,
|
||||
p.is_active AS provider_is_active,
|
||||
pe.id AS endpoint_id,
|
||||
pe.api_format AS endpoint_api_format,
|
||||
pe.api_family AS endpoint_api_family,
|
||||
pe.endpoint_kind AS endpoint_kind,
|
||||
pe.is_active AS endpoint_is_active,
|
||||
pak.id AS key_id,
|
||||
pak.name AS key_name,
|
||||
pak.auth_type AS key_auth_type,
|
||||
pak.is_active AS key_is_active,
|
||||
pak.api_formats AS key_api_formats,
|
||||
pak.allowed_models AS key_allowed_models,
|
||||
pak.capabilities AS key_capabilities,
|
||||
pak.internal_priority AS key_internal_priority,
|
||||
pak.global_priority_by_format AS key_global_priority_by_format,
|
||||
m.id AS model_id,
|
||||
m.global_model_id AS global_model_id,
|
||||
gm.name AS global_model_name,
|
||||
CASE
|
||||
WHEN gm.config IS NOT NULL THEN gm.config -> 'model_mappings'
|
||||
ELSE NULL
|
||||
END AS global_model_mappings,
|
||||
CASE
|
||||
WHEN gm.config IS NOT NULL AND gm.config ? 'streaming'
|
||||
THEN (gm.config ->> 'streaming')::BOOLEAN
|
||||
ELSE NULL
|
||||
END AS global_model_supports_streaming,
|
||||
m.provider_model_name AS model_provider_model_name,
|
||||
m.provider_model_mappings AS model_provider_model_mappings,
|
||||
m.supports_streaming AS model_supports_streaming,
|
||||
m.is_active AS model_is_active,
|
||||
m.is_available AS model_is_available
|
||||
FROM providers p
|
||||
INNER JOIN provider_endpoints pe
|
||||
ON pe.provider_id = p.id
|
||||
INNER JOIN provider_api_keys pak
|
||||
ON pak.provider_id = p.id
|
||||
INNER JOIN models m
|
||||
ON m.provider_id = p.id
|
||||
INNER JOIN global_models gm
|
||||
ON gm.id = m.global_model_id
|
||||
WHERE p.is_active = TRUE
|
||||
AND pe.is_active = TRUE
|
||||
AND pak.is_active = TRUE
|
||||
AND m.is_active = TRUE
|
||||
AND m.is_available = TRUE
|
||||
AND gm.is_active = TRUE
|
||||
AND LOWER(pe.api_format) = LOWER($1)
|
||||
AND gm.name = $2
|
||||
AND (
|
||||
pak.api_formats IS NULL
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM json_array_elements_text(pak.api_formats) AS fmt(value)
|
||||
WHERE LOWER(fmt.value) = LOWER($1)
|
||||
)
|
||||
)
|
||||
ORDER BY
|
||||
p.provider_priority ASC,
|
||||
pak.internal_priority ASC,
|
||||
p.id ASC,
|
||||
pe.id ASC,
|
||||
pak.id ASC,
|
||||
m.id ASC
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxMinimalCandidateSelectionReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxMinimalCandidateSelectionReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub async fn list_for_exact_api_format(
|
||||
&self,
|
||||
api_format: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
let rows = sqlx::query(LIST_FOR_EXACT_API_FORMAT_SQL)
|
||||
.bind(api_format)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_candidate_selection_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_for_exact_api_format_and_global_model(
|
||||
&self,
|
||||
api_format: &str,
|
||||
global_model_name: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
let rows = sqlx::query(LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL)
|
||||
.bind(api_format)
|
||||
.bind(global_model_name)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_candidate_selection_row).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl MinimalCandidateSelectionReadRepository for SqlxMinimalCandidateSelectionReadRepository {
|
||||
async fn list_for_exact_api_format(
|
||||
&self,
|
||||
api_format: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Self::list_for_exact_api_format(self, api_format).await
|
||||
}
|
||||
|
||||
async fn list_for_exact_api_format_and_global_model(
|
||||
&self,
|
||||
api_format: &str,
|
||||
global_model_name: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Self::list_for_exact_api_format_and_global_model(self, api_format, global_model_name).await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_candidate_selection_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> Result<StoredMinimalCandidateSelectionRow, DataLayerError> {
|
||||
Ok(StoredMinimalCandidateSelectionRow {
|
||||
provider_id: row.try_get("provider_id")?,
|
||||
provider_name: row.try_get("provider_name")?,
|
||||
provider_type: row.try_get("provider_type")?,
|
||||
provider_priority: row.try_get("provider_priority")?,
|
||||
provider_is_active: row.try_get("provider_is_active")?,
|
||||
endpoint_id: row.try_get("endpoint_id")?,
|
||||
endpoint_api_format: row.try_get("endpoint_api_format")?,
|
||||
endpoint_api_family: row.try_get("endpoint_api_family")?,
|
||||
endpoint_kind: row.try_get("endpoint_kind")?,
|
||||
endpoint_is_active: row.try_get("endpoint_is_active")?,
|
||||
key_id: row.try_get("key_id")?,
|
||||
key_name: row.try_get("key_name")?,
|
||||
key_auth_type: row.try_get("key_auth_type")?,
|
||||
key_is_active: row.try_get("key_is_active")?,
|
||||
key_api_formats: parse_string_list(
|
||||
row.try_get("key_api_formats")?,
|
||||
"provider_api_keys.api_formats",
|
||||
)?,
|
||||
key_allowed_models: parse_string_list(
|
||||
row.try_get("key_allowed_models")?,
|
||||
"provider_api_keys.allowed_models",
|
||||
)?,
|
||||
key_capabilities: row.try_get("key_capabilities")?,
|
||||
key_internal_priority: row.try_get("key_internal_priority")?,
|
||||
key_global_priority_by_format: row.try_get("key_global_priority_by_format")?,
|
||||
model_id: row.try_get("model_id")?,
|
||||
global_model_id: row.try_get("global_model_id")?,
|
||||
global_model_name: row.try_get("global_model_name")?,
|
||||
global_model_mappings: parse_string_list(
|
||||
row.try_get("global_model_mappings")?,
|
||||
"global_models.config.model_mappings",
|
||||
)?,
|
||||
global_model_supports_streaming: row.try_get("global_model_supports_streaming")?,
|
||||
model_provider_model_name: row.try_get("model_provider_model_name")?,
|
||||
model_provider_model_mappings: parse_provider_model_mappings(
|
||||
row.try_get("model_provider_model_mappings")?,
|
||||
)?,
|
||||
model_supports_streaming: row.try_get("model_supports_streaming")?,
|
||||
model_is_active: row.try_get("model_is_active")?,
|
||||
model_is_available: row.try_get("model_is_available")?,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_string_list(
|
||||
value: Option<serde_json::Value>,
|
||||
field_name: &str,
|
||||
) -> Result<Option<Vec<String>>, DataLayerError> {
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
parse_string_list_value(&value, field_name)
|
||||
}
|
||||
|
||||
fn parse_string_list_value(
|
||||
value: &serde_json::Value,
|
||||
field_name: &str,
|
||||
) -> Result<Option<Vec<String>>, DataLayerError> {
|
||||
match value {
|
||||
serde_json::Value::Null => Ok(None),
|
||||
serde_json::Value::Array(array) => parse_string_list_array(array, field_name).map(Some),
|
||||
serde_json::Value::String(raw) => parse_embedded_string_list(raw, field_name),
|
||||
_ => Err(DataLayerError::UnexpectedValue(format!(
|
||||
"{field_name} is not a JSON array"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_embedded_string_list(
|
||||
raw: &str,
|
||||
field_name: &str,
|
||||
) -> Result<Option<Vec<String>>, DataLayerError> {
|
||||
let raw = raw.trim();
|
||||
if raw.is_empty() || raw.eq_ignore_ascii_case("null") {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if let Ok(decoded) = serde_json::from_str::<serde_json::Value>(raw) {
|
||||
return parse_string_list_value(&decoded, field_name);
|
||||
}
|
||||
|
||||
Ok(Some(vec![raw.to_string()]))
|
||||
}
|
||||
|
||||
fn parse_string_list_array(
|
||||
array: &[serde_json::Value],
|
||||
field_name: &str,
|
||||
) -> Result<Vec<String>, DataLayerError> {
|
||||
let mut items = Vec::with_capacity(array.len());
|
||||
for item in array {
|
||||
let Some(item) = item.as_str() else {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"{field_name} contains a non-string item"
|
||||
)));
|
||||
};
|
||||
let item = item.trim();
|
||||
if !item.is_empty() {
|
||||
items.push(item.to_string());
|
||||
}
|
||||
}
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
fn parse_provider_model_mappings(
|
||||
value: Option<serde_json::Value>,
|
||||
) -> Result<Option<Vec<StoredProviderModelMapping>>, DataLayerError> {
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
parse_provider_model_mappings_value(&value)
|
||||
}
|
||||
|
||||
fn parse_provider_model_mappings_value(
|
||||
value: &serde_json::Value,
|
||||
) -> Result<Option<Vec<StoredProviderModelMapping>>, DataLayerError> {
|
||||
match value {
|
||||
serde_json::Value::Null => Ok(None),
|
||||
serde_json::Value::Array(array) => parse_provider_model_mappings_array(array),
|
||||
serde_json::Value::Object(object) => {
|
||||
parse_provider_model_mapping_object(object).map(|mapping| Some(vec![mapping]))
|
||||
}
|
||||
serde_json::Value::String(raw) => parse_embedded_provider_model_mappings(raw),
|
||||
_ => Err(DataLayerError::UnexpectedValue(
|
||||
"models.provider_model_mappings is not a JSON array".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_embedded_provider_model_mappings(
|
||||
raw: &str,
|
||||
) -> Result<Option<Vec<StoredProviderModelMapping>>, DataLayerError> {
|
||||
let raw = raw.trim();
|
||||
if raw.is_empty() || raw.eq_ignore_ascii_case("null") {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if let Ok(decoded) = serde_json::from_str::<serde_json::Value>(raw) {
|
||||
return parse_provider_model_mappings_value(&decoded);
|
||||
}
|
||||
|
||||
Ok(Some(vec![StoredProviderModelMapping {
|
||||
name: raw.to_string(),
|
||||
priority: 1,
|
||||
api_formats: None,
|
||||
}]))
|
||||
}
|
||||
|
||||
fn parse_provider_model_mappings_array(
|
||||
array: &[serde_json::Value],
|
||||
) -> Result<Option<Vec<StoredProviderModelMapping>>, DataLayerError> {
|
||||
let mut mappings = Vec::with_capacity(array.len());
|
||||
for raw in array {
|
||||
match raw {
|
||||
serde_json::Value::Object(object) => {
|
||||
if let Some(mapping) = parse_provider_model_mapping_object_lenient(object)? {
|
||||
mappings.push(mapping);
|
||||
}
|
||||
}
|
||||
serde_json::Value::String(raw) => {
|
||||
let raw = raw.trim();
|
||||
if !raw.is_empty() {
|
||||
mappings.push(StoredProviderModelMapping {
|
||||
name: raw.to_string(),
|
||||
priority: 1,
|
||||
api_formats: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
serde_json::Value::Null => {}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if mappings.is_empty() {
|
||||
Ok(None)
|
||||
} else {
|
||||
Ok(Some(mappings))
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_provider_model_mapping_object(
|
||||
object: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> Result<StoredProviderModelMapping, DataLayerError> {
|
||||
parse_provider_model_mapping_object_lenient(object)?.ok_or_else(|| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"models.provider_model_mappings item is missing a valid name".to_string(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_provider_model_mapping_object_lenient(
|
||||
object: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> Result<Option<StoredProviderModelMapping>, DataLayerError> {
|
||||
let Some(name) = object
|
||||
.get("name")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let priority = object
|
||||
.get("priority")
|
||||
.and_then(serde_json::Value::as_i64)
|
||||
.unwrap_or(1)
|
||||
.max(1);
|
||||
let api_formats = parse_string_list(
|
||||
object.get("api_formats").cloned(),
|
||||
"models.provider_model_mappings.api_formats",
|
||||
)?;
|
||||
|
||||
Ok(Some(StoredProviderModelMapping {
|
||||
name: name.to_string(),
|
||||
priority: i32::try_from(priority).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"invalid models.provider_model_mappings.priority: {priority}"
|
||||
))
|
||||
})?,
|
||||
api_formats,
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
parse_provider_model_mappings, parse_string_list,
|
||||
SqlxMinimalCandidateSelectionReadRepository,
|
||||
};
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
use crate::repository::candidate_selection::StoredProviderModelMapping;
|
||||
|
||||
#[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 = SqlxMinimalCandidateSelectionReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_string_list_accepts_stringified_array() {
|
||||
let parsed = parse_string_list(
|
||||
Some(json!("[\"gpt-5.2\", \"gpt-5\"]")),
|
||||
"provider_api_keys.allowed_models",
|
||||
)
|
||||
.expect("stringified array should parse");
|
||||
|
||||
assert_eq!(
|
||||
parsed,
|
||||
Some(vec!["gpt-5.2".to_string(), "gpt-5".to_string()])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_string_list_accepts_single_string() {
|
||||
let parsed = parse_string_list(Some(json!("gpt-5.2")), "provider_api_keys.allowed_models")
|
||||
.expect("single string should parse");
|
||||
|
||||
assert_eq!(parsed, Some(vec!["gpt-5.2".to_string()]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_provider_model_mappings_accepts_stringified_array() {
|
||||
let parsed = parse_provider_model_mappings(Some(json!(
|
||||
"[{\"name\":\"gpt-5.2\",\"priority\":2,\"api_formats\":[\"openai:chat\"]}]"
|
||||
)))
|
||||
.expect("stringified provider_model_mappings should parse");
|
||||
|
||||
assert_eq!(
|
||||
parsed,
|
||||
Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5.2".to_string(),
|
||||
priority: 2,
|
||||
api_formats: Some(vec!["openai:chat".to_string()]),
|
||||
}])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_provider_model_mappings_accepts_single_string_alias() {
|
||||
let parsed = parse_provider_model_mappings(Some(json!("gpt-5.2")))
|
||||
.expect("single-string provider_model_mappings should parse");
|
||||
|
||||
assert_eq!(
|
||||
parsed,
|
||||
Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5.2".to_string(),
|
||||
priority: 1,
|
||||
api_formats: None,
|
||||
}])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_provider_model_mappings_skips_invalid_array_items() {
|
||||
let parsed = parse_provider_model_mappings(Some(json!([
|
||||
{"name": "gpt-5.2", "priority": 1},
|
||||
{"priority": 2},
|
||||
3,
|
||||
null,
|
||||
"gpt-5.2-mini"
|
||||
])))
|
||||
.expect("mixed provider_model_mappings should parse");
|
||||
|
||||
assert_eq!(
|
||||
parsed,
|
||||
Some(vec![
|
||||
StoredProviderModelMapping {
|
||||
name: "gpt-5.2".to_string(),
|
||||
priority: 1,
|
||||
api_formats: None,
|
||||
},
|
||||
StoredProviderModelMapping {
|
||||
name: "gpt-5.2-mini".to_string(),
|
||||
priority: 1,
|
||||
api_formats: None,
|
||||
}
|
||||
])
|
||||
);
|
||||
}
|
||||
}
|
||||
144
crates/aether-data/src/repository/candidate_selection/types.rs
Normal file
144
crates/aether-data/src/repository/candidate_selection/types.rs
Normal file
@@ -0,0 +1,144 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderModelMapping {
|
||||
pub name: String,
|
||||
pub priority: i32,
|
||||
pub api_formats: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredMinimalCandidateSelectionRow {
|
||||
pub provider_id: String,
|
||||
pub provider_name: String,
|
||||
pub provider_type: String,
|
||||
pub provider_priority: i32,
|
||||
pub provider_is_active: bool,
|
||||
pub endpoint_id: String,
|
||||
pub endpoint_api_format: String,
|
||||
pub endpoint_api_family: Option<String>,
|
||||
pub endpoint_kind: Option<String>,
|
||||
pub endpoint_is_active: bool,
|
||||
pub key_id: String,
|
||||
pub key_name: String,
|
||||
pub key_auth_type: String,
|
||||
pub key_is_active: bool,
|
||||
pub key_api_formats: Option<Vec<String>>,
|
||||
pub key_allowed_models: Option<Vec<String>>,
|
||||
pub key_capabilities: Option<serde_json::Value>,
|
||||
pub key_internal_priority: i32,
|
||||
pub key_global_priority_by_format: Option<serde_json::Value>,
|
||||
pub model_id: String,
|
||||
pub global_model_id: String,
|
||||
pub global_model_name: String,
|
||||
pub global_model_mappings: Option<Vec<String>>,
|
||||
pub global_model_supports_streaming: Option<bool>,
|
||||
pub model_provider_model_name: String,
|
||||
pub model_provider_model_mappings: Option<Vec<StoredProviderModelMapping>>,
|
||||
pub model_supports_streaming: Option<bool>,
|
||||
pub model_is_active: bool,
|
||||
pub model_is_available: bool,
|
||||
}
|
||||
|
||||
impl StoredMinimalCandidateSelectionRow {
|
||||
pub fn supports_streaming(&self) -> bool {
|
||||
self.model_supports_streaming
|
||||
.or(self.global_model_supports_streaming)
|
||||
.unwrap_or(true)
|
||||
}
|
||||
|
||||
pub fn key_supports_api_format(&self, api_format: &str) -> bool {
|
||||
let target = api_format.trim();
|
||||
match self.key_api_formats.as_deref() {
|
||||
None => true,
|
||||
Some(formats) => formats
|
||||
.iter()
|
||||
.any(|value| value.eq_ignore_ascii_case(target)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait MinimalCandidateSelectionReadRepository: Send + Sync {
|
||||
async fn list_for_exact_api_format(
|
||||
&self,
|
||||
api_format: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, crate::DataLayerError>;
|
||||
|
||||
async fn list_for_exact_api_format_and_global_model(
|
||||
&self,
|
||||
api_format: &str,
|
||||
global_model_name: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait MinimalCandidateSelectionRepository:
|
||||
MinimalCandidateSelectionReadRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> MinimalCandidateSelectionRepository for T where
|
||||
T: MinimalCandidateSelectionReadRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{StoredMinimalCandidateSelectionRow, StoredProviderModelMapping};
|
||||
|
||||
fn sample_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-1".to_string(),
|
||||
provider_name: "OpenAI".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
endpoint_api_format: "openai:chat".to_string(),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["openai:chat".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 50,
|
||||
key_global_priority_by_format: None,
|
||||
model_id: "model-1".to_string(),
|
||||
global_model_id: "global-model-1".to_string(),
|
||||
global_model_name: "gpt-4.1".to_string(),
|
||||
global_model_mappings: Some(vec!["gpt-4\\.1-.*".to_string()]),
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gpt-4.1-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-4.1-canary".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:chat".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: None,
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn defaults_streaming_support_to_true() {
|
||||
let mut row = sample_row();
|
||||
row.model_supports_streaming = None;
|
||||
row.global_model_supports_streaming = None;
|
||||
|
||||
assert!(row.supports_streaming());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn key_api_formats_none_means_support_all_formats() {
|
||||
let mut row = sample_row();
|
||||
row.key_api_formats = None;
|
||||
|
||||
assert!(row.key_supports_api_format("openai:chat"));
|
||||
assert!(row.key_supports_api_format("openai:responses"));
|
||||
}
|
||||
}
|
||||
@@ -1,9 +1,13 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{RequestCandidateReadRepository, StoredRequestCandidate};
|
||||
use super::types::{
|
||||
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
|
||||
RequestCandidateStatus, RequestCandidateWriteRepository, StoredRequestCandidate,
|
||||
UpsertRequestCandidateRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
@@ -68,13 +72,340 @@ impl RequestCandidateReadRepository for InMemoryRequestCandidateRepository {
|
||||
rows.truncate(limit);
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
async fn list_by_provider_id(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut rows = self
|
||||
.by_id
|
||||
.read()
|
||||
.expect("request candidate repository lock")
|
||||
.values()
|
||||
.filter(|row| row.provider_id.as_deref() == Some(provider_id))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
rows.sort_by(|left, right| right.created_at_unix_secs.cmp(&left.created_at_unix_secs));
|
||||
rows.truncate(limit);
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
async fn list_finalized_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
if endpoint_ids.is_empty() || limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let endpoint_ids = endpoint_ids.iter().cloned().collect::<BTreeSet<_>>();
|
||||
let mut rows = self
|
||||
.by_id
|
||||
.read()
|
||||
.expect("request candidate repository lock")
|
||||
.values()
|
||||
.filter(|row| {
|
||||
row.endpoint_id
|
||||
.as_ref()
|
||||
.is_some_and(|endpoint_id| endpoint_ids.contains(endpoint_id))
|
||||
&& row.created_at_unix_secs >= since_unix_secs
|
||||
&& matches!(
|
||||
row.status,
|
||||
RequestCandidateStatus::Success
|
||||
| RequestCandidateStatus::Failed
|
||||
| RequestCandidateStatus::Skipped
|
||||
)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
rows.sort_by(|left, right| right.created_at_unix_secs.cmp(&left.created_at_unix_secs));
|
||||
rows.truncate(limit);
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
async fn count_finalized_statuses_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
) -> Result<Vec<PublicHealthStatusCount>, DataLayerError> {
|
||||
if endpoint_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let endpoint_ids = endpoint_ids.iter().cloned().collect::<BTreeSet<_>>();
|
||||
let mut counts = BTreeMap::<(String, &'static str), u64>::new();
|
||||
for row in self
|
||||
.by_id
|
||||
.read()
|
||||
.expect("request candidate repository lock")
|
||||
.values()
|
||||
{
|
||||
let Some(endpoint_id) = row.endpoint_id.as_ref() else {
|
||||
continue;
|
||||
};
|
||||
if !endpoint_ids.contains(endpoint_id) || row.created_at_unix_secs < since_unix_secs {
|
||||
continue;
|
||||
}
|
||||
if !matches!(
|
||||
row.status,
|
||||
RequestCandidateStatus::Success
|
||||
| RequestCandidateStatus::Failed
|
||||
| RequestCandidateStatus::Skipped
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
let status_key = match row.status {
|
||||
RequestCandidateStatus::Success => "success",
|
||||
RequestCandidateStatus::Failed => "failed",
|
||||
RequestCandidateStatus::Skipped => "skipped",
|
||||
_ => continue,
|
||||
};
|
||||
*counts.entry((endpoint_id.clone(), status_key)).or_insert(0) += 1;
|
||||
}
|
||||
|
||||
Ok(counts
|
||||
.into_iter()
|
||||
.map(|((endpoint_id, status_key), count)| {
|
||||
let status = match status_key {
|
||||
"success" => RequestCandidateStatus::Success,
|
||||
"failed" => RequestCandidateStatus::Failed,
|
||||
"skipped" => RequestCandidateStatus::Skipped,
|
||||
_ => unreachable!("filtered status should stay finalized"),
|
||||
};
|
||||
PublicHealthStatusCount {
|
||||
endpoint_id,
|
||||
status,
|
||||
count,
|
||||
}
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn aggregate_finalized_timeline_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
until_unix_secs: u64,
|
||||
segments: u32,
|
||||
) -> Result<Vec<PublicHealthTimelineBucket>, DataLayerError> {
|
||||
if endpoint_ids.is_empty() || segments == 0 || until_unix_secs < since_unix_secs {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let endpoint_ids = endpoint_ids.iter().cloned().collect::<BTreeSet<_>>();
|
||||
let span = until_unix_secs.saturating_sub(since_unix_secs);
|
||||
let mut buckets = BTreeMap::<(String, u32), PublicHealthTimelineBucket>::new();
|
||||
|
||||
for row in self
|
||||
.by_id
|
||||
.read()
|
||||
.expect("request candidate repository lock")
|
||||
.values()
|
||||
{
|
||||
let Some(endpoint_id) = row.endpoint_id.as_ref() else {
|
||||
continue;
|
||||
};
|
||||
if !endpoint_ids.contains(endpoint_id)
|
||||
|| row.created_at_unix_secs < since_unix_secs
|
||||
|| row.created_at_unix_secs > until_unix_secs
|
||||
{
|
||||
continue;
|
||||
}
|
||||
if !matches!(
|
||||
row.status,
|
||||
RequestCandidateStatus::Success
|
||||
| RequestCandidateStatus::Failed
|
||||
| RequestCandidateStatus::Skipped
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let segment_idx = if span == 0 {
|
||||
0
|
||||
} else {
|
||||
let offset = row.created_at_unix_secs.saturating_sub(since_unix_secs);
|
||||
let idx = ((offset as u128) * (segments as u128) / (span as u128)) as u32;
|
||||
idx.min(segments.saturating_sub(1))
|
||||
};
|
||||
let bucket = buckets
|
||||
.entry((endpoint_id.clone(), segment_idx))
|
||||
.or_insert_with(|| PublicHealthTimelineBucket {
|
||||
endpoint_id: endpoint_id.clone(),
|
||||
segment_idx,
|
||||
total_count: 0,
|
||||
success_count: 0,
|
||||
failed_count: 0,
|
||||
min_created_at_unix_secs: None,
|
||||
max_created_at_unix_secs: None,
|
||||
});
|
||||
bucket.total_count += 1;
|
||||
if row.status == RequestCandidateStatus::Success {
|
||||
bucket.success_count += 1;
|
||||
} else if row.status == RequestCandidateStatus::Failed {
|
||||
bucket.failed_count += 1;
|
||||
}
|
||||
bucket.min_created_at_unix_secs = Some(
|
||||
bucket
|
||||
.min_created_at_unix_secs
|
||||
.map(|value| value.min(row.created_at_unix_secs))
|
||||
.unwrap_or(row.created_at_unix_secs),
|
||||
);
|
||||
bucket.max_created_at_unix_secs = Some(
|
||||
bucket
|
||||
.max_created_at_unix_secs
|
||||
.map(|value| value.max(row.created_at_unix_secs))
|
||||
.unwrap_or(row.created_at_unix_secs),
|
||||
);
|
||||
}
|
||||
|
||||
Ok(buckets.into_values().collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RequestCandidateWriteRepository for InMemoryRequestCandidateRepository {
|
||||
async fn upsert(
|
||||
&self,
|
||||
candidate: UpsertRequestCandidateRecord,
|
||||
) -> Result<StoredRequestCandidate, DataLayerError> {
|
||||
candidate.validate()?;
|
||||
|
||||
let mut by_id = self
|
||||
.by_id
|
||||
.write()
|
||||
.expect("request candidate repository lock");
|
||||
let existing = by_id
|
||||
.values()
|
||||
.find(|row| {
|
||||
row.request_id == candidate.request_id
|
||||
&& row.candidate_index == candidate.candidate_index
|
||||
&& row.retry_index == candidate.retry_index
|
||||
})
|
||||
.cloned();
|
||||
|
||||
let created_at_unix_secs = existing
|
||||
.as_ref()
|
||||
.map(|row| row.created_at_unix_secs)
|
||||
.or(candidate.created_at_unix_secs)
|
||||
.or(candidate.started_at_unix_secs)
|
||||
.or(candidate.finished_at_unix_secs)
|
||||
.unwrap_or_default();
|
||||
|
||||
let stored = StoredRequestCandidate {
|
||||
id: existing
|
||||
.as_ref()
|
||||
.map(|row| row.id.clone())
|
||||
.unwrap_or_else(|| candidate.id.clone()),
|
||||
request_id: candidate.request_id.clone(),
|
||||
user_id: candidate
|
||||
.user_id
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.user_id.clone())),
|
||||
api_key_id: candidate
|
||||
.api_key_id
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.api_key_id.clone())),
|
||||
username: candidate
|
||||
.username
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.username.clone())),
|
||||
api_key_name: candidate
|
||||
.api_key_name
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.api_key_name.clone())),
|
||||
candidate_index: candidate.candidate_index,
|
||||
retry_index: candidate.retry_index,
|
||||
provider_id: candidate
|
||||
.provider_id
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.provider_id.clone())),
|
||||
endpoint_id: candidate
|
||||
.endpoint_id
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.endpoint_id.clone())),
|
||||
key_id: candidate
|
||||
.key_id
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.key_id.clone())),
|
||||
status: candidate.status,
|
||||
skip_reason: candidate
|
||||
.skip_reason
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.skip_reason.clone())),
|
||||
is_cached: candidate
|
||||
.is_cached
|
||||
.unwrap_or_else(|| existing.as_ref().map(|row| row.is_cached).unwrap_or(false)),
|
||||
status_code: candidate
|
||||
.status_code
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.status_code)),
|
||||
error_type: candidate
|
||||
.error_type
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.error_type.clone())),
|
||||
error_message: candidate
|
||||
.error_message
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.error_message.clone())),
|
||||
latency_ms: candidate
|
||||
.latency_ms
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.latency_ms)),
|
||||
concurrent_requests: candidate
|
||||
.concurrent_requests
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.concurrent_requests)),
|
||||
extra_data: candidate
|
||||
.extra_data
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.extra_data.clone())),
|
||||
required_capabilities: candidate.required_capabilities.or_else(|| {
|
||||
existing
|
||||
.as_ref()
|
||||
.and_then(|row| row.required_capabilities.clone())
|
||||
}),
|
||||
created_at_unix_secs,
|
||||
started_at_unix_secs: candidate
|
||||
.started_at_unix_secs
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.started_at_unix_secs)),
|
||||
finished_at_unix_secs: candidate
|
||||
.finished_at_unix_secs
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.finished_at_unix_secs)),
|
||||
};
|
||||
|
||||
by_id.insert(stored.id.clone(), stored.clone());
|
||||
Ok(stored)
|
||||
}
|
||||
|
||||
async fn delete_created_before(
|
||||
&self,
|
||||
created_before_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let mut by_id = self
|
||||
.by_id
|
||||
.write()
|
||||
.expect("request candidate repository lock");
|
||||
let mut ids = by_id
|
||||
.values()
|
||||
.filter(|row| row.created_at_unix_secs < created_before_unix_secs)
|
||||
.map(|row| (row.created_at_unix_secs, row.id.clone()))
|
||||
.collect::<Vec<_>>();
|
||||
ids.sort_by(|left, right| left.cmp(right));
|
||||
|
||||
let mut deleted = 0usize;
|
||||
for (_, id) in ids.into_iter().take(limit) {
|
||||
if by_id.remove(&id).is_some() {
|
||||
deleted += 1;
|
||||
}
|
||||
}
|
||||
Ok(deleted)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryRequestCandidateRepository;
|
||||
use crate::repository::candidates::{
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, StoredRequestCandidate,
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
|
||||
StoredRequestCandidate, UpsertRequestCandidateRecord,
|
||||
};
|
||||
|
||||
fn sample_candidate(
|
||||
@@ -145,4 +476,133 @@ mod tests {
|
||||
assert_eq!(rows[0].id, "cand-2");
|
||||
assert_eq!(rows[1].id, "cand-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn aggregates_finalized_health_data_by_endpoint_ids() {
|
||||
let repository = InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_candidate("cand-1", "req-1", 100),
|
||||
sample_candidate("cand-2", "req-2", 200),
|
||||
]);
|
||||
|
||||
let counts = repository
|
||||
.count_finalized_statuses_by_endpoint_ids_since(&["endpoint-1".to_string()], 0)
|
||||
.await
|
||||
.expect("count should succeed");
|
||||
assert_eq!(counts.len(), 1);
|
||||
assert_eq!(counts[0].endpoint_id, "endpoint-1");
|
||||
assert_eq!(counts[0].status, RequestCandidateStatus::Success);
|
||||
assert_eq!(counts[0].count, 2);
|
||||
|
||||
let timeline = repository
|
||||
.aggregate_finalized_timeline_by_endpoint_ids_since(
|
||||
&["endpoint-1".to_string()],
|
||||
0,
|
||||
300,
|
||||
3,
|
||||
)
|
||||
.await
|
||||
.expect("timeline should succeed");
|
||||
assert_eq!(timeline.len(), 2);
|
||||
|
||||
let attempts = repository
|
||||
.list_finalized_by_endpoint_ids_since(&["endpoint-1".to_string()], 0, 1)
|
||||
.await
|
||||
.expect("attempt list should succeed");
|
||||
assert_eq!(attempts.len(), 1);
|
||||
assert_eq!(attempts[0].id, "cand-2");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upsert_writes_and_updates_request_candidate() {
|
||||
let repository = InMemoryRequestCandidateRepository::default();
|
||||
let created = repository
|
||||
.upsert(UpsertRequestCandidateRecord {
|
||||
id: "cand-1".to_string(),
|
||||
request_id: "req-1".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("api-key-1".to_string()),
|
||||
username: Some("alice".to_string()),
|
||||
api_key_name: Some("default".to_string()),
|
||||
candidate_index: 0,
|
||||
retry_index: 0,
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
endpoint_id: Some("endpoint-1".to_string()),
|
||||
key_id: Some("key-1".to_string()),
|
||||
status: RequestCandidateStatus::Available,
|
||||
skip_reason: None,
|
||||
is_cached: Some(false),
|
||||
status_code: None,
|
||||
error_type: None,
|
||||
error_message: None,
|
||||
latency_ms: None,
|
||||
concurrent_requests: None,
|
||||
extra_data: None,
|
||||
required_capabilities: None,
|
||||
created_at_unix_secs: Some(100),
|
||||
started_at_unix_secs: None,
|
||||
finished_at_unix_secs: None,
|
||||
})
|
||||
.await
|
||||
.expect("create should succeed");
|
||||
assert_eq!(created.id, "cand-1");
|
||||
assert_eq!(created.status, RequestCandidateStatus::Available);
|
||||
|
||||
let updated = repository
|
||||
.upsert(UpsertRequestCandidateRecord {
|
||||
id: "cand-1-replacement".to_string(),
|
||||
request_id: "req-1".to_string(),
|
||||
user_id: None,
|
||||
api_key_id: None,
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
candidate_index: 0,
|
||||
retry_index: 0,
|
||||
provider_id: None,
|
||||
endpoint_id: None,
|
||||
key_id: None,
|
||||
status: RequestCandidateStatus::Success,
|
||||
skip_reason: None,
|
||||
is_cached: None,
|
||||
status_code: Some(200),
|
||||
error_type: None,
|
||||
error_message: None,
|
||||
latency_ms: Some(25),
|
||||
concurrent_requests: Some(2),
|
||||
extra_data: None,
|
||||
required_capabilities: None,
|
||||
created_at_unix_secs: None,
|
||||
started_at_unix_secs: Some(101),
|
||||
finished_at_unix_secs: Some(102),
|
||||
})
|
||||
.await
|
||||
.expect("update should succeed");
|
||||
assert_eq!(updated.id, "cand-1");
|
||||
assert_eq!(updated.status, RequestCandidateStatus::Success);
|
||||
assert_eq!(updated.status_code, Some(200));
|
||||
assert_eq!(updated.latency_ms, Some(25));
|
||||
assert_eq!(updated.started_at_unix_secs, Some(101));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn delete_created_before_removes_oldest_matching_rows_up_to_limit() {
|
||||
let repository = InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_candidate("cand-1", "req-1", 100),
|
||||
sample_candidate("cand-2", "req-2", 200),
|
||||
sample_candidate("cand-3", "req-3", 400),
|
||||
]);
|
||||
|
||||
let deleted = repository
|
||||
.delete_created_before(350, 1)
|
||||
.await
|
||||
.expect("delete should succeed");
|
||||
assert_eq!(deleted, 1);
|
||||
|
||||
let rows = repository
|
||||
.list_recent(10)
|
||||
.await
|
||||
.expect("list recent should succeed");
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(rows[0].id, "cand-3");
|
||||
assert_eq!(rows[1].id, "cand-2");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ mod types;
|
||||
pub use memory::InMemoryRequestCandidateRepository;
|
||||
pub use sql::SqlxRequestCandidateReadRepository;
|
||||
pub use types::{
|
||||
RequestCandidateReadRepository, RequestCandidateRepository, RequestCandidateStatus,
|
||||
StoredRequestCandidate,
|
||||
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
|
||||
RequestCandidateRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
|
||||
StoredRequestCandidate, UpsertRequestCandidateRecord,
|
||||
};
|
||||
|
||||
@@ -1,9 +1,14 @@
|
||||
use async_trait::async_trait;
|
||||
use futures_util::future::BoxFuture;
|
||||
use sqlx::{PgPool, Row};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::types::{
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, StoredRequestCandidate,
|
||||
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
|
||||
RequestCandidateStatus, RequestCandidateWriteRepository, StoredRequestCandidate,
|
||||
UpsertRequestCandidateRecord,
|
||||
};
|
||||
use crate::postgres::PostgresTransactionRunner;
|
||||
use crate::DataLayerError;
|
||||
|
||||
const LIST_BY_REQUEST_ID_SQL: &str = r#"
|
||||
@@ -68,20 +73,235 @@ ORDER BY created_at DESC
|
||||
LIMIT $1
|
||||
"#;
|
||||
|
||||
const LIST_BY_PROVIDER_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
status,
|
||||
skip_reason,
|
||||
is_cached,
|
||||
status_code,
|
||||
error_type,
|
||||
error_message,
|
||||
latency_ms,
|
||||
concurrent_requests,
|
||||
extra_data,
|
||||
required_capabilities,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM started_at) AS BIGINT) AS started_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finished_at) AS BIGINT) AS finished_at_unix_secs
|
||||
FROM request_candidates
|
||||
WHERE provider_id = $1
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $2
|
||||
"#;
|
||||
|
||||
const LIST_FINALIZED_BY_ENDPOINT_IDS_SINCE_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
status,
|
||||
skip_reason,
|
||||
is_cached,
|
||||
status_code,
|
||||
error_type,
|
||||
error_message,
|
||||
latency_ms,
|
||||
concurrent_requests,
|
||||
extra_data,
|
||||
required_capabilities,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM started_at) AS BIGINT) AS started_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finished_at) AS BIGINT) AS finished_at_unix_secs
|
||||
FROM request_candidates
|
||||
WHERE endpoint_id = ANY($1)
|
||||
AND created_at >= TO_TIMESTAMP($2)
|
||||
AND status IN ('success', 'failed', 'skipped')
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $3
|
||||
"#;
|
||||
|
||||
const COUNT_FINALIZED_STATUSES_BY_ENDPOINT_IDS_SINCE_SQL: &str = r#"
|
||||
SELECT
|
||||
endpoint_id,
|
||||
status,
|
||||
COUNT(id) AS count
|
||||
FROM request_candidates
|
||||
WHERE endpoint_id = ANY($1)
|
||||
AND created_at >= TO_TIMESTAMP($2)
|
||||
AND status IN ('success', 'failed', 'skipped')
|
||||
GROUP BY endpoint_id, status
|
||||
"#;
|
||||
|
||||
const AGGREGATE_FINALIZED_TIMELINE_BY_ENDPOINT_IDS_SINCE_SQL: &str = r#"
|
||||
SELECT
|
||||
endpoint_id,
|
||||
FLOOR(EXTRACT(EPOCH FROM (created_at - TO_TIMESTAMP($2))) / $4)::BIGINT AS segment_idx,
|
||||
COUNT(id) AS total_count,
|
||||
SUM(CASE WHEN status = 'success' THEN 1 ELSE 0 END) AS success_count,
|
||||
SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END) AS failed_count,
|
||||
CAST(EXTRACT(EPOCH FROM MIN(created_at)) AS BIGINT) AS min_created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM MAX(created_at)) AS BIGINT) AS max_created_at_unix_secs
|
||||
FROM request_candidates
|
||||
WHERE endpoint_id = ANY($1)
|
||||
AND created_at >= TO_TIMESTAMP($2)
|
||||
AND created_at <= TO_TIMESTAMP($3)
|
||||
AND status IN ('success', 'failed', 'skipped')
|
||||
GROUP BY
|
||||
endpoint_id,
|
||||
FLOOR(EXTRACT(EPOCH FROM (created_at - TO_TIMESTAMP($2))) / $4)::BIGINT
|
||||
"#;
|
||||
|
||||
const UPSERT_SQL: &str = r#"
|
||||
INSERT INTO request_candidates (
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
status,
|
||||
skip_reason,
|
||||
is_cached,
|
||||
status_code,
|
||||
error_type,
|
||||
error_message,
|
||||
latency_ms,
|
||||
concurrent_requests,
|
||||
extra_data,
|
||||
required_capabilities,
|
||||
created_at,
|
||||
started_at,
|
||||
finished_at
|
||||
)
|
||||
VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4,
|
||||
$5,
|
||||
$6,
|
||||
$7,
|
||||
$8,
|
||||
$9,
|
||||
$10,
|
||||
$11,
|
||||
$12,
|
||||
$13,
|
||||
COALESCE($14, false),
|
||||
$15,
|
||||
$16,
|
||||
$17,
|
||||
$18,
|
||||
$19,
|
||||
$20,
|
||||
$21,
|
||||
TO_TIMESTAMP(COALESCE($22, 0)),
|
||||
TO_TIMESTAMP($23),
|
||||
TO_TIMESTAMP($24)
|
||||
)
|
||||
ON CONFLICT (request_id, candidate_index, retry_index)
|
||||
DO UPDATE SET
|
||||
user_id = COALESCE(EXCLUDED.user_id, request_candidates.user_id),
|
||||
api_key_id = COALESCE(EXCLUDED.api_key_id, request_candidates.api_key_id),
|
||||
username = COALESCE(EXCLUDED.username, request_candidates.username),
|
||||
api_key_name = COALESCE(EXCLUDED.api_key_name, request_candidates.api_key_name),
|
||||
provider_id = COALESCE(EXCLUDED.provider_id, request_candidates.provider_id),
|
||||
endpoint_id = COALESCE(EXCLUDED.endpoint_id, request_candidates.endpoint_id),
|
||||
key_id = COALESCE(EXCLUDED.key_id, request_candidates.key_id),
|
||||
status = EXCLUDED.status,
|
||||
skip_reason = COALESCE(EXCLUDED.skip_reason, request_candidates.skip_reason),
|
||||
is_cached = COALESCE($14, request_candidates.is_cached),
|
||||
status_code = COALESCE(EXCLUDED.status_code, request_candidates.status_code),
|
||||
error_type = COALESCE(EXCLUDED.error_type, request_candidates.error_type),
|
||||
error_message = COALESCE(EXCLUDED.error_message, request_candidates.error_message),
|
||||
latency_ms = COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms),
|
||||
concurrent_requests = COALESCE(EXCLUDED.concurrent_requests, request_candidates.concurrent_requests),
|
||||
extra_data = COALESCE(EXCLUDED.extra_data, request_candidates.extra_data),
|
||||
required_capabilities = COALESCE(EXCLUDED.required_capabilities, request_candidates.required_capabilities),
|
||||
started_at = COALESCE(EXCLUDED.started_at, request_candidates.started_at),
|
||||
finished_at = COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
|
||||
RETURNING
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
status,
|
||||
skip_reason,
|
||||
is_cached,
|
||||
status_code,
|
||||
error_type,
|
||||
error_message,
|
||||
latency_ms,
|
||||
concurrent_requests,
|
||||
extra_data,
|
||||
required_capabilities,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM started_at) AS BIGINT) AS started_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finished_at) AS BIGINT) AS finished_at_unix_secs
|
||||
"#;
|
||||
|
||||
const DELETE_CREATED_BEFORE_SQL: &str = r#"
|
||||
DELETE FROM request_candidates
|
||||
WHERE id IN (
|
||||
SELECT id
|
||||
FROM request_candidates
|
||||
WHERE created_at < TO_TIMESTAMP($1)
|
||||
ORDER BY created_at ASC, id ASC
|
||||
LIMIT $2
|
||||
)
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxRequestCandidateReadRepository {
|
||||
pool: PgPool,
|
||||
tx_runner: PostgresTransactionRunner,
|
||||
}
|
||||
|
||||
impl SqlxRequestCandidateReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
let tx_runner = PostgresTransactionRunner::new(pool.clone());
|
||||
Self { pool, tx_runner }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub fn transaction_runner(&self) -> &PostgresTransactionRunner {
|
||||
&self.tx_runner
|
||||
}
|
||||
|
||||
pub async fn list_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
@@ -111,6 +331,240 @@ impl SqlxRequestCandidateReadRepository {
|
||||
.await?;
|
||||
rows.iter().map(map_request_candidate_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_by_provider_id(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let limit_value = i64::try_from(limit).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"invalid provider request candidate limit: {limit}"
|
||||
))
|
||||
})?;
|
||||
|
||||
let rows = sqlx::query(LIST_BY_PROVIDER_ID_SQL)
|
||||
.bind(provider_id)
|
||||
.bind(limit_value)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_request_candidate_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_finalized_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
if endpoint_ids.is_empty() || limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = sqlx::query(LIST_FINALIZED_BY_ENDPOINT_IDS_SINCE_SQL)
|
||||
.bind(endpoint_ids)
|
||||
.bind(since_unix_secs as f64)
|
||||
.bind(i64::try_from(limit).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"invalid finalized request candidate limit: {limit}"
|
||||
))
|
||||
})?)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_request_candidate_row).collect()
|
||||
}
|
||||
|
||||
pub async fn count_finalized_statuses_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
) -> Result<Vec<PublicHealthStatusCount>, DataLayerError> {
|
||||
if endpoint_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = sqlx::query(COUNT_FINALIZED_STATUSES_BY_ENDPOINT_IDS_SINCE_SQL)
|
||||
.bind(endpoint_ids)
|
||||
.bind(since_unix_secs as f64)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter()
|
||||
.map(|row| {
|
||||
let status = RequestCandidateStatus::from_database(
|
||||
row.try_get::<String, _>("status")?.as_str(),
|
||||
)?;
|
||||
Ok(PublicHealthStatusCount {
|
||||
endpoint_id: row.try_get("endpoint_id")?,
|
||||
status,
|
||||
count: u64::try_from(row.try_get::<i64, _>("count")?).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"public health status count out of range".to_string(),
|
||||
)
|
||||
})?,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn aggregate_finalized_timeline_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
until_unix_secs: u64,
|
||||
segments: u32,
|
||||
) -> Result<Vec<PublicHealthTimelineBucket>, DataLayerError> {
|
||||
if endpoint_ids.is_empty() || segments == 0 || until_unix_secs < since_unix_secs {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let span_seconds = until_unix_secs.saturating_sub(since_unix_secs);
|
||||
let segment_seconds = if span_seconds == 0 {
|
||||
1.0
|
||||
} else {
|
||||
(span_seconds as f64) / (segments as f64)
|
||||
};
|
||||
|
||||
let rows = sqlx::query(AGGREGATE_FINALIZED_TIMELINE_BY_ENDPOINT_IDS_SINCE_SQL)
|
||||
.bind(endpoint_ids)
|
||||
.bind(since_unix_secs as f64)
|
||||
.bind(until_unix_secs as f64)
|
||||
.bind(segment_seconds)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter()
|
||||
.map(|row| {
|
||||
let raw_segment_idx = row.try_get::<i64, _>("segment_idx")?;
|
||||
let segment_idx = if raw_segment_idx < 0 {
|
||||
0
|
||||
} else {
|
||||
u32::try_from(raw_segment_idx).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"public health segment idx out of range: {raw_segment_idx}"
|
||||
))
|
||||
})?
|
||||
}
|
||||
.min(segments.saturating_sub(1));
|
||||
|
||||
Ok(PublicHealthTimelineBucket {
|
||||
endpoint_id: row.try_get("endpoint_id")?,
|
||||
segment_idx,
|
||||
total_count: u64::try_from(row.try_get::<i64, _>("total_count")?).map_err(
|
||||
|_| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"public health total_count out of range".to_string(),
|
||||
)
|
||||
},
|
||||
)?,
|
||||
success_count: u64::try_from(row.try_get::<i64, _>("success_count")?).map_err(
|
||||
|_| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"public health success_count out of range".to_string(),
|
||||
)
|
||||
},
|
||||
)?,
|
||||
failed_count: u64::try_from(row.try_get::<i64, _>("failed_count")?).map_err(
|
||||
|_| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"public health failed_count out of range".to_string(),
|
||||
)
|
||||
},
|
||||
)?,
|
||||
min_created_at_unix_secs: row
|
||||
.try_get::<Option<i64>, _>("min_created_at_unix_secs")?
|
||||
.map(|value| {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"public health min_created_at_unix_secs out of range: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?,
|
||||
max_created_at_unix_secs: row
|
||||
.try_get::<Option<i64>, _>("max_created_at_unix_secs")?
|
||||
.map(|value| {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"public health max_created_at_unix_secs out of range: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn upsert(
|
||||
&self,
|
||||
candidate: UpsertRequestCandidateRecord,
|
||||
) -> Result<StoredRequestCandidate, DataLayerError> {
|
||||
candidate.validate()?;
|
||||
self.tx_runner
|
||||
.run_read_write(|tx| {
|
||||
Box::pin(async move {
|
||||
let row = sqlx::query(UPSERT_SQL)
|
||||
.bind(if candidate.id.trim().is_empty() {
|
||||
Uuid::new_v4().to_string()
|
||||
} else {
|
||||
candidate.id.clone()
|
||||
})
|
||||
.bind(&candidate.request_id)
|
||||
.bind(&candidate.user_id)
|
||||
.bind(&candidate.api_key_id)
|
||||
.bind(&candidate.username)
|
||||
.bind(&candidate.api_key_name)
|
||||
.bind(to_i32(candidate.candidate_index)?)
|
||||
.bind(to_i32(candidate.retry_index)?)
|
||||
.bind(&candidate.provider_id)
|
||||
.bind(&candidate.endpoint_id)
|
||||
.bind(&candidate.key_id)
|
||||
.bind(status_to_database(candidate.status))
|
||||
.bind(&candidate.skip_reason)
|
||||
.bind(candidate.is_cached)
|
||||
.bind(candidate.status_code.map(i32::from))
|
||||
.bind(&candidate.error_type)
|
||||
.bind(&candidate.error_message)
|
||||
.bind(candidate.latency_ms.map(to_i32_u64).transpose()?)
|
||||
.bind(candidate.concurrent_requests.map(to_i32).transpose()?)
|
||||
.bind(&candidate.extra_data)
|
||||
.bind(&candidate.required_capabilities)
|
||||
.bind(candidate.created_at_unix_secs.map(|value| value as f64))
|
||||
.bind(candidate.started_at_unix_secs.map(|value| value as f64))
|
||||
.bind(candidate.finished_at_unix_secs.map(|value| value as f64))
|
||||
.fetch_one(&mut **tx)
|
||||
.await?;
|
||||
map_request_candidate_row(&row)
|
||||
}) as BoxFuture<'_, Result<StoredRequestCandidate, DataLayerError>>
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn delete_created_before(
|
||||
&self,
|
||||
created_before_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let result = sqlx::query(DELETE_CREATED_BEFORE_SQL)
|
||||
.bind(created_before_unix_secs as f64)
|
||||
.bind(i64::try_from(limit).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request candidate delete limit: {limit}"
|
||||
))
|
||||
})?)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected() as usize)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -128,6 +582,67 @@ impl RequestCandidateReadRepository for SqlxRequestCandidateReadRepository {
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
Self::list_recent(self, limit).await
|
||||
}
|
||||
|
||||
async fn list_finalized_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
Self::list_finalized_by_endpoint_ids_since(self, endpoint_ids, since_unix_secs, limit).await
|
||||
}
|
||||
|
||||
async fn list_by_provider_id(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
Self::list_by_provider_id(self, provider_id, limit).await
|
||||
}
|
||||
|
||||
async fn count_finalized_statuses_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
) -> Result<Vec<PublicHealthStatusCount>, DataLayerError> {
|
||||
Self::count_finalized_statuses_by_endpoint_ids_since(self, endpoint_ids, since_unix_secs)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn aggregate_finalized_timeline_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
until_unix_secs: u64,
|
||||
segments: u32,
|
||||
) -> Result<Vec<PublicHealthTimelineBucket>, DataLayerError> {
|
||||
Self::aggregate_finalized_timeline_by_endpoint_ids_since(
|
||||
self,
|
||||
endpoint_ids,
|
||||
since_unix_secs,
|
||||
until_unix_secs,
|
||||
segments,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RequestCandidateWriteRepository for SqlxRequestCandidateReadRepository {
|
||||
async fn upsert(
|
||||
&self,
|
||||
candidate: UpsertRequestCandidateRecord,
|
||||
) -> Result<StoredRequestCandidate, DataLayerError> {
|
||||
Self::upsert(self, candidate).await
|
||||
}
|
||||
|
||||
async fn delete_created_before(
|
||||
&self,
|
||||
created_before_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
Self::delete_created_before(self, created_before_unix_secs, limit).await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_request_candidate_row(
|
||||
@@ -163,6 +678,31 @@ fn map_request_candidate_row(
|
||||
)
|
||||
}
|
||||
|
||||
fn status_to_database(status: RequestCandidateStatus) -> &'static str {
|
||||
match status {
|
||||
RequestCandidateStatus::Available => "available",
|
||||
RequestCandidateStatus::Unused => "unused",
|
||||
RequestCandidateStatus::Pending => "pending",
|
||||
RequestCandidateStatus::Streaming => "streaming",
|
||||
RequestCandidateStatus::Success => "success",
|
||||
RequestCandidateStatus::Failed => "failed",
|
||||
RequestCandidateStatus::Cancelled => "cancelled",
|
||||
RequestCandidateStatus::Skipped => "skipped",
|
||||
}
|
||||
}
|
||||
|
||||
fn to_i32(value: u32) -> Result<i32, DataLayerError> {
|
||||
i32::try_from(value).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!("request candidate value out of range: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn to_i32_u64(value: u64) -> Result<i32, DataLayerError> {
|
||||
i32::try_from(value).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!("request candidate value out of range: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxRequestCandidateReadRepository;
|
||||
@@ -185,5 +725,6 @@ mod tests {
|
||||
let pool = factory.connect_lazy().expect("pool should build");
|
||||
let repository = SqlxRequestCandidateReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
let _ = repository.transaction_runner();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -185,6 +185,24 @@ impl StoredRequestCandidate {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct PublicHealthStatusCount {
|
||||
pub endpoint_id: String,
|
||||
pub status: RequestCandidateStatus,
|
||||
pub count: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct PublicHealthTimelineBucket {
|
||||
pub endpoint_id: String,
|
||||
pub segment_idx: u32,
|
||||
pub total_count: u64,
|
||||
pub success_count: u64,
|
||||
pub failed_count: u64,
|
||||
pub min_created_at_unix_secs: Option<u64>,
|
||||
pub max_created_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait RequestCandidateReadRepository: Send + Sync {
|
||||
async fn list_by_request_id(
|
||||
@@ -196,15 +214,106 @@ pub trait RequestCandidateReadRepository: Send + Sync {
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
|
||||
|
||||
async fn list_by_provider_id(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
|
||||
|
||||
async fn list_finalized_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
|
||||
|
||||
async fn count_finalized_statuses_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
) -> Result<Vec<PublicHealthStatusCount>, crate::DataLayerError>;
|
||||
|
||||
async fn aggregate_finalized_timeline_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
until_unix_secs: u64,
|
||||
segments: u32,
|
||||
) -> Result<Vec<PublicHealthTimelineBucket>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait RequestCandidateRepository: RequestCandidateReadRepository + Send + Sync {}
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct UpsertRequestCandidateRecord {
|
||||
pub id: String,
|
||||
pub request_id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub api_key_id: Option<String>,
|
||||
pub username: Option<String>,
|
||||
pub api_key_name: Option<String>,
|
||||
pub candidate_index: u32,
|
||||
pub retry_index: u32,
|
||||
pub provider_id: Option<String>,
|
||||
pub endpoint_id: Option<String>,
|
||||
pub key_id: Option<String>,
|
||||
pub status: RequestCandidateStatus,
|
||||
pub skip_reason: Option<String>,
|
||||
pub is_cached: Option<bool>,
|
||||
pub status_code: Option<u16>,
|
||||
pub error_type: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
pub latency_ms: Option<u64>,
|
||||
pub concurrent_requests: Option<u32>,
|
||||
pub extra_data: Option<serde_json::Value>,
|
||||
pub required_capabilities: Option<serde_json::Value>,
|
||||
pub created_at_unix_secs: Option<u64>,
|
||||
pub started_at_unix_secs: Option<u64>,
|
||||
pub finished_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl<T> RequestCandidateRepository for T where T: RequestCandidateReadRepository + Send + Sync {}
|
||||
impl UpsertRequestCandidateRecord {
|
||||
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||
if self.id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"request candidate upsert id cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.request_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"request candidate upsert request_id cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait RequestCandidateWriteRepository: Send + Sync {
|
||||
async fn upsert(
|
||||
&self,
|
||||
candidate: UpsertRequestCandidateRecord,
|
||||
) -> Result<StoredRequestCandidate, crate::DataLayerError>;
|
||||
|
||||
async fn delete_created_before(
|
||||
&self,
|
||||
created_before_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<usize, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait RequestCandidateRepository:
|
||||
RequestCandidateReadRepository + RequestCandidateWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> RequestCandidateRepository for T where
|
||||
T: RequestCandidateReadRepository + RequestCandidateWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{RequestCandidateStatus, StoredRequestCandidate};
|
||||
use super::{RequestCandidateStatus, StoredRequestCandidate, UpsertRequestCandidateRecord};
|
||||
|
||||
#[test]
|
||||
fn parses_status_from_database_text() {
|
||||
@@ -286,4 +395,36 @@ mod tests {
|
||||
assert!(!RequestCandidateStatus::Pending.is_attempted(None));
|
||||
assert!(RequestCandidateStatus::Pending.is_attempted(Some(1)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_upsert_payload() {
|
||||
assert!(UpsertRequestCandidateRecord {
|
||||
id: "".to_string(),
|
||||
request_id: "".to_string(),
|
||||
user_id: None,
|
||||
api_key_id: None,
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
candidate_index: 0,
|
||||
retry_index: 0,
|
||||
provider_id: None,
|
||||
endpoint_id: None,
|
||||
key_id: None,
|
||||
status: RequestCandidateStatus::Available,
|
||||
skip_reason: None,
|
||||
is_cached: None,
|
||||
status_code: None,
|
||||
error_type: None,
|
||||
error_message: None,
|
||||
latency_ms: None,
|
||||
concurrent_requests: None,
|
||||
extra_data: None,
|
||||
required_capabilities: None,
|
||||
created_at_unix_secs: None,
|
||||
started_at_unix_secs: None,
|
||||
finished_at_unix_secs: None,
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
}
|
||||
}
|
||||
|
||||
321
crates/aether-data/src/repository/gemini_file_mappings/memory.rs
Normal file
321
crates/aether-data/src/repository/gemini_file_mappings/memory.rs
Normal file
@@ -0,0 +1,321 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
GeminiFileMappingListQuery, GeminiFileMappingMimeTypeCount, GeminiFileMappingReadRepository,
|
||||
GeminiFileMappingStats, GeminiFileMappingWriteRepository, StoredGeminiFileMapping,
|
||||
StoredGeminiFileMappingListPage, UpsertGeminiFileMappingRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct InMemoryGeminiFileMappingRepository {
|
||||
by_file: RwLock<BTreeMap<String, StoredGeminiFileMapping>>,
|
||||
}
|
||||
|
||||
impl InMemoryGeminiFileMappingRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredGeminiFileMapping>,
|
||||
{
|
||||
let mut by_file = BTreeMap::new();
|
||||
for item in items {
|
||||
by_file.insert(item.file_name.clone(), item);
|
||||
}
|
||||
Self {
|
||||
by_file: RwLock::new(by_file),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl GeminiFileMappingReadRepository for InMemoryGeminiFileMappingRepository {
|
||||
async fn find_by_file_name(
|
||||
&self,
|
||||
file_name: &str,
|
||||
) -> Result<Option<StoredGeminiFileMapping>, DataLayerError> {
|
||||
let guard = self.by_file.read().expect("gemini mapping repository lock");
|
||||
Ok(guard.get(file_name).cloned())
|
||||
}
|
||||
|
||||
async fn list_mappings(
|
||||
&self,
|
||||
query: &GeminiFileMappingListQuery,
|
||||
) -> Result<StoredGeminiFileMappingListPage, DataLayerError> {
|
||||
let guard = self.by_file.read().expect("gemini mapping repository lock");
|
||||
let search = query
|
||||
.search
|
||||
.as_deref()
|
||||
.map(|value| value.to_ascii_lowercase());
|
||||
let mut items = guard
|
||||
.values()
|
||||
.filter(|item| query.include_expired || item.expires_at_unix_secs > query.now_unix_secs)
|
||||
.filter(|item| {
|
||||
search.as_deref().is_none_or(|needle| {
|
||||
item.file_name.to_ascii_lowercase().contains(needle)
|
||||
|| item
|
||||
.display_name
|
||||
.as_deref()
|
||||
.map(|value| value.to_ascii_lowercase().contains(needle))
|
||||
.unwrap_or(false)
|
||||
})
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
items.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.cmp(&left.created_at_unix_secs)
|
||||
.then_with(|| left.file_name.cmp(&right.file_name))
|
||||
});
|
||||
let total = items.len();
|
||||
let page_items = items
|
||||
.into_iter()
|
||||
.skip(query.offset)
|
||||
.take(query.limit)
|
||||
.collect::<Vec<_>>();
|
||||
Ok(StoredGeminiFileMappingListPage {
|
||||
items: page_items,
|
||||
total,
|
||||
})
|
||||
}
|
||||
|
||||
async fn summarize_mappings(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<GeminiFileMappingStats, DataLayerError> {
|
||||
let guard = self.by_file.read().expect("gemini mapping repository lock");
|
||||
let total_mappings = guard.len();
|
||||
let active_items = guard
|
||||
.values()
|
||||
.filter(|item| item.expires_at_unix_secs > now_unix_secs)
|
||||
.collect::<Vec<_>>();
|
||||
let active_mappings = active_items.len();
|
||||
let expired_mappings = total_mappings.saturating_sub(active_mappings);
|
||||
let mut by_mime_type = BTreeMap::<String, usize>::new();
|
||||
for item in active_items {
|
||||
let mime_type = item
|
||||
.mime_type
|
||||
.clone()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
*by_mime_type.entry(mime_type).or_default() += 1;
|
||||
}
|
||||
Ok(GeminiFileMappingStats {
|
||||
total_mappings,
|
||||
active_mappings,
|
||||
expired_mappings,
|
||||
by_mime_type: by_mime_type
|
||||
.into_iter()
|
||||
.map(|(mime_type, count)| GeminiFileMappingMimeTypeCount { mime_type, count })
|
||||
.collect(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl GeminiFileMappingWriteRepository for InMemoryGeminiFileMappingRepository {
|
||||
async fn upsert(
|
||||
&self,
|
||||
record: UpsertGeminiFileMappingRecord,
|
||||
) -> Result<StoredGeminiFileMapping, DataLayerError> {
|
||||
record.validate()?;
|
||||
let mut guard = self
|
||||
.by_file
|
||||
.write()
|
||||
.expect("gemini mapping repository lock");
|
||||
let created_at_unix_secs = guard
|
||||
.get(&record.file_name)
|
||||
.map(|existing| existing.created_at_unix_secs)
|
||||
.unwrap_or_else(current_unix_secs);
|
||||
let mapping = StoredGeminiFileMapping {
|
||||
id: record.id.clone(),
|
||||
file_name: record.file_name.clone(),
|
||||
key_id: record.key_id.clone(),
|
||||
user_id: record.user_id.clone(),
|
||||
display_name: record.display_name.clone(),
|
||||
mime_type: record.mime_type.clone(),
|
||||
source_hash: record.source_hash.clone(),
|
||||
created_at_unix_secs,
|
||||
expires_at_unix_secs: record.expires_at_unix_secs,
|
||||
};
|
||||
guard.insert(record.file_name.clone(), mapping.clone());
|
||||
Ok(mapping)
|
||||
}
|
||||
|
||||
async fn delete_by_file_name(&self, file_name: &str) -> Result<bool, DataLayerError> {
|
||||
let mut guard = self
|
||||
.by_file
|
||||
.write()
|
||||
.expect("gemini mapping repository lock");
|
||||
Ok(guard.remove(file_name).is_some())
|
||||
}
|
||||
|
||||
async fn delete_by_id(
|
||||
&self,
|
||||
mapping_id: &str,
|
||||
) -> Result<Option<StoredGeminiFileMapping>, DataLayerError> {
|
||||
let mut guard = self
|
||||
.by_file
|
||||
.write()
|
||||
.expect("gemini mapping repository lock");
|
||||
let Some(file_name) = guard
|
||||
.iter()
|
||||
.find_map(|(file_name, item)| (item.id == mapping_id).then(|| file_name.clone()))
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(guard.remove(&file_name))
|
||||
}
|
||||
|
||||
async fn delete_expired_before(&self, now_unix_secs: u64) -> Result<usize, DataLayerError> {
|
||||
let mut guard = self
|
||||
.by_file
|
||||
.write()
|
||||
.expect("gemini mapping repository lock");
|
||||
let before = guard.len();
|
||||
guard.retain(|_, item| item.expires_at_unix_secs > now_unix_secs);
|
||||
Ok(before.saturating_sub(guard.len()))
|
||||
}
|
||||
}
|
||||
|
||||
fn current_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::repository::gemini_file_mappings::{
|
||||
GeminiFileMappingListQuery, GeminiFileMappingReadRepository,
|
||||
GeminiFileMappingWriteRepository,
|
||||
};
|
||||
|
||||
use super::{InMemoryGeminiFileMappingRepository, UpsertGeminiFileMappingRecord};
|
||||
use crate::DataLayerError;
|
||||
|
||||
fn sample_record(id: &str, file_name: &str) -> UpsertGeminiFileMappingRecord {
|
||||
UpsertGeminiFileMappingRecord {
|
||||
id: id.to_string(),
|
||||
file_name: file_name.to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
display_name: Some("display".to_string()),
|
||||
mime_type: Some("image/png".to_string()),
|
||||
source_hash: Some("hash-1".to_string()),
|
||||
expires_at_unix_secs: 4_102_444_800,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upsert_and_find() -> Result<(), DataLayerError> {
|
||||
let repo = InMemoryGeminiFileMappingRepository::default();
|
||||
let record = sample_record("id-1", "files/abc");
|
||||
let stored = repo.upsert(record.clone()).await?;
|
||||
assert_eq!(stored.file_name, "files/abc");
|
||||
|
||||
let fetched = repo.find_by_file_name("files/abc").await?;
|
||||
assert_eq!(fetched.unwrap().key_id, "key-1");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn delete_removes_entry() -> Result<(), DataLayerError> {
|
||||
let repo = InMemoryGeminiFileMappingRepository::default();
|
||||
let record = sample_record("id-2", "files/def");
|
||||
let _stored = repo.upsert(record).await?;
|
||||
assert!(repo.find_by_file_name("files/def").await?.is_some());
|
||||
assert!(repo.delete_by_file_name("files/def").await?);
|
||||
assert!(repo.find_by_file_name("files/def").await?.is_none());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upsert_preserves_created_at_when_replacing_existing_file_name(
|
||||
) -> Result<(), DataLayerError> {
|
||||
let repo = InMemoryGeminiFileMappingRepository::default();
|
||||
let first = repo.upsert(sample_record("id-1", "files/same")).await?;
|
||||
let replaced = repo.upsert(sample_record("id-2", "files/same")).await?;
|
||||
|
||||
assert_eq!(replaced.created_at_unix_secs, first.created_at_unix_secs);
|
||||
assert_eq!(replaced.id, "id-2");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_and_summarize_mappings() -> Result<(), DataLayerError> {
|
||||
let repo = InMemoryGeminiFileMappingRepository::seed(vec![
|
||||
repo_item("id-1", "files/alpha", "image/png", 10, 200),
|
||||
repo_item("id-2", "files/beta", "video/mp4", 20, 50),
|
||||
repo_item("id-3", "files/gamma", "", 30, 220),
|
||||
]);
|
||||
|
||||
let page = repo
|
||||
.list_mappings(&GeminiFileMappingListQuery {
|
||||
include_expired: false,
|
||||
search: Some("ga".to_string()),
|
||||
offset: 0,
|
||||
limit: 10,
|
||||
now_unix_secs: 100,
|
||||
})
|
||||
.await?;
|
||||
assert_eq!(page.total, 1);
|
||||
assert_eq!(page.items[0].id, "id-3");
|
||||
|
||||
let stats = repo.summarize_mappings(100).await?;
|
||||
assert_eq!(stats.total_mappings, 3);
|
||||
assert_eq!(stats.active_mappings, 2);
|
||||
assert_eq!(stats.expired_mappings, 1);
|
||||
assert_eq!(stats.by_mime_type.len(), 2);
|
||||
assert_eq!(stats.by_mime_type[0].mime_type, "image/png");
|
||||
assert_eq!(stats.by_mime_type[0].count, 1);
|
||||
assert_eq!(stats.by_mime_type[1].mime_type, "unknown");
|
||||
assert_eq!(stats.by_mime_type[1].count, 1);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn delete_by_id_and_cleanup_expired() -> Result<(), DataLayerError> {
|
||||
let repo = InMemoryGeminiFileMappingRepository::seed(vec![
|
||||
repo_item("id-1", "files/alpha", "image/png", 10, 200),
|
||||
repo_item("id-2", "files/beta", "video/mp4", 20, 50),
|
||||
]);
|
||||
|
||||
let deleted = repo.delete_by_id("id-1").await?;
|
||||
assert_eq!(
|
||||
deleted.as_ref().map(|item| item.file_name.as_str()),
|
||||
Some("files/alpha")
|
||||
);
|
||||
assert!(repo.find_by_file_name("files/alpha").await?.is_none());
|
||||
|
||||
let deleted_count = repo.delete_expired_before(100).await?;
|
||||
assert_eq!(deleted_count, 1);
|
||||
assert!(repo.find_by_file_name("files/beta").await?.is_none());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn repo_item(
|
||||
id: &str,
|
||||
file_name: &str,
|
||||
mime_type: &str,
|
||||
created_at_unix_secs: u64,
|
||||
expires_at_unix_secs: u64,
|
||||
) -> crate::repository::gemini_file_mappings::StoredGeminiFileMapping {
|
||||
crate::repository::gemini_file_mappings::StoredGeminiFileMapping {
|
||||
id: id.to_string(),
|
||||
file_name: file_name.to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
display_name: Some(format!("display-{id}")),
|
||||
mime_type: (!mime_type.is_empty()).then(|| mime_type.to_string()),
|
||||
source_hash: Some(format!("hash-{id}")),
|
||||
created_at_unix_secs,
|
||||
expires_at_unix_secs,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
pub mod memory;
|
||||
pub mod sql;
|
||||
pub mod types;
|
||||
|
||||
pub use memory::InMemoryGeminiFileMappingRepository;
|
||||
pub use sql::SqlxGeminiFileMappingRepository;
|
||||
pub use types::{
|
||||
GeminiFileMappingListQuery, GeminiFileMappingMimeTypeCount, GeminiFileMappingReadRepository,
|
||||
GeminiFileMappingRepository, GeminiFileMappingStats, GeminiFileMappingWriteRepository,
|
||||
StoredGeminiFileMapping, StoredGeminiFileMappingListPage, UpsertGeminiFileMappingRecord,
|
||||
};
|
||||
319
crates/aether-data/src/repository/gemini_file_mappings/sql.rs
Normal file
319
crates/aether-data/src/repository/gemini_file_mappings/sql.rs
Normal file
@@ -0,0 +1,319 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
||||
|
||||
use super::types::{
|
||||
GeminiFileMappingListQuery, GeminiFileMappingMimeTypeCount, GeminiFileMappingReadRepository,
|
||||
GeminiFileMappingStats, GeminiFileMappingWriteRepository, StoredGeminiFileMapping,
|
||||
StoredGeminiFileMappingListPage, UpsertGeminiFileMappingRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxGeminiFileMappingRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxGeminiFileMappingRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
fn map_row(row: &PgRow) -> Result<StoredGeminiFileMapping, DataLayerError> {
|
||||
Ok(StoredGeminiFileMapping {
|
||||
id: row.try_get("id")?,
|
||||
file_name: row.try_get("file_name")?,
|
||||
key_id: row.try_get("key_id")?,
|
||||
user_id: row.try_get("user_id").ok().flatten(),
|
||||
display_name: row.try_get("display_name").ok().flatten(),
|
||||
mime_type: row.try_get("mime_type").ok().flatten(),
|
||||
source_hash: row.try_get("source_hash").ok().flatten(),
|
||||
created_at_unix_secs: u64::try_from(row.try_get::<i64, _>("created_at_unix_secs")?)
|
||||
.map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"gemini_file_mappings.created_at is invalid".to_string(),
|
||||
)
|
||||
})?,
|
||||
expires_at_unix_secs: u64::try_from(row.try_get::<i64, _>("expires_at_unix_secs")?)
|
||||
.map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"gemini_file_mappings.expires_at is invalid".to_string(),
|
||||
)
|
||||
})?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl GeminiFileMappingReadRepository for SqlxGeminiFileMappingRepository {
|
||||
async fn find_by_file_name(
|
||||
&self,
|
||||
file_name: &str,
|
||||
) -> Result<Option<StoredGeminiFileMapping>, DataLayerError> {
|
||||
let row = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
file_name,
|
||||
key_id,
|
||||
user_id,
|
||||
display_name,
|
||||
mime_type,
|
||||
source_hash,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs
|
||||
FROM gemini_file_mappings
|
||||
WHERE file_name = $1
|
||||
"#,
|
||||
)
|
||||
.bind(file_name)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
match row {
|
||||
Some(row) => Ok(Some(Self::map_row(&row)?)),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_mappings(
|
||||
&self,
|
||||
query: &GeminiFileMappingListQuery,
|
||||
) -> Result<StoredGeminiFileMappingListPage, DataLayerError> {
|
||||
let total = build_list_count_query(query)
|
||||
.build_query_scalar::<i64>()
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
let rows = build_list_rows_query(query)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
Ok(StoredGeminiFileMappingListPage {
|
||||
items: rows
|
||||
.iter()
|
||||
.map(Self::map_row)
|
||||
.collect::<Result<Vec<_>, _>>()?,
|
||||
total: usize::try_from(total).unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn summarize_mappings(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<GeminiFileMappingStats, DataLayerError> {
|
||||
let totals = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
COUNT(*)::bigint AS total_mappings,
|
||||
COUNT(*) FILTER (WHERE expires_at > TO_TIMESTAMP($1::double precision))::bigint AS active_mappings
|
||||
FROM gemini_file_mappings
|
||||
"#,
|
||||
)
|
||||
.bind(now_unix_secs as f64)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
let total_mappings =
|
||||
usize::try_from(totals.try_get::<i64, _>("total_mappings")?).unwrap_or_default();
|
||||
let active_mappings =
|
||||
usize::try_from(totals.try_get::<i64, _>("active_mappings")?).unwrap_or_default();
|
||||
let by_mime_type_rows = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
COALESCE(NULLIF(TRIM(mime_type), ''), 'unknown') AS mime_type,
|
||||
COUNT(*)::bigint AS count
|
||||
FROM gemini_file_mappings
|
||||
WHERE expires_at > TO_TIMESTAMP($1::double precision)
|
||||
GROUP BY COALESCE(NULLIF(TRIM(mime_type), ''), 'unknown')
|
||||
ORDER BY mime_type ASC
|
||||
"#,
|
||||
)
|
||||
.bind(now_unix_secs as f64)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
Ok(GeminiFileMappingStats {
|
||||
total_mappings,
|
||||
active_mappings,
|
||||
expired_mappings: total_mappings.saturating_sub(active_mappings),
|
||||
by_mime_type: by_mime_type_rows
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
Ok(GeminiFileMappingMimeTypeCount {
|
||||
mime_type: row.try_get("mime_type")?,
|
||||
count: usize::try_from(row.try_get::<i64, _>("count")?).unwrap_or_default(),
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, DataLayerError>>()?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl GeminiFileMappingWriteRepository for SqlxGeminiFileMappingRepository {
|
||||
async fn upsert(
|
||||
&self,
|
||||
record: UpsertGeminiFileMappingRecord,
|
||||
) -> Result<StoredGeminiFileMapping, DataLayerError> {
|
||||
record.validate()?;
|
||||
let row = sqlx::query(
|
||||
r#"
|
||||
INSERT INTO gemini_file_mappings (
|
||||
id,
|
||||
file_name,
|
||||
key_id,
|
||||
user_id,
|
||||
display_name,
|
||||
mime_type,
|
||||
source_hash,
|
||||
created_at,
|
||||
expires_at
|
||||
)
|
||||
VALUES ($1,$2,$3,$4,$5,$6,$7,NOW(),TO_TIMESTAMP($8::double precision))
|
||||
ON CONFLICT (file_name)
|
||||
DO UPDATE
|
||||
SET
|
||||
key_id = EXCLUDED.key_id,
|
||||
user_id = EXCLUDED.user_id,
|
||||
display_name = EXCLUDED.display_name,
|
||||
mime_type = EXCLUDED.mime_type,
|
||||
source_hash = EXCLUDED.source_hash,
|
||||
expires_at = EXCLUDED.expires_at
|
||||
RETURNING
|
||||
id,
|
||||
file_name,
|
||||
key_id,
|
||||
user_id,
|
||||
display_name,
|
||||
mime_type,
|
||||
source_hash,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs
|
||||
"#,
|
||||
)
|
||||
.bind(record.id.clone())
|
||||
.bind(record.file_name.clone())
|
||||
.bind(record.key_id.clone())
|
||||
.bind(record.user_id.clone())
|
||||
.bind(record.display_name.clone())
|
||||
.bind(record.mime_type.clone())
|
||||
.bind(record.source_hash.clone())
|
||||
.bind(record.expires_at_unix_secs as f64)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
Self::map_row(&row)
|
||||
}
|
||||
|
||||
async fn delete_by_file_name(&self, file_name: &str) -> Result<bool, DataLayerError> {
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
DELETE FROM gemini_file_mappings
|
||||
WHERE file_name = $1
|
||||
#"#,
|
||||
)
|
||||
.bind(file_name)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
async fn delete_by_id(
|
||||
&self,
|
||||
mapping_id: &str,
|
||||
) -> Result<Option<StoredGeminiFileMapping>, DataLayerError> {
|
||||
let row = sqlx::query(
|
||||
r#"
|
||||
DELETE FROM gemini_file_mappings
|
||||
WHERE id = $1
|
||||
RETURNING
|
||||
id,
|
||||
file_name,
|
||||
key_id,
|
||||
user_id,
|
||||
display_name,
|
||||
mime_type,
|
||||
source_hash,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs
|
||||
"#,
|
||||
)
|
||||
.bind(mapping_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
match row {
|
||||
Some(row) => Ok(Some(Self::map_row(&row)?)),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_expired_before(&self, now_unix_secs: u64) -> Result<usize, DataLayerError> {
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
DELETE FROM gemini_file_mappings
|
||||
WHERE expires_at <= TO_TIMESTAMP($1::double precision)
|
||||
"#,
|
||||
)
|
||||
.bind(now_unix_secs as f64)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(usize::try_from(result.rows_affected()).unwrap_or_default())
|
||||
}
|
||||
}
|
||||
|
||||
fn build_list_count_query(query: &GeminiFileMappingListQuery) -> QueryBuilder<'_, Postgres> {
|
||||
let mut builder = QueryBuilder::<Postgres>::new(
|
||||
"SELECT COUNT(*)::bigint AS total FROM gemini_file_mappings WHERE 1=1",
|
||||
);
|
||||
apply_list_filters(&mut builder, query);
|
||||
builder
|
||||
}
|
||||
|
||||
fn build_list_rows_query(query: &GeminiFileMappingListQuery) -> QueryBuilder<'_, Postgres> {
|
||||
let mut builder = QueryBuilder::<Postgres>::new(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
file_name,
|
||||
key_id,
|
||||
user_id,
|
||||
display_name,
|
||||
mime_type,
|
||||
source_hash,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs
|
||||
FROM gemini_file_mappings
|
||||
WHERE 1=1
|
||||
"#,
|
||||
);
|
||||
apply_list_filters(&mut builder, query);
|
||||
builder.push(" ORDER BY created_at DESC, file_name ASC LIMIT ");
|
||||
builder.push_bind(i64::try_from(query.limit).unwrap_or(i64::MAX));
|
||||
builder.push(" OFFSET ");
|
||||
builder.push_bind(i64::try_from(query.offset).unwrap_or(i64::MAX));
|
||||
builder
|
||||
}
|
||||
|
||||
fn apply_list_filters(
|
||||
builder: &mut QueryBuilder<'_, Postgres>,
|
||||
query: &GeminiFileMappingListQuery,
|
||||
) {
|
||||
if !query.include_expired {
|
||||
builder.push(" AND expires_at > TO_TIMESTAMP(");
|
||||
builder.push_bind(query.now_unix_secs as f64);
|
||||
builder.push("::double precision)");
|
||||
}
|
||||
if let Some(search) = query
|
||||
.search
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
let pattern = format!("%{search}%");
|
||||
builder.push(" AND (file_name ILIKE ");
|
||||
builder.push_bind(pattern.clone());
|
||||
builder.push(" OR COALESCE(display_name, '') ILIKE ");
|
||||
builder.push_bind(pattern);
|
||||
builder.push(")");
|
||||
}
|
||||
}
|
||||
165
crates/aether-data/src/repository/gemini_file_mappings/types.rs
Normal file
165
crates/aether-data/src/repository/gemini_file_mappings/types.rs
Normal file
@@ -0,0 +1,165 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct GeminiFileMappingListQuery {
|
||||
pub include_expired: bool,
|
||||
pub search: Option<String>,
|
||||
pub offset: usize,
|
||||
pub limit: usize,
|
||||
pub now_unix_secs: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct StoredGeminiFileMappingListPage {
|
||||
pub items: Vec<StoredGeminiFileMapping>,
|
||||
pub total: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct GeminiFileMappingMimeTypeCount {
|
||||
pub mime_type: String,
|
||||
pub count: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct GeminiFileMappingStats {
|
||||
pub total_mappings: usize,
|
||||
pub active_mappings: usize,
|
||||
pub expired_mappings: usize,
|
||||
pub by_mime_type: Vec<GeminiFileMappingMimeTypeCount>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct StoredGeminiFileMapping {
|
||||
pub id: String,
|
||||
pub file_name: String,
|
||||
pub key_id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub display_name: Option<String>,
|
||||
pub mime_type: Option<String>,
|
||||
pub source_hash: Option<String>,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub expires_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
impl StoredGeminiFileMapping {
|
||||
pub fn new(
|
||||
id: String,
|
||||
file_name: String,
|
||||
key_id: String,
|
||||
created_at_unix_secs: i64,
|
||||
expires_at_unix_secs: i64,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if file_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"gemini_file_mappings.file_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if key_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"gemini_file_mappings.key_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid gemini_file_mappings.created_at: {created_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
let expires_at_unix_secs = u64::try_from(expires_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid gemini_file_mappings.expires_at: {expires_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
Ok(Self {
|
||||
id,
|
||||
file_name,
|
||||
key_id,
|
||||
user_id: None,
|
||||
display_name: None,
|
||||
mime_type: None,
|
||||
source_hash: None,
|
||||
created_at_unix_secs,
|
||||
expires_at_unix_secs,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpsertGeminiFileMappingRecord {
|
||||
pub id: String,
|
||||
pub file_name: String,
|
||||
pub key_id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub display_name: Option<String>,
|
||||
pub mime_type: Option<String>,
|
||||
pub source_hash: Option<String>,
|
||||
pub expires_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
impl UpsertGeminiFileMappingRecord {
|
||||
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||
if self.file_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"gemini_file_mappings.file_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.key_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"gemini_file_mappings.key_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.expires_at_unix_secs == 0 {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"gemini_file_mappings.expires_at is empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait GeminiFileMappingReadRepository: Send + Sync {
|
||||
async fn find_by_file_name(
|
||||
&self,
|
||||
file_name: &str,
|
||||
) -> Result<Option<StoredGeminiFileMapping>, crate::DataLayerError>;
|
||||
|
||||
async fn list_mappings(
|
||||
&self,
|
||||
query: &GeminiFileMappingListQuery,
|
||||
) -> Result<StoredGeminiFileMappingListPage, crate::DataLayerError>;
|
||||
|
||||
async fn summarize_mappings(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<GeminiFileMappingStats, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait GeminiFileMappingWriteRepository: Send + Sync {
|
||||
async fn upsert(
|
||||
&self,
|
||||
record: UpsertGeminiFileMappingRecord,
|
||||
) -> Result<StoredGeminiFileMapping, crate::DataLayerError>;
|
||||
async fn delete_by_file_name(&self, file_name: &str) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn delete_by_id(
|
||||
&self,
|
||||
mapping_id: &str,
|
||||
) -> Result<Option<StoredGeminiFileMapping>, crate::DataLayerError>;
|
||||
|
||||
async fn delete_expired_before(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<usize, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait GeminiFileMappingRepository:
|
||||
GeminiFileMappingReadRepository + GeminiFileMappingWriteRepository
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> GeminiFileMappingRepository for T where
|
||||
T: GeminiFileMappingReadRepository + GeminiFileMappingWriteRepository
|
||||
{
|
||||
}
|
||||
789
crates/aether-data/src/repository/global_models/memory.rs
Normal file
789
crates/aether-data/src/repository/global_models/memory.rs
Normal file
@@ -0,0 +1,789 @@
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
|
||||
GlobalModelReadRepository, GlobalModelWriteRepository, PublicCatalogModelListQuery,
|
||||
PublicCatalogModelSearchQuery, PublicGlobalModelQuery, StoredAdminGlobalModel,
|
||||
StoredAdminGlobalModelPage, StoredAdminProviderModel, StoredProviderActiveGlobalModel,
|
||||
StoredProviderModelStats, StoredPublicCatalogModel, StoredPublicGlobalModel,
|
||||
StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryGlobalModelReadRepository {
|
||||
items: RwLock<Vec<StoredPublicGlobalModel>>,
|
||||
admin_global_model_items: RwLock<Vec<StoredAdminGlobalModel>>,
|
||||
public_catalog_items: RwLock<Vec<StoredPublicCatalogModel>>,
|
||||
admin_provider_model_items: RwLock<Vec<StoredAdminProviderModel>>,
|
||||
provider_model_stats: RwLock<Vec<StoredProviderModelStats>>,
|
||||
active_global_model_refs: RwLock<Vec<StoredProviderActiveGlobalModel>>,
|
||||
}
|
||||
|
||||
impl InMemoryGlobalModelReadRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredPublicGlobalModel>,
|
||||
{
|
||||
Self {
|
||||
items: RwLock::new(items.into_iter().collect()),
|
||||
admin_global_model_items: RwLock::new(Vec::new()),
|
||||
public_catalog_items: RwLock::new(Vec::new()),
|
||||
admin_provider_model_items: RwLock::new(Vec::new()),
|
||||
provider_model_stats: RwLock::new(Vec::new()),
|
||||
active_global_model_refs: RwLock::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_public_catalog_models<I>(self, items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredPublicCatalogModel>,
|
||||
{
|
||||
*self
|
||||
.public_catalog_items
|
||||
.write()
|
||||
.expect("public catalog model repository lock") = items.into_iter().collect();
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_provider_model_stats<I>(self, items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredProviderModelStats>,
|
||||
{
|
||||
*self
|
||||
.provider_model_stats
|
||||
.write()
|
||||
.expect("provider model stats repository lock") = items.into_iter().collect();
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_admin_provider_models<I>(self, items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredAdminProviderModel>,
|
||||
{
|
||||
*self
|
||||
.admin_provider_model_items
|
||||
.write()
|
||||
.expect("admin provider model repository lock") = items.into_iter().collect();
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_active_global_model_refs<I>(self, items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredProviderActiveGlobalModel>,
|
||||
{
|
||||
*self
|
||||
.active_global_model_refs
|
||||
.write()
|
||||
.expect("active global model repository lock") = items.into_iter().collect();
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_admin_global_models<I>(self, items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredAdminGlobalModel>,
|
||||
{
|
||||
*self
|
||||
.admin_global_model_items
|
||||
.write()
|
||||
.expect("admin global model repository lock") = items.into_iter().collect();
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl GlobalModelReadRepository for InMemoryGlobalModelReadRepository {
|
||||
async fn list_public_models(
|
||||
&self,
|
||||
query: &PublicGlobalModelQuery,
|
||||
) -> Result<StoredPublicGlobalModelPage, DataLayerError> {
|
||||
let items = self.items.read().expect("global model repository lock");
|
||||
let search = query
|
||||
.search
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|value| value.to_ascii_lowercase());
|
||||
|
||||
let mut filtered = items
|
||||
.iter()
|
||||
.filter(|item| match query.is_active {
|
||||
Some(is_active) => item.is_active == is_active,
|
||||
None => item.is_active,
|
||||
})
|
||||
.filter(|item| {
|
||||
let Some(search) = search.as_deref() else {
|
||||
return true;
|
||||
};
|
||||
item.name.to_ascii_lowercase().contains(search)
|
||||
|| item
|
||||
.display_name
|
||||
.as_deref()
|
||||
.map(|value| value.to_ascii_lowercase().contains(search))
|
||||
.unwrap_or(false)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
filtered.sort_by(|left, right| left.name.cmp(&right.name));
|
||||
let total = filtered.len();
|
||||
let items = filtered
|
||||
.into_iter()
|
||||
.skip(query.offset)
|
||||
.take(query.limit)
|
||||
.collect();
|
||||
Ok(StoredPublicGlobalModelPage { items, total })
|
||||
}
|
||||
|
||||
async fn get_public_model_by_name(
|
||||
&self,
|
||||
model_name: &str,
|
||||
) -> Result<Option<StoredPublicGlobalModel>, DataLayerError> {
|
||||
let items = self.items.read().expect("global model repository lock");
|
||||
Ok(items
|
||||
.iter()
|
||||
.find(|item| item.is_active && item.name == model_name)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn list_public_catalog_models(
|
||||
&self,
|
||||
query: &PublicCatalogModelListQuery,
|
||||
) -> Result<Vec<StoredPublicCatalogModel>, DataLayerError> {
|
||||
let items = self
|
||||
.public_catalog_items
|
||||
.read()
|
||||
.expect("public catalog model repository lock");
|
||||
let provider_id = query
|
||||
.provider_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
|
||||
let mut filtered = items
|
||||
.iter()
|
||||
.filter(|item| item.is_active)
|
||||
.filter(|item| match provider_id {
|
||||
Some(provider_id) => item.provider_id == provider_id,
|
||||
None => true,
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
filtered.sort_by(|left, right| {
|
||||
left.provider_name
|
||||
.cmp(&right.provider_name)
|
||||
.then_with(|| left.name.cmp(&right.name))
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
Ok(filtered
|
||||
.into_iter()
|
||||
.skip(query.offset)
|
||||
.take(query.limit)
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn search_public_catalog_models(
|
||||
&self,
|
||||
query: &PublicCatalogModelSearchQuery,
|
||||
) -> Result<Vec<StoredPublicCatalogModel>, DataLayerError> {
|
||||
let items = self
|
||||
.public_catalog_items
|
||||
.read()
|
||||
.expect("public catalog model repository lock");
|
||||
let provider_id = query
|
||||
.provider_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
let search = query.search.trim().to_ascii_lowercase();
|
||||
|
||||
let mut filtered = items
|
||||
.iter()
|
||||
.filter(|item| item.is_active)
|
||||
.filter(|item| match provider_id {
|
||||
Some(provider_id) => item.provider_id == provider_id,
|
||||
None => true,
|
||||
})
|
||||
.filter(|item| {
|
||||
item.provider_model_name
|
||||
.to_ascii_lowercase()
|
||||
.contains(&search)
|
||||
|| item.name.to_ascii_lowercase().contains(&search)
|
||||
|| item.display_name.to_ascii_lowercase().contains(&search)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
filtered.sort_by(|left, right| {
|
||||
left.provider_name
|
||||
.cmp(&right.provider_name)
|
||||
.then_with(|| left.name.cmp(&right.name))
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
filtered.truncate(query.limit);
|
||||
Ok(filtered)
|
||||
}
|
||||
|
||||
async fn list_admin_global_models(
|
||||
&self,
|
||||
query: &AdminGlobalModelListQuery,
|
||||
) -> Result<StoredAdminGlobalModelPage, DataLayerError> {
|
||||
let items = self
|
||||
.admin_global_model_items
|
||||
.read()
|
||||
.expect("admin global model repository lock");
|
||||
let search = query
|
||||
.search
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|value| value.to_ascii_lowercase());
|
||||
let mut filtered = items
|
||||
.iter()
|
||||
.filter(|item| match query.is_active {
|
||||
Some(is_active) => item.is_active == is_active,
|
||||
None => true,
|
||||
})
|
||||
.filter(|item| {
|
||||
let Some(search) = search.as_deref() else {
|
||||
return true;
|
||||
};
|
||||
item.name.to_ascii_lowercase().contains(search)
|
||||
|| item.display_name.to_ascii_lowercase().contains(search)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
filtered.sort_by(|left, right| left.name.cmp(&right.name));
|
||||
let total = filtered.len();
|
||||
let items = filtered
|
||||
.into_iter()
|
||||
.skip(query.offset)
|
||||
.take(query.limit)
|
||||
.collect();
|
||||
Ok(StoredAdminGlobalModelPage { items, total })
|
||||
}
|
||||
|
||||
async fn list_admin_provider_models(
|
||||
&self,
|
||||
query: &AdminProviderModelListQuery,
|
||||
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
||||
let items = self
|
||||
.admin_provider_model_items
|
||||
.read()
|
||||
.expect("admin provider model repository lock");
|
||||
let mut filtered = items
|
||||
.iter()
|
||||
.filter(|item| item.provider_id == query.provider_id)
|
||||
.filter(|item| match query.is_active {
|
||||
Some(is_active) => item.is_active == is_active,
|
||||
None => true,
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
filtered.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.unwrap_or_default()
|
||||
.cmp(&left.created_at_unix_secs.unwrap_or_default())
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
Ok(filtered
|
||||
.into_iter()
|
||||
.skip(query.offset)
|
||||
.take(query.limit)
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn get_admin_provider_model(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
model_id: &str,
|
||||
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
|
||||
Ok(self
|
||||
.admin_provider_model_items
|
||||
.read()
|
||||
.expect("admin provider model repository lock")
|
||||
.iter()
|
||||
.find(|item| item.provider_id == provider_id && item.id == model_id)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn list_admin_provider_available_source_models(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
||||
let items = self
|
||||
.admin_provider_model_items
|
||||
.read()
|
||||
.expect("admin provider model repository lock");
|
||||
let active_globals = self
|
||||
.admin_global_model_items
|
||||
.read()
|
||||
.expect("admin global model repository lock");
|
||||
|
||||
let mut filtered = items
|
||||
.iter()
|
||||
.filter(|item| item.provider_id == provider_id && item.is_active)
|
||||
.filter(|item| {
|
||||
active_globals
|
||||
.iter()
|
||||
.find(|global| global.id == item.global_model_id)
|
||||
.map(|global| global.is_active)
|
||||
.unwrap_or(false)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
filtered.sort_by(|left, right| {
|
||||
left.global_model_name
|
||||
.cmp(&right.global_model_name)
|
||||
.then_with(|| right.created_at_unix_secs.cmp(&left.created_at_unix_secs))
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
Ok(filtered)
|
||||
}
|
||||
|
||||
async fn get_admin_global_model_by_id(
|
||||
&self,
|
||||
global_model_id: &str,
|
||||
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
||||
let items = self
|
||||
.admin_global_model_items
|
||||
.read()
|
||||
.expect("admin global model repository lock");
|
||||
Ok(items
|
||||
.iter()
|
||||
.find(|item| item.id == global_model_id)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn get_admin_global_model_by_name(
|
||||
&self,
|
||||
model_name: &str,
|
||||
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
||||
let items = self
|
||||
.admin_global_model_items
|
||||
.read()
|
||||
.expect("admin global model repository lock");
|
||||
Ok(items.iter().find(|item| item.name == model_name).cloned())
|
||||
}
|
||||
|
||||
async fn list_admin_provider_models_by_global_model_id(
|
||||
&self,
|
||||
global_model_id: &str,
|
||||
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
||||
let items = self
|
||||
.admin_provider_model_items
|
||||
.read()
|
||||
.expect("admin provider model repository lock");
|
||||
let mut filtered = items
|
||||
.iter()
|
||||
.filter(|item| item.global_model_id == global_model_id)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
filtered.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.unwrap_or_default()
|
||||
.cmp(&left.created_at_unix_secs.unwrap_or_default())
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
Ok(filtered)
|
||||
}
|
||||
|
||||
async fn list_provider_model_stats(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderModelStats>, DataLayerError> {
|
||||
let provider_ids = provider_ids
|
||||
.iter()
|
||||
.cloned()
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
Ok(self
|
||||
.provider_model_stats
|
||||
.read()
|
||||
.expect("provider model stats repository lock")
|
||||
.iter()
|
||||
.filter(|item| provider_ids.contains(&item.provider_id))
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_active_global_model_ids_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderActiveGlobalModel>, DataLayerError> {
|
||||
let provider_ids = provider_ids
|
||||
.iter()
|
||||
.cloned()
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
Ok(self
|
||||
.active_global_model_refs
|
||||
.read()
|
||||
.expect("active global model repository lock")
|
||||
.iter()
|
||||
.filter(|item| provider_ids.contains(&item.provider_id))
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl GlobalModelWriteRepository for InMemoryGlobalModelReadRepository {
|
||||
async fn create_admin_provider_model(
|
||||
&self,
|
||||
record: &UpsertAdminProviderModelRecord,
|
||||
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
|
||||
let global_model = self
|
||||
.get_admin_global_model_by_id(&record.global_model_id)
|
||||
.await?
|
||||
.ok_or_else(|| DataLayerError::UnexpectedValue("global model not found".to_string()))?;
|
||||
|
||||
let stored = StoredAdminProviderModel::new(
|
||||
record.id.clone(),
|
||||
record.provider_id.clone(),
|
||||
record.global_model_id.clone(),
|
||||
record.provider_model_name.clone(),
|
||||
record.provider_model_mappings.clone(),
|
||||
record.price_per_request,
|
||||
record.tiered_pricing.clone(),
|
||||
record.supports_vision,
|
||||
record.supports_function_calling,
|
||||
record.supports_streaming,
|
||||
record.supports_extended_thinking,
|
||||
record.supports_image_generation,
|
||||
record.is_active,
|
||||
record.is_available,
|
||||
record.config.clone(),
|
||||
Some(1_711_000_000),
|
||||
Some(1_711_000_000),
|
||||
Some(global_model.name.clone()),
|
||||
Some(global_model.display_name.clone()),
|
||||
global_model.default_price_per_request,
|
||||
global_model.default_tiered_pricing.clone(),
|
||||
global_model.config.clone(),
|
||||
)?;
|
||||
self.admin_provider_model_items
|
||||
.write()
|
||||
.expect("admin provider model repository lock")
|
||||
.push(stored.clone());
|
||||
Ok(Some(stored))
|
||||
}
|
||||
|
||||
async fn update_admin_provider_model(
|
||||
&self,
|
||||
record: &UpsertAdminProviderModelRecord,
|
||||
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
|
||||
let global_model = self
|
||||
.get_admin_global_model_by_id(&record.global_model_id)
|
||||
.await?
|
||||
.ok_or_else(|| DataLayerError::UnexpectedValue("global model not found".to_string()))?;
|
||||
let mut items = self
|
||||
.admin_provider_model_items
|
||||
.write()
|
||||
.expect("admin provider model repository lock");
|
||||
let Some(existing) = items
|
||||
.iter_mut()
|
||||
.find(|item| item.id == record.id && item.provider_id == record.provider_id)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
existing.global_model_id = record.global_model_id.clone();
|
||||
existing.provider_model_name = record.provider_model_name.clone();
|
||||
existing.provider_model_mappings = record.provider_model_mappings.clone();
|
||||
existing.price_per_request = record.price_per_request;
|
||||
existing.tiered_pricing = record.tiered_pricing.clone();
|
||||
existing.supports_vision = record.supports_vision;
|
||||
existing.supports_function_calling = record.supports_function_calling;
|
||||
existing.supports_streaming = record.supports_streaming;
|
||||
existing.supports_extended_thinking = record.supports_extended_thinking;
|
||||
existing.supports_image_generation = record.supports_image_generation;
|
||||
existing.is_active = record.is_active;
|
||||
existing.is_available = record.is_available;
|
||||
existing.config = record.config.clone();
|
||||
existing.updated_at_unix_secs = Some(1_711_000_100);
|
||||
existing.global_model_name = Some(global_model.name.clone());
|
||||
existing.global_model_display_name = Some(global_model.display_name.clone());
|
||||
existing.global_model_default_price_per_request = global_model.default_price_per_request;
|
||||
existing.global_model_default_tiered_pricing = global_model.default_tiered_pricing.clone();
|
||||
existing.global_model_config = global_model.config.clone();
|
||||
Ok(Some(existing.clone()))
|
||||
}
|
||||
|
||||
async fn delete_admin_provider_model(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
model_id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let mut items = self
|
||||
.admin_provider_model_items
|
||||
.write()
|
||||
.expect("admin provider model repository lock");
|
||||
let original_len = items.len();
|
||||
items.retain(|item| !(item.provider_id == provider_id && item.id == model_id));
|
||||
Ok(items.len() != original_len)
|
||||
}
|
||||
|
||||
async fn create_admin_global_model(
|
||||
&self,
|
||||
record: &CreateAdminGlobalModelRecord,
|
||||
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
||||
let stored = StoredAdminGlobalModel::new(
|
||||
record.id.clone(),
|
||||
record.name.clone(),
|
||||
record.display_name.clone(),
|
||||
record.is_active,
|
||||
record.default_price_per_request,
|
||||
record.default_tiered_pricing.clone(),
|
||||
record.supported_capabilities.clone(),
|
||||
record.config.clone(),
|
||||
Some(1_711_000_000),
|
||||
Some(1_711_000_000),
|
||||
)?;
|
||||
self.admin_global_model_items
|
||||
.write()
|
||||
.expect("admin global model repository lock")
|
||||
.push(stored.clone());
|
||||
Ok(Some(stored))
|
||||
}
|
||||
|
||||
async fn update_admin_global_model(
|
||||
&self,
|
||||
record: &UpdateAdminGlobalModelRecord,
|
||||
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
||||
let mut items = self
|
||||
.admin_global_model_items
|
||||
.write()
|
||||
.expect("admin global model repository lock");
|
||||
let Some(existing) = items.iter_mut().find(|item| item.id == record.id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
existing.display_name = record.display_name.clone();
|
||||
existing.is_active = record.is_active;
|
||||
existing.default_price_per_request = record.default_price_per_request;
|
||||
existing.default_tiered_pricing = record.default_tiered_pricing.clone();
|
||||
existing.supported_capabilities = record.supported_capabilities.clone();
|
||||
existing.config = record.config.clone();
|
||||
existing.updated_at_unix_secs = Some(1_711_000_100);
|
||||
Ok(Some(existing.clone()))
|
||||
}
|
||||
|
||||
async fn delete_admin_global_model(
|
||||
&self,
|
||||
global_model_id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let mut globals = self
|
||||
.admin_global_model_items
|
||||
.write()
|
||||
.expect("admin global model repository lock");
|
||||
let original_len = globals.len();
|
||||
globals.retain(|item| item.id != global_model_id);
|
||||
drop(globals);
|
||||
self.admin_provider_model_items
|
||||
.write()
|
||||
.expect("admin provider model repository lock")
|
||||
.retain(|item| item.global_model_id != global_model_id);
|
||||
Ok(original_len
|
||||
!= self
|
||||
.admin_global_model_items
|
||||
.read()
|
||||
.expect("admin global model repository lock")
|
||||
.len())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::InMemoryGlobalModelReadRepository;
|
||||
use crate::repository::global_models::{
|
||||
GlobalModelReadRepository, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery,
|
||||
PublicGlobalModelQuery, StoredPublicCatalogModel, StoredPublicGlobalModel,
|
||||
};
|
||||
|
||||
fn sample_model(
|
||||
id: &str,
|
||||
name: &str,
|
||||
display_name: &str,
|
||||
is_active: bool,
|
||||
) -> StoredPublicGlobalModel {
|
||||
StoredPublicGlobalModel::new(
|
||||
id.to_string(),
|
||||
name.to_string(),
|
||||
Some(display_name.to_string()),
|
||||
is_active,
|
||||
Some(0.02),
|
||||
Some(json!({"tiers":[{"up_to": null, "input_price_per_1m": 3.0, "output_price_per_1m": 15.0}]})),
|
||||
Some(json!(["vision"])),
|
||||
Some(json!({"family": "test"})),
|
||||
0,
|
||||
)
|
||||
.expect("global model should build")
|
||||
}
|
||||
|
||||
fn sample_public_catalog_model(
|
||||
id: &str,
|
||||
provider_id: &str,
|
||||
provider_name: &str,
|
||||
provider_model_name: &str,
|
||||
name: &str,
|
||||
display_name: &str,
|
||||
) -> StoredPublicCatalogModel {
|
||||
StoredPublicCatalogModel::new(
|
||||
id.to_string(),
|
||||
provider_id.to_string(),
|
||||
provider_name.to_string(),
|
||||
provider_model_name.to_string(),
|
||||
name.to_string(),
|
||||
display_name.to_string(),
|
||||
Some(format!("{display_name} description")),
|
||||
Some(format!("https://cdn.example/{name}.png")),
|
||||
Some(3.0),
|
||||
Some(15.0),
|
||||
Some(1.5),
|
||||
Some(0.3),
|
||||
Some(true),
|
||||
Some(true),
|
||||
Some(true),
|
||||
true,
|
||||
)
|
||||
.expect("public catalog model should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn defaults_to_active_models_only() {
|
||||
let repository = InMemoryGlobalModelReadRepository::seed(vec![
|
||||
sample_model("gm-1", "claude-sonnet-4-5", "Claude Sonnet 4.5", true),
|
||||
sample_model("gm-2", "legacy-model", "Legacy Model", false),
|
||||
]);
|
||||
|
||||
let page = repository
|
||||
.list_public_models(&PublicGlobalModelQuery {
|
||||
offset: 0,
|
||||
limit: 50,
|
||||
is_active: None,
|
||||
search: None,
|
||||
})
|
||||
.await
|
||||
.expect("list should succeed");
|
||||
|
||||
assert_eq!(page.total, 1);
|
||||
assert_eq!(page.items[0].name, "claude-sonnet-4-5");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn search_matches_name_and_display_name() {
|
||||
let repository = InMemoryGlobalModelReadRepository::seed(vec![
|
||||
sample_model("gm-1", "gpt-5", "GPT 5", true),
|
||||
sample_model("gm-2", "claude-sonnet-4-5", "Claude Sonnet 4.5", true),
|
||||
]);
|
||||
|
||||
let page = repository
|
||||
.list_public_models(&PublicGlobalModelQuery {
|
||||
offset: 0,
|
||||
limit: 50,
|
||||
is_active: None,
|
||||
search: Some("sonnet".to_string()),
|
||||
})
|
||||
.await
|
||||
.expect("list should succeed");
|
||||
|
||||
assert_eq!(page.total, 1);
|
||||
assert_eq!(page.items[0].name, "claude-sonnet-4-5");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_public_model_by_name_only_returns_active_exact_match() {
|
||||
let repository = InMemoryGlobalModelReadRepository::seed(vec![
|
||||
sample_model("gm-1", "gpt-5", "GPT 5", true),
|
||||
sample_model("gm-2", "gpt-5-old", "GPT 5 Old", false),
|
||||
]);
|
||||
|
||||
let model = repository
|
||||
.get_public_model_by_name("gpt-5")
|
||||
.await
|
||||
.expect("lookup should succeed");
|
||||
assert_eq!(model.expect("model should exist").name, "gpt-5");
|
||||
|
||||
let missing = repository
|
||||
.get_public_model_by_name("gpt-5-old")
|
||||
.await
|
||||
.expect("lookup should succeed");
|
||||
assert!(missing.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lists_public_catalog_models_with_provider_filter() {
|
||||
let repository =
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::<StoredPublicGlobalModel>::new())
|
||||
.with_public_catalog_models(vec![
|
||||
sample_public_catalog_model(
|
||||
"model-1",
|
||||
"provider-openai",
|
||||
"openai",
|
||||
"gpt-5-preview",
|
||||
"gpt-5",
|
||||
"GPT 5",
|
||||
),
|
||||
sample_public_catalog_model(
|
||||
"model-2",
|
||||
"provider-claude",
|
||||
"claude",
|
||||
"claude-3-7-sonnet",
|
||||
"claude-3-7-sonnet",
|
||||
"Claude 3.7 Sonnet",
|
||||
),
|
||||
]);
|
||||
|
||||
let items = repository
|
||||
.list_public_catalog_models(&PublicCatalogModelListQuery {
|
||||
provider_id: Some("provider-openai".to_string()),
|
||||
offset: 0,
|
||||
limit: 50,
|
||||
})
|
||||
.await
|
||||
.expect("list should succeed");
|
||||
|
||||
assert_eq!(items.len(), 1);
|
||||
assert_eq!(items[0].provider_id, "provider-openai");
|
||||
assert_eq!(items[0].name, "gpt-5");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn searches_public_catalog_models_by_provider_and_display_name() {
|
||||
let repository =
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::<StoredPublicGlobalModel>::new())
|
||||
.with_public_catalog_models(vec![
|
||||
sample_public_catalog_model(
|
||||
"model-1",
|
||||
"provider-openai",
|
||||
"openai",
|
||||
"gpt-5-preview",
|
||||
"gpt-5",
|
||||
"GPT 5",
|
||||
),
|
||||
sample_public_catalog_model(
|
||||
"model-2",
|
||||
"provider-claude",
|
||||
"claude",
|
||||
"claude-3-7-sonnet",
|
||||
"claude-3-7-sonnet",
|
||||
"Claude 3.7 Sonnet",
|
||||
),
|
||||
]);
|
||||
|
||||
let items = repository
|
||||
.search_public_catalog_models(&PublicCatalogModelSearchQuery {
|
||||
search: "sonnet".to_string(),
|
||||
provider_id: Some("provider-claude".to_string()),
|
||||
limit: 20,
|
||||
})
|
||||
.await
|
||||
.expect("search should succeed");
|
||||
|
||||
assert_eq!(items.len(), 1);
|
||||
assert_eq!(items[0].provider_name, "claude");
|
||||
assert_eq!(items[0].display_name, "Claude 3.7 Sonnet");
|
||||
}
|
||||
}
|
||||
14
crates/aether-data/src/repository/global_models/mod.rs
Normal file
14
crates/aether-data/src/repository/global_models/mod.rs
Normal file
@@ -0,0 +1,14 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryGlobalModelReadRepository;
|
||||
pub use sql::SqlxGlobalModelReadRepository;
|
||||
pub use types::{
|
||||
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
|
||||
GlobalModelReadRepository, GlobalModelWriteRepository, PublicCatalogModelListQuery,
|
||||
PublicCatalogModelSearchQuery, PublicGlobalModelQuery, StoredAdminGlobalModel,
|
||||
StoredAdminGlobalModelPage, StoredAdminProviderModel, StoredProviderActiveGlobalModel,
|
||||
StoredProviderModelStats, StoredPublicCatalogModel, StoredPublicGlobalModel,
|
||||
StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
|
||||
};
|
||||
1074
crates/aether-data/src/repository/global_models/sql.rs
Normal file
1074
crates/aether-data/src/repository/global_models/sql.rs
Normal file
File diff suppressed because it is too large
Load Diff
688
crates/aether-data/src/repository/global_models/types.rs
Normal file
688
crates/aether-data/src/repository/global_models/types.rs
Normal file
@@ -0,0 +1,688 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredPublicGlobalModel {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub display_name: Option<String>,
|
||||
pub is_active: bool,
|
||||
pub default_price_per_request: Option<f64>,
|
||||
pub default_tiered_pricing: Option<Value>,
|
||||
pub supported_capabilities: Option<Value>,
|
||||
pub config: Option<Value>,
|
||||
pub usage_count: u64,
|
||||
}
|
||||
|
||||
impl StoredPublicGlobalModel {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
name: String,
|
||||
display_name: Option<String>,
|
||||
is_active: bool,
|
||||
default_price_per_request: Option<f64>,
|
||||
default_tiered_pricing: Option<Value>,
|
||||
supported_capabilities: Option<Value>,
|
||||
config: Option<Value>,
|
||||
usage_count: u64,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"global_models.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"global_models.name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
name,
|
||||
display_name,
|
||||
is_active,
|
||||
default_price_per_request,
|
||||
default_tiered_pricing,
|
||||
supported_capabilities,
|
||||
config,
|
||||
usage_count,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct PublicGlobalModelQuery {
|
||||
pub offset: usize,
|
||||
pub limit: usize,
|
||||
pub is_active: Option<bool>,
|
||||
pub search: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredPublicCatalogModel {
|
||||
pub id: String,
|
||||
pub provider_id: String,
|
||||
pub provider_name: String,
|
||||
pub provider_model_name: String,
|
||||
pub name: String,
|
||||
pub display_name: String,
|
||||
pub description: Option<String>,
|
||||
pub icon_url: Option<String>,
|
||||
pub input_price_per_1m: Option<f64>,
|
||||
pub output_price_per_1m: Option<f64>,
|
||||
pub cache_creation_price_per_1m: Option<f64>,
|
||||
pub cache_read_price_per_1m: Option<f64>,
|
||||
pub supports_vision: Option<bool>,
|
||||
pub supports_function_calling: Option<bool>,
|
||||
pub supports_streaming: Option<bool>,
|
||||
pub is_active: bool,
|
||||
}
|
||||
|
||||
impl StoredPublicCatalogModel {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
provider_id: String,
|
||||
provider_name: String,
|
||||
provider_model_name: String,
|
||||
name: String,
|
||||
display_name: String,
|
||||
description: Option<String>,
|
||||
icon_url: Option<String>,
|
||||
input_price_per_1m: Option<f64>,
|
||||
output_price_per_1m: Option<f64>,
|
||||
cache_creation_price_per_1m: Option<f64>,
|
||||
cache_read_price_per_1m: Option<f64>,
|
||||
supports_vision: Option<bool>,
|
||||
supports_function_calling: Option<bool>,
|
||||
supports_streaming: Option<bool>,
|
||||
is_active: bool,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"models.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if provider_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"models.provider_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if provider_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"providers.name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if provider_model_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"models.provider_model_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"public model name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if display_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"public model display_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
provider_id,
|
||||
provider_name,
|
||||
provider_model_name,
|
||||
name,
|
||||
display_name,
|
||||
description,
|
||||
icon_url,
|
||||
input_price_per_1m,
|
||||
output_price_per_1m,
|
||||
cache_creation_price_per_1m,
|
||||
cache_read_price_per_1m,
|
||||
supports_vision,
|
||||
supports_function_calling,
|
||||
supports_streaming,
|
||||
is_active,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct PublicCatalogModelListQuery {
|
||||
pub provider_id: Option<String>,
|
||||
pub offset: usize,
|
||||
pub limit: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PublicCatalogModelSearchQuery {
|
||||
pub search: String,
|
||||
pub provider_id: Option<String>,
|
||||
pub limit: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AdminProviderModelListQuery {
|
||||
pub provider_id: String,
|
||||
pub is_active: Option<bool>,
|
||||
pub offset: usize,
|
||||
pub limit: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredAdminGlobalModel {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub display_name: String,
|
||||
pub is_active: bool,
|
||||
pub default_price_per_request: Option<f64>,
|
||||
pub default_tiered_pricing: Option<Value>,
|
||||
pub supported_capabilities: Option<Value>,
|
||||
pub config: Option<Value>,
|
||||
pub created_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl StoredAdminGlobalModel {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
name: String,
|
||||
display_name: String,
|
||||
is_active: bool,
|
||||
default_price_per_request: Option<f64>,
|
||||
default_tiered_pricing: Option<Value>,
|
||||
supported_capabilities: Option<Value>,
|
||||
config: Option<Value>,
|
||||
created_at_unix_secs: Option<u64>,
|
||||
updated_at_unix_secs: Option<u64>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"global_models.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"global_models.name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if display_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"global_models.display_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
name,
|
||||
display_name,
|
||||
is_active,
|
||||
default_price_per_request,
|
||||
default_tiered_pricing,
|
||||
supported_capabilities,
|
||||
config,
|
||||
created_at_unix_secs,
|
||||
updated_at_unix_secs,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct AdminGlobalModelListQuery {
|
||||
pub offset: usize,
|
||||
pub limit: usize,
|
||||
pub is_active: Option<bool>,
|
||||
pub search: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredAdminProviderModel {
|
||||
pub id: String,
|
||||
pub provider_id: String,
|
||||
pub global_model_id: String,
|
||||
pub provider_model_name: String,
|
||||
pub provider_model_mappings: Option<Value>,
|
||||
pub price_per_request: Option<f64>,
|
||||
pub tiered_pricing: Option<Value>,
|
||||
pub supports_vision: Option<bool>,
|
||||
pub supports_function_calling: Option<bool>,
|
||||
pub supports_streaming: Option<bool>,
|
||||
pub supports_extended_thinking: Option<bool>,
|
||||
pub supports_image_generation: Option<bool>,
|
||||
pub is_active: bool,
|
||||
pub is_available: bool,
|
||||
pub config: Option<Value>,
|
||||
pub created_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: Option<u64>,
|
||||
pub global_model_name: Option<String>,
|
||||
pub global_model_display_name: Option<String>,
|
||||
pub global_model_default_price_per_request: Option<f64>,
|
||||
pub global_model_default_tiered_pricing: Option<Value>,
|
||||
pub global_model_config: Option<Value>,
|
||||
}
|
||||
|
||||
impl StoredAdminProviderModel {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
provider_id: String,
|
||||
global_model_id: String,
|
||||
provider_model_name: String,
|
||||
provider_model_mappings: Option<Value>,
|
||||
price_per_request: Option<f64>,
|
||||
tiered_pricing: Option<Value>,
|
||||
supports_vision: Option<bool>,
|
||||
supports_function_calling: Option<bool>,
|
||||
supports_streaming: Option<bool>,
|
||||
supports_extended_thinking: Option<bool>,
|
||||
supports_image_generation: Option<bool>,
|
||||
is_active: bool,
|
||||
is_available: bool,
|
||||
config: Option<Value>,
|
||||
created_at_unix_secs: Option<u64>,
|
||||
updated_at_unix_secs: Option<u64>,
|
||||
global_model_name: Option<String>,
|
||||
global_model_display_name: Option<String>,
|
||||
global_model_default_price_per_request: Option<f64>,
|
||||
global_model_default_tiered_pricing: Option<Value>,
|
||||
global_model_config: Option<Value>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"models.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if provider_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"models.provider_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if global_model_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"models.global_model_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if provider_model_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"models.provider_model_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
provider_id,
|
||||
global_model_id,
|
||||
provider_model_name,
|
||||
provider_model_mappings,
|
||||
price_per_request,
|
||||
tiered_pricing,
|
||||
supports_vision,
|
||||
supports_function_calling,
|
||||
supports_streaming,
|
||||
supports_extended_thinking,
|
||||
supports_image_generation,
|
||||
is_active,
|
||||
is_available,
|
||||
config,
|
||||
created_at_unix_secs,
|
||||
updated_at_unix_secs,
|
||||
global_model_name,
|
||||
global_model_display_name,
|
||||
global_model_default_price_per_request,
|
||||
global_model_default_tiered_pricing,
|
||||
global_model_config,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct UpsertAdminProviderModelRecord {
|
||||
pub id: String,
|
||||
pub provider_id: String,
|
||||
pub global_model_id: String,
|
||||
pub provider_model_name: String,
|
||||
pub provider_model_mappings: Option<Value>,
|
||||
pub price_per_request: Option<f64>,
|
||||
pub tiered_pricing: Option<Value>,
|
||||
pub supports_vision: Option<bool>,
|
||||
pub supports_function_calling: Option<bool>,
|
||||
pub supports_streaming: Option<bool>,
|
||||
pub supports_extended_thinking: Option<bool>,
|
||||
pub supports_image_generation: Option<bool>,
|
||||
pub is_active: bool,
|
||||
pub is_available: bool,
|
||||
pub config: Option<Value>,
|
||||
}
|
||||
|
||||
impl UpsertAdminProviderModelRecord {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
provider_id: String,
|
||||
global_model_id: String,
|
||||
provider_model_name: String,
|
||||
provider_model_mappings: Option<Value>,
|
||||
price_per_request: Option<f64>,
|
||||
tiered_pricing: Option<Value>,
|
||||
supports_vision: Option<bool>,
|
||||
supports_function_calling: Option<bool>,
|
||||
supports_streaming: Option<bool>,
|
||||
supports_extended_thinking: Option<bool>,
|
||||
supports_image_generation: Option<bool>,
|
||||
is_active: bool,
|
||||
is_available: bool,
|
||||
config: Option<Value>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"models.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if provider_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"models.provider_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if global_model_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"models.global_model_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if provider_model_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"models.provider_model_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
provider_id,
|
||||
global_model_id,
|
||||
provider_model_name,
|
||||
provider_model_mappings,
|
||||
price_per_request,
|
||||
tiered_pricing,
|
||||
supports_vision,
|
||||
supports_function_calling,
|
||||
supports_streaming,
|
||||
supports_extended_thinking,
|
||||
supports_image_generation,
|
||||
is_active,
|
||||
is_available,
|
||||
config,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct CreateAdminGlobalModelRecord {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub display_name: String,
|
||||
pub is_active: bool,
|
||||
pub default_price_per_request: Option<f64>,
|
||||
pub default_tiered_pricing: Option<Value>,
|
||||
pub supported_capabilities: Option<Value>,
|
||||
pub config: Option<Value>,
|
||||
}
|
||||
|
||||
impl CreateAdminGlobalModelRecord {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
name: String,
|
||||
display_name: String,
|
||||
is_active: bool,
|
||||
default_price_per_request: Option<f64>,
|
||||
default_tiered_pricing: Option<Value>,
|
||||
supported_capabilities: Option<Value>,
|
||||
config: Option<Value>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"global_models.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"global_models.name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if display_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"global_models.display_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
name,
|
||||
display_name,
|
||||
is_active,
|
||||
default_price_per_request,
|
||||
default_tiered_pricing,
|
||||
supported_capabilities,
|
||||
config,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct UpdateAdminGlobalModelRecord {
|
||||
pub id: String,
|
||||
pub display_name: String,
|
||||
pub is_active: bool,
|
||||
pub default_price_per_request: Option<f64>,
|
||||
pub default_tiered_pricing: Option<Value>,
|
||||
pub supported_capabilities: Option<Value>,
|
||||
pub config: Option<Value>,
|
||||
}
|
||||
|
||||
impl UpdateAdminGlobalModelRecord {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
display_name: String,
|
||||
is_active: bool,
|
||||
default_price_per_request: Option<f64>,
|
||||
default_tiered_pricing: Option<Value>,
|
||||
supported_capabilities: Option<Value>,
|
||||
config: Option<Value>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"global_models.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if display_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"global_models.display_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
display_name,
|
||||
is_active,
|
||||
default_price_per_request,
|
||||
default_tiered_pricing,
|
||||
supported_capabilities,
|
||||
config,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredPublicGlobalModelPage {
|
||||
pub items: Vec<StoredPublicGlobalModel>,
|
||||
pub total: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredAdminGlobalModelPage {
|
||||
pub items: Vec<StoredAdminGlobalModel>,
|
||||
pub total: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderModelStats {
|
||||
pub provider_id: String,
|
||||
pub total_models: u64,
|
||||
pub active_models: u64,
|
||||
}
|
||||
|
||||
impl StoredProviderModelStats {
|
||||
pub fn new(
|
||||
provider_id: String,
|
||||
total_models: i64,
|
||||
active_models: i64,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if provider_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider model stats provider_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if total_models < 0 || active_models < 0 {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider model stats count is negative".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(Self {
|
||||
provider_id,
|
||||
total_models: total_models as u64,
|
||||
active_models: active_models as u64,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderActiveGlobalModel {
|
||||
pub provider_id: String,
|
||||
pub global_model_id: String,
|
||||
}
|
||||
|
||||
impl StoredProviderActiveGlobalModel {
|
||||
pub fn new(
|
||||
provider_id: String,
|
||||
global_model_id: String,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if provider_id.trim().is_empty() || global_model_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider active global model identity is empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(Self {
|
||||
provider_id,
|
||||
global_model_id,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait GlobalModelReadRepository: Send + Sync {
|
||||
async fn list_public_models(
|
||||
&self,
|
||||
query: &PublicGlobalModelQuery,
|
||||
) -> Result<StoredPublicGlobalModelPage, crate::DataLayerError>;
|
||||
|
||||
async fn get_public_model_by_name(
|
||||
&self,
|
||||
model_name: &str,
|
||||
) -> Result<Option<StoredPublicGlobalModel>, crate::DataLayerError>;
|
||||
|
||||
async fn list_public_catalog_models(
|
||||
&self,
|
||||
query: &PublicCatalogModelListQuery,
|
||||
) -> Result<Vec<StoredPublicCatalogModel>, crate::DataLayerError>;
|
||||
|
||||
async fn search_public_catalog_models(
|
||||
&self,
|
||||
query: &PublicCatalogModelSearchQuery,
|
||||
) -> Result<Vec<StoredPublicCatalogModel>, crate::DataLayerError>;
|
||||
|
||||
async fn list_admin_global_models(
|
||||
&self,
|
||||
query: &AdminGlobalModelListQuery,
|
||||
) -> Result<StoredAdminGlobalModelPage, crate::DataLayerError>;
|
||||
|
||||
async fn list_admin_provider_models(
|
||||
&self,
|
||||
query: &AdminProviderModelListQuery,
|
||||
) -> Result<Vec<StoredAdminProviderModel>, crate::DataLayerError>;
|
||||
|
||||
async fn list_admin_provider_available_source_models(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
) -> Result<Vec<StoredAdminProviderModel>, crate::DataLayerError>;
|
||||
|
||||
async fn get_admin_provider_model(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
model_id: &str,
|
||||
) -> Result<Option<StoredAdminProviderModel>, crate::DataLayerError>;
|
||||
|
||||
async fn get_admin_global_model_by_id(
|
||||
&self,
|
||||
global_model_id: &str,
|
||||
) -> Result<Option<StoredAdminGlobalModel>, crate::DataLayerError>;
|
||||
|
||||
async fn get_admin_global_model_by_name(
|
||||
&self,
|
||||
model_name: &str,
|
||||
) -> Result<Option<StoredAdminGlobalModel>, crate::DataLayerError>;
|
||||
|
||||
async fn list_admin_provider_models_by_global_model_id(
|
||||
&self,
|
||||
global_model_id: &str,
|
||||
) -> Result<Vec<StoredAdminProviderModel>, crate::DataLayerError>;
|
||||
|
||||
async fn list_provider_model_stats(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderModelStats>, crate::DataLayerError>;
|
||||
|
||||
async fn list_active_global_model_ids_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderActiveGlobalModel>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait GlobalModelWriteRepository: Send + Sync {
|
||||
async fn create_admin_provider_model(
|
||||
&self,
|
||||
record: &UpsertAdminProviderModelRecord,
|
||||
) -> Result<Option<StoredAdminProviderModel>, crate::DataLayerError>;
|
||||
|
||||
async fn update_admin_provider_model(
|
||||
&self,
|
||||
record: &UpsertAdminProviderModelRecord,
|
||||
) -> Result<Option<StoredAdminProviderModel>, crate::DataLayerError>;
|
||||
|
||||
async fn delete_admin_provider_model(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
model_id: &str,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn create_admin_global_model(
|
||||
&self,
|
||||
record: &CreateAdminGlobalModelRecord,
|
||||
) -> Result<Option<StoredAdminGlobalModel>, crate::DataLayerError>;
|
||||
|
||||
async fn update_admin_global_model(
|
||||
&self,
|
||||
record: &UpdateAdminGlobalModelRecord,
|
||||
) -> Result<Option<StoredAdminGlobalModel>, crate::DataLayerError>;
|
||||
|
||||
async fn delete_admin_global_model(
|
||||
&self,
|
||||
global_model_id: &str,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
}
|
||||
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>;
|
||||
}
|
||||
@@ -1,6 +1,18 @@
|
||||
pub mod announcements;
|
||||
pub mod auth;
|
||||
pub mod auth_modules;
|
||||
pub mod billing;
|
||||
pub mod candidate_selection;
|
||||
pub mod candidates;
|
||||
pub mod gemini_file_mappings;
|
||||
pub mod global_models;
|
||||
pub mod management_tokens;
|
||||
pub mod oauth_providers;
|
||||
pub mod provider_catalog;
|
||||
pub mod proxy_nodes;
|
||||
pub mod quota;
|
||||
pub mod shadow_results;
|
||||
pub mod usage;
|
||||
pub mod users;
|
||||
pub mod video_tasks;
|
||||
pub mod wallet;
|
||||
|
||||
196
crates/aether-data/src/repository/oauth_providers/memory.rs
Normal file
196
crates/aether-data/src/repository/oauth_providers/memory.rs
Normal file
@@ -0,0 +1,196 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
EncryptedSecretUpdate, OAuthProviderReadRepository, OAuthProviderWriteRepository,
|
||||
StoredOAuthProviderConfig, UpsertOAuthProviderConfigRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryOAuthProviderRepository {
|
||||
items: RwLock<BTreeMap<String, StoredOAuthProviderConfig>>,
|
||||
}
|
||||
|
||||
impl InMemoryOAuthProviderRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredOAuthProviderConfig>,
|
||||
{
|
||||
let items = items
|
||||
.into_iter()
|
||||
.map(|item| (item.provider_type.clone(), item))
|
||||
.collect();
|
||||
Self {
|
||||
items: RwLock::new(items),
|
||||
}
|
||||
}
|
||||
|
||||
fn now_unix_secs() -> Option<u64> {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl OAuthProviderReadRepository for InMemoryOAuthProviderRepository {
|
||||
async fn list_oauth_provider_configs(
|
||||
&self,
|
||||
) -> Result<Vec<StoredOAuthProviderConfig>, DataLayerError> {
|
||||
let items = self.items.read().expect("oauth provider repository lock");
|
||||
Ok(items.values().cloned().collect())
|
||||
}
|
||||
|
||||
async fn get_oauth_provider_config(
|
||||
&self,
|
||||
provider_type: &str,
|
||||
) -> Result<Option<StoredOAuthProviderConfig>, DataLayerError> {
|
||||
let items = self.items.read().expect("oauth provider repository lock");
|
||||
Ok(items.get(provider_type).cloned())
|
||||
}
|
||||
|
||||
async fn count_locked_users_if_provider_disabled(
|
||||
&self,
|
||||
_provider_type: &str,
|
||||
_ldap_exclusive: bool,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
Ok(0)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl OAuthProviderWriteRepository for InMemoryOAuthProviderRepository {
|
||||
async fn upsert_oauth_provider_config(
|
||||
&self,
|
||||
record: &UpsertOAuthProviderConfigRecord,
|
||||
) -> Result<StoredOAuthProviderConfig, DataLayerError> {
|
||||
record.validate()?;
|
||||
|
||||
let mut items = self.items.write().expect("oauth provider repository lock");
|
||||
let now = Self::now_unix_secs();
|
||||
let existing = items.get(&record.provider_type).cloned();
|
||||
let created_at = existing
|
||||
.as_ref()
|
||||
.and_then(|item| item.created_at_unix_secs)
|
||||
.or(now);
|
||||
let client_secret_encrypted = match (&record.client_secret_encrypted, existing.as_ref()) {
|
||||
(EncryptedSecretUpdate::Preserve, Some(item)) => item.client_secret_encrypted.clone(),
|
||||
(EncryptedSecretUpdate::Preserve, None) => None,
|
||||
(EncryptedSecretUpdate::Clear, _) => None,
|
||||
(EncryptedSecretUpdate::Set(value), _) => Some(value.clone()),
|
||||
};
|
||||
|
||||
let item = StoredOAuthProviderConfig::new(
|
||||
record.provider_type.clone(),
|
||||
record.display_name.clone(),
|
||||
record.client_id.clone(),
|
||||
record.redirect_uri.clone(),
|
||||
record.frontend_callback_url.clone(),
|
||||
)?
|
||||
.with_config_fields(
|
||||
client_secret_encrypted,
|
||||
record.authorization_url_override.clone(),
|
||||
record.token_url_override.clone(),
|
||||
record.userinfo_url_override.clone(),
|
||||
record.scopes.clone(),
|
||||
record.attribute_mapping.clone(),
|
||||
record.extra_config.clone(),
|
||||
record.is_enabled,
|
||||
)
|
||||
.with_timestamps(created_at, now);
|
||||
|
||||
items.insert(record.provider_type.clone(), item.clone());
|
||||
Ok(item)
|
||||
}
|
||||
|
||||
async fn delete_oauth_provider_config(
|
||||
&self,
|
||||
provider_type: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let mut items = self.items.write().expect("oauth provider repository lock");
|
||||
Ok(items.remove(provider_type).is_some())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryOAuthProviderRepository;
|
||||
use crate::repository::oauth_providers::{
|
||||
EncryptedSecretUpdate, OAuthProviderReadRepository, OAuthProviderWriteRepository,
|
||||
StoredOAuthProviderConfig, UpsertOAuthProviderConfigRecord,
|
||||
};
|
||||
|
||||
fn sample_provider(provider_type: &str) -> StoredOAuthProviderConfig {
|
||||
StoredOAuthProviderConfig::new(
|
||||
provider_type.to_string(),
|
||||
format!("{provider_type} display"),
|
||||
format!("{provider_type}-client"),
|
||||
format!("https://{provider_type}.example.com/redirect"),
|
||||
"https://frontend.example.com/auth/callback".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
}
|
||||
|
||||
fn sample_upsert(provider_type: &str) -> UpsertOAuthProviderConfigRecord {
|
||||
UpsertOAuthProviderConfigRecord {
|
||||
provider_type: provider_type.to_string(),
|
||||
display_name: format!("{provider_type} display"),
|
||||
client_id: format!("{provider_type}-client"),
|
||||
client_secret_encrypted: EncryptedSecretUpdate::Preserve,
|
||||
authorization_url_override: Some(format!("https://{provider_type}.example.com/auth")),
|
||||
token_url_override: Some(format!("https://{provider_type}.example.com/token")),
|
||||
userinfo_url_override: None,
|
||||
scopes: Some(vec!["openid".to_string(), "profile".to_string()]),
|
||||
redirect_uri: format!("https://{provider_type}.example.com/redirect"),
|
||||
frontend_callback_url: "https://frontend.example.com/auth/callback".to_string(),
|
||||
attribute_mapping: Some(serde_json::json!({"email": "email"})),
|
||||
extra_config: Some(serde_json::json!({"team": true})),
|
||||
is_enabled: true,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_and_mutates_oauth_provider_configs() {
|
||||
let repository = InMemoryOAuthProviderRepository::seed(vec![
|
||||
sample_provider("linuxdo"),
|
||||
sample_provider("github"),
|
||||
]);
|
||||
|
||||
let listed = repository
|
||||
.list_oauth_provider_configs()
|
||||
.await
|
||||
.expect("list should succeed");
|
||||
assert_eq!(listed.len(), 2);
|
||||
assert_eq!(listed[0].provider_type, "github");
|
||||
assert_eq!(listed[1].provider_type, "linuxdo");
|
||||
|
||||
let created = repository
|
||||
.upsert_oauth_provider_config(&UpsertOAuthProviderConfigRecord {
|
||||
client_secret_encrypted: EncryptedSecretUpdate::Set("secret-1".to_string()),
|
||||
..sample_upsert("google")
|
||||
})
|
||||
.await
|
||||
.expect("create should succeed");
|
||||
assert_eq!(created.client_secret_encrypted.as_deref(), Some("secret-1"));
|
||||
|
||||
let updated = repository
|
||||
.upsert_oauth_provider_config(&UpsertOAuthProviderConfigRecord {
|
||||
client_secret_encrypted: EncryptedSecretUpdate::Clear,
|
||||
..sample_upsert("google")
|
||||
})
|
||||
.await
|
||||
.expect("update should succeed");
|
||||
assert!(updated.client_secret_encrypted.is_none());
|
||||
|
||||
let deleted = repository
|
||||
.delete_oauth_provider_config("google")
|
||||
.await
|
||||
.expect("delete should succeed");
|
||||
assert!(deleted);
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/oauth_providers/mod.rs
Normal file
10
crates/aether-data/src/repository/oauth_providers/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryOAuthProviderRepository;
|
||||
pub use sql::SqlxOAuthProviderRepository;
|
||||
pub use types::{
|
||||
EncryptedSecretUpdate, OAuthProviderReadRepository, OAuthProviderRepository,
|
||||
OAuthProviderWriteRepository, StoredOAuthProviderConfig, UpsertOAuthProviderConfigRecord,
|
||||
};
|
||||
341
crates/aether-data/src/repository/oauth_providers/sql.rs
Normal file
341
crates/aether-data/src/repository/oauth_providers/sql.rs
Normal file
@@ -0,0 +1,341 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
OAuthProviderReadRepository, OAuthProviderWriteRepository, StoredOAuthProviderConfig,
|
||||
UpsertOAuthProviderConfigRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const LIST_OAUTH_PROVIDER_CONFIGS_SQL: &str = r#"
|
||||
SELECT
|
||||
provider_type,
|
||||
display_name,
|
||||
client_id,
|
||||
client_secret_encrypted,
|
||||
authorization_url_override,
|
||||
token_url_override,
|
||||
userinfo_url_override,
|
||||
scopes,
|
||||
redirect_uri,
|
||||
frontend_callback_url,
|
||||
attribute_mapping,
|
||||
extra_config,
|
||||
is_enabled,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
FROM oauth_providers
|
||||
ORDER BY provider_type ASC
|
||||
"#;
|
||||
|
||||
const GET_OAUTH_PROVIDER_CONFIG_SQL: &str = r#"
|
||||
SELECT
|
||||
provider_type,
|
||||
display_name,
|
||||
client_id,
|
||||
client_secret_encrypted,
|
||||
authorization_url_override,
|
||||
token_url_override,
|
||||
userinfo_url_override,
|
||||
scopes,
|
||||
redirect_uri,
|
||||
frontend_callback_url,
|
||||
attribute_mapping,
|
||||
extra_config,
|
||||
is_enabled,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
FROM oauth_providers
|
||||
WHERE provider_type = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const COUNT_LOCKED_USERS_IF_PROVIDER_DISABLED_SQL: &str = r#"
|
||||
WITH affected_users AS (
|
||||
SELECT DISTINCT
|
||||
users.id,
|
||||
users.auth_source,
|
||||
users.role,
|
||||
(
|
||||
SELECT COUNT(*)
|
||||
FROM user_oauth_links other_links
|
||||
JOIN oauth_providers other_provider
|
||||
ON other_links.provider_type = other_provider.provider_type
|
||||
WHERE other_links.user_id = users.id
|
||||
AND other_links.provider_type <> $1
|
||||
AND other_provider.is_enabled IS TRUE
|
||||
) AS other_enabled_count
|
||||
FROM users
|
||||
JOIN user_oauth_links
|
||||
ON users.id = user_oauth_links.user_id
|
||||
WHERE users.is_active IS TRUE
|
||||
AND users.is_deleted IS FALSE
|
||||
AND user_oauth_links.provider_type = $1
|
||||
)
|
||||
SELECT COUNT(*)::bigint AS locked_count
|
||||
FROM affected_users
|
||||
WHERE (
|
||||
auth_source = 'oauth'
|
||||
AND other_enabled_count = 0
|
||||
) OR (
|
||||
$2::boolean IS TRUE
|
||||
AND auth_source = 'local'
|
||||
AND role <> 'admin'
|
||||
AND other_enabled_count = 0
|
||||
)
|
||||
"#;
|
||||
|
||||
const UPSERT_OAUTH_PROVIDER_CONFIG_SQL: &str = r#"
|
||||
INSERT INTO oauth_providers (
|
||||
provider_type,
|
||||
display_name,
|
||||
client_id,
|
||||
client_secret_encrypted,
|
||||
authorization_url_override,
|
||||
token_url_override,
|
||||
userinfo_url_override,
|
||||
scopes,
|
||||
redirect_uri,
|
||||
frontend_callback_url,
|
||||
attribute_mapping,
|
||||
extra_config,
|
||||
is_enabled,
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
CASE $4
|
||||
WHEN 'set' THEN $5
|
||||
WHEN 'clear' THEN NULL
|
||||
ELSE NULL
|
||||
END,
|
||||
$6,
|
||||
$7,
|
||||
$8,
|
||||
$9,
|
||||
$10,
|
||||
$11,
|
||||
$12,
|
||||
$13,
|
||||
$14,
|
||||
NOW(),
|
||||
NOW()
|
||||
)
|
||||
ON CONFLICT (provider_type) DO UPDATE
|
||||
SET display_name = EXCLUDED.display_name,
|
||||
client_id = EXCLUDED.client_id,
|
||||
client_secret_encrypted = CASE $4
|
||||
WHEN 'set' THEN $5
|
||||
WHEN 'clear' THEN NULL
|
||||
ELSE oauth_providers.client_secret_encrypted
|
||||
END,
|
||||
authorization_url_override = EXCLUDED.authorization_url_override,
|
||||
token_url_override = EXCLUDED.token_url_override,
|
||||
userinfo_url_override = EXCLUDED.userinfo_url_override,
|
||||
scopes = EXCLUDED.scopes,
|
||||
redirect_uri = EXCLUDED.redirect_uri,
|
||||
frontend_callback_url = EXCLUDED.frontend_callback_url,
|
||||
attribute_mapping = EXCLUDED.attribute_mapping,
|
||||
extra_config = EXCLUDED.extra_config,
|
||||
is_enabled = EXCLUDED.is_enabled,
|
||||
updated_at = NOW()
|
||||
RETURNING
|
||||
provider_type,
|
||||
display_name,
|
||||
client_id,
|
||||
client_secret_encrypted,
|
||||
authorization_url_override,
|
||||
token_url_override,
|
||||
userinfo_url_override,
|
||||
scopes,
|
||||
redirect_uri,
|
||||
frontend_callback_url,
|
||||
attribute_mapping,
|
||||
extra_config,
|
||||
is_enabled,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const DELETE_OAUTH_PROVIDER_CONFIG_SQL: &str = r#"
|
||||
DELETE FROM oauth_providers
|
||||
WHERE provider_type = $1
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxOAuthProviderRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxOAuthProviderRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl OAuthProviderReadRepository for SqlxOAuthProviderRepository {
|
||||
async fn list_oauth_provider_configs(
|
||||
&self,
|
||||
) -> Result<Vec<StoredOAuthProviderConfig>, DataLayerError> {
|
||||
let rows = sqlx::query(LIST_OAUTH_PROVIDER_CONFIGS_SQL)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_oauth_provider_row).collect()
|
||||
}
|
||||
|
||||
async fn get_oauth_provider_config(
|
||||
&self,
|
||||
provider_type: &str,
|
||||
) -> Result<Option<StoredOAuthProviderConfig>, DataLayerError> {
|
||||
let row = sqlx::query(GET_OAUTH_PROVIDER_CONFIG_SQL)
|
||||
.bind(provider_type)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_oauth_provider_row).transpose()
|
||||
}
|
||||
|
||||
async fn count_locked_users_if_provider_disabled(
|
||||
&self,
|
||||
provider_type: &str,
|
||||
ldap_exclusive: bool,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
let locked_count: i64 = sqlx::query_scalar(COUNT_LOCKED_USERS_IF_PROVIDER_DISABLED_SQL)
|
||||
.bind(provider_type)
|
||||
.bind(ldap_exclusive)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
usize::try_from(locked_count).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"oauth_providers.locked_user_count is negative".to_string(),
|
||||
)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl OAuthProviderWriteRepository for SqlxOAuthProviderRepository {
|
||||
async fn upsert_oauth_provider_config(
|
||||
&self,
|
||||
record: &UpsertOAuthProviderConfigRecord,
|
||||
) -> Result<StoredOAuthProviderConfig, DataLayerError> {
|
||||
record.validate()?;
|
||||
let row = sqlx::query(UPSERT_OAUTH_PROVIDER_CONFIG_SQL)
|
||||
.bind(&record.provider_type)
|
||||
.bind(&record.display_name)
|
||||
.bind(&record.client_id)
|
||||
.bind(record.client_secret_encrypted.mode_name())
|
||||
.bind(record.client_secret_encrypted.value())
|
||||
.bind(record.authorization_url_override.as_deref())
|
||||
.bind(record.token_url_override.as_deref())
|
||||
.bind(record.userinfo_url_override.as_deref())
|
||||
.bind(scopes_to_json(record.scopes.as_ref()))
|
||||
.bind(&record.redirect_uri)
|
||||
.bind(&record.frontend_callback_url)
|
||||
.bind(record.attribute_mapping.as_ref())
|
||||
.bind(record.extra_config.as_ref())
|
||||
.bind(record.is_enabled)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
map_oauth_provider_row(&row)
|
||||
}
|
||||
|
||||
async fn delete_oauth_provider_config(
|
||||
&self,
|
||||
provider_type: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let result = sqlx::query(DELETE_OAUTH_PROVIDER_CONFIG_SQL)
|
||||
.bind(provider_type)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_unix_secs(value: Option<i64>) -> Option<u64> {
|
||||
value.and_then(|value| u64::try_from(value).ok())
|
||||
}
|
||||
|
||||
fn scopes_to_json(scopes: Option<&Vec<String>>) -> Option<serde_json::Value> {
|
||||
scopes.map(|items| {
|
||||
serde_json::Value::Array(
|
||||
items
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(serde_json::Value::String)
|
||||
.collect(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_scopes(value: Option<serde_json::Value>) -> Result<Option<Vec<String>>, DataLayerError> {
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
let serde_json::Value::Array(items) = value else {
|
||||
return Err(DataLayerError::UnexpectedValue(
|
||||
"oauth_providers.scopes is not a JSON array".to_string(),
|
||||
));
|
||||
};
|
||||
let mut scopes = Vec::with_capacity(items.len());
|
||||
for item in items {
|
||||
let serde_json::Value::String(scope) = item else {
|
||||
return Err(DataLayerError::UnexpectedValue(
|
||||
"oauth_providers.scopes contains non-string value".to_string(),
|
||||
));
|
||||
};
|
||||
scopes.push(scope);
|
||||
}
|
||||
Ok(Some(scopes))
|
||||
}
|
||||
|
||||
fn map_oauth_provider_row(row: &PgRow) -> Result<StoredOAuthProviderConfig, DataLayerError> {
|
||||
Ok(StoredOAuthProviderConfig::new(
|
||||
row.try_get("provider_type")?,
|
||||
row.try_get("display_name")?,
|
||||
row.try_get("client_id")?,
|
||||
row.try_get("redirect_uri")?,
|
||||
row.try_get("frontend_callback_url")?,
|
||||
)?
|
||||
.with_config_fields(
|
||||
row.try_get("client_secret_encrypted")?,
|
||||
row.try_get("authorization_url_override")?,
|
||||
row.try_get("token_url_override")?,
|
||||
row.try_get("userinfo_url_override")?,
|
||||
parse_scopes(row.try_get("scopes")?)?,
|
||||
row.try_get("attribute_mapping")?,
|
||||
row.try_get("extra_config")?,
|
||||
row.try_get("is_enabled")?,
|
||||
)
|
||||
.with_timestamps(
|
||||
optional_unix_secs(row.try_get("created_at_unix_secs")?),
|
||||
optional_unix_secs(row.try_get("updated_at_unix_secs")?),
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxOAuthProviderRepository;
|
||||
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 = SqlxOAuthProviderRepository::new(pool);
|
||||
}
|
||||
}
|
||||
230
crates/aether-data/src/repository/oauth_providers/types.rs
Normal file
230
crates/aether-data/src/repository/oauth_providers/types.rs
Normal file
@@ -0,0 +1,230 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredOAuthProviderConfig {
|
||||
pub provider_type: String,
|
||||
pub display_name: String,
|
||||
pub client_id: String,
|
||||
pub client_secret_encrypted: Option<String>,
|
||||
pub authorization_url_override: Option<String>,
|
||||
pub token_url_override: Option<String>,
|
||||
pub userinfo_url_override: Option<String>,
|
||||
pub scopes: Option<Vec<String>>,
|
||||
pub redirect_uri: String,
|
||||
pub frontend_callback_url: String,
|
||||
pub attribute_mapping: Option<serde_json::Value>,
|
||||
pub extra_config: Option<serde_json::Value>,
|
||||
pub is_enabled: bool,
|
||||
pub created_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl StoredOAuthProviderConfig {
|
||||
pub fn new(
|
||||
provider_type: String,
|
||||
display_name: String,
|
||||
client_id: String,
|
||||
redirect_uri: String,
|
||||
frontend_callback_url: String,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if provider_type.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"oauth_providers.provider_type is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if display_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"oauth_providers.display_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if client_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"oauth_providers.client_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if redirect_uri.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"oauth_providers.redirect_uri is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if frontend_callback_url.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"oauth_providers.frontend_callback_url is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
provider_type,
|
||||
display_name,
|
||||
client_id,
|
||||
client_secret_encrypted: None,
|
||||
authorization_url_override: None,
|
||||
token_url_override: None,
|
||||
userinfo_url_override: None,
|
||||
scopes: None,
|
||||
redirect_uri,
|
||||
frontend_callback_url,
|
||||
attribute_mapping: None,
|
||||
extra_config: None,
|
||||
is_enabled: false,
|
||||
created_at_unix_secs: None,
|
||||
updated_at_unix_secs: None,
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn with_config_fields(
|
||||
mut self,
|
||||
client_secret_encrypted: Option<String>,
|
||||
authorization_url_override: Option<String>,
|
||||
token_url_override: Option<String>,
|
||||
userinfo_url_override: Option<String>,
|
||||
scopes: Option<Vec<String>>,
|
||||
attribute_mapping: Option<serde_json::Value>,
|
||||
extra_config: Option<serde_json::Value>,
|
||||
is_enabled: bool,
|
||||
) -> Self {
|
||||
self.client_secret_encrypted = client_secret_encrypted;
|
||||
self.authorization_url_override = authorization_url_override;
|
||||
self.token_url_override = token_url_override;
|
||||
self.userinfo_url_override = userinfo_url_override;
|
||||
self.scopes = scopes;
|
||||
self.attribute_mapping = attribute_mapping;
|
||||
self.extra_config = extra_config;
|
||||
self.is_enabled = is_enabled;
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, Default)]
|
||||
pub enum EncryptedSecretUpdate {
|
||||
#[default]
|
||||
Preserve,
|
||||
Clear,
|
||||
Set(String),
|
||||
}
|
||||
|
||||
impl EncryptedSecretUpdate {
|
||||
pub fn mode_name(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Preserve => "preserve",
|
||||
Self::Clear => "clear",
|
||||
Self::Set(_) => "set",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn value(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::Set(value) => Some(value.as_str()),
|
||||
Self::Preserve | Self::Clear => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct UpsertOAuthProviderConfigRecord {
|
||||
pub provider_type: String,
|
||||
pub display_name: String,
|
||||
pub client_id: String,
|
||||
pub client_secret_encrypted: EncryptedSecretUpdate,
|
||||
pub authorization_url_override: Option<String>,
|
||||
pub token_url_override: Option<String>,
|
||||
pub userinfo_url_override: Option<String>,
|
||||
pub scopes: Option<Vec<String>>,
|
||||
pub redirect_uri: String,
|
||||
pub frontend_callback_url: String,
|
||||
pub attribute_mapping: Option<serde_json::Value>,
|
||||
pub extra_config: Option<serde_json::Value>,
|
||||
pub is_enabled: bool,
|
||||
}
|
||||
|
||||
impl UpsertOAuthProviderConfigRecord {
|
||||
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||
if self.provider_type.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"provider_type is required".to_string(),
|
||||
));
|
||||
}
|
||||
if self.display_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"display_name is required".to_string(),
|
||||
));
|
||||
}
|
||||
if self.client_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"client_id is required".to_string(),
|
||||
));
|
||||
}
|
||||
if self.redirect_uri.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"redirect_uri is required".to_string(),
|
||||
));
|
||||
}
|
||||
if self.frontend_callback_url.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"frontend_callback_url is required".to_string(),
|
||||
));
|
||||
}
|
||||
if let Some(scopes) = &self.scopes {
|
||||
for scope in scopes {
|
||||
if scope.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"scopes must not contain empty values".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait OAuthProviderReadRepository: Send + Sync {
|
||||
async fn list_oauth_provider_configs(
|
||||
&self,
|
||||
) -> Result<Vec<StoredOAuthProviderConfig>, crate::DataLayerError>;
|
||||
|
||||
async fn get_oauth_provider_config(
|
||||
&self,
|
||||
provider_type: &str,
|
||||
) -> Result<Option<StoredOAuthProviderConfig>, crate::DataLayerError>;
|
||||
|
||||
async fn count_locked_users_if_provider_disabled(
|
||||
&self,
|
||||
provider_type: &str,
|
||||
ldap_exclusive: bool,
|
||||
) -> Result<usize, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait OAuthProviderWriteRepository: Send + Sync {
|
||||
async fn upsert_oauth_provider_config(
|
||||
&self,
|
||||
record: &UpsertOAuthProviderConfigRecord,
|
||||
) -> Result<StoredOAuthProviderConfig, crate::DataLayerError>;
|
||||
|
||||
async fn delete_oauth_provider_config(
|
||||
&self,
|
||||
provider_type: &str,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait OAuthProviderRepository:
|
||||
OAuthProviderReadRepository + OAuthProviderWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> OAuthProviderRepository for T where
|
||||
T: OAuthProviderReadRepository + OAuthProviderWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
@@ -4,8 +4,9 @@ use std::sync::RwLock;
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
|
||||
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -45,6 +46,25 @@ impl InMemoryProviderCatalogReadRepository {
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
|
||||
async fn list_providers(
|
||||
&self,
|
||||
active_only: bool,
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
let index = self.index.read().expect("provider catalog repository lock");
|
||||
let mut providers = index
|
||||
.providers
|
||||
.values()
|
||||
.filter(|provider| !active_only || provider.is_active)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
providers.sort_by(|left, right| {
|
||||
left.provider_priority
|
||||
.cmp(&right.provider_priority)
|
||||
.then(left.name.cmp(&right.name))
|
||||
});
|
||||
Ok(providers)
|
||||
}
|
||||
|
||||
async fn list_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
@@ -67,6 +87,30 @@ impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_endpoints_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
let index = self.index.read().expect("provider catalog repository lock");
|
||||
let mut endpoints = index
|
||||
.endpoints
|
||||
.values()
|
||||
.filter(|endpoint| {
|
||||
provider_ids
|
||||
.iter()
|
||||
.any(|provider_id| provider_id == &endpoint.provider_id)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
endpoints.sort_by(|left, right| {
|
||||
left.provider_id
|
||||
.cmp(&right.provider_id)
|
||||
.then(left.api_format.cmp(&right.api_format))
|
||||
.then(left.id.cmp(&right.id))
|
||||
});
|
||||
Ok(endpoints)
|
||||
}
|
||||
|
||||
async fn list_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
@@ -77,14 +121,302 @@ impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
|
||||
.filter_map(|id| index.keys.get(id).cloned())
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_keys_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
let index = self.index.read().expect("provider catalog repository lock");
|
||||
let mut keys = index
|
||||
.keys
|
||||
.values()
|
||||
.filter(|key| {
|
||||
provider_ids
|
||||
.iter()
|
||||
.any(|provider_id| provider_id == &key.provider_id)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
keys.sort_by(|left, right| {
|
||||
left.provider_id
|
||||
.cmp(&right.provider_id)
|
||||
.then(left.name.cmp(&right.name))
|
||||
.then(left.id.cmp(&right.id))
|
||||
});
|
||||
Ok(keys)
|
||||
}
|
||||
|
||||
async fn list_keys_page(
|
||||
&self,
|
||||
query: &ProviderCatalogKeyListQuery,
|
||||
) -> Result<StoredProviderCatalogKeyPage, DataLayerError> {
|
||||
let index = self.index.read().expect("provider catalog repository lock");
|
||||
let mut keys = index
|
||||
.keys
|
||||
.values()
|
||||
.filter(|key| key.provider_id == query.provider_id)
|
||||
.filter(|key| {
|
||||
query.search.as_ref().is_none_or(|keyword| {
|
||||
let keyword = keyword.trim().to_ascii_lowercase();
|
||||
keyword.is_empty()
|
||||
|| key.name.to_ascii_lowercase().contains(&keyword)
|
||||
|| key.id.to_ascii_lowercase().contains(&keyword)
|
||||
})
|
||||
})
|
||||
.filter(|key| {
|
||||
query
|
||||
.is_active
|
||||
.is_none_or(|is_active| key.is_active == is_active)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
keys.sort_by(|left, right| {
|
||||
left.internal_priority
|
||||
.cmp(&right.internal_priority)
|
||||
.then(left.name.cmp(&right.name))
|
||||
.then(left.id.cmp(&right.id))
|
||||
});
|
||||
let total = keys.len();
|
||||
let items = keys
|
||||
.into_iter()
|
||||
.skip(query.offset)
|
||||
.take(query.limit)
|
||||
.collect();
|
||||
Ok(StoredProviderCatalogKeyPage { items, total })
|
||||
}
|
||||
|
||||
async fn list_key_stats_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKeyStats>, DataLayerError> {
|
||||
let index = self.index.read().expect("provider catalog repository lock");
|
||||
let mut stats = provider_ids
|
||||
.iter()
|
||||
.map(|provider_id| {
|
||||
let total_keys = index
|
||||
.keys
|
||||
.values()
|
||||
.filter(|key| &key.provider_id == provider_id)
|
||||
.count() as i64;
|
||||
let active_keys = index
|
||||
.keys
|
||||
.values()
|
||||
.filter(|key| &key.provider_id == provider_id && key.is_active)
|
||||
.count() as i64;
|
||||
StoredProviderCatalogKeyStats::new(provider_id.clone(), total_keys, active_keys)
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
stats.retain(|item| item.total_keys > 0);
|
||||
Ok(stats)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository {
|
||||
async fn create_provider(
|
||||
&self,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
shift_existing_priorities_from: Option<i32>,
|
||||
) -> Result<StoredProviderCatalogProvider, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("provider catalog repository lock");
|
||||
if let Some(target_priority) = shift_existing_priorities_from {
|
||||
for existing in index.providers.values_mut() {
|
||||
if existing.provider_priority >= target_priority {
|
||||
existing.provider_priority += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
index
|
||||
.providers
|
||||
.insert(provider.id.clone(), provider.clone());
|
||||
Ok(provider.clone())
|
||||
}
|
||||
|
||||
async fn update_provider(
|
||||
&self,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
) -> Result<StoredProviderCatalogProvider, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("provider catalog repository lock");
|
||||
let Some(stored) = index.providers.get_mut(&provider.id) else {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"provider catalog provider {} not found",
|
||||
provider.id
|
||||
)));
|
||||
};
|
||||
*stored = provider.clone();
|
||||
Ok(stored.clone())
|
||||
}
|
||||
|
||||
async fn delete_provider(&self, provider_id: &str) -> Result<bool, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("provider catalog repository lock");
|
||||
Ok(index.providers.remove(provider_id).is_some())
|
||||
}
|
||||
|
||||
async fn cleanup_deleted_provider_refs(
|
||||
&self,
|
||||
_provider_id: &str,
|
||||
_endpoint_ids: &[String],
|
||||
_key_ids: &[String],
|
||||
) -> Result<(), DataLayerError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_endpoint(
|
||||
&self,
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
) -> Result<StoredProviderCatalogEndpoint, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("provider catalog repository lock");
|
||||
index
|
||||
.endpoints
|
||||
.insert(endpoint.id.clone(), endpoint.clone());
|
||||
Ok(endpoint.clone())
|
||||
}
|
||||
|
||||
async fn update_endpoint(
|
||||
&self,
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
) -> Result<StoredProviderCatalogEndpoint, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("provider catalog repository lock");
|
||||
let Some(stored) = index.endpoints.get_mut(&endpoint.id) else {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"provider catalog endpoint {} not found",
|
||||
endpoint.id
|
||||
)));
|
||||
};
|
||||
*stored = endpoint.clone();
|
||||
Ok(stored.clone())
|
||||
}
|
||||
|
||||
async fn delete_endpoint(&self, endpoint_id: &str) -> Result<bool, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("provider catalog repository lock");
|
||||
Ok(index.endpoints.remove(endpoint_id).is_some())
|
||||
}
|
||||
|
||||
async fn create_key(
|
||||
&self,
|
||||
key: &StoredProviderCatalogKey,
|
||||
) -> Result<StoredProviderCatalogKey, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("provider catalog repository lock");
|
||||
index.keys.insert(key.id.clone(), key.clone());
|
||||
Ok(key.clone())
|
||||
}
|
||||
|
||||
async fn update_key(
|
||||
&self,
|
||||
key: &StoredProviderCatalogKey,
|
||||
) -> Result<StoredProviderCatalogKey, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("provider catalog repository lock");
|
||||
let Some(stored) = index.keys.get_mut(&key.id) else {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"provider catalog key {} not found",
|
||||
key.id
|
||||
)));
|
||||
};
|
||||
*stored = key.clone();
|
||||
Ok(stored.clone())
|
||||
}
|
||||
|
||||
async fn delete_key(&self, key_id: &str) -> Result<bool, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("provider catalog repository lock");
|
||||
Ok(index.keys.remove(key_id).is_some())
|
||||
}
|
||||
|
||||
async fn clear_key_oauth_invalid_marker(&self, key_id: &str) -> Result<bool, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("provider catalog repository lock");
|
||||
let Some(key) = index.keys.get_mut(key_id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
|
||||
key.oauth_invalid_at_unix_secs = None;
|
||||
key.oauth_invalid_reason = None;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn update_key_oauth_credentials(
|
||||
&self,
|
||||
key_id: &str,
|
||||
encrypted_api_key: &str,
|
||||
encrypted_auth_config: Option<&str>,
|
||||
expires_at_unix_secs: Option<u64>,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
if encrypted_api_key.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"provider catalog oauth api_key is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("provider catalog repository lock");
|
||||
let Some(key) = index.keys.get_mut(key_id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
|
||||
key.encrypted_api_key = encrypted_api_key.to_string();
|
||||
key.encrypted_auth_config = encrypted_auth_config.map(ToOwned::to_owned);
|
||||
key.expires_at_unix_secs = expires_at_unix_secs;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn update_key_health_state(
|
||||
&self,
|
||||
key_id: &str,
|
||||
is_active: bool,
|
||||
health_by_format: Option<&serde_json::Value>,
|
||||
circuit_breaker_by_format: Option<&serde_json::Value>,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("provider catalog repository lock");
|
||||
let Some(key) = index.keys.get_mut(key_id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
|
||||
key.is_active = is_active;
|
||||
key.health_by_format = health_by_format.cloned();
|
||||
key.circuit_breaker_by_format = circuit_breaker_by_format.cloned();
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryProviderCatalogReadRepository;
|
||||
use crate::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
|
||||
fn sample_provider(id: &str) -> StoredProviderCatalogProvider {
|
||||
@@ -107,6 +439,7 @@ mod tests {
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_health_score(0.9)
|
||||
}
|
||||
|
||||
fn sample_key(id: &str, provider_id: &str) -> StoredProviderCatalogKey {
|
||||
@@ -154,4 +487,286 @@ mod tests {
|
||||
1
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lists_active_providers_in_priority_order() {
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider("provider-2").with_routing_fields(20),
|
||||
sample_provider("provider-1").with_routing_fields(10),
|
||||
sample_provider("provider-3")
|
||||
.with_routing_fields(5)
|
||||
.with_transport_fields(false, false, false, None, None, None, None, None, None),
|
||||
],
|
||||
vec![],
|
||||
vec![],
|
||||
);
|
||||
|
||||
let providers = repository
|
||||
.list_providers(true)
|
||||
.await
|
||||
.expect("providers should list");
|
||||
assert_eq!(
|
||||
providers
|
||||
.iter()
|
||||
.map(|provider| provider.id.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
vec!["provider-1", "provider-2"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn updates_oauth_credentials_for_existing_key() {
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1")],
|
||||
vec![sample_endpoint("endpoint-1", "provider-1")],
|
||||
vec![sample_key("key-1", "provider-1")
|
||||
.with_transport_fields(
|
||||
None,
|
||||
"ciphertext-placeholder".to_string(),
|
||||
Some("ciphertext-auth-1".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")],
|
||||
);
|
||||
|
||||
assert!(repository
|
||||
.update_key_oauth_credentials(
|
||||
"key-1",
|
||||
"ciphertext-updated-token",
|
||||
Some("ciphertext-auth-2"),
|
||||
Some(4_102_444_800),
|
||||
)
|
||||
.await
|
||||
.expect("update should succeed"));
|
||||
|
||||
let stored = repository
|
||||
.list_keys_by_ids(&["key-1".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert_eq!(stored.len(), 1);
|
||||
assert_eq!(stored[0].encrypted_api_key, "ciphertext-updated-token");
|
||||
assert_eq!(
|
||||
stored[0].encrypted_auth_config.as_deref(),
|
||||
Some("ciphertext-auth-2")
|
||||
);
|
||||
assert_eq!(stored[0].expires_at_unix_secs, Some(4_102_444_800));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn paginates_provider_keys_with_search_and_active_filter() {
|
||||
let mut alpha = sample_key("key-1", "provider-1");
|
||||
alpha.name = "alpha".to_string();
|
||||
alpha.internal_priority = 20;
|
||||
let mut beta = sample_key("key-2", "provider-1");
|
||||
beta.name = "beta".to_string();
|
||||
beta.internal_priority = 10;
|
||||
let mut gamma = sample_key("key-3", "provider-1");
|
||||
gamma.name = "gamma".to_string();
|
||||
gamma.internal_priority = 30;
|
||||
gamma.is_active = false;
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1"), sample_provider("provider-2")],
|
||||
vec![],
|
||||
vec![alpha, beta, gamma, sample_key("key-4", "provider-2")],
|
||||
);
|
||||
|
||||
let page = repository
|
||||
.list_keys_page(&ProviderCatalogKeyListQuery {
|
||||
provider_id: "provider-1".to_string(),
|
||||
search: Some("a".to_string()),
|
||||
is_active: Some(true),
|
||||
offset: 0,
|
||||
limit: 10,
|
||||
})
|
||||
.await
|
||||
.expect("keys should page");
|
||||
|
||||
assert_eq!(page.total, 2);
|
||||
assert_eq!(page.items.len(), 2);
|
||||
assert_eq!(
|
||||
page.items
|
||||
.iter()
|
||||
.map(|item| item.name.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
vec!["beta", "alpha"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn summarizes_provider_key_stats() {
|
||||
let mut inactive = sample_key("key-2", "provider-1");
|
||||
inactive.is_active = false;
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1"), sample_provider("provider-2")],
|
||||
vec![],
|
||||
vec![
|
||||
sample_key("key-1", "provider-1"),
|
||||
inactive,
|
||||
sample_key("key-3", "provider-2"),
|
||||
],
|
||||
);
|
||||
|
||||
let stats = repository
|
||||
.list_key_stats_by_provider_ids(&["provider-1".to_string(), "provider-2".to_string()])
|
||||
.await
|
||||
.expect("stats should list");
|
||||
assert_eq!(stats.len(), 2);
|
||||
assert_eq!(stats[0].provider_id, "provider-1");
|
||||
assert_eq!(stats[0].total_keys, 2);
|
||||
assert_eq!(stats[0].active_keys, 1);
|
||||
assert_eq!(stats[1].provider_id, "provider-2");
|
||||
assert_eq!(stats[1].total_keys, 1);
|
||||
assert_eq!(stats[1].active_keys, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn creates_key() {
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1")],
|
||||
vec![],
|
||||
vec![],
|
||||
);
|
||||
let key = sample_key("key-1", "provider-1");
|
||||
|
||||
let created = repository
|
||||
.create_key(&key)
|
||||
.await
|
||||
.expect("key should create");
|
||||
|
||||
assert_eq!(created.id, "key-1");
|
||||
let stored = repository
|
||||
.list_keys_by_ids(&["key-1".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert_eq!(stored.len(), 1);
|
||||
assert_eq!(stored[0].provider_id, "provider-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn creates_endpoint() {
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1")],
|
||||
vec![],
|
||||
vec![],
|
||||
);
|
||||
let endpoint = sample_endpoint("endpoint-1", "provider-1");
|
||||
|
||||
let created = repository
|
||||
.create_endpoint(&endpoint)
|
||||
.await
|
||||
.expect("endpoint should create");
|
||||
|
||||
assert_eq!(created.id, "endpoint-1");
|
||||
let stored = repository
|
||||
.list_endpoints_by_ids(&["endpoint-1".to_string()])
|
||||
.await
|
||||
.expect("endpoints should read");
|
||||
assert_eq!(stored.len(), 1);
|
||||
assert_eq!(stored[0].provider_id, "provider-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn updates_key() {
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1")],
|
||||
vec![],
|
||||
vec![sample_key("key-1", "provider-1")],
|
||||
);
|
||||
let mut updated = sample_key("key-1", "provider-1");
|
||||
updated.name = "updated".to_string();
|
||||
updated.internal_priority = 7;
|
||||
|
||||
let stored = repository
|
||||
.update_key(&updated)
|
||||
.await
|
||||
.expect("key should update");
|
||||
|
||||
assert_eq!(stored.name, "updated");
|
||||
assert_eq!(stored.internal_priority, 7);
|
||||
let reloaded = repository
|
||||
.list_keys_by_ids(&["key-1".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert_eq!(reloaded[0].name, "updated");
|
||||
assert_eq!(reloaded[0].internal_priority, 7);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn updates_endpoint() {
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1")],
|
||||
vec![sample_endpoint("endpoint-1", "provider-1")],
|
||||
vec![],
|
||||
);
|
||||
let updated = sample_endpoint("endpoint-1", "provider-1")
|
||||
.with_transport_fields(
|
||||
"https://updated.example".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(5),
|
||||
Some("/v1/chat/completions".to_string()),
|
||||
Some(serde_json::json!({"foo":"bar"})),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build");
|
||||
|
||||
let stored = repository
|
||||
.update_endpoint(&updated)
|
||||
.await
|
||||
.expect("endpoint should update");
|
||||
|
||||
assert_eq!(stored.base_url, "https://updated.example");
|
||||
assert_eq!(stored.max_retries, Some(5));
|
||||
let reloaded = repository
|
||||
.list_endpoints_by_ids(&["endpoint-1".to_string()])
|
||||
.await
|
||||
.expect("endpoints should read");
|
||||
assert_eq!(reloaded[0].base_url, "https://updated.example");
|
||||
assert_eq!(reloaded[0].max_retries, Some(5));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deletes_key() {
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1")],
|
||||
vec![],
|
||||
vec![sample_key("key-1", "provider-1")],
|
||||
);
|
||||
|
||||
assert!(repository
|
||||
.delete_key("key-1")
|
||||
.await
|
||||
.expect("delete should succeed"));
|
||||
let reloaded = repository
|
||||
.list_keys_by_ids(&["key-1".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert!(reloaded.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deletes_endpoint() {
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1")],
|
||||
vec![sample_endpoint("endpoint-1", "provider-1")],
|
||||
vec![],
|
||||
);
|
||||
|
||||
assert!(repository
|
||||
.delete_endpoint("endpoint-1")
|
||||
.await
|
||||
.expect("delete should succeed"));
|
||||
let reloaded = repository
|
||||
.list_endpoints_by_ids(&["endpoint-1".to_string()])
|
||||
.await
|
||||
.expect("endpoints should read");
|
||||
assert!(reloaded.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ mod types;
|
||||
pub use memory::InMemoryProviderCatalogReadRepository;
|
||||
pub use sql::SqlxProviderCatalogReadRepository;
|
||||
pub use types::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
|
||||
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||
};
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,11 +1,30 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderCatalogProvider {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub website: Option<String>,
|
||||
pub provider_type: String,
|
||||
pub billing_type: Option<String>,
|
||||
pub monthly_quota_usd: Option<f64>,
|
||||
pub monthly_used_usd: Option<f64>,
|
||||
pub quota_reset_day: Option<u64>,
|
||||
pub quota_last_reset_at_unix_secs: Option<u64>,
|
||||
pub quota_expires_at_unix_secs: Option<u64>,
|
||||
pub provider_priority: i32,
|
||||
pub is_active: bool,
|
||||
pub keep_priority_on_conversion: bool,
|
||||
pub enable_format_conversion: bool,
|
||||
pub concurrent_limit: Option<i32>,
|
||||
pub max_retries: Option<i32>,
|
||||
pub proxy: Option<serde_json::Value>,
|
||||
pub request_timeout_secs: Option<f64>,
|
||||
pub stream_first_byte_timeout_secs: Option<f64>,
|
||||
pub config: Option<serde_json::Value>,
|
||||
pub created_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl StoredProviderCatalogProvider {
|
||||
@@ -29,13 +48,96 @@ impl StoredProviderCatalogProvider {
|
||||
Ok(Self {
|
||||
id,
|
||||
name,
|
||||
description: None,
|
||||
website,
|
||||
provider_type,
|
||||
billing_type: None,
|
||||
monthly_quota_usd: None,
|
||||
monthly_used_usd: None,
|
||||
quota_reset_day: None,
|
||||
quota_last_reset_at_unix_secs: None,
|
||||
quota_expires_at_unix_secs: None,
|
||||
provider_priority: 0,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: false,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
created_at_unix_secs: None,
|
||||
updated_at_unix_secs: None,
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn with_transport_fields(
|
||||
mut self,
|
||||
is_active: bool,
|
||||
keep_priority_on_conversion: bool,
|
||||
enable_format_conversion: bool,
|
||||
concurrent_limit: Option<i32>,
|
||||
max_retries: Option<i32>,
|
||||
proxy: Option<serde_json::Value>,
|
||||
request_timeout_secs: Option<f64>,
|
||||
stream_first_byte_timeout_secs: Option<f64>,
|
||||
config: Option<serde_json::Value>,
|
||||
) -> Self {
|
||||
self.is_active = is_active;
|
||||
self.keep_priority_on_conversion = keep_priority_on_conversion;
|
||||
self.enable_format_conversion = enable_format_conversion;
|
||||
self.concurrent_limit = concurrent_limit;
|
||||
self.max_retries = max_retries;
|
||||
self.proxy = proxy;
|
||||
self.request_timeout_secs = request_timeout_secs;
|
||||
self.stream_first_byte_timeout_secs = stream_first_byte_timeout_secs;
|
||||
self.config = config;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_description(mut self, description: Option<String>) -> Self {
|
||||
self.description = description;
|
||||
self
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn with_billing_fields(
|
||||
mut self,
|
||||
billing_type: Option<String>,
|
||||
monthly_quota_usd: Option<f64>,
|
||||
monthly_used_usd: Option<f64>,
|
||||
quota_reset_day: Option<u64>,
|
||||
quota_last_reset_at_unix_secs: Option<u64>,
|
||||
quota_expires_at_unix_secs: Option<u64>,
|
||||
) -> Self {
|
||||
self.billing_type = billing_type;
|
||||
self.monthly_quota_usd = monthly_quota_usd;
|
||||
self.monthly_used_usd = monthly_used_usd;
|
||||
self.quota_reset_day = quota_reset_day;
|
||||
self.quota_last_reset_at_unix_secs = quota_last_reset_at_unix_secs;
|
||||
self.quota_expires_at_unix_secs = quota_expires_at_unix_secs;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_routing_fields(mut self, provider_priority: i32) -> Self {
|
||||
self.provider_priority = provider_priority;
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderCatalogEndpoint {
|
||||
pub id: String,
|
||||
pub provider_id: String,
|
||||
@@ -43,6 +145,17 @@ pub struct StoredProviderCatalogEndpoint {
|
||||
pub api_family: Option<String>,
|
||||
pub endpoint_kind: Option<String>,
|
||||
pub is_active: bool,
|
||||
pub health_score: f64,
|
||||
pub base_url: String,
|
||||
pub header_rules: Option<serde_json::Value>,
|
||||
pub body_rules: Option<serde_json::Value>,
|
||||
pub max_retries: Option<i32>,
|
||||
pub custom_path: Option<String>,
|
||||
pub config: Option<serde_json::Value>,
|
||||
pub format_acceptance_config: Option<serde_json::Value>,
|
||||
pub proxy: Option<serde_json::Value>,
|
||||
pub created_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl StoredProviderCatalogEndpoint {
|
||||
@@ -67,11 +180,66 @@ impl StoredProviderCatalogEndpoint {
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
is_active,
|
||||
health_score: 1.0,
|
||||
base_url: String::new(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
created_at_unix_secs: None,
|
||||
updated_at_unix_secs: None,
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn with_transport_fields(
|
||||
mut self,
|
||||
base_url: String,
|
||||
header_rules: Option<serde_json::Value>,
|
||||
body_rules: Option<serde_json::Value>,
|
||||
max_retries: Option<i32>,
|
||||
custom_path: Option<String>,
|
||||
config: Option<serde_json::Value>,
|
||||
format_acceptance_config: Option<serde_json::Value>,
|
||||
proxy: Option<serde_json::Value>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if base_url.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider_endpoints.base_url is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
self.base_url = base_url;
|
||||
self.header_rules = header_rules;
|
||||
self.body_rules = body_rules;
|
||||
self.max_retries = max_retries;
|
||||
self.custom_path = custom_path;
|
||||
self.config = config;
|
||||
self.format_acceptance_config = format_acceptance_config;
|
||||
self.proxy = proxy;
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
pub fn with_health_score(mut self, health_score: f64) -> Self {
|
||||
self.health_score = health_score;
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderCatalogKey {
|
||||
pub id: String,
|
||||
pub provider_id: String,
|
||||
@@ -79,6 +247,47 @@ pub struct StoredProviderCatalogKey {
|
||||
pub auth_type: String,
|
||||
pub capabilities: Option<serde_json::Value>,
|
||||
pub is_active: bool,
|
||||
pub api_formats: Option<serde_json::Value>,
|
||||
pub encrypted_api_key: String,
|
||||
pub encrypted_auth_config: Option<String>,
|
||||
pub note: Option<String>,
|
||||
pub internal_priority: i32,
|
||||
pub rate_multipliers: Option<serde_json::Value>,
|
||||
pub global_priority_by_format: Option<serde_json::Value>,
|
||||
pub allowed_models: Option<serde_json::Value>,
|
||||
pub expires_at_unix_secs: Option<u64>,
|
||||
pub cache_ttl_minutes: i32,
|
||||
pub max_probe_interval_minutes: i32,
|
||||
pub proxy: Option<serde_json::Value>,
|
||||
pub fingerprint: Option<serde_json::Value>,
|
||||
pub rpm_limit: Option<u32>,
|
||||
pub learned_rpm_limit: Option<u32>,
|
||||
pub concurrent_429_count: Option<u32>,
|
||||
pub rpm_429_count: Option<u32>,
|
||||
pub last_429_at_unix_secs: Option<u64>,
|
||||
pub last_429_type: Option<String>,
|
||||
pub adjustment_history: Option<serde_json::Value>,
|
||||
pub utilization_samples: Option<serde_json::Value>,
|
||||
pub last_probe_increase_at_unix_secs: Option<u64>,
|
||||
pub request_count: Option<u32>,
|
||||
pub success_count: Option<u32>,
|
||||
pub error_count: Option<u32>,
|
||||
pub total_response_time_ms: Option<u32>,
|
||||
pub last_used_at_unix_secs: Option<u64>,
|
||||
pub auto_fetch_models: bool,
|
||||
pub last_models_fetch_at_unix_secs: Option<u64>,
|
||||
pub last_models_fetch_error: Option<String>,
|
||||
pub locked_models: Option<serde_json::Value>,
|
||||
pub model_include_patterns: Option<serde_json::Value>,
|
||||
pub model_exclude_patterns: Option<serde_json::Value>,
|
||||
pub upstream_metadata: Option<serde_json::Value>,
|
||||
pub oauth_invalid_at_unix_secs: Option<u64>,
|
||||
pub oauth_invalid_reason: Option<String>,
|
||||
pub status_snapshot: Option<serde_json::Value>,
|
||||
pub created_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: Option<u64>,
|
||||
pub health_by_format: Option<serde_json::Value>,
|
||||
pub circuit_breaker_by_format: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
impl StoredProviderCatalogKey {
|
||||
@@ -108,12 +317,179 @@ impl StoredProviderCatalogKey {
|
||||
auth_type,
|
||||
capabilities,
|
||||
is_active,
|
||||
api_formats: None,
|
||||
encrypted_api_key: String::new(),
|
||||
encrypted_auth_config: None,
|
||||
note: None,
|
||||
internal_priority: 50,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
allowed_models: None,
|
||||
expires_at_unix_secs: None,
|
||||
cache_ttl_minutes: 5,
|
||||
max_probe_interval_minutes: 32,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
rpm_limit: None,
|
||||
learned_rpm_limit: None,
|
||||
concurrent_429_count: None,
|
||||
rpm_429_count: None,
|
||||
last_429_at_unix_secs: None,
|
||||
last_429_type: None,
|
||||
adjustment_history: None,
|
||||
utilization_samples: None,
|
||||
last_probe_increase_at_unix_secs: None,
|
||||
request_count: None,
|
||||
success_count: None,
|
||||
error_count: None,
|
||||
total_response_time_ms: None,
|
||||
last_used_at_unix_secs: None,
|
||||
auto_fetch_models: false,
|
||||
last_models_fetch_at_unix_secs: None,
|
||||
last_models_fetch_error: None,
|
||||
locked_models: None,
|
||||
model_include_patterns: None,
|
||||
model_exclude_patterns: None,
|
||||
upstream_metadata: None,
|
||||
oauth_invalid_at_unix_secs: None,
|
||||
oauth_invalid_reason: None,
|
||||
status_snapshot: None,
|
||||
created_at_unix_secs: None,
|
||||
updated_at_unix_secs: None,
|
||||
health_by_format: None,
|
||||
circuit_breaker_by_format: None,
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn with_transport_fields(
|
||||
mut self,
|
||||
api_formats: Option<serde_json::Value>,
|
||||
encrypted_api_key: String,
|
||||
encrypted_auth_config: Option<String>,
|
||||
rate_multipliers: Option<serde_json::Value>,
|
||||
global_priority_by_format: Option<serde_json::Value>,
|
||||
allowed_models: Option<serde_json::Value>,
|
||||
expires_at_unix_secs: Option<u64>,
|
||||
proxy: Option<serde_json::Value>,
|
||||
fingerprint: Option<serde_json::Value>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if encrypted_api_key.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider_api_keys.api_key is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
self.api_formats = api_formats;
|
||||
self.encrypted_api_key = encrypted_api_key;
|
||||
self.encrypted_auth_config = encrypted_auth_config;
|
||||
self.rate_multipliers = rate_multipliers;
|
||||
self.global_priority_by_format = global_priority_by_format;
|
||||
self.allowed_models = allowed_models;
|
||||
self.expires_at_unix_secs = expires_at_unix_secs;
|
||||
self.proxy = proxy;
|
||||
self.fingerprint = fingerprint;
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn with_rate_limit_fields(
|
||||
mut self,
|
||||
rpm_limit: Option<u32>,
|
||||
learned_rpm_limit: Option<u32>,
|
||||
concurrent_429_count: Option<u32>,
|
||||
rpm_429_count: Option<u32>,
|
||||
last_429_at_unix_secs: Option<u64>,
|
||||
adjustment_history: Option<serde_json::Value>,
|
||||
request_count: Option<u32>,
|
||||
success_count: Option<u32>,
|
||||
) -> Self {
|
||||
self.rpm_limit = rpm_limit;
|
||||
self.learned_rpm_limit = learned_rpm_limit;
|
||||
self.concurrent_429_count = concurrent_429_count;
|
||||
self.rpm_429_count = rpm_429_count;
|
||||
self.last_429_at_unix_secs = last_429_at_unix_secs;
|
||||
self.adjustment_history = adjustment_history;
|
||||
self.request_count = request_count;
|
||||
self.success_count = success_count;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_usage_fields(
|
||||
mut self,
|
||||
error_count: Option<u32>,
|
||||
total_response_time_ms: Option<u32>,
|
||||
) -> Self {
|
||||
self.error_count = error_count;
|
||||
self.total_response_time_ms = total_response_time_ms;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_health_fields(
|
||||
mut self,
|
||||
health_by_format: Option<serde_json::Value>,
|
||||
circuit_breaker_by_format: Option<serde_json::Value>,
|
||||
) -> Self {
|
||||
self.health_by_format = health_by_format;
|
||||
self.circuit_breaker_by_format = circuit_breaker_by_format;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct ProviderCatalogKeyListQuery {
|
||||
pub provider_id: String,
|
||||
pub search: Option<String>,
|
||||
pub is_active: Option<bool>,
|
||||
pub offset: usize,
|
||||
pub limit: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderCatalogKeyPage {
|
||||
pub items: Vec<StoredProviderCatalogKey>,
|
||||
pub total: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderCatalogKeyStats {
|
||||
pub provider_id: String,
|
||||
pub total_keys: u64,
|
||||
pub active_keys: u64,
|
||||
}
|
||||
|
||||
impl StoredProviderCatalogKeyStats {
|
||||
pub fn new(
|
||||
provider_id: String,
|
||||
total_keys: i64,
|
||||
active_keys: i64,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if provider_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider key stats provider_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if total_keys < 0 || active_keys < 0 {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider key stats count is negative".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
provider_id,
|
||||
total_keys: total_keys as u64,
|
||||
active_keys: active_keys as u64,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProviderCatalogReadRepository: Send + Sync {
|
||||
async fn list_providers(
|
||||
&self,
|
||||
active_only: bool,
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, crate::DataLayerError>;
|
||||
|
||||
async fn list_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
@@ -124,10 +500,98 @@ pub trait ProviderCatalogReadRepository: Send + Sync {
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, crate::DataLayerError>;
|
||||
|
||||
async fn list_endpoints_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, crate::DataLayerError>;
|
||||
|
||||
async fn list_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, crate::DataLayerError>;
|
||||
|
||||
async fn list_keys_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, crate::DataLayerError>;
|
||||
|
||||
async fn list_keys_page(
|
||||
&self,
|
||||
query: &ProviderCatalogKeyListQuery,
|
||||
) -> Result<StoredProviderCatalogKeyPage, crate::DataLayerError>;
|
||||
|
||||
async fn list_key_stats_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKeyStats>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProviderCatalogWriteRepository: Send + Sync {
|
||||
async fn create_provider(
|
||||
&self,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
shift_existing_priorities_from: Option<i32>,
|
||||
) -> Result<StoredProviderCatalogProvider, crate::DataLayerError>;
|
||||
|
||||
async fn update_provider(
|
||||
&self,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
) -> Result<StoredProviderCatalogProvider, crate::DataLayerError>;
|
||||
|
||||
async fn delete_provider(&self, provider_id: &str) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn cleanup_deleted_provider_refs(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
endpoint_ids: &[String],
|
||||
key_ids: &[String],
|
||||
) -> Result<(), crate::DataLayerError>;
|
||||
|
||||
async fn create_endpoint(
|
||||
&self,
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
) -> Result<StoredProviderCatalogEndpoint, crate::DataLayerError>;
|
||||
|
||||
async fn update_endpoint(
|
||||
&self,
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
) -> Result<StoredProviderCatalogEndpoint, crate::DataLayerError>;
|
||||
|
||||
async fn delete_endpoint(&self, endpoint_id: &str) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn create_key(
|
||||
&self,
|
||||
key: &StoredProviderCatalogKey,
|
||||
) -> Result<StoredProviderCatalogKey, crate::DataLayerError>;
|
||||
|
||||
async fn update_key(
|
||||
&self,
|
||||
key: &StoredProviderCatalogKey,
|
||||
) -> Result<StoredProviderCatalogKey, crate::DataLayerError>;
|
||||
|
||||
async fn delete_key(&self, key_id: &str) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn clear_key_oauth_invalid_marker(
|
||||
&self,
|
||||
key_id: &str,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn update_key_oauth_credentials(
|
||||
&self,
|
||||
key_id: &str,
|
||||
encrypted_api_key: &str,
|
||||
encrypted_auth_config: Option<&str>,
|
||||
expires_at_unix_secs: Option<u64>,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn update_key_health_state(
|
||||
&self,
|
||||
key_id: &str,
|
||||
is_active: bool,
|
||||
health_by_format: Option<&serde_json::Value>,
|
||||
circuit_breaker_by_format: Option<&serde_json::Value>,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -160,6 +624,22 @@ mod tests {
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_endpoint_base_url() {
|
||||
let endpoint = StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"openai:chat".to_string(),
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build");
|
||||
assert!(endpoint
|
||||
.with_transport_fields("".to_string(), None, None, None, None, None, None, None,)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_key_auth_type() {
|
||||
assert!(StoredProviderCatalogKey::new(
|
||||
@@ -172,4 +652,87 @@ mod tests {
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_encrypted_api_key() {
|
||||
let key = StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"default".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build");
|
||||
assert!(key
|
||||
.with_transport_fields(
|
||||
None,
|
||||
"".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stores_key_rate_limit_fields() {
|
||||
let key = StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"default".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_rate_limit_fields(
|
||||
Some(100),
|
||||
Some(80),
|
||||
Some(2),
|
||||
Some(3),
|
||||
Some(1_700_000_000),
|
||||
Some(serde_json::json!([{"new_limit": 80}])),
|
||||
Some(120),
|
||||
Some(110),
|
||||
);
|
||||
|
||||
assert_eq!(key.rpm_limit, Some(100));
|
||||
assert_eq!(key.learned_rpm_limit, Some(80));
|
||||
assert_eq!(key.concurrent_429_count, Some(2));
|
||||
assert_eq!(key.rpm_429_count, Some(3));
|
||||
assert_eq!(key.last_429_at_unix_secs, Some(1_700_000_000));
|
||||
assert_eq!(key.request_count, Some(120));
|
||||
assert_eq!(key.success_count, Some(110));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stores_key_health_fields() {
|
||||
let key = StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"default".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_health_fields(
|
||||
Some(serde_json::json!({"openai:chat": {"health_score": 0.4}})),
|
||||
Some(serde_json::json!({"openai:chat": {"open": true}})),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
key.health_by_format,
|
||||
Some(serde_json::json!({"openai:chat": {"health_score": 0.4}}))
|
||||
);
|
||||
assert_eq!(
|
||||
key.circuit_breaker_by_format,
|
||||
Some(serde_json::json!({"openai:chat": {"open": true}}))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
389
crates/aether-data/src/repository/proxy_nodes/memory.rs
Normal file
389
crates/aether-data/src/repository/proxy_nodes/memory.rs
Normal file
@@ -0,0 +1,389 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
normalize_proxy_metadata, ProxyNodeHeartbeatMutation, ProxyNodeReadRepository,
|
||||
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryProxyNodeRepository {
|
||||
nodes: RwLock<BTreeMap<String, StoredProxyNode>>,
|
||||
events: RwLock<Vec<StoredProxyNodeEvent>>,
|
||||
}
|
||||
|
||||
impl InMemoryProxyNodeRepository {
|
||||
pub fn seed<I>(nodes: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredProxyNode>,
|
||||
{
|
||||
Self {
|
||||
nodes: RwLock::new(
|
||||
nodes
|
||||
.into_iter()
|
||||
.map(|node| (node.id.clone(), node))
|
||||
.collect(),
|
||||
),
|
||||
events: RwLock::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn seed_with_events<I, J>(nodes: I, events: J) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredProxyNode>,
|
||||
J: IntoIterator<Item = StoredProxyNodeEvent>,
|
||||
{
|
||||
Self {
|
||||
nodes: RwLock::new(
|
||||
nodes
|
||||
.into_iter()
|
||||
.map(|node| (node.id.clone(), node))
|
||||
.collect(),
|
||||
),
|
||||
events: RwLock::new(events.into_iter().collect()),
|
||||
}
|
||||
}
|
||||
|
||||
fn now_unix_secs() -> Option<u64> {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
}
|
||||
|
||||
fn next_event_id(events: &[StoredProxyNodeEvent]) -> i64 {
|
||||
events.iter().map(|event| event.id).max().unwrap_or(0) + 1
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyNodeReadRepository for InMemoryProxyNodeRepository {
|
||||
async fn list_proxy_nodes(&self) -> Result<Vec<StoredProxyNode>, DataLayerError> {
|
||||
let nodes = self.nodes.read().expect("proxy node repository lock");
|
||||
let mut items = nodes.values().cloned().collect::<Vec<_>>();
|
||||
items.sort_by(|left, right| left.name.cmp(&right.name).then(left.id.cmp(&right.id)));
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
async fn find_proxy_node(
|
||||
&self,
|
||||
node_id: &str,
|
||||
) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
||||
let nodes = self.nodes.read().expect("proxy node repository lock");
|
||||
Ok(nodes.get(node_id).cloned())
|
||||
}
|
||||
|
||||
async fn list_proxy_node_events(
|
||||
&self,
|
||||
node_id: &str,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredProxyNodeEvent>, DataLayerError> {
|
||||
let events = self.events.read().expect("proxy node repository lock");
|
||||
let mut items = events
|
||||
.iter()
|
||||
.filter(|event| event.node_id == node_id)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
items.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.unwrap_or(0)
|
||||
.cmp(&left.created_at_unix_secs.unwrap_or(0))
|
||||
.then(right.id.cmp(&left.id))
|
||||
});
|
||||
items.truncate(limit);
|
||||
Ok(items)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
|
||||
async fn apply_heartbeat(
|
||||
&self,
|
||||
mutation: &ProxyNodeHeartbeatMutation,
|
||||
) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
||||
let mut nodes = self.nodes.write().expect("proxy node repository lock");
|
||||
let Some(node) = nodes.get_mut(&mutation.node_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !node.tunnel_mode {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"non-tunnel mode is no longer supported, please upgrade aether-proxy to use tunnel mode"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let now = Self::now_unix_secs();
|
||||
node.last_heartbeat_at_unix_secs = now;
|
||||
if node.status != "online" || !node.tunnel_connected {
|
||||
node.status = "online".to_string();
|
||||
node.tunnel_connected = true;
|
||||
node.tunnel_connected_at_unix_secs = now;
|
||||
node.updated_at_unix_secs = now;
|
||||
}
|
||||
|
||||
if let Some(value) = mutation.heartbeat_interval {
|
||||
node.heartbeat_interval = value;
|
||||
}
|
||||
if let Some(value) = mutation.active_connections {
|
||||
node.active_connections = value;
|
||||
}
|
||||
if let Some(value) = mutation.avg_latency_ms {
|
||||
node.avg_latency_ms = Some(value);
|
||||
}
|
||||
let normalized_proxy_metadata = normalize_proxy_metadata(
|
||||
mutation.proxy_metadata.as_ref(),
|
||||
mutation.proxy_version.as_deref(),
|
||||
);
|
||||
if let Some(value) = normalized_proxy_metadata {
|
||||
node.proxy_metadata = Some(value);
|
||||
}
|
||||
if let Some(value) = mutation.total_requests_delta.filter(|value| *value > 0) {
|
||||
node.total_requests += value;
|
||||
}
|
||||
if let Some(value) = mutation.failed_requests_delta.filter(|value| *value > 0) {
|
||||
node.failed_requests += value;
|
||||
}
|
||||
if let Some(value) = mutation.dns_failures_delta.filter(|value| *value > 0) {
|
||||
node.dns_failures += value;
|
||||
}
|
||||
if let Some(value) = mutation.stream_errors_delta.filter(|value| *value > 0) {
|
||||
node.stream_errors += value;
|
||||
}
|
||||
|
||||
Ok(Some(node.clone()))
|
||||
}
|
||||
|
||||
async fn update_tunnel_status(
|
||||
&self,
|
||||
mutation: &ProxyNodeTunnelStatusMutation,
|
||||
) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
||||
let mut nodes = self.nodes.write().expect("proxy node repository lock");
|
||||
let Some(node) = nodes.get_mut(&mutation.node_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let event_time = mutation
|
||||
.observed_at_unix_secs
|
||||
.or_else(Self::now_unix_secs)
|
||||
.unwrap_or(0);
|
||||
let event_type = if mutation.connected {
|
||||
"connected"
|
||||
} else {
|
||||
"disconnected"
|
||||
};
|
||||
let event_detail = mutation.detail.clone().unwrap_or_else(|| {
|
||||
format!(
|
||||
"[hub_node_status] conn_count={}",
|
||||
i32::max(mutation.conn_count, 0)
|
||||
)
|
||||
});
|
||||
let mut events = self.events.write().expect("proxy node repository lock");
|
||||
if let Some(last_transition) = node.tunnel_connected_at_unix_secs {
|
||||
if event_time < last_transition {
|
||||
let event_id = Self::next_event_id(&events);
|
||||
events.push(StoredProxyNodeEvent {
|
||||
id: event_id,
|
||||
node_id: mutation.node_id.clone(),
|
||||
event_type: event_type.to_string(),
|
||||
detail: Some(format!("[stale_ignored] {event_detail}")),
|
||||
created_at_unix_secs: Self::now_unix_secs(),
|
||||
});
|
||||
return Ok(Some(node.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
node.tunnel_connected = mutation.connected;
|
||||
node.tunnel_connected_at_unix_secs = Some(event_time);
|
||||
node.status = if mutation.connected {
|
||||
"online".to_string()
|
||||
} else {
|
||||
"offline".to_string()
|
||||
};
|
||||
node.updated_at_unix_secs = Some(event_time);
|
||||
let event_id = Self::next_event_id(&events);
|
||||
events.push(StoredProxyNodeEvent {
|
||||
id: event_id,
|
||||
node_id: mutation.node_id.clone(),
|
||||
event_type: event_type.to_string(),
|
||||
detail: Some(event_detail),
|
||||
created_at_unix_secs: Some(event_time),
|
||||
});
|
||||
Ok(Some(node.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryProxyNodeRepository;
|
||||
use crate::repository::proxy_nodes::{
|
||||
ProxyNodeHeartbeatMutation, ProxyNodeReadRepository, ProxyNodeTunnelStatusMutation,
|
||||
ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
fn sample_node() -> StoredProxyNode {
|
||||
StoredProxyNode::new(
|
||||
"node-1".to_string(),
|
||||
"proxy-1".to_string(),
|
||||
"127.0.0.1".to_string(),
|
||||
0,
|
||||
false,
|
||||
"offline".to_string(),
|
||||
30,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
true,
|
||||
false,
|
||||
2,
|
||||
)
|
||||
.expect("node should build")
|
||||
.with_runtime_fields(
|
||||
Some("test".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(json!({"allowed_ports": [443]})),
|
||||
Some(1_700_000_000),
|
||||
Some(1_700_000_001),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn applies_heartbeat_and_tunnel_status_mutations() {
|
||||
let repository = InMemoryProxyNodeRepository::seed(vec![sample_node()]);
|
||||
|
||||
let heartbeat = repository
|
||||
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
||||
node_id: "node-1".to_string(),
|
||||
heartbeat_interval: Some(45),
|
||||
active_connections: Some(5),
|
||||
total_requests_delta: Some(8),
|
||||
avg_latency_ms: Some(12.5),
|
||||
failed_requests_delta: Some(2),
|
||||
dns_failures_delta: Some(1),
|
||||
stream_errors_delta: Some(3),
|
||||
proxy_metadata: Some(json!({"arch": "arm64"})),
|
||||
proxy_version: Some("1.2.3".to_string()),
|
||||
})
|
||||
.await
|
||||
.expect("heartbeat should succeed")
|
||||
.expect("node should exist");
|
||||
|
||||
assert_eq!(heartbeat.status, "online");
|
||||
assert_eq!(heartbeat.heartbeat_interval, 45);
|
||||
assert_eq!(heartbeat.active_connections, 5);
|
||||
assert_eq!(heartbeat.total_requests, 8);
|
||||
assert_eq!(heartbeat.failed_requests, 2);
|
||||
assert_eq!(heartbeat.dns_failures, 1);
|
||||
assert_eq!(heartbeat.stream_errors, 3);
|
||||
assert_eq!(
|
||||
heartbeat
|
||||
.proxy_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("version"))
|
||||
.and_then(|value| value.as_str()),
|
||||
Some("1.2.3")
|
||||
);
|
||||
|
||||
let stale = repository
|
||||
.update_tunnel_status(&ProxyNodeTunnelStatusMutation {
|
||||
node_id: "node-1".to_string(),
|
||||
connected: false,
|
||||
conn_count: 0,
|
||||
detail: None,
|
||||
observed_at_unix_secs: Some(1),
|
||||
})
|
||||
.await
|
||||
.expect("status update should succeed")
|
||||
.expect("node should exist");
|
||||
assert_eq!(stale.status, "online");
|
||||
|
||||
let stale_events = repository
|
||||
.list_proxy_node_events("node-1", 10)
|
||||
.await
|
||||
.expect("list events should succeed");
|
||||
assert_eq!(stale_events.len(), 1);
|
||||
assert_eq!(stale_events[0].event_type, "disconnected");
|
||||
assert_eq!(
|
||||
stale_events[0].detail.as_deref(),
|
||||
Some("[stale_ignored] [hub_node_status] conn_count=0")
|
||||
);
|
||||
|
||||
let updated = repository
|
||||
.update_tunnel_status(&ProxyNodeTunnelStatusMutation {
|
||||
node_id: "node-1".to_string(),
|
||||
connected: false,
|
||||
conn_count: 0,
|
||||
detail: None,
|
||||
observed_at_unix_secs: Some(1_800_000_000),
|
||||
})
|
||||
.await
|
||||
.expect("status update should succeed")
|
||||
.expect("node should exist");
|
||||
assert_eq!(updated.status, "offline");
|
||||
assert!(!updated.tunnel_connected);
|
||||
|
||||
let events = repository
|
||||
.list_proxy_node_events("node-1", 10)
|
||||
.await
|
||||
.expect("list events should succeed");
|
||||
assert_eq!(events.len(), 2);
|
||||
assert_eq!(events[0].event_type, "disconnected");
|
||||
assert_eq!(events[0].created_at_unix_secs, Some(1_800_000_000));
|
||||
assert_eq!(
|
||||
events[0].detail.as_deref(),
|
||||
Some("[hub_node_status] conn_count=0")
|
||||
);
|
||||
|
||||
let found = repository
|
||||
.find_proxy_node("node-1")
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.expect("node should exist");
|
||||
assert_eq!(found.status, "offline");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lists_seeded_proxy_node_events_in_descending_order() {
|
||||
let repository = InMemoryProxyNodeRepository::seed_with_events(
|
||||
vec![sample_node()],
|
||||
vec![
|
||||
StoredProxyNodeEvent {
|
||||
id: 1,
|
||||
node_id: "node-1".to_string(),
|
||||
event_type: "connected".to_string(),
|
||||
detail: Some("older".to_string()),
|
||||
created_at_unix_secs: Some(1_710_000_000),
|
||||
},
|
||||
StoredProxyNodeEvent {
|
||||
id: 2,
|
||||
node_id: "node-1".to_string(),
|
||||
event_type: "disconnected".to_string(),
|
||||
detail: Some("newer".to_string()),
|
||||
created_at_unix_secs: Some(1_710_000_100),
|
||||
},
|
||||
],
|
||||
);
|
||||
|
||||
let events = repository
|
||||
.list_proxy_node_events("node-1", 1)
|
||||
.await
|
||||
.expect("list events should succeed");
|
||||
|
||||
assert_eq!(events.len(), 1);
|
||||
assert_eq!(events[0].id, 2);
|
||||
assert_eq!(events[0].detail.as_deref(), Some("newer"));
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/proxy_nodes/mod.rs
Normal file
10
crates/aether-data/src/repository/proxy_nodes/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryProxyNodeRepository;
|
||||
pub use sql::SqlxProxyNodeRepository;
|
||||
pub use types::{
|
||||
ProxyNodeHeartbeatMutation, ProxyNodeReadRepository, ProxyNodeTunnelStatusMutation,
|
||||
ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
};
|
||||
363
crates/aether-data/src/repository/proxy_nodes/sql.rs
Normal file
363
crates/aether-data/src/repository/proxy_nodes/sql.rs
Normal file
@@ -0,0 +1,363 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
normalize_proxy_metadata, ProxyNodeHeartbeatMutation, ProxyNodeReadRepository,
|
||||
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_PROXY_NODE_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
name,
|
||||
ip,
|
||||
port,
|
||||
region,
|
||||
is_manual,
|
||||
proxy_url,
|
||||
proxy_username,
|
||||
proxy_password,
|
||||
CAST(status AS TEXT) AS status,
|
||||
registered_by,
|
||||
EXTRACT(EPOCH FROM last_heartbeat_at)::bigint AS last_heartbeat_at_unix_secs,
|
||||
heartbeat_interval,
|
||||
active_connections,
|
||||
total_requests,
|
||||
CAST(avg_latency_ms AS DOUBLE PRECISION) AS avg_latency_ms,
|
||||
failed_requests,
|
||||
dns_failures,
|
||||
stream_errors,
|
||||
proxy_metadata,
|
||||
hardware_info,
|
||||
estimated_max_concurrency,
|
||||
tunnel_mode,
|
||||
tunnel_connected,
|
||||
EXTRACT(EPOCH FROM tunnel_connected_at)::bigint AS tunnel_connected_at_unix_secs,
|
||||
remote_config,
|
||||
config_version,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
FROM proxy_nodes
|
||||
WHERE id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const LIST_PROXY_NODES_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
name,
|
||||
ip,
|
||||
port,
|
||||
region,
|
||||
is_manual,
|
||||
proxy_url,
|
||||
proxy_username,
|
||||
proxy_password,
|
||||
CAST(status AS TEXT) AS status,
|
||||
registered_by,
|
||||
EXTRACT(EPOCH FROM last_heartbeat_at)::bigint AS last_heartbeat_at_unix_secs,
|
||||
heartbeat_interval,
|
||||
active_connections,
|
||||
total_requests,
|
||||
CAST(avg_latency_ms AS DOUBLE PRECISION) AS avg_latency_ms,
|
||||
failed_requests,
|
||||
dns_failures,
|
||||
stream_errors,
|
||||
proxy_metadata,
|
||||
hardware_info,
|
||||
estimated_max_concurrency,
|
||||
tunnel_mode,
|
||||
tunnel_connected,
|
||||
EXTRACT(EPOCH FROM tunnel_connected_at)::bigint AS tunnel_connected_at_unix_secs,
|
||||
remote_config,
|
||||
config_version,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
FROM proxy_nodes
|
||||
ORDER BY name ASC, id ASC
|
||||
"#;
|
||||
|
||||
const LIST_PROXY_NODE_EVENTS_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
node_id,
|
||||
CAST(event_type AS TEXT) AS event_type,
|
||||
detail,
|
||||
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs
|
||||
FROM proxy_node_events
|
||||
WHERE node_id = $1
|
||||
ORDER BY created_at DESC, id DESC
|
||||
LIMIT $2
|
||||
"#;
|
||||
|
||||
const APPLY_HEARTBEAT_SQL: &str = r#"
|
||||
UPDATE proxy_nodes
|
||||
SET
|
||||
last_heartbeat_at = NOW(),
|
||||
status = CASE
|
||||
WHEN status <> 'online'::proxynodestatus OR tunnel_connected = FALSE
|
||||
THEN 'online'::proxynodestatus
|
||||
ELSE status
|
||||
END,
|
||||
tunnel_connected = CASE
|
||||
WHEN status <> 'online'::proxynodestatus OR tunnel_connected = FALSE
|
||||
THEN TRUE
|
||||
ELSE tunnel_connected
|
||||
END,
|
||||
tunnel_connected_at = CASE
|
||||
WHEN status <> 'online'::proxynodestatus OR tunnel_connected = FALSE
|
||||
THEN NOW()
|
||||
ELSE tunnel_connected_at
|
||||
END,
|
||||
updated_at = CASE
|
||||
WHEN status <> 'online'::proxynodestatus OR tunnel_connected = FALSE
|
||||
THEN NOW()
|
||||
ELSE updated_at
|
||||
END,
|
||||
heartbeat_interval = COALESCE($2, heartbeat_interval),
|
||||
active_connections = COALESCE($3, active_connections),
|
||||
avg_latency_ms = COALESCE($4, avg_latency_ms),
|
||||
proxy_metadata = COALESCE($5, proxy_metadata),
|
||||
total_requests = total_requests + GREATEST(COALESCE($6, 0), 0),
|
||||
failed_requests = failed_requests + GREATEST(COALESCE($7, 0), 0),
|
||||
dns_failures = dns_failures + GREATEST(COALESCE($8, 0), 0),
|
||||
stream_errors = stream_errors + GREATEST(COALESCE($9, 0), 0)
|
||||
WHERE id = $1
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxProxyNodeRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxProxyNodeRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
fn optional_unix_secs(value: Option<i64>) -> Option<u64> {
|
||||
value.and_then(|value| u64::try_from(value).ok())
|
||||
}
|
||||
|
||||
fn row_to_stored(row: &PgRow) -> Result<StoredProxyNode, DataLayerError> {
|
||||
Ok(StoredProxyNode::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("name")?,
|
||||
row.try_get("ip")?,
|
||||
row.try_get("port")?,
|
||||
row.try_get("is_manual")?,
|
||||
row.try_get("status")?,
|
||||
row.try_get("heartbeat_interval")?,
|
||||
row.try_get("active_connections")?,
|
||||
row.try_get("total_requests")?,
|
||||
row.try_get("failed_requests")?,
|
||||
row.try_get("dns_failures")?,
|
||||
row.try_get("stream_errors")?,
|
||||
row.try_get("tunnel_mode")?,
|
||||
row.try_get("tunnel_connected")?,
|
||||
row.try_get("config_version")?,
|
||||
)?
|
||||
.with_manual_proxy_fields(
|
||||
row.try_get("proxy_url")?,
|
||||
row.try_get("proxy_username")?,
|
||||
row.try_get("proxy_password")?,
|
||||
)
|
||||
.with_runtime_fields(
|
||||
row.try_get("region")?,
|
||||
row.try_get("registered_by")?,
|
||||
Self::optional_unix_secs(row.try_get("last_heartbeat_at_unix_secs")?),
|
||||
row.try_get("avg_latency_ms")?,
|
||||
row.try_get("proxy_metadata")?,
|
||||
row.try_get("hardware_info")?,
|
||||
row.try_get("estimated_max_concurrency")?,
|
||||
Self::optional_unix_secs(row.try_get("tunnel_connected_at_unix_secs")?),
|
||||
row.try_get("remote_config")?,
|
||||
Self::optional_unix_secs(row.try_get("created_at_unix_secs")?),
|
||||
Self::optional_unix_secs(row.try_get("updated_at_unix_secs")?),
|
||||
))
|
||||
}
|
||||
|
||||
fn row_to_event(row: &PgRow) -> Result<StoredProxyNodeEvent, DataLayerError> {
|
||||
Ok(StoredProxyNodeEvent {
|
||||
id: row.try_get("id")?,
|
||||
node_id: row.try_get("node_id")?,
|
||||
event_type: row.try_get("event_type")?,
|
||||
detail: row.try_get("detail")?,
|
||||
created_at_unix_secs: Self::optional_unix_secs(row.try_get("created_at_unix_secs")?),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyNodeReadRepository for SqlxProxyNodeRepository {
|
||||
async fn list_proxy_nodes(&self) -> Result<Vec<StoredProxyNode>, DataLayerError> {
|
||||
let rows = sqlx::query(LIST_PROXY_NODES_SQL)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(Self::row_to_stored).collect()
|
||||
}
|
||||
|
||||
async fn find_proxy_node(
|
||||
&self,
|
||||
node_id: &str,
|
||||
) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_PROXY_NODE_SQL)
|
||||
.bind(node_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.map(|row| Self::row_to_stored(&row)).transpose()
|
||||
}
|
||||
|
||||
async fn list_proxy_node_events(
|
||||
&self,
|
||||
node_id: &str,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredProxyNodeEvent>, DataLayerError> {
|
||||
let rows = sqlx::query(LIST_PROXY_NODE_EVENTS_SQL)
|
||||
.bind(node_id)
|
||||
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(Self::row_to_event).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyNodeWriteRepository for SqlxProxyNodeRepository {
|
||||
async fn apply_heartbeat(
|
||||
&self,
|
||||
mutation: &ProxyNodeHeartbeatMutation,
|
||||
) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
||||
let existing = self.find_proxy_node(&mutation.node_id).await?;
|
||||
let Some(existing) = existing else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !existing.tunnel_mode {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"non-tunnel mode is no longer supported, please upgrade aether-proxy to use tunnel mode"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let normalized_proxy_metadata = normalize_proxy_metadata(
|
||||
mutation.proxy_metadata.as_ref(),
|
||||
mutation.proxy_version.as_deref(),
|
||||
);
|
||||
|
||||
sqlx::query(APPLY_HEARTBEAT_SQL)
|
||||
.bind(&mutation.node_id)
|
||||
.bind(mutation.heartbeat_interval)
|
||||
.bind(mutation.active_connections)
|
||||
.bind(mutation.avg_latency_ms)
|
||||
.bind(normalized_proxy_metadata)
|
||||
.bind(mutation.total_requests_delta)
|
||||
.bind(mutation.failed_requests_delta)
|
||||
.bind(mutation.dns_failures_delta)
|
||||
.bind(mutation.stream_errors_delta)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
self.find_proxy_node(&mutation.node_id).await
|
||||
}
|
||||
|
||||
async fn update_tunnel_status(
|
||||
&self,
|
||||
mutation: &ProxyNodeTunnelStatusMutation,
|
||||
) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
||||
let existing = self.find_proxy_node(&mutation.node_id).await?;
|
||||
let Some(existing) = existing else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let observed_at_unix_secs = mutation.observed_at_unix_secs;
|
||||
let event_type = if mutation.connected {
|
||||
"connected"
|
||||
} else {
|
||||
"disconnected"
|
||||
};
|
||||
let event_detail = mutation.detail.clone().unwrap_or_else(|| {
|
||||
format!(
|
||||
"[hub_node_status] conn_count={}",
|
||||
i32::max(mutation.conn_count, 0)
|
||||
)
|
||||
});
|
||||
|
||||
let mut tx = self.pool.begin().await?;
|
||||
|
||||
if existing
|
||||
.tunnel_connected_at_unix_secs
|
||||
.zip(observed_at_unix_secs)
|
||||
.is_some_and(|(last_transition, observed_at)| observed_at < last_transition)
|
||||
{
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO proxy_node_events (node_id, event_type, detail, created_at)
|
||||
VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
NOW()
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(&mutation.node_id)
|
||||
.bind(event_type)
|
||||
.bind(format!("[stale_ignored] {event_detail}"))
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
tx.commit().await?;
|
||||
return self.find_proxy_node(&mutation.node_id).await;
|
||||
}
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE proxy_nodes
|
||||
SET
|
||||
tunnel_connected = $2,
|
||||
tunnel_connected_at = CASE
|
||||
WHEN $3::double precision IS NULL THEN NOW()
|
||||
ELSE TO_TIMESTAMP($3::double precision)
|
||||
END,
|
||||
status = CASE
|
||||
WHEN $2 THEN 'online'::proxynodestatus
|
||||
ELSE 'offline'::proxynodestatus
|
||||
END,
|
||||
updated_at = CASE
|
||||
WHEN $3::double precision IS NULL THEN NOW()
|
||||
ELSE TO_TIMESTAMP($3::double precision)
|
||||
END
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(&mutation.node_id)
|
||||
.bind(mutation.connected)
|
||||
.bind(observed_at_unix_secs.map(|value| value as f64))
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO proxy_node_events (node_id, event_type, detail, created_at)
|
||||
VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
CASE
|
||||
WHEN $4::double precision IS NULL THEN NOW()
|
||||
ELSE TO_TIMESTAMP($4::double precision)
|
||||
END
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(&mutation.node_id)
|
||||
.bind(event_type)
|
||||
.bind(event_detail)
|
||||
.bind(observed_at_unix_secs.map(|value| value as f64))
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
tx.commit().await?;
|
||||
self.find_proxy_node(&mutation.node_id).await
|
||||
}
|
||||
}
|
||||
244
crates/aether-data/src/repository/proxy_nodes/types.rs
Normal file
244
crates/aether-data/src/repository/proxy_nodes/types.rs
Normal file
@@ -0,0 +1,244 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProxyNode {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub ip: String,
|
||||
pub port: i32,
|
||||
pub region: Option<String>,
|
||||
pub is_manual: bool,
|
||||
pub proxy_url: Option<String>,
|
||||
pub proxy_username: Option<String>,
|
||||
pub proxy_password: Option<String>,
|
||||
pub status: String,
|
||||
pub registered_by: Option<String>,
|
||||
pub last_heartbeat_at_unix_secs: Option<u64>,
|
||||
pub heartbeat_interval: i32,
|
||||
pub active_connections: i32,
|
||||
pub total_requests: i64,
|
||||
pub avg_latency_ms: Option<f64>,
|
||||
pub failed_requests: i64,
|
||||
pub dns_failures: i64,
|
||||
pub stream_errors: i64,
|
||||
pub proxy_metadata: Option<serde_json::Value>,
|
||||
pub hardware_info: Option<serde_json::Value>,
|
||||
pub estimated_max_concurrency: Option<i32>,
|
||||
pub tunnel_mode: bool,
|
||||
pub tunnel_connected: bool,
|
||||
pub tunnel_connected_at_unix_secs: Option<u64>,
|
||||
pub remote_config: Option<serde_json::Value>,
|
||||
pub config_version: i32,
|
||||
pub created_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl StoredProxyNode {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
name: String,
|
||||
ip: String,
|
||||
port: i32,
|
||||
is_manual: bool,
|
||||
status: String,
|
||||
heartbeat_interval: i32,
|
||||
active_connections: i32,
|
||||
total_requests: i64,
|
||||
failed_requests: i64,
|
||||
dns_failures: i64,
|
||||
stream_errors: i64,
|
||||
tunnel_mode: bool,
|
||||
tunnel_connected: bool,
|
||||
config_version: i32,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"proxy_nodes.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"proxy_nodes.name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if ip.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"proxy_nodes.ip is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if status.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"proxy_nodes.status is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
name,
|
||||
ip,
|
||||
port,
|
||||
region: None,
|
||||
is_manual,
|
||||
proxy_url: None,
|
||||
proxy_username: None,
|
||||
proxy_password: None,
|
||||
status,
|
||||
registered_by: None,
|
||||
last_heartbeat_at_unix_secs: None,
|
||||
heartbeat_interval,
|
||||
active_connections,
|
||||
total_requests,
|
||||
avg_latency_ms: None,
|
||||
failed_requests,
|
||||
dns_failures,
|
||||
stream_errors,
|
||||
proxy_metadata: None,
|
||||
hardware_info: None,
|
||||
estimated_max_concurrency: None,
|
||||
tunnel_mode,
|
||||
tunnel_connected,
|
||||
tunnel_connected_at_unix_secs: None,
|
||||
remote_config: None,
|
||||
config_version,
|
||||
created_at_unix_secs: None,
|
||||
updated_at_unix_secs: None,
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn with_runtime_fields(
|
||||
mut self,
|
||||
region: Option<String>,
|
||||
registered_by: Option<String>,
|
||||
last_heartbeat_at_unix_secs: Option<u64>,
|
||||
avg_latency_ms: Option<f64>,
|
||||
proxy_metadata: Option<serde_json::Value>,
|
||||
hardware_info: Option<serde_json::Value>,
|
||||
estimated_max_concurrency: Option<i32>,
|
||||
tunnel_connected_at_unix_secs: Option<u64>,
|
||||
remote_config: Option<serde_json::Value>,
|
||||
created_at_unix_secs: Option<u64>,
|
||||
updated_at_unix_secs: Option<u64>,
|
||||
) -> Self {
|
||||
self.region = region;
|
||||
self.registered_by = registered_by;
|
||||
self.last_heartbeat_at_unix_secs = last_heartbeat_at_unix_secs;
|
||||
self.avg_latency_ms = avg_latency_ms;
|
||||
self.proxy_metadata = proxy_metadata;
|
||||
self.hardware_info = hardware_info;
|
||||
self.estimated_max_concurrency = estimated_max_concurrency;
|
||||
self.tunnel_connected_at_unix_secs = tunnel_connected_at_unix_secs;
|
||||
self.remote_config = remote_config;
|
||||
self.created_at_unix_secs = created_at_unix_secs;
|
||||
self.updated_at_unix_secs = updated_at_unix_secs;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_manual_proxy_fields(
|
||||
mut self,
|
||||
proxy_url: Option<String>,
|
||||
proxy_username: Option<String>,
|
||||
proxy_password: Option<String>,
|
||||
) -> Self {
|
||||
self.proxy_url = proxy_url;
|
||||
self.proxy_username = proxy_username;
|
||||
self.proxy_password = proxy_password;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ProxyNodeHeartbeatMutation {
|
||||
pub node_id: String,
|
||||
pub heartbeat_interval: Option<i32>,
|
||||
pub active_connections: Option<i32>,
|
||||
pub total_requests_delta: Option<i64>,
|
||||
pub avg_latency_ms: Option<f64>,
|
||||
pub failed_requests_delta: Option<i64>,
|
||||
pub dns_failures_delta: Option<i64>,
|
||||
pub stream_errors_delta: Option<i64>,
|
||||
pub proxy_metadata: Option<serde_json::Value>,
|
||||
pub proxy_version: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ProxyNodeTunnelStatusMutation {
|
||||
pub node_id: String,
|
||||
pub connected: bool,
|
||||
pub conn_count: i32,
|
||||
pub detail: Option<String>,
|
||||
pub observed_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProxyNodeEvent {
|
||||
pub id: i64,
|
||||
pub node_id: String,
|
||||
pub event_type: String,
|
||||
pub detail: Option<String>,
|
||||
pub created_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
pub fn normalize_proxy_metadata(
|
||||
proxy_metadata: Option<&serde_json::Value>,
|
||||
proxy_version: Option<&str>,
|
||||
) -> Option<serde_json::Value> {
|
||||
let mut normalized = match proxy_metadata {
|
||||
Some(serde_json::Value::Object(map)) => map.clone(),
|
||||
Some(_) | None => serde_json::Map::new(),
|
||||
};
|
||||
|
||||
let raw_version = normalized
|
||||
.remove("version")
|
||||
.and_then(|value| value.as_str().map(str::to_string));
|
||||
let version = proxy_version
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|value| value.chars().take(20).collect::<String>())
|
||||
.or_else(|| {
|
||||
raw_version
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|value| value.chars().take(20).collect::<String>())
|
||||
});
|
||||
if let Some(version) = version {
|
||||
normalized.insert("version".to_string(), serde_json::Value::String(version));
|
||||
}
|
||||
|
||||
if normalized.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(serde_json::Value::Object(normalized))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProxyNodeReadRepository: Send + Sync {
|
||||
async fn list_proxy_nodes(&self) -> Result<Vec<StoredProxyNode>, crate::DataLayerError>;
|
||||
|
||||
async fn find_proxy_node(
|
||||
&self,
|
||||
node_id: &str,
|
||||
) -> Result<Option<StoredProxyNode>, crate::DataLayerError>;
|
||||
|
||||
async fn list_proxy_node_events(
|
||||
&self,
|
||||
node_id: &str,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredProxyNodeEvent>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProxyNodeWriteRepository: Send + Sync {
|
||||
async fn apply_heartbeat(
|
||||
&self,
|
||||
mutation: &ProxyNodeHeartbeatMutation,
|
||||
) -> Result<Option<StoredProxyNode>, crate::DataLayerError>;
|
||||
|
||||
async fn update_tunnel_status(
|
||||
&self,
|
||||
mutation: &ProxyNodeTunnelStatusMutation,
|
||||
) -> Result<Option<StoredProxyNode>, crate::DataLayerError>;
|
||||
}
|
||||
109
crates/aether-data/src/repository/quota/memory.rs
Normal file
109
crates/aether-data/src/repository/quota/memory.rs
Normal file
@@ -0,0 +1,109 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
use aether_wallet::{ProviderBillingType, ProviderQuotaSnapshot};
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryProviderQuotaRepository {
|
||||
by_provider_id: RwLock<BTreeMap<String, StoredProviderQuotaSnapshot>>,
|
||||
}
|
||||
|
||||
impl InMemoryProviderQuotaRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredProviderQuotaSnapshot>,
|
||||
{
|
||||
let mut by_provider_id = BTreeMap::new();
|
||||
for item in items {
|
||||
by_provider_id.insert(item.provider_id.clone(), item);
|
||||
}
|
||||
Self {
|
||||
by_provider_id: RwLock::new(by_provider_id),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderQuotaReadRepository for InMemoryProviderQuotaRepository {
|
||||
async fn find_by_provider_id(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
) -> Result<Option<StoredProviderQuotaSnapshot>, DataLayerError> {
|
||||
Ok(self
|
||||
.by_provider_id
|
||||
.read()
|
||||
.expect("quota repository lock")
|
||||
.get(provider_id)
|
||||
.cloned())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderQuotaWriteRepository for InMemoryProviderQuotaRepository {
|
||||
async fn reset_due(&self, now_unix_secs: u64) -> Result<usize, DataLayerError> {
|
||||
let mut count = 0usize;
|
||||
let mut quotas = self.by_provider_id.write().expect("quota repository lock");
|
||||
for quota in quotas.values_mut() {
|
||||
let snapshot = ProviderQuotaSnapshot {
|
||||
provider_id: quota.provider_id.clone(),
|
||||
billing_type: ProviderBillingType::parse("a.billing_type),
|
||||
monthly_quota_usd: quota.monthly_quota_usd,
|
||||
monthly_used_usd: quota.monthly_used_usd,
|
||||
quota_reset_day: quota.quota_reset_day,
|
||||
quota_last_reset_at_unix_secs: quota.quota_last_reset_at_unix_secs,
|
||||
quota_expires_at_unix_secs: quota.quota_expires_at_unix_secs,
|
||||
is_active: quota.is_active,
|
||||
};
|
||||
if snapshot.should_reset(now_unix_secs) {
|
||||
quota.monthly_used_usd = 0.0;
|
||||
quota.quota_last_reset_at_unix_secs = Some(now_unix_secs);
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
Ok(count)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryProviderQuotaRepository;
|
||||
use crate::repository::quota::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot,
|
||||
};
|
||||
|
||||
fn sample_quota() -> StoredProviderQuotaSnapshot {
|
||||
StoredProviderQuotaSnapshot::new(
|
||||
"provider-1".to_string(),
|
||||
"monthly_quota".to_string(),
|
||||
Some(20.0),
|
||||
5.0,
|
||||
Some(7),
|
||||
Some(1_000),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("quota should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resets_due_monthly_quota() {
|
||||
let repository = InMemoryProviderQuotaRepository::seed(vec![sample_quota()]);
|
||||
let reset = repository
|
||||
.reset_due(1_000 + 7 * 24 * 60 * 60)
|
||||
.await
|
||||
.expect("reset should succeed");
|
||||
assert_eq!(reset, 1);
|
||||
let stored = repository
|
||||
.find_by_provider_id("provider-1")
|
||||
.await
|
||||
.expect("lookup should succeed")
|
||||
.expect("quota should exist");
|
||||
assert_eq!(stored.monthly_used_usd, 0.0);
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/quota/mod.rs
Normal file
10
crates/aether-data/src/repository/quota/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryProviderQuotaRepository;
|
||||
pub use sql::SqlxProviderQuotaRepository;
|
||||
pub use types::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaRepository, ProviderQuotaWriteRepository,
|
||||
StoredProviderQuotaSnapshot,
|
||||
};
|
||||
112
crates/aether-data/src/repository/quota/sql.rs
Normal file
112
crates/aether-data/src/repository/quota/sql.rs
Normal file
@@ -0,0 +1,112 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_BY_PROVIDER_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id AS provider_id,
|
||||
CAST(billing_type AS TEXT) AS billing_type,
|
||||
CAST(monthly_quota_usd AS DOUBLE PRECISION) AS monthly_quota_usd,
|
||||
CAST(COALESCE(monthly_used_usd, 0) AS DOUBLE PRECISION) AS monthly_used_usd,
|
||||
quota_reset_day,
|
||||
CAST(EXTRACT(EPOCH FROM quota_last_reset_at) AS BIGINT) AS quota_last_reset_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM quota_expires_at) AS BIGINT) AS quota_expires_at_unix_secs,
|
||||
is_active
|
||||
FROM providers
|
||||
WHERE id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const RESET_DUE_SQL: &str = r#"
|
||||
UPDATE providers
|
||||
SET
|
||||
monthly_used_usd = 0,
|
||||
quota_last_reset_at = TO_TIMESTAMP($1::double precision),
|
||||
updated_at = NOW()
|
||||
WHERE
|
||||
billing_type = 'monthly_quota'
|
||||
AND is_active = TRUE
|
||||
AND (
|
||||
quota_last_reset_at IS NULL
|
||||
OR (EXTRACT(EPOCH FROM TO_TIMESTAMP($1::double precision)) - EXTRACT(EPOCH FROM quota_last_reset_at)) >= (quota_reset_day * 86400)
|
||||
)
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxProviderQuotaRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxProviderQuotaRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderQuotaReadRepository for SqlxProviderQuotaRepository {
|
||||
async fn find_by_provider_id(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
) -> Result<Option<StoredProviderQuotaSnapshot>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_PROVIDER_ID_SQL)
|
||||
.bind(provider_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderQuotaWriteRepository for SqlxProviderQuotaRepository {
|
||||
async fn reset_due(&self, now_unix_secs: u64) -> Result<usize, DataLayerError> {
|
||||
let result = sqlx::query(RESET_DUE_SQL)
|
||||
.bind(i64::try_from(now_unix_secs).map_err(|_| {
|
||||
DataLayerError::InvalidInput("provider quota reset timestamp overflow".to_string())
|
||||
})?)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected() as usize)
|
||||
}
|
||||
}
|
||||
|
||||
fn map_row(row: &sqlx::postgres::PgRow) -> Result<StoredProviderQuotaSnapshot, DataLayerError> {
|
||||
StoredProviderQuotaSnapshot::new(
|
||||
row.try_get("provider_id")?,
|
||||
row.try_get("billing_type")?,
|
||||
row.try_get("monthly_quota_usd")?,
|
||||
row.try_get("monthly_used_usd")?,
|
||||
row.try_get("quota_reset_day")?,
|
||||
row.try_get("quota_last_reset_at_unix_secs")?,
|
||||
row.try_get("quota_expires_at_unix_secs")?,
|
||||
row.try_get("is_active")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxProviderQuotaRepository;
|
||||
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 = SqlxProviderQuotaRepository::new(pool);
|
||||
}
|
||||
}
|
||||
71
crates/aether-data/src/repository/quota/types.rs
Normal file
71
crates/aether-data/src/repository/quota/types.rs
Normal file
@@ -0,0 +1,71 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderQuotaSnapshot {
|
||||
pub provider_id: String,
|
||||
pub billing_type: String,
|
||||
pub monthly_quota_usd: Option<f64>,
|
||||
pub monthly_used_usd: f64,
|
||||
pub quota_reset_day: Option<u64>,
|
||||
pub quota_last_reset_at_unix_secs: Option<u64>,
|
||||
pub quota_expires_at_unix_secs: Option<u64>,
|
||||
pub is_active: bool,
|
||||
}
|
||||
|
||||
impl StoredProviderQuotaSnapshot {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
provider_id: String,
|
||||
billing_type: String,
|
||||
monthly_quota_usd: Option<f64>,
|
||||
monthly_used_usd: f64,
|
||||
quota_reset_day: Option<i32>,
|
||||
quota_last_reset_at_unix_secs: Option<i64>,
|
||||
quota_expires_at_unix_secs: Option<i64>,
|
||||
is_active: bool,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if provider_id.trim().is_empty() || billing_type.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider quota identity is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if !monthly_used_usd.is_finite() || monthly_quota_usd.is_some_and(|v| !v.is_finite()) {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider quota value is not finite".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(Self {
|
||||
provider_id,
|
||||
billing_type,
|
||||
monthly_quota_usd,
|
||||
monthly_used_usd,
|
||||
quota_reset_day: quota_reset_day.map(|value| value as u64),
|
||||
quota_last_reset_at_unix_secs: quota_last_reset_at_unix_secs.map(|value| value as u64),
|
||||
quota_expires_at_unix_secs: quota_expires_at_unix_secs.map(|value| value as u64),
|
||||
is_active,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProviderQuotaReadRepository: Send + Sync {
|
||||
async fn find_by_provider_id(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
) -> Result<Option<StoredProviderQuotaSnapshot>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProviderQuotaWriteRepository: Send + Sync {
|
||||
async fn reset_due(&self, now_unix_secs: u64) -> Result<usize, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait ProviderQuotaRepository:
|
||||
ProviderQuotaReadRepository + ProviderQuotaWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> ProviderQuotaRepository for T where
|
||||
T: ProviderQuotaReadRepository + ProviderQuotaWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
@@ -3,12 +3,16 @@ use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
use super::types::{
|
||||
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
|
||||
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageWriteRepository,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryUsageReadRepository {
|
||||
by_request_id: RwLock<BTreeMap<String, StoredRequestUsageAudit>>,
|
||||
provider_usage_windows: RwLock<Vec<StoredProviderUsageWindow>>,
|
||||
}
|
||||
|
||||
impl InMemoryUsageReadRepository {
|
||||
@@ -22,12 +26,36 @@ impl InMemoryUsageReadRepository {
|
||||
}
|
||||
Self {
|
||||
by_request_id: RwLock::new(by_request_id),
|
||||
provider_usage_windows: RwLock::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_provider_usage_windows<I>(self, items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredProviderUsageWindow>,
|
||||
{
|
||||
Self {
|
||||
by_request_id: self.by_request_id,
|
||||
provider_usage_windows: RwLock::new(items.into_iter().collect()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
async fn find_by_id(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Ok(self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.values()
|
||||
.find(|item| item.id == id)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
@@ -39,12 +67,250 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
.get(request_id)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn list_usage_audits(
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let mut items: Vec<_> = self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.values()
|
||||
.filter(|item| {
|
||||
if let Some(created_from_unix_secs) = query.created_from_unix_secs {
|
||||
if item.created_at_unix_secs < created_from_unix_secs {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(created_until_unix_secs) = query.created_until_unix_secs {
|
||||
if item.created_at_unix_secs >= created_until_unix_secs {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(user_id) = query.user_id.as_deref() {
|
||||
if item.user_id.as_deref() != Some(user_id) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(provider_name) = query.provider_name.as_deref() {
|
||||
if item.provider_name != provider_name {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(model) = query.model.as_deref() {
|
||||
if item.model != model {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
})
|
||||
.cloned()
|
||||
.collect();
|
||||
items.sort_by(|left, right| {
|
||||
left.created_at_unix_secs
|
||||
.cmp(&right.created_at_unix_secs)
|
||||
.then_with(|| left.request_id.cmp(&right.request_id))
|
||||
});
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
async fn list_recent_usage_audits(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let mut items: Vec<_> = self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.values()
|
||||
.filter(|item| match user_id {
|
||||
Some(user_id) => item.user_id.as_deref() == Some(user_id),
|
||||
None => true,
|
||||
})
|
||||
.cloned()
|
||||
.collect();
|
||||
items.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.cmp(&left.created_at_unix_secs)
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
items.truncate(limit);
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
async fn summarize_total_tokens_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<BTreeMap<String, u64>, DataLayerError> {
|
||||
let api_key_id_set = api_key_ids.iter().map(String::as_str).collect::<Vec<_>>();
|
||||
let mut totals = BTreeMap::<String, u64>::new();
|
||||
for item in self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.values()
|
||||
{
|
||||
let Some(api_key_id) = item.api_key_id.as_deref() else {
|
||||
continue;
|
||||
};
|
||||
if !api_key_id_set.contains(&api_key_id) {
|
||||
continue;
|
||||
}
|
||||
let entry = totals.entry(api_key_id.to_string()).or_insert(0);
|
||||
*entry = (*entry).saturating_add(item.total_tokens);
|
||||
}
|
||||
Ok(totals)
|
||||
}
|
||||
|
||||
async fn summarize_provider_usage_since(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
since_unix_secs: u64,
|
||||
) -> Result<StoredProviderUsageSummary, DataLayerError> {
|
||||
let windows = self
|
||||
.provider_usage_windows
|
||||
.read()
|
||||
.expect("provider usage repository lock");
|
||||
|
||||
let mut summary = StoredProviderUsageSummary::default();
|
||||
let mut response_time_samples = 0u64;
|
||||
for window in windows.iter().filter(|window| {
|
||||
window.provider_id == provider_id && window.window_start_unix_secs >= since_unix_secs
|
||||
}) {
|
||||
summary.total_requests = summary.total_requests.saturating_add(window.total_requests);
|
||||
summary.successful_requests = summary
|
||||
.successful_requests
|
||||
.saturating_add(window.successful_requests);
|
||||
summary.failed_requests = summary
|
||||
.failed_requests
|
||||
.saturating_add(window.failed_requests);
|
||||
summary.total_cost_usd += window.total_cost_usd;
|
||||
summary.avg_response_time_ms += window.avg_response_time_ms;
|
||||
response_time_samples = response_time_samples.saturating_add(1);
|
||||
}
|
||||
|
||||
if response_time_samples > 0 {
|
||||
summary.avg_response_time_ms /= response_time_samples as f64;
|
||||
}
|
||||
|
||||
Ok(summary)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageWriteRepository for InMemoryUsageReadRepository {
|
||||
async fn upsert(
|
||||
&self,
|
||||
usage: UpsertUsageRecord,
|
||||
) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
||||
usage.validate()?;
|
||||
let mut by_request_id = self.by_request_id.write().expect("usage repository lock");
|
||||
|
||||
let created_at_unix_secs = by_request_id
|
||||
.get(&usage.request_id)
|
||||
.map(|existing| existing.created_at_unix_secs)
|
||||
.or(usage.created_at_unix_secs)
|
||||
.unwrap_or(usage.updated_at_unix_secs);
|
||||
|
||||
let total_tokens = usage
|
||||
.total_tokens
|
||||
.or_else(|| {
|
||||
Some(
|
||||
usage.input_tokens.unwrap_or_default()
|
||||
+ usage.output_tokens.unwrap_or_default(),
|
||||
)
|
||||
})
|
||||
.unwrap_or_default();
|
||||
let existing = by_request_id.get(&usage.request_id);
|
||||
|
||||
let stored = StoredRequestUsageAudit {
|
||||
id: existing
|
||||
.map(|existing| existing.id.clone())
|
||||
.unwrap_or_else(|| format!("usage-{}", usage.request_id)),
|
||||
request_id: usage.request_id.clone(),
|
||||
user_id: usage.user_id,
|
||||
api_key_id: usage.api_key_id,
|
||||
username: usage.username,
|
||||
api_key_name: usage.api_key_name,
|
||||
provider_name: usage.provider_name,
|
||||
model: usage.model,
|
||||
target_model: usage.target_model,
|
||||
provider_id: usage.provider_id,
|
||||
provider_endpoint_id: usage.provider_endpoint_id,
|
||||
provider_api_key_id: usage.provider_api_key_id,
|
||||
request_type: usage.request_type,
|
||||
api_format: usage.api_format,
|
||||
api_family: usage.api_family,
|
||||
endpoint_kind: usage.endpoint_kind,
|
||||
endpoint_api_format: usage.endpoint_api_format,
|
||||
provider_api_family: usage.provider_api_family,
|
||||
provider_endpoint_kind: usage.provider_endpoint_kind,
|
||||
has_format_conversion: usage.has_format_conversion.unwrap_or(false),
|
||||
is_stream: usage.is_stream.unwrap_or(false),
|
||||
input_tokens: usage.input_tokens.unwrap_or_default(),
|
||||
output_tokens: usage.output_tokens.unwrap_or_default(),
|
||||
total_tokens,
|
||||
cache_creation_input_tokens: usage.cache_creation_input_tokens.unwrap_or_else(|| {
|
||||
existing
|
||||
.map(|existing| existing.cache_creation_input_tokens)
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
cache_read_input_tokens: usage.cache_read_input_tokens.unwrap_or_else(|| {
|
||||
existing
|
||||
.map(|existing| existing.cache_read_input_tokens)
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
cache_creation_cost_usd: usage.cache_creation_cost_usd.unwrap_or_else(|| {
|
||||
existing
|
||||
.map(|existing| existing.cache_creation_cost_usd)
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
cache_read_cost_usd: usage.cache_read_cost_usd.unwrap_or_else(|| {
|
||||
existing
|
||||
.map(|existing| existing.cache_read_cost_usd)
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
output_price_per_1m: usage
|
||||
.output_price_per_1m
|
||||
.or_else(|| existing.and_then(|existing| existing.output_price_per_1m)),
|
||||
total_cost_usd: usage.total_cost_usd.unwrap_or_else(|| {
|
||||
existing
|
||||
.map(|existing| existing.total_cost_usd)
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
actual_total_cost_usd: usage.actual_total_cost_usd.unwrap_or_else(|| {
|
||||
existing
|
||||
.map(|existing| existing.actual_total_cost_usd)
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
status_code: usage.status_code,
|
||||
error_message: usage.error_message,
|
||||
error_category: usage.error_category,
|
||||
response_time_ms: usage.response_time_ms,
|
||||
first_byte_time_ms: usage.first_byte_time_ms,
|
||||
status: usage.status,
|
||||
billing_status: usage.billing_status,
|
||||
created_at_unix_secs,
|
||||
updated_at_unix_secs: usage.updated_at_unix_secs,
|
||||
finalized_at_unix_secs: usage.finalized_at_unix_secs,
|
||||
};
|
||||
|
||||
by_request_id.insert(stored.request_id.clone(), stored.clone());
|
||||
Ok(stored)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryUsageReadRepository;
|
||||
use crate::repository::usage::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
use crate::repository::usage::{
|
||||
StoredProviderUsageWindow, StoredRequestUsageAudit, UpsertUsageRecord, UsageReadRepository,
|
||||
UsageWriteRepository,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
fn sample_usage(request_id: &str, created_at_unix_secs: i64) -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
@@ -104,4 +370,124 @@ mod tests {
|
||||
assert_eq!(usage.request_id, "req-2");
|
||||
assert_eq!(usage.total_tokens, 150);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upsert_writes_usage_record() {
|
||||
let repository = InMemoryUsageReadRepository::default();
|
||||
let stored = repository
|
||||
.upsert(UpsertUsageRecord {
|
||||
request_id: "req-upsert-1".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("key-1".to_string()),
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
target_model: Some("gpt-5-mini".to_string()),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||||
provider_api_key_id: Some("provider-key-1".to_string()),
|
||||
request_type: Some("chat".to_string()),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_api_format: Some("openai:chat".to_string()),
|
||||
provider_api_family: Some("openai".to_string()),
|
||||
provider_endpoint_kind: Some("chat".to_string()),
|
||||
has_format_conversion: Some(false),
|
||||
is_stream: Some(true),
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(20),
|
||||
total_tokens: None,
|
||||
cache_creation_input_tokens: None,
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_cost_usd: None,
|
||||
cache_read_cost_usd: None,
|
||||
output_price_per_1m: None,
|
||||
total_cost_usd: Some(0.25),
|
||||
actual_total_cost_usd: Some(0.15),
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
error_category: None,
|
||||
response_time_ms: Some(300),
|
||||
first_byte_time_ms: Some(120),
|
||||
status: "completed".to_string(),
|
||||
billing_status: "pending".to_string(),
|
||||
request_headers: Some(json!({"authorization": "Bearer test"})),
|
||||
request_body: Some(json!({"model": "gpt-5"})),
|
||||
provider_request_headers: None,
|
||||
provider_request_body: None,
|
||||
response_headers: None,
|
||||
response_body: None,
|
||||
client_response_headers: None,
|
||||
client_response_body: None,
|
||||
request_metadata: None,
|
||||
finalized_at_unix_secs: None,
|
||||
created_at_unix_secs: Some(100),
|
||||
updated_at_unix_secs: 101,
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
assert_eq!(stored.request_id, "req-upsert-1");
|
||||
assert_eq!(stored.total_tokens, 30);
|
||||
assert_eq!(stored.total_cost_usd, 0.25);
|
||||
assert_eq!(stored.actual_total_cost_usd, 0.15);
|
||||
assert_eq!(
|
||||
repository
|
||||
.find_by_request_id("req-upsert-1")
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.expect("usage should exist")
|
||||
.model,
|
||||
"gpt-5"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn summarizes_provider_usage_windows_since_timestamp() {
|
||||
let repository = InMemoryUsageReadRepository::default().with_provider_usage_windows(vec![
|
||||
StoredProviderUsageWindow::new(
|
||||
"provider-1".to_string(),
|
||||
1_700_000_000,
|
||||
10,
|
||||
9,
|
||||
1,
|
||||
120.0,
|
||||
1.25,
|
||||
)
|
||||
.expect("window should build"),
|
||||
StoredProviderUsageWindow::new(
|
||||
"provider-1".to_string(),
|
||||
1_700_003_600,
|
||||
6,
|
||||
5,
|
||||
1,
|
||||
180.0,
|
||||
0.75,
|
||||
)
|
||||
.expect("window should build"),
|
||||
StoredProviderUsageWindow::new(
|
||||
"provider-2".to_string(),
|
||||
1_700_003_600,
|
||||
99,
|
||||
99,
|
||||
0,
|
||||
50.0,
|
||||
5.0,
|
||||
)
|
||||
.expect("window should build"),
|
||||
]);
|
||||
|
||||
let summary = repository
|
||||
.summarize_provider_usage_since("provider-1", 1_700_000_100)
|
||||
.await
|
||||
.expect("summary should succeed");
|
||||
|
||||
assert_eq!(summary.total_requests, 6);
|
||||
assert_eq!(summary.successful_requests, 5);
|
||||
assert_eq!(summary.failed_requests, 1);
|
||||
assert_eq!(summary.avg_response_time_ms, 180.0);
|
||||
assert_eq!(summary.total_cost_usd, 0.75);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,4 +4,8 @@ mod types;
|
||||
|
||||
pub use memory::InMemoryUsageReadRepository;
|
||||
pub use sql::SqlxUsageReadRepository;
|
||||
pub use types::{StoredRequestUsageAudit, UsageReadRepository, UsageRepository};
|
||||
pub use types::{
|
||||
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
|
||||
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageRepository,
|
||||
UsageWriteRepository,
|
||||
};
|
||||
|
||||
@@ -1,7 +1,13 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
use futures_util::future::BoxFuture;
|
||||
use sqlx::{PgPool, Postgres, QueryBuilder, Row};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::types::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
use super::types::{
|
||||
StoredProviderUsageSummary, StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery,
|
||||
UsageReadRepository, UsageWriteRepository,
|
||||
};
|
||||
use crate::postgres::PostgresTransactionRunner;
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_BY_REQUEST_ID_SQL: &str = r#"
|
||||
@@ -30,6 +36,11 @@ SELECT
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
COALESCE(cache_creation_input_tokens, 0) AS cache_creation_input_tokens,
|
||||
COALESCE(cache_read_input_tokens, 0) AS cache_read_input_tokens,
|
||||
COALESCE(CAST(cache_creation_cost_usd AS DOUBLE PRECISION), 0) AS cache_creation_cost_usd,
|
||||
COALESCE(CAST(cache_read_cost_usd AS DOUBLE PRECISION), 0) AS cache_read_cost_usd,
|
||||
CAST(output_price_per_1m AS DOUBLE PRECISION) AS output_price_per_1m,
|
||||
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
|
||||
COALESCE(CAST(actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
|
||||
status_code,
|
||||
@@ -40,27 +51,399 @@ SELECT
|
||||
status,
|
||||
billing_status,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
FROM "usage"
|
||||
WHERE request_id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const FIND_BY_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
COALESCE(has_format_conversion, FALSE) AS has_format_conversion,
|
||||
COALESCE(is_stream, FALSE) AS is_stream,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
COALESCE(cache_creation_input_tokens, 0) AS cache_creation_input_tokens,
|
||||
COALESCE(cache_read_input_tokens, 0) AS cache_read_input_tokens,
|
||||
COALESCE(CAST(cache_creation_cost_usd AS DOUBLE PRECISION), 0) AS cache_creation_cost_usd,
|
||||
COALESCE(CAST(cache_read_cost_usd AS DOUBLE PRECISION), 0) AS cache_read_cost_usd,
|
||||
CAST(output_price_per_1m AS DOUBLE PRECISION) AS output_price_per_1m,
|
||||
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
|
||||
COALESCE(CAST(actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
|
||||
status_code,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms,
|
||||
first_byte_time_ms,
|
||||
status,
|
||||
billing_status,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
FROM "usage"
|
||||
WHERE id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const SUMMARIZE_PROVIDER_USAGE_SINCE_SQL: &str = r#"
|
||||
SELECT
|
||||
COALESCE(SUM(total_requests), 0) AS total_requests,
|
||||
COALESCE(SUM(successful_requests), 0) AS successful_requests,
|
||||
COALESCE(SUM(failed_requests), 0) AS failed_requests,
|
||||
COALESCE(AVG(avg_response_time_ms), 0) AS avg_response_time_ms,
|
||||
COALESCE(SUM(total_cost_usd), 0) AS total_cost_usd
|
||||
FROM provider_usage_tracking
|
||||
WHERE provider_id = $1
|
||||
AND window_start >= TO_TIMESTAMP($2::double precision)
|
||||
"#;
|
||||
|
||||
const SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL: &str = r#"
|
||||
SELECT
|
||||
api_key_id,
|
||||
COALESCE(
|
||||
SUM(
|
||||
COALESCE(
|
||||
total_tokens,
|
||||
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)
|
||||
)
|
||||
),
|
||||
0
|
||||
) AS total_tokens
|
||||
FROM "usage"
|
||||
WHERE api_key_id = ANY($1::TEXT[])
|
||||
GROUP BY api_key_id
|
||||
ORDER BY api_key_id ASC
|
||||
"#;
|
||||
|
||||
const LIST_USAGE_AUDITS_PREFIX: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
COALESCE(has_format_conversion, FALSE) AS has_format_conversion,
|
||||
COALESCE(is_stream, FALSE) AS is_stream,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
COALESCE(cache_creation_input_tokens, 0) AS cache_creation_input_tokens,
|
||||
COALESCE(cache_read_input_tokens, 0) AS cache_read_input_tokens,
|
||||
COALESCE(CAST(cache_creation_cost_usd AS DOUBLE PRECISION), 0) AS cache_creation_cost_usd,
|
||||
COALESCE(CAST(cache_read_cost_usd AS DOUBLE PRECISION), 0) AS cache_read_cost_usd,
|
||||
CAST(output_price_per_1m AS DOUBLE PRECISION) AS output_price_per_1m,
|
||||
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
|
||||
COALESCE(CAST(actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
|
||||
status_code,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms,
|
||||
first_byte_time_ms,
|
||||
status,
|
||||
billing_status,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
FROM "usage"
|
||||
"#;
|
||||
|
||||
const LIST_RECENT_USAGE_AUDITS_PREFIX: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
COALESCE(has_format_conversion, FALSE) AS has_format_conversion,
|
||||
COALESCE(is_stream, FALSE) AS is_stream,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
COALESCE(cache_creation_input_tokens, 0) AS cache_creation_input_tokens,
|
||||
COALESCE(cache_read_input_tokens, 0) AS cache_read_input_tokens,
|
||||
COALESCE(CAST(cache_creation_cost_usd AS DOUBLE PRECISION), 0) AS cache_creation_cost_usd,
|
||||
COALESCE(CAST(cache_read_cost_usd AS DOUBLE PRECISION), 0) AS cache_read_cost_usd,
|
||||
CAST(output_price_per_1m AS DOUBLE PRECISION) AS output_price_per_1m,
|
||||
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
|
||||
COALESCE(CAST(actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
|
||||
status_code,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms,
|
||||
first_byte_time_ms,
|
||||
status,
|
||||
billing_status,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
FROM "usage"
|
||||
"#;
|
||||
|
||||
const UPSERT_SQL: &str = r#"
|
||||
INSERT INTO "usage" (
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
has_format_conversion,
|
||||
is_stream,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
cache_creation_input_tokens,
|
||||
cache_read_input_tokens,
|
||||
cache_creation_cost_usd,
|
||||
cache_read_cost_usd,
|
||||
output_price_per_1m,
|
||||
total_cost_usd,
|
||||
actual_total_cost_usd,
|
||||
status_code,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms,
|
||||
first_byte_time_ms,
|
||||
status,
|
||||
billing_status,
|
||||
request_headers,
|
||||
request_body,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
response_headers,
|
||||
response_body,
|
||||
client_response_headers,
|
||||
client_response_body,
|
||||
request_metadata,
|
||||
finalized_at,
|
||||
created_at
|
||||
) VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4,
|
||||
$5,
|
||||
$6,
|
||||
$7,
|
||||
$8,
|
||||
$9,
|
||||
$10,
|
||||
$11,
|
||||
$12,
|
||||
$13,
|
||||
$14,
|
||||
$15,
|
||||
$16,
|
||||
$17,
|
||||
$18,
|
||||
$19,
|
||||
COALESCE($20, FALSE),
|
||||
COALESCE($21, FALSE),
|
||||
COALESCE($22, 0),
|
||||
COALESCE($23, 0),
|
||||
COALESCE($24, COALESCE($22, 0) + COALESCE($23, 0)),
|
||||
COALESCE($25, 0),
|
||||
COALESCE($26, 0),
|
||||
COALESCE($27, 0),
|
||||
COALESCE($28, 0),
|
||||
$29,
|
||||
COALESCE($30, 0),
|
||||
COALESCE($31, 0),
|
||||
$32,
|
||||
$33,
|
||||
$34,
|
||||
$35,
|
||||
$36,
|
||||
$37,
|
||||
$38,
|
||||
$39,
|
||||
$40,
|
||||
$41,
|
||||
$42,
|
||||
$43,
|
||||
$44,
|
||||
$45,
|
||||
$46,
|
||||
CASE
|
||||
WHEN $47 IS NULL THEN NULL
|
||||
ELSE TO_TIMESTAMP($47::double precision)
|
||||
END,
|
||||
COALESCE(TO_TIMESTAMP($48::double precision), NOW())
|
||||
)
|
||||
ON CONFLICT (request_id)
|
||||
DO UPDATE SET
|
||||
user_id = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.user_id, "usage".user_id) ELSE "usage".user_id END,
|
||||
api_key_id = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.api_key_id, "usage".api_key_id) ELSE "usage".api_key_id END,
|
||||
username = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.username, "usage".username) ELSE "usage".username END,
|
||||
api_key_name = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.api_key_name, "usage".api_key_name) ELSE "usage".api_key_name END,
|
||||
provider_name = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_name, "usage".provider_name) ELSE "usage".provider_name END,
|
||||
model = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.model, "usage".model) ELSE "usage".model END,
|
||||
target_model = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.target_model, "usage".target_model) ELSE "usage".target_model END,
|
||||
provider_id = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_id, "usage".provider_id) ELSE "usage".provider_id END,
|
||||
provider_endpoint_id = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_endpoint_id, "usage".provider_endpoint_id) ELSE "usage".provider_endpoint_id END,
|
||||
provider_api_key_id = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_api_key_id, "usage".provider_api_key_id) ELSE "usage".provider_api_key_id END,
|
||||
request_type = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.request_type, "usage".request_type) ELSE "usage".request_type END,
|
||||
api_format = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.api_format, "usage".api_format) ELSE "usage".api_format END,
|
||||
api_family = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.api_family, "usage".api_family) ELSE "usage".api_family END,
|
||||
endpoint_kind = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.endpoint_kind, "usage".endpoint_kind) ELSE "usage".endpoint_kind END,
|
||||
endpoint_api_format = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.endpoint_api_format, "usage".endpoint_api_format) ELSE "usage".endpoint_api_format END,
|
||||
provider_api_family = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_api_family, "usage".provider_api_family) ELSE "usage".provider_api_family END,
|
||||
provider_endpoint_kind = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_endpoint_kind, "usage".provider_endpoint_kind) ELSE "usage".provider_endpoint_kind END,
|
||||
has_format_conversion = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.has_format_conversion, "usage".has_format_conversion) ELSE "usage".has_format_conversion END,
|
||||
is_stream = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.is_stream, "usage".is_stream) ELSE "usage".is_stream END,
|
||||
input_tokens = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.input_tokens, "usage".input_tokens) ELSE "usage".input_tokens END,
|
||||
output_tokens = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.output_tokens, "usage".output_tokens) ELSE "usage".output_tokens END,
|
||||
total_tokens = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.total_tokens, "usage".total_tokens) ELSE "usage".total_tokens END,
|
||||
cache_creation_input_tokens = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.cache_creation_input_tokens, "usage".cache_creation_input_tokens) ELSE "usage".cache_creation_input_tokens END,
|
||||
cache_read_input_tokens = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.cache_read_input_tokens, "usage".cache_read_input_tokens) ELSE "usage".cache_read_input_tokens END,
|
||||
cache_creation_cost_usd = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.cache_creation_cost_usd, "usage".cache_creation_cost_usd) ELSE "usage".cache_creation_cost_usd END,
|
||||
cache_read_cost_usd = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.cache_read_cost_usd, "usage".cache_read_cost_usd) ELSE "usage".cache_read_cost_usd END,
|
||||
output_price_per_1m = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.output_price_per_1m, "usage".output_price_per_1m) ELSE "usage".output_price_per_1m END,
|
||||
total_cost_usd = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.total_cost_usd, "usage".total_cost_usd) ELSE "usage".total_cost_usd END,
|
||||
actual_total_cost_usd = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.actual_total_cost_usd, "usage".actual_total_cost_usd) ELSE "usage".actual_total_cost_usd END,
|
||||
status_code = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.status_code, "usage".status_code) ELSE "usage".status_code END,
|
||||
error_message = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.error_message, "usage".error_message) ELSE "usage".error_message END,
|
||||
error_category = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.error_category, "usage".error_category) ELSE "usage".error_category END,
|
||||
response_time_ms = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.response_time_ms, "usage".response_time_ms) ELSE "usage".response_time_ms END,
|
||||
first_byte_time_ms = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.first_byte_time_ms, "usage".first_byte_time_ms) ELSE "usage".first_byte_time_ms END,
|
||||
status = CASE WHEN "usage".billing_status = 'pending' THEN EXCLUDED.status ELSE "usage".status END,
|
||||
billing_status = CASE WHEN "usage".billing_status = 'pending' THEN EXCLUDED.billing_status ELSE "usage".billing_status END,
|
||||
request_headers = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.request_headers, "usage".request_headers) ELSE "usage".request_headers END,
|
||||
request_body = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.request_body, "usage".request_body) ELSE "usage".request_body END,
|
||||
provider_request_headers = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_request_headers, "usage".provider_request_headers) ELSE "usage".provider_request_headers END,
|
||||
provider_request_body = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_request_body, "usage".provider_request_body) ELSE "usage".provider_request_body END,
|
||||
response_headers = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.response_headers, "usage".response_headers) ELSE "usage".response_headers END,
|
||||
response_body = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.response_body, "usage".response_body) ELSE "usage".response_body END,
|
||||
client_response_headers = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.client_response_headers, "usage".client_response_headers) ELSE "usage".client_response_headers END,
|
||||
client_response_body = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.client_response_body, "usage".client_response_body) ELSE "usage".client_response_body END,
|
||||
request_metadata = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.request_metadata, "usage".request_metadata) ELSE "usage".request_metadata END,
|
||||
finalized_at = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.finalized_at, "usage".finalized_at) ELSE "usage".finalized_at END
|
||||
RETURNING
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
COALESCE(has_format_conversion, FALSE) AS has_format_conversion,
|
||||
COALESCE(is_stream, FALSE) AS is_stream,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
COALESCE(cache_creation_input_tokens, 0) AS cache_creation_input_tokens,
|
||||
COALESCE(cache_read_input_tokens, 0) AS cache_read_input_tokens,
|
||||
COALESCE(CAST(cache_creation_cost_usd AS DOUBLE PRECISION), 0) AS cache_creation_cost_usd,
|
||||
COALESCE(CAST(cache_read_cost_usd AS DOUBLE PRECISION), 0) AS cache_read_cost_usd,
|
||||
CAST(output_price_per_1m AS DOUBLE PRECISION) AS output_price_per_1m,
|
||||
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
|
||||
COALESCE(CAST(actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
|
||||
status_code,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms,
|
||||
first_byte_time_ms,
|
||||
status,
|
||||
billing_status,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxUsageReadRepository {
|
||||
pool: PgPool,
|
||||
tx_runner: PostgresTransactionRunner,
|
||||
}
|
||||
|
||||
impl SqlxUsageReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
let tx_runner = PostgresTransactionRunner::new(pool.clone());
|
||||
Self { pool, tx_runner }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub fn transaction_runner(&self) -> &PostgresTransactionRunner {
|
||||
&self.tx_runner
|
||||
}
|
||||
|
||||
pub async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
@@ -71,20 +454,262 @@ impl SqlxUsageReadRepository {
|
||||
.await?;
|
||||
row.as_ref().map(map_usage_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn find_by_id(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_ID_SQL)
|
||||
.bind(id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_usage_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn summarize_provider_usage_since(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
since_unix_secs: u64,
|
||||
) -> Result<StoredProviderUsageSummary, DataLayerError> {
|
||||
let row = sqlx::query(SUMMARIZE_PROVIDER_USAGE_SINCE_SQL)
|
||||
.bind(provider_id)
|
||||
.bind(since_unix_secs as f64)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(StoredProviderUsageSummary {
|
||||
total_requests: row.try_get::<i64, _>("total_requests")?.max(0) as u64,
|
||||
successful_requests: row.try_get::<i64, _>("successful_requests")?.max(0) as u64,
|
||||
failed_requests: row.try_get::<i64, _>("failed_requests")?.max(0) as u64,
|
||||
avg_response_time_ms: row.try_get::<f64, _>("avg_response_time_ms")?,
|
||||
total_cost_usd: row.try_get::<f64, _>("total_cost_usd")?,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn list_usage_audits(
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let mut builder = QueryBuilder::<Postgres>::new(LIST_USAGE_AUDITS_PREFIX);
|
||||
let mut has_where = false;
|
||||
|
||||
if let Some(created_from_unix_secs) = query.created_from_unix_secs {
|
||||
builder.push(if has_where { " AND " } else { " WHERE " });
|
||||
has_where = true;
|
||||
builder
|
||||
.push("created_at >= TO_TIMESTAMP(")
|
||||
.push_bind(created_from_unix_secs as f64)
|
||||
.push("::double precision)");
|
||||
}
|
||||
if let Some(created_until_unix_secs) = query.created_until_unix_secs {
|
||||
builder.push(if has_where { " AND " } else { " WHERE " });
|
||||
has_where = true;
|
||||
builder
|
||||
.push("created_at < TO_TIMESTAMP(")
|
||||
.push_bind(created_until_unix_secs as f64)
|
||||
.push("::double precision)");
|
||||
}
|
||||
if let Some(user_id) = query.user_id.as_deref() {
|
||||
builder.push(if has_where { " AND " } else { " WHERE " });
|
||||
has_where = true;
|
||||
builder.push("user_id = ").push_bind(user_id.to_string());
|
||||
}
|
||||
if let Some(provider_name) = query.provider_name.as_deref() {
|
||||
builder.push(if has_where { " AND " } else { " WHERE " });
|
||||
has_where = true;
|
||||
builder
|
||||
.push("provider_name = ")
|
||||
.push_bind(provider_name.to_string());
|
||||
}
|
||||
if let Some(model) = query.model.as_deref() {
|
||||
builder.push(if has_where { " AND " } else { " WHERE " });
|
||||
builder.push("model = ").push_bind(model.to_string());
|
||||
}
|
||||
|
||||
builder.push(" ORDER BY created_at ASC, request_id ASC");
|
||||
let rows = builder.build().fetch_all(&self.pool).await?;
|
||||
rows.iter().map(map_usage_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_recent_usage_audits(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let mut builder = QueryBuilder::<Postgres>::new(LIST_RECENT_USAGE_AUDITS_PREFIX);
|
||||
if let Some(user_id) = user_id {
|
||||
builder
|
||||
.push(" WHERE user_id = ")
|
||||
.push_bind(user_id.to_string());
|
||||
}
|
||||
builder
|
||||
.push(" ORDER BY created_at DESC, id ASC LIMIT ")
|
||||
.push_bind(i64::try_from(limit).map_err(|_| {
|
||||
DataLayerError::InvalidInput(format!("invalid recent usage limit: {limit}"))
|
||||
})?);
|
||||
let rows = builder.build().fetch_all(&self.pool).await?;
|
||||
rows.iter().map(map_usage_row).collect()
|
||||
}
|
||||
|
||||
pub async fn summarize_total_tokens_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<std::collections::BTreeMap<String, u64>, DataLayerError> {
|
||||
if api_key_ids.is_empty() {
|
||||
return Ok(std::collections::BTreeMap::new());
|
||||
}
|
||||
|
||||
let rows = sqlx::query(SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL)
|
||||
.bind(api_key_ids)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
let mut totals = std::collections::BTreeMap::new();
|
||||
for row in rows {
|
||||
let api_key_id: String = row.try_get("api_key_id")?;
|
||||
let total_tokens = row.try_get::<i64, _>("total_tokens")?.max(0) as u64;
|
||||
totals.insert(api_key_id, total_tokens);
|
||||
}
|
||||
Ok(totals)
|
||||
}
|
||||
|
||||
pub async fn upsert(
|
||||
&self,
|
||||
usage: UpsertUsageRecord,
|
||||
) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
||||
usage.validate()?;
|
||||
self.tx_runner
|
||||
.run_read_write(|tx| {
|
||||
Box::pin(async move {
|
||||
let row = sqlx::query(UPSERT_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(&usage.request_id)
|
||||
.bind(&usage.user_id)
|
||||
.bind(&usage.api_key_id)
|
||||
.bind(&usage.username)
|
||||
.bind(&usage.api_key_name)
|
||||
.bind(&usage.provider_name)
|
||||
.bind(&usage.model)
|
||||
.bind(&usage.target_model)
|
||||
.bind(&usage.provider_id)
|
||||
.bind(&usage.provider_endpoint_id)
|
||||
.bind(&usage.provider_api_key_id)
|
||||
.bind(&usage.request_type)
|
||||
.bind(&usage.api_format)
|
||||
.bind(&usage.api_family)
|
||||
.bind(&usage.endpoint_kind)
|
||||
.bind(&usage.endpoint_api_format)
|
||||
.bind(&usage.provider_api_family)
|
||||
.bind(&usage.provider_endpoint_kind)
|
||||
.bind(usage.has_format_conversion)
|
||||
.bind(usage.is_stream)
|
||||
.bind(usage.input_tokens.map(to_i32).transpose()?)
|
||||
.bind(usage.output_tokens.map(to_i32).transpose()?)
|
||||
.bind(
|
||||
usage
|
||||
.total_tokens
|
||||
.or_else(|| {
|
||||
Some(
|
||||
usage.input_tokens.unwrap_or_default()
|
||||
+ usage.output_tokens.unwrap_or_default(),
|
||||
)
|
||||
})
|
||||
.map(to_i32)
|
||||
.transpose()?,
|
||||
)
|
||||
.bind(usage.cache_creation_input_tokens.map(to_i32).transpose()?)
|
||||
.bind(usage.cache_read_input_tokens.map(to_i32).transpose()?)
|
||||
.bind(usage.cache_creation_cost_usd)
|
||||
.bind(usage.cache_read_cost_usd)
|
||||
.bind(usage.output_price_per_1m)
|
||||
.bind(usage.total_cost_usd)
|
||||
.bind(usage.actual_total_cost_usd)
|
||||
.bind(usage.status_code.map(i32::from))
|
||||
.bind(&usage.error_message)
|
||||
.bind(&usage.error_category)
|
||||
.bind(usage.response_time_ms.map(to_i32).transpose()?)
|
||||
.bind(usage.first_byte_time_ms.map(to_i32).transpose()?)
|
||||
.bind(&usage.status)
|
||||
.bind(&usage.billing_status)
|
||||
.bind(&usage.request_headers)
|
||||
.bind(&usage.request_body)
|
||||
.bind(&usage.provider_request_headers)
|
||||
.bind(&usage.provider_request_body)
|
||||
.bind(&usage.response_headers)
|
||||
.bind(&usage.response_body)
|
||||
.bind(&usage.client_response_headers)
|
||||
.bind(&usage.client_response_body)
|
||||
.bind(&usage.request_metadata)
|
||||
.bind(usage.finalized_at_unix_secs.map(|value| value as f64))
|
||||
.bind(usage.created_at_unix_secs.map(|value| value as f64))
|
||||
.fetch_one(&mut **tx)
|
||||
.await?;
|
||||
map_usage_row(&row)
|
||||
}) as BoxFuture<'_, Result<StoredRequestUsageAudit, DataLayerError>>
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageReadRepository for SqlxUsageReadRepository {
|
||||
async fn find_by_id(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Self::find_by_id(self, id).await
|
||||
}
|
||||
|
||||
async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Self::find_by_request_id(self, request_id).await
|
||||
}
|
||||
|
||||
async fn list_usage_audits(
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Self::list_usage_audits(self, query).await
|
||||
}
|
||||
|
||||
async fn list_recent_usage_audits(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Self::list_recent_usage_audits(self, user_id, limit).await
|
||||
}
|
||||
|
||||
async fn summarize_total_tokens_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<std::collections::BTreeMap<String, u64>, DataLayerError> {
|
||||
Self::summarize_total_tokens_by_api_key_ids(self, api_key_ids).await
|
||||
}
|
||||
|
||||
async fn summarize_provider_usage_since(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
since_unix_secs: u64,
|
||||
) -> Result<StoredProviderUsageSummary, DataLayerError> {
|
||||
Self::summarize_provider_usage_since(self, provider_id, since_unix_secs).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageWriteRepository for SqlxUsageReadRepository {
|
||||
async fn upsert(
|
||||
&self,
|
||||
usage: UpsertUsageRecord,
|
||||
) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
||||
Self::upsert(self, usage).await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_usage_row(row: &sqlx::postgres::PgRow) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
||||
StoredRequestUsageAudit::new(
|
||||
let mut usage = StoredRequestUsageAudit::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("request_id")?,
|
||||
row.try_get("user_id")?,
|
||||
@@ -121,13 +746,39 @@ fn map_usage_row(row: &sqlx::postgres::PgRow) -> Result<StoredRequestUsageAudit,
|
||||
row.try_get("created_at_unix_secs")?,
|
||||
row.try_get("updated_at_unix_secs")?,
|
||||
row.try_get("finalized_at_unix_secs")?,
|
||||
)
|
||||
)?;
|
||||
usage.cache_creation_input_tokens = row
|
||||
.try_get::<Option<i32>, _>("cache_creation_input_tokens")?
|
||||
.map(|value| to_u64(value, "usage.cache_creation_input_tokens"))
|
||||
.transpose()?
|
||||
.unwrap_or_default();
|
||||
usage.cache_read_input_tokens = row
|
||||
.try_get::<Option<i32>, _>("cache_read_input_tokens")?
|
||||
.map(|value| to_u64(value, "usage.cache_read_input_tokens"))
|
||||
.transpose()?
|
||||
.unwrap_or_default();
|
||||
usage.cache_creation_cost_usd = row.try_get::<f64, _>("cache_creation_cost_usd")?;
|
||||
usage.cache_read_cost_usd = row.try_get::<f64, _>("cache_read_cost_usd")?;
|
||||
usage.output_price_per_1m = row.try_get("output_price_per_1m")?;
|
||||
Ok(usage)
|
||||
}
|
||||
|
||||
fn to_i32(value: u64) -> Result<i32, DataLayerError> {
|
||||
i32::try_from(value).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!("invalid usage integer value: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn to_u64(value: i32, field_name: &str) -> Result<u64, DataLayerError> {
|
||||
u64::try_from(value)
|
||||
.map_err(|_| DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}")))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxUsageReadRepository;
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
use crate::repository::usage::UpsertUsageRecord;
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_constructs_from_lazy_pool() {
|
||||
@@ -146,5 +797,98 @@ mod tests {
|
||||
let pool = factory.connect_lazy().expect("pool should build");
|
||||
let repository = SqlxUsageReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
let _ = repository.transaction_runner();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn validates_upsert_before_hitting_database() {
|
||||
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 = SqlxUsageReadRepository::new(pool);
|
||||
let result = repository
|
||||
.upsert(UpsertUsageRecord {
|
||||
request_id: "".to_string(),
|
||||
user_id: None,
|
||||
api_key_id: None,
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
provider_name: "openai".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
target_model: None,
|
||||
provider_id: None,
|
||||
provider_endpoint_id: None,
|
||||
provider_api_key_id: None,
|
||||
request_type: Some("chat".to_string()),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_api_format: Some("openai:chat".to_string()),
|
||||
provider_api_family: Some("openai".to_string()),
|
||||
provider_endpoint_kind: Some("chat".to_string()),
|
||||
has_format_conversion: Some(false),
|
||||
is_stream: Some(false),
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(20),
|
||||
total_tokens: Some(30),
|
||||
cache_creation_input_tokens: None,
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_cost_usd: None,
|
||||
cache_read_cost_usd: None,
|
||||
output_price_per_1m: None,
|
||||
total_cost_usd: None,
|
||||
actual_total_cost_usd: None,
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
error_category: None,
|
||||
response_time_ms: Some(100),
|
||||
first_byte_time_ms: None,
|
||||
status: "completed".to_string(),
|
||||
billing_status: "pending".to_string(),
|
||||
request_headers: None,
|
||||
request_body: None,
|
||||
provider_request_headers: None,
|
||||
provider_request_body: None,
|
||||
response_headers: None,
|
||||
response_body: None,
|
||||
client_response_headers: None,
|
||||
client_response_body: None,
|
||||
request_metadata: None,
|
||||
finalized_at_unix_secs: None,
|
||||
created_at_unix_secs: Some(100),
|
||||
updated_at_unix_secs: 101,
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_sql_does_not_require_updated_at_column() {
|
||||
assert!(!super::FIND_BY_REQUEST_ID_SQL.contains("COALESCE(updated_at, created_at)"));
|
||||
assert!(!super::LIST_USAGE_AUDITS_PREFIX.contains("COALESCE(updated_at, created_at)"));
|
||||
assert!(!super::UPSERT_SQL.contains("\n updated_at\n"));
|
||||
assert!(!super::UPSERT_SQL.contains("updated_at = CASE"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_sql_summarizes_tokens_by_api_key_ids_in_database() {
|
||||
assert!(super::SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL.contains("GROUP BY api_key_id"));
|
||||
assert!(super::SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL.contains("ANY($1::TEXT[])"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_sql_supports_recent_usage_audits_query() {
|
||||
assert!(super::LIST_RECENT_USAGE_AUDITS_PREFIX.contains("FROM \"usage\""));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredRequestUsageAudit {
|
||||
@@ -26,6 +27,11 @@ pub struct StoredRequestUsageAudit {
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub cache_creation_input_tokens: u64,
|
||||
pub cache_read_input_tokens: u64,
|
||||
pub cache_creation_cost_usd: f64,
|
||||
pub cache_read_cost_usd: f64,
|
||||
pub output_price_per_1m: Option<f64>,
|
||||
pub total_cost_usd: f64,
|
||||
pub actual_total_cost_usd: f64,
|
||||
pub status_code: Option<u16>,
|
||||
@@ -141,6 +147,11 @@ impl StoredRequestUsageAudit {
|
||||
input_tokens: parse_u64(input_tokens, "usage.input_tokens")?,
|
||||
output_tokens: parse_u64(output_tokens, "usage.output_tokens")?,
|
||||
total_tokens: parse_u64(total_tokens, "usage.total_tokens")?,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
cache_creation_cost_usd: 0.0,
|
||||
cache_read_cost_usd: 0.0,
|
||||
output_price_per_1m: None,
|
||||
total_cost_usd,
|
||||
actual_total_cost_usd,
|
||||
status_code: parse_u16(status_code, "usage.status_code")?,
|
||||
@@ -163,19 +174,258 @@ impl StoredRequestUsageAudit {
|
||||
.transpose()?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn with_cache_input_tokens(
|
||||
mut self,
|
||||
cache_creation_input_tokens: u64,
|
||||
cache_read_input_tokens: u64,
|
||||
) -> Self {
|
||||
self.cache_creation_input_tokens = cache_creation_input_tokens;
|
||||
self.cache_read_input_tokens = cache_read_input_tokens;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderUsageWindow {
|
||||
pub provider_id: String,
|
||||
pub window_start_unix_secs: u64,
|
||||
pub total_requests: u64,
|
||||
pub successful_requests: u64,
|
||||
pub failed_requests: u64,
|
||||
pub avg_response_time_ms: f64,
|
||||
pub total_cost_usd: f64,
|
||||
}
|
||||
|
||||
impl StoredProviderUsageWindow {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
provider_id: String,
|
||||
window_start_unix_secs: i64,
|
||||
total_requests: i64,
|
||||
successful_requests: i64,
|
||||
failed_requests: i64,
|
||||
avg_response_time_ms: f64,
|
||||
total_cost_usd: f64,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if provider_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider usage window provider_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if !avg_response_time_ms.is_finite() || !total_cost_usd.is_finite() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider usage window value is not finite".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
provider_id,
|
||||
window_start_unix_secs: parse_timestamp(
|
||||
window_start_unix_secs,
|
||||
"provider_usage_tracking.window_start_unix_secs",
|
||||
)?,
|
||||
total_requests: parse_timestamp(
|
||||
total_requests,
|
||||
"provider_usage_tracking.total_requests",
|
||||
)?,
|
||||
successful_requests: parse_timestamp(
|
||||
successful_requests,
|
||||
"provider_usage_tracking.successful_requests",
|
||||
)?,
|
||||
failed_requests: parse_timestamp(
|
||||
failed_requests,
|
||||
"provider_usage_tracking.failed_requests",
|
||||
)?,
|
||||
avg_response_time_ms,
|
||||
total_cost_usd,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderUsageSummary {
|
||||
pub total_requests: u64,
|
||||
pub successful_requests: u64,
|
||||
pub failed_requests: u64,
|
||||
pub avg_response_time_ms: f64,
|
||||
pub total_cost_usd: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct UsageAuditListQuery {
|
||||
pub created_from_unix_secs: Option<u64>,
|
||||
pub created_until_unix_secs: Option<u64>,
|
||||
pub user_id: Option<String>,
|
||||
pub provider_name: Option<String>,
|
||||
pub model: Option<String>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait UsageReadRepository: Send + Sync {
|
||||
async fn find_by_id(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, crate::DataLayerError>;
|
||||
|
||||
async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, crate::DataLayerError>;
|
||||
|
||||
async fn list_usage_audits(
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, crate::DataLayerError>;
|
||||
|
||||
async fn list_recent_usage_audits(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, crate::DataLayerError>;
|
||||
|
||||
async fn summarize_total_tokens_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<std::collections::BTreeMap<String, u64>, crate::DataLayerError>;
|
||||
|
||||
async fn summarize_provider_usage_since(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
since_unix_secs: u64,
|
||||
) -> Result<StoredProviderUsageSummary, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait UsageRepository: UsageReadRepository + Send + Sync {}
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct UpsertUsageRecord {
|
||||
pub request_id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub api_key_id: Option<String>,
|
||||
pub username: Option<String>,
|
||||
pub api_key_name: Option<String>,
|
||||
pub provider_name: String,
|
||||
pub model: String,
|
||||
pub target_model: Option<String>,
|
||||
pub provider_id: Option<String>,
|
||||
pub provider_endpoint_id: Option<String>,
|
||||
pub provider_api_key_id: Option<String>,
|
||||
pub request_type: Option<String>,
|
||||
pub api_format: Option<String>,
|
||||
pub api_family: Option<String>,
|
||||
pub endpoint_kind: Option<String>,
|
||||
pub endpoint_api_format: Option<String>,
|
||||
pub provider_api_family: Option<String>,
|
||||
pub provider_endpoint_kind: Option<String>,
|
||||
pub has_format_conversion: Option<bool>,
|
||||
pub is_stream: Option<bool>,
|
||||
pub input_tokens: Option<u64>,
|
||||
pub output_tokens: Option<u64>,
|
||||
pub total_tokens: Option<u64>,
|
||||
pub cache_creation_input_tokens: Option<u64>,
|
||||
pub cache_read_input_tokens: Option<u64>,
|
||||
pub cache_creation_cost_usd: Option<f64>,
|
||||
pub cache_read_cost_usd: Option<f64>,
|
||||
pub output_price_per_1m: Option<f64>,
|
||||
pub total_cost_usd: Option<f64>,
|
||||
pub actual_total_cost_usd: Option<f64>,
|
||||
pub status_code: Option<u16>,
|
||||
pub error_message: Option<String>,
|
||||
pub error_category: Option<String>,
|
||||
pub response_time_ms: Option<u64>,
|
||||
pub first_byte_time_ms: Option<u64>,
|
||||
pub status: String,
|
||||
pub billing_status: String,
|
||||
pub request_headers: Option<Value>,
|
||||
pub request_body: Option<Value>,
|
||||
pub provider_request_headers: Option<Value>,
|
||||
pub provider_request_body: Option<Value>,
|
||||
pub response_headers: Option<Value>,
|
||||
pub response_body: Option<Value>,
|
||||
pub client_response_headers: Option<Value>,
|
||||
pub client_response_body: Option<Value>,
|
||||
pub request_metadata: Option<Value>,
|
||||
pub finalized_at_unix_secs: Option<u64>,
|
||||
pub created_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
impl<T> UsageRepository for T where T: UsageReadRepository + Send + Sync {}
|
||||
impl UpsertUsageRecord {
|
||||
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||
if self.request_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert request_id cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.provider_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert provider_name cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.model.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert model cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.status.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert status cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.billing_status.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert billing_status cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if let Some(value) = self.total_cost_usd {
|
||||
if !value.is_finite() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert total_cost_usd must be finite".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
if let Some(value) = self.cache_creation_cost_usd {
|
||||
if !value.is_finite() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert cache_creation_cost_usd must be finite".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
if let Some(value) = self.cache_read_cost_usd {
|
||||
if !value.is_finite() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert cache_read_cost_usd must be finite".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
if let Some(value) = self.output_price_per_1m {
|
||||
if !value.is_finite() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert output_price_per_1m must be finite".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
if let Some(value) = self.actual_total_cost_usd {
|
||||
if !value.is_finite() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert actual_total_cost_usd must be finite".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait UsageWriteRepository: Send + Sync {
|
||||
async fn upsert(
|
||||
&self,
|
||||
usage: UpsertUsageRecord,
|
||||
) -> Result<StoredRequestUsageAudit, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait UsageRepository: UsageReadRepository + UsageWriteRepository + Send + Sync {}
|
||||
|
||||
impl<T> UsageRepository for T where T: UsageReadRepository + UsageWriteRepository + Send + Sync {}
|
||||
|
||||
fn parse_u64(value: i32, field_name: &str) -> Result<u64, crate::DataLayerError> {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
@@ -214,7 +464,8 @@ fn parse_timestamp(value: i64, field_name: &str) -> Result<u64, crate::DataLayer
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::StoredRequestUsageAudit;
|
||||
use super::{StoredRequestUsageAudit, UpsertUsageRecord};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_request_id() {
|
||||
@@ -301,4 +552,61 @@ mod tests {
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_upsert_payload() {
|
||||
let record = UpsertUsageRecord {
|
||||
request_id: "".to_string(),
|
||||
user_id: None,
|
||||
api_key_id: None,
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
provider_name: "openai".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
target_model: None,
|
||||
provider_id: None,
|
||||
provider_endpoint_id: None,
|
||||
provider_api_key_id: None,
|
||||
request_type: Some("chat".to_string()),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_api_format: Some("openai:chat".to_string()),
|
||||
provider_api_family: Some("openai".to_string()),
|
||||
provider_endpoint_kind: Some("chat".to_string()),
|
||||
has_format_conversion: Some(false),
|
||||
is_stream: Some(false),
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(20),
|
||||
total_tokens: Some(30),
|
||||
cache_creation_input_tokens: None,
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_cost_usd: None,
|
||||
cache_read_cost_usd: None,
|
||||
output_price_per_1m: None,
|
||||
total_cost_usd: None,
|
||||
actual_total_cost_usd: None,
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
error_category: None,
|
||||
response_time_ms: Some(120),
|
||||
first_byte_time_ms: None,
|
||||
status: "completed".to_string(),
|
||||
billing_status: "pending".to_string(),
|
||||
request_headers: Some(json!({"authorization": "Bearer test"})),
|
||||
request_body: Some(json!({"model": "gpt-5"})),
|
||||
provider_request_headers: None,
|
||||
provider_request_body: None,
|
||||
response_headers: None,
|
||||
response_body: None,
|
||||
client_response_headers: None,
|
||||
client_response_body: None,
|
||||
request_metadata: None,
|
||||
finalized_at_unix_secs: None,
|
||||
created_at_unix_secs: Some(100),
|
||||
updated_at_unix_secs: 101,
|
||||
};
|
||||
|
||||
assert!(record.validate().is_err());
|
||||
}
|
||||
}
|
||||
|
||||
363
crates/aether-data/src/repository/users/memory.rs
Normal file
363
crates/aether-data/src/repository/users/memory.rs
Normal file
@@ -0,0 +1,363 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
StoredUserAuthRecord, StoredUserExportRow, StoredUserSummary, UserExportListQuery,
|
||||
UserExportSummary, UserReadRepository,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryUserReadRepository {
|
||||
by_id: RwLock<BTreeMap<String, StoredUserSummary>>,
|
||||
auth_by_id: RwLock<BTreeMap<String, StoredUserAuthRecord>>,
|
||||
auth_by_identifier: RwLock<BTreeMap<String, String>>,
|
||||
export_rows: RwLock<Vec<StoredUserExportRow>>,
|
||||
}
|
||||
|
||||
impl InMemoryUserReadRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredUserSummary>,
|
||||
{
|
||||
let mut by_id = BTreeMap::new();
|
||||
for item in items {
|
||||
by_id.insert(item.id.clone(), item);
|
||||
}
|
||||
Self {
|
||||
by_id: RwLock::new(by_id),
|
||||
auth_by_id: RwLock::new(BTreeMap::new()),
|
||||
auth_by_identifier: RwLock::new(BTreeMap::new()),
|
||||
export_rows: RwLock::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn seed_auth_users<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredUserAuthRecord>,
|
||||
{
|
||||
let mut by_id = BTreeMap::new();
|
||||
let mut auth_by_id = BTreeMap::new();
|
||||
let mut auth_by_identifier = BTreeMap::new();
|
||||
for item in items {
|
||||
let summary = item
|
||||
.to_summary()
|
||||
.expect("in-memory auth user should convert to summary");
|
||||
by_id.insert(summary.id.clone(), summary);
|
||||
auth_by_identifier.insert(item.username.clone(), item.id.clone());
|
||||
if let Some(email) = item.email.as_ref() {
|
||||
auth_by_identifier.insert(email.clone(), item.id.clone());
|
||||
}
|
||||
auth_by_id.insert(item.id.clone(), item);
|
||||
}
|
||||
Self {
|
||||
by_id: RwLock::new(by_id),
|
||||
auth_by_id: RwLock::new(auth_by_id),
|
||||
auth_by_identifier: RwLock::new(auth_by_identifier),
|
||||
export_rows: RwLock::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn seed_export_users<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredUserExportRow>,
|
||||
{
|
||||
Self {
|
||||
by_id: RwLock::new(BTreeMap::new()),
|
||||
auth_by_id: RwLock::new(BTreeMap::new()),
|
||||
auth_by_identifier: RwLock::new(BTreeMap::new()),
|
||||
export_rows: RwLock::new(items.into_iter().collect()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_export_users<I>(self, items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredUserExportRow>,
|
||||
{
|
||||
let rows = items.into_iter().collect();
|
||||
*self.export_rows.write().expect("user repository lock") = rows;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UserReadRepository for InMemoryUserReadRepository {
|
||||
async fn list_users_by_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredUserSummary>, DataLayerError> {
|
||||
let index = self.by_id.read().expect("user repository lock");
|
||||
Ok(user_ids
|
||||
.iter()
|
||||
.filter_map(|user_id| index.get(user_id).cloned())
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_non_admin_export_users(
|
||||
&self,
|
||||
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
||||
Ok(self
|
||||
.export_rows
|
||||
.read()
|
||||
.expect("user repository lock")
|
||||
.iter()
|
||||
.filter(|row| !row.role.eq_ignore_ascii_case("admin"))
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_export_users(&self) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
||||
Ok(self
|
||||
.export_rows
|
||||
.read()
|
||||
.expect("user repository lock")
|
||||
.clone())
|
||||
}
|
||||
|
||||
async fn list_export_users_page(
|
||||
&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);
|
||||
}
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.skip(query.skip)
|
||||
.take(query.limit)
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn summarize_export_users(&self) -> Result<UserExportSummary, DataLayerError> {
|
||||
let rows = self.export_rows.read().expect("user repository lock");
|
||||
Ok(UserExportSummary {
|
||||
total: rows.len() as u64,
|
||||
active: rows.iter().filter(|row| row.is_active).count() as u64,
|
||||
})
|
||||
}
|
||||
|
||||
async fn find_export_user_by_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Option<StoredUserExportRow>, DataLayerError> {
|
||||
Ok(self
|
||||
.export_rows
|
||||
.read()
|
||||
.expect("user repository lock")
|
||||
.iter()
|
||||
.find(|row| row.id == user_id)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn find_user_auth_by_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Option<StoredUserAuthRecord>, DataLayerError> {
|
||||
Ok(self
|
||||
.auth_by_id
|
||||
.read()
|
||||
.expect("user repository lock")
|
||||
.get(user_id)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn list_user_auth_by_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredUserAuthRecord>, DataLayerError> {
|
||||
let auth_by_id = self.auth_by_id.read().expect("user repository lock");
|
||||
Ok(user_ids
|
||||
.iter()
|
||||
.filter_map(|user_id| auth_by_id.get(user_id).cloned())
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn find_user_auth_by_identifier(
|
||||
&self,
|
||||
identifier: &str,
|
||||
) -> Result<Option<StoredUserAuthRecord>, DataLayerError> {
|
||||
let auth_by_identifier = self
|
||||
.auth_by_identifier
|
||||
.read()
|
||||
.expect("user repository lock");
|
||||
let Some(user_id) = auth_by_identifier.get(identifier) else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(self
|
||||
.auth_by_id
|
||||
.read()
|
||||
.expect("user repository lock")
|
||||
.get(user_id)
|
||||
.cloned())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::repository::users::{UserExportListQuery, UserReadRepository};
|
||||
|
||||
#[tokio::test]
|
||||
async fn lists_seeded_users() {
|
||||
let user = StoredUserSummary::new(
|
||||
"user-1".to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
true,
|
||||
false,
|
||||
)
|
||||
.expect("user should build");
|
||||
let repository = InMemoryUserReadRepository::seed(vec![user.clone()]);
|
||||
let rows = repository
|
||||
.list_users_by_ids(&["user-1".to_string()])
|
||||
.await
|
||||
.expect("lookup should succeed");
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(rows[0], user);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lists_seeded_non_admin_export_users() {
|
||||
let user = StoredUserExportRow::new(
|
||||
"user-1".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
true,
|
||||
"alice".to_string(),
|
||||
Some("hash".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
Some(60),
|
||||
Some(serde_json::json!({"gpt-4.1": {"cache_1h": true}})),
|
||||
true,
|
||||
)
|
||||
.expect("user export row should build");
|
||||
let repository = InMemoryUserReadRepository::seed_export_users(vec![user.clone()]);
|
||||
|
||||
let rows = repository
|
||||
.list_non_admin_export_users()
|
||||
.await
|
||||
.expect("export should succeed");
|
||||
|
||||
assert_eq!(rows, vec![user]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn finds_seeded_auth_user_by_id_and_identifier() {
|
||||
let user = StoredUserAuthRecord::new(
|
||||
"user-1".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
true,
|
||||
"alice".to_string(),
|
||||
Some("hash".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("auth user should build");
|
||||
let repository = InMemoryUserReadRepository::seed_auth_users(vec![user.clone()]);
|
||||
|
||||
let by_id = repository
|
||||
.find_user_auth_by_id("user-1")
|
||||
.await
|
||||
.expect("lookup by id should succeed");
|
||||
let by_email = repository
|
||||
.find_user_auth_by_identifier("alice@example.com")
|
||||
.await
|
||||
.expect("lookup by email should succeed");
|
||||
let by_username = repository
|
||||
.find_user_auth_by_identifier("alice")
|
||||
.await
|
||||
.expect("lookup by username should succeed");
|
||||
|
||||
assert_eq!(by_id, Some(user.clone()));
|
||||
assert_eq!(by_email, Some(user.clone()));
|
||||
assert_eq!(by_username, Some(user));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn paginates_export_users_in_memory() {
|
||||
let repository = InMemoryUserReadRepository::seed_export_users(vec![
|
||||
StoredUserExportRow::new(
|
||||
"user-1".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
true,
|
||||
"alice".to_string(),
|
||||
Some("hash".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(60),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("user export row should build"),
|
||||
StoredUserExportRow::new(
|
||||
"user-2".to_string(),
|
||||
Some("bob@example.com".to_string()),
|
||||
true,
|
||||
"bob".to_string(),
|
||||
Some("hash".to_string()),
|
||||
"admin".to_string(),
|
||||
"local".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(30),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("user export row should build"),
|
||||
StoredUserExportRow::new(
|
||||
"user-3".to_string(),
|
||||
Some("carol@example.com".to_string()),
|
||||
true,
|
||||
"carol".to_string(),
|
||||
Some("hash".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(10),
|
||||
None,
|
||||
false,
|
||||
)
|
||||
.expect("user export row should build"),
|
||||
]);
|
||||
|
||||
let rows = repository
|
||||
.list_export_users_page(&UserExportListQuery {
|
||||
skip: 0,
|
||||
limit: 10,
|
||||
role: Some("user".to_string()),
|
||||
is_active: Some(true),
|
||||
})
|
||||
.await
|
||||
.expect("paged export should succeed");
|
||||
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(rows[0].id, "user-1");
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/users/mod.rs
Normal file
10
crates/aether-data/src/repository/users/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryUserReadRepository;
|
||||
pub use sql::SqlxUserReadRepository;
|
||||
pub use types::{
|
||||
StoredUserAuthRecord, StoredUserExportRow, StoredUserSummary, UserExportListQuery,
|
||||
UserExportSummary, UserReadRepository,
|
||||
};
|
||||
408
crates/aether-data/src/repository/users/sql.rs
Normal file
408
crates/aether-data/src/repository/users/sql.rs
Normal file
@@ -0,0 +1,408 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Postgres, QueryBuilder, Row};
|
||||
|
||||
use super::types::{
|
||||
StoredUserAuthRecord, StoredUserExportRow, StoredUserSummary, UserExportListQuery,
|
||||
UserExportSummary, UserReadRepository,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const LIST_USERS_BY_IDS_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
username,
|
||||
email,
|
||||
role::text AS role,
|
||||
is_active,
|
||||
is_deleted
|
||||
FROM users
|
||||
WHERE id = ANY($1::text[])
|
||||
ORDER BY id ASC
|
||||
"#;
|
||||
|
||||
const LIST_NON_ADMIN_EXPORT_USERS_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
email,
|
||||
email_verified,
|
||||
username,
|
||||
password_hash,
|
||||
role::text AS role,
|
||||
auth_source::text AS auth_source,
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
rate_limit,
|
||||
model_capability_settings,
|
||||
is_active
|
||||
FROM users
|
||||
WHERE is_deleted IS FALSE
|
||||
AND role::text != 'admin'
|
||||
ORDER BY id ASC
|
||||
"#;
|
||||
|
||||
const LIST_EXPORT_USERS_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
email,
|
||||
email_verified,
|
||||
username,
|
||||
password_hash,
|
||||
role::text AS role,
|
||||
auth_source::text AS auth_source,
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
rate_limit,
|
||||
model_capability_settings,
|
||||
is_active
|
||||
FROM users
|
||||
WHERE is_deleted IS FALSE
|
||||
ORDER BY id ASC
|
||||
"#;
|
||||
|
||||
const LIST_EXPORT_USERS_PAGE_PREFIX: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
email,
|
||||
email_verified,
|
||||
username,
|
||||
password_hash,
|
||||
role::text AS role,
|
||||
auth_source::text AS auth_source,
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
rate_limit,
|
||||
model_capability_settings,
|
||||
is_active
|
||||
FROM users
|
||||
WHERE is_deleted IS FALSE
|
||||
"#;
|
||||
|
||||
const SUMMARIZE_EXPORT_USERS_SQL: &str = r#"
|
||||
SELECT
|
||||
COUNT(*)::BIGINT AS total,
|
||||
COUNT(*) FILTER (WHERE is_active = TRUE)::BIGINT AS active
|
||||
FROM users
|
||||
WHERE is_deleted IS FALSE
|
||||
"#;
|
||||
|
||||
const FIND_EXPORT_USER_BY_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
email,
|
||||
email_verified,
|
||||
username,
|
||||
password_hash,
|
||||
role::text AS role,
|
||||
auth_source::text AS auth_source,
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
rate_limit,
|
||||
model_capability_settings,
|
||||
is_active
|
||||
FROM users
|
||||
WHERE is_deleted IS FALSE
|
||||
AND id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const FIND_USER_AUTH_BY_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
email,
|
||||
email_verified,
|
||||
username,
|
||||
password_hash,
|
||||
role::text AS role,
|
||||
auth_source::text AS auth_source,
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
is_active,
|
||||
is_deleted,
|
||||
created_at,
|
||||
last_login_at
|
||||
FROM users
|
||||
WHERE id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const LIST_USER_AUTH_BY_IDS_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
email,
|
||||
email_verified,
|
||||
username,
|
||||
password_hash,
|
||||
role::text AS role,
|
||||
auth_source::text AS auth_source,
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
is_active,
|
||||
is_deleted,
|
||||
created_at,
|
||||
last_login_at
|
||||
FROM users
|
||||
WHERE id = ANY($1::text[])
|
||||
ORDER BY id ASC
|
||||
"#;
|
||||
|
||||
const FIND_USER_AUTH_BY_IDENTIFIER_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
email,
|
||||
email_verified,
|
||||
username,
|
||||
password_hash,
|
||||
role::text AS role,
|
||||
auth_source::text AS auth_source,
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
is_active,
|
||||
is_deleted,
|
||||
created_at,
|
||||
last_login_at
|
||||
FROM users
|
||||
WHERE email = $1 OR username = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxUserReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxUserReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub async fn list_users_by_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredUserSummary>, DataLayerError> {
|
||||
if user_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let rows = sqlx::query(LIST_USERS_BY_IDS_SQL)
|
||||
.bind(user_ids)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_user_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_non_admin_export_users(
|
||||
&self,
|
||||
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
||||
let rows = sqlx::query(LIST_NON_ADMIN_EXPORT_USERS_SQL)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_user_export_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_export_users(&self) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
||||
let rows = sqlx::query(LIST_EXPORT_USERS_SQL)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_user_export_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_export_users_page(
|
||||
&self,
|
||||
query: &UserExportListQuery,
|
||||
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
||||
let mut builder = QueryBuilder::<Postgres>::new(LIST_EXPORT_USERS_PAGE_PREFIX);
|
||||
|
||||
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);
|
||||
}
|
||||
|
||||
builder
|
||||
.push(" ORDER BY id ASC OFFSET ")
|
||||
.push_bind(i64::try_from(query.skip).map_err(|_| {
|
||||
DataLayerError::InvalidInput(format!("invalid user export skip: {}", query.skip))
|
||||
})?)
|
||||
.push(" LIMIT ")
|
||||
.push_bind(i64::try_from(query.limit).map_err(|_| {
|
||||
DataLayerError::InvalidInput(format!("invalid user export limit: {}", query.limit))
|
||||
})?);
|
||||
|
||||
let rows = builder.build().fetch_all(&self.pool).await?;
|
||||
rows.iter().map(map_user_export_row).collect()
|
||||
}
|
||||
|
||||
pub async fn summarize_export_users(&self) -> Result<UserExportSummary, DataLayerError> {
|
||||
let row = sqlx::query(SUMMARIZE_EXPORT_USERS_SQL)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
Ok(UserExportSummary {
|
||||
total: row.try_get::<i64, _>("total")?.max(0) as u64,
|
||||
active: row.try_get::<i64, _>("active")?.max(0) as u64,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn find_export_user_by_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Option<StoredUserExportRow>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_EXPORT_USER_BY_ID_SQL)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_user_export_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn list_user_auth_by_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredUserAuthRecord>, DataLayerError> {
|
||||
if user_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = sqlx::query(LIST_USER_AUTH_BY_IDS_SQL)
|
||||
.bind(user_ids)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_user_auth_row).collect()
|
||||
}
|
||||
|
||||
pub async fn find_user_auth_by_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Option<StoredUserAuthRecord>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_USER_AUTH_BY_ID_SQL)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_user_auth_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn find_user_auth_by_identifier(
|
||||
&self,
|
||||
identifier: &str,
|
||||
) -> Result<Option<StoredUserAuthRecord>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_USER_AUTH_BY_IDENTIFIER_SQL)
|
||||
.bind(identifier)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_user_auth_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
fn map_user_row(row: &sqlx::postgres::PgRow) -> Result<StoredUserSummary, DataLayerError> {
|
||||
StoredUserSummary::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("username")?,
|
||||
row.try_get("email")?,
|
||||
row.try_get("role")?,
|
||||
row.try_get("is_active")?,
|
||||
row.try_get("is_deleted")?,
|
||||
)
|
||||
}
|
||||
|
||||
fn map_user_export_row(row: &sqlx::postgres::PgRow) -> Result<StoredUserExportRow, DataLayerError> {
|
||||
StoredUserExportRow::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("email")?,
|
||||
row.try_get("email_verified")?,
|
||||
row.try_get("username")?,
|
||||
row.try_get("password_hash")?,
|
||||
row.try_get("role")?,
|
||||
row.try_get("auth_source")?,
|
||||
row.try_get("allowed_providers")?,
|
||||
row.try_get("allowed_api_formats")?,
|
||||
row.try_get("allowed_models")?,
|
||||
row.try_get("rate_limit")?,
|
||||
row.try_get("model_capability_settings")?,
|
||||
row.try_get("is_active")?,
|
||||
)
|
||||
}
|
||||
|
||||
fn map_user_auth_row(row: &sqlx::postgres::PgRow) -> Result<StoredUserAuthRecord, DataLayerError> {
|
||||
StoredUserAuthRecord::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("email")?,
|
||||
row.try_get("email_verified")?,
|
||||
row.try_get("username")?,
|
||||
row.try_get("password_hash")?,
|
||||
row.try_get("role")?,
|
||||
row.try_get("auth_source")?,
|
||||
row.try_get("allowed_providers")?,
|
||||
row.try_get("allowed_api_formats")?,
|
||||
row.try_get("allowed_models")?,
|
||||
row.try_get("is_active")?,
|
||||
row.try_get("is_deleted")?,
|
||||
row.try_get("created_at")?,
|
||||
row.try_get("last_login_at")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UserReadRepository for SqlxUserReadRepository {
|
||||
async fn list_users_by_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredUserSummary>, DataLayerError> {
|
||||
self.list_users_by_ids(user_ids).await
|
||||
}
|
||||
|
||||
async fn list_non_admin_export_users(
|
||||
&self,
|
||||
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
||||
self.list_non_admin_export_users().await
|
||||
}
|
||||
|
||||
async fn list_export_users(&self) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
||||
self.list_export_users().await
|
||||
}
|
||||
|
||||
async fn list_export_users_page(
|
||||
&self,
|
||||
query: &UserExportListQuery,
|
||||
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
||||
self.list_export_users_page(query).await
|
||||
}
|
||||
|
||||
async fn summarize_export_users(&self) -> Result<UserExportSummary, DataLayerError> {
|
||||
self.summarize_export_users().await
|
||||
}
|
||||
|
||||
async fn find_export_user_by_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Option<StoredUserExportRow>, DataLayerError> {
|
||||
self.find_export_user_by_id(user_id).await
|
||||
}
|
||||
|
||||
async fn find_user_auth_by_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Option<StoredUserAuthRecord>, DataLayerError> {
|
||||
self.find_user_auth_by_id(user_id).await
|
||||
}
|
||||
|
||||
async fn list_user_auth_by_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredUserAuthRecord>, DataLayerError> {
|
||||
self.list_user_auth_by_ids(user_ids).await
|
||||
}
|
||||
|
||||
async fn find_user_auth_by_identifier(
|
||||
&self,
|
||||
identifier: &str,
|
||||
) -> Result<Option<StoredUserAuthRecord>, DataLayerError> {
|
||||
self.find_user_auth_by_identifier(identifier).await
|
||||
}
|
||||
}
|
||||
450
crates/aether-data/src/repository/users/types.rs
Normal file
450
crates/aether-data/src/repository/users/types.rs
Normal file
@@ -0,0 +1,450 @@
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredUserSummary {
|
||||
pub id: String,
|
||||
pub username: String,
|
||||
pub email: Option<String>,
|
||||
pub role: String,
|
||||
pub is_active: bool,
|
||||
pub is_deleted: bool,
|
||||
}
|
||||
|
||||
impl StoredUserSummary {
|
||||
pub fn new(
|
||||
id: String,
|
||||
username: String,
|
||||
email: Option<String>,
|
||||
role: String,
|
||||
is_active: bool,
|
||||
is_deleted: bool,
|
||||
) -> 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,
|
||||
username,
|
||||
email,
|
||||
role,
|
||||
is_active,
|
||||
is_deleted,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredUserAuthRecord {
|
||||
pub id: String,
|
||||
pub email: Option<String>,
|
||||
pub email_verified: bool,
|
||||
pub username: String,
|
||||
pub password_hash: Option<String>,
|
||||
pub role: String,
|
||||
pub auth_source: String,
|
||||
pub allowed_providers: Option<Vec<String>>,
|
||||
pub allowed_api_formats: Option<Vec<String>>,
|
||||
pub allowed_models: Option<Vec<String>>,
|
||||
pub is_active: bool,
|
||||
pub is_deleted: bool,
|
||||
pub created_at: Option<DateTime<Utc>>,
|
||||
pub last_login_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
impl StoredUserAuthRecord {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
email: Option<String>,
|
||||
email_verified: bool,
|
||||
username: String,
|
||||
password_hash: Option<String>,
|
||||
role: String,
|
||||
auth_source: String,
|
||||
allowed_providers: Option<Value>,
|
||||
allowed_api_formats: Option<Value>,
|
||||
allowed_models: Option<Value>,
|
||||
is_active: bool,
|
||||
is_deleted: bool,
|
||||
created_at: Option<DateTime<Utc>>,
|
||||
last_login_at: Option<DateTime<Utc>>,
|
||||
) -> 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(),
|
||||
));
|
||||
}
|
||||
if auth_source.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"users.auth_source is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
email,
|
||||
email_verified,
|
||||
username,
|
||||
password_hash,
|
||||
role,
|
||||
auth_source,
|
||||
allowed_providers: parse_string_list(allowed_providers, "users.allowed_providers")?,
|
||||
allowed_api_formats: parse_string_list(
|
||||
allowed_api_formats,
|
||||
"users.allowed_api_formats",
|
||||
)?,
|
||||
allowed_models: parse_string_list(allowed_models, "users.allowed_models")?,
|
||||
is_active,
|
||||
is_deleted,
|
||||
created_at,
|
||||
last_login_at,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn to_summary(&self) -> Result<StoredUserSummary, crate::DataLayerError> {
|
||||
StoredUserSummary::new(
|
||||
self.id.clone(),
|
||||
self.username.clone(),
|
||||
self.email.clone(),
|
||||
self.role.clone(),
|
||||
self.is_active,
|
||||
self.is_deleted,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredUserExportRow {
|
||||
pub id: String,
|
||||
pub email: Option<String>,
|
||||
pub email_verified: bool,
|
||||
pub username: String,
|
||||
pub password_hash: Option<String>,
|
||||
pub role: String,
|
||||
pub auth_source: String,
|
||||
pub allowed_providers: Option<Vec<String>>,
|
||||
pub allowed_api_formats: Option<Vec<String>>,
|
||||
pub allowed_models: Option<Vec<String>>,
|
||||
pub rate_limit: Option<i32>,
|
||||
pub model_capability_settings: Option<Value>,
|
||||
pub is_active: bool,
|
||||
}
|
||||
|
||||
impl StoredUserExportRow {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
email: Option<String>,
|
||||
email_verified: bool,
|
||||
username: String,
|
||||
password_hash: Option<String>,
|
||||
role: String,
|
||||
auth_source: String,
|
||||
allowed_providers: Option<Value>,
|
||||
allowed_api_formats: Option<Value>,
|
||||
allowed_models: Option<Value>,
|
||||
rate_limit: Option<i32>,
|
||||
model_capability_settings: Option<Value>,
|
||||
is_active: bool,
|
||||
) -> 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(),
|
||||
));
|
||||
}
|
||||
if auth_source.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"users.auth_source is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
email,
|
||||
email_verified,
|
||||
username,
|
||||
password_hash,
|
||||
role,
|
||||
auth_source,
|
||||
allowed_providers: parse_string_list(allowed_providers, "users.allowed_providers")?,
|
||||
allowed_api_formats: parse_string_list(
|
||||
allowed_api_formats,
|
||||
"users.allowed_api_formats",
|
||||
)?,
|
||||
allowed_models: parse_string_list(allowed_models, "users.allowed_models")?,
|
||||
rate_limit,
|
||||
model_capability_settings: normalize_optional_json(model_capability_settings),
|
||||
is_active,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct UserExportListQuery {
|
||||
pub skip: usize,
|
||||
pub limit: usize,
|
||||
pub role: Option<String>,
|
||||
pub is_active: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct UserExportSummary {
|
||||
pub total: u64,
|
||||
pub active: u64,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait UserReadRepository: Send + Sync {
|
||||
async fn list_users_by_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredUserSummary>, crate::DataLayerError>;
|
||||
|
||||
async fn list_export_users(&self) -> Result<Vec<StoredUserExportRow>, crate::DataLayerError>;
|
||||
|
||||
async fn list_export_users_page(
|
||||
&self,
|
||||
query: &UserExportListQuery,
|
||||
) -> Result<Vec<StoredUserExportRow>, crate::DataLayerError>;
|
||||
|
||||
async fn summarize_export_users(&self) -> Result<UserExportSummary, crate::DataLayerError>;
|
||||
|
||||
async fn find_export_user_by_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Option<StoredUserExportRow>, crate::DataLayerError>;
|
||||
|
||||
async fn list_non_admin_export_users(
|
||||
&self,
|
||||
) -> Result<Vec<StoredUserExportRow>, crate::DataLayerError>;
|
||||
|
||||
async fn find_user_auth_by_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Option<StoredUserAuthRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn list_user_auth_by_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredUserAuthRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn find_user_auth_by_identifier(
|
||||
&self,
|
||||
identifier: &str,
|
||||
) -> Result<Option<StoredUserAuthRecord>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
fn normalize_optional_json(value: Option<Value>) -> Option<Value> {
|
||||
match value {
|
||||
Some(Value::Null) | None => None,
|
||||
Some(value) => Some(value),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_string_list(
|
||||
value: Option<Value>,
|
||||
field_name: &str,
|
||||
) -> Result<Option<Vec<String>>, crate::DataLayerError> {
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
parse_string_list_value(&value, field_name)
|
||||
}
|
||||
|
||||
fn parse_string_list_value(
|
||||
value: &Value,
|
||||
field_name: &str,
|
||||
) -> Result<Option<Vec<String>>, crate::DataLayerError> {
|
||||
match value {
|
||||
Value::Null => Ok(None),
|
||||
Value::Array(array) => parse_string_list_array(array, field_name).map(Some),
|
||||
Value::String(raw) => parse_embedded_string_list(raw, field_name),
|
||||
_ => Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"{field_name} is not a JSON array"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_embedded_string_list(
|
||||
raw: &str,
|
||||
field_name: &str,
|
||||
) -> Result<Option<Vec<String>>, crate::DataLayerError> {
|
||||
let raw = raw.trim();
|
||||
if raw.is_empty() || raw.eq_ignore_ascii_case("null") {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if let Ok(decoded) = serde_json::from_str::<Value>(raw) {
|
||||
return parse_string_list_value(&decoded, field_name);
|
||||
}
|
||||
|
||||
Ok(Some(vec![raw.to_string()]))
|
||||
}
|
||||
|
||||
fn parse_string_list_array(
|
||||
array: &[Value],
|
||||
field_name: &str,
|
||||
) -> Result<Vec<String>, crate::DataLayerError> {
|
||||
let mut items = Vec::with_capacity(array.len());
|
||||
for item in array {
|
||||
let Some(item) = item.as_str() else {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"{field_name} contains a non-string item"
|
||||
)));
|
||||
};
|
||||
let item = item.trim();
|
||||
if !item.is_empty() {
|
||||
items.push(item.to_string());
|
||||
}
|
||||
}
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{StoredUserAuthRecord, StoredUserExportRow};
|
||||
|
||||
#[test]
|
||||
fn builds_user_export_row_with_allowed_lists() {
|
||||
let row = StoredUserExportRow::new(
|
||||
"user-1".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
true,
|
||||
"alice".to_string(),
|
||||
Some("hash".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
Some(serde_json::json!(["openai", "anthropic"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
Some(60),
|
||||
Some(serde_json::json!({"gpt-4.1": {"cache_1h": true}})),
|
||||
true,
|
||||
)
|
||||
.expect("row should build");
|
||||
|
||||
assert_eq!(
|
||||
row.allowed_providers,
|
||||
Some(vec!["openai".to_string(), "anthropic".to_string()])
|
||||
);
|
||||
assert_eq!(
|
||||
row.allowed_api_formats,
|
||||
Some(vec!["openai:chat".to_string()])
|
||||
);
|
||||
assert_eq!(row.allowed_models, Some(vec!["gpt-4.1".to_string()]));
|
||||
assert_eq!(
|
||||
row.model_capability_settings,
|
||||
Some(serde_json::json!({"gpt-4.1": {"cache_1h": true}}))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_embedded_string_lists_for_user_export_row() {
|
||||
let row = StoredUserExportRow::new(
|
||||
"user-1".to_string(),
|
||||
None,
|
||||
false,
|
||||
"alice".to_string(),
|
||||
None,
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
Some(serde_json::json!("[\"openai\"]")),
|
||||
Some(serde_json::json!("null")),
|
||||
Some(serde_json::json!("gpt-4.1")),
|
||||
None,
|
||||
Some(Value::Null),
|
||||
true,
|
||||
)
|
||||
.expect("row should build");
|
||||
|
||||
assert_eq!(row.allowed_providers, Some(vec!["openai".to_string()]));
|
||||
assert_eq!(row.allowed_api_formats, None);
|
||||
assert_eq!(row.allowed_models, Some(vec!["gpt-4.1".to_string()]));
|
||||
assert_eq!(row.model_capability_settings, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_object_allowed_providers_for_user_export_row() {
|
||||
let result = StoredUserExportRow::new(
|
||||
"user-1".to_string(),
|
||||
None,
|
||||
false,
|
||||
"alice".to_string(),
|
||||
None,
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
Some(serde_json::json!({"bad": true})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
);
|
||||
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_user_auth_record_with_allowed_lists() {
|
||||
let row = StoredUserAuthRecord::new(
|
||||
"user-1".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
true,
|
||||
"alice".to_string(),
|
||||
Some("hash".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("auth row should build");
|
||||
|
||||
assert_eq!(row.allowed_providers, Some(vec!["openai".to_string()]));
|
||||
assert_eq!(
|
||||
row.allowed_api_formats,
|
||||
Some(vec!["openai:chat".to_string()])
|
||||
);
|
||||
assert_eq!(row.allowed_models, Some(vec!["gpt-4.1".to_string()]));
|
||||
}
|
||||
}
|
||||
@@ -4,7 +4,8 @@ use std::sync::RwLock;
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
|
||||
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount,
|
||||
VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskStatus, VideoTaskStatusCount,
|
||||
VideoTaskWriteRepository,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
@@ -47,6 +48,34 @@ impl InMemoryVideoTaskRepository {
|
||||
|
||||
task
|
||||
}
|
||||
|
||||
fn matches_filter(task: &StoredVideoTask, filter: &VideoTaskQueryFilter) -> bool {
|
||||
if let Some(user_id) = filter.user_id.as_deref() {
|
||||
if task.user_id.as_deref() != Some(user_id) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(status) = filter.status {
|
||||
if task.status != status {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(model_substring) = filter.model_substring.as_deref() {
|
||||
let needle = model_substring.trim().to_ascii_lowercase();
|
||||
let Some(model) = task.model.as_deref() else {
|
||||
return false;
|
||||
};
|
||||
if !model.to_ascii_lowercase().contains(&needle) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(client_api_format) = filter.client_api_format.as_deref() {
|
||||
if task.client_api_format.as_deref() != Some(client_api_format) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -92,6 +121,178 @@ impl VideoTaskReadRepository for InMemoryVideoTaskRepository {
|
||||
tasks.truncate(limit);
|
||||
Ok(tasks)
|
||||
}
|
||||
|
||||
async fn list_due(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut tasks = self
|
||||
.index
|
||||
.read()
|
||||
.expect("video task repository lock")
|
||||
.by_id
|
||||
.values()
|
||||
.filter(|task| {
|
||||
matches!(
|
||||
task.status,
|
||||
super::types::VideoTaskStatus::Submitted
|
||||
| super::types::VideoTaskStatus::Queued
|
||||
| super::types::VideoTaskStatus::Processing
|
||||
) && task.poll_count < task.max_poll_count
|
||||
&& task
|
||||
.next_poll_at_unix_secs
|
||||
.is_some_and(|value| value <= now_unix_secs)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
tasks.sort_by(|left, right| {
|
||||
left.next_poll_at_unix_secs
|
||||
.cmp(&right.next_poll_at_unix_secs)
|
||||
.then_with(|| left.updated_at_unix_secs.cmp(&right.updated_at_unix_secs))
|
||||
});
|
||||
tasks.truncate(limit);
|
||||
Ok(tasks)
|
||||
}
|
||||
|
||||
async fn list_page(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut tasks = self
|
||||
.index
|
||||
.read()
|
||||
.expect("video task repository lock")
|
||||
.by_id
|
||||
.values()
|
||||
.filter(|task| Self::matches_filter(task, filter))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
tasks.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.cmp(&left.created_at_unix_secs)
|
||||
.then_with(|| right.updated_at_unix_secs.cmp(&left.updated_at_unix_secs))
|
||||
});
|
||||
Ok(tasks.into_iter().skip(offset).take(limit).collect())
|
||||
}
|
||||
|
||||
async fn count(&self, filter: &VideoTaskQueryFilter) -> Result<u64, DataLayerError> {
|
||||
Ok(self
|
||||
.index
|
||||
.read()
|
||||
.expect("video task repository lock")
|
||||
.by_id
|
||||
.values()
|
||||
.filter(|task| Self::matches_filter(task, filter))
|
||||
.count() as u64)
|
||||
}
|
||||
|
||||
async fn count_by_status(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
) -> Result<Vec<VideoTaskStatusCount>, DataLayerError> {
|
||||
let mut counts = BTreeMap::<VideoTaskStatus, u64>::new();
|
||||
for task in self
|
||||
.index
|
||||
.read()
|
||||
.expect("video task repository lock")
|
||||
.by_id
|
||||
.values()
|
||||
.filter(|task| Self::matches_filter(task, filter))
|
||||
{
|
||||
*counts.entry(task.status).or_default() += 1;
|
||||
}
|
||||
Ok(counts
|
||||
.into_iter()
|
||||
.map(|(status, count)| VideoTaskStatusCount { status, count })
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn count_distinct_users(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
let index = self.index.read().expect("video task repository lock");
|
||||
let users = index
|
||||
.by_id
|
||||
.values()
|
||||
.filter(|task| Self::matches_filter(task, filter))
|
||||
.filter_map(|task| task.user_id.as_deref())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
Ok(users.len() as u64)
|
||||
}
|
||||
|
||||
async fn top_models(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
limit: usize,
|
||||
) -> Result<Vec<VideoTaskModelCount>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut counts = BTreeMap::<String, u64>::new();
|
||||
for task in self
|
||||
.index
|
||||
.read()
|
||||
.expect("video task repository lock")
|
||||
.by_id
|
||||
.values()
|
||||
.filter(|task| Self::matches_filter(task, filter))
|
||||
{
|
||||
let Some(model) = task.model.as_deref() else {
|
||||
continue;
|
||||
};
|
||||
if model.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
*counts.entry(model.to_string()).or_default() += 1;
|
||||
}
|
||||
|
||||
let mut models = counts
|
||||
.into_iter()
|
||||
.map(|(model, count)| VideoTaskModelCount { model, count })
|
||||
.collect::<Vec<_>>();
|
||||
models.sort_by(|left, right| {
|
||||
right
|
||||
.count
|
||||
.cmp(&left.count)
|
||||
.then_with(|| left.model.cmp(&right.model))
|
||||
});
|
||||
models.truncate(limit);
|
||||
Ok(models)
|
||||
}
|
||||
|
||||
async fn count_created_since(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
created_since_unix_secs: u64,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
Ok(self
|
||||
.index
|
||||
.read()
|
||||
.expect("video task repository lock")
|
||||
.by_id
|
||||
.values()
|
||||
.filter(|task| {
|
||||
Self::matches_filter(task, filter)
|
||||
&& task.created_at_unix_secs >= created_since_unix_secs
|
||||
})
|
||||
.count() as u64)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -100,14 +301,76 @@ impl VideoTaskWriteRepository for InMemoryVideoTaskRepository {
|
||||
let mut index = self.index.write().expect("video task repository lock");
|
||||
Ok(Self::store_locked(&mut index, task.into_stored()))
|
||||
}
|
||||
|
||||
async fn update_if_active(
|
||||
&self,
|
||||
task: UpsertVideoTask,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
let mut index = self.index.write().expect("video task repository lock");
|
||||
let Some(existing) = index.by_id.get(&task.id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !existing.status.is_active() {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(Self::store_locked(&mut index, task.into_stored())))
|
||||
}
|
||||
|
||||
async fn claim_due(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
claim_until_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut index = self.index.write().expect("video task repository lock");
|
||||
let mut due_ids = index
|
||||
.by_id
|
||||
.values()
|
||||
.filter(|task| {
|
||||
matches!(
|
||||
task.status,
|
||||
VideoTaskStatus::Submitted
|
||||
| VideoTaskStatus::Queued
|
||||
| VideoTaskStatus::Processing
|
||||
) && task.poll_count < task.max_poll_count
|
||||
&& task
|
||||
.next_poll_at_unix_secs
|
||||
.is_some_and(|value| value <= now_unix_secs)
|
||||
})
|
||||
.map(|task| task.id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
due_ids.sort_by(|left_id, right_id| {
|
||||
let left = index.by_id.get(left_id).expect("task should exist");
|
||||
let right = index.by_id.get(right_id).expect("task should exist");
|
||||
left.next_poll_at_unix_secs
|
||||
.cmp(&right.next_poll_at_unix_secs)
|
||||
.then_with(|| left.updated_at_unix_secs.cmp(&right.updated_at_unix_secs))
|
||||
});
|
||||
due_ids.truncate(limit);
|
||||
|
||||
let mut claimed = Vec::with_capacity(due_ids.len());
|
||||
for id in due_ids {
|
||||
let Some(task) = index.by_id.get_mut(&id) else {
|
||||
continue;
|
||||
};
|
||||
task.next_poll_at_unix_secs = Some(claim_until_unix_secs);
|
||||
task.updated_at_unix_secs = now_unix_secs.max(task.updated_at_unix_secs);
|
||||
claimed.push(task.clone());
|
||||
}
|
||||
Ok(claimed)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryVideoTaskRepository;
|
||||
use crate::repository::video_tasks::{
|
||||
UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository, VideoTaskStatus,
|
||||
VideoTaskWriteRepository,
|
||||
UpsertVideoTask, VideoTaskLookupKey, VideoTaskQueryFilter, VideoTaskReadRepository,
|
||||
VideoTaskStatus, VideoTaskWriteRepository,
|
||||
};
|
||||
|
||||
fn sample_task(
|
||||
@@ -118,19 +381,41 @@ mod tests {
|
||||
UpsertVideoTask {
|
||||
id: id.to_string(),
|
||||
short_id: Some(format!("short-{id}")),
|
||||
request_id: format!("request-{id}"),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("api-key-1".to_string()),
|
||||
username: Some("user".to_string()),
|
||||
api_key_name: Some("primary".to_string()),
|
||||
external_task_id: Some(format!("ext-{id}")),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
endpoint_id: Some("endpoint-1".to_string()),
|
||||
key_id: Some("provider-key-1".to_string()),
|
||||
client_api_format: Some("openai:video".to_string()),
|
||||
provider_api_format: Some("openai:video".to_string()),
|
||||
format_converted: false,
|
||||
model: Some("sora-2".to_string()),
|
||||
prompt: Some("hello".to_string()),
|
||||
original_request_body: Some(serde_json::json!({"prompt": "hello"})),
|
||||
duration_seconds: Some(4),
|
||||
resolution: Some("720p".to_string()),
|
||||
aspect_ratio: Some("16:9".to_string()),
|
||||
size: Some("1280x720".to_string()),
|
||||
status,
|
||||
progress_percent: 0,
|
||||
progress_message: None,
|
||||
retry_count: 0,
|
||||
poll_interval_seconds: 10,
|
||||
next_poll_at_unix_secs: Some(updated_at_unix_secs),
|
||||
poll_count: 0,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: updated_at_unix_secs.saturating_sub(10),
|
||||
submitted_at_unix_secs: Some(updated_at_unix_secs.saturating_sub(10)),
|
||||
completed_at_unix_secs: None,
|
||||
updated_at_unix_secs,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
request_metadata: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -193,19 +478,41 @@ mod tests {
|
||||
repo.upsert(UpsertVideoTask {
|
||||
id: "task-1".to_string(),
|
||||
short_id: Some("short-task-1b".to_string()),
|
||||
request_id: "request-task-1b".to_string(),
|
||||
user_id: Some("user-2".to_string()),
|
||||
api_key_id: Some("api-key-2".to_string()),
|
||||
username: Some("user-2".to_string()),
|
||||
api_key_name: Some("secondary".to_string()),
|
||||
external_task_id: Some("ext-task-1b".to_string()),
|
||||
provider_id: Some("provider-2".to_string()),
|
||||
endpoint_id: Some("endpoint-2".to_string()),
|
||||
key_id: Some("provider-key-2".to_string()),
|
||||
client_api_format: Some("gemini:video".to_string()),
|
||||
provider_api_format: Some("gemini:video".to_string()),
|
||||
format_converted: false,
|
||||
model: Some("veo-3".to_string()),
|
||||
prompt: Some("remix".to_string()),
|
||||
original_request_body: Some(serde_json::json!({"prompt": "remix"})),
|
||||
duration_seconds: Some(8),
|
||||
resolution: Some("1080p".to_string()),
|
||||
aspect_ratio: Some("16:9".to_string()),
|
||||
size: Some("720p".to_string()),
|
||||
status: VideoTaskStatus::Processing,
|
||||
progress_percent: 50,
|
||||
progress_message: Some("processing".to_string()),
|
||||
retry_count: 1,
|
||||
poll_interval_seconds: 10,
|
||||
next_poll_at_unix_secs: Some(200),
|
||||
poll_count: 2,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: 150,
|
||||
submitted_at_unix_secs: Some(150),
|
||||
completed_at_unix_secs: None,
|
||||
updated_at_unix_secs: 200,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
request_metadata: None,
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
@@ -229,4 +536,141 @@ mod tests {
|
||||
.expect("find should succeed")
|
||||
.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_due_returns_due_active_tasks_in_next_poll_order() {
|
||||
let repo = InMemoryVideoTaskRepository::default();
|
||||
repo.upsert(sample_task("task-1", VideoTaskStatus::Submitted, 300))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
repo.upsert(sample_task("task-2", VideoTaskStatus::Processing, 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
repo.upsert(UpsertVideoTask {
|
||||
next_poll_at_unix_secs: Some(500),
|
||||
..sample_task("task-3", VideoTaskStatus::Queued, 200)
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let due = repo
|
||||
.list_due(300, 10)
|
||||
.await
|
||||
.expect("list due should succeed");
|
||||
assert_eq!(due.len(), 2);
|
||||
assert_eq!(due[0].id, "task-2");
|
||||
assert_eq!(due[1].id, "task-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn update_if_active_skips_terminal_tasks() {
|
||||
let repo = InMemoryVideoTaskRepository::default();
|
||||
repo.upsert(sample_task("task-1", VideoTaskStatus::Completed, 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let updated = repo
|
||||
.update_if_active(UpsertVideoTask {
|
||||
progress_percent: 100,
|
||||
..sample_task("task-1", VideoTaskStatus::Completed, 200)
|
||||
})
|
||||
.await
|
||||
.expect("update should succeed");
|
||||
|
||||
assert!(updated.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_page_and_stats_apply_filters() {
|
||||
let repo = InMemoryVideoTaskRepository::default();
|
||||
repo.upsert(sample_task("task-1", VideoTaskStatus::Submitted, 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
repo.upsert(UpsertVideoTask {
|
||||
model: Some("veo-3-fast".to_string()),
|
||||
user_id: Some("user-2".to_string()),
|
||||
client_api_format: Some("gemini:video".to_string()),
|
||||
created_at_unix_secs: 250,
|
||||
updated_at_unix_secs: 250,
|
||||
..sample_task("task-2", VideoTaskStatus::Completed, 250)
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
repo.upsert(UpsertVideoTask {
|
||||
model: Some("veo-3-fast".to_string()),
|
||||
user_id: Some("user-2".to_string()),
|
||||
client_api_format: Some("gemini:video".to_string()),
|
||||
created_at_unix_secs: 260,
|
||||
updated_at_unix_secs: 260,
|
||||
..sample_task("task-3", VideoTaskStatus::Completed, 260)
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let filter = VideoTaskQueryFilter {
|
||||
user_id: Some("user-2".to_string()),
|
||||
status: Some(VideoTaskStatus::Completed),
|
||||
model_substring: Some("veo".to_string()),
|
||||
client_api_format: Some("gemini:video".to_string()),
|
||||
};
|
||||
|
||||
let page = repo
|
||||
.list_page(&filter, 0, 10)
|
||||
.await
|
||||
.expect("list page should succeed");
|
||||
assert_eq!(page.len(), 2);
|
||||
assert_eq!(page[0].id, "task-3");
|
||||
assert_eq!(page[1].id, "task-2");
|
||||
|
||||
let count = repo.count(&filter).await.expect("count should succeed");
|
||||
assert_eq!(count, 2);
|
||||
|
||||
let by_status = repo
|
||||
.count_by_status(&filter)
|
||||
.await
|
||||
.expect("status count should succeed");
|
||||
assert_eq!(by_status.len(), 1);
|
||||
assert_eq!(by_status[0].status, VideoTaskStatus::Completed);
|
||||
assert_eq!(by_status[0].count, 2);
|
||||
|
||||
let top_models = repo
|
||||
.top_models(&filter, 10)
|
||||
.await
|
||||
.expect("top models should succeed");
|
||||
assert_eq!(top_models.len(), 1);
|
||||
assert_eq!(top_models[0].model, "veo-3-fast");
|
||||
assert_eq!(top_models[0].count, 2);
|
||||
|
||||
let today_count = repo
|
||||
.count_created_since(&filter, 255)
|
||||
.await
|
||||
.expect("today count should succeed");
|
||||
assert_eq!(today_count, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn claim_due_advances_claimed_tasks_until_claim_deadline() {
|
||||
let repo = InMemoryVideoTaskRepository::default();
|
||||
repo.upsert(sample_task("task-1", VideoTaskStatus::Submitted, 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
repo.upsert(sample_task("task-2", VideoTaskStatus::Processing, 90))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let claimed = repo
|
||||
.claim_due(100, 130, 1)
|
||||
.await
|
||||
.expect("claim should succeed");
|
||||
assert_eq!(claimed.len(), 1);
|
||||
assert_eq!(claimed[0].id, "task-2");
|
||||
assert_eq!(claimed[0].next_poll_at_unix_secs, Some(130));
|
||||
|
||||
let remaining = repo
|
||||
.list_due(100, 10)
|
||||
.await
|
||||
.expect("list due should succeed");
|
||||
assert_eq!(remaining.len(), 1);
|
||||
assert_eq!(remaining[0].id, "task-1");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,8 +3,9 @@ mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryVideoTaskRepository;
|
||||
pub use sql::SqlxVideoTaskReadRepository;
|
||||
pub use sql::{SqlxVideoTaskReadRepository, SqlxVideoTaskRepository};
|
||||
pub use types::{
|
||||
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
|
||||
VideoTaskRepository, VideoTaskStatus, VideoTaskWriteRepository,
|
||||
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount,
|
||||
VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskRepository, VideoTaskStatus,
|
||||
VideoTaskStatusCount, VideoTaskWriteRepository,
|
||||
};
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,9 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[derive(
|
||||
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
|
||||
)]
|
||||
pub enum VideoTaskStatus {
|
||||
Pending,
|
||||
Submitted,
|
||||
@@ -43,71 +46,166 @@ impl VideoTaskStatus {
|
||||
pub struct StoredVideoTask {
|
||||
pub id: String,
|
||||
pub short_id: Option<String>,
|
||||
pub request_id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub api_key_id: Option<String>,
|
||||
pub username: Option<String>,
|
||||
pub api_key_name: Option<String>,
|
||||
pub external_task_id: Option<String>,
|
||||
pub provider_id: Option<String>,
|
||||
pub endpoint_id: Option<String>,
|
||||
pub key_id: Option<String>,
|
||||
pub client_api_format: Option<String>,
|
||||
pub provider_api_format: Option<String>,
|
||||
pub format_converted: bool,
|
||||
pub model: Option<String>,
|
||||
pub prompt: Option<String>,
|
||||
pub original_request_body: Option<Value>,
|
||||
pub duration_seconds: Option<u32>,
|
||||
pub resolution: Option<String>,
|
||||
pub aspect_ratio: Option<String>,
|
||||
pub size: Option<String>,
|
||||
pub status: VideoTaskStatus,
|
||||
pub progress_percent: u16,
|
||||
pub progress_message: Option<String>,
|
||||
pub retry_count: u32,
|
||||
pub poll_interval_seconds: u32,
|
||||
pub next_poll_at_unix_secs: Option<u64>,
|
||||
pub poll_count: u32,
|
||||
pub max_poll_count: u32,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub submitted_at_unix_secs: Option<u64>,
|
||||
pub completed_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: u64,
|
||||
pub error_code: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
pub video_url: Option<String>,
|
||||
pub request_metadata: Option<Value>,
|
||||
}
|
||||
|
||||
impl StoredVideoTask {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
short_id: Option<String>,
|
||||
request_id: String,
|
||||
user_id: Option<String>,
|
||||
api_key_id: Option<String>,
|
||||
username: Option<String>,
|
||||
api_key_name: Option<String>,
|
||||
external_task_id: Option<String>,
|
||||
provider_id: Option<String>,
|
||||
endpoint_id: Option<String>,
|
||||
key_id: Option<String>,
|
||||
client_api_format: Option<String>,
|
||||
provider_api_format: Option<String>,
|
||||
format_converted: bool,
|
||||
model: Option<String>,
|
||||
prompt: Option<String>,
|
||||
original_request_body: Option<Value>,
|
||||
duration_seconds: Option<i32>,
|
||||
resolution: Option<String>,
|
||||
aspect_ratio: Option<String>,
|
||||
size: Option<String>,
|
||||
status: VideoTaskStatus,
|
||||
progress_percent: i32,
|
||||
progress_message: Option<String>,
|
||||
retry_count: i32,
|
||||
poll_interval_seconds: i32,
|
||||
next_poll_at_unix_secs: Option<i64>,
|
||||
poll_count: i32,
|
||||
max_poll_count: i32,
|
||||
created_at_unix_secs: i64,
|
||||
submitted_at_unix_secs: Option<i64>,
|
||||
completed_at_unix_secs: Option<i64>,
|
||||
updated_at_unix_secs: i64,
|
||||
error_code: Option<String>,
|
||||
error_message: Option<String>,
|
||||
video_url: Option<String>,
|
||||
request_metadata: Option<Value>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
let progress_percent = u16::try_from(progress_percent).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid progress_percent: {progress_percent}"
|
||||
))
|
||||
})?;
|
||||
let retry_count = u32::try_from(retry_count).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid retry_count: {retry_count}"))
|
||||
})?;
|
||||
let poll_interval_seconds = u32::try_from(poll_interval_seconds).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid poll_interval_seconds: {poll_interval_seconds}"
|
||||
))
|
||||
})?;
|
||||
let next_poll_at_unix_secs =
|
||||
coerce_optional_unix_secs(next_poll_at_unix_secs, "next_poll_at_unix_secs")?;
|
||||
let poll_count = u32::try_from(poll_count).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid poll_count: {poll_count}"))
|
||||
})?;
|
||||
let max_poll_count = u32::try_from(max_poll_count).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid max_poll_count: {max_poll_count}"
|
||||
))
|
||||
})?;
|
||||
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid created_at_unix_secs: {created_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
let submitted_at_unix_secs =
|
||||
coerce_optional_unix_secs(submitted_at_unix_secs, "submitted_at_unix_secs")?;
|
||||
let completed_at_unix_secs =
|
||||
coerce_optional_unix_secs(completed_at_unix_secs, "completed_at_unix_secs")?;
|
||||
let updated_at_unix_secs = u64::try_from(updated_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid updated_at_unix_secs: {updated_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
let duration_seconds = match duration_seconds {
|
||||
Some(value) => Some(u32::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid duration_seconds: {value}"))
|
||||
})?),
|
||||
None => None,
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
short_id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
external_task_id,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
format_converted,
|
||||
model,
|
||||
prompt,
|
||||
original_request_body,
|
||||
duration_seconds,
|
||||
resolution,
|
||||
aspect_ratio,
|
||||
size,
|
||||
status,
|
||||
progress_percent,
|
||||
progress_message,
|
||||
retry_count,
|
||||
poll_interval_seconds,
|
||||
next_poll_at_unix_secs,
|
||||
poll_count,
|
||||
max_poll_count,
|
||||
created_at_unix_secs,
|
||||
submitted_at_unix_secs,
|
||||
completed_at_unix_secs,
|
||||
updated_at_unix_secs,
|
||||
error_code,
|
||||
error_message,
|
||||
video_url,
|
||||
request_metadata,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -116,19 +214,41 @@ impl StoredVideoTask {
|
||||
pub struct UpsertVideoTask {
|
||||
pub id: String,
|
||||
pub short_id: Option<String>,
|
||||
pub request_id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub api_key_id: Option<String>,
|
||||
pub username: Option<String>,
|
||||
pub api_key_name: Option<String>,
|
||||
pub external_task_id: Option<String>,
|
||||
pub provider_id: Option<String>,
|
||||
pub endpoint_id: Option<String>,
|
||||
pub key_id: Option<String>,
|
||||
pub client_api_format: Option<String>,
|
||||
pub provider_api_format: Option<String>,
|
||||
pub format_converted: bool,
|
||||
pub model: Option<String>,
|
||||
pub prompt: Option<String>,
|
||||
pub original_request_body: Option<Value>,
|
||||
pub duration_seconds: Option<u32>,
|
||||
pub resolution: Option<String>,
|
||||
pub aspect_ratio: Option<String>,
|
||||
pub size: Option<String>,
|
||||
pub status: VideoTaskStatus,
|
||||
pub progress_percent: u16,
|
||||
pub progress_message: Option<String>,
|
||||
pub retry_count: u32,
|
||||
pub poll_interval_seconds: u32,
|
||||
pub next_poll_at_unix_secs: Option<u64>,
|
||||
pub poll_count: u32,
|
||||
pub max_poll_count: u32,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub submitted_at_unix_secs: Option<u64>,
|
||||
pub completed_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: u64,
|
||||
pub error_code: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
pub video_url: Option<String>,
|
||||
pub request_metadata: Option<Value>,
|
||||
}
|
||||
|
||||
impl UpsertVideoTask {
|
||||
@@ -136,19 +256,85 @@ impl UpsertVideoTask {
|
||||
StoredVideoTask {
|
||||
id: self.id,
|
||||
short_id: self.short_id,
|
||||
request_id: self.request_id,
|
||||
user_id: self.user_id,
|
||||
api_key_id: self.api_key_id,
|
||||
username: self.username,
|
||||
api_key_name: self.api_key_name,
|
||||
external_task_id: self.external_task_id,
|
||||
provider_id: self.provider_id,
|
||||
endpoint_id: self.endpoint_id,
|
||||
key_id: self.key_id,
|
||||
client_api_format: self.client_api_format,
|
||||
provider_api_format: self.provider_api_format,
|
||||
format_converted: self.format_converted,
|
||||
model: self.model,
|
||||
prompt: self.prompt,
|
||||
original_request_body: self.original_request_body,
|
||||
duration_seconds: self.duration_seconds,
|
||||
resolution: self.resolution,
|
||||
aspect_ratio: self.aspect_ratio,
|
||||
size: self.size,
|
||||
status: self.status,
|
||||
progress_percent: self.progress_percent,
|
||||
progress_message: self.progress_message,
|
||||
retry_count: self.retry_count,
|
||||
poll_interval_seconds: self.poll_interval_seconds,
|
||||
next_poll_at_unix_secs: self.next_poll_at_unix_secs,
|
||||
poll_count: self.poll_count,
|
||||
max_poll_count: self.max_poll_count,
|
||||
created_at_unix_secs: self.created_at_unix_secs,
|
||||
submitted_at_unix_secs: self.submitted_at_unix_secs,
|
||||
completed_at_unix_secs: self.completed_at_unix_secs,
|
||||
updated_at_unix_secs: self.updated_at_unix_secs,
|
||||
error_code: self.error_code,
|
||||
error_message: self.error_message,
|
||||
video_url: self.video_url,
|
||||
request_metadata: self.request_metadata,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<StoredVideoTask> for UpsertVideoTask {
|
||||
fn from(task: StoredVideoTask) -> Self {
|
||||
Self {
|
||||
id: task.id,
|
||||
short_id: task.short_id,
|
||||
request_id: task.request_id,
|
||||
user_id: task.user_id,
|
||||
api_key_id: task.api_key_id,
|
||||
username: task.username,
|
||||
api_key_name: task.api_key_name,
|
||||
external_task_id: task.external_task_id,
|
||||
provider_id: task.provider_id,
|
||||
endpoint_id: task.endpoint_id,
|
||||
key_id: task.key_id,
|
||||
client_api_format: task.client_api_format,
|
||||
provider_api_format: task.provider_api_format,
|
||||
format_converted: task.format_converted,
|
||||
model: task.model,
|
||||
prompt: task.prompt,
|
||||
original_request_body: task.original_request_body,
|
||||
duration_seconds: task.duration_seconds,
|
||||
resolution: task.resolution,
|
||||
aspect_ratio: task.aspect_ratio,
|
||||
size: task.size,
|
||||
status: task.status,
|
||||
progress_percent: task.progress_percent,
|
||||
progress_message: task.progress_message,
|
||||
retry_count: task.retry_count,
|
||||
poll_interval_seconds: task.poll_interval_seconds,
|
||||
next_poll_at_unix_secs: task.next_poll_at_unix_secs,
|
||||
poll_count: task.poll_count,
|
||||
max_poll_count: task.max_poll_count,
|
||||
created_at_unix_secs: task.created_at_unix_secs,
|
||||
submitted_at_unix_secs: task.submitted_at_unix_secs,
|
||||
completed_at_unix_secs: task.completed_at_unix_secs,
|
||||
updated_at_unix_secs: task.updated_at_unix_secs,
|
||||
error_code: task.error_code,
|
||||
error_message: task.error_message,
|
||||
video_url: task.video_url,
|
||||
request_metadata: task.request_metadata,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -163,6 +349,26 @@ pub enum VideoTaskLookupKey<'a> {
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct VideoTaskQueryFilter {
|
||||
pub user_id: Option<String>,
|
||||
pub status: Option<VideoTaskStatus>,
|
||||
pub model_substring: Option<String>,
|
||||
pub client_api_format: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct VideoTaskStatusCount {
|
||||
pub status: VideoTaskStatus,
|
||||
pub count: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct VideoTaskModelCount {
|
||||
pub model: String,
|
||||
pub count: u64,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait VideoTaskReadRepository: Send + Sync {
|
||||
async fn find(
|
||||
@@ -174,12 +380,61 @@ pub trait VideoTaskReadRepository: Send + Sync {
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
|
||||
|
||||
async fn list_due(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
|
||||
|
||||
async fn list_page(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
|
||||
|
||||
async fn count(&self, filter: &VideoTaskQueryFilter) -> Result<u64, crate::DataLayerError>;
|
||||
|
||||
async fn count_by_status(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
) -> Result<Vec<VideoTaskStatusCount>, crate::DataLayerError>;
|
||||
|
||||
async fn count_distinct_users(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
) -> Result<u64, crate::DataLayerError>;
|
||||
|
||||
async fn top_models(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
limit: usize,
|
||||
) -> Result<Vec<VideoTaskModelCount>, crate::DataLayerError>;
|
||||
|
||||
async fn count_created_since(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
created_since_unix_secs: u64,
|
||||
) -> Result<u64, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait VideoTaskWriteRepository: Send + Sync {
|
||||
async fn upsert(&self, task: UpsertVideoTask)
|
||||
-> Result<StoredVideoTask, crate::DataLayerError>;
|
||||
|
||||
async fn update_if_active(
|
||||
&self,
|
||||
task: UpsertVideoTask,
|
||||
) -> Result<Option<StoredVideoTask>, crate::DataLayerError>;
|
||||
|
||||
async fn claim_due(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
claim_until_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait VideoTaskRepository:
|
||||
@@ -192,10 +447,103 @@ impl<T> VideoTaskRepository for T where
|
||||
{
|
||||
}
|
||||
|
||||
fn coerce_optional_unix_secs(
|
||||
value: Option<i64>,
|
||||
field: &str,
|
||||
) -> Result<Option<u64>, crate::DataLayerError> {
|
||||
match value {
|
||||
Some(value) => Ok(Some(u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid {field}: {value}"))
|
||||
})?)),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{StoredVideoTask, VideoTaskStatus};
|
||||
|
||||
#[allow(clippy::type_complexity)]
|
||||
fn base_new_args() -> (
|
||||
String,
|
||||
Option<String>,
|
||||
String,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
bool,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<serde_json::Value>,
|
||||
Option<i32>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
VideoTaskStatus,
|
||||
i32,
|
||||
Option<String>,
|
||||
i32,
|
||||
i32,
|
||||
Option<i64>,
|
||||
i32,
|
||||
i32,
|
||||
i64,
|
||||
Option<i64>,
|
||||
Option<i64>,
|
||||
i64,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<serde_json::Value>,
|
||||
) {
|
||||
(
|
||||
"task-1".to_string(),
|
||||
None,
|
||||
"request-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
VideoTaskStatus::Submitted,
|
||||
10,
|
||||
None,
|
||||
0,
|
||||
10,
|
||||
Some(1),
|
||||
0,
|
||||
360,
|
||||
1,
|
||||
None,
|
||||
None,
|
||||
1,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_status_from_database_text() {
|
||||
assert_eq!(
|
||||
@@ -211,66 +559,52 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_numeric_fields() {
|
||||
let mut args = base_new_args();
|
||||
args.22 = -1;
|
||||
assert!(StoredVideoTask::new(
|
||||
"task-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
VideoTaskStatus::Submitted,
|
||||
-1,
|
||||
1,
|
||||
1,
|
||||
None,
|
||||
None,
|
||||
None
|
||||
args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7, args.8, args.9,
|
||||
args.10, args.11, args.12, args.13, args.14, args.15, args.16, args.17, args.18,
|
||||
args.19, args.20, args.21, args.22, args.23, args.24, args.25, args.26, args.27,
|
||||
args.28, args.29, args.30, args.31, args.32, args.33, args.34, args.35, args.36,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_updated_at_values() {
|
||||
let mut args = base_new_args();
|
||||
args.32 = -1;
|
||||
assert!(StoredVideoTask::new(
|
||||
"task-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
VideoTaskStatus::Submitted,
|
||||
10,
|
||||
1,
|
||||
-1,
|
||||
None,
|
||||
None,
|
||||
None
|
||||
args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7, args.8, args.9,
|
||||
args.10, args.11, args.12, args.13, args.14, args.15, args.16, args.17, args.18,
|
||||
args.19, args.20, args.21, args.22, args.23, args.24, args.25, args.26, args.27,
|
||||
args.28, args.29, args.30, args.31, args.32, args.33, args.34, args.35, args.36,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_created_at_values() {
|
||||
let mut args = base_new_args();
|
||||
args.29 = -1;
|
||||
assert!(StoredVideoTask::new(
|
||||
"task-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
VideoTaskStatus::Submitted,
|
||||
10,
|
||||
-1,
|
||||
1,
|
||||
None,
|
||||
None,
|
||||
None
|
||||
args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7, args.8, args.9,
|
||||
args.10, args.11, args.12, args.13, args.14, args.15, args.16, args.17, args.18,
|
||||
args.19, args.20, args.21, args.22, args.23, args.24, args.25, args.26, args.27,
|
||||
args.28, args.29, args.30, args.31, args.32, args.33, args.34, args.35, args.36,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_optional_completed_at_values() {
|
||||
let mut args = base_new_args();
|
||||
args.31 = Some(-1);
|
||||
assert!(StoredVideoTask::new(
|
||||
args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7, args.8, args.9,
|
||||
args.10, args.11, args.12, args.13, args.14, args.15, args.16, args.17, args.18,
|
||||
args.19, args.20, args.21, args.22, args.23, args.24, args.25, args.26, args.27,
|
||||
args.28, args.29, args.30, args.31, args.32, args.33, args.34, args.35, args.36,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
270
crates/aether-data/src/repository/wallet/memory.rs
Normal file
270
crates/aether-data/src/repository/wallet/memory.rs
Normal file
@@ -0,0 +1,270 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
StoredUsageSettlement, StoredWalletSnapshot, UsageSettlementInput, WalletLookupKey,
|
||||
WalletReadRepository, WalletWriteRepository,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryWalletRepository {
|
||||
wallets_by_id: RwLock<BTreeMap<String, StoredWalletSnapshot>>,
|
||||
provider_monthly_used: RwLock<BTreeMap<String, f64>>,
|
||||
}
|
||||
|
||||
impl InMemoryWalletRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredWalletSnapshot>,
|
||||
{
|
||||
let mut wallets_by_id = BTreeMap::new();
|
||||
for item in items {
|
||||
wallets_by_id.insert(item.id.clone(), item);
|
||||
}
|
||||
Self {
|
||||
wallets_by_id: RwLock::new(wallets_by_id),
|
||||
provider_monthly_used: RwLock::new(BTreeMap::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WalletReadRepository for InMemoryWalletRepository {
|
||||
async fn find(
|
||||
&self,
|
||||
key: WalletLookupKey<'_>,
|
||||
) -> Result<Option<StoredWalletSnapshot>, DataLayerError> {
|
||||
let wallets = self.wallets_by_id.read().expect("wallet repo lock");
|
||||
Ok(match key {
|
||||
WalletLookupKey::WalletId(wallet_id) => wallets.get(wallet_id).cloned(),
|
||||
WalletLookupKey::UserId(user_id) => wallets
|
||||
.values()
|
||||
.find(|wallet| wallet.user_id.as_deref() == Some(user_id))
|
||||
.cloned(),
|
||||
WalletLookupKey::ApiKeyId(api_key_id) => wallets
|
||||
.values()
|
||||
.find(|wallet| wallet.api_key_id.as_deref() == Some(api_key_id))
|
||||
.cloned(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_wallets_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredWalletSnapshot>, DataLayerError> {
|
||||
if user_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let user_set: std::collections::BTreeSet<&str> =
|
||||
user_ids.iter().map(String::as_str).collect();
|
||||
let wallets = self.wallets_by_id.read().expect("wallet repo lock");
|
||||
Ok(wallets
|
||||
.values()
|
||||
.filter(|wallet| {
|
||||
wallet
|
||||
.user_id
|
||||
.as_deref()
|
||||
.map(|user_id| user_set.contains(user_id))
|
||||
.unwrap_or(false)
|
||||
})
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_wallets_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredWalletSnapshot>, DataLayerError> {
|
||||
if api_key_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let key_set: std::collections::BTreeSet<&str> =
|
||||
api_key_ids.iter().map(String::as_str).collect();
|
||||
let wallets = self.wallets_by_id.read().expect("wallet repo lock");
|
||||
Ok(wallets
|
||||
.values()
|
||||
.filter(|wallet| {
|
||||
wallet
|
||||
.api_key_id
|
||||
.as_deref()
|
||||
.map(|api_key_id| key_set.contains(api_key_id))
|
||||
.unwrap_or(false)
|
||||
})
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WalletWriteRepository for InMemoryWalletRepository {
|
||||
async fn settle_usage(
|
||||
&self,
|
||||
input: UsageSettlementInput,
|
||||
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
|
||||
input.validate()?;
|
||||
if input.billing_status != "pending" {
|
||||
return Ok(Some(StoredUsageSettlement {
|
||||
request_id: input.request_id,
|
||||
wallet_id: None,
|
||||
billing_status: input.billing_status,
|
||||
wallet_balance_before: None,
|
||||
wallet_balance_after: None,
|
||||
wallet_recharge_balance_before: None,
|
||||
wallet_recharge_balance_after: None,
|
||||
wallet_gift_balance_before: None,
|
||||
wallet_gift_balance_after: None,
|
||||
provider_monthly_used_usd: None,
|
||||
finalized_at_unix_secs: input.finalized_at_unix_secs,
|
||||
}));
|
||||
}
|
||||
|
||||
let mut wallets = self.wallets_by_id.write().expect("wallet repo lock");
|
||||
let wallet_id = input
|
||||
.api_key_id
|
||||
.as_deref()
|
||||
.and_then(|api_key_id| {
|
||||
wallets
|
||||
.values()
|
||||
.find(|wallet| wallet.api_key_id.as_deref() == Some(api_key_id))
|
||||
.map(|wallet| wallet.id.clone())
|
||||
})
|
||||
.or_else(|| {
|
||||
input.user_id.as_deref().and_then(|user_id| {
|
||||
wallets
|
||||
.values()
|
||||
.find(|wallet| wallet.user_id.as_deref() == Some(user_id))
|
||||
.map(|wallet| wallet.id.clone())
|
||||
})
|
||||
});
|
||||
let wallet = wallet_id
|
||||
.as_deref()
|
||||
.and_then(|wallet_id| wallets.get_mut(wallet_id));
|
||||
|
||||
let final_billing_status = if input.status == "completed" {
|
||||
"settled"
|
||||
} else {
|
||||
"void"
|
||||
};
|
||||
|
||||
let mut settlement = StoredUsageSettlement {
|
||||
request_id: input.request_id,
|
||||
wallet_id: None,
|
||||
billing_status: final_billing_status.to_string(),
|
||||
wallet_balance_before: None,
|
||||
wallet_balance_after: None,
|
||||
wallet_recharge_balance_before: None,
|
||||
wallet_recharge_balance_after: None,
|
||||
wallet_gift_balance_before: None,
|
||||
wallet_gift_balance_after: None,
|
||||
provider_monthly_used_usd: None,
|
||||
finalized_at_unix_secs: input.finalized_at_unix_secs,
|
||||
};
|
||||
|
||||
if let Some(wallet) = wallet {
|
||||
let before_recharge = wallet.balance;
|
||||
let before_gift = wallet.gift_balance;
|
||||
let before_total = before_recharge + before_gift;
|
||||
settlement.wallet_id = Some(wallet.id.clone());
|
||||
settlement.wallet_balance_before = Some(before_total);
|
||||
settlement.wallet_recharge_balance_before = Some(before_recharge);
|
||||
settlement.wallet_gift_balance_before = Some(before_gift);
|
||||
|
||||
if final_billing_status == "settled" {
|
||||
if wallet.limit_mode.eq_ignore_ascii_case("unlimited") {
|
||||
wallet.total_consumed += input.total_cost_usd;
|
||||
} else {
|
||||
let gift_deduction = before_gift.max(0.0).min(input.total_cost_usd);
|
||||
let recharge_deduction = input.total_cost_usd - gift_deduction;
|
||||
wallet.gift_balance = before_gift - gift_deduction;
|
||||
wallet.balance = before_recharge - recharge_deduction;
|
||||
wallet.total_consumed += input.total_cost_usd;
|
||||
}
|
||||
}
|
||||
|
||||
settlement.wallet_recharge_balance_after = Some(wallet.balance);
|
||||
settlement.wallet_gift_balance_after = Some(wallet.gift_balance);
|
||||
settlement.wallet_balance_after = Some(wallet.balance + wallet.gift_balance);
|
||||
}
|
||||
|
||||
if final_billing_status == "settled" {
|
||||
if let Some(provider_id) = input.provider_id {
|
||||
let mut quotas = self
|
||||
.provider_monthly_used
|
||||
.write()
|
||||
.expect("provider quota lock");
|
||||
let value = quotas.entry(provider_id).or_insert(0.0);
|
||||
*value += input.actual_total_cost_usd;
|
||||
settlement.provider_monthly_used_usd = Some(*value);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Some(settlement))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryWalletRepository;
|
||||
use crate::repository::wallet::{
|
||||
StoredWalletSnapshot, UsageSettlementInput, WalletLookupKey, WalletReadRepository,
|
||||
WalletWriteRepository,
|
||||
};
|
||||
|
||||
fn sample_wallet() -> StoredWalletSnapshot {
|
||||
StoredWalletSnapshot::new(
|
||||
"wallet-1".to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("key-1".to_string()),
|
||||
10.0,
|
||||
2.0,
|
||||
"finite".to_string(),
|
||||
"USD".to_string(),
|
||||
"active".to_string(),
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
100,
|
||||
)
|
||||
.expect("wallet should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn finds_wallet_by_owner() {
|
||||
let repository = InMemoryWalletRepository::seed(vec![sample_wallet()]);
|
||||
let wallet = repository
|
||||
.find(WalletLookupKey::UserId("user-1"))
|
||||
.await
|
||||
.expect("lookup should succeed")
|
||||
.expect("wallet should exist");
|
||||
assert_eq!(wallet.id, "wallet-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn settles_usage_against_wallet_and_provider_quota() {
|
||||
let repository = InMemoryWalletRepository::seed(vec![sample_wallet()]);
|
||||
let settlement = repository
|
||||
.settle_usage(UsageSettlementInput {
|
||||
request_id: "req-1".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("key-1".to_string()),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
status: "completed".to_string(),
|
||||
billing_status: "pending".to_string(),
|
||||
total_cost_usd: 3.0,
|
||||
actual_total_cost_usd: 1.5,
|
||||
finalized_at_unix_secs: Some(200),
|
||||
})
|
||||
.await
|
||||
.expect("settlement should succeed")
|
||||
.expect("settlement should exist");
|
||||
|
||||
assert_eq!(settlement.billing_status, "settled");
|
||||
assert_eq!(settlement.wallet_balance_before, Some(12.0));
|
||||
assert_eq!(settlement.wallet_balance_after, Some(9.0));
|
||||
assert_eq!(settlement.provider_monthly_used_usd, Some(1.5));
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/wallet/mod.rs
Normal file
10
crates/aether-data/src/repository/wallet/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryWalletRepository;
|
||||
pub use sql::SqlxWalletRepository;
|
||||
pub use types::{
|
||||
StoredUsageSettlement, StoredWalletSnapshot, UsageSettlementInput, WalletLookupKey,
|
||||
WalletReadRepository, WalletRepository, WalletWriteRepository,
|
||||
};
|
||||
485
crates/aether-data/src/repository/wallet/sql.rs
Normal file
485
crates/aether-data/src/repository/wallet/sql.rs
Normal file
@@ -0,0 +1,485 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
StoredUsageSettlement, StoredWalletSnapshot, UsageSettlementInput, WalletLookupKey,
|
||||
WalletReadRepository, WalletWriteRepository,
|
||||
};
|
||||
use crate::postgres::PostgresTransactionRunner;
|
||||
use crate::DataLayerError;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
const FIND_BY_WALLET_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
CAST(balance AS DOUBLE PRECISION) AS balance,
|
||||
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
CAST(total_recharged AS DOUBLE PRECISION) AS total_recharged,
|
||||
CAST(total_consumed AS DOUBLE PRECISION) AS total_consumed,
|
||||
CAST(total_refunded AS DOUBLE PRECISION) AS total_refunded,
|
||||
CAST(total_adjusted AS DOUBLE PRECISION) AS total_adjusted,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
FROM wallets
|
||||
WHERE id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const FIND_BY_USER_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
CAST(balance AS DOUBLE PRECISION) AS balance,
|
||||
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
CAST(total_recharged AS DOUBLE PRECISION) AS total_recharged,
|
||||
CAST(total_consumed AS DOUBLE PRECISION) AS total_consumed,
|
||||
CAST(total_refunded AS DOUBLE PRECISION) AS total_refunded,
|
||||
CAST(total_adjusted AS DOUBLE PRECISION) AS total_adjusted,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
FROM wallets
|
||||
WHERE user_id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const FIND_BY_API_KEY_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
CAST(balance AS DOUBLE PRECISION) AS balance,
|
||||
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
CAST(total_recharged AS DOUBLE PRECISION) AS total_recharged,
|
||||
CAST(total_consumed AS DOUBLE PRECISION) AS total_consumed,
|
||||
CAST(total_refunded AS DOUBLE PRECISION) AS total_refunded,
|
||||
CAST(total_adjusted AS DOUBLE PRECISION) AS total_adjusted,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
FROM wallets
|
||||
WHERE api_key_id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const LIST_BY_USER_IDS_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
CAST(balance AS DOUBLE PRECISION) AS balance,
|
||||
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
CAST(total_recharged AS DOUBLE PRECISION) AS total_recharged,
|
||||
CAST(total_consumed AS DOUBLE PRECISION) AS total_consumed,
|
||||
CAST(total_refunded AS DOUBLE PRECISION) AS total_refunded,
|
||||
CAST(total_adjusted AS DOUBLE PRECISION) AS total_adjusted,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
FROM wallets
|
||||
WHERE user_id = ANY($1)
|
||||
"#;
|
||||
|
||||
const LIST_BY_API_KEY_IDS_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
CAST(balance AS DOUBLE PRECISION) AS balance,
|
||||
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
CAST(total_recharged AS DOUBLE PRECISION) AS total_recharged,
|
||||
CAST(total_consumed AS DOUBLE PRECISION) AS total_consumed,
|
||||
CAST(total_refunded AS DOUBLE PRECISION) AS total_refunded,
|
||||
CAST(total_adjusted AS DOUBLE PRECISION) AS total_adjusted,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
FROM wallets
|
||||
WHERE api_key_id = ANY($1)
|
||||
"#;
|
||||
|
||||
const FINALIZE_USAGE_BILLING_SQL: &str = r#"
|
||||
UPDATE "usage"
|
||||
SET
|
||||
billing_status = $2,
|
||||
finalized_at = TO_TIMESTAMP($3::double precision)
|
||||
WHERE request_id = $1
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxWalletRepository {
|
||||
pool: PgPool,
|
||||
tx_runner: PostgresTransactionRunner,
|
||||
}
|
||||
|
||||
impl SqlxWalletRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
let tx_runner = PostgresTransactionRunner::new(pool.clone());
|
||||
Self { pool, tx_runner }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WalletReadRepository for SqlxWalletRepository {
|
||||
async fn find(
|
||||
&self,
|
||||
key: WalletLookupKey<'_>,
|
||||
) -> Result<Option<StoredWalletSnapshot>, DataLayerError> {
|
||||
let query = match key {
|
||||
WalletLookupKey::WalletId(_) => FIND_BY_WALLET_ID_SQL,
|
||||
WalletLookupKey::UserId(_) => FIND_BY_USER_ID_SQL,
|
||||
WalletLookupKey::ApiKeyId(_) => FIND_BY_API_KEY_ID_SQL,
|
||||
};
|
||||
let bind = match key {
|
||||
WalletLookupKey::WalletId(value)
|
||||
| WalletLookupKey::UserId(value)
|
||||
| WalletLookupKey::ApiKeyId(value) => value,
|
||||
};
|
||||
let row = sqlx::query(query)
|
||||
.bind(bind)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_wallet_row).transpose()
|
||||
}
|
||||
async fn list_wallets_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredWalletSnapshot>, DataLayerError> {
|
||||
if user_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut ids_map = BTreeMap::new();
|
||||
for (index, id) in user_ids.iter().enumerate() {
|
||||
ids_map.entry(id).or_insert_with(Vec::new).push(index);
|
||||
}
|
||||
let rows = sqlx::query(LIST_BY_USER_IDS_SQL)
|
||||
.bind(user_ids)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
let mut wallets = Vec::with_capacity(rows.len());
|
||||
for row in rows {
|
||||
let wallet = map_wallet_row(&row)?;
|
||||
wallets.push(wallet);
|
||||
}
|
||||
Ok(wallets)
|
||||
}
|
||||
async fn list_wallets_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredWalletSnapshot>, DataLayerError> {
|
||||
if api_key_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let rows = sqlx::query(LIST_BY_API_KEY_IDS_SQL)
|
||||
.bind(api_key_ids)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
let mut wallets = Vec::with_capacity(rows.len());
|
||||
for row in rows {
|
||||
let wallet = map_wallet_row(&row)?;
|
||||
wallets.push(wallet);
|
||||
}
|
||||
Ok(wallets)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WalletWriteRepository for SqlxWalletRepository {
|
||||
async fn settle_usage(
|
||||
&self,
|
||||
input: UsageSettlementInput,
|
||||
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
|
||||
input.validate()?;
|
||||
self.tx_runner
|
||||
.run_read_write(|tx| {
|
||||
Box::pin(async move {
|
||||
let row = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
request_id,
|
||||
wallet_id,
|
||||
billing_status,
|
||||
CAST(wallet_balance_before AS DOUBLE PRECISION) AS wallet_balance_before,
|
||||
CAST(wallet_balance_after AS DOUBLE PRECISION) AS wallet_balance_after,
|
||||
CAST(wallet_recharge_balance_before AS DOUBLE PRECISION) AS wallet_recharge_balance_before,
|
||||
CAST(wallet_recharge_balance_after AS DOUBLE PRECISION) AS wallet_recharge_balance_after,
|
||||
CAST(wallet_gift_balance_before AS DOUBLE PRECISION) AS wallet_gift_balance_before,
|
||||
CAST(wallet_gift_balance_after AS DOUBLE PRECISION) AS wallet_gift_balance_after,
|
||||
provider_id,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
FROM "usage"
|
||||
WHERE request_id = $1
|
||||
FOR UPDATE
|
||||
"#,
|
||||
)
|
||||
.bind(&input.request_id)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await?;
|
||||
|
||||
let Some(usage_row) = row else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let current_billing_status: String = usage_row.try_get("billing_status")?;
|
||||
if current_billing_status == "settled" || current_billing_status == "void" {
|
||||
return Ok(Some(StoredUsageSettlement {
|
||||
request_id: usage_row.try_get("request_id")?,
|
||||
wallet_id: usage_row.try_get("wallet_id")?,
|
||||
billing_status: current_billing_status,
|
||||
wallet_balance_before: usage_row.try_get("wallet_balance_before")?,
|
||||
wallet_balance_after: usage_row.try_get("wallet_balance_after")?,
|
||||
wallet_recharge_balance_before: usage_row
|
||||
.try_get("wallet_recharge_balance_before")?,
|
||||
wallet_recharge_balance_after: usage_row
|
||||
.try_get("wallet_recharge_balance_after")?,
|
||||
wallet_gift_balance_before: usage_row
|
||||
.try_get("wallet_gift_balance_before")?,
|
||||
wallet_gift_balance_after: usage_row
|
||||
.try_get("wallet_gift_balance_after")?,
|
||||
provider_monthly_used_usd: None,
|
||||
finalized_at_unix_secs: usage_row
|
||||
.try_get::<Option<i64>, _>("finalized_at_unix_secs")?
|
||||
.map(|value| value as u64),
|
||||
}));
|
||||
}
|
||||
|
||||
let final_billing_status = if input.status == "completed" {
|
||||
"settled"
|
||||
} else {
|
||||
"void"
|
||||
};
|
||||
let finalized_at =
|
||||
i64::try_from(input.finalized_at_unix_secs.unwrap_or_else(|| {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}))
|
||||
.map_err(|_| {
|
||||
DataLayerError::InvalidInput("finalized_at overflow".to_string())
|
||||
})?;
|
||||
|
||||
let mut settlement = StoredUsageSettlement {
|
||||
request_id: input.request_id.clone(),
|
||||
wallet_id: None,
|
||||
billing_status: final_billing_status.to_string(),
|
||||
wallet_balance_before: None,
|
||||
wallet_balance_after: None,
|
||||
wallet_recharge_balance_before: None,
|
||||
wallet_recharge_balance_after: None,
|
||||
wallet_gift_balance_before: None,
|
||||
wallet_gift_balance_after: None,
|
||||
provider_monthly_used_usd: None,
|
||||
finalized_at_unix_secs: Some(finalized_at as u64),
|
||||
};
|
||||
|
||||
if final_billing_status == "settled" {
|
||||
let wallet_row = if let Some(api_key_id) = input
|
||||
.api_key_id
|
||||
.as_deref()
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
CAST(balance AS DOUBLE PRECISION) AS balance,
|
||||
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
|
||||
limit_mode
|
||||
FROM wallets
|
||||
WHERE api_key_id = $1
|
||||
FOR UPDATE
|
||||
LIMIT 1
|
||||
"#,
|
||||
)
|
||||
.bind(api_key_id)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let wallet_row = if wallet_row.is_some() {
|
||||
wallet_row
|
||||
} else if let Some(user_id) =
|
||||
input.user_id.as_deref().filter(|value| !value.is_empty())
|
||||
{
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
CAST(balance AS DOUBLE PRECISION) AS balance,
|
||||
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
|
||||
limit_mode
|
||||
FROM wallets
|
||||
WHERE user_id = $1
|
||||
FOR UPDATE
|
||||
LIMIT 1
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
if let Some(wallet_row) = wallet_row {
|
||||
let wallet_id: String = wallet_row.try_get("id")?;
|
||||
let before_recharge: f64 = wallet_row.try_get("balance")?;
|
||||
let before_gift: f64 = wallet_row.try_get("gift_balance")?;
|
||||
let limit_mode: String = wallet_row.try_get("limit_mode")?;
|
||||
let before_total = before_recharge + before_gift;
|
||||
let mut after_recharge = before_recharge;
|
||||
let mut after_gift = before_gift;
|
||||
if !limit_mode.eq_ignore_ascii_case("unlimited") {
|
||||
let gift_deduction = before_gift.max(0.0).min(input.total_cost_usd);
|
||||
let recharge_deduction = input.total_cost_usd - gift_deduction;
|
||||
after_gift = before_gift - gift_deduction;
|
||||
after_recharge = before_recharge - recharge_deduction;
|
||||
}
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE wallets
|
||||
SET
|
||||
balance = $2,
|
||||
gift_balance = $3,
|
||||
total_consumed = CAST(total_consumed AS DOUBLE PRECISION) + $4,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(&wallet_id)
|
||||
.bind(after_recharge)
|
||||
.bind(after_gift)
|
||||
.bind(input.total_cost_usd)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
|
||||
settlement.wallet_id = Some(wallet_id.clone());
|
||||
settlement.wallet_balance_before = Some(before_total);
|
||||
settlement.wallet_balance_after = Some(after_recharge + after_gift);
|
||||
settlement.wallet_recharge_balance_before = Some(before_recharge);
|
||||
settlement.wallet_recharge_balance_after = Some(after_recharge);
|
||||
settlement.wallet_gift_balance_before = Some(before_gift);
|
||||
settlement.wallet_gift_balance_after = Some(after_gift);
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE "usage"
|
||||
SET
|
||||
wallet_id = $2,
|
||||
wallet_balance_before = $3,
|
||||
wallet_balance_after = $4,
|
||||
wallet_recharge_balance_before = $5,
|
||||
wallet_recharge_balance_after = $6,
|
||||
wallet_gift_balance_before = $7,
|
||||
wallet_gift_balance_after = $8
|
||||
WHERE request_id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(&input.request_id)
|
||||
.bind(&wallet_id)
|
||||
.bind(before_total)
|
||||
.bind(after_recharge + after_gift)
|
||||
.bind(before_recharge)
|
||||
.bind(after_recharge)
|
||||
.bind(before_gift)
|
||||
.bind(after_gift)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
}
|
||||
|
||||
if let Some(provider_id) = input
|
||||
.provider_id
|
||||
.as_deref()
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
let quota_row = sqlx::query(
|
||||
r#"
|
||||
UPDATE providers
|
||||
SET
|
||||
monthly_used_usd = COALESCE(monthly_used_usd, 0) + $2,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
RETURNING CAST(monthly_used_usd AS DOUBLE PRECISION) AS monthly_used_usd
|
||||
"#,
|
||||
)
|
||||
.bind(provider_id)
|
||||
.bind(input.actual_total_cost_usd)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await?;
|
||||
settlement.provider_monthly_used_usd =
|
||||
quota_row.and_then(|row| row.try_get("monthly_used_usd").ok());
|
||||
}
|
||||
}
|
||||
|
||||
sqlx::query(FINALIZE_USAGE_BILLING_SQL)
|
||||
.bind(&input.request_id)
|
||||
.bind(final_billing_status)
|
||||
.bind(finalized_at)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
|
||||
Ok(Some(settlement))
|
||||
})
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_wallet_row(row: &sqlx::postgres::PgRow) -> Result<StoredWalletSnapshot, DataLayerError> {
|
||||
StoredWalletSnapshot::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("user_id")?,
|
||||
row.try_get("api_key_id")?,
|
||||
row.try_get("balance")?,
|
||||
row.try_get("gift_balance")?,
|
||||
row.try_get("limit_mode")?,
|
||||
row.try_get("currency")?,
|
||||
row.try_get("status")?,
|
||||
row.try_get("total_recharged")?,
|
||||
row.try_get("total_consumed")?,
|
||||
row.try_get("total_refunded")?,
|
||||
row.try_get("total_adjusted")?,
|
||||
row.try_get("updated_at_unix_secs")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxWalletRepository;
|
||||
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 = SqlxWalletRepository::new(pool);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wallet_usage_finalize_sql_does_not_require_usage_updated_at_column() {
|
||||
assert!(!super::FINALIZE_USAGE_BILLING_SQL.contains("updated_at"));
|
||||
}
|
||||
}
|
||||
215
crates/aether-data/src/repository/wallet/types.rs
Normal file
215
crates/aether-data/src/repository/wallet/types.rs
Normal file
@@ -0,0 +1,215 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum WalletLookupKey<'a> {
|
||||
WalletId(&'a str),
|
||||
UserId(&'a str),
|
||||
ApiKeyId(&'a str),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredWalletSnapshot {
|
||||
pub id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub api_key_id: Option<String>,
|
||||
pub balance: f64,
|
||||
pub gift_balance: f64,
|
||||
pub limit_mode: String,
|
||||
pub currency: String,
|
||||
pub status: String,
|
||||
pub total_recharged: f64,
|
||||
pub total_consumed: f64,
|
||||
pub total_refunded: f64,
|
||||
pub total_adjusted: f64,
|
||||
pub updated_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
impl StoredWalletSnapshot {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
user_id: Option<String>,
|
||||
api_key_id: Option<String>,
|
||||
balance: f64,
|
||||
gift_balance: f64,
|
||||
limit_mode: String,
|
||||
currency: String,
|
||||
status: String,
|
||||
total_recharged: f64,
|
||||
total_consumed: f64,
|
||||
total_refunded: f64,
|
||||
total_adjusted: f64,
|
||||
updated_at_unix_secs: i64,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"wallet.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if limit_mode.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"wallet.limit_mode is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if currency.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"wallet.currency is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if status.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"wallet.status is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if !balance.is_finite()
|
||||
|| !gift_balance.is_finite()
|
||||
|| !total_recharged.is_finite()
|
||||
|| !total_consumed.is_finite()
|
||||
|| !total_refunded.is_finite()
|
||||
|| !total_adjusted.is_finite()
|
||||
{
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"wallet numeric value is not finite".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(Self {
|
||||
id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
balance,
|
||||
gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
total_recharged,
|
||||
total_consumed,
|
||||
total_refunded,
|
||||
total_adjusted,
|
||||
updated_at_unix_secs: u64::try_from(updated_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(
|
||||
"wallet.updated_at_unix_secs is negative".to_string(),
|
||||
)
|
||||
})?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct UsageSettlementInput {
|
||||
pub request_id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub api_key_id: Option<String>,
|
||||
pub provider_id: Option<String>,
|
||||
pub status: String,
|
||||
pub billing_status: String,
|
||||
pub total_cost_usd: f64,
|
||||
pub actual_total_cost_usd: f64,
|
||||
pub finalized_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl UsageSettlementInput {
|
||||
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||
if self.request_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"wallet settlement request_id cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.status.trim().is_empty() || self.billing_status.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"wallet settlement status cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if !self.total_cost_usd.is_finite() || !self.actual_total_cost_usd.is_finite() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"wallet settlement cost must be finite".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredUsageSettlement {
|
||||
pub request_id: String,
|
||||
pub wallet_id: Option<String>,
|
||||
pub billing_status: String,
|
||||
pub wallet_balance_before: Option<f64>,
|
||||
pub wallet_balance_after: Option<f64>,
|
||||
pub wallet_recharge_balance_before: Option<f64>,
|
||||
pub wallet_recharge_balance_after: Option<f64>,
|
||||
pub wallet_gift_balance_before: Option<f64>,
|
||||
pub wallet_gift_balance_after: Option<f64>,
|
||||
pub provider_monthly_used_usd: Option<f64>,
|
||||
pub finalized_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait WalletReadRepository: Send + Sync {
|
||||
async fn find(
|
||||
&self,
|
||||
key: WalletLookupKey<'_>,
|
||||
) -> Result<Option<StoredWalletSnapshot>, crate::DataLayerError>;
|
||||
|
||||
async fn list_wallets_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredWalletSnapshot>, crate::DataLayerError>;
|
||||
|
||||
async fn list_wallets_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredWalletSnapshot>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait WalletWriteRepository: Send + Sync {
|
||||
async fn settle_usage(
|
||||
&self,
|
||||
input: UsageSettlementInput,
|
||||
) -> Result<Option<StoredUsageSettlement>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait WalletRepository: WalletReadRepository + WalletWriteRepository + Send + Sync {}
|
||||
|
||||
impl<T> WalletRepository for T where T: WalletReadRepository + WalletWriteRepository + Send + Sync {}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{StoredWalletSnapshot, UsageSettlementInput};
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_wallet_snapshot() {
|
||||
assert!(StoredWalletSnapshot::new(
|
||||
"".to_string(),
|
||||
None,
|
||||
None,
|
||||
1.0,
|
||||
0.0,
|
||||
"finite".to_string(),
|
||||
"USD".to_string(),
|
||||
"active".to_string(),
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
1,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_settlement_input() {
|
||||
let input = UsageSettlementInput {
|
||||
request_id: "".to_string(),
|
||||
user_id: None,
|
||||
api_key_id: None,
|
||||
provider_id: None,
|
||||
status: "completed".to_string(),
|
||||
billing_status: "pending".to_string(),
|
||||
total_cost_usd: 0.1,
|
||||
actual_total_cost_usd: 0.1,
|
||||
finalized_at_unix_secs: None,
|
||||
};
|
||||
assert!(input.validate().is_err());
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,54 @@
|
||||
use std::error::Error as _;
|
||||
|
||||
use http::method::InvalidMethod;
|
||||
use thiserror::Error;
|
||||
|
||||
pub(crate) fn format_upstream_request_error(err: &reqwest::Error) -> String {
|
||||
let mut kinds = Vec::new();
|
||||
if err.is_connect() {
|
||||
kinds.push("connect");
|
||||
}
|
||||
if err.is_timeout() {
|
||||
kinds.push("timeout");
|
||||
}
|
||||
if err.is_redirect() {
|
||||
kinds.push("redirect");
|
||||
}
|
||||
if err.is_body() {
|
||||
kinds.push("body");
|
||||
}
|
||||
if err.is_decode() {
|
||||
kinds.push("decode");
|
||||
}
|
||||
if err.is_request() {
|
||||
kinds.push("request");
|
||||
}
|
||||
|
||||
let mut detail = err.to_string();
|
||||
let mut source = err.source();
|
||||
while let Some(cause) = source {
|
||||
let cause_text = cause.to_string();
|
||||
if !cause_text.is_empty() && !detail.contains(&cause_text) {
|
||||
detail.push_str(": ");
|
||||
detail.push_str(&cause_text);
|
||||
}
|
||||
source = cause.source();
|
||||
}
|
||||
|
||||
if let Some(url) = err.url() {
|
||||
detail.push_str(" [url=");
|
||||
detail.push_str(url.as_str());
|
||||
detail.push(']');
|
||||
}
|
||||
if !kinds.is_empty() {
|
||||
detail.push_str(" [kind=");
|
||||
detail.push_str(&kinds.join(","));
|
||||
detail.push(']');
|
||||
}
|
||||
|
||||
detail
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ExecutorClientError {
|
||||
#[error("executor endpoint is not configured")]
|
||||
@@ -46,7 +94,7 @@ pub enum ExecutorServiceError {
|
||||
#[error("executor overloaded: gate {gate} saturated at {limit}")]
|
||||
Overloaded { gate: &'static str, limit: usize },
|
||||
#[error("failed to execute upstream request: {0}")]
|
||||
UpstreamRequest(reqwest::Error),
|
||||
UpstreamRequest(String),
|
||||
#[error("hub relay request failed: {0}")]
|
||||
RelayError(String),
|
||||
#[error("upstream response is not valid JSON: {0}")]
|
||||
|
||||
@@ -14,7 +14,7 @@ use reqwest::tls::Version;
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::ExecutorServiceError;
|
||||
use crate::{error::format_upstream_request_error, ExecutorServiceError};
|
||||
|
||||
const HUB_RELAY_CONTENT_TYPE: &str = "application/vnd.aether.tunnel-envelope";
|
||||
const HUB_RELAY_ERROR_HEADER: &str = "x-aether-tunnel-error";
|
||||
@@ -57,10 +57,9 @@ impl SyncExecutor {
|
||||
let response = send_request(&plan, body_bytes).await?;
|
||||
let status_code = response.status().as_u16();
|
||||
let headers = collect_response_headers(response.headers());
|
||||
let body_bytes = response
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(ExecutorServiceError::UpstreamRequest)?;
|
||||
let body_bytes = response.bytes().await.map_err(|err| {
|
||||
ExecutorServiceError::UpstreamRequest(format_upstream_request_error(&err))
|
||||
})?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
let upstream_bytes = body_bytes.len() as u64;
|
||||
|
||||
@@ -160,7 +159,7 @@ async fn send_request(
|
||||
request
|
||||
.send()
|
||||
.await
|
||||
.map_err(ExecutorServiceError::UpstreamRequest)
|
||||
.map_err(|err| ExecutorServiceError::UpstreamRequest(format_upstream_request_error(&err)))
|
||||
}
|
||||
|
||||
async fn send_via_tunnel_relay(
|
||||
|
||||
@@ -7,24 +7,40 @@ repository.workspace = true
|
||||
description = "Rust ingress gateway for Aether phase 3a transparent proxy"
|
||||
|
||||
[dependencies]
|
||||
aether-billing.workspace = true
|
||||
aether-cache.workspace = true
|
||||
aether-contracts.workspace = true
|
||||
aether-crypto.workspace = true
|
||||
aether-data.workspace = true
|
||||
aether-http.workspace = true
|
||||
aether-runtime.workspace = true
|
||||
aether-wallet.workspace = true
|
||||
async-stream.workspace = true
|
||||
async-trait.workspace = true
|
||||
axum = { version = "0.8" }
|
||||
base64.workspace = true
|
||||
bcrypt.workspace = true
|
||||
bytes.workspace = true
|
||||
chrono.workspace = true
|
||||
chrono-tz.workspace = true
|
||||
clap = { version = "4", features = ["derive", "env"] }
|
||||
flate2.workspace = true
|
||||
futures-util.workspace = true
|
||||
hmac.workspace = true
|
||||
http.workspace = true
|
||||
ldap3 = "0.11"
|
||||
regex.workspace = true
|
||||
redis.workspace = true
|
||||
reqwest.workspace = true
|
||||
rustls.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
sqlx.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
tokio-util.workspace = true
|
||||
tracing.workspace = true
|
||||
url.workspace = true
|
||||
uuid.workspace = true
|
||||
webpki-roots.workspace = true
|
||||
|
||||
14
crates/aether-gateway/src/api/ai/claude.rs
Normal file
14
crates/aether-gateway/src/api/ai/claude.rs
Normal file
@@ -0,0 +1,14 @@
|
||||
pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> {
|
||||
match api_format {
|
||||
"claude:chat" => Some("claude:chat"),
|
||||
"claude:cli" => Some("claude:cli"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn local_path(api_format: &str) -> Option<&'static str> {
|
||||
match api_format {
|
||||
"claude" | "claude:chat" | "claude:cli" => Some("/v1/messages"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
16
crates/aether-gateway/src/api/ai/gemini.rs
Normal file
16
crates/aether-gateway/src/api/ai/gemini.rs
Normal file
@@ -0,0 +1,16 @@
|
||||
pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> {
|
||||
match api_format {
|
||||
"gemini:chat" => Some("gemini:chat"),
|
||||
"gemini:cli" => Some("gemini:cli"),
|
||||
"gemini:video" => Some("gemini:video"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn local_path(api_format: &str) -> Option<&'static str> {
|
||||
match api_format {
|
||||
"gemini" | "gemini:chat" | "gemini:cli" => Some("/v1beta/models/{model}:{action}"),
|
||||
"gemini:video" => Some("/v1beta/models/{model}:predictLongRunning"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
115
crates/aether-gateway/src/api/ai/mod.rs
Normal file
115
crates/aether-gateway/src/api/ai/mod.rs
Normal file
@@ -0,0 +1,115 @@
|
||||
mod claude;
|
||||
mod gemini;
|
||||
mod openai;
|
||||
|
||||
use axum::routing::{any, post};
|
||||
use axum::Router;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::gateway::{proxy_request, AppState};
|
||||
|
||||
pub(crate) fn mount_ai_routes(router: Router<AppState>) -> Router<AppState> {
|
||||
router
|
||||
.route("/v1/chat/completions", post(proxy_request))
|
||||
.route("/v1/messages", post(proxy_request))
|
||||
.route("/v1/messages/count_tokens", post(proxy_request))
|
||||
.route("/v1/responses", post(proxy_request))
|
||||
.route("/v1/responses/compact", post(proxy_request))
|
||||
.route("/v1/models/{*gemini_path}", any(proxy_request))
|
||||
.route("/v1beta/models/{*gemini_path}", any(proxy_request))
|
||||
.route("/v1beta/operations", any(proxy_request))
|
||||
.route("/v1beta/operations/{*operation_path}", any(proxy_request))
|
||||
.route("/v1/videos", any(proxy_request))
|
||||
.route("/v1/videos/{*video_path}", any(proxy_request))
|
||||
.route("/upload/v1beta/files", any(proxy_request))
|
||||
.route("/v1beta/files", any(proxy_request))
|
||||
.route("/v1beta/files/{*file_path}", any(proxy_request))
|
||||
}
|
||||
|
||||
pub(crate) fn public_api_format_local_path(api_format: &str) -> &'static str {
|
||||
let normalized = api_format.trim().to_ascii_lowercase();
|
||||
openai::local_path(&normalized)
|
||||
.or_else(|| claude::local_path(&normalized))
|
||||
.or_else(|| gemini::local_path(&normalized))
|
||||
.unwrap_or("/")
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_admin_endpoint_signature(api_format: &str) -> Option<&'static str> {
|
||||
let normalized = api_format.trim().to_ascii_lowercase();
|
||||
openai::normalized_signature(&normalized)
|
||||
.or_else(|| claude::normalized_signature(&normalized))
|
||||
.or_else(|| gemini::normalized_signature(&normalized))
|
||||
}
|
||||
|
||||
pub(crate) fn admin_endpoint_signature_parts(
|
||||
api_format: &str,
|
||||
) -> Option<(&'static str, &'static str, &'static str)> {
|
||||
let normalized = normalize_admin_endpoint_signature(api_format)?;
|
||||
let (api_family, endpoint_kind) = normalized.split_once(':')?;
|
||||
Some((normalized, api_family, endpoint_kind))
|
||||
}
|
||||
|
||||
pub(crate) fn provider_type_is_fixed(provider_type: &str) -> bool {
|
||||
matches!(
|
||||
provider_type.trim().to_ascii_lowercase().as_str(),
|
||||
"claude_code" | "kiro" | "codex" | "gemini_cli" | "antigravity" | "vertex_ai"
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn provider_type_enables_format_conversion_by_default(provider_type: &str) -> bool {
|
||||
matches!(
|
||||
provider_type.trim().to_ascii_lowercase().as_str(),
|
||||
"claude_code" | "kiro" | "codex" | "antigravity" | "vertex_ai"
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn fixed_provider_template(
|
||||
provider_type: &str,
|
||||
) -> Option<(&'static str, &'static [&'static str])> {
|
||||
match provider_type.trim().to_ascii_lowercase().as_str() {
|
||||
"claude_code" => Some(("https://api.anthropic.com", &["claude:cli"])),
|
||||
"codex" => Some((
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
&["openai:cli", "openai:compact"],
|
||||
)),
|
||||
"kiro" => Some(("https://q.{region}.amazonaws.com", &["claude:cli"])),
|
||||
"gemini_cli" => Some(("https://cloudcode-pa.googleapis.com", &["gemini:cli"])),
|
||||
"vertex_ai" => Some((
|
||||
"https://aiplatform.googleapis.com",
|
||||
&["gemini:chat", "claude:chat"],
|
||||
)),
|
||||
"antigravity" => Some(("https://cloudcode-pa.googleapis.com", &["gemini:chat"])),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn codex_default_body_rules() -> Vec<serde_json::Value> {
|
||||
vec![
|
||||
json!({"action": "drop", "path": "max_output_tokens"}),
|
||||
json!({"action": "drop", "path": "temperature"}),
|
||||
json!({"action": "drop", "path": "top_p"}),
|
||||
json!({"action": "set", "path": "store", "value": false}),
|
||||
json!({
|
||||
"action": "set",
|
||||
"path": "instructions",
|
||||
"value": "You are GPT-5.",
|
||||
"condition": {"path": "instructions", "op": "not_exists"},
|
||||
}),
|
||||
]
|
||||
}
|
||||
|
||||
pub(crate) fn admin_default_body_rules_for_signature(
|
||||
api_format: &str,
|
||||
provider_type: Option<&str>,
|
||||
) -> Option<(String, Vec<serde_json::Value>)> {
|
||||
let normalized_api_format = normalize_admin_endpoint_signature(api_format)?.to_string();
|
||||
let provider_type = provider_type.map(|value| value.trim().to_ascii_lowercase());
|
||||
let body_rules = if normalized_api_format == "openai:compact"
|
||||
|| (normalized_api_format == "openai:cli" && provider_type.as_deref() == Some("codex"))
|
||||
{
|
||||
codex_default_body_rules()
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
Some((normalized_api_format, body_rules))
|
||||
}
|
||||
19
crates/aether-gateway/src/api/ai/openai.rs
Normal file
19
crates/aether-gateway/src/api/ai/openai.rs
Normal file
@@ -0,0 +1,19 @@
|
||||
pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> {
|
||||
match api_format {
|
||||
"openai:chat" => Some("openai:chat"),
|
||||
"openai:cli" => Some("openai:cli"),
|
||||
"openai:compact" => Some("openai:compact"),
|
||||
"openai:video" => Some("openai:video"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn local_path(api_format: &str) -> Option<&'static str> {
|
||||
match api_format {
|
||||
"openai" | "openai:chat" => Some("/v1/chat/completions"),
|
||||
"openai:cli" => Some("/v1/responses"),
|
||||
"openai:compact" => Some("/v1/responses/compact"),
|
||||
"openai:video" => Some("/v1/videos"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user