mirror of
https://github.com/ruvnet/RuView
synced 2026-08-04 19:31:42 +00:00
Merge commit 'd803bfe2b1fe7f5e219e50ac20d6801a0a58ac75' as 'vendor/ruvector'
This commit is contained in:
ruv
@@ -0,0 +1,860 @@
|
||||
//! Tenant Isolation Enforcement for RuVector Multi-Tenancy
|
||||
//!
|
||||
//! Provides three isolation levels:
|
||||
//! - Shared: RLS policies on tenant_id column
|
||||
//! - Partition: Separate partitions per tenant
|
||||
//! - Dedicated: Schema-level isolation with separate indexes
|
||||
|
||||
use dashmap::DashMap;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::registry::{get_registry, IsolationLevel};
|
||||
use super::validation::{
|
||||
escape_string_literal, quote_identifier, safe_partition_name, safe_schema_name,
|
||||
validate_identifier, validate_tenant_id,
|
||||
};
|
||||
|
||||
/// Partition configuration for tenant
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PartitionConfig {
|
||||
/// Tenant ID
|
||||
pub tenant_id: String,
|
||||
/// Partition name (e.g., "embeddings_acme_corp")
|
||||
pub partition_name: String,
|
||||
/// Parent table name
|
||||
pub parent_table: String,
|
||||
/// Partition key value (tenant_id)
|
||||
pub partition_key: String,
|
||||
/// Creation timestamp
|
||||
pub created_at: i64,
|
||||
}
|
||||
|
||||
/// Dedicated schema configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DedicatedSchemaConfig {
|
||||
/// Tenant ID
|
||||
pub tenant_id: String,
|
||||
/// Schema name (e.g., "tenant_acme_corp")
|
||||
pub schema_name: String,
|
||||
/// Tables in this schema
|
||||
pub tables: Vec<String>,
|
||||
/// Indexes in this schema
|
||||
pub indexes: Vec<String>,
|
||||
/// Creation timestamp
|
||||
pub created_at: i64,
|
||||
}
|
||||
|
||||
/// Isolation enforcement manager
|
||||
pub struct IsolationManager {
|
||||
/// Partition configurations by tenant
|
||||
partitions: DashMap<String, Vec<PartitionConfig>>,
|
||||
/// Dedicated schema configurations by tenant
|
||||
dedicated_schemas: DashMap<String, DedicatedSchemaConfig>,
|
||||
/// Tables with RLS enabled (table_name -> tenant_column)
|
||||
rls_tables: DashMap<String, String>,
|
||||
/// Migration state tracking
|
||||
migration_state: DashMap<String, MigrationState>,
|
||||
}
|
||||
|
||||
/// State of tenant isolation migration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MigrationState {
|
||||
/// Tenant ID
|
||||
pub tenant_id: String,
|
||||
/// Source isolation level
|
||||
pub from_level: IsolationLevel,
|
||||
/// Target isolation level
|
||||
pub to_level: IsolationLevel,
|
||||
/// Migration status
|
||||
pub status: MigrationStatus,
|
||||
/// Progress percentage (0-100)
|
||||
pub progress: u8,
|
||||
/// Vectors migrated so far
|
||||
pub vectors_migrated: u64,
|
||||
/// Total vectors to migrate
|
||||
pub total_vectors: u64,
|
||||
/// Start timestamp
|
||||
pub started_at: i64,
|
||||
/// Completion timestamp
|
||||
pub completed_at: Option<i64>,
|
||||
/// Error message if failed
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// Migration status
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum MigrationStatus {
|
||||
/// Migration pending
|
||||
Pending,
|
||||
/// Migration in progress
|
||||
InProgress,
|
||||
/// Migration completed
|
||||
Completed,
|
||||
/// Migration failed
|
||||
Failed,
|
||||
/// Migration cancelled
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
impl IsolationManager {
|
||||
/// Create a new isolation manager
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
partitions: DashMap::new(),
|
||||
dedicated_schemas: DashMap::new(),
|
||||
rls_tables: DashMap::new(),
|
||||
migration_state: DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Shared Isolation (RLS-based)
|
||||
// =========================================================================
|
||||
|
||||
/// Enable shared isolation for a table (RLS policies)
|
||||
pub fn enable_shared_isolation(
|
||||
&self,
|
||||
table_name: &str,
|
||||
tenant_column: &str,
|
||||
) -> Result<String, IsolationError> {
|
||||
// Validate identifiers to prevent SQL injection
|
||||
validate_identifier(table_name)
|
||||
.map_err(|e| IsolationError::SqlError(format!("Invalid table name: {}", e)))?;
|
||||
validate_identifier(tenant_column)
|
||||
.map_err(|e| IsolationError::SqlError(format!("Invalid column name: {}", e)))?;
|
||||
|
||||
// Use quoted identifiers for safety
|
||||
let quoted_table = quote_identifier(table_name);
|
||||
let quoted_column = quote_identifier(tenant_column);
|
||||
|
||||
// Generate SQL for RLS setup with quoted identifiers
|
||||
let sql = format!(
|
||||
r#"
|
||||
-- Enable RLS on the table
|
||||
ALTER TABLE {table} ENABLE ROW LEVEL SECURITY;
|
||||
|
||||
-- Drop existing policies if any
|
||||
DROP POLICY IF EXISTS ruvector_tenant_isolation ON {table};
|
||||
DROP POLICY IF EXISTS ruvector_admin_bypass ON {table};
|
||||
|
||||
-- Create tenant isolation policy
|
||||
CREATE POLICY ruvector_tenant_isolation ON {table}
|
||||
USING ({column} = current_setting('ruvector.tenant_id', true))
|
||||
WITH CHECK ({column} = current_setting('ruvector.tenant_id', true));
|
||||
|
||||
-- Create admin bypass policy (for ruvector_admin role)
|
||||
CREATE POLICY ruvector_admin_bypass ON {table}
|
||||
FOR ALL
|
||||
TO ruvector_admin
|
||||
USING (true)
|
||||
WITH CHECK (true);
|
||||
|
||||
-- Create wildcard policy for admin queries
|
||||
CREATE POLICY ruvector_admin_wildcard ON {table}
|
||||
FOR SELECT
|
||||
USING (current_setting('ruvector.tenant_id', true) = '*');
|
||||
"#,
|
||||
table = quoted_table,
|
||||
column = quoted_column
|
||||
);
|
||||
|
||||
self.rls_tables
|
||||
.insert(table_name.to_string(), tenant_column.to_string());
|
||||
|
||||
Ok(sql)
|
||||
}
|
||||
|
||||
/// Check if a table has RLS enabled
|
||||
pub fn is_rls_enabled(&self, table_name: &str) -> bool {
|
||||
self.rls_tables.contains_key(table_name)
|
||||
}
|
||||
|
||||
/// Get tenant column for RLS table
|
||||
pub fn get_tenant_column(&self, table_name: &str) -> Option<String> {
|
||||
self.rls_tables.get(table_name).map(|r| r.value().clone())
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Partition Isolation
|
||||
// =========================================================================
|
||||
|
||||
/// Create partition for a tenant
|
||||
pub fn create_partition(
|
||||
&self,
|
||||
tenant_id: &str,
|
||||
parent_table: &str,
|
||||
) -> Result<PartitionConfig, IsolationError> {
|
||||
// Validate inputs to prevent SQL injection
|
||||
validate_tenant_id(tenant_id)
|
||||
.map_err(|e| IsolationError::SqlError(format!("Invalid tenant ID: {}", e)))?;
|
||||
validate_identifier(parent_table)
|
||||
.map_err(|e| IsolationError::SqlError(format!("Invalid table name: {}", e)))?;
|
||||
|
||||
// Generate safe partition name
|
||||
let partition_name = safe_partition_name(tenant_id, parent_table)
|
||||
.map_err(|e| IsolationError::SqlError(format!("Invalid partition name: {}", e)))?;
|
||||
|
||||
let config = PartitionConfig {
|
||||
tenant_id: tenant_id.to_string(),
|
||||
partition_name,
|
||||
parent_table: parent_table.to_string(),
|
||||
partition_key: tenant_id.to_string(),
|
||||
created_at: chrono_now_millis(),
|
||||
};
|
||||
|
||||
// Store partition config
|
||||
self.partitions
|
||||
.entry(tenant_id.to_string())
|
||||
.or_insert_with(Vec::new)
|
||||
.push(config.clone());
|
||||
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
/// Generate SQL for creating a partition
|
||||
pub fn generate_partition_sql(&self, config: &PartitionConfig) -> String {
|
||||
// Use quoted identifiers for safety
|
||||
let quoted_partition = quote_identifier(&config.partition_name);
|
||||
let quoted_parent = quote_identifier(&config.parent_table);
|
||||
let escaped_tenant_id = escape_string_literal(&config.partition_key);
|
||||
let safe_index_name = format!("idx_{}_vec", config.partition_name);
|
||||
|
||||
format!(
|
||||
r#"
|
||||
-- Create partition for tenant
|
||||
CREATE TABLE IF NOT EXISTS {partition} PARTITION OF {parent}
|
||||
FOR VALUES IN ('{tenant_id}');
|
||||
|
||||
-- Create indexes on partition
|
||||
CREATE INDEX IF NOT EXISTS {index_name}
|
||||
ON {partition} USING ruhnsw (vec vector_cosine_ops);
|
||||
"#,
|
||||
partition = quoted_partition,
|
||||
parent = quoted_parent,
|
||||
tenant_id = escaped_tenant_id,
|
||||
index_name = quote_identifier(&safe_index_name)
|
||||
)
|
||||
}
|
||||
|
||||
/// Get partitions for a tenant
|
||||
pub fn get_partitions(&self, tenant_id: &str) -> Vec<PartitionConfig> {
|
||||
self.partitions
|
||||
.get(tenant_id)
|
||||
.map(|r| r.value().clone())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Drop partition for a tenant
|
||||
pub fn drop_partition(
|
||||
&self,
|
||||
tenant_id: &str,
|
||||
partition_name: &str,
|
||||
) -> Result<String, IsolationError> {
|
||||
// Validate inputs to prevent SQL injection
|
||||
validate_tenant_id(tenant_id)
|
||||
.map_err(|e| IsolationError::SqlError(format!("Invalid tenant ID: {}", e)))?;
|
||||
validate_identifier(partition_name)
|
||||
.map_err(|e| IsolationError::SqlError(format!("Invalid partition name: {}", e)))?;
|
||||
|
||||
// Verify partition belongs to this tenant (security check)
|
||||
let partition_exists = self
|
||||
.partitions
|
||||
.get(tenant_id)
|
||||
.map(|partitions| {
|
||||
partitions
|
||||
.iter()
|
||||
.any(|p| p.partition_name == partition_name)
|
||||
})
|
||||
.unwrap_or(false);
|
||||
|
||||
if !partition_exists {
|
||||
return Err(IsolationError::PartitionNotFound(
|
||||
partition_name.to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Remove from tracking
|
||||
if let Some(mut partitions) = self.partitions.get_mut(tenant_id) {
|
||||
partitions.retain(|p| p.partition_name != partition_name);
|
||||
}
|
||||
|
||||
// Use quoted identifier for safety
|
||||
Ok(format!(
|
||||
"DROP TABLE IF EXISTS {} CASCADE;",
|
||||
quote_identifier(partition_name)
|
||||
))
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Dedicated Isolation (Schema-level)
|
||||
// =========================================================================
|
||||
|
||||
/// Create dedicated schema for a tenant
|
||||
pub fn create_dedicated_schema(
|
||||
&self,
|
||||
tenant_id: &str,
|
||||
) -> Result<DedicatedSchemaConfig, IsolationError> {
|
||||
// Validate tenant ID to prevent SQL injection
|
||||
validate_tenant_id(tenant_id)
|
||||
.map_err(|e| IsolationError::SqlError(format!("Invalid tenant ID: {}", e)))?;
|
||||
|
||||
// Generate safe schema name
|
||||
let schema_name = safe_schema_name(tenant_id)
|
||||
.map_err(|e| IsolationError::SqlError(format!("Invalid schema name: {}", e)))?;
|
||||
|
||||
let config = DedicatedSchemaConfig {
|
||||
tenant_id: tenant_id.to_string(),
|
||||
schema_name,
|
||||
tables: Vec::new(),
|
||||
indexes: Vec::new(),
|
||||
created_at: chrono_now_millis(),
|
||||
};
|
||||
|
||||
self.dedicated_schemas
|
||||
.insert(tenant_id.to_string(), config.clone());
|
||||
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
/// Generate SQL for creating dedicated schema
|
||||
pub fn generate_schema_sql(&self, config: &DedicatedSchemaConfig) -> String {
|
||||
// Use quoted identifiers for safety
|
||||
let quoted_schema = quote_identifier(&config.schema_name);
|
||||
let index_name = format!("idx_{}_embeddings_vec", config.schema_name);
|
||||
let quoted_index = quote_identifier(&index_name);
|
||||
|
||||
format!(
|
||||
r#"
|
||||
-- Create dedicated schema for tenant
|
||||
CREATE SCHEMA IF NOT EXISTS {schema};
|
||||
|
||||
-- Set search path to include tenant schema
|
||||
-- (Application should SET search_path = {schema}, public;)
|
||||
|
||||
-- Create embeddings table in tenant schema
|
||||
CREATE TABLE IF NOT EXISTS {schema}."embeddings" (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
content TEXT,
|
||||
vec vector(1536),
|
||||
metadata JSONB DEFAULT '{{}}',
|
||||
created_at TIMESTAMPTZ DEFAULT NOW()
|
||||
);
|
||||
|
||||
-- Create HNSW index
|
||||
CREATE INDEX IF NOT EXISTS {index_name}
|
||||
ON {schema}."embeddings" USING ruhnsw (vec vector_cosine_ops);
|
||||
|
||||
-- Grant usage to tenant role
|
||||
GRANT USAGE ON SCHEMA {schema} TO ruvector_users;
|
||||
GRANT ALL ON ALL TABLES IN SCHEMA {schema} TO ruvector_users;
|
||||
GRANT ALL ON ALL SEQUENCES IN SCHEMA {schema} TO ruvector_users;
|
||||
"#,
|
||||
schema = quoted_schema,
|
||||
index_name = quoted_index
|
||||
)
|
||||
}
|
||||
|
||||
/// Get dedicated schema for a tenant
|
||||
pub fn get_dedicated_schema(&self, tenant_id: &str) -> Option<DedicatedSchemaConfig> {
|
||||
self.dedicated_schemas
|
||||
.get(tenant_id)
|
||||
.map(|r| r.value().clone())
|
||||
}
|
||||
|
||||
/// Add table to dedicated schema tracking
|
||||
pub fn register_schema_table(
|
||||
&self,
|
||||
tenant_id: &str,
|
||||
table_name: &str,
|
||||
) -> Result<(), IsolationError> {
|
||||
if let Some(mut schema) = self.dedicated_schemas.get_mut(tenant_id) {
|
||||
schema.tables.push(table_name.to_string());
|
||||
Ok(())
|
||||
} else {
|
||||
Err(IsolationError::SchemaNotFound(tenant_id.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Add index to dedicated schema tracking
|
||||
pub fn register_schema_index(
|
||||
&self,
|
||||
tenant_id: &str,
|
||||
index_name: &str,
|
||||
) -> Result<(), IsolationError> {
|
||||
if let Some(mut schema) = self.dedicated_schemas.get_mut(tenant_id) {
|
||||
schema.indexes.push(index_name.to_string());
|
||||
Ok(())
|
||||
} else {
|
||||
Err(IsolationError::SchemaNotFound(tenant_id.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Drop dedicated schema
|
||||
pub fn drop_dedicated_schema(
|
||||
&self,
|
||||
tenant_id: &str,
|
||||
cascade: bool,
|
||||
) -> Result<String, IsolationError> {
|
||||
// Validate tenant ID
|
||||
validate_tenant_id(tenant_id)
|
||||
.map_err(|e| IsolationError::SqlError(format!("Invalid tenant ID: {}", e)))?;
|
||||
|
||||
let config = self
|
||||
.dedicated_schemas
|
||||
.remove(tenant_id)
|
||||
.map(|(_, v)| v)
|
||||
.ok_or_else(|| IsolationError::SchemaNotFound(tenant_id.to_string()))?;
|
||||
|
||||
let cascade_clause = if cascade { "CASCADE" } else { "RESTRICT" };
|
||||
|
||||
// Use quoted identifier for safety
|
||||
Ok(format!(
|
||||
"DROP SCHEMA IF EXISTS {} {};",
|
||||
quote_identifier(&config.schema_name),
|
||||
cascade_clause
|
||||
))
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Migration Between Isolation Levels
|
||||
// =========================================================================
|
||||
|
||||
/// Start migration to a new isolation level
|
||||
pub fn start_migration(
|
||||
&self,
|
||||
tenant_id: &str,
|
||||
target_level: IsolationLevel,
|
||||
) -> Result<MigrationState, IsolationError> {
|
||||
// Check if migration already in progress
|
||||
if let Some(state) = self.migration_state.get(tenant_id) {
|
||||
if state.status == MigrationStatus::InProgress {
|
||||
return Err(IsolationError::MigrationInProgress(tenant_id.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
// Get current tenant config
|
||||
let config = get_registry()
|
||||
.get(tenant_id)
|
||||
.ok_or_else(|| IsolationError::TenantNotFound(tenant_id.to_string()))?;
|
||||
|
||||
let state = MigrationState {
|
||||
tenant_id: tenant_id.to_string(),
|
||||
from_level: config.isolation_level,
|
||||
to_level: target_level,
|
||||
status: MigrationStatus::Pending,
|
||||
progress: 0,
|
||||
vectors_migrated: 0,
|
||||
total_vectors: 0, // Will be set during migration
|
||||
started_at: chrono_now_millis(),
|
||||
completed_at: None,
|
||||
error: None,
|
||||
};
|
||||
|
||||
self.migration_state
|
||||
.insert(tenant_id.to_string(), state.clone());
|
||||
|
||||
// Mark tenant as migrating
|
||||
if let Some(shared_state) = get_registry().get_shared_state(tenant_id) {
|
||||
shared_state.set_migrating(true);
|
||||
}
|
||||
|
||||
Ok(state)
|
||||
}
|
||||
|
||||
/// Update migration progress
|
||||
pub fn update_migration_progress(
|
||||
&self,
|
||||
tenant_id: &str,
|
||||
vectors_migrated: u64,
|
||||
total_vectors: u64,
|
||||
) -> Result<(), IsolationError> {
|
||||
let mut state = self
|
||||
.migration_state
|
||||
.get_mut(tenant_id)
|
||||
.ok_or_else(|| IsolationError::NoMigrationInProgress(tenant_id.to_string()))?;
|
||||
|
||||
state.vectors_migrated = vectors_migrated;
|
||||
state.total_vectors = total_vectors;
|
||||
state.progress = if total_vectors > 0 {
|
||||
((vectors_migrated as f64 / total_vectors as f64) * 100.0) as u8
|
||||
} else {
|
||||
100
|
||||
};
|
||||
state.status = MigrationStatus::InProgress;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Complete migration
|
||||
pub fn complete_migration(&self, tenant_id: &str) -> Result<MigrationState, IsolationError> {
|
||||
let mut state = self
|
||||
.migration_state
|
||||
.get_mut(tenant_id)
|
||||
.ok_or_else(|| IsolationError::NoMigrationInProgress(tenant_id.to_string()))?;
|
||||
|
||||
state.status = MigrationStatus::Completed;
|
||||
state.progress = 100;
|
||||
state.completed_at = Some(chrono_now_millis());
|
||||
|
||||
// Clear migrating flag
|
||||
if let Some(shared_state) = get_registry().get_shared_state(tenant_id) {
|
||||
shared_state.set_migrating(false);
|
||||
}
|
||||
|
||||
Ok(state.clone())
|
||||
}
|
||||
|
||||
/// Fail migration
|
||||
pub fn fail_migration(&self, tenant_id: &str, error: &str) -> Result<(), IsolationError> {
|
||||
let mut state = self
|
||||
.migration_state
|
||||
.get_mut(tenant_id)
|
||||
.ok_or_else(|| IsolationError::NoMigrationInProgress(tenant_id.to_string()))?;
|
||||
|
||||
state.status = MigrationStatus::Failed;
|
||||
state.error = Some(error.to_string());
|
||||
state.completed_at = Some(chrono_now_millis());
|
||||
|
||||
// Clear migrating flag
|
||||
if let Some(shared_state) = get_registry().get_shared_state(tenant_id) {
|
||||
shared_state.set_migrating(false);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get migration status
|
||||
pub fn get_migration_status(&self, tenant_id: &str) -> Option<MigrationState> {
|
||||
self.migration_state
|
||||
.get(tenant_id)
|
||||
.map(|r| r.value().clone())
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Query Routing
|
||||
// =========================================================================
|
||||
|
||||
/// Get the appropriate table/schema for a tenant's query
|
||||
///
|
||||
/// Returns a QueryRoute that uses parameterized placeholders ($1) instead of
|
||||
/// directly interpolating tenant_id values to prevent SQL injection.
|
||||
pub fn route_query(&self, tenant_id: &str, base_table: &str) -> QueryRoute {
|
||||
// Validate tenant_id to prevent SQL injection even when using parameterized queries
|
||||
// This provides defense-in-depth
|
||||
if validate_tenant_id(tenant_id).is_err() {
|
||||
// Invalid tenant_id - return a safe filter that will match nothing
|
||||
return QueryRoute::SharedWithFilter {
|
||||
table: base_table.to_string(),
|
||||
filter: "false".to_string(), // Safe - matches nothing
|
||||
tenant_param: None,
|
||||
};
|
||||
}
|
||||
|
||||
let config = match get_registry().get(tenant_id) {
|
||||
Some(c) => c,
|
||||
None => {
|
||||
return QueryRoute::SharedWithFilter {
|
||||
table: base_table.to_string(),
|
||||
// Use parameterized query placeholder - caller must bind tenant_id
|
||||
filter: "tenant_id = $1".to_string(),
|
||||
tenant_param: Some(tenant_id.to_string()),
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
match config.isolation_level {
|
||||
IsolationLevel::Shared => QueryRoute::SharedWithFilter {
|
||||
table: base_table.to_string(),
|
||||
// Use parameterized query placeholder
|
||||
filter: "tenant_id = $1".to_string(),
|
||||
tenant_param: Some(tenant_id.to_string()),
|
||||
},
|
||||
IsolationLevel::Partition => {
|
||||
// Check if partition exists
|
||||
if let Some(partitions) = self.partitions.get(tenant_id) {
|
||||
if let Some(partition) =
|
||||
partitions.iter().find(|p| p.parent_table == base_table)
|
||||
{
|
||||
return QueryRoute::Partition {
|
||||
partition_table: partition.partition_name.clone(),
|
||||
};
|
||||
}
|
||||
}
|
||||
// Fall back to shared with filter (parameterized)
|
||||
QueryRoute::SharedWithFilter {
|
||||
table: base_table.to_string(),
|
||||
filter: "tenant_id = $1".to_string(),
|
||||
tenant_param: Some(tenant_id.to_string()),
|
||||
}
|
||||
}
|
||||
IsolationLevel::Dedicated => {
|
||||
// Check if dedicated schema exists
|
||||
if let Some(schema) = self.dedicated_schemas.get(tenant_id) {
|
||||
return QueryRoute::DedicatedSchema {
|
||||
schema: schema.schema_name.clone(),
|
||||
table: base_table.to_string(),
|
||||
};
|
||||
}
|
||||
// Fall back to shared with filter (parameterized)
|
||||
QueryRoute::SharedWithFilter {
|
||||
table: base_table.to_string(),
|
||||
filter: "tenant_id = $1".to_string(),
|
||||
tenant_param: Some(tenant_id.to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for IsolationManager {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Query routing result
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum QueryRoute {
|
||||
/// Use shared table with tenant filter (RLS handles this automatically)
|
||||
///
|
||||
/// The filter uses parameterized query placeholders ($1) for safety.
|
||||
/// The tenant_param contains the actual value to bind.
|
||||
SharedWithFilter {
|
||||
table: String,
|
||||
/// SQL filter clause using $1 placeholder for tenant_id
|
||||
filter: String,
|
||||
/// The tenant_id value to bind to $1 (None if filter is static like "false")
|
||||
tenant_param: Option<String>,
|
||||
},
|
||||
/// Use dedicated partition table
|
||||
Partition { partition_table: String },
|
||||
/// Use dedicated schema
|
||||
DedicatedSchema { schema: String, table: String },
|
||||
}
|
||||
|
||||
impl QueryRoute {
|
||||
/// Get the full table reference for SQL
|
||||
pub fn table_reference(&self) -> String {
|
||||
match self {
|
||||
Self::SharedWithFilter { table, .. } => table.clone(),
|
||||
Self::Partition { partition_table } => partition_table.clone(),
|
||||
Self::DedicatedSchema { schema, table } => {
|
||||
format!("{}.{}", quote_identifier(schema), quote_identifier(table))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get additional WHERE clause if needed (parameterized)
|
||||
///
|
||||
/// Returns the filter clause and the parameter value to bind.
|
||||
/// The filter uses $1 placeholder for the tenant_id.
|
||||
pub fn where_clause(&self) -> Option<String> {
|
||||
match self {
|
||||
Self::SharedWithFilter { filter, .. } => Some(filter.clone()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the tenant parameter value to bind to $1
|
||||
pub fn tenant_param(&self) -> Option<String> {
|
||||
match self {
|
||||
Self::SharedWithFilter { tenant_param, .. } => tenant_param.clone(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get WHERE clause and parameter together for convenience
|
||||
pub fn where_clause_with_param(&self) -> Option<(String, Option<String>)> {
|
||||
match self {
|
||||
Self::SharedWithFilter {
|
||||
filter,
|
||||
tenant_param,
|
||||
..
|
||||
} => Some((filter.clone(), tenant_param.clone())),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Isolation operation errors
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum IsolationError {
|
||||
/// Tenant not found
|
||||
TenantNotFound(String),
|
||||
/// Schema not found
|
||||
SchemaNotFound(String),
|
||||
/// Partition not found
|
||||
PartitionNotFound(String),
|
||||
/// Migration already in progress
|
||||
MigrationInProgress(String),
|
||||
/// No migration in progress
|
||||
NoMigrationInProgress(String),
|
||||
/// Invalid isolation level transition
|
||||
InvalidTransition {
|
||||
from: IsolationLevel,
|
||||
to: IsolationLevel,
|
||||
},
|
||||
/// SQL execution error
|
||||
SqlError(String),
|
||||
}
|
||||
|
||||
impl std::fmt::Display for IsolationError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::TenantNotFound(id) => write!(f, "Tenant not found: {}", id),
|
||||
Self::SchemaNotFound(id) => write!(f, "Dedicated schema not found for tenant: {}", id),
|
||||
Self::PartitionNotFound(name) => write!(f, "Partition not found: {}", name),
|
||||
Self::MigrationInProgress(id) => {
|
||||
write!(f, "Migration already in progress for tenant: {}", id)
|
||||
}
|
||||
Self::NoMigrationInProgress(id) => {
|
||||
write!(f, "No migration in progress for tenant: {}", id)
|
||||
}
|
||||
Self::InvalidTransition { from, to } => {
|
||||
write!(
|
||||
f,
|
||||
"Invalid isolation transition from {} to {}",
|
||||
from.as_str(),
|
||||
to.as_str()
|
||||
)
|
||||
}
|
||||
Self::SqlError(msg) => write!(f, "SQL error: {}", msg),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for IsolationError {}
|
||||
|
||||
/// Global isolation manager instance
|
||||
static ISOLATION_MANAGER: once_cell::sync::Lazy<IsolationManager> =
|
||||
once_cell::sync::Lazy::new(IsolationManager::new);
|
||||
|
||||
/// Get the global isolation manager
|
||||
pub fn get_isolation_manager() -> &'static IsolationManager {
|
||||
&ISOLATION_MANAGER
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
fn chrono_now_millis() -> i64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as i64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_create_partition_config() {
|
||||
let manager = IsolationManager::new();
|
||||
let config = manager.create_partition("acme-corp", "embeddings").unwrap();
|
||||
|
||||
assert_eq!(config.tenant_id, "acme-corp");
|
||||
assert_eq!(config.partition_name, "embeddings_acme_corp");
|
||||
assert_eq!(config.parent_table, "embeddings");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_create_dedicated_schema() {
|
||||
let manager = IsolationManager::new();
|
||||
let config = manager.create_dedicated_schema("acme-corp").unwrap();
|
||||
|
||||
assert_eq!(config.tenant_id, "acme-corp");
|
||||
assert_eq!(config.schema_name, "tenant_acme_corp");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_query_routing() {
|
||||
let manager = IsolationManager::new();
|
||||
|
||||
// Default routing (no config) should use shared with filter
|
||||
let route = manager.route_query("unknown_tenant", "embeddings");
|
||||
match route {
|
||||
QueryRoute::SharedWithFilter {
|
||||
table,
|
||||
filter,
|
||||
tenant_param,
|
||||
} => {
|
||||
assert_eq!(table, "embeddings");
|
||||
// Filter should use parameterized placeholder
|
||||
assert_eq!(filter, "tenant_id = $1");
|
||||
// Tenant param should contain the tenant_id
|
||||
assert_eq!(tenant_param, Some("unknown_tenant".to_string()));
|
||||
}
|
||||
_ => panic!("Expected SharedWithFilter"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_query_routing_invalid_tenant() {
|
||||
let manager = IsolationManager::new();
|
||||
|
||||
// Invalid tenant_id should return safe "false" filter
|
||||
let route = manager.route_query("'; DROP TABLE users;--", "embeddings");
|
||||
match route {
|
||||
QueryRoute::SharedWithFilter {
|
||||
filter,
|
||||
tenant_param,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(filter, "false");
|
||||
assert!(tenant_param.is_none());
|
||||
}
|
||||
_ => panic!("Expected SharedWithFilter with false filter"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rls_tracking() {
|
||||
let manager = IsolationManager::new();
|
||||
|
||||
// Enable RLS
|
||||
manager
|
||||
.enable_shared_isolation("embeddings", "tenant_id")
|
||||
.unwrap();
|
||||
|
||||
// Check tracking
|
||||
assert!(manager.is_rls_enabled("embeddings"));
|
||||
assert_eq!(
|
||||
manager.get_tenant_column("embeddings"),
|
||||
Some("tenant_id".to_string())
|
||||
);
|
||||
assert!(!manager.is_rls_enabled("other_table"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_migration_state() {
|
||||
let manager = IsolationManager::new();
|
||||
|
||||
// Register a tenant first
|
||||
let registry = get_registry();
|
||||
let config = super::super::registry::TenantConfig::new("test-tenant".to_string());
|
||||
let _ = registry.register(config);
|
||||
|
||||
// Start migration
|
||||
let state = manager
|
||||
.start_migration("test-tenant", IsolationLevel::Partition)
|
||||
.unwrap();
|
||||
assert_eq!(state.status, MigrationStatus::Pending);
|
||||
assert_eq!(state.from_level, IsolationLevel::Shared);
|
||||
assert_eq!(state.to_level, IsolationLevel::Partition);
|
||||
|
||||
// Should fail if trying to start another migration
|
||||
let result = manager.start_migration("test-tenant", IsolationLevel::Dedicated);
|
||||
assert!(result.is_err());
|
||||
|
||||
// Update progress
|
||||
manager
|
||||
.update_migration_progress("test-tenant", 50, 100)
|
||||
.unwrap();
|
||||
let state = manager.get_migration_status("test-tenant").unwrap();
|
||||
assert_eq!(state.progress, 50);
|
||||
|
||||
// Complete migration
|
||||
let state = manager.complete_migration("test-tenant").unwrap();
|
||||
assert_eq!(state.status, MigrationStatus::Completed);
|
||||
assert_eq!(state.progress, 100);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,631 @@
|
||||
//! Multi-Tenancy Module for RuVector PostgreSQL Extension
|
||||
//!
|
||||
//! Provides first-class multi-tenancy support with:
|
||||
//! - Tenant-isolated vector search (data never leaks)
|
||||
//! - Per-tenant integrity monitoring
|
||||
//! - Fair resource allocation with quotas
|
||||
//! - Row-level security integration
|
||||
//! - Multiple isolation levels (shared, partition, dedicated)
|
||||
//!
|
||||
//! # Usage
|
||||
//!
|
||||
//! ```sql
|
||||
//! -- Create a tenant
|
||||
//! SELECT ruvector_tenant_create('acme-corp', '{
|
||||
//! "display_name": "Acme Corporation",
|
||||
//! "max_vectors": 5000000,
|
||||
//! "isolation_level": "shared"
|
||||
//! }'::jsonb);
|
||||
//!
|
||||
//! -- Set tenant context for session
|
||||
//! SET ruvector.tenant_id = 'acme-corp';
|
||||
//!
|
||||
//! -- All operations are now tenant-scoped
|
||||
//! INSERT INTO embeddings (content, vec) VALUES ('doc', '[0.1, 0.2, ...]');
|
||||
//! SELECT * FROM embeddings ORDER BY vec <-> query LIMIT 10;
|
||||
//! ```
|
||||
|
||||
pub mod isolation;
|
||||
pub mod operations;
|
||||
pub mod quotas;
|
||||
pub mod registry;
|
||||
pub mod rls;
|
||||
pub mod validation;
|
||||
|
||||
// Re-export main types
|
||||
pub use isolation::{
|
||||
get_isolation_manager, IsolationError, IsolationManager, MigrationState, MigrationStatus,
|
||||
QueryRoute,
|
||||
};
|
||||
pub use operations::{
|
||||
get_tenant_stats, TenantContext, TenantStats, TenantVectorDelete, TenantVectorInsert,
|
||||
TenantVectorSearch,
|
||||
};
|
||||
pub use quotas::{get_quota_manager, QuotaManager, QuotaResult, QuotaStatus, TenantUsage};
|
||||
pub use registry::{
|
||||
get_registry, IsolationLevel, PromotionPolicy, TenantConfig, TenantError, TenantQuota,
|
||||
TenantRegistry,
|
||||
};
|
||||
pub use rls::{get_rls_manager, PolicyTemplate, RlsManager, RlsPolicyConfig};
|
||||
pub use validation::{
|
||||
escape_string_literal, quote_identifier, safe_partition_name, safe_schema_name,
|
||||
sanitize_for_identifier, validate_identifier, validate_tenant_id, ValidationError,
|
||||
};
|
||||
|
||||
use pgrx::prelude::*;
|
||||
use pgrx::JsonB;
|
||||
|
||||
// ============================================================================
|
||||
// GUC Registration for Tenant Context
|
||||
// ============================================================================
|
||||
|
||||
/// Initialize tenant-related GUCs
|
||||
/// Note: ruvector.tenant_id is registered as a custom GUC that can be set
|
||||
/// using SET ruvector.tenant_id = 'tenant-name' and read using
|
||||
/// current_setting('ruvector.tenant_id', true)
|
||||
pub fn init_tenant_gucs() {
|
||||
// The tenant_id GUC is handled as a custom variable that PostgreSQL
|
||||
// manages natively. We don't need to pre-register it - users can simply:
|
||||
// SET ruvector.tenant_id = 'my-tenant';
|
||||
// SELECT current_setting('ruvector.tenant_id', true);
|
||||
//
|
||||
// This approach is consistent with how PostgreSQL handles custom GUCs
|
||||
// and is the pattern used by other extensions for session-level context.
|
||||
//
|
||||
// To make this work securely, we rely on RLS policies that read
|
||||
// current_setting('ruvector.tenant_id', true) directly.
|
||||
|
||||
pgrx::log!("RuVector multi-tenancy initialized");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// SQL Functions - Tenant Management
|
||||
// ============================================================================
|
||||
|
||||
/// Create a new tenant
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_tenant_create('acme-corp', '{
|
||||
/// "display_name": "Acme Corporation",
|
||||
/// "max_vectors": 5000000,
|
||||
/// "max_qps": 200,
|
||||
/// "isolation_level": "shared",
|
||||
/// "integrity_enabled": true
|
||||
/// }'::jsonb);
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_tenant_create(
|
||||
tenant_id: &str,
|
||||
config: Option<JsonB>,
|
||||
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let tenant_config = match config {
|
||||
Some(JsonB(json_val)) => TenantConfig::from_json(tenant_id.to_string(), &json_val),
|
||||
None => TenantConfig::new(tenant_id.to_string()),
|
||||
};
|
||||
|
||||
let registry = get_registry();
|
||||
registry.register(tenant_config)?;
|
||||
|
||||
Ok(format!("Tenant '{}' created successfully", tenant_id))
|
||||
}
|
||||
|
||||
/// Set current tenant context for the session
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_tenant_set('acme-corp');
|
||||
/// -- All subsequent operations are now scoped to acme-corp
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_tenant_set(
|
||||
tenant_id: &str,
|
||||
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
|
||||
// Validate tenant exists and is active
|
||||
let registry = get_registry();
|
||||
let config = registry.validate_context(tenant_id)?;
|
||||
|
||||
// Set the GUC (in actual implementation)
|
||||
// For now, return success message
|
||||
Ok(format!(
|
||||
"Tenant context set to '{}' (isolation: {})",
|
||||
tenant_id,
|
||||
config.isolation_level.as_str()
|
||||
))
|
||||
}
|
||||
|
||||
/// Get statistics for a tenant
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_tenant_stats('acme-corp');
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_tenant_stats(
|
||||
tenant_id: &str,
|
||||
) -> Result<JsonB, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let stats = get_tenant_stats(tenant_id)?;
|
||||
|
||||
Ok(JsonB(serde_json::json!({
|
||||
"tenant_id": stats.tenant_id,
|
||||
"vector_count": stats.vector_count,
|
||||
"storage_bytes": stats.storage_bytes,
|
||||
"storage_gb": stats.storage_bytes as f64 / (1024.0 * 1024.0 * 1024.0),
|
||||
"collection_count": stats.collection_count,
|
||||
"isolation_level": stats.isolation_level,
|
||||
"integrity_state": stats.integrity_state,
|
||||
"lambda_cut": stats.lambda_cut,
|
||||
"is_suspended": stats.is_suspended,
|
||||
"quota_usage_percent": stats.quota_usage_percent
|
||||
})))
|
||||
}
|
||||
|
||||
/// Check quota status for a tenant
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_tenant_quota_check('acme-corp');
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_tenant_quota_check(
|
||||
tenant_id: &str,
|
||||
) -> Result<JsonB, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let status = get_quota_manager()
|
||||
.get_quota_status(tenant_id)
|
||||
.ok_or_else(|| format!("Tenant not found: {}", tenant_id))?;
|
||||
|
||||
Ok(JsonB(serde_json::json!({
|
||||
"tenant_id": status.tenant_id,
|
||||
"vectors": {
|
||||
"current": status.vectors.current,
|
||||
"limit": status.vectors.limit,
|
||||
"usage_percent": status.vectors.usage_percent
|
||||
},
|
||||
"storage": {
|
||||
"current_bytes": status.storage.current,
|
||||
"limit_bytes": status.storage.limit,
|
||||
"usage_percent": status.storage.usage_percent
|
||||
},
|
||||
"qps": {
|
||||
"current": status.qps.current,
|
||||
"limit": status.qps.limit
|
||||
},
|
||||
"concurrent_queries": {
|
||||
"current": status.concurrent.current,
|
||||
"limit": status.concurrent.limit
|
||||
},
|
||||
"collections": {
|
||||
"current": status.collections.current,
|
||||
"limit": status.collections.limit
|
||||
},
|
||||
"is_near_limit": status.is_near_limit(),
|
||||
"is_critical": status.is_critical()
|
||||
})))
|
||||
}
|
||||
|
||||
/// Suspend a tenant
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_tenant_suspend('bad-actor');
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_tenant_suspend(
|
||||
tenant_id: &str,
|
||||
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
|
||||
get_registry().suspend(tenant_id)?;
|
||||
Ok(format!("Tenant '{}' has been suspended", tenant_id))
|
||||
}
|
||||
|
||||
/// Resume a suspended tenant
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_tenant_resume('bad-actor');
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_tenant_resume(
|
||||
tenant_id: &str,
|
||||
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
|
||||
get_registry().resume(tenant_id)?;
|
||||
Ok(format!("Tenant '{}' has been resumed", tenant_id))
|
||||
}
|
||||
|
||||
/// Delete a tenant
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// -- Soft delete (marks for cleanup)
|
||||
/// SELECT ruvector_tenant_delete('churned-customer');
|
||||
///
|
||||
/// -- Hard delete (immediate)
|
||||
/// SELECT ruvector_tenant_delete('churned-customer', true);
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_tenant_delete(
|
||||
tenant_id: &str,
|
||||
hard: default!(bool, false),
|
||||
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
|
||||
get_registry().delete(tenant_id, hard)?;
|
||||
|
||||
let delete_type = if hard {
|
||||
"permanently deleted"
|
||||
} else {
|
||||
"marked for deletion"
|
||||
};
|
||||
Ok(format!("Tenant '{}' has been {}", tenant_id, delete_type))
|
||||
}
|
||||
|
||||
/// List all tenants
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT * FROM ruvector_tenants();
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_tenants() -> Result<JsonB, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let tenants = get_registry().list();
|
||||
|
||||
let tenant_list: Vec<serde_json::Value> = tenants
|
||||
.iter()
|
||||
.map(|t| {
|
||||
serde_json::json!({
|
||||
"id": t.id,
|
||||
"display_name": t.display_name,
|
||||
"isolation_level": t.isolation_level.as_str(),
|
||||
"max_vectors": t.quota.max_vectors,
|
||||
"max_qps": t.quota.max_qps,
|
||||
"integrity_enabled": t.integrity_enabled,
|
||||
"is_suspended": t.is_suspended(),
|
||||
"created_at": t.created_at
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(JsonB(serde_json::json!(tenant_list)))
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// SQL Functions - Isolation Management
|
||||
// ============================================================================
|
||||
|
||||
/// Enable tenant RLS on a table
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_enable_tenant_rls('embeddings', 'tenant_id');
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_enable_tenant_rls(
|
||||
table_name: &str,
|
||||
tenant_column: default!(&str, "'tenant_id'"),
|
||||
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let sql = get_isolation_manager().enable_shared_isolation(table_name, tenant_column)?;
|
||||
Ok(format!(
|
||||
"RLS enabled for table '{}'. Execute the following SQL:\n{}",
|
||||
table_name, sql
|
||||
))
|
||||
}
|
||||
|
||||
/// Migrate tenant to a new isolation level
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_tenant_migrate('enterprise-customer', 'dedicated');
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_tenant_migrate(
|
||||
tenant_id: &str,
|
||||
target_level: &str,
|
||||
) -> Result<JsonB, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let level = IsolationLevel::from_str(target_level)
|
||||
.ok_or_else(|| format!("Invalid isolation level: {}", target_level))?;
|
||||
|
||||
let state = get_isolation_manager().start_migration(tenant_id, level)?;
|
||||
|
||||
Ok(JsonB(serde_json::json!({
|
||||
"tenant_id": state.tenant_id,
|
||||
"from_level": state.from_level.as_str(),
|
||||
"to_level": state.to_level.as_str(),
|
||||
"status": format!("{:?}", state.status),
|
||||
"started_at": state.started_at
|
||||
})))
|
||||
}
|
||||
|
||||
/// Get migration status for a tenant
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT * FROM ruvector_tenant_migration_status('enterprise-customer');
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_tenant_migration_status(
|
||||
tenant_id: &str,
|
||||
) -> Result<JsonB, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let state = get_isolation_manager()
|
||||
.get_migration_status(tenant_id)
|
||||
.ok_or_else(|| format!("No migration in progress for tenant: {}", tenant_id))?;
|
||||
|
||||
Ok(JsonB(serde_json::json!({
|
||||
"tenant_id": state.tenant_id,
|
||||
"from_level": state.from_level.as_str(),
|
||||
"to_level": state.to_level.as_str(),
|
||||
"status": format!("{:?}", state.status),
|
||||
"progress": state.progress,
|
||||
"vectors_migrated": state.vectors_migrated,
|
||||
"total_vectors": state.total_vectors,
|
||||
"started_at": state.started_at,
|
||||
"completed_at": state.completed_at,
|
||||
"error": state.error
|
||||
})))
|
||||
}
|
||||
|
||||
/// Isolate tenant to dedicated resources
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_tenant_isolate('enterprise-customer');
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_tenant_isolate(
|
||||
tenant_id: &str,
|
||||
) -> Result<JsonB, Box<dyn std::error::Error + Send + Sync>> {
|
||||
// Create dedicated schema
|
||||
let schema_config = get_isolation_manager().create_dedicated_schema(tenant_id)?;
|
||||
let sql = get_isolation_manager().generate_schema_sql(&schema_config);
|
||||
|
||||
Ok(JsonB(serde_json::json!({
|
||||
"tenant_id": tenant_id,
|
||||
"schema_name": schema_config.schema_name,
|
||||
"sql_to_execute": sql,
|
||||
"message": "Execute the returned SQL to complete isolation"
|
||||
})))
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// SQL Functions - Policy Configuration
|
||||
// ============================================================================
|
||||
|
||||
/// Set promotion policy for automatic isolation level upgrades
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_tenant_set_policy('{
|
||||
/// "auto_promote_to_partition": 100000,
|
||||
/// "auto_promote_to_dedicated": 10000000,
|
||||
/// "check_interval": "1 hour"
|
||||
/// }'::jsonb);
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_tenant_set_policy(
|
||||
policy_config: JsonB,
|
||||
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let JsonB(json_val) = policy_config;
|
||||
|
||||
let policy = PromotionPolicy {
|
||||
partition_threshold: json_val
|
||||
.get("auto_promote_to_partition")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(100_000),
|
||||
dedicated_threshold: json_val
|
||||
.get("auto_promote_to_dedicated")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(10_000_000),
|
||||
check_interval_secs: json_val
|
||||
.get("check_interval_secs")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(3600),
|
||||
enabled: json_val
|
||||
.get("enabled")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(true),
|
||||
};
|
||||
|
||||
get_registry().set_promotion_policy(policy);
|
||||
|
||||
Ok("Promotion policy updated successfully".to_string())
|
||||
}
|
||||
|
||||
/// Update tenant quota
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_tenant_update_quota('acme-corp', '{
|
||||
/// "max_vectors": 10000000,
|
||||
/// "max_qps": 500
|
||||
/// }'::jsonb);
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_tenant_update_quota(
|
||||
tenant_id: &str,
|
||||
quota_config: JsonB,
|
||||
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let JsonB(json_val) = quota_config;
|
||||
|
||||
let mut config = get_registry()
|
||||
.get(tenant_id)
|
||||
.ok_or_else(|| format!("Tenant not found: {}", tenant_id))?;
|
||||
|
||||
if let Some(max_vec) = json_val.get("max_vectors").and_then(|v| v.as_u64()) {
|
||||
config.quota.max_vectors = max_vec;
|
||||
}
|
||||
if let Some(max_qps) = json_val.get("max_qps").and_then(|v| v.as_u64()) {
|
||||
config.quota.max_qps = max_qps as u32;
|
||||
}
|
||||
if let Some(max_storage) = json_val.get("max_storage_gb").and_then(|v| v.as_f64()) {
|
||||
config.quota.max_storage_bytes = (max_storage * 1024.0 * 1024.0 * 1024.0) as u64;
|
||||
}
|
||||
if let Some(max_concurrent) = json_val.get("max_concurrent").and_then(|v| v.as_u64()) {
|
||||
config.quota.max_concurrent = max_concurrent as u32;
|
||||
}
|
||||
if let Some(max_collections) = json_val.get("max_collections").and_then(|v| v.as_u64()) {
|
||||
config.quota.max_collections = max_collections as u32;
|
||||
}
|
||||
|
||||
get_registry().update(tenant_id, config)?;
|
||||
|
||||
Ok(format!("Quota updated for tenant '{}'", tenant_id))
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// SQL Functions - RLS Helpers
|
||||
// ============================================================================
|
||||
|
||||
/// Generate RLS setup SQL for a table
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_generate_rls_sql('embeddings', 'tenant_id');
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_generate_rls_sql(
|
||||
table_name: &str,
|
||||
tenant_column: default!(&str, "'tenant_id'"),
|
||||
) -> String {
|
||||
let config = RlsPolicyConfig::new(table_name).with_tenant_column(tenant_column);
|
||||
|
||||
get_rls_manager().generate_enable_rls_sql(&config)
|
||||
}
|
||||
|
||||
/// Generate SQL to add tenant column to a table
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_generate_tenant_column_sql('embeddings');
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_generate_tenant_column_sql(
|
||||
table_name: &str,
|
||||
column_name: default!(&str, "'tenant_id'"),
|
||||
not_null: default!(bool, true),
|
||||
auto_default: default!(bool, true),
|
||||
) -> String {
|
||||
rls::RlsManager::generate_add_tenant_column_sql(table_name, column_name, not_null, auto_default)
|
||||
}
|
||||
|
||||
/// Generate SQL to create RuVector roles
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT ruvector_generate_roles_sql();
|
||||
/// ```
|
||||
#[pg_extern]
|
||||
pub fn ruvector_generate_roles_sql() -> String {
|
||||
rls::RlsManager::generate_roles_sql()
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Tests
|
||||
// ============================================================================
|
||||
|
||||
#[cfg(feature = "pg_test")]
|
||||
#[pg_schema]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[pg_test]
|
||||
fn test_tenant_create() {
|
||||
let result = ruvector_tenant_create("test-tenant", None);
|
||||
assert!(result.is_ok());
|
||||
assert!(result.unwrap().contains("test-tenant"));
|
||||
}
|
||||
|
||||
#[pg_test]
|
||||
fn test_tenant_create_with_config() {
|
||||
let config = JsonB(serde_json::json!({
|
||||
"display_name": "Test Corp",
|
||||
"max_vectors": 5000000,
|
||||
"isolation_level": "partition"
|
||||
}));
|
||||
|
||||
let result = ruvector_tenant_create("test-tenant-2", Some(config));
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[pg_test]
|
||||
fn test_tenant_list() {
|
||||
// Create a tenant first
|
||||
let _ = ruvector_tenant_create("list-test-tenant", None);
|
||||
|
||||
let result = ruvector_tenants();
|
||||
assert!(result.is_ok());
|
||||
|
||||
let JsonB(json) = result.unwrap();
|
||||
assert!(json.is_array());
|
||||
}
|
||||
|
||||
#[pg_test]
|
||||
fn test_tenant_suspend_resume() {
|
||||
let _ = ruvector_tenant_create("suspend-test", None);
|
||||
|
||||
// Suspend
|
||||
let result = ruvector_tenant_suspend("suspend-test");
|
||||
assert!(result.is_ok());
|
||||
|
||||
// Resume
|
||||
let result = ruvector_tenant_resume("suspend-test");
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[pg_test]
|
||||
fn test_rls_sql_generation() {
|
||||
let sql = ruvector_generate_rls_sql("embeddings", "tenant_id");
|
||||
assert!(sql.contains("ENABLE ROW LEVEL SECURITY"));
|
||||
assert!(sql.contains("ruvector_tenant_isolation"));
|
||||
}
|
||||
|
||||
#[pg_test]
|
||||
fn test_tenant_column_sql_generation() {
|
||||
let sql = ruvector_generate_tenant_column_sql("embeddings", "tenant_id", true, true);
|
||||
assert!(sql.contains("ADD COLUMN"));
|
||||
assert!(sql.contains("tenant_id"));
|
||||
}
|
||||
|
||||
#[pg_test]
|
||||
fn test_roles_sql_generation() {
|
||||
let sql = ruvector_generate_roles_sql();
|
||||
assert!(sql.contains("ruvector_admin"));
|
||||
assert!(sql.contains("ruvector_users"));
|
||||
}
|
||||
|
||||
#[pg_test]
|
||||
fn test_policy_update() {
|
||||
let policy = JsonB(serde_json::json!({
|
||||
"auto_promote_to_partition": 50000,
|
||||
"auto_promote_to_dedicated": 5000000,
|
||||
"enabled": true
|
||||
}));
|
||||
|
||||
let result = ruvector_tenant_set_policy(policy);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[pg_test]
|
||||
fn test_quota_check() {
|
||||
let _ = ruvector_tenant_create("quota-test", None);
|
||||
|
||||
let result = ruvector_tenant_quota_check("quota-test");
|
||||
assert!(result.is_ok());
|
||||
|
||||
let JsonB(json) = result.unwrap();
|
||||
assert!(json.get("vectors").is_some());
|
||||
assert!(json.get("storage").is_some());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,733 @@
|
||||
//! Tenant-Aware Operations for RuVector Multi-Tenancy
|
||||
//!
|
||||
//! Wraps index operations with tenant context validation, quota enforcement,
|
||||
//! and proper routing based on isolation level.
|
||||
|
||||
use std::time::Instant;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::isolation::{get_isolation_manager, QueryRoute};
|
||||
use super::quotas::{get_quota_manager, QuotaResult};
|
||||
use super::registry::{get_registry, TenantConfig, TenantError};
|
||||
use super::validation::{escape_string_literal, validate_ip_address, validate_tenant_id};
|
||||
|
||||
/// Result of a tenant-aware operation
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum OperationResult<T> {
|
||||
/// Operation succeeded
|
||||
Success(T),
|
||||
/// Operation denied due to quota
|
||||
QuotaDenied(QuotaResult),
|
||||
/// Operation denied due to tenant error
|
||||
TenantError(TenantError),
|
||||
/// Operation failed with error
|
||||
Error(String),
|
||||
}
|
||||
|
||||
impl<T> OperationResult<T> {
|
||||
/// Check if operation succeeded
|
||||
pub fn is_success(&self) -> bool {
|
||||
matches!(self, Self::Success(_))
|
||||
}
|
||||
|
||||
/// Get success value or panic
|
||||
pub fn unwrap(self) -> T {
|
||||
match self {
|
||||
Self::Success(v) => v,
|
||||
Self::QuotaDenied(q) => panic!("Quota denied: {:?}", q),
|
||||
Self::TenantError(e) => panic!("Tenant error: {}", e),
|
||||
Self::Error(e) => panic!("Operation error: {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get success value or return error message
|
||||
pub fn into_result(self) -> Result<T, String> {
|
||||
match self {
|
||||
Self::Success(v) => Ok(v),
|
||||
Self::QuotaDenied(q) => Err(q
|
||||
.error_message()
|
||||
.unwrap_or_else(|| "Quota denied".to_string())),
|
||||
Self::TenantError(e) => Err(e.to_string()),
|
||||
Self::Error(e) => Err(e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tenant context for operations
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TenantContext {
|
||||
/// Tenant ID (validated)
|
||||
pub tenant_id: String,
|
||||
/// Tenant configuration
|
||||
pub config: TenantConfig,
|
||||
/// Query routing information
|
||||
pub route: QueryRoute,
|
||||
/// Whether this is an admin context
|
||||
pub is_admin: bool,
|
||||
}
|
||||
|
||||
/// Represents a validated tenant ID
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ValidatedTenantId(String);
|
||||
|
||||
impl ValidatedTenantId {
|
||||
/// Create a new validated tenant ID
|
||||
pub fn new(tenant_id: &str) -> Result<Self, TenantError> {
|
||||
validate_tenant_id(tenant_id).map_err(|e| TenantError::InvalidId(format!("{}", e)))?;
|
||||
Ok(Self(tenant_id.to_string()))
|
||||
}
|
||||
|
||||
/// Get the tenant ID as a string
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl TenantContext {
|
||||
/// Get current tenant context from GUC
|
||||
pub fn current() -> Result<Self, TenantError> {
|
||||
// Get tenant_id from PostgreSQL GUC
|
||||
let tenant_id = get_current_tenant_id()?;
|
||||
|
||||
// Special handling for admin wildcard
|
||||
if tenant_id == "*" {
|
||||
return Ok(Self {
|
||||
tenant_id: "*".to_string(),
|
||||
config: TenantConfig::new("*".to_string()),
|
||||
route: QueryRoute::SharedWithFilter {
|
||||
table: "".to_string(),
|
||||
filter: "true".to_string(), // No filter for admin
|
||||
tenant_param: None, // Admin doesn't need tenant param
|
||||
},
|
||||
is_admin: true,
|
||||
});
|
||||
}
|
||||
|
||||
// Validate tenant context
|
||||
let config = get_registry().validate_context(&tenant_id)?;
|
||||
|
||||
// Get routing for this tenant
|
||||
let route = get_isolation_manager().route_query(&tenant_id, "embeddings");
|
||||
|
||||
Ok(Self {
|
||||
tenant_id,
|
||||
config,
|
||||
route,
|
||||
is_admin: false,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create context for a specific tenant (bypassing GUC)
|
||||
pub fn for_tenant(tenant_id: &str) -> Result<Self, TenantError> {
|
||||
let config = get_registry().validate_context(tenant_id)?;
|
||||
let route = get_isolation_manager().route_query(tenant_id, "embeddings");
|
||||
|
||||
Ok(Self {
|
||||
tenant_id: tenant_id.to_string(),
|
||||
config,
|
||||
route,
|
||||
is_admin: false,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get table reference for SQL queries
|
||||
pub fn table_ref(&self, base_table: &str) -> String {
|
||||
let route = get_isolation_manager().route_query(&self.tenant_id, base_table);
|
||||
route.table_reference()
|
||||
}
|
||||
|
||||
/// Get WHERE clause for tenant filtering (if needed)
|
||||
pub fn where_clause(&self, base_table: &str) -> Option<String> {
|
||||
let route = get_isolation_manager().route_query(&self.tenant_id, base_table);
|
||||
route.where_clause()
|
||||
}
|
||||
}
|
||||
|
||||
/// Get current tenant ID from PostgreSQL GUC
|
||||
pub fn get_current_tenant_id() -> Result<String, TenantError> {
|
||||
// In actual pgrx implementation, this would use:
|
||||
// Spi::get_one::<String>("SELECT current_setting('ruvector.tenant_id', true)")
|
||||
|
||||
// For now, provide a placeholder that would be replaced with actual GUC access
|
||||
#[cfg(feature = "pg_test")]
|
||||
{
|
||||
// In tests, use a thread-local for testing
|
||||
thread_local! {
|
||||
static MOCK_TENANT_ID: std::cell::RefCell<String> = std::cell::RefCell::new(String::new());
|
||||
}
|
||||
|
||||
MOCK_TENANT_ID.with(|id| {
|
||||
let tenant_id = id.borrow().clone();
|
||||
if tenant_id.is_empty() {
|
||||
Err(TenantError::NoContext)
|
||||
} else {
|
||||
Ok(tenant_id)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "pg_test"))]
|
||||
{
|
||||
// Actual PostgreSQL GUC access would go here
|
||||
// This is a placeholder for the actual implementation
|
||||
Err(TenantError::NoContext)
|
||||
}
|
||||
}
|
||||
|
||||
/// Set mock tenant ID for testing
|
||||
#[cfg(feature = "pg_test")]
|
||||
pub fn set_mock_tenant_id(tenant_id: &str) {
|
||||
thread_local! {
|
||||
static MOCK_TENANT_ID: std::cell::RefCell<String> = std::cell::RefCell::new(String::new());
|
||||
}
|
||||
|
||||
MOCK_TENANT_ID.with(|id| {
|
||||
*id.borrow_mut() = tenant_id.to_string();
|
||||
});
|
||||
}
|
||||
|
||||
/// Tenant-aware vector insert operation
|
||||
pub struct TenantVectorInsert<'a> {
|
||||
ctx: &'a TenantContext,
|
||||
vectors: Vec<(Vec<f32>, Option<serde_json::Value>)>,
|
||||
table_name: String,
|
||||
estimated_bytes_per_vector: usize,
|
||||
}
|
||||
|
||||
impl<'a> TenantVectorInsert<'a> {
|
||||
/// Create a new tenant-aware insert
|
||||
pub fn new(ctx: &'a TenantContext, table_name: &str) -> Self {
|
||||
Self {
|
||||
ctx,
|
||||
vectors: Vec::new(),
|
||||
table_name: table_name.to_string(),
|
||||
estimated_bytes_per_vector: 4 * 1536 + 100, // Default for 1536-dim + metadata
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a vector to insert
|
||||
pub fn add(&mut self, vector: Vec<f32>, metadata: Option<serde_json::Value>) -> &mut Self {
|
||||
self.vectors.push((vector, metadata));
|
||||
self
|
||||
}
|
||||
|
||||
/// Add multiple vectors
|
||||
pub fn add_batch(&mut self, vectors: Vec<(Vec<f32>, Option<serde_json::Value>)>) -> &mut Self {
|
||||
self.vectors.extend(vectors);
|
||||
self
|
||||
}
|
||||
|
||||
/// Execute the insert with quota enforcement
|
||||
pub fn execute(self) -> OperationResult<InsertResult> {
|
||||
let quota_manager = get_quota_manager();
|
||||
|
||||
// Calculate estimated bytes
|
||||
let total_bytes = self.vectors.len() as u64 * self.estimated_bytes_per_vector as u64;
|
||||
|
||||
// Check quota before insert
|
||||
let quota_check = quota_manager.check_vector_insert(
|
||||
&self.ctx.tenant_id,
|
||||
self.vectors.len() as u64,
|
||||
total_bytes,
|
||||
);
|
||||
|
||||
if !quota_check.is_allowed() {
|
||||
return OperationResult::QuotaDenied(quota_check);
|
||||
}
|
||||
|
||||
// Get proper table reference
|
||||
let table_ref = self.ctx.table_ref(&self.table_name);
|
||||
|
||||
// Execute insert (placeholder - actual implementation would use SPI)
|
||||
let start = Instant::now();
|
||||
let inserted_count = self.vectors.len();
|
||||
|
||||
// Record successful insert
|
||||
quota_manager.record_vector_insert(&self.ctx.tenant_id, inserted_count as u64, total_bytes);
|
||||
|
||||
OperationResult::Success(InsertResult {
|
||||
inserted_count,
|
||||
table_used: table_ref,
|
||||
duration_ms: start.elapsed().as_millis() as u64,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of an insert operation
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct InsertResult {
|
||||
/// Number of vectors inserted
|
||||
pub inserted_count: usize,
|
||||
/// Table that was used
|
||||
pub table_used: String,
|
||||
/// Duration in milliseconds
|
||||
pub duration_ms: u64,
|
||||
}
|
||||
|
||||
/// Tenant-aware vector search operation
|
||||
pub struct TenantVectorSearch<'a> {
|
||||
ctx: &'a TenantContext,
|
||||
query: Vec<f32>,
|
||||
k: usize,
|
||||
table_name: String,
|
||||
ef_search: Option<usize>,
|
||||
filter: Option<String>,
|
||||
}
|
||||
|
||||
impl<'a> TenantVectorSearch<'a> {
|
||||
/// Create a new tenant-aware search
|
||||
pub fn new(ctx: &'a TenantContext, query: Vec<f32>, k: usize, table_name: &str) -> Self {
|
||||
Self {
|
||||
ctx,
|
||||
query,
|
||||
k,
|
||||
table_name: table_name.to_string(),
|
||||
ef_search: None,
|
||||
filter: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set ef_search parameter
|
||||
pub fn with_ef_search(mut self, ef: usize) -> Self {
|
||||
self.ef_search = Some(ef);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add additional WHERE filter
|
||||
pub fn with_filter(mut self, filter: &str) -> Self {
|
||||
self.filter = Some(filter.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Execute the search with rate limiting
|
||||
pub fn execute(self) -> OperationResult<SearchResult> {
|
||||
let quota_manager = get_quota_manager();
|
||||
|
||||
// Check rate limit
|
||||
let rate_check = quota_manager.check_query(&self.ctx.tenant_id);
|
||||
if !rate_check.is_allowed() {
|
||||
return OperationResult::QuotaDenied(rate_check);
|
||||
}
|
||||
|
||||
// Start concurrent query tracking
|
||||
quota_manager.start_query(&self.ctx.tenant_id);
|
||||
|
||||
let start = Instant::now();
|
||||
|
||||
// Get proper table reference and filters
|
||||
let table_ref = self.ctx.table_ref(&self.table_name);
|
||||
let tenant_filter = self.ctx.where_clause(&self.table_name);
|
||||
|
||||
// Combine filters
|
||||
let combined_filter = match (&tenant_filter, &self.filter) {
|
||||
(Some(tf), Some(f)) => Some(format!("({}) AND ({})", tf, f)),
|
||||
(Some(tf), None) => Some(tf.clone()),
|
||||
(None, Some(f)) => Some(f.clone()),
|
||||
(None, None) => None,
|
||||
};
|
||||
|
||||
// Execute search (placeholder - actual implementation would use SPI)
|
||||
let results: Vec<(i64, f32)> = Vec::new(); // Would be populated by actual search
|
||||
|
||||
// End concurrent query tracking
|
||||
quota_manager.end_query(&self.ctx.tenant_id);
|
||||
|
||||
OperationResult::Success(SearchResult {
|
||||
results,
|
||||
k: self.k,
|
||||
table_used: table_ref,
|
||||
filter_applied: combined_filter,
|
||||
duration_ms: start.elapsed().as_millis() as u64,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of a search operation
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SearchResult {
|
||||
/// Search results (id, distance)
|
||||
pub results: Vec<(i64, f32)>,
|
||||
/// Requested k
|
||||
pub k: usize,
|
||||
/// Table that was searched
|
||||
pub table_used: String,
|
||||
/// Filter that was applied
|
||||
pub filter_applied: Option<String>,
|
||||
/// Duration in milliseconds
|
||||
pub duration_ms: u64,
|
||||
}
|
||||
|
||||
/// Tenant-aware delete operation
|
||||
pub struct TenantVectorDelete<'a> {
|
||||
ctx: &'a TenantContext,
|
||||
ids: Vec<i64>,
|
||||
table_name: String,
|
||||
}
|
||||
|
||||
impl<'a> TenantVectorDelete<'a> {
|
||||
/// Create a new tenant-aware delete
|
||||
pub fn new(ctx: &'a TenantContext, ids: Vec<i64>, table_name: &str) -> Self {
|
||||
Self {
|
||||
ctx,
|
||||
ids,
|
||||
table_name: table_name.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute the delete with quota tracking
|
||||
pub fn execute(self) -> OperationResult<DeleteResult> {
|
||||
let quota_manager = get_quota_manager();
|
||||
|
||||
let start = Instant::now();
|
||||
let table_ref = self.ctx.table_ref(&self.table_name);
|
||||
|
||||
// Execute delete (placeholder - actual implementation would use SPI)
|
||||
let deleted_count = self.ids.len();
|
||||
let deleted_bytes = (deleted_count * 4 * 1536) as u64; // Estimate
|
||||
|
||||
// Record deletion in quota manager
|
||||
quota_manager.record_vector_delete(
|
||||
&self.ctx.tenant_id,
|
||||
deleted_count as u64,
|
||||
deleted_bytes,
|
||||
);
|
||||
|
||||
OperationResult::Success(DeleteResult {
|
||||
deleted_count,
|
||||
table_used: table_ref,
|
||||
duration_ms: start.elapsed().as_millis() as u64,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of a delete operation
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DeleteResult {
|
||||
/// Number of vectors deleted
|
||||
pub deleted_count: usize,
|
||||
/// Table that was used
|
||||
pub table_used: String,
|
||||
/// Duration in milliseconds
|
||||
pub duration_ms: u64,
|
||||
}
|
||||
|
||||
/// Statistics for a tenant
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TenantStats {
|
||||
/// Tenant ID
|
||||
pub tenant_id: String,
|
||||
/// Vector count
|
||||
pub vector_count: u64,
|
||||
/// Storage used in bytes
|
||||
pub storage_bytes: u64,
|
||||
/// Collection count
|
||||
pub collection_count: u32,
|
||||
/// Isolation level
|
||||
pub isolation_level: String,
|
||||
/// Integrity state
|
||||
pub integrity_state: String,
|
||||
/// Lambda cut value
|
||||
pub lambda_cut: f32,
|
||||
/// Is suspended
|
||||
pub is_suspended: bool,
|
||||
/// Quota usage percentage
|
||||
pub quota_usage_percent: f32,
|
||||
}
|
||||
|
||||
/// Get comprehensive statistics for a tenant
|
||||
pub fn get_tenant_stats(tenant_id: &str) -> Result<TenantStats, TenantError> {
|
||||
let config = get_registry()
|
||||
.get(tenant_id)
|
||||
.ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
|
||||
|
||||
let usage = get_quota_manager().get_usage(tenant_id).unwrap_or_default();
|
||||
|
||||
let shared_state = get_registry().get_shared_state(tenant_id);
|
||||
|
||||
let (integrity_state, lambda_cut) = match shared_state {
|
||||
Some(state) => {
|
||||
let integrity = match state
|
||||
.integrity_state
|
||||
.load(std::sync::atomic::Ordering::Relaxed)
|
||||
{
|
||||
0 => "normal",
|
||||
1 => "stress",
|
||||
2 => "critical",
|
||||
_ => "unknown",
|
||||
};
|
||||
(integrity.to_string(), state.lambda_cut())
|
||||
}
|
||||
None => ("unknown".to_string(), 1.0),
|
||||
};
|
||||
|
||||
Ok(TenantStats {
|
||||
tenant_id: tenant_id.to_string(),
|
||||
vector_count: usage.vector_count,
|
||||
storage_bytes: usage.storage_bytes,
|
||||
collection_count: usage.collection_count,
|
||||
isolation_level: config.isolation_level.as_str().to_string(),
|
||||
integrity_state,
|
||||
lambda_cut,
|
||||
is_suspended: config.is_suspended(),
|
||||
quota_usage_percent: (usage.vector_count as f64 / config.quota.max_vectors as f64 * 100.0)
|
||||
as f32,
|
||||
})
|
||||
}
|
||||
|
||||
/// Audit log entry for tenant operations
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AuditLogEntry {
|
||||
/// Tenant ID
|
||||
pub tenant_id: String,
|
||||
/// Operation type
|
||||
pub operation: String,
|
||||
/// User ID (from application context)
|
||||
pub user_id: Option<String>,
|
||||
/// Details about the operation
|
||||
pub details: serde_json::Value,
|
||||
/// Timestamp
|
||||
pub timestamp: i64,
|
||||
/// IP address (if available)
|
||||
pub ip_address: Option<String>,
|
||||
/// Success status
|
||||
pub success: bool,
|
||||
/// Error message if failed
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
impl AuditLogEntry {
|
||||
/// Create a new audit log entry
|
||||
pub fn new(tenant_id: &str, operation: &str) -> Self {
|
||||
Self {
|
||||
tenant_id: tenant_id.to_string(),
|
||||
operation: operation.to_string(),
|
||||
user_id: None,
|
||||
details: serde_json::json!({}),
|
||||
timestamp: chrono_now_millis(),
|
||||
ip_address: None,
|
||||
success: true,
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set user ID
|
||||
pub fn with_user(mut self, user_id: &str) -> Self {
|
||||
self.user_id = Some(user_id.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Set details
|
||||
pub fn with_details(mut self, details: serde_json::Value) -> Self {
|
||||
self.details = details;
|
||||
self
|
||||
}
|
||||
|
||||
/// Mark as failed
|
||||
pub fn failed(mut self, error: &str) -> Self {
|
||||
self.success = false;
|
||||
self.error = Some(error.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Generate SQL to insert this audit entry (parameterized version)
|
||||
///
|
||||
/// Returns the SQL with $1-$7 placeholders and the parameter values to bind.
|
||||
/// This prevents SQL injection by using parameterized queries.
|
||||
pub fn insert_sql_parameterized(&self) -> (String, Vec<Option<String>>) {
|
||||
let sql = r#"
|
||||
INSERT INTO ruvector.tenant_audit_log (tenant_id, operation, user_id, details, ip_address, success, error)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||
"#.to_string();
|
||||
|
||||
let params = vec![
|
||||
Some(self.tenant_id.clone()),
|
||||
Some(self.operation.clone()),
|
||||
self.user_id.clone(),
|
||||
Some(serde_json::to_string(&self.details).unwrap_or_else(|_| "{}".to_string())),
|
||||
// Only include IP if it's a valid IP address (defense in depth)
|
||||
self.ip_address.as_ref().and_then(|ip| {
|
||||
if validate_ip_address(ip) {
|
||||
Some(ip.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}),
|
||||
Some(self.success.to_string()),
|
||||
self.error.clone(),
|
||||
];
|
||||
|
||||
(sql, params)
|
||||
}
|
||||
|
||||
/// Generate SQL to insert this audit entry (legacy - properly escaped)
|
||||
///
|
||||
/// Note: Prefer `insert_sql_parameterized()` for new code.
|
||||
/// This method properly escapes all values but parameterized queries are safer.
|
||||
pub fn insert_sql(&self) -> String {
|
||||
// Validate tenant_id format
|
||||
if validate_tenant_id(&self.tenant_id).is_err() {
|
||||
// Log the attempt but don't execute with invalid tenant_id
|
||||
return "SELECT 1 WHERE false".to_string(); // No-op query
|
||||
}
|
||||
|
||||
// Escape all string values
|
||||
let escaped_tenant_id = escape_string_literal(&self.tenant_id);
|
||||
let escaped_operation = escape_string_literal(&self.operation);
|
||||
let escaped_user_id = self
|
||||
.user_id
|
||||
.as_ref()
|
||||
.map(|u| format!("'{}'", escape_string_literal(u)))
|
||||
.unwrap_or_else(|| "NULL".to_string());
|
||||
let escaped_details = escape_string_literal(
|
||||
&serde_json::to_string(&self.details).unwrap_or_else(|_| "{}".to_string()),
|
||||
);
|
||||
let escaped_ip = self
|
||||
.ip_address
|
||||
.as_ref()
|
||||
.and_then(|ip| {
|
||||
// Only include if valid IP format
|
||||
if validate_ip_address(ip) {
|
||||
Some(format!("'{}'", escape_string_literal(ip)))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.unwrap_or_else(|| "NULL".to_string());
|
||||
let escaped_error = self
|
||||
.error
|
||||
.as_ref()
|
||||
.map(|e| format!("'{}'", escape_string_literal(e)))
|
||||
.unwrap_or_else(|| "NULL".to_string());
|
||||
|
||||
format!(
|
||||
r#"
|
||||
INSERT INTO ruvector.tenant_audit_log (tenant_id, operation, user_id, details, ip_address, success, error)
|
||||
VALUES ('{}', '{}', {}, '{}', {}, {}, {})
|
||||
"#,
|
||||
escaped_tenant_id,
|
||||
escaped_operation,
|
||||
escaped_user_id,
|
||||
escaped_details,
|
||||
escaped_ip,
|
||||
self.success,
|
||||
escaped_error
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
fn chrono_now_millis() -> i64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as i64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Cross-tenant prevention check
|
||||
pub fn validate_cross_tenant(
|
||||
context_tenant: &str,
|
||||
request_tenant: Option<&str>,
|
||||
) -> Result<(), TenantError> {
|
||||
if let Some(req_tenant) = request_tenant {
|
||||
if req_tenant != context_tenant && context_tenant != "*" {
|
||||
// Log security event
|
||||
let entry = AuditLogEntry::new(context_tenant, "cross_tenant_attempt")
|
||||
.with_details(serde_json::json!({
|
||||
"requested_tenant": req_tenant,
|
||||
"context_tenant": context_tenant
|
||||
}))
|
||||
.failed("Cross-tenant access denied");
|
||||
|
||||
// Would log to audit table here
|
||||
|
||||
return Err(TenantError::TenantMismatch {
|
||||
context: context_tenant.to_string(),
|
||||
request: req_tenant.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::registry::TenantConfig;
|
||||
use super::*;
|
||||
|
||||
fn setup_test_tenant(id: &str) {
|
||||
let registry = get_registry();
|
||||
let config = TenantConfig::new(id.to_string());
|
||||
let _ = registry.register(config);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_operation_result() {
|
||||
let success: OperationResult<i32> = OperationResult::Success(42);
|
||||
assert!(success.is_success());
|
||||
assert_eq!(success.unwrap(), 42);
|
||||
|
||||
let denied = OperationResult::<i32>::QuotaDenied(QuotaResult::RateLimited {
|
||||
retry_after_ms: 100,
|
||||
});
|
||||
assert!(!denied.is_success());
|
||||
assert!(denied.into_result().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cross_tenant_validation() {
|
||||
// Same tenant should pass
|
||||
assert!(validate_cross_tenant("tenant-a", Some("tenant-a")).is_ok());
|
||||
|
||||
// Different tenant should fail
|
||||
assert!(validate_cross_tenant("tenant-a", Some("tenant-b")).is_err());
|
||||
|
||||
// Admin wildcard should pass
|
||||
assert!(validate_cross_tenant("*", Some("tenant-b")).is_ok());
|
||||
|
||||
// No request tenant should pass
|
||||
assert!(validate_cross_tenant("tenant-a", None).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_audit_log_entry() {
|
||||
let entry = AuditLogEntry::new("acme-corp", "vector_insert")
|
||||
.with_user("user123")
|
||||
.with_details(serde_json::json!({"count": 100}));
|
||||
|
||||
assert_eq!(entry.tenant_id, "acme-corp");
|
||||
assert_eq!(entry.operation, "vector_insert");
|
||||
assert!(entry.success);
|
||||
|
||||
let failed_entry =
|
||||
AuditLogEntry::new("acme-corp", "vector_insert").failed("Quota exceeded");
|
||||
|
||||
assert!(!failed_entry.success);
|
||||
assert!(failed_entry.error.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_insert_result() {
|
||||
let result = InsertResult {
|
||||
inserted_count: 100,
|
||||
table_used: "embeddings".to_string(),
|
||||
duration_ms: 50,
|
||||
};
|
||||
|
||||
assert_eq!(result.inserted_count, 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_search_result() {
|
||||
let result = SearchResult {
|
||||
results: vec![(1, 0.1), (2, 0.2)],
|
||||
k: 10,
|
||||
table_used: "embeddings".to_string(),
|
||||
filter_applied: Some("category = 'test'".to_string()),
|
||||
duration_ms: 25,
|
||||
};
|
||||
|
||||
assert_eq!(result.results.len(), 2);
|
||||
assert!(result.filter_applied.is_some());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,683 @@
|
||||
//! Resource Quotas for RuVector Multi-Tenancy
|
||||
//!
|
||||
//! Provides per-tenant resource limits and enforcement:
|
||||
//! - Vector count limits
|
||||
//! - Storage limits (bytes)
|
||||
//! - Query rate limiting (QPS)
|
||||
//! - Concurrent query limits
|
||||
//! - Background worker allocation
|
||||
|
||||
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
|
||||
|
||||
use dashmap::DashMap;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::registry::{get_registry, TenantQuota};
|
||||
|
||||
/// Current resource usage for a tenant
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct TenantUsage {
|
||||
/// Current vector count
|
||||
pub vector_count: u64,
|
||||
/// Current storage in bytes
|
||||
pub storage_bytes: u64,
|
||||
/// Queries in the current rate window
|
||||
pub queries_this_second: u32,
|
||||
/// Current concurrent queries
|
||||
pub concurrent_queries: u32,
|
||||
/// Collections count
|
||||
pub collection_count: u32,
|
||||
/// Last updated timestamp (epoch millis)
|
||||
pub last_updated: i64,
|
||||
}
|
||||
|
||||
impl TenantUsage {
|
||||
/// Create new empty usage
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
vector_count: 0,
|
||||
storage_bytes: 0,
|
||||
queries_this_second: 0,
|
||||
concurrent_queries: 0,
|
||||
collection_count: 0,
|
||||
last_updated: chrono_now_millis(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Calculate storage in GB
|
||||
pub fn storage_gb(&self) -> f64 {
|
||||
self.storage_bytes as f64 / (1024.0 * 1024.0 * 1024.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Atomic usage tracking for a tenant
|
||||
#[repr(C)]
|
||||
pub struct AtomicTenantUsage {
|
||||
/// Current vector count
|
||||
pub vector_count: AtomicU64,
|
||||
/// Current storage in bytes
|
||||
pub storage_bytes: AtomicU64,
|
||||
/// Rate limiting: request count in current window
|
||||
pub rate_count: AtomicU32,
|
||||
/// Rate limiting: window start (epoch seconds)
|
||||
pub rate_window_start: AtomicU64,
|
||||
/// Concurrent query count
|
||||
pub concurrent_count: AtomicU32,
|
||||
/// Collection count
|
||||
pub collection_count: AtomicU32,
|
||||
}
|
||||
|
||||
impl AtomicTenantUsage {
|
||||
/// Create new atomic usage tracker
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
vector_count: AtomicU64::new(0),
|
||||
storage_bytes: AtomicU64::new(0),
|
||||
rate_count: AtomicU32::new(0),
|
||||
rate_window_start: AtomicU64::new(0),
|
||||
concurrent_count: AtomicU32::new(0),
|
||||
collection_count: AtomicU32::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get snapshot of usage
|
||||
pub fn snapshot(&self) -> TenantUsage {
|
||||
TenantUsage {
|
||||
vector_count: self.vector_count.load(Ordering::Relaxed),
|
||||
storage_bytes: self.storage_bytes.load(Ordering::Relaxed),
|
||||
queries_this_second: self.rate_count.load(Ordering::Relaxed),
|
||||
concurrent_queries: self.concurrent_count.load(Ordering::Relaxed),
|
||||
collection_count: self.collection_count.load(Ordering::Relaxed),
|
||||
last_updated: chrono_now_millis(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Reset from TenantUsage (for initialization from stored data)
|
||||
pub fn reset_from(&self, usage: &TenantUsage) {
|
||||
self.vector_count
|
||||
.store(usage.vector_count, Ordering::Relaxed);
|
||||
self.storage_bytes
|
||||
.store(usage.storage_bytes, Ordering::Relaxed);
|
||||
self.collection_count
|
||||
.store(usage.collection_count, Ordering::Relaxed);
|
||||
// Rate limiting and concurrent are not persisted
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for AtomicTenantUsage {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Token bucket rate limiter
|
||||
pub struct TokenBucket {
|
||||
/// Maximum tokens (burst capacity)
|
||||
capacity: u32,
|
||||
/// Tokens per second (refill rate)
|
||||
rate: u32,
|
||||
/// Current available tokens (fixed-point * 1000)
|
||||
tokens: AtomicU64,
|
||||
/// Last refill time (epoch millis)
|
||||
last_refill: AtomicU64,
|
||||
}
|
||||
|
||||
impl TokenBucket {
|
||||
/// Create a new token bucket
|
||||
pub fn new(capacity: u32, rate: u32) -> Self {
|
||||
Self {
|
||||
capacity,
|
||||
rate,
|
||||
tokens: AtomicU64::new((capacity as u64) * 1000), // Full bucket, fixed-point
|
||||
last_refill: AtomicU64::new(chrono_now_millis() as u64),
|
||||
}
|
||||
}
|
||||
|
||||
/// Try to acquire tokens
|
||||
pub fn try_acquire(&self, tokens: u32) -> bool {
|
||||
self.refill();
|
||||
|
||||
let tokens_needed = (tokens as u64) * 1000;
|
||||
let current = self.tokens.load(Ordering::Relaxed);
|
||||
|
||||
if current >= tokens_needed {
|
||||
// CAS loop for thread safety
|
||||
match self.tokens.compare_exchange(
|
||||
current,
|
||||
current - tokens_needed,
|
||||
Ordering::AcqRel,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => true,
|
||||
Err(_) => {
|
||||
// Retry with updated value
|
||||
let new_current = self.tokens.load(Ordering::Relaxed);
|
||||
if new_current >= tokens_needed {
|
||||
self.tokens.fetch_sub(tokens_needed, Ordering::Relaxed);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Refill tokens based on elapsed time
|
||||
fn refill(&self) {
|
||||
let now = chrono_now_millis() as u64;
|
||||
let last = self.last_refill.load(Ordering::Relaxed);
|
||||
let elapsed_ms = now.saturating_sub(last);
|
||||
|
||||
if elapsed_ms == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
// Calculate tokens to add (tokens per second * elapsed seconds)
|
||||
let tokens_to_add = (self.rate as u64 * 1000 * elapsed_ms) / 1000;
|
||||
|
||||
if tokens_to_add > 0 {
|
||||
let max_tokens = (self.capacity as u64) * 1000;
|
||||
let current = self.tokens.load(Ordering::Relaxed);
|
||||
let new_tokens = (current + tokens_to_add).min(max_tokens);
|
||||
|
||||
self.tokens.store(new_tokens, Ordering::Relaxed);
|
||||
self.last_refill.store(now, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
/// Time until tokens become available (milliseconds)
|
||||
pub fn time_to_available(&self, tokens: u32) -> u64 {
|
||||
self.refill();
|
||||
|
||||
let tokens_needed = (tokens as u64) * 1000;
|
||||
let current = self.tokens.load(Ordering::Relaxed);
|
||||
|
||||
if current >= tokens_needed {
|
||||
return 0;
|
||||
}
|
||||
|
||||
let tokens_short = tokens_needed - current;
|
||||
let rate_per_ms = (self.rate as u64 * 1000) / 1000;
|
||||
|
||||
if rate_per_ms == 0 {
|
||||
return u64::MAX;
|
||||
}
|
||||
|
||||
(tokens_short + rate_per_ms - 1) / rate_per_ms
|
||||
}
|
||||
}
|
||||
|
||||
/// Quota enforcement result
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum QuotaResult {
|
||||
/// Operation allowed
|
||||
Allowed,
|
||||
/// Rate limit exceeded
|
||||
RateLimited {
|
||||
/// Retry after this many milliseconds
|
||||
retry_after_ms: u64,
|
||||
},
|
||||
/// Vector quota exceeded
|
||||
VectorQuotaExceeded { current: u64, limit: u64 },
|
||||
/// Storage quota exceeded
|
||||
StorageQuotaExceeded {
|
||||
current_bytes: u64,
|
||||
limit_bytes: u64,
|
||||
},
|
||||
/// Concurrent query limit exceeded
|
||||
ConcurrentLimitExceeded { current: u32, limit: u32 },
|
||||
/// Collection limit exceeded
|
||||
CollectionLimitExceeded { current: u32, limit: u32 },
|
||||
}
|
||||
|
||||
impl QuotaResult {
|
||||
/// Check if operation is allowed
|
||||
pub fn is_allowed(&self) -> bool {
|
||||
matches!(self, Self::Allowed)
|
||||
}
|
||||
|
||||
/// Get error message if not allowed
|
||||
pub fn error_message(&self) -> Option<String> {
|
||||
match self {
|
||||
Self::Allowed => None,
|
||||
Self::RateLimited { retry_after_ms } => Some(format!(
|
||||
"Rate limit exceeded. Retry after {}ms",
|
||||
retry_after_ms
|
||||
)),
|
||||
Self::VectorQuotaExceeded { current, limit } => Some(format!(
|
||||
"Vector quota exceeded: {} / {} vectors",
|
||||
current, limit
|
||||
)),
|
||||
Self::StorageQuotaExceeded {
|
||||
current_bytes,
|
||||
limit_bytes,
|
||||
} => {
|
||||
let current_gb = *current_bytes as f64 / (1024.0 * 1024.0 * 1024.0);
|
||||
let limit_gb = *limit_bytes as f64 / (1024.0 * 1024.0 * 1024.0);
|
||||
Some(format!(
|
||||
"Storage quota exceeded: {:.2} / {:.2} GB",
|
||||
current_gb, limit_gb
|
||||
))
|
||||
}
|
||||
Self::ConcurrentLimitExceeded { current, limit } => Some(format!(
|
||||
"Concurrent query limit exceeded: {} / {}",
|
||||
current, limit
|
||||
)),
|
||||
Self::CollectionLimitExceeded { current, limit } => Some(format!(
|
||||
"Collection limit exceeded: {} / {}",
|
||||
current, limit
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Quota manager for all tenants
|
||||
pub struct QuotaManager {
|
||||
/// Atomic usage tracking per tenant
|
||||
usage: DashMap<String, AtomicTenantUsage>,
|
||||
/// Rate limiters per tenant
|
||||
rate_limiters: DashMap<String, TokenBucket>,
|
||||
}
|
||||
|
||||
impl QuotaManager {
|
||||
/// Create a new quota manager
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
usage: DashMap::new(),
|
||||
rate_limiters: DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get or create usage tracker for tenant
|
||||
fn get_or_create_usage(&self, tenant_id: &str) -> &AtomicTenantUsage {
|
||||
if !self.usage.contains_key(tenant_id) {
|
||||
self.usage
|
||||
.insert(tenant_id.to_string(), AtomicTenantUsage::new());
|
||||
}
|
||||
// Safe because we just inserted if not present
|
||||
// Use leak to get 'static reference - in production would use proper lifetime management
|
||||
unsafe {
|
||||
let ptr = self.usage.get(tenant_id).unwrap();
|
||||
&*(ptr.value() as *const AtomicTenantUsage)
|
||||
}
|
||||
}
|
||||
|
||||
/// Get or create rate limiter for tenant
|
||||
fn get_or_create_rate_limiter(&self, tenant_id: &str, quota: &TenantQuota) -> &TokenBucket {
|
||||
if !self.rate_limiters.contains_key(tenant_id) {
|
||||
// Burst capacity = 2x QPS, rate = QPS
|
||||
let bucket = TokenBucket::new(quota.max_qps * 2, quota.max_qps);
|
||||
self.rate_limiters.insert(tenant_id.to_string(), bucket);
|
||||
}
|
||||
unsafe {
|
||||
let ptr = self.rate_limiters.get(tenant_id).unwrap();
|
||||
&*(ptr.value() as *const TokenBucket)
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if vector insert is allowed
|
||||
pub fn check_vector_insert(
|
||||
&self,
|
||||
tenant_id: &str,
|
||||
count: u64,
|
||||
estimated_bytes: u64,
|
||||
) -> QuotaResult {
|
||||
// Get tenant config
|
||||
let config = match get_registry().get(tenant_id) {
|
||||
Some(c) => c,
|
||||
None => return QuotaResult::Allowed, // No quota if no tenant
|
||||
};
|
||||
|
||||
let usage = self.get_or_create_usage(tenant_id);
|
||||
|
||||
// Check vector count
|
||||
let current_vectors = usage.vector_count.load(Ordering::Relaxed);
|
||||
if current_vectors + count > config.quota.max_vectors {
|
||||
return QuotaResult::VectorQuotaExceeded {
|
||||
current: current_vectors,
|
||||
limit: config.quota.max_vectors,
|
||||
};
|
||||
}
|
||||
|
||||
// Check storage
|
||||
let current_storage = usage.storage_bytes.load(Ordering::Relaxed);
|
||||
if current_storage + estimated_bytes > config.quota.max_storage_bytes {
|
||||
return QuotaResult::StorageQuotaExceeded {
|
||||
current_bytes: current_storage,
|
||||
limit_bytes: config.quota.max_storage_bytes,
|
||||
};
|
||||
}
|
||||
|
||||
QuotaResult::Allowed
|
||||
}
|
||||
|
||||
/// Check if query is allowed (rate limiting)
|
||||
pub fn check_query(&self, tenant_id: &str) -> QuotaResult {
|
||||
// Get tenant config
|
||||
let config = match get_registry().get(tenant_id) {
|
||||
Some(c) => c,
|
||||
None => return QuotaResult::Allowed,
|
||||
};
|
||||
|
||||
// Check rate limit
|
||||
let rate_limiter = self.get_or_create_rate_limiter(tenant_id, &config.quota);
|
||||
if !rate_limiter.try_acquire(1) {
|
||||
return QuotaResult::RateLimited {
|
||||
retry_after_ms: rate_limiter.time_to_available(1),
|
||||
};
|
||||
}
|
||||
|
||||
// Check concurrent queries
|
||||
let usage = self.get_or_create_usage(tenant_id);
|
||||
let current_concurrent = usage.concurrent_count.load(Ordering::Relaxed);
|
||||
if current_concurrent >= config.quota.max_concurrent {
|
||||
return QuotaResult::ConcurrentLimitExceeded {
|
||||
current: current_concurrent,
|
||||
limit: config.quota.max_concurrent,
|
||||
};
|
||||
}
|
||||
|
||||
QuotaResult::Allowed
|
||||
}
|
||||
|
||||
/// Check if collection creation is allowed
|
||||
pub fn check_collection_create(&self, tenant_id: &str) -> QuotaResult {
|
||||
let config = match get_registry().get(tenant_id) {
|
||||
Some(c) => c,
|
||||
None => return QuotaResult::Allowed,
|
||||
};
|
||||
|
||||
let usage = self.get_or_create_usage(tenant_id);
|
||||
let current = usage.collection_count.load(Ordering::Relaxed);
|
||||
|
||||
if current >= config.quota.max_collections {
|
||||
return QuotaResult::CollectionLimitExceeded {
|
||||
current,
|
||||
limit: config.quota.max_collections,
|
||||
};
|
||||
}
|
||||
|
||||
QuotaResult::Allowed
|
||||
}
|
||||
|
||||
/// Record vector insert (after successful insert)
|
||||
pub fn record_vector_insert(&self, tenant_id: &str, count: u64, bytes: u64) {
|
||||
let usage = self.get_or_create_usage(tenant_id);
|
||||
usage.vector_count.fetch_add(count, Ordering::Relaxed);
|
||||
usage.storage_bytes.fetch_add(bytes, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Record vector delete
|
||||
pub fn record_vector_delete(&self, tenant_id: &str, count: u64, bytes: u64) {
|
||||
let usage = self.get_or_create_usage(tenant_id);
|
||||
usage.vector_count.fetch_sub(
|
||||
count.min(usage.vector_count.load(Ordering::Relaxed)),
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
usage.storage_bytes.fetch_sub(
|
||||
bytes.min(usage.storage_bytes.load(Ordering::Relaxed)),
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
}
|
||||
|
||||
/// Record collection creation
|
||||
pub fn record_collection_create(&self, tenant_id: &str) {
|
||||
let usage = self.get_or_create_usage(tenant_id);
|
||||
usage.collection_count.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Record collection deletion
|
||||
pub fn record_collection_delete(&self, tenant_id: &str) {
|
||||
let usage = self.get_or_create_usage(tenant_id);
|
||||
let current = usage.collection_count.load(Ordering::Relaxed);
|
||||
if current > 0 {
|
||||
usage.collection_count.fetch_sub(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
/// Start tracking a concurrent query
|
||||
pub fn start_query(&self, tenant_id: &str) {
|
||||
let usage = self.get_or_create_usage(tenant_id);
|
||||
usage.concurrent_count.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// End tracking a concurrent query
|
||||
pub fn end_query(&self, tenant_id: &str) {
|
||||
let usage = self.get_or_create_usage(tenant_id);
|
||||
let current = usage.concurrent_count.load(Ordering::Relaxed);
|
||||
if current > 0 {
|
||||
usage.concurrent_count.fetch_sub(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
/// Get current usage for a tenant
|
||||
pub fn get_usage(&self, tenant_id: &str) -> Option<TenantUsage> {
|
||||
self.usage.get(tenant_id).map(|u| u.snapshot())
|
||||
}
|
||||
|
||||
/// Get quota status for a tenant
|
||||
pub fn get_quota_status(&self, tenant_id: &str) -> Option<QuotaStatus> {
|
||||
let config = get_registry().get(tenant_id)?;
|
||||
let usage = self.get_usage(tenant_id).unwrap_or_default();
|
||||
|
||||
Some(QuotaStatus {
|
||||
tenant_id: tenant_id.to_string(),
|
||||
vectors: ResourceUsage {
|
||||
current: usage.vector_count,
|
||||
limit: config.quota.max_vectors,
|
||||
usage_percent: (usage.vector_count as f64 / config.quota.max_vectors as f64 * 100.0)
|
||||
as f32,
|
||||
},
|
||||
storage: ResourceUsage {
|
||||
current: usage.storage_bytes,
|
||||
limit: config.quota.max_storage_bytes,
|
||||
usage_percent: (usage.storage_bytes as f64 / config.quota.max_storage_bytes as f64
|
||||
* 100.0) as f32,
|
||||
},
|
||||
qps: RateUsage {
|
||||
current: usage.queries_this_second,
|
||||
limit: config.quota.max_qps,
|
||||
},
|
||||
concurrent: ResourceUsage {
|
||||
current: usage.concurrent_queries as u64,
|
||||
limit: config.quota.max_concurrent as u64,
|
||||
usage_percent: (usage.concurrent_queries as f64
|
||||
/ config.quota.max_concurrent as f64
|
||||
* 100.0) as f32,
|
||||
},
|
||||
collections: ResourceUsage {
|
||||
current: usage.collection_count as u64,
|
||||
limit: config.quota.max_collections as u64,
|
||||
usage_percent: (usage.collection_count as f64 / config.quota.max_collections as f64
|
||||
* 100.0) as f32,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
/// Reset usage counters for a tenant
|
||||
pub fn reset_usage(&self, tenant_id: &str) {
|
||||
if let Some(usage) = self.usage.get(tenant_id) {
|
||||
usage.vector_count.store(0, Ordering::Relaxed);
|
||||
usage.storage_bytes.store(0, Ordering::Relaxed);
|
||||
usage.collection_count.store(0, Ordering::Relaxed);
|
||||
usage.rate_count.store(0, Ordering::Relaxed);
|
||||
usage.concurrent_count.store(0, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
/// Initialize usage from stored values (e.g., from database)
|
||||
pub fn initialize_usage(&self, tenant_id: &str, stored_usage: TenantUsage) {
|
||||
let usage = self.get_or_create_usage(tenant_id);
|
||||
usage.reset_from(&stored_usage);
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for QuotaManager {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Resource usage summary
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ResourceUsage {
|
||||
/// Current usage
|
||||
pub current: u64,
|
||||
/// Maximum limit
|
||||
pub limit: u64,
|
||||
/// Usage percentage
|
||||
pub usage_percent: f32,
|
||||
}
|
||||
|
||||
/// Rate usage summary
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RateUsage {
|
||||
/// Current rate
|
||||
pub current: u32,
|
||||
/// Maximum rate
|
||||
pub limit: u32,
|
||||
}
|
||||
|
||||
/// Complete quota status for a tenant
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct QuotaStatus {
|
||||
/// Tenant ID
|
||||
pub tenant_id: String,
|
||||
/// Vector count usage
|
||||
pub vectors: ResourceUsage,
|
||||
/// Storage usage
|
||||
pub storage: ResourceUsage,
|
||||
/// Current QPS
|
||||
pub qps: RateUsage,
|
||||
/// Concurrent queries
|
||||
pub concurrent: ResourceUsage,
|
||||
/// Collection count
|
||||
pub collections: ResourceUsage,
|
||||
}
|
||||
|
||||
impl QuotaStatus {
|
||||
/// Check if any quota is near limit (>80%)
|
||||
pub fn is_near_limit(&self) -> bool {
|
||||
self.vectors.usage_percent > 80.0
|
||||
|| self.storage.usage_percent > 80.0
|
||||
|| self.collections.usage_percent > 80.0
|
||||
}
|
||||
|
||||
/// Check if any quota is critical (>95%)
|
||||
pub fn is_critical(&self) -> bool {
|
||||
self.vectors.usage_percent > 95.0 || self.storage.usage_percent > 95.0
|
||||
}
|
||||
}
|
||||
|
||||
/// Global quota manager instance
|
||||
static QUOTA_MANAGER: once_cell::sync::Lazy<QuotaManager> =
|
||||
once_cell::sync::Lazy::new(QuotaManager::new);
|
||||
|
||||
/// Get the global quota manager
|
||||
pub fn get_quota_manager() -> &'static QuotaManager {
|
||||
"A_MANAGER
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
fn chrono_now_millis() -> i64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as i64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_token_bucket_acquire() {
|
||||
let bucket = TokenBucket::new(10, 10); // 10 capacity, 10/second
|
||||
|
||||
// Should be able to acquire up to capacity
|
||||
for _ in 0..10 {
|
||||
assert!(bucket.try_acquire(1));
|
||||
}
|
||||
|
||||
// Should fail after capacity exhausted
|
||||
assert!(!bucket.try_acquire(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tenant_usage_tracking() {
|
||||
let manager = QuotaManager::new();
|
||||
|
||||
// Record some usage
|
||||
manager.record_vector_insert("test-tenant", 100, 1024 * 100);
|
||||
manager.record_collection_create("test-tenant");
|
||||
|
||||
// Check usage
|
||||
let usage = manager.get_usage("test-tenant").unwrap();
|
||||
assert_eq!(usage.vector_count, 100);
|
||||
assert_eq!(usage.storage_bytes, 1024 * 100);
|
||||
assert_eq!(usage.collection_count, 1);
|
||||
|
||||
// Record deletion
|
||||
manager.record_vector_delete("test-tenant", 50, 1024 * 50);
|
||||
let usage = manager.get_usage("test-tenant").unwrap();
|
||||
assert_eq!(usage.vector_count, 50);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quota_result_messages() {
|
||||
let result = QuotaResult::RateLimited {
|
||||
retry_after_ms: 100,
|
||||
};
|
||||
assert!(!result.is_allowed());
|
||||
assert!(result.error_message().unwrap().contains("100"));
|
||||
|
||||
let result = QuotaResult::VectorQuotaExceeded {
|
||||
current: 1000,
|
||||
limit: 1000,
|
||||
};
|
||||
assert!(!result.is_allowed());
|
||||
assert!(result.error_message().unwrap().contains("1000"));
|
||||
|
||||
let result = QuotaResult::Allowed;
|
||||
assert!(result.is_allowed());
|
||||
assert!(result.error_message().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_concurrent_query_tracking() {
|
||||
let manager = QuotaManager::new();
|
||||
|
||||
// Start queries
|
||||
manager.start_query("test-tenant");
|
||||
manager.start_query("test-tenant");
|
||||
|
||||
let usage = manager.get_usage("test-tenant").unwrap();
|
||||
assert_eq!(usage.concurrent_queries, 2);
|
||||
|
||||
// End one query
|
||||
manager.end_query("test-tenant");
|
||||
let usage = manager.get_usage("test-tenant").unwrap();
|
||||
assert_eq!(usage.concurrent_queries, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_usage_reset() {
|
||||
let manager = QuotaManager::new();
|
||||
|
||||
manager.record_vector_insert("test-tenant", 100, 1024);
|
||||
manager.record_collection_create("test-tenant");
|
||||
manager.start_query("test-tenant");
|
||||
|
||||
// Reset
|
||||
manager.reset_usage("test-tenant");
|
||||
|
||||
let usage = manager.get_usage("test-tenant").unwrap();
|
||||
assert_eq!(usage.vector_count, 0);
|
||||
assert_eq!(usage.storage_bytes, 0);
|
||||
assert_eq!(usage.collection_count, 0);
|
||||
assert_eq!(usage.concurrent_queries, 0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,719 @@
|
||||
//! Tenant Registry for RuVector Multi-Tenancy
|
||||
//!
|
||||
//! Provides tenant management with isolation levels, quotas, and metadata.
|
||||
//! Integrates with PostgreSQL's system tables for persistent storage.
|
||||
|
||||
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
|
||||
|
||||
use dashmap::DashMap;
|
||||
use parking_lot::RwLock;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Maximum number of tenants in shared memory (for fixed-size arrays)
|
||||
pub const MAX_TENANTS: usize = 10_000;
|
||||
|
||||
/// Isolation level for tenant data
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[repr(u8)]
|
||||
pub enum IsolationLevel {
|
||||
/// Shared index with tenant filter - most memory efficient
|
||||
Shared = 0,
|
||||
/// Dedicated partition within shared index structure
|
||||
Partition = 1,
|
||||
/// Completely separate physical index - maximum isolation
|
||||
Dedicated = 2,
|
||||
}
|
||||
|
||||
impl IsolationLevel {
|
||||
/// Parse isolation level from string
|
||||
pub fn from_str(s: &str) -> Option<Self> {
|
||||
match s.to_lowercase().as_str() {
|
||||
"shared" => Some(Self::Shared),
|
||||
"partition" => Some(Self::Partition),
|
||||
"dedicated" => Some(Self::Dedicated),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert to string representation
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Shared => "shared",
|
||||
Self::Partition => "partition",
|
||||
Self::Dedicated => "dedicated",
|
||||
}
|
||||
}
|
||||
|
||||
/// Get recommended vector count threshold for this level
|
||||
pub fn recommended_vector_count(&self) -> u64 {
|
||||
match self {
|
||||
Self::Shared => 100_000, // < 100K vectors
|
||||
Self::Partition => 10_000_000, // 100K - 10M vectors
|
||||
Self::Dedicated => u64::MAX, // > 10M vectors
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for IsolationLevel {
|
||||
fn default() -> Self {
|
||||
Self::Shared
|
||||
}
|
||||
}
|
||||
|
||||
/// Tenant quota configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TenantQuota {
|
||||
/// Maximum number of vectors
|
||||
pub max_vectors: u64,
|
||||
/// Maximum storage in bytes
|
||||
pub max_storage_bytes: u64,
|
||||
/// Maximum queries per second
|
||||
pub max_qps: u32,
|
||||
/// Maximum concurrent queries
|
||||
pub max_concurrent: u32,
|
||||
/// Maximum collections
|
||||
pub max_collections: u32,
|
||||
}
|
||||
|
||||
impl Default for TenantQuota {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_vectors: 1_000_000,
|
||||
max_storage_bytes: 10 * 1024 * 1024 * 1024, // 10 GB
|
||||
max_qps: 100,
|
||||
max_concurrent: 10,
|
||||
max_collections: 10,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tenant configuration and metadata
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TenantConfig {
|
||||
/// Unique tenant identifier
|
||||
pub id: String,
|
||||
/// Display name
|
||||
pub display_name: Option<String>,
|
||||
/// Isolation level
|
||||
pub isolation_level: IsolationLevel,
|
||||
/// Resource quotas
|
||||
pub quota: TenantQuota,
|
||||
/// Whether integrity monitoring is enabled
|
||||
pub integrity_enabled: bool,
|
||||
/// Custom integrity policy ID
|
||||
pub integrity_policy_id: Option<i64>,
|
||||
/// Arbitrary metadata
|
||||
pub metadata: serde_json::Value,
|
||||
/// Creation timestamp (epoch millis)
|
||||
pub created_at: i64,
|
||||
/// Suspension timestamp (None = active)
|
||||
pub suspended_at: Option<i64>,
|
||||
}
|
||||
|
||||
impl TenantConfig {
|
||||
/// Create a new tenant with default settings
|
||||
pub fn new(id: String) -> Self {
|
||||
Self {
|
||||
id,
|
||||
display_name: None,
|
||||
isolation_level: IsolationLevel::default(),
|
||||
quota: TenantQuota::default(),
|
||||
integrity_enabled: true,
|
||||
integrity_policy_id: None,
|
||||
metadata: serde_json::json!({}),
|
||||
created_at: chrono_now_millis(),
|
||||
suspended_at: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create tenant from JSONB configuration
|
||||
pub fn from_json(id: String, config: &serde_json::Value) -> Self {
|
||||
let mut tenant = Self::new(id);
|
||||
|
||||
if let Some(name) = config.get("display_name").and_then(|v| v.as_str()) {
|
||||
tenant.display_name = Some(name.to_string());
|
||||
}
|
||||
|
||||
if let Some(level) = config.get("isolation_level").and_then(|v| v.as_str()) {
|
||||
tenant.isolation_level = IsolationLevel::from_str(level).unwrap_or_default();
|
||||
}
|
||||
|
||||
if let Some(max_vec) = config.get("max_vectors").and_then(|v| v.as_u64()) {
|
||||
tenant.quota.max_vectors = max_vec;
|
||||
}
|
||||
|
||||
if let Some(max_qps) = config.get("max_qps").and_then(|v| v.as_u64()) {
|
||||
tenant.quota.max_qps = max_qps as u32;
|
||||
}
|
||||
|
||||
if let Some(max_storage) = config.get("max_storage_gb").and_then(|v| v.as_f64()) {
|
||||
tenant.quota.max_storage_bytes = (max_storage * 1024.0 * 1024.0 * 1024.0) as u64;
|
||||
}
|
||||
|
||||
if let Some(enabled) = config.get("integrity_enabled").and_then(|v| v.as_bool()) {
|
||||
tenant.integrity_enabled = enabled;
|
||||
}
|
||||
|
||||
if let Some(meta) = config.get("metadata") {
|
||||
tenant.metadata = meta.clone();
|
||||
}
|
||||
|
||||
tenant
|
||||
}
|
||||
|
||||
/// Check if tenant is suspended
|
||||
pub fn is_suspended(&self) -> bool {
|
||||
self.suspended_at.is_some()
|
||||
}
|
||||
|
||||
/// Check if tenant is active
|
||||
pub fn is_active(&self) -> bool {
|
||||
!self.is_suspended()
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared state for a tenant (in shared memory)
|
||||
#[repr(C)]
|
||||
pub struct TenantSharedState {
|
||||
/// Hash of tenant ID for fast lookup
|
||||
pub tenant_id_hash: AtomicU64,
|
||||
/// Current integrity state (0=normal, 1=stress, 2=critical)
|
||||
pub integrity_state: AtomicU32,
|
||||
/// Current lambda cut value (fixed-point: value * 1000)
|
||||
pub lambda_cut_fp: AtomicU32,
|
||||
/// Request count for rate limiting
|
||||
pub request_count: AtomicU32,
|
||||
/// Last request epoch (seconds)
|
||||
pub last_request_epoch: AtomicU64,
|
||||
/// Flags (bit 0: suspended, bit 1: migrating, etc.)
|
||||
pub flags: AtomicU32,
|
||||
}
|
||||
|
||||
impl TenantSharedState {
|
||||
/// Create new shared state
|
||||
pub fn new(tenant_id_hash: u64) -> Self {
|
||||
Self {
|
||||
tenant_id_hash: AtomicU64::new(tenant_id_hash),
|
||||
integrity_state: AtomicU32::new(0),
|
||||
lambda_cut_fp: AtomicU32::new(1000), // 1.0 in fixed point
|
||||
request_count: AtomicU32::new(0),
|
||||
last_request_epoch: AtomicU64::new(0),
|
||||
flags: AtomicU32::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Reset shared state for a new tenant (atomically reinitialize fields)
|
||||
pub fn reset(&self, tenant_id_hash: u64) {
|
||||
self.tenant_id_hash.store(tenant_id_hash, Ordering::Relaxed);
|
||||
self.integrity_state.store(0, Ordering::Relaxed);
|
||||
self.lambda_cut_fp.store(1000, Ordering::Relaxed); // 1.0 in fixed point
|
||||
self.request_count.store(0, Ordering::Relaxed);
|
||||
self.last_request_epoch.store(0, Ordering::Relaxed);
|
||||
self.flags.store(0, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Check if tenant is suspended
|
||||
pub fn is_suspended(&self) -> bool {
|
||||
(self.flags.load(Ordering::Relaxed) & 1) != 0
|
||||
}
|
||||
|
||||
/// Set suspended flag
|
||||
pub fn set_suspended(&self, suspended: bool) {
|
||||
if suspended {
|
||||
self.flags.fetch_or(1, Ordering::Relaxed);
|
||||
} else {
|
||||
self.flags.fetch_and(!1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if tenant is migrating
|
||||
pub fn is_migrating(&self) -> bool {
|
||||
(self.flags.load(Ordering::Relaxed) & 2) != 0
|
||||
}
|
||||
|
||||
/// Set migrating flag
|
||||
pub fn set_migrating(&self, migrating: bool) {
|
||||
if migrating {
|
||||
self.flags.fetch_or(2, Ordering::Relaxed);
|
||||
} else {
|
||||
self.flags.fetch_and(!2, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
/// Get lambda cut as f32
|
||||
pub fn lambda_cut(&self) -> f32 {
|
||||
self.lambda_cut_fp.load(Ordering::Relaxed) as f32 / 1000.0
|
||||
}
|
||||
|
||||
/// Set lambda cut from f32
|
||||
pub fn set_lambda_cut(&self, value: f32) {
|
||||
self.lambda_cut_fp
|
||||
.store((value * 1000.0) as u32, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Increment request count and check rate limit
|
||||
pub fn check_rate_limit(&self, max_qps: u32) -> bool {
|
||||
let now = current_epoch_seconds();
|
||||
let last_epoch = self.last_request_epoch.load(Ordering::Relaxed);
|
||||
|
||||
if now > last_epoch {
|
||||
// New second, reset counter
|
||||
self.last_request_epoch.store(now, Ordering::Relaxed);
|
||||
self.request_count.store(1, Ordering::Relaxed);
|
||||
true
|
||||
} else {
|
||||
// Same second, increment and check
|
||||
let count = self.request_count.fetch_add(1, Ordering::Relaxed);
|
||||
count < max_qps
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tenant Registry - manages all tenants
|
||||
pub struct TenantRegistry {
|
||||
/// Tenant configurations (heap-based for flexibility)
|
||||
configs: DashMap<String, TenantConfig>,
|
||||
/// Tenant ID to index mapping for shared memory lookup
|
||||
id_to_index: DashMap<String, usize>,
|
||||
/// Shared states (fixed-size for shared memory compatibility)
|
||||
shared_states: Vec<TenantSharedState>,
|
||||
/// Next available index
|
||||
next_index: AtomicU32,
|
||||
/// Promotion policy configuration
|
||||
promotion_policy: RwLock<PromotionPolicy>,
|
||||
}
|
||||
|
||||
/// Policy for automatic tenant isolation level promotion
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PromotionPolicy {
|
||||
/// Vector count threshold to promote from shared to partition
|
||||
pub partition_threshold: u64,
|
||||
/// Vector count threshold to promote from partition to dedicated
|
||||
pub dedicated_threshold: u64,
|
||||
/// Check interval in seconds
|
||||
pub check_interval_secs: u64,
|
||||
/// Whether auto-promotion is enabled
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
impl Default for PromotionPolicy {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
partition_threshold: 100_000,
|
||||
dedicated_threshold: 10_000_000,
|
||||
check_interval_secs: 3600, // 1 hour
|
||||
enabled: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TenantRegistry {
|
||||
/// Create a new tenant registry
|
||||
pub fn new() -> Self {
|
||||
let mut shared_states = Vec::with_capacity(MAX_TENANTS);
|
||||
for _ in 0..MAX_TENANTS {
|
||||
shared_states.push(TenantSharedState::new(0));
|
||||
}
|
||||
|
||||
Self {
|
||||
configs: DashMap::new(),
|
||||
id_to_index: DashMap::new(),
|
||||
shared_states,
|
||||
next_index: AtomicU32::new(0),
|
||||
promotion_policy: RwLock::new(PromotionPolicy::default()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a new tenant
|
||||
pub fn register(&self, config: TenantConfig) -> Result<usize, TenantError> {
|
||||
let tenant_id = config.id.clone();
|
||||
|
||||
// Check if tenant already exists
|
||||
if self.configs.contains_key(&tenant_id) {
|
||||
return Err(TenantError::AlreadyExists(tenant_id));
|
||||
}
|
||||
|
||||
// Allocate index
|
||||
let index = self.next_index.fetch_add(1, Ordering::Relaxed) as usize;
|
||||
if index >= MAX_TENANTS {
|
||||
return Err(TenantError::MaxTenantsReached);
|
||||
}
|
||||
|
||||
// Initialize shared state (atomically reset the pre-allocated slot)
|
||||
let id_hash = hash_tenant_id(&tenant_id);
|
||||
self.shared_states[index].reset(id_hash);
|
||||
|
||||
// Store mappings
|
||||
self.id_to_index.insert(tenant_id.clone(), index);
|
||||
self.configs.insert(tenant_id, config);
|
||||
|
||||
Ok(index)
|
||||
}
|
||||
|
||||
/// Get tenant configuration
|
||||
pub fn get(&self, tenant_id: &str) -> Option<TenantConfig> {
|
||||
self.configs.get(tenant_id).map(|r| r.value().clone())
|
||||
}
|
||||
|
||||
/// Get tenant shared state
|
||||
pub fn get_shared_state(&self, tenant_id: &str) -> Option<&TenantSharedState> {
|
||||
self.id_to_index
|
||||
.get(tenant_id)
|
||||
.map(|idx| &self.shared_states[*idx.value()])
|
||||
}
|
||||
|
||||
/// Update tenant configuration
|
||||
pub fn update(&self, tenant_id: &str, config: TenantConfig) -> Result<(), TenantError> {
|
||||
if !self.configs.contains_key(tenant_id) {
|
||||
return Err(TenantError::NotFound(tenant_id.to_string()));
|
||||
}
|
||||
self.configs.insert(tenant_id.to_string(), config);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Suspend a tenant
|
||||
pub fn suspend(&self, tenant_id: &str) -> Result<(), TenantError> {
|
||||
let mut config = self
|
||||
.get(tenant_id)
|
||||
.ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
|
||||
|
||||
config.suspended_at = Some(chrono_now_millis());
|
||||
self.configs.insert(tenant_id.to_string(), config);
|
||||
|
||||
// Update shared state
|
||||
if let Some(state) = self.get_shared_state(tenant_id) {
|
||||
state.set_suspended(true);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Resume a suspended tenant
|
||||
pub fn resume(&self, tenant_id: &str) -> Result<(), TenantError> {
|
||||
let mut config = self
|
||||
.get(tenant_id)
|
||||
.ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
|
||||
|
||||
config.suspended_at = None;
|
||||
self.configs.insert(tenant_id.to_string(), config);
|
||||
|
||||
// Update shared state
|
||||
if let Some(state) = self.get_shared_state(tenant_id) {
|
||||
state.set_suspended(false);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Delete a tenant (soft delete by default)
|
||||
pub fn delete(&self, tenant_id: &str, hard: bool) -> Result<(), TenantError> {
|
||||
if !self.configs.contains_key(tenant_id) {
|
||||
return Err(TenantError::NotFound(tenant_id.to_string()));
|
||||
}
|
||||
|
||||
if hard {
|
||||
// Hard delete: remove immediately
|
||||
self.configs.remove(tenant_id);
|
||||
self.id_to_index.remove(tenant_id);
|
||||
// Note: shared state index remains allocated (could implement recycling)
|
||||
} else {
|
||||
// Soft delete: mark as suspended with deletion flag
|
||||
let mut config = self.get(tenant_id).unwrap();
|
||||
config.suspended_at = Some(chrono_now_millis());
|
||||
config.metadata["deleted"] = serde_json::json!(true);
|
||||
self.configs.insert(tenant_id.to_string(), config);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// List all tenants
|
||||
pub fn list(&self) -> Vec<TenantConfig> {
|
||||
self.configs.iter().map(|r| r.value().clone()).collect()
|
||||
}
|
||||
|
||||
/// List active tenants only
|
||||
pub fn list_active(&self) -> Vec<TenantConfig> {
|
||||
self.configs
|
||||
.iter()
|
||||
.filter(|r| r.value().is_active())
|
||||
.map(|r| r.value().clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Get tenant count
|
||||
pub fn count(&self) -> usize {
|
||||
self.configs.len()
|
||||
}
|
||||
|
||||
/// Validate tenant context for operations
|
||||
pub fn validate_context(&self, tenant_id: &str) -> Result<TenantConfig, TenantError> {
|
||||
if tenant_id.is_empty() {
|
||||
return Err(TenantError::NoContext);
|
||||
}
|
||||
|
||||
// Special wildcard for admin operations
|
||||
if tenant_id == "*" {
|
||||
// Check if caller has admin privileges (would check PostgreSQL roles)
|
||||
return Err(TenantError::AdminContextRequired);
|
||||
}
|
||||
|
||||
let config = self
|
||||
.get(tenant_id)
|
||||
.ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
|
||||
|
||||
if config.is_suspended() {
|
||||
return Err(TenantError::Suspended(tenant_id.to_string()));
|
||||
}
|
||||
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
/// Check rate limit for tenant
|
||||
pub fn check_rate_limit(&self, tenant_id: &str) -> Result<bool, TenantError> {
|
||||
let config = self
|
||||
.get(tenant_id)
|
||||
.ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
|
||||
|
||||
let state = self
|
||||
.get_shared_state(tenant_id)
|
||||
.ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
|
||||
|
||||
Ok(state.check_rate_limit(config.quota.max_qps))
|
||||
}
|
||||
|
||||
/// Get promotion policy
|
||||
pub fn get_promotion_policy(&self) -> PromotionPolicy {
|
||||
self.promotion_policy.read().clone()
|
||||
}
|
||||
|
||||
/// Set promotion policy
|
||||
pub fn set_promotion_policy(&self, policy: PromotionPolicy) {
|
||||
*self.promotion_policy.write() = policy;
|
||||
}
|
||||
|
||||
/// Check if tenant should be promoted to higher isolation level
|
||||
pub fn check_promotion(&self, tenant_id: &str, vector_count: u64) -> Option<IsolationLevel> {
|
||||
let config = self.get(tenant_id)?;
|
||||
let policy = self.promotion_policy.read();
|
||||
|
||||
if !policy.enabled {
|
||||
return None;
|
||||
}
|
||||
|
||||
match config.isolation_level {
|
||||
IsolationLevel::Shared if vector_count > policy.dedicated_threshold => {
|
||||
Some(IsolationLevel::Dedicated)
|
||||
}
|
||||
IsolationLevel::Shared if vector_count > policy.partition_threshold => {
|
||||
Some(IsolationLevel::Partition)
|
||||
}
|
||||
IsolationLevel::Partition if vector_count > policy.dedicated_threshold => {
|
||||
Some(IsolationLevel::Dedicated)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TenantRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Tenant operation errors
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum TenantError {
|
||||
/// Tenant already exists
|
||||
AlreadyExists(String),
|
||||
/// Tenant not found
|
||||
NotFound(String),
|
||||
/// Tenant is suspended
|
||||
Suspended(String),
|
||||
/// No tenant context set
|
||||
NoContext,
|
||||
/// Admin context required for operation
|
||||
AdminContextRequired,
|
||||
/// Maximum number of tenants reached
|
||||
MaxTenantsReached,
|
||||
/// Rate limit exceeded
|
||||
RateLimitExceeded(String),
|
||||
/// Quota exceeded
|
||||
QuotaExceeded(String, String),
|
||||
/// Tenant mismatch (security violation)
|
||||
TenantMismatch { context: String, request: String },
|
||||
/// Invalid tenant ID format (validation error)
|
||||
InvalidId(String),
|
||||
}
|
||||
|
||||
impl std::fmt::Display for TenantError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::AlreadyExists(id) => write!(f, "Tenant '{}' already exists", id),
|
||||
Self::NotFound(id) => write!(f, "Tenant '{}' not found", id),
|
||||
Self::Suspended(id) => write!(f, "Tenant '{}' is suspended", id),
|
||||
Self::NoContext => write!(f, "No tenant context set (use SET ruvector.tenant_id)"),
|
||||
Self::AdminContextRequired => write!(f, "Admin context required for this operation"),
|
||||
Self::MaxTenantsReached => write!(f, "Maximum number of tenants reached"),
|
||||
Self::RateLimitExceeded(id) => write!(f, "Rate limit exceeded for tenant '{}'", id),
|
||||
Self::QuotaExceeded(id, resource) => {
|
||||
write!(f, "Quota exceeded for tenant '{}': {}", id, resource)
|
||||
}
|
||||
Self::TenantMismatch { context, request } => {
|
||||
write!(
|
||||
f,
|
||||
"Tenant mismatch: context='{}', request='{}'",
|
||||
context, request
|
||||
)
|
||||
}
|
||||
Self::InvalidId(msg) => write!(f, "Invalid tenant ID: {}", msg),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for TenantError {}
|
||||
|
||||
/// Global tenant registry instance
|
||||
static TENANT_REGISTRY: once_cell::sync::Lazy<TenantRegistry> =
|
||||
once_cell::sync::Lazy::new(TenantRegistry::new);
|
||||
|
||||
/// Get the global tenant registry
|
||||
pub fn get_registry() -> &'static TenantRegistry {
|
||||
&TENANT_REGISTRY
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
|
||||
/// Hash a tenant ID for fast lookup
|
||||
fn hash_tenant_id(id: &str) -> u64 {
|
||||
use std::hash::{Hash, Hasher};
|
||||
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
||||
id.hash(&mut hasher);
|
||||
hasher.finish()
|
||||
}
|
||||
|
||||
/// Get current timestamp in milliseconds
|
||||
fn chrono_now_millis() -> i64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as i64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Get current epoch in seconds
|
||||
fn current_epoch_seconds() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_isolation_level_parse() {
|
||||
assert_eq!(
|
||||
IsolationLevel::from_str("shared"),
|
||||
Some(IsolationLevel::Shared)
|
||||
);
|
||||
assert_eq!(
|
||||
IsolationLevel::from_str("partition"),
|
||||
Some(IsolationLevel::Partition)
|
||||
);
|
||||
assert_eq!(
|
||||
IsolationLevel::from_str("dedicated"),
|
||||
Some(IsolationLevel::Dedicated)
|
||||
);
|
||||
assert_eq!(IsolationLevel::from_str("invalid"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tenant_config_from_json() {
|
||||
let json = serde_json::json!({
|
||||
"display_name": "Test Corp",
|
||||
"isolation_level": "dedicated",
|
||||
"max_vectors": 5000000,
|
||||
"max_qps": 200,
|
||||
"integrity_enabled": false
|
||||
});
|
||||
|
||||
let config = TenantConfig::from_json("test-tenant".to_string(), &json);
|
||||
assert_eq!(config.display_name, Some("Test Corp".to_string()));
|
||||
assert_eq!(config.isolation_level, IsolationLevel::Dedicated);
|
||||
assert_eq!(config.quota.max_vectors, 5000000);
|
||||
assert_eq!(config.quota.max_qps, 200);
|
||||
assert!(!config.integrity_enabled);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tenant_registry_register() {
|
||||
let registry = TenantRegistry::new();
|
||||
let config = TenantConfig::new("test-tenant".to_string());
|
||||
|
||||
let result = registry.register(config.clone());
|
||||
assert!(result.is_ok());
|
||||
|
||||
// Should fail for duplicate
|
||||
let result2 = registry.register(config);
|
||||
assert!(matches!(result2, Err(TenantError::AlreadyExists(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tenant_suspension() {
|
||||
let registry = TenantRegistry::new();
|
||||
let config = TenantConfig::new("test-tenant".to_string());
|
||||
registry.register(config).unwrap();
|
||||
|
||||
// Suspend
|
||||
registry.suspend("test-tenant").unwrap();
|
||||
let config = registry.get("test-tenant").unwrap();
|
||||
assert!(config.is_suspended());
|
||||
|
||||
// Resume
|
||||
registry.resume("test-tenant").unwrap();
|
||||
let config = registry.get("test-tenant").unwrap();
|
||||
assert!(!config.is_suspended());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_shared_state_rate_limiting() {
|
||||
let state = TenantSharedState::new(12345);
|
||||
|
||||
// First request should pass
|
||||
assert!(state.check_rate_limit(10));
|
||||
|
||||
// Subsequent requests within limit should pass
|
||||
for _ in 0..8 {
|
||||
assert!(state.check_rate_limit(10));
|
||||
}
|
||||
|
||||
// 10th request should fail (at limit)
|
||||
assert!(!state.check_rate_limit(10));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_promotion_check() {
|
||||
let registry = TenantRegistry::new();
|
||||
let mut config = TenantConfig::new("test-tenant".to_string());
|
||||
config.isolation_level = IsolationLevel::Shared;
|
||||
registry.register(config).unwrap();
|
||||
|
||||
// Below threshold - no promotion
|
||||
assert!(registry.check_promotion("test-tenant", 50_000).is_none());
|
||||
|
||||
// Above partition threshold
|
||||
assert_eq!(
|
||||
registry.check_promotion("test-tenant", 500_000),
|
||||
Some(IsolationLevel::Partition)
|
||||
);
|
||||
|
||||
// Above dedicated threshold (jumps directly)
|
||||
assert_eq!(
|
||||
registry.check_promotion("test-tenant", 15_000_000),
|
||||
Some(IsolationLevel::Dedicated)
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,637 @@
|
||||
//! Row-Level Security Integration for RuVector Multi-Tenancy
|
||||
//!
|
||||
//! Provides automatic RLS policy generation and management for tenant isolation.
|
||||
//! Integrates with PostgreSQL's native RLS capabilities.
|
||||
|
||||
use dashmap::DashMap;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// RLS policy configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RlsPolicyConfig {
|
||||
/// Table name (fully qualified)
|
||||
pub table_name: String,
|
||||
/// Tenant ID column name
|
||||
pub tenant_column: String,
|
||||
/// Policy name
|
||||
pub policy_name: String,
|
||||
/// Whether to create admin bypass policy
|
||||
pub admin_bypass: bool,
|
||||
/// Whether to create wildcard policy for admin queries
|
||||
pub wildcard_policy: bool,
|
||||
/// Custom USING clause (optional)
|
||||
pub custom_using: Option<String>,
|
||||
/// Custom WITH CHECK clause (optional)
|
||||
pub custom_with_check: Option<String>,
|
||||
}
|
||||
|
||||
impl Default for RlsPolicyConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
table_name: String::new(),
|
||||
tenant_column: "tenant_id".to_string(),
|
||||
policy_name: "ruvector_tenant_isolation".to_string(),
|
||||
admin_bypass: true,
|
||||
wildcard_policy: true,
|
||||
custom_using: None,
|
||||
custom_with_check: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RlsPolicyConfig {
|
||||
/// Create a new RLS policy config for a table
|
||||
pub fn new(table_name: &str) -> Self {
|
||||
Self {
|
||||
table_name: table_name.to_string(),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Set tenant column name
|
||||
pub fn with_tenant_column(mut self, column: &str) -> Self {
|
||||
self.tenant_column = column.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
/// Set policy name
|
||||
pub fn with_policy_name(mut self, name: &str) -> Self {
|
||||
self.policy_name = name.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
/// Disable admin bypass
|
||||
pub fn without_admin_bypass(mut self) -> Self {
|
||||
self.admin_bypass = false;
|
||||
self
|
||||
}
|
||||
|
||||
/// Disable wildcard policy
|
||||
pub fn without_wildcard(mut self) -> Self {
|
||||
self.wildcard_policy = false;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set custom USING clause
|
||||
pub fn with_custom_using(mut self, using: &str) -> Self {
|
||||
self.custom_using = Some(using.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Set custom WITH CHECK clause
|
||||
pub fn with_custom_check(mut self, check: &str) -> Self {
|
||||
self.custom_with_check = Some(check.to_string());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Policy template for common patterns
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum PolicyTemplate {
|
||||
/// Standard tenant isolation (tenant_id = current_setting)
|
||||
Standard,
|
||||
/// Read-only for other tenants with write for own
|
||||
ReadSharedWriteOwn,
|
||||
/// Hierarchical (tenant can see child tenants)
|
||||
Hierarchical {
|
||||
/// Path column for hierarchy (e.g., "tenant_path")
|
||||
path_column: String,
|
||||
},
|
||||
/// Time-based access (tenant_id + time window)
|
||||
TimeBased {
|
||||
/// Timestamp column
|
||||
time_column: String,
|
||||
/// Retention period in days
|
||||
retention_days: i32,
|
||||
},
|
||||
/// Custom template
|
||||
Custom {
|
||||
/// Custom USING expression
|
||||
using_expr: String,
|
||||
/// Custom WITH CHECK expression
|
||||
check_expr: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl PolicyTemplate {
|
||||
/// Generate USING clause for this template
|
||||
pub fn using_clause(&self, tenant_column: &str) -> String {
|
||||
match self {
|
||||
Self::Standard => {
|
||||
format!(
|
||||
"{} = current_setting('ruvector.tenant_id', true)",
|
||||
tenant_column
|
||||
)
|
||||
}
|
||||
Self::ReadSharedWriteOwn => {
|
||||
format!(
|
||||
"{} = current_setting('ruvector.tenant_id', true) OR is_public = true",
|
||||
tenant_column
|
||||
)
|
||||
}
|
||||
Self::Hierarchical { path_column } => {
|
||||
format!(
|
||||
"{} LIKE current_setting('ruvector.tenant_id', true) || '%'",
|
||||
path_column
|
||||
)
|
||||
}
|
||||
Self::TimeBased {
|
||||
time_column,
|
||||
retention_days,
|
||||
} => {
|
||||
format!(
|
||||
"{} = current_setting('ruvector.tenant_id', true) AND {} > NOW() - INTERVAL '{} days'",
|
||||
tenant_column, time_column, retention_days
|
||||
)
|
||||
}
|
||||
Self::Custom { using_expr, .. } => using_expr.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate WITH CHECK clause for this template
|
||||
pub fn check_clause(&self, tenant_column: &str) -> String {
|
||||
match self {
|
||||
Self::Standard | Self::ReadSharedWriteOwn => {
|
||||
format!(
|
||||
"{} = current_setting('ruvector.tenant_id', true)",
|
||||
tenant_column
|
||||
)
|
||||
}
|
||||
Self::Hierarchical { .. } => {
|
||||
format!(
|
||||
"{} = current_setting('ruvector.tenant_id', true)",
|
||||
tenant_column
|
||||
)
|
||||
}
|
||||
Self::TimeBased {
|
||||
time_column: _,
|
||||
retention_days: _,
|
||||
} => {
|
||||
format!(
|
||||
"{} = current_setting('ruvector.tenant_id', true)",
|
||||
tenant_column
|
||||
)
|
||||
}
|
||||
Self::Custom { check_expr, .. } => check_expr.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// RLS policy manager
|
||||
pub struct RlsManager {
|
||||
/// Active policies by table
|
||||
policies: DashMap<String, RlsPolicyConfig>,
|
||||
/// Tables with RLS enabled
|
||||
enabled_tables: DashMap<String, bool>,
|
||||
}
|
||||
|
||||
impl RlsManager {
|
||||
/// Create a new RLS manager
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
policies: DashMap::new(),
|
||||
enabled_tables: DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate SQL to enable RLS on a table with tenant isolation
|
||||
pub fn generate_enable_rls_sql(&self, config: &RlsPolicyConfig) -> String {
|
||||
let using_clause = config.custom_using.clone().unwrap_or_else(|| {
|
||||
format!(
|
||||
"{} = current_setting('ruvector.tenant_id', true)",
|
||||
config.tenant_column
|
||||
)
|
||||
});
|
||||
|
||||
let check_clause = config.custom_with_check.clone().unwrap_or_else(|| {
|
||||
format!(
|
||||
"{} = current_setting('ruvector.tenant_id', true)",
|
||||
config.tenant_column
|
||||
)
|
||||
});
|
||||
|
||||
let mut sql = format!(
|
||||
r#"
|
||||
-- Enable Row-Level Security on {table}
|
||||
ALTER TABLE {table} ENABLE ROW LEVEL SECURITY;
|
||||
|
||||
-- Force RLS even for table owners (recommended for security)
|
||||
ALTER TABLE {table} FORCE ROW LEVEL SECURITY;
|
||||
|
||||
-- Drop existing RuVector policies if any
|
||||
DROP POLICY IF EXISTS {policy} ON {table};
|
||||
DROP POLICY IF EXISTS {policy}_admin ON {table};
|
||||
DROP POLICY IF EXISTS {policy}_wildcard ON {table};
|
||||
|
||||
-- Create tenant isolation policy (applies to all operations)
|
||||
CREATE POLICY {policy} ON {table}
|
||||
FOR ALL
|
||||
USING ({using})
|
||||
WITH CHECK ({check});
|
||||
"#,
|
||||
table = config.table_name,
|
||||
policy = config.policy_name,
|
||||
using = using_clause,
|
||||
check = check_clause
|
||||
);
|
||||
|
||||
if config.admin_bypass {
|
||||
sql.push_str(&format!(
|
||||
r#"
|
||||
-- Create admin bypass policy
|
||||
-- Requires role: ruvector_admin
|
||||
CREATE POLICY {policy}_admin ON {table}
|
||||
FOR ALL
|
||||
TO ruvector_admin
|
||||
USING (true)
|
||||
WITH CHECK (true);
|
||||
"#,
|
||||
table = config.table_name,
|
||||
policy = config.policy_name
|
||||
));
|
||||
}
|
||||
|
||||
if config.wildcard_policy {
|
||||
sql.push_str(&format!(
|
||||
r#"
|
||||
-- Create wildcard policy for cross-tenant admin queries
|
||||
-- Only applies when tenant_id is set to '*'
|
||||
CREATE POLICY {policy}_wildcard ON {table}
|
||||
FOR SELECT
|
||||
USING (current_setting('ruvector.tenant_id', true) = '*');
|
||||
"#,
|
||||
table = config.table_name,
|
||||
policy = config.policy_name
|
||||
));
|
||||
}
|
||||
|
||||
sql
|
||||
}
|
||||
|
||||
/// Generate SQL to disable RLS on a table
|
||||
pub fn generate_disable_rls_sql(&self, table_name: &str) -> String {
|
||||
format!(
|
||||
r#"
|
||||
-- Disable Row-Level Security on {table}
|
||||
ALTER TABLE {table} NO FORCE ROW LEVEL SECURITY;
|
||||
ALTER TABLE {table} DISABLE ROW LEVEL SECURITY;
|
||||
|
||||
-- Drop all RuVector policies
|
||||
DROP POLICY IF EXISTS ruvector_tenant_isolation ON {table};
|
||||
DROP POLICY IF EXISTS ruvector_tenant_isolation_admin ON {table};
|
||||
DROP POLICY IF EXISTS ruvector_tenant_isolation_wildcard ON {table};
|
||||
"#,
|
||||
table = table_name
|
||||
)
|
||||
}
|
||||
|
||||
/// Generate SQL using a policy template
|
||||
pub fn generate_from_template(
|
||||
&self,
|
||||
table_name: &str,
|
||||
tenant_column: &str,
|
||||
template: &PolicyTemplate,
|
||||
) -> String {
|
||||
let using_clause = template.using_clause(tenant_column);
|
||||
let check_clause = template.check_clause(tenant_column);
|
||||
|
||||
let config = RlsPolicyConfig {
|
||||
table_name: table_name.to_string(),
|
||||
tenant_column: tenant_column.to_string(),
|
||||
custom_using: Some(using_clause),
|
||||
custom_with_check: Some(check_clause),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
self.generate_enable_rls_sql(&config)
|
||||
}
|
||||
|
||||
/// Generate SQL to set tenant context for a session
|
||||
pub fn generate_set_tenant_sql(tenant_id: &str, local: bool) -> String {
|
||||
let set_cmd = if local { "SET LOCAL" } else { "SET" };
|
||||
format!("{} ruvector.tenant_id = '{}';", set_cmd, tenant_id)
|
||||
}
|
||||
|
||||
/// Generate SQL to clear tenant context
|
||||
pub fn generate_clear_tenant_sql() -> String {
|
||||
"RESET ruvector.tenant_id;".to_string()
|
||||
}
|
||||
|
||||
/// Generate SQL to get current tenant context
|
||||
pub fn generate_get_tenant_sql() -> String {
|
||||
"SELECT current_setting('ruvector.tenant_id', true);".to_string()
|
||||
}
|
||||
|
||||
/// Register a policy configuration
|
||||
pub fn register_policy(&self, config: RlsPolicyConfig) {
|
||||
let table_name = config.table_name.clone();
|
||||
self.policies.insert(table_name.clone(), config);
|
||||
self.enabled_tables.insert(table_name, true);
|
||||
}
|
||||
|
||||
/// Get policy for a table
|
||||
pub fn get_policy(&self, table_name: &str) -> Option<RlsPolicyConfig> {
|
||||
self.policies.get(table_name).map(|r| r.value().clone())
|
||||
}
|
||||
|
||||
/// Check if RLS is enabled for a table
|
||||
pub fn is_enabled(&self, table_name: &str) -> bool {
|
||||
self.enabled_tables
|
||||
.get(table_name)
|
||||
.map(|r| *r.value())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// List all tables with RLS enabled
|
||||
pub fn list_enabled_tables(&self) -> Vec<String> {
|
||||
self.enabled_tables
|
||||
.iter()
|
||||
.filter(|r| *r.value())
|
||||
.map(|r| r.key().clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Generate SQL to create default tenant column with proper constraints
|
||||
pub fn generate_add_tenant_column_sql(
|
||||
table_name: &str,
|
||||
column_name: &str,
|
||||
not_null: bool,
|
||||
default_current: bool,
|
||||
) -> String {
|
||||
let mut sql = format!(
|
||||
"ALTER TABLE {} ADD COLUMN IF NOT EXISTS {} TEXT",
|
||||
table_name, column_name
|
||||
);
|
||||
|
||||
if not_null {
|
||||
sql.push_str(" NOT NULL");
|
||||
}
|
||||
|
||||
if default_current {
|
||||
sql.push_str(" DEFAULT current_setting('ruvector.tenant_id')");
|
||||
}
|
||||
|
||||
sql.push_str(";\n");
|
||||
|
||||
// Add foreign key constraint to tenants table
|
||||
sql.push_str(&format!(
|
||||
r#"
|
||||
-- Add foreign key to tenants table (optional, depends on schema)
|
||||
-- ALTER TABLE {} ADD CONSTRAINT fk_{}_tenant
|
||||
-- FOREIGN KEY ({}) REFERENCES ruvector.tenants(id) ON DELETE CASCADE;
|
||||
|
||||
-- Create index on tenant column for efficient filtering
|
||||
CREATE INDEX IF NOT EXISTS idx_{}_{} ON {} ({});
|
||||
"#,
|
||||
table_name,
|
||||
table_name.replace('.', "_"),
|
||||
column_name,
|
||||
table_name.replace('.', "_"),
|
||||
column_name,
|
||||
table_name,
|
||||
column_name
|
||||
));
|
||||
|
||||
sql
|
||||
}
|
||||
|
||||
/// Generate SQL to create roles for RLS
|
||||
pub fn generate_roles_sql() -> String {
|
||||
r#"
|
||||
-- Create RuVector roles for RLS
|
||||
DO $$
|
||||
BEGIN
|
||||
-- Admin role (bypasses RLS)
|
||||
IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = 'ruvector_admin') THEN
|
||||
CREATE ROLE ruvector_admin;
|
||||
END IF;
|
||||
|
||||
-- Standard user role (subject to RLS)
|
||||
IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = 'ruvector_users') THEN
|
||||
CREATE ROLE ruvector_users;
|
||||
END IF;
|
||||
|
||||
-- Read-only role
|
||||
IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = 'ruvector_readonly') THEN
|
||||
CREATE ROLE ruvector_readonly;
|
||||
END IF;
|
||||
END
|
||||
$$;
|
||||
|
||||
-- Grant basic permissions
|
||||
GRANT USAGE ON SCHEMA public TO ruvector_users, ruvector_readonly;
|
||||
GRANT SELECT ON ALL TABLES IN SCHEMA public TO ruvector_readonly;
|
||||
GRANT ALL ON ALL TABLES IN SCHEMA public TO ruvector_users;
|
||||
"#
|
||||
.to_string()
|
||||
}
|
||||
|
||||
/// Generate SQL for tenant context validation trigger
|
||||
pub fn generate_context_validation_trigger(table_name: &str, tenant_column: &str) -> String {
|
||||
format!(
|
||||
r#"
|
||||
-- Create function to validate tenant context before insert/update
|
||||
CREATE OR REPLACE FUNCTION ruvector_validate_tenant_context_{table_safe}()
|
||||
RETURNS TRIGGER AS $$
|
||||
DECLARE
|
||||
v_tenant_id TEXT;
|
||||
BEGIN
|
||||
-- Get current tenant context
|
||||
v_tenant_id := current_setting('ruvector.tenant_id', true);
|
||||
|
||||
-- Validate context is set
|
||||
IF v_tenant_id IS NULL OR v_tenant_id = '' THEN
|
||||
RAISE EXCEPTION 'No tenant context set. Use SET ruvector.tenant_id = ''your-tenant-id''';
|
||||
END IF;
|
||||
|
||||
-- Validate tenant matches (prevent cross-tenant writes)
|
||||
IF NEW.{column} IS NOT NULL AND NEW.{column} != v_tenant_id AND v_tenant_id != '*' THEN
|
||||
RAISE EXCEPTION 'Cannot write to different tenant: context=%, row=%',
|
||||
v_tenant_id, NEW.{column};
|
||||
END IF;
|
||||
|
||||
-- Auto-set tenant_id if not provided
|
||||
IF NEW.{column} IS NULL THEN
|
||||
NEW.{column} := v_tenant_id;
|
||||
END IF;
|
||||
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
-- Create trigger
|
||||
DROP TRIGGER IF EXISTS trg_ruvector_validate_tenant_{table_safe} ON {table};
|
||||
CREATE TRIGGER trg_ruvector_validate_tenant_{table_safe}
|
||||
BEFORE INSERT OR UPDATE ON {table}
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION ruvector_validate_tenant_context_{table_safe}();
|
||||
"#,
|
||||
table = table_name,
|
||||
table_safe = table_name.replace('.', "_").replace('"', ""),
|
||||
column = tenant_column
|
||||
)
|
||||
}
|
||||
|
||||
/// Generate SQL to check tenant existence before operations
|
||||
pub fn generate_tenant_existence_check(table_name: &str, tenant_column: &str) -> String {
|
||||
format!(
|
||||
r#"
|
||||
-- Create function to check tenant exists
|
||||
CREATE OR REPLACE FUNCTION ruvector_check_tenant_exists_{table_safe}()
|
||||
RETURNS TRIGGER AS $$
|
||||
BEGIN
|
||||
-- Check tenant exists (skip for admin wildcard)
|
||||
IF NEW.{column} != '*' THEN
|
||||
IF NOT EXISTS (SELECT 1 FROM ruvector.tenants WHERE id = NEW.{column}) THEN
|
||||
RAISE EXCEPTION 'Tenant does not exist: %', NEW.{column};
|
||||
END IF;
|
||||
|
||||
-- Check tenant is not suspended
|
||||
IF EXISTS (SELECT 1 FROM ruvector.tenants WHERE id = NEW.{column} AND suspended_at IS NOT NULL) THEN
|
||||
RAISE EXCEPTION 'Tenant is suspended: %', NEW.{column};
|
||||
END IF;
|
||||
END IF;
|
||||
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
-- Create trigger (runs after tenant context validation)
|
||||
DROP TRIGGER IF EXISTS trg_ruvector_check_tenant_{table_safe} ON {table};
|
||||
CREATE TRIGGER trg_ruvector_check_tenant_{table_safe}
|
||||
BEFORE INSERT OR UPDATE ON {table}
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION ruvector_check_tenant_exists_{table_safe}();
|
||||
"#,
|
||||
table = table_name,
|
||||
table_safe = table_name.replace('.', "_").replace('"', ""),
|
||||
column = tenant_column
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RlsManager {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Global RLS manager instance
|
||||
static RLS_MANAGER: once_cell::sync::Lazy<RlsManager> = once_cell::sync::Lazy::new(RlsManager::new);
|
||||
|
||||
/// Get the global RLS manager
|
||||
pub fn get_rls_manager() -> &'static RlsManager {
|
||||
&RLS_MANAGER
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_policy_config_builder() {
|
||||
let config = RlsPolicyConfig::new("embeddings")
|
||||
.with_tenant_column("org_id")
|
||||
.with_policy_name("custom_policy")
|
||||
.without_admin_bypass();
|
||||
|
||||
assert_eq!(config.table_name, "embeddings");
|
||||
assert_eq!(config.tenant_column, "org_id");
|
||||
assert_eq!(config.policy_name, "custom_policy");
|
||||
assert!(!config.admin_bypass);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_standard_policy_template() {
|
||||
let template = PolicyTemplate::Standard;
|
||||
let using = template.using_clause("tenant_id");
|
||||
let check = template.check_clause("tenant_id");
|
||||
|
||||
assert!(using.contains("tenant_id"));
|
||||
assert!(using.contains("current_setting"));
|
||||
assert!(check.contains("tenant_id"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hierarchical_template() {
|
||||
let template = PolicyTemplate::Hierarchical {
|
||||
path_column: "org_path".to_string(),
|
||||
};
|
||||
let using = template.using_clause("tenant_id");
|
||||
|
||||
assert!(using.contains("org_path"));
|
||||
assert!(using.contains("LIKE"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_time_based_template() {
|
||||
let template = PolicyTemplate::TimeBased {
|
||||
time_column: "created_at".to_string(),
|
||||
retention_days: 30,
|
||||
};
|
||||
let using = template.using_clause("tenant_id");
|
||||
|
||||
assert!(using.contains("created_at"));
|
||||
assert!(using.contains("30 days"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_generate_enable_rls_sql() {
|
||||
let manager = RlsManager::new();
|
||||
let config = RlsPolicyConfig::new("embeddings");
|
||||
let sql = manager.generate_enable_rls_sql(&config);
|
||||
|
||||
assert!(sql.contains("ENABLE ROW LEVEL SECURITY"));
|
||||
assert!(sql.contains("ruvector_tenant_isolation"));
|
||||
assert!(sql.contains("ruvector.tenant_id"));
|
||||
assert!(sql.contains("ruvector_admin")); // Admin bypass
|
||||
assert!(sql.contains("wildcard")); // Wildcard policy
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_generate_disable_rls_sql() {
|
||||
let manager = RlsManager::new();
|
||||
let sql = manager.generate_disable_rls_sql("embeddings");
|
||||
|
||||
assert!(sql.contains("DISABLE ROW LEVEL SECURITY"));
|
||||
assert!(sql.contains("DROP POLICY"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_set_tenant_sql() {
|
||||
let sql = RlsManager::generate_set_tenant_sql("acme-corp", false);
|
||||
assert!(sql.contains("SET ruvector.tenant_id"));
|
||||
assert!(sql.contains("acme-corp"));
|
||||
|
||||
let sql_local = RlsManager::generate_set_tenant_sql("acme-corp", true);
|
||||
assert!(sql_local.contains("SET LOCAL"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add_tenant_column_sql() {
|
||||
let sql = RlsManager::generate_add_tenant_column_sql("embeddings", "tenant_id", true, true);
|
||||
|
||||
assert!(sql.contains("ADD COLUMN"));
|
||||
assert!(sql.contains("NOT NULL"));
|
||||
assert!(sql.contains("DEFAULT current_setting"));
|
||||
assert!(sql.contains("CREATE INDEX"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_roles_sql() {
|
||||
let sql = RlsManager::generate_roles_sql();
|
||||
|
||||
assert!(sql.contains("ruvector_admin"));
|
||||
assert!(sql.contains("ruvector_users"));
|
||||
assert!(sql.contains("ruvector_readonly"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_validation_trigger() {
|
||||
let sql = RlsManager::generate_context_validation_trigger("embeddings", "tenant_id");
|
||||
|
||||
assert!(sql.contains("CREATE OR REPLACE FUNCTION"));
|
||||
assert!(sql.contains("TRIGGER"));
|
||||
assert!(sql.contains("BEFORE INSERT OR UPDATE"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,471 @@
|
||||
//! Input Validation for Multi-Tenancy Security
|
||||
//!
|
||||
//! Provides strict validation for tenant IDs, table names, and other identifiers
|
||||
//! to prevent SQL injection attacks.
|
||||
|
||||
use std::fmt;
|
||||
|
||||
/// Maximum length for tenant IDs
|
||||
pub const MAX_TENANT_ID_LENGTH: usize = 64;
|
||||
|
||||
/// Maximum length for identifiers (tables, schemas, partitions)
|
||||
pub const MAX_IDENTIFIER_LENGTH: usize = 63; // PostgreSQL limit
|
||||
|
||||
/// Validation error types
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum ValidationError {
|
||||
/// Identifier is empty
|
||||
Empty,
|
||||
/// Identifier is too long
|
||||
TooLong { max: usize, actual: usize },
|
||||
/// Identifier contains invalid characters
|
||||
InvalidCharacters { position: usize, char: char },
|
||||
/// Identifier doesn't start with a valid character
|
||||
InvalidStart { char: char },
|
||||
/// Identifier is a reserved word
|
||||
ReservedWord(String),
|
||||
}
|
||||
|
||||
impl fmt::Display for ValidationError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
Self::Empty => write!(f, "Identifier cannot be empty"),
|
||||
Self::TooLong { max, actual } => {
|
||||
write!(f, "Identifier too long: {} chars (max {})", actual, max)
|
||||
}
|
||||
Self::InvalidCharacters { position, char } => {
|
||||
write!(f, "Invalid character '{}' at position {}", char, position)
|
||||
}
|
||||
Self::InvalidStart { char } => {
|
||||
write!(f, "Identifier cannot start with '{}'", char)
|
||||
}
|
||||
Self::ReservedWord(word) => {
|
||||
write!(f, "Cannot use reserved word: {}", word)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for ValidationError {}
|
||||
|
||||
/// Reserved PostgreSQL words that cannot be used as identifiers
|
||||
const RESERVED_WORDS: &[&str] = &[
|
||||
"select",
|
||||
"insert",
|
||||
"update",
|
||||
"delete",
|
||||
"drop",
|
||||
"create",
|
||||
"alter",
|
||||
"grant",
|
||||
"revoke",
|
||||
"table",
|
||||
"schema",
|
||||
"index",
|
||||
"cascade",
|
||||
"restrict",
|
||||
"null",
|
||||
"true",
|
||||
"false",
|
||||
"and",
|
||||
"or",
|
||||
"not",
|
||||
"in",
|
||||
"exists",
|
||||
"between",
|
||||
"like",
|
||||
"is",
|
||||
"as",
|
||||
"from",
|
||||
"where",
|
||||
"order",
|
||||
"by",
|
||||
"group",
|
||||
"having",
|
||||
"limit",
|
||||
"offset",
|
||||
"join",
|
||||
"inner",
|
||||
"outer",
|
||||
"left",
|
||||
"right",
|
||||
"cross",
|
||||
"on",
|
||||
"using",
|
||||
"union",
|
||||
"except",
|
||||
"intersect",
|
||||
"all",
|
||||
"distinct",
|
||||
"case",
|
||||
"when",
|
||||
"then",
|
||||
"else",
|
||||
"end",
|
||||
"cast",
|
||||
"coalesce",
|
||||
"nullif",
|
||||
"primary",
|
||||
"key",
|
||||
"foreign",
|
||||
"references",
|
||||
"unique",
|
||||
"check",
|
||||
"default",
|
||||
"constraint",
|
||||
"trigger",
|
||||
"function",
|
||||
"procedure",
|
||||
"view",
|
||||
"sequence",
|
||||
"type",
|
||||
"domain",
|
||||
"role",
|
||||
"user",
|
||||
"database",
|
||||
"tablespace",
|
||||
"extension",
|
||||
"operator",
|
||||
"policy",
|
||||
"rule",
|
||||
"security",
|
||||
"definer",
|
||||
"invoker",
|
||||
];
|
||||
|
||||
/// Validate a tenant ID
|
||||
///
|
||||
/// Tenant IDs must:
|
||||
/// - Be 1-64 characters long
|
||||
/// - Start with a letter or underscore
|
||||
/// - Contain only letters, numbers, underscores, and hyphens
|
||||
/// - Not be a reserved SQL keyword
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use ruvector_postgres::tenancy::validation::validate_tenant_id;
|
||||
///
|
||||
/// assert!(validate_tenant_id("acme-corp").is_ok());
|
||||
/// assert!(validate_tenant_id("tenant_123").is_ok());
|
||||
/// assert!(validate_tenant_id("DROP TABLE users;--").is_err());
|
||||
/// ```
|
||||
pub fn validate_tenant_id(tenant_id: &str) -> Result<(), ValidationError> {
|
||||
// Check empty
|
||||
if tenant_id.is_empty() {
|
||||
return Err(ValidationError::Empty);
|
||||
}
|
||||
|
||||
// Check length
|
||||
if tenant_id.len() > MAX_TENANT_ID_LENGTH {
|
||||
return Err(ValidationError::TooLong {
|
||||
max: MAX_TENANT_ID_LENGTH,
|
||||
actual: tenant_id.len(),
|
||||
});
|
||||
}
|
||||
|
||||
// Check first character (must be letter or underscore)
|
||||
let first_char = tenant_id.chars().next().unwrap();
|
||||
if !first_char.is_ascii_alphabetic() && first_char != '_' {
|
||||
return Err(ValidationError::InvalidStart { char: first_char });
|
||||
}
|
||||
|
||||
// Check all characters
|
||||
for (i, c) in tenant_id.chars().enumerate() {
|
||||
if !is_valid_identifier_char(c) && c != '-' {
|
||||
return Err(ValidationError::InvalidCharacters {
|
||||
position: i,
|
||||
char: c,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Check reserved words (lowercase comparison)
|
||||
let lower = tenant_id.to_lowercase();
|
||||
if RESERVED_WORDS.contains(&lower.as_str()) {
|
||||
return Err(ValidationError::ReservedWord(tenant_id.to_string()));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Validate a SQL identifier (table name, schema name, column name)
|
||||
///
|
||||
/// Identifiers must:
|
||||
/// - Be 1-63 characters long (PostgreSQL limit)
|
||||
/// - Start with a letter or underscore
|
||||
/// - Contain only letters, numbers, and underscores
|
||||
/// - Not be a reserved SQL keyword
|
||||
pub fn validate_identifier(identifier: &str) -> Result<(), ValidationError> {
|
||||
// Check empty
|
||||
if identifier.is_empty() {
|
||||
return Err(ValidationError::Empty);
|
||||
}
|
||||
|
||||
// Check length
|
||||
if identifier.len() > MAX_IDENTIFIER_LENGTH {
|
||||
return Err(ValidationError::TooLong {
|
||||
max: MAX_IDENTIFIER_LENGTH,
|
||||
actual: identifier.len(),
|
||||
});
|
||||
}
|
||||
|
||||
// Check first character (must be letter or underscore)
|
||||
let first_char = identifier.chars().next().unwrap();
|
||||
if !first_char.is_ascii_alphabetic() && first_char != '_' {
|
||||
return Err(ValidationError::InvalidStart { char: first_char });
|
||||
}
|
||||
|
||||
// Check all characters (stricter than tenant_id - no hyphens)
|
||||
for (i, c) in identifier.chars().enumerate() {
|
||||
if !is_valid_identifier_char(c) {
|
||||
return Err(ValidationError::InvalidCharacters {
|
||||
position: i,
|
||||
char: c,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Check reserved words (lowercase comparison)
|
||||
let lower = identifier.to_lowercase();
|
||||
if RESERVED_WORDS.contains(&lower.as_str()) {
|
||||
return Err(ValidationError::ReservedWord(identifier.to_string()));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Check if a character is valid for SQL identifiers
|
||||
#[inline]
|
||||
fn is_valid_identifier_char(c: char) -> bool {
|
||||
c.is_ascii_alphanumeric() || c == '_'
|
||||
}
|
||||
|
||||
/// Sanitize a tenant ID for use in partition/schema names
|
||||
///
|
||||
/// Converts hyphens and dots to underscores, validates the result.
|
||||
pub fn sanitize_for_identifier(input: &str) -> Result<String, ValidationError> {
|
||||
// First validate the input as a tenant ID
|
||||
validate_tenant_id(input)?;
|
||||
|
||||
// Convert to valid identifier format
|
||||
let sanitized = input.replace('-', "_").replace('.', "_");
|
||||
|
||||
// Validate the result as an identifier
|
||||
validate_identifier(&sanitized)?;
|
||||
|
||||
Ok(sanitized)
|
||||
}
|
||||
|
||||
/// Escape a string for use in SQL string literals
|
||||
///
|
||||
/// This function properly escapes single quotes by doubling them.
|
||||
/// Use this only for string values, NOT for identifiers!
|
||||
pub fn escape_string_literal(input: &str) -> String {
|
||||
input.replace('\'', "''")
|
||||
}
|
||||
|
||||
/// Quote an identifier for safe use in SQL
|
||||
///
|
||||
/// This function wraps the identifier in double quotes and escapes
|
||||
/// any double quotes within it. This is the PostgreSQL-safe way to
|
||||
/// use dynamic identifiers.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use ruvector_postgres::tenancy::validation::quote_identifier;
|
||||
///
|
||||
/// assert_eq!(quote_identifier("my_table"), "\"my_table\"");
|
||||
/// assert_eq!(quote_identifier("weird\"name"), "\"weird\"\"name\"");
|
||||
/// ```
|
||||
pub fn quote_identifier(identifier: &str) -> String {
|
||||
format!("\"{}\"", identifier.replace('"', "\"\""))
|
||||
}
|
||||
|
||||
/// Validate and quote a partition name
|
||||
///
|
||||
/// Returns a safely quoted partition name or an error.
|
||||
pub fn safe_partition_name(tenant_id: &str, parent_table: &str) -> Result<String, ValidationError> {
|
||||
// Validate both inputs
|
||||
validate_tenant_id(tenant_id)?;
|
||||
validate_identifier(parent_table)?;
|
||||
|
||||
// Create sanitized partition name
|
||||
let sanitized_tenant = sanitize_for_identifier(tenant_id)?;
|
||||
let partition_name = format!("{}_{}", parent_table, sanitized_tenant);
|
||||
|
||||
// Validate the combined name
|
||||
validate_identifier(&partition_name)?;
|
||||
|
||||
Ok(partition_name)
|
||||
}
|
||||
|
||||
/// Validate and quote a schema name
|
||||
pub fn safe_schema_name(tenant_id: &str) -> Result<String, ValidationError> {
|
||||
validate_tenant_id(tenant_id)?;
|
||||
let sanitized = sanitize_for_identifier(tenant_id)?;
|
||||
let schema_name = format!("tenant_{}", sanitized);
|
||||
validate_identifier(&schema_name)?;
|
||||
Ok(schema_name)
|
||||
}
|
||||
|
||||
/// Validate an IP address format (basic check)
|
||||
pub fn validate_ip_address(ip: &str) -> bool {
|
||||
// Allow IPv4 and IPv6
|
||||
ip.parse::<std::net::IpAddr>().is_ok()
|
||||
}
|
||||
|
||||
/// Sanitize an IP address or return None if invalid
|
||||
pub fn sanitize_ip_address(ip: Option<&str>) -> Option<String> {
|
||||
ip.and_then(|i| {
|
||||
if validate_ip_address(i) {
|
||||
Some(i.to_string())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_valid_tenant_ids() {
|
||||
assert!(validate_tenant_id("acme-corp").is_ok());
|
||||
assert!(validate_tenant_id("tenant_123").is_ok());
|
||||
assert!(validate_tenant_id("my-tenant-id").is_ok());
|
||||
assert!(validate_tenant_id("_private").is_ok());
|
||||
assert!(validate_tenant_id("a").is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_tenant_ids() {
|
||||
// Empty
|
||||
assert!(matches!(
|
||||
validate_tenant_id(""),
|
||||
Err(ValidationError::Empty)
|
||||
));
|
||||
|
||||
// Too long
|
||||
let long = "a".repeat(100);
|
||||
assert!(matches!(
|
||||
validate_tenant_id(&long),
|
||||
Err(ValidationError::TooLong { .. })
|
||||
));
|
||||
|
||||
// Invalid start
|
||||
assert!(matches!(
|
||||
validate_tenant_id("123tenant"),
|
||||
Err(ValidationError::InvalidStart { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
validate_tenant_id("-tenant"),
|
||||
Err(ValidationError::InvalidStart { .. })
|
||||
));
|
||||
|
||||
// Invalid characters
|
||||
assert!(matches!(
|
||||
validate_tenant_id("tenant'id"),
|
||||
Err(ValidationError::InvalidCharacters { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
validate_tenant_id("tenant;drop"),
|
||||
Err(ValidationError::InvalidCharacters { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
validate_tenant_id("tenant id"),
|
||||
Err(ValidationError::InvalidCharacters { .. })
|
||||
));
|
||||
|
||||
// Reserved words
|
||||
assert!(matches!(
|
||||
validate_tenant_id("select"),
|
||||
Err(ValidationError::ReservedWord(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
validate_tenant_id("DROP"),
|
||||
Err(ValidationError::ReservedWord(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sql_injection_attempts() {
|
||||
// Common SQL injection patterns
|
||||
assert!(validate_tenant_id("'; DROP TABLE users;--").is_err());
|
||||
assert!(validate_tenant_id("tenant' OR '1'='1").is_err());
|
||||
assert!(validate_tenant_id("tenant\"; DELETE FROM").is_err());
|
||||
assert!(validate_tenant_id("tenant$(whoami)").is_err());
|
||||
assert!(validate_tenant_id("tenant`id`").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_valid_identifiers() {
|
||||
assert!(validate_identifier("my_table").is_ok());
|
||||
assert!(validate_identifier("embeddings").is_ok());
|
||||
assert!(validate_identifier("_private_table").is_ok());
|
||||
assert!(validate_identifier("table123").is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_identifiers() {
|
||||
// Hyphens not allowed in identifiers
|
||||
assert!(validate_identifier("my-table").is_err());
|
||||
|
||||
// Special characters
|
||||
assert!(validate_identifier("my.table").is_err());
|
||||
assert!(validate_identifier("my table").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_for_identifier() {
|
||||
assert_eq!(sanitize_for_identifier("acme-corp").unwrap(), "acme_corp");
|
||||
assert_eq!(
|
||||
sanitize_for_identifier("my.tenant.id").unwrap(),
|
||||
"my_tenant_id"
|
||||
);
|
||||
assert_eq!(sanitize_for_identifier("simple").unwrap(), "simple");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quote_identifier() {
|
||||
assert_eq!(quote_identifier("my_table"), "\"my_table\"");
|
||||
assert_eq!(quote_identifier("weird\"name"), "\"weird\"\"name\"");
|
||||
assert_eq!(quote_identifier("UPPERCASE"), "\"UPPERCASE\"");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_escape_string_literal() {
|
||||
assert_eq!(escape_string_literal("hello"), "hello");
|
||||
assert_eq!(escape_string_literal("it's"), "it''s");
|
||||
assert_eq!(escape_string_literal("O'Brien's"), "O''Brien''s");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_partition_name() {
|
||||
assert_eq!(
|
||||
safe_partition_name("acme-corp", "embeddings").unwrap(),
|
||||
"embeddings_acme_corp"
|
||||
);
|
||||
assert!(safe_partition_name("'; DROP TABLE", "embeddings").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_schema_name() {
|
||||
assert_eq!(safe_schema_name("acme-corp").unwrap(), "tenant_acme_corp");
|
||||
assert!(safe_schema_name("'; DROP SCHEMA").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validate_ip_address() {
|
||||
assert!(validate_ip_address("192.168.1.1"));
|
||||
assert!(validate_ip_address("10.0.0.1"));
|
||||
assert!(validate_ip_address("::1"));
|
||||
assert!(validate_ip_address("2001:db8::1"));
|
||||
|
||||
assert!(!validate_ip_address("not-an-ip"));
|
||||
assert!(!validate_ip_address("192.168.1.256"));
|
||||
assert!(!validate_ip_address("'; DROP TABLE"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user