1use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use thiserror::Error;
7
8pub type MultiTenancyResult<T> = Result<T, MultiTenancyError>;
10
11#[derive(Debug, Error, Clone, Serialize, Deserialize)]
13pub enum MultiTenancyError {
14 #[error("Tenant not found: {tenant_id}")]
15 TenantNotFound { tenant_id: String },
16
17 #[error("Tenant already exists: {tenant_id}")]
18 TenantAlreadyExists { tenant_id: String },
19
20 #[error("Quota exceeded for tenant {tenant_id}: {resource}")]
21 QuotaExceeded { tenant_id: String, resource: String },
22
23 #[error("Rate limit exceeded for tenant {tenant_id}")]
24 RateLimitExceeded { tenant_id: String },
25
26 #[error("Access denied for tenant {tenant_id}: {reason}")]
27 AccessDenied { tenant_id: String, reason: String },
28
29 #[error("Tenant is suspended: {tenant_id}")]
30 TenantSuspended { tenant_id: String },
31
32 #[error("Invalid tenant configuration: {message}")]
33 InvalidConfiguration { message: String },
34
35 #[error("Isolation violation: {message}")]
36 IsolationViolation { message: String },
37
38 #[error("Billing error: {message}")]
39 BillingError { message: String },
40
41 #[error("Internal error: {message}")]
42 InternalError { message: String },
43}
44
45#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct TenantContext {
48 pub tenant_id: String,
50
51 pub timestamp: DateTime<Utc>,
53
54 pub metadata: HashMap<String, String>,
56
57 pub auth_token: Option<String>,
59
60 pub client_ip: Option<String>,
62
63 pub user_agent: Option<String>,
65}
66
67impl TenantContext {
68 pub fn new(tenant_id: impl Into<String>) -> Self {
70 Self {
71 tenant_id: tenant_id.into(),
72 timestamp: Utc::now(),
73 metadata: HashMap::new(),
74 auth_token: None,
75 client_ip: None,
76 user_agent: None,
77 }
78 }
79
80 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
82 self.metadata.insert(key.into(), value.into());
83 self
84 }
85
86 pub fn with_auth_token(mut self, token: impl Into<String>) -> Self {
88 self.auth_token = Some(token.into());
89 self
90 }
91
92 pub fn with_client_ip(mut self, ip: impl Into<String>) -> Self {
94 self.client_ip = Some(ip.into());
95 self
96 }
97
98 pub fn with_user_agent(mut self, user_agent: impl Into<String>) -> Self {
100 self.user_agent = Some(user_agent.into());
101 self
102 }
103}
104
105#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
107pub enum TenantOperation {
108 VectorInsert,
110
111 VectorSearch,
113
114 VectorUpdate,
116
117 VectorDelete,
119
120 IndexBuild,
122
123 BatchOperation,
125
126 EmbeddingGeneration,
128
129 Reranking,
131
132 Custom(u32),
134}
135
136impl TenantOperation {
137 pub fn name(&self) -> &'static str {
139 match self {
140 Self::VectorInsert => "vector_insert",
141 Self::VectorSearch => "vector_search",
142 Self::VectorUpdate => "vector_update",
143 Self::VectorDelete => "vector_delete",
144 Self::IndexBuild => "index_build",
145 Self::BatchOperation => "batch_operation",
146 Self::EmbeddingGeneration => "embedding_generation",
147 Self::Reranking => "reranking",
148 Self::Custom(_) => "custom",
149 }
150 }
151
152 pub fn default_cost_weight(&self) -> f64 {
154 match self {
155 Self::VectorInsert => 1.0,
156 Self::VectorSearch => 2.0,
157 Self::VectorUpdate => 1.5,
158 Self::VectorDelete => 0.5,
159 Self::IndexBuild => 10.0,
160 Self::BatchOperation => 5.0,
161 Self::EmbeddingGeneration => 3.0,
162 Self::Reranking => 4.0,
163 Self::Custom(_) => 1.0,
164 }
165 }
166}
167
168#[derive(Debug, Clone, Serialize, Deserialize)]
170pub struct TenantStatistics {
171 pub tenant_id: String,
173
174 pub total_vectors: usize,
176
177 pub total_queries: u64,
179
180 pub avg_query_latency_ms: f64,
182
183 pub peak_qps: f64,
185
186 pub storage_bytes: u64,
188
189 pub memory_bytes: u64,
191
192 pub api_calls: u64,
194
195 pub error_count: u64,
197
198 pub last_activity: DateTime<Utc>,
200
201 pub operation_counts: HashMap<String, u64>,
203
204 pub custom_metrics: HashMap<String, f64>,
206}
207
208impl TenantStatistics {
209 pub fn new(tenant_id: impl Into<String>) -> Self {
211 Self {
212 tenant_id: tenant_id.into(),
213 total_vectors: 0,
214 total_queries: 0,
215 avg_query_latency_ms: 0.0,
216 peak_qps: 0.0,
217 storage_bytes: 0,
218 memory_bytes: 0,
219 api_calls: 0,
220 error_count: 0,
221 last_activity: Utc::now(),
222 operation_counts: HashMap::new(),
223 custom_metrics: HashMap::new(),
224 }
225 }
226
227 pub fn record_operation(&mut self, operation: TenantOperation) {
229 let op_name = operation.name().to_string();
230 *self.operation_counts.entry(op_name).or_insert(0) += 1;
231 self.api_calls += 1;
232 self.last_activity = Utc::now();
233 }
234
235 pub fn record_query(&mut self, latency_ms: f64) {
237 self.total_queries += 1;
238 self.record_operation(TenantOperation::VectorSearch);
239
240 let alpha = 0.1;
242 self.avg_query_latency_ms = alpha * latency_ms + (1.0 - alpha) * self.avg_query_latency_ms;
243 }
244
245 pub fn record_error(&mut self) {
247 self.error_count += 1;
248 self.last_activity = Utc::now();
249 }
250
251 pub fn update_storage(&mut self, bytes: u64) {
253 self.storage_bytes = bytes;
254 }
255
256 pub fn update_memory(&mut self, bytes: u64) {
258 self.memory_bytes = bytes;
259 }
260
261 pub fn set_custom_metric(&mut self, key: impl Into<String>, value: f64) {
263 self.custom_metrics.insert(key.into(), value);
264 }
265
266 pub fn error_rate(&self) -> f64 {
268 if self.api_calls == 0 {
269 0.0
270 } else {
271 self.error_count as f64 / self.api_calls as f64
272 }
273 }
274}
275
276#[cfg(test)]
277mod tests {
278 use super::*;
279
280 #[test]
281 fn test_tenant_context_creation() {
282 let ctx = TenantContext::new("tenant1")
283 .with_metadata("region", "us-west")
284 .with_auth_token("token123")
285 .with_client_ip("192.168.1.1");
286
287 assert_eq!(ctx.tenant_id, "tenant1");
288 assert_eq!(ctx.metadata.get("region"), Some(&"us-west".to_string()));
289 assert_eq!(ctx.auth_token, Some("token123".to_string()));
290 assert_eq!(ctx.client_ip, Some("192.168.1.1".to_string()));
291 }
292
293 #[test]
294 fn test_tenant_operation_cost_weights() {
295 assert_eq!(TenantOperation::VectorInsert.default_cost_weight(), 1.0);
296 assert_eq!(TenantOperation::VectorSearch.default_cost_weight(), 2.0);
297 assert_eq!(TenantOperation::IndexBuild.default_cost_weight(), 10.0);
298 }
299
300 #[test]
301 fn test_tenant_statistics() {
302 let mut stats = TenantStatistics::new("tenant1");
303
304 assert_eq!(stats.total_queries, 0);
305 assert_eq!(stats.api_calls, 0);
306
307 stats.record_query(100.0);
308 assert_eq!(stats.total_queries, 1);
309 assert_eq!(stats.api_calls, 1);
310 assert!((stats.avg_query_latency_ms - 10.0).abs() < 1.0); stats.record_operation(TenantOperation::VectorInsert);
313 assert_eq!(stats.api_calls, 2);
314
315 stats.record_error();
316 assert_eq!(stats.error_count, 1);
317 assert!((stats.error_rate() - 0.5).abs() < 0.01); }
319
320 #[test]
321 fn test_multitenance_error_display() {
322 let error = MultiTenancyError::TenantNotFound {
323 tenant_id: "tenant1".to_string(),
324 };
325 assert!(error.to_string().contains("tenant1"));
326
327 let error = MultiTenancyError::QuotaExceeded {
328 tenant_id: "tenant2".to_string(),
329 resource: "storage".to_string(),
330 };
331 assert!(error.to_string().contains("tenant2"));
332 assert!(error.to_string().contains("storage"));
333 }
334}