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:
fawney19
2026-03-31 19:19:04 +08:00
parent b5a0070023
commit ddf18fed9a
690 changed files with 235087 additions and 16301 deletions

View 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

View 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()
}

View 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()]);
}
}

View 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};

View 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"))
);
}
}

View 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);
}
}

View 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,
}

View 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))
);
}
}

View 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);
}
}

View 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
);
}
}

View 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);
}
}

View File

@@ -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);

View 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

View 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,
};

View 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(&timestamp.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");
}
}

View File

@@ -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

View File

@@ -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());
}
}

View File

@@ -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();
}
}

View File

@@ -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());
}
}

View File

@@ -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");
}
}

View File

@@ -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());
}
}

View 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());
}
}

View File

@@ -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::{

View File

@@ -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,

View 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);
}
}

View 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,
};

View 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);
}
}

View 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>;
}

View File

@@ -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
);
}
}

View File

@@ -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

View File

@@ -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);
}
}

View 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"
);
}
}

View 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,
};

View 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);
}
}

View 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>;
}

View 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");
}
}

View File

@@ -0,0 +1,7 @@
mod memory;
mod sql;
mod types;
pub use memory::InMemoryBillingReadRepository;
pub use sql::SqlxBillingReadRepository;
pub use types::{BillingReadRepository, StoredBillingModelContext};

View 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);
}
}

View 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>;
}

View 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");
}
}

View 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,
};

View 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,
}
])
);
}
}

View 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"));
}
}

View File

@@ -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");
}
}

View File

@@ -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,
};

View File

@@ -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();
}
}

View File

@@ -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());
}
}

View 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,
}
}
}

View File

@@ -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,
};

View 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(")");
}
}

View 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
{
}

View 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");
}
}

View 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,
};

File diff suppressed because it is too large Load Diff

View 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>;
}

View 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);
}
}

View 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,
};

View 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);
}
}

View 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>;
}

View File

@@ -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;

View 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);
}
}

View 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,
};

View 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);
}
}

View 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
{
}

View File

@@ -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());
}
}

View File

@@ -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

View File

@@ -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}}))
);
}
}

View 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"));
}
}

View 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,
};

View 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
}
}

View 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>;
}

View 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(&quota.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);
}
}

View 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,
};

View 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);
}
}

View 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
{
}

View File

@@ -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);
}
}

View File

@@ -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,
};

View File

@@ -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\""));
}
}

View File

@@ -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());
}
}

View 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");
}
}

View 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,
};

View 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
}
}

View 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()]));
}
}

View File

@@ -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");
}
}

View File

@@ -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

View File

@@ -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());
}

View 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));
}
}

View 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,
};

View 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"));
}
}

View 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());
}
}

View File

@@ -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}")]

View File

@@ -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(

View File

@@ -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

View 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,
}
}

View 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,
}
}

View 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))
}

View 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