oxirs_vec/multi_tenancy/
types.rs

1//! Core types for multi-tenancy support
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use thiserror::Error;
7
8/// Result type for multi-tenancy operations
9pub type MultiTenancyResult<T> = Result<T, MultiTenancyError>;
10
11/// Errors that can occur in multi-tenancy operations
12#[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/// Context for tenant operations
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct TenantContext {
48    /// Tenant identifier
49    pub tenant_id: String,
50
51    /// Request timestamp
52    pub timestamp: DateTime<Utc>,
53
54    /// Request metadata
55    pub metadata: HashMap<String, String>,
56
57    /// Authentication token (optional)
58    pub auth_token: Option<String>,
59
60    /// Client IP address (optional)
61    pub client_ip: Option<String>,
62
63    /// User agent (optional)
64    pub user_agent: Option<String>,
65}
66
67impl TenantContext {
68    /// Create a new tenant context
69    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    /// Add metadata to the context
81    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    /// Set authentication token
87    pub fn with_auth_token(mut self, token: impl Into<String>) -> Self {
88        self.auth_token = Some(token.into());
89        self
90    }
91
92    /// Set client IP
93    pub fn with_client_ip(mut self, ip: impl Into<String>) -> Self {
94        self.client_ip = Some(ip.into());
95        self
96    }
97
98    /// Set user agent
99    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/// Type of tenant operation for billing and metering
106#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
107pub enum TenantOperation {
108    /// Vector insertion
109    VectorInsert,
110
111    /// Vector search/query
112    VectorSearch,
113
114    /// Vector update
115    VectorUpdate,
116
117    /// Vector deletion
118    VectorDelete,
119
120    /// Index build/rebuild
121    IndexBuild,
122
123    /// Batch operation
124    BatchOperation,
125
126    /// Embedding generation
127    EmbeddingGeneration,
128
129    /// Cross-encoder re-ranking
130    Reranking,
131
132    /// Custom operation
133    Custom(u32),
134}
135
136impl TenantOperation {
137    /// Get operation name
138    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    /// Get default cost weight for billing
153    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/// Tenant statistics for monitoring and analytics
169#[derive(Debug, Clone, Serialize, Deserialize)]
170pub struct TenantStatistics {
171    /// Tenant ID
172    pub tenant_id: String,
173
174    /// Total vectors stored
175    pub total_vectors: usize,
176
177    /// Total queries executed
178    pub total_queries: u64,
179
180    /// Average query latency (milliseconds)
181    pub avg_query_latency_ms: f64,
182
183    /// Peak queries per second
184    pub peak_qps: f64,
185
186    /// Current storage usage (bytes)
187    pub storage_bytes: u64,
188
189    /// Current memory usage (bytes)
190    pub memory_bytes: u64,
191
192    /// Total API calls
193    pub api_calls: u64,
194
195    /// Error count
196    pub error_count: u64,
197
198    /// Last activity timestamp
199    pub last_activity: DateTime<Utc>,
200
201    /// Operation counts by type
202    pub operation_counts: HashMap<String, u64>,
203
204    /// Custom metrics
205    pub custom_metrics: HashMap<String, f64>,
206}
207
208impl TenantStatistics {
209    /// Create new statistics for a tenant
210    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    /// Record an operation
228    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    /// Record a query with latency
236    pub fn record_query(&mut self, latency_ms: f64) {
237        self.total_queries += 1;
238        self.record_operation(TenantOperation::VectorSearch);
239
240        // Update average latency using exponential moving average
241        let alpha = 0.1;
242        self.avg_query_latency_ms = alpha * latency_ms + (1.0 - alpha) * self.avg_query_latency_ms;
243    }
244
245    /// Record an error
246    pub fn record_error(&mut self) {
247        self.error_count += 1;
248        self.last_activity = Utc::now();
249    }
250
251    /// Update storage usage
252    pub fn update_storage(&mut self, bytes: u64) {
253        self.storage_bytes = bytes;
254    }
255
256    /// Update memory usage
257    pub fn update_memory(&mut self, bytes: u64) {
258        self.memory_bytes = bytes;
259    }
260
261    /// Set custom metric
262    pub fn set_custom_metric(&mut self, key: impl Into<String>, value: f64) {
263        self.custom_metrics.insert(key.into(), value);
264    }
265
266    /// Get error rate
267    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); // EMA starting from 0
311
312        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); // 1 error out of 2 calls
318    }
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}