Skip to main content

hyperinfer_core/traits/
database.rs

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    // --- Deployment CRUD ---
62    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    // --- Routing Config ---
77    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}