prax_query/tenant/
strategy.rs

1//! Tenant isolation strategies.
2
3use std::collections::HashSet;
4
5use crate::error::{QueryError, QueryResult};
6use crate::sql::is_valid_sql_identifier;
7
8/// The isolation strategy for multi-tenancy.
9#[derive(Debug, Clone)]
10pub enum IsolationStrategy {
11    /// Row-level security: all tenants share tables, filtered by column.
12    RowLevel(RowLevelConfig),
13    /// Schema-based: each tenant has their own schema.
14    Schema(SchemaConfig),
15    /// Database-based: each tenant has their own database.
16    Database(DatabaseConfig),
17    /// Hybrid: combination of strategies (e.g., schema + row-level).
18    Hybrid(Box<IsolationStrategy>, Box<IsolationStrategy>),
19}
20
21impl IsolationStrategy {
22    /// Create a row-level isolation strategy.
23    pub fn row_level(column: impl Into<String>) -> Self {
24        Self::RowLevel(RowLevelConfig::new(column))
25    }
26
27    /// Create a schema-based isolation strategy.
28    pub fn schema_based() -> Self {
29        Self::Schema(SchemaConfig::default())
30    }
31
32    /// Create a database-based isolation strategy.
33    pub fn database_based() -> Self {
34        Self::Database(DatabaseConfig::default())
35    }
36
37    /// Check if this is row-level isolation.
38    pub fn is_row_level(&self) -> bool {
39        matches!(self, Self::RowLevel(_))
40    }
41
42    /// Check if this is schema-based isolation.
43    pub fn is_schema_based(&self) -> bool {
44        matches!(self, Self::Schema(_))
45    }
46
47    /// Check if this is database-based isolation.
48    pub fn is_database_based(&self) -> bool {
49        matches!(self, Self::Database(_))
50    }
51
52    /// Get the row-level config if applicable.
53    pub fn row_level_config(&self) -> Option<&RowLevelConfig> {
54        match self {
55            Self::RowLevel(config) => Some(config),
56            Self::Hybrid(a, b) => a.row_level_config().or_else(|| b.row_level_config()),
57            _ => None,
58        }
59    }
60
61    /// Get the schema config if applicable.
62    pub fn schema_config(&self) -> Option<&SchemaConfig> {
63        match self {
64            Self::Schema(config) => Some(config),
65            Self::Hybrid(a, b) => a.schema_config().or_else(|| b.schema_config()),
66            _ => None,
67        }
68    }
69
70    /// Get the database config if applicable.
71    pub fn database_config(&self) -> Option<&DatabaseConfig> {
72        match self {
73            Self::Database(config) => Some(config),
74            Self::Hybrid(a, b) => a.database_config().or_else(|| b.database_config()),
75            _ => None,
76        }
77    }
78}
79
80/// Configuration for row-level tenant isolation.
81#[derive(Debug, Clone)]
82pub struct RowLevelConfig {
83    /// The column name that stores the tenant ID.
84    pub column: String,
85    /// The column type (for type-safe comparisons).
86    pub column_type: ColumnType,
87    /// Tables that should be excluded from tenant filtering.
88    pub excluded_tables: HashSet<String>,
89    /// Tables that are shared across all tenants.
90    pub shared_tables: HashSet<String>,
91    /// Whether to automatically add tenant_id to INSERT statements.
92    pub auto_insert: bool,
93    /// Whether to validate tenant_id on UPDATE/DELETE.
94    pub validate_writes: bool,
95    /// Whether to use database-level RLS (PostgreSQL).
96    pub use_database_rls: bool,
97}
98
99impl RowLevelConfig {
100    /// Create a new row-level config with the given column name.
101    pub fn new(column: impl Into<String>) -> Self {
102        Self {
103            column: column.into(),
104            column_type: ColumnType::String,
105            excluded_tables: HashSet::new(),
106            shared_tables: HashSet::new(),
107            auto_insert: true,
108            validate_writes: true,
109            use_database_rls: false,
110        }
111    }
112
113    /// Set the column type.
114    pub fn with_column_type(mut self, column_type: ColumnType) -> Self {
115        self.column_type = column_type;
116        self
117    }
118
119    /// Exclude a table from tenant filtering.
120    pub fn exclude_table(mut self, table: impl Into<String>) -> Self {
121        self.excluded_tables.insert(table.into());
122        self
123    }
124
125    /// Mark a table as shared (no tenant filtering).
126    pub fn shared_table(mut self, table: impl Into<String>) -> Self {
127        self.shared_tables.insert(table.into());
128        self
129    }
130
131    /// Disable automatic tenant_id insertion.
132    pub fn without_auto_insert(mut self) -> Self {
133        self.auto_insert = false;
134        self
135    }
136
137    /// Disable write validation.
138    pub fn without_write_validation(mut self) -> Self {
139        self.validate_writes = false;
140        self
141    }
142
143    /// Enable PostgreSQL database-level RLS.
144    pub fn with_database_rls(mut self) -> Self {
145        self.use_database_rls = true;
146        self
147    }
148
149    /// Check if a table should be filtered.
150    pub fn should_filter(&self, table: &str) -> bool {
151        !self.excluded_tables.contains(table) && !self.shared_tables.contains(table)
152    }
153}
154
155/// The type of the tenant column.
156#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
157pub enum ColumnType {
158    /// String/VARCHAR/TEXT column.
159    #[default]
160    String,
161    /// UUID column.
162    Uuid,
163    /// Integer column.
164    Integer,
165    /// BigInt column.
166    BigInt,
167}
168
169impl ColumnType {
170    /// Get the SQL placeholder for this column type.
171    pub fn placeholder(&self, index: usize) -> String {
172        format!("${}", index)
173    }
174
175    /// Format a value for this column type.
176    ///
177    /// # Hazard
178    ///
179    /// `Uuid`, `Integer`, and `BigInt` values are interpolated into the
180    /// resulting SQL fragment **without any validation or escaping**. A
181    /// malicious value (e.g. a raw HTTP header used as a tenant id) can
182    /// inject arbitrary SQL. Use [`Self::try_format_value`] instead; the
183    /// tenant middleware also validates tenant ids at the middleware
184    /// boundary.
185    #[deprecated(
186        since = "0.11.0",
187        note = "interpolates `value` into SQL without validation; use `try_format_value`, \
188                which rejects invalid UUID/integer values with a `QueryError`"
189    )]
190    pub fn format_value(&self, value: &str) -> String {
191        match self {
192            Self::String => format!("'{}'", value.replace('\'', "''")),
193            Self::Uuid => format!("'{}'::uuid", value),
194            Self::Integer | Self::BigInt => value.to_string(),
195        }
196    }
197
198    /// Format a value for this column type, validating it first.
199    ///
200    /// - `String` values must match the conservative tenant id whitelist
201    ///   `^[A-Za-z0-9_\-\:.@]+$` — `''`-doubling alone is not sufficient on
202    ///   backends that honor backslash escapes (MySQL), so string ids are
203    ///   fail-closed instead of escaped.
204    /// - `Uuid` values must parse via [`uuid::Uuid::parse_str`].
205    /// - `Integer`/`BigInt` values must parse as an `i64`.
206    ///
207    /// Invalid input is rejected with a [`QueryError`] instead of being
208    /// interpolated into SQL.
209    pub fn try_format_value(&self, value: &str) -> QueryResult<String> {
210        match self {
211            Self::String => {
212                if is_valid_string_tenant_id(value) {
213                    Ok(format!("'{}'", value))
214                } else {
215                    Err(QueryError::invalid_input(
216                        "tenant_id",
217                        "value is outside the string tenant id whitelist \
218                         (allowed: letters, digits, `_`, `-`, `:`, `.`, `@`)",
219                    ))
220                }
221            }
222            Self::Uuid => uuid::Uuid::parse_str(value)
223                .map(|uuid| format!("'{}'::uuid", uuid))
224                .map_err(|_| QueryError::invalid_input("tenant_id", "value is not a valid UUID")),
225            Self::Integer | Self::BigInt => {
226                value.parse::<i64>().map(|n| n.to_string()).map_err(|_| {
227                    QueryError::invalid_input("tenant_id", "value is not a valid integer")
228                })
229            }
230        }
231    }
232}
233
234/// Configuration for schema-based tenant isolation.
235#[derive(Debug, Clone, Default)]
236pub struct SchemaConfig {
237    /// Prefix for tenant schema names (e.g., "tenant_" -> "tenant_acme").
238    pub schema_prefix: Option<String>,
239    /// Suffix for tenant schema names.
240    pub schema_suffix: Option<String>,
241    /// Name of the shared schema for common tables.
242    pub shared_schema: Option<String>,
243    /// Whether to create schemas automatically.
244    pub auto_create: bool,
245    /// Default schema for new tenants.
246    pub default_schema: Option<String>,
247    /// Schema search path format.
248    pub search_path_format: SearchPathFormat,
249}
250
251impl SchemaConfig {
252    /// Set the schema prefix.
253    pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
254        self.schema_prefix = Some(prefix.into());
255        self
256    }
257
258    /// Set the schema suffix.
259    pub fn with_suffix(mut self, suffix: impl Into<String>) -> Self {
260        self.schema_suffix = Some(suffix.into());
261        self
262    }
263
264    /// Set the shared schema name.
265    pub fn with_shared_schema(mut self, schema: impl Into<String>) -> Self {
266        self.shared_schema = Some(schema.into());
267        self
268    }
269
270    /// Enable auto-creation of schemas.
271    pub fn with_auto_create(mut self) -> Self {
272        self.auto_create = true;
273        self
274    }
275
276    /// Set the default schema.
277    pub fn with_default_schema(mut self, schema: impl Into<String>) -> Self {
278        self.default_schema = Some(schema.into());
279        self
280    }
281
282    /// Set the search path format.
283    pub fn with_search_path(mut self, format: SearchPathFormat) -> Self {
284        self.search_path_format = format;
285        self
286    }
287
288    /// Generate the schema name for a tenant.
289    ///
290    /// This is a pure string composition; no validation is performed. Use
291    /// [`Self::try_schema_name`] to reject tenant ids that would compose
292    /// into an invalid SQL identifier.
293    pub fn schema_name(&self, tenant_id: &str) -> String {
294        let mut name = String::new();
295        if let Some(prefix) = &self.schema_prefix {
296            name.push_str(prefix);
297        }
298        name.push_str(tenant_id);
299        if let Some(suffix) = &self.schema_suffix {
300            name.push_str(suffix);
301        }
302        name
303    }
304
305    /// Generate the schema name for a tenant, validating the composed name.
306    ///
307    /// The composed schema name must match the strict SQL identifier charset
308    /// `^[A-Za-z_][A-Za-z0-9_]*$`; otherwise a [`QueryError`] is returned so
309    /// a malicious tenant id is rejected before it can reach a SQL
310    /// statement.
311    pub fn try_schema_name(&self, tenant_id: &str) -> QueryResult<String> {
312        let name = self.schema_name(tenant_id);
313        if is_valid_schema_ident(&name) {
314            Ok(name)
315        } else {
316            Err(QueryError::invalid_input(
317                "tenant_id",
318                format!("schema name `{}` is not a valid SQL identifier", name),
319            ))
320        }
321    }
322
323    /// Generate the search_path SQL for a tenant.
324    ///
325    /// The composed tenant schema name is validated against the strict SQL
326    /// identifier charset `^[A-Za-z_][A-Za-z0-9_]*$` at the point of SQL
327    /// generation. A name outside that charset is emitted as a double-quoted
328    /// identifier (with embedded `"` doubled) so a malicious tenant id
329    /// cannot break out of the identifier and inject SQL. Use
330    /// [`Self::try_search_path`] to reject invalid tenant ids with an error
331    /// instead.
332    pub fn search_path(&self, tenant_id: &str) -> String {
333        let tenant_schema = safe_schema_ident(&self.schema_name(tenant_id));
334        let shared_schema = self.shared_schema.as_deref().map(safe_schema_ident);
335        self.search_path_sql(&tenant_schema, shared_schema.as_deref())
336    }
337
338    /// Generate the search_path SQL for a tenant, validating all schema
339    /// names first.
340    ///
341    /// Returns a [`QueryError`] if the composed tenant schema name or the
342    /// configured shared schema is not a valid SQL identifier
343    /// (`^[A-Za-z_][A-Za-z0-9_]*$`).
344    pub fn try_search_path(&self, tenant_id: &str) -> QueryResult<String> {
345        let tenant_schema = self.try_schema_name(tenant_id)?;
346        if let Some(shared) = &self.shared_schema
347            && !is_valid_schema_ident(shared)
348        {
349            return Err(QueryError::invalid_input(
350                "shared_schema",
351                format!("schema name `{}` is not a valid SQL identifier", shared),
352            ));
353        }
354        Ok(self.search_path_sql(&tenant_schema, self.shared_schema.as_deref()))
355    }
356
357    /// Build the search_path SQL from pre-validated (or safely quoted)
358    /// schema names.
359    fn search_path_sql(&self, tenant_schema: &str, shared_schema: Option<&str>) -> String {
360        match self.search_path_format {
361            SearchPathFormat::TenantOnly => {
362                format!("SET search_path TO {}", tenant_schema)
363            }
364            SearchPathFormat::TenantFirst => {
365                if let Some(shared) = shared_schema {
366                    format!("SET search_path TO {}, {}", tenant_schema, shared)
367                } else {
368                    format!("SET search_path TO {}, public", tenant_schema)
369                }
370            }
371            SearchPathFormat::SharedFirst => {
372                if let Some(shared) = shared_schema {
373                    format!("SET search_path TO {}, {}", shared, tenant_schema)
374                } else {
375                    format!("SET search_path TO public, {}", tenant_schema)
376                }
377            }
378        }
379    }
380}
381
382/// Check whether a string tenant id matches the conservative whitelist
383/// `^[A-Za-z0-9_\-\:.@]+$`. Quote-doubling alone is not sufficient
384/// validation on backends that honor backslash escapes (MySQL), so string
385/// tenant ids are restricted to characters that are safe inside a
386/// single-quoted SQL literal on every backend.
387fn is_valid_string_tenant_id(value: &str) -> bool {
388    !value.is_empty()
389        && value
390            .bytes()
391            .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'_' | b'-' | b':' | b'.' | b'@'))
392}
393
394/// Check whether a composed schema name is a strict SQL identifier
395/// (`^[A-Za-z_][A-Za-z0-9_]*$`).
396fn is_valid_schema_ident(name: &str) -> bool {
397    is_valid_sql_identifier(name)
398}
399
400/// Return the schema name unchanged when it is a strict SQL identifier;
401/// otherwise emit it as a double-quoted identifier (with embedded `"`
402/// doubled) so it cannot inject SQL into a `SET search_path` statement.
403fn safe_schema_ident(name: &str) -> String {
404    if is_valid_schema_ident(name) {
405        name.to_string()
406    } else {
407        format!("\"{}\"", name.replace('"', "\"\""))
408    }
409}
410
411/// Format for the schema search path.
412#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
413pub enum SearchPathFormat {
414    /// Only include the tenant schema.
415    TenantOnly,
416    /// Tenant schema first, then shared.
417    #[default]
418    TenantFirst,
419    /// Shared schema first, then tenant.
420    SharedFirst,
421}
422
423/// Configuration for database-based tenant isolation.
424#[derive(Debug, Clone, Default)]
425pub struct DatabaseConfig {
426    /// Prefix for tenant database names.
427    pub database_prefix: Option<String>,
428    /// Suffix for tenant database names.
429    pub database_suffix: Option<String>,
430    /// Whether to create databases automatically.
431    pub auto_create: bool,
432    /// Template database for new tenant databases.
433    pub template_database: Option<String>,
434    /// Connection pool size per tenant.
435    pub pool_size_per_tenant: usize,
436    /// Maximum number of tenant connections to keep.
437    pub max_tenant_connections: usize,
438}
439
440impl DatabaseConfig {
441    /// Set the database prefix.
442    pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
443        self.database_prefix = Some(prefix.into());
444        self
445    }
446
447    /// Set the database suffix.
448    pub fn with_suffix(mut self, suffix: impl Into<String>) -> Self {
449        self.database_suffix = Some(suffix.into());
450        self
451    }
452
453    /// Enable auto-creation of databases.
454    pub fn with_auto_create(mut self) -> Self {
455        self.auto_create = true;
456        self
457    }
458
459    /// Set the template database.
460    pub fn with_template(mut self, template: impl Into<String>) -> Self {
461        self.template_database = Some(template.into());
462        self
463    }
464
465    /// Set the pool size per tenant.
466    pub fn with_pool_size(mut self, size: usize) -> Self {
467        self.pool_size_per_tenant = size;
468        self
469    }
470
471    /// Set the maximum tenant connections.
472    pub fn with_max_connections(mut self, max: usize) -> Self {
473        self.max_tenant_connections = max;
474        self
475    }
476
477    /// Generate the database name for a tenant.
478    pub fn database_name(&self, tenant_id: &str) -> String {
479        let mut name = String::new();
480        if let Some(prefix) = &self.database_prefix {
481            name.push_str(prefix);
482        }
483        name.push_str(tenant_id);
484        if let Some(suffix) = &self.database_suffix {
485            name.push_str(suffix);
486        }
487        name
488    }
489}
490
491#[cfg(test)]
492mod tests {
493    use super::*;
494
495    #[test]
496    fn test_row_level_config() {
497        let config = RowLevelConfig::new("tenant_id")
498            .with_column_type(ColumnType::Uuid)
499            .exclude_table("audit_logs")
500            .shared_table("plans");
501
502        assert_eq!(config.column, "tenant_id");
503        assert_eq!(config.column_type, ColumnType::Uuid);
504        assert!(config.should_filter("users"));
505        assert!(!config.should_filter("audit_logs"));
506        assert!(!config.should_filter("plans"));
507    }
508
509    #[test]
510    fn test_schema_config() {
511        let config = SchemaConfig::default()
512            .with_prefix("tenant_")
513            .with_shared_schema("shared");
514
515        assert_eq!(config.schema_name("acme"), "tenant_acme");
516        assert!(config.search_path("acme").contains("tenant_acme"));
517        assert!(config.search_path("acme").contains("shared"));
518    }
519
520    #[test]
521    fn test_database_config() {
522        let config = DatabaseConfig::default()
523            .with_prefix("prax_")
524            .with_suffix("_db");
525
526        assert_eq!(config.database_name("acme"), "prax_acme_db");
527    }
528
529    #[test]
530    #[allow(deprecated)]
531    fn test_column_type_format() {
532        assert_eq!(ColumnType::String.format_value("test"), "'test'");
533        assert_eq!(
534            ColumnType::Uuid.format_value("123e4567-e89b-12d3-a456-426614174000"),
535            "'123e4567-e89b-12d3-a456-426614174000'::uuid"
536        );
537        assert_eq!(ColumnType::Integer.format_value("42"), "42");
538    }
539
540    #[test]
541    fn test_schema_config_rejects_malicious_tenant_ids() {
542        let config = SchemaConfig::default().with_prefix("tenant_");
543
544        for bad in ["acme'; DROP", "1 OR true--", "a b"] {
545            assert!(config.try_schema_name(bad).is_err());
546            assert!(config.try_search_path(bad).is_err());
547        }
548
549        assert_eq!(config.try_schema_name("acme").unwrap(), "tenant_acme");
550        assert_eq!(
551            config.try_search_path("acme").unwrap(),
552            "SET search_path TO tenant_acme, public"
553        );
554    }
555
556    #[test]
557    fn test_search_path_neutralizes_malicious_tenant_ids() {
558        let config = SchemaConfig::default().with_prefix("tenant_");
559
560        for bad in ["acme'; DROP", "1 OR true--", "a b"] {
561            // The infallible path cannot reject, so the composed name is
562            // emitted as a single double-quoted identifier and cannot break
563            // out of the statement to inject SQL.
564            assert_eq!(
565                config.search_path(bad),
566                format!(
567                    "SET search_path TO \"tenant_{}\", public",
568                    bad.replace('"', "\"\"")
569                )
570            );
571        }
572    }
573
574    #[test]
575    fn test_try_format_value_rejects_malicious_values() {
576        for bad in [
577            "acme'; DROP",
578            "1 OR true--",
579            "a b",
580            "' OR 1=1-- ",
581            "\\' OR 1=1-- ",
582            "",
583        ] {
584            // String tenant ids are fail-closed: anything outside the
585            // whitelist is rejected, not escaped.
586            assert!(ColumnType::String.try_format_value(bad).is_err());
587            assert!(ColumnType::Uuid.try_format_value(bad).is_err());
588            assert!(ColumnType::Integer.try_format_value(bad).is_err());
589            assert!(ColumnType::BigInt.try_format_value(bad).is_err());
590        }
591    }
592
593    #[test]
594    fn test_try_format_value_valid_values() {
595        // Whitelisted string tenant ids: letters, digits, `_`, `-`, `:`,
596        // `.`, `@`.
597        for good in ["test", "tenant-123", "a_b-c:d.e@f", "Tenant_01", "x"] {
598            assert_eq!(
599                ColumnType::String.try_format_value(good).unwrap(),
600                format!("'{good}'")
601            );
602        }
603        assert_eq!(
604            ColumnType::Uuid
605                .try_format_value("123e4567-e89b-12d3-a456-426614174000")
606                .unwrap(),
607            "'123e4567-e89b-12d3-a456-426614174000'::uuid"
608        );
609        assert_eq!(ColumnType::Integer.try_format_value("42").unwrap(), "42");
610        assert_eq!(ColumnType::BigInt.try_format_value("-42").unwrap(), "-42");
611    }
612}
prax_query/tenant/strategy.rs

prax_query/tenant/
strategy.rs
