1use async_trait::async_trait;
2use chrono::{DateTime, Utc};
3use serde::{Deserialize, Serialize};
4
5use crate::error::DbError;
6use crate::types::{
7 CreateDeploymentRequest, Deployment, RoutingConfig, UpdateRoutingConfigRequest,
8};
9
10#[async_trait]
11pub trait Database: Clone + Send + Sync + 'static {
12 async fn get_team(&self, id: &str) -> Result<Option<Team>, DbError>;
13 async fn create_team(&self, name: &str, budget_cents: i64) -> Result<Team, DbError>;
14 async fn get_user(&self, id: &str) -> Result<Option<User>, DbError>;
15 async fn get_user_by_email(&self, email: &str) -> Result<Option<User>, DbError>;
16 async fn create_user(
17 &self,
18 team_id: &str,
19 email: &str,
20 role: &str,
21 password_hash: Option<String>,
22 ) -> Result<User, DbError>;
23 async fn update_password_hash(&self, user_id: &str, password_hash: &str)
24 -> Result<(), DbError>;
25 async fn count_users_by_role(&self, role: &str) -> Result<i64, DbError>;
26 async fn get_api_key(&self, id: &str) -> Result<Option<ApiKey>, DbError>;
27 async fn get_api_key_by_hash(&self, key_hash: &str) -> Result<Option<ApiKey>, DbError>;
28 async fn create_api_key(
29 &self,
30 key_hash: &str,
31 user_id: &str,
32 team_id: &str,
33 name: Option<String>,
34 ) -> Result<ApiKey, DbError>;
35 async fn deactivate_api_key(&self, id: &str) -> Result<ApiKey, DbError>;
36 async fn get_model_alias(&self, id: &str) -> Result<Option<ModelAlias>, DbError>;
37 async fn create_model_alias(
38 &self,
39 team_id: &str,
40 alias: &str,
41 target_model: &str,
42 provider: &str,
43 ) -> Result<ModelAlias, DbError>;
44 async fn get_quota(&self, team_id: &str) -> Result<Option<Quota>, DbError>;
45 async fn create_quota(
46 &self,
47 team_id: &str,
48 rpm_limit: i32,
49 tpm_limit: i32,
50 ) -> Result<Quota, DbError>;
51 async fn record_usage(
52 &self,
53 team_id: &str,
54 api_key_id: &str,
55 model: &str,
56 input_tokens: i32,
57 output_tokens: i32,
58 response_time_ms: i64,
59 ) -> Result<UsageLog, DbError>;
60
61 async fn list_deployments(
63 &self,
64 model: &str,
65 is_active: Option<bool>,
66 ) -> Result<Vec<Deployment>, DbError>;
67 async fn get_deployment(&self, id: &str) -> Result<Option<Deployment>, DbError>;
68 async fn create_deployment(&self, req: CreateDeploymentRequest) -> Result<Deployment, DbError>;
69 async fn update_deployment(
70 &self,
71 id: &str,
72 req: CreateDeploymentRequest,
73 ) -> Result<Deployment, DbError>;
74 async fn delete_deployment(&self, id: &str) -> Result<(), DbError>;
75
76 async fn get_routing_config(&self) -> Result<Option<RoutingConfig>, DbError>;
78 async fn update_routing_config(
79 &self,
80 req: UpdateRoutingConfigRequest,
81 ) -> Result<RoutingConfig, DbError>;
82
83 async fn ping(&self) -> Result<(), DbError>;
84}
85
86#[derive(Debug, Clone, Serialize, Deserialize)]
87pub struct Team {
88 pub id: String,
89 pub name: String,
90 pub budget_cents: i64,
91 pub created_at: DateTime<Utc>,
92 pub updated_at: DateTime<Utc>,
93}
94
95#[derive(Debug, Clone, Serialize, Deserialize)]
96pub struct User {
97 pub id: String,
98 pub team_id: String,
99 pub email: String,
100 pub role: String,
101 #[serde(skip_serializing)]
102 pub password_hash: Option<String>,
103 pub created_at: DateTime<Utc>,
104}
105
106#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct ApiKey {
108 pub id: String,
109 pub key_hash: String,
110 pub user_id: String,
111 pub team_id: String,
112 pub name: Option<String>,
113 pub is_active: bool,
114 pub created_at: DateTime<Utc>,
115 pub expires_at: Option<DateTime<Utc>>,
116}
117
118#[derive(Debug, Clone, Serialize, Deserialize)]
119pub struct ModelAlias {
120 pub id: String,
121 pub team_id: String,
122 pub alias: String,
123 pub target_model: String,
124 pub provider: String,
125 pub created_at: DateTime<Utc>,
126}
127
128#[derive(Debug, Clone, Serialize, Deserialize)]
129pub struct Quota {
130 pub id: String,
131 pub team_id: String,
132 pub rpm_limit: i32,
133 pub tpm_limit: i32,
134 pub updated_at: DateTime<Utc>,
135}
136
137#[derive(Debug, Clone, Serialize, Deserialize)]
138pub struct UsageLog {
139 pub id: String,
140 pub team_id: String,
141 pub api_key_id: String,
142 pub model: String,
143 pub input_tokens: i32,
144 pub output_tokens: i32,
145 pub response_time_ms: i64,
146 pub recorded_at: DateTime<Utc>,
147}