prax_query/
partition.rs

1//! Table partitioning and sharding support.
2//!
3//! This module provides types for defining and managing table partitions
4//! across different database backends.
5//!
6//! # Supported Features
7//!
8//! | Feature            | PostgreSQL | MySQL | SQLite | MSSQL | MongoDB |
9//! |--------------------|------------|-------|--------|-------|---------|
10//! | Range Partitioning | ✅         | ✅    | ❌     | ✅    | ✅*     |
11//! | List Partitioning  | ✅         | ✅    | ❌     | ✅    | ❌      |
12//! | Hash Partitioning  | ✅         | ✅    | ❌     | ✅    | ✅*     |
13//! | Partition Pruning  | ✅         | ✅    | ❌     | ✅    | ✅      |
14//! | Zone Sharding      | ❌         | ❌    | ❌     | ❌    | ✅      |
15//!
16//! > *MongoDB uses sharding with range or hashed shard keys
17//!
18//! # Example Usage
19//!
20//! ```rust,ignore
21//! use prax_query::partition::{Partition, PartitionType, RangeBound};
22//!
23//! // Define a range-partitioned table
24//! let partition = Partition::builder("orders")
25//!     .range_partition()
26//!     .column("created_at")
27//!     .add_range("orders_2024_q1", RangeBound::from("2024-01-01"), RangeBound::to("2024-04-01"))
28//!     .add_range("orders_2024_q2", RangeBound::from("2024-04-01"), RangeBound::to("2024-07-01"))
29//!     .build();
30//!
31//! // Generate SQL
32//! let sql = partition.to_postgres_sql()?;
33//! ```
34
35use std::borrow::Cow;
36
37use serde::{Deserialize, Serialize};
38
39use crate::error::{QueryError, QueryResult};
40use crate::sql::DatabaseType;
41
42/// The type of partitioning strategy.
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
44pub enum PartitionType {
45    /// Partition by value ranges (e.g., date ranges).
46    Range,
47    /// Partition by specific values (e.g., country codes).
48    List,
49    /// Partition by hash of column values.
50    Hash,
51}
52
53impl PartitionType {
54    /// Convert to SQL keyword for PostgreSQL.
55    pub fn to_postgres_sql(&self) -> &'static str {
56        match self {
57            Self::Range => "RANGE",
58            Self::List => "LIST",
59            Self::Hash => "HASH",
60        }
61    }
62
63    /// Convert to SQL keyword for MySQL.
64    pub fn to_mysql_sql(&self) -> &'static str {
65        match self {
66            Self::Range => "RANGE",
67            Self::List => "LIST",
68            Self::Hash => "HASH",
69        }
70    }
71
72    /// Convert to SQL keyword for MSSQL.
73    pub fn to_mssql_sql(&self) -> &'static str {
74        match self {
75            Self::Range => "RANGE",
76            Self::List => "LIST",
77            Self::Hash => "HASH",
78        }
79    }
80}
81
82/// A bound value for range partitions.
83#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
84pub enum RangeBound {
85    /// Minimum value (MINVALUE in PostgreSQL).
86    MinValue,
87    /// Maximum value (MAXVALUE in PostgreSQL).
88    MaxValue,
89    /// A specific value.
90    Value(String),
91    /// A date value.
92    Date(String),
93    /// An integer value.
94    Int(i64),
95}
96
97impl RangeBound {
98    /// Create a specific value bound.
99    pub fn value(v: impl Into<String>) -> Self {
100        Self::Value(v.into())
101    }
102
103    /// Create a date bound.
104    pub fn date(d: impl Into<String>) -> Self {
105        Self::Date(d.into())
106    }
107
108    /// Create an integer bound.
109    pub fn int(i: i64) -> Self {
110        Self::Int(i)
111    }
112
113    /// Convert to SQL expression.
114    pub fn to_sql(&self) -> Cow<'static, str> {
115        match self {
116            Self::MinValue => Cow::Borrowed("MINVALUE"),
117            Self::MaxValue => Cow::Borrowed("MAXVALUE"),
118            Self::Value(v) => Cow::Owned(format!("'{}'", v)),
119            Self::Date(d) => Cow::Owned(format!("'{}'", d)),
120            Self::Int(i) => Cow::Owned(i.to_string()),
121        }
122    }
123}
124
125/// A range partition definition.
126#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
127pub struct RangePartitionDef {
128    /// Partition name.
129    pub name: String,
130    /// Lower bound (inclusive).
131    pub from: RangeBound,
132    /// Upper bound (exclusive).
133    pub to: RangeBound,
134    /// Tablespace (optional).
135    pub tablespace: Option<String>,
136}
137
138impl RangePartitionDef {
139    /// Create a new range partition definition.
140    pub fn new(name: impl Into<String>, from: RangeBound, to: RangeBound) -> Self {
141        Self {
142            name: name.into(),
143            from,
144            to,
145            tablespace: None,
146        }
147    }
148
149    /// Set the tablespace.
150    pub fn tablespace(mut self, tablespace: impl Into<String>) -> Self {
151        self.tablespace = Some(tablespace.into());
152        self
153    }
154}
155
156/// A list partition definition.
157#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
158pub struct ListPartitionDef {
159    /// Partition name.
160    pub name: String,
161    /// Values in this partition.
162    pub values: Vec<String>,
163    /// Tablespace (optional).
164    pub tablespace: Option<String>,
165}
166
167impl ListPartitionDef {
168    /// Create a new list partition definition.
169    pub fn new(
170        name: impl Into<String>,
171        values: impl IntoIterator<Item = impl Into<String>>,
172    ) -> Self {
173        Self {
174            name: name.into(),
175            values: values.into_iter().map(Into::into).collect(),
176            tablespace: None,
177        }
178    }
179
180    /// Set the tablespace.
181    pub fn tablespace(mut self, tablespace: impl Into<String>) -> Self {
182        self.tablespace = Some(tablespace.into());
183        self
184    }
185}
186
187/// A hash partition definition.
188#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
189pub struct HashPartitionDef {
190    /// Partition name.
191    pub name: String,
192    /// Modulus for hash partitioning.
193    pub modulus: u32,
194    /// Remainder for this partition.
195    pub remainder: u32,
196    /// Tablespace (optional).
197    pub tablespace: Option<String>,
198}
199
200impl HashPartitionDef {
201    /// Create a new hash partition definition.
202    pub fn new(name: impl Into<String>, modulus: u32, remainder: u32) -> Self {
203        Self {
204            name: name.into(),
205            modulus,
206            remainder,
207            tablespace: None,
208        }
209    }
210
211    /// Set the tablespace.
212    pub fn tablespace(mut self, tablespace: impl Into<String>) -> Self {
213        self.tablespace = Some(tablespace.into());
214        self
215    }
216}
217
218/// Partition definitions based on partition type.
219#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
220pub enum PartitionDef {
221    /// Range partitions.
222    Range(Vec<RangePartitionDef>),
223    /// List partitions.
224    List(Vec<ListPartitionDef>),
225    /// Hash partitions.
226    Hash(Vec<HashPartitionDef>),
227}
228
229/// A table partition specification.
230#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
231pub struct Partition {
232    /// Table name.
233    pub table: String,
234    /// Schema name (optional).
235    pub schema: Option<String>,
236    /// Partition type.
237    pub partition_type: PartitionType,
238    /// Partition columns.
239    pub columns: Vec<String>,
240    /// Partition definitions.
241    pub partitions: PartitionDef,
242    /// Optional comment.
243    pub comment: Option<String>,
244}
245
246impl Partition {
247    /// Create a new partition builder.
248    pub fn builder(table: impl Into<String>) -> PartitionBuilder {
249        PartitionBuilder::new(table)
250    }
251
252    /// Get the fully qualified table name.
253    pub fn qualified_table(&self) -> Cow<'_, str> {
254        match &self.schema {
255            Some(schema) => Cow::Owned(format!("{}.{}", schema, self.table)),
256            None => Cow::Borrowed(&self.table),
257        }
258    }
259
260    /// Generate PostgreSQL CREATE TABLE with partitioning.
261    pub fn to_postgres_partition_clause(&self) -> String {
262        format!(
263            "PARTITION BY {} ({})",
264            self.partition_type.to_postgres_sql(),
265            self.columns.join(", ")
266        )
267    }
268
269    /// Generate PostgreSQL CREATE TABLE for a child partition.
270    pub fn to_postgres_create_partition(&self, def: &RangePartitionDef) -> String {
271        let mut sql = format!(
272            "CREATE TABLE {} PARTITION OF {}\n    FOR VALUES FROM ({}) TO ({})",
273            def.name,
274            self.qualified_table(),
275            def.from.to_sql(),
276            def.to.to_sql()
277        );
278
279        if let Some(ref ts) = def.tablespace {
280            sql.push_str(&format!("\n    TABLESPACE {}", ts));
281        }
282
283        sql.push(';');
284        sql
285    }
286
287    /// Generate PostgreSQL CREATE TABLE for a list partition.
288    pub fn to_postgres_create_list_partition(&self, def: &ListPartitionDef) -> String {
289        let values: Vec<String> = def.values.iter().map(|v| format!("'{}'", v)).collect();
290
291        let mut sql = format!(
292            "CREATE TABLE {} PARTITION OF {}\n    FOR VALUES IN ({})",
293            def.name,
294            self.qualified_table(),
295            values.join(", ")
296        );
297
298        if let Some(ref ts) = def.tablespace {
299            sql.push_str(&format!("\n    TABLESPACE {}", ts));
300        }
301
302        sql.push(';');
303        sql
304    }
305
306    /// Generate PostgreSQL CREATE TABLE for a hash partition.
307    pub fn to_postgres_create_hash_partition(&self, def: &HashPartitionDef) -> String {
308        let mut sql = format!(
309            "CREATE TABLE {} PARTITION OF {}\n    FOR VALUES WITH (MODULUS {}, REMAINDER {})",
310            def.name,
311            self.qualified_table(),
312            def.modulus,
313            def.remainder
314        );
315
316        if let Some(ref ts) = def.tablespace {
317            sql.push_str(&format!("\n    TABLESPACE {}", ts));
318        }
319
320        sql.push(';');
321        sql
322    }
323
324    /// Generate all PostgreSQL partition creation SQL.
325    pub fn to_postgres_create_all_partitions(&self) -> Vec<String> {
326        match &self.partitions {
327            PartitionDef::Range(ranges) => ranges
328                .iter()
329                .map(|r| self.to_postgres_create_partition(r))
330                .collect(),
331            PartitionDef::List(lists) => lists
332                .iter()
333                .map(|l| self.to_postgres_create_list_partition(l))
334                .collect(),
335            PartitionDef::Hash(hashes) => hashes
336                .iter()
337                .map(|h| self.to_postgres_create_hash_partition(h))
338                .collect(),
339        }
340    }
341
342    /// Generate MySQL PARTITION BY clause.
343    pub fn to_mysql_partition_clause(&self) -> String {
344        let columns_expr = if self.columns.len() == 1 {
345            self.columns[0].clone()
346        } else {
347            format!("({})", self.columns.join(", "))
348        };
349
350        let mut sql = format!(
351            "PARTITION BY {} ({})",
352            self.partition_type.to_mysql_sql(),
353            columns_expr
354        );
355
356        // Add partition definitions inline for MySQL
357        match &self.partitions {
358            PartitionDef::Range(ranges) => {
359                sql.push_str(" (\n");
360                for (i, r) in ranges.iter().enumerate() {
361                    if i > 0 {
362                        sql.push_str(",\n");
363                    }
364                    sql.push_str(&format!(
365                        "    PARTITION {} VALUES LESS THAN ({})",
366                        r.name,
367                        r.to.to_sql()
368                    ));
369                }
370                sql.push_str("\n)");
371            }
372            PartitionDef::List(lists) => {
373                sql.push_str(" (\n");
374                for (i, l) in lists.iter().enumerate() {
375                    if i > 0 {
376                        sql.push_str(",\n");
377                    }
378                    let values: Vec<String> = l.values.iter().map(|v| format!("'{}'", v)).collect();
379                    sql.push_str(&format!(
380                        "    PARTITION {} VALUES IN ({})",
381                        l.name,
382                        values.join(", ")
383                    ));
384                }
385                sql.push_str("\n)");
386            }
387            PartitionDef::Hash(hashes) => {
388                sql.push_str(&format!(" PARTITIONS {}", hashes.len()));
389            }
390        }
391
392        sql
393    }
394
395    /// Generate MSSQL partition function and scheme.
396    pub fn to_mssql_partition_sql(&self) -> QueryResult<Vec<String>> {
397        match &self.partitions {
398            PartitionDef::Range(ranges) => {
399                let mut sqls = Vec::new();
400
401                // Create partition function
402                let boundaries: Vec<String> = ranges
403                    .iter()
404                    .filter(|r| !matches!(r.to, RangeBound::MaxValue))
405                    .map(|r| r.to.to_sql().into_owned())
406                    .collect();
407
408                let func_name = format!("{}_pf", self.table);
409                sqls.push(format!(
410                    "CREATE PARTITION FUNCTION {}(datetime2)\nAS RANGE RIGHT FOR VALUES ({});",
411                    func_name,
412                    boundaries.join(", ")
413                ));
414
415                // Create partition scheme
416                let scheme_name = format!("{}_ps", self.table);
417                let filegroups: Vec<String> =
418                    ranges.iter().map(|_| "PRIMARY".to_string()).collect();
419                sqls.push(format!(
420                    "CREATE PARTITION SCHEME {}\nAS PARTITION {}\nTO ({});",
421                    scheme_name,
422                    func_name,
423                    filegroups.join(", ")
424                ));
425
426                Ok(sqls)
427            }
428            PartitionDef::List(_) => Err(QueryError::unsupported(
429                "MSSQL uses partition functions differently for list partitioning. Consider using range partitioning.",
430            )),
431            PartitionDef::Hash(_) => Err(QueryError::unsupported(
432                "MSSQL does not directly support hash partitioning. Use a computed column with range partitioning.",
433            )),
434        }
435    }
436
437    /// Generate SQL for attaching a partition.
438    pub fn attach_partition_sql(
439        &self,
440        partition_name: &str,
441        db_type: DatabaseType,
442    ) -> QueryResult<String> {
443        match db_type {
444            DatabaseType::PostgreSQL => Ok(format!(
445                "ALTER TABLE {} ATTACH PARTITION {};",
446                self.qualified_table(),
447                partition_name
448            )),
449            DatabaseType::MySQL => Err(QueryError::unsupported(
450                "MySQL does not support ATTACH PARTITION. Use ALTER TABLE ... REORGANIZE PARTITION.",
451            )),
452            DatabaseType::SQLite => Err(QueryError::unsupported(
453                "SQLite does not support table partitioning.",
454            )),
455            DatabaseType::MSSQL => Err(QueryError::unsupported(
456                "MSSQL uses SWITCH to move partitions. Use partition switching instead.",
457            )),
458        }
459    }
460
461    /// Generate SQL for detaching a partition.
462    pub fn detach_partition_sql(
463        &self,
464        partition_name: &str,
465        db_type: DatabaseType,
466    ) -> QueryResult<String> {
467        match db_type {
468            DatabaseType::PostgreSQL => Ok(format!(
469                "ALTER TABLE {} DETACH PARTITION {};",
470                self.qualified_table(),
471                partition_name
472            )),
473            DatabaseType::MySQL => Err(QueryError::unsupported(
474                "MySQL does not support DETACH PARTITION. Drop and recreate the partition.",
475            )),
476            DatabaseType::SQLite => Err(QueryError::unsupported(
477                "SQLite does not support table partitioning.",
478            )),
479            DatabaseType::MSSQL => Err(QueryError::unsupported(
480                "MSSQL uses SWITCH to move partitions. Use partition switching instead.",
481            )),
482        }
483    }
484
485    /// Generate SQL for dropping a partition.
486    pub fn drop_partition_sql(
487        &self,
488        partition_name: &str,
489        db_type: DatabaseType,
490    ) -> QueryResult<String> {
491        match db_type {
492            DatabaseType::PostgreSQL => Ok(format!("DROP TABLE IF EXISTS {};", partition_name)),
493            DatabaseType::MySQL => Ok(format!(
494                "ALTER TABLE {} DROP PARTITION {};",
495                self.qualified_table(),
496                partition_name
497            )),
498            DatabaseType::SQLite => Err(QueryError::unsupported(
499                "SQLite does not support table partitioning.",
500            )),
501            DatabaseType::MSSQL => Ok(format!(
502                "ALTER TABLE {} DROP PARTITION {};",
503                self.qualified_table(),
504                partition_name
505            )),
506        }
507    }
508}
509
510/// Builder for creating partition specifications.
511#[derive(Debug, Clone)]
512pub struct PartitionBuilder {
513    table: String,
514    schema: Option<String>,
515    partition_type: Option<PartitionType>,
516    columns: Vec<String>,
517    range_partitions: Vec<RangePartitionDef>,
518    list_partitions: Vec<ListPartitionDef>,
519    hash_partitions: Vec<HashPartitionDef>,
520    comment: Option<String>,
521}
522
523impl PartitionBuilder {
524    /// Create a new partition builder.
525    pub fn new(table: impl Into<String>) -> Self {
526        Self {
527            table: table.into(),
528            schema: None,
529            partition_type: None,
530            columns: Vec::new(),
531            range_partitions: Vec::new(),
532            list_partitions: Vec::new(),
533            hash_partitions: Vec::new(),
534            comment: None,
535        }
536    }
537
538    /// Set the schema.
539    pub fn schema(mut self, schema: impl Into<String>) -> Self {
540        self.schema = Some(schema.into());
541        self
542    }
543
544    /// Set to range partitioning.
545    pub fn range_partition(mut self) -> Self {
546        self.partition_type = Some(PartitionType::Range);
547        self
548    }
549
550    /// Set to list partitioning.
551    pub fn list_partition(mut self) -> Self {
552        self.partition_type = Some(PartitionType::List);
553        self
554    }
555
556    /// Set to hash partitioning.
557    pub fn hash_partition(mut self) -> Self {
558        self.partition_type = Some(PartitionType::Hash);
559        self
560    }
561
562    /// Add a partition column.
563    pub fn column(mut self, column: impl Into<String>) -> Self {
564        self.columns.push(column.into());
565        self
566    }
567
568    /// Add multiple partition columns.
569    pub fn columns(mut self, columns: impl IntoIterator<Item = impl Into<String>>) -> Self {
570        self.columns.extend(columns.into_iter().map(Into::into));
571        self
572    }
573
574    /// Add a range partition.
575    pub fn add_range(mut self, name: impl Into<String>, from: RangeBound, to: RangeBound) -> Self {
576        self.range_partitions
577            .push(RangePartitionDef::new(name, from, to));
578        self
579    }
580
581    /// Add a range partition with tablespace.
582    pub fn add_range_with_tablespace(
583        mut self,
584        name: impl Into<String>,
585        from: RangeBound,
586        to: RangeBound,
587        tablespace: impl Into<String>,
588    ) -> Self {
589        self.range_partitions
590            .push(RangePartitionDef::new(name, from, to).tablespace(tablespace));
591        self
592    }
593
594    /// Add a list partition.
595    pub fn add_list(
596        mut self,
597        name: impl Into<String>,
598        values: impl IntoIterator<Item = impl Into<String>>,
599    ) -> Self {
600        self.list_partitions
601            .push(ListPartitionDef::new(name, values));
602        self
603    }
604
605    /// Add a hash partition.
606    pub fn add_hash(mut self, name: impl Into<String>, modulus: u32, remainder: u32) -> Self {
607        self.hash_partitions
608            .push(HashPartitionDef::new(name, modulus, remainder));
609        self
610    }
611
612    /// Add multiple hash partitions automatically.
613    pub fn add_hash_partitions(mut self, count: u32, name_prefix: impl Into<String>) -> Self {
614        let prefix = name_prefix.into();
615        for i in 0..count {
616            self.hash_partitions
617                .push(HashPartitionDef::new(format!("{}_{}", prefix, i), count, i));
618        }
619        self
620    }
621
622    /// Add a comment.
623    pub fn comment(mut self, comment: impl Into<String>) -> Self {
624        self.comment = Some(comment.into());
625        self
626    }
627
628    /// Build the partition specification.
629    pub fn build(self) -> QueryResult<Partition> {
630        let partition_type = self.partition_type.ok_or_else(|| {
631            QueryError::invalid_input(
632                "partition_type",
633                "Must specify partition type (range_partition, list_partition, or hash_partition)",
634            )
635        })?;
636
637        if self.columns.is_empty() {
638            return Err(QueryError::invalid_input(
639                "columns",
640                "Must specify at least one partition column",
641            ));
642        }
643
644        let partitions = match partition_type {
645            PartitionType::Range => {
646                if self.range_partitions.is_empty() {
647                    return Err(QueryError::invalid_input(
648                        "partitions",
649                        "Must define at least one range partition with add_range()",
650                    ));
651                }
652                PartitionDef::Range(self.range_partitions)
653            }
654            PartitionType::List => {
655                if self.list_partitions.is_empty() {
656                    return Err(QueryError::invalid_input(
657                        "partitions",
658                        "Must define at least one list partition with add_list()",
659                    ));
660                }
661                PartitionDef::List(self.list_partitions)
662            }
663            PartitionType::Hash => {
664                if self.hash_partitions.is_empty() {
665                    return Err(QueryError::invalid_input(
666                        "partitions",
667                        "Must define at least one hash partition with add_hash() or add_hash_partitions()",
668                    ));
669                }
670                PartitionDef::Hash(self.hash_partitions)
671            }
672        };
673
674        Ok(Partition {
675            table: self.table,
676            schema: self.schema,
677            partition_type,
678            columns: self.columns,
679            partitions,
680            comment: self.comment,
681        })
682    }
683}
684
685/// Time-based partition generation helpers.
686pub mod time_partitions {
687    use super::*;
688
689    /// Generate monthly partitions for a date range.
690    pub fn monthly_partitions(
691        table: &str,
692        column: &str,
693        start_year: i32,
694        start_month: u32,
695        count: u32,
696    ) -> PartitionBuilder {
697        let mut builder = Partition::builder(table).range_partition().column(column);
698
699        let mut year = start_year;
700        let mut month = start_month;
701
702        for _ in 0..count {
703            let from_date = format!("{:04}-{:02}-01", year, month);
704
705            // Calculate next month
706            let (next_year, next_month) = if month == 12 {
707                (year + 1, 1)
708            } else {
709                (year, month + 1)
710            };
711            let to_date = format!("{:04}-{:02}-01", next_year, next_month);
712
713            let partition_name = format!("{}_{:04}_{:02}", table, year, month);
714
715            builder = builder.add_range(
716                partition_name,
717                RangeBound::date(from_date),
718                RangeBound::date(to_date),
719            );
720
721            year = next_year;
722            month = next_month;
723        }
724
725        builder
726    }
727
728    /// Generate quarterly partitions for a date range.
729    pub fn quarterly_partitions(
730        table: &str,
731        column: &str,
732        start_year: i32,
733        count: u32,
734    ) -> PartitionBuilder {
735        let mut builder = Partition::builder(table).range_partition().column(column);
736
737        let mut year = start_year;
738        let mut quarter = 1;
739
740        for _ in 0..count {
741            let from_month = (quarter - 1) * 3 + 1;
742            let from_date = format!("{:04}-{:02}-01", year, from_month);
743
744            let (next_year, next_quarter) = if quarter == 4 {
745                (year + 1, 1)
746            } else {
747                (year, quarter + 1)
748            };
749            let to_month = (next_quarter - 1) * 3 + 1;
750            let to_date = format!("{:04}-{:02}-01", next_year, to_month);
751
752            let partition_name = format!("{}_{}q{}", table, year, quarter);
753
754            builder = builder.add_range(
755                partition_name,
756                RangeBound::date(from_date),
757                RangeBound::date(to_date),
758            );
759
760            year = next_year;
761            quarter = next_quarter;
762        }
763
764        builder
765    }
766
767    /// Generate yearly partitions.
768    pub fn yearly_partitions(
769        table: &str,
770        column: &str,
771        start_year: i32,
772        count: u32,
773    ) -> PartitionBuilder {
774        let mut builder = Partition::builder(table).range_partition().column(column);
775
776        for i in 0..count {
777            let year = start_year + i as i32;
778            let from_date = format!("{:04}-01-01", year);
779            let to_date = format!("{:04}-01-01", year + 1);
780            let partition_name = format!("{}_{}", table, year);
781
782            builder = builder.add_range(
783                partition_name,
784                RangeBound::date(from_date),
785                RangeBound::date(to_date),
786            );
787        }
788
789        builder
790    }
791}
792
793/// MongoDB sharding support.
794pub mod mongodb {
795    use serde::{Deserialize, Serialize};
796
797    /// Shard key type for MongoDB.
798    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
799    pub enum ShardKeyType {
800        /// Range-based sharding (good for range queries).
801        Range,
802        /// Hash-based sharding (better distribution).
803        Hashed,
804    }
805
806    impl ShardKeyType {
807        /// Get the MongoDB index specification value.
808        pub fn as_index_value(&self) -> serde_json::Value {
809            match self {
810                Self::Range => serde_json::json!(1),
811                Self::Hashed => serde_json::json!("hashed"),
812            }
813        }
814    }
815
816    /// A shard key definition.
817    #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
818    pub struct ShardKey {
819        /// Fields in the shard key.
820        pub fields: Vec<(String, ShardKeyType)>,
821        /// Whether the collection is unique on the shard key.
822        pub unique: bool,
823    }
824
825    impl ShardKey {
826        /// Create a new shard key builder.
827        pub fn builder() -> ShardKeyBuilder {
828            ShardKeyBuilder::new()
829        }
830
831        /// Get the shardCollection command.
832        pub fn shard_collection_command(
833            &self,
834            database: &str,
835            collection: &str,
836        ) -> serde_json::Value {
837            let mut key = serde_json::Map::new();
838            for (field, key_type) in &self.fields {
839                key.insert(field.clone(), key_type.as_index_value());
840            }
841
842            serde_json::json!({
843                "shardCollection": format!("{}.{}", database, collection),
844                "key": key,
845                "unique": self.unique
846            })
847        }
848
849        /// Get the index specification for the shard key.
850        pub fn index_spec(&self) -> serde_json::Value {
851            let mut spec = serde_json::Map::new();
852            for (field, key_type) in &self.fields {
853                spec.insert(field.clone(), key_type.as_index_value());
854            }
855            serde_json::Value::Object(spec)
856        }
857    }
858
859    /// Builder for shard keys.
860    #[derive(Debug, Clone, Default)]
861    pub struct ShardKeyBuilder {
862        fields: Vec<(String, ShardKeyType)>,
863        unique: bool,
864    }
865
866    impl ShardKeyBuilder {
867        /// Create a new shard key builder.
868        pub fn new() -> Self {
869            Self::default()
870        }
871
872        /// Add a range field to the shard key.
873        pub fn range_field(mut self, field: impl Into<String>) -> Self {
874            self.fields.push((field.into(), ShardKeyType::Range));
875            self
876        }
877
878        /// Add a hashed field to the shard key.
879        pub fn hashed_field(mut self, field: impl Into<String>) -> Self {
880            self.fields.push((field.into(), ShardKeyType::Hashed));
881            self
882        }
883
884        /// Set whether the shard key should enforce uniqueness.
885        pub fn unique(mut self, unique: bool) -> Self {
886            self.unique = unique;
887            self
888        }
889
890        /// Build the shard key.
891        pub fn build(self) -> ShardKey {
892            ShardKey {
893                fields: self.fields,
894                unique: self.unique,
895            }
896        }
897    }
898
899    /// Zone (tag-based) sharding configuration.
900    #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
901    pub struct ShardZone {
902        /// Zone name.
903        pub name: String,
904        /// Minimum shard key value.
905        pub min: serde_json::Value,
906        /// Maximum shard key value.
907        pub max: serde_json::Value,
908    }
909
910    impl ShardZone {
911        /// Create a new shard zone.
912        pub fn new(
913            name: impl Into<String>,
914            min: serde_json::Value,
915            max: serde_json::Value,
916        ) -> Self {
917            Self {
918                name: name.into(),
919                min,
920                max,
921            }
922        }
923
924        /// Get the updateZoneKeyRange command.
925        pub fn update_zone_key_range_command(&self, namespace: &str) -> serde_json::Value {
926            serde_json::json!({
927                "updateZoneKeyRange": namespace,
928                "min": self.min,
929                "max": self.max,
930                "zone": self.name
931            })
932        }
933
934        /// Get the addShardToZone command.
935        pub fn add_shard_to_zone_command(&self, shard: &str) -> serde_json::Value {
936            serde_json::json!({
937                "addShardToZone": shard,
938                "zone": self.name
939            })
940        }
941    }
942
943    /// Builder for zone sharding configuration.
944    #[derive(Debug, Clone, Default)]
945    pub struct ZoneShardingBuilder {
946        zones: Vec<ShardZone>,
947        shard_assignments: Vec<(String, String)>, // (shard_name, zone_name)
948    }
949
950    impl ZoneShardingBuilder {
951        /// Create a new zone sharding builder.
952        pub fn new() -> Self {
953            Self::default()
954        }
955
956        /// Add a zone.
957        pub fn add_zone(
958            mut self,
959            name: impl Into<String>,
960            min: serde_json::Value,
961            max: serde_json::Value,
962        ) -> Self {
963            self.zones.push(ShardZone::new(name, min, max));
964            self
965        }
966
967        /// Assign a shard to a zone.
968        pub fn assign_shard(mut self, shard: impl Into<String>, zone: impl Into<String>) -> Self {
969            self.shard_assignments.push((shard.into(), zone.into()));
970            self
971        }
972
973        /// Get all configuration commands.
974        pub fn build_commands(&self, namespace: &str) -> Vec<serde_json::Value> {
975            let mut commands = Vec::new();
976
977            // Add shard to zone assignments
978            for (shard, zone) in &self.shard_assignments {
979                commands.push(serde_json::json!({
980                    "addShardToZone": shard,
981                    "zone": zone
982                }));
983            }
984
985            // Add zone key ranges
986            for zone in &self.zones {
987                commands.push(zone.update_zone_key_range_command(namespace));
988            }
989
990            commands
991        }
992    }
993
994    /// Helper to create a shard key.
995    pub fn shard_key() -> ShardKeyBuilder {
996        ShardKeyBuilder::new()
997    }
998
999    /// Helper to create zone sharding configuration.
1000    pub fn zone_sharding() -> ZoneShardingBuilder {
1001        ZoneShardingBuilder::new()
1002    }
1003}
1004
1005#[cfg(test)]
1006mod tests {
1007    use super::*;
1008
1009    #[test]
1010    fn test_range_partition_builder() {
1011        let partition = Partition::builder("orders")
1012            .schema("sales")
1013            .range_partition()
1014            .column("created_at")
1015            .add_range(
1016                "orders_2024_q1",
1017                RangeBound::date("2024-01-01"),
1018                RangeBound::date("2024-04-01"),
1019            )
1020            .add_range(
1021                "orders_2024_q2",
1022                RangeBound::date("2024-04-01"),
1023                RangeBound::date("2024-07-01"),
1024            )
1025            .build()
1026            .unwrap();
1027
1028        assert_eq!(partition.table, "orders");
1029        assert_eq!(partition.partition_type, PartitionType::Range);
1030        assert_eq!(partition.columns, vec!["created_at"]);
1031    }
1032
1033    #[test]
1034    fn test_postgres_partition_clause() {
1035        let partition = Partition::builder("orders")
1036            .range_partition()
1037            .column("created_at")
1038            .add_range(
1039                "orders_2024",
1040                RangeBound::date("2024-01-01"),
1041                RangeBound::date("2025-01-01"),
1042            )
1043            .build()
1044            .unwrap();
1045
1046        let clause = partition.to_postgres_partition_clause();
1047        assert_eq!(clause, "PARTITION BY RANGE (created_at)");
1048    }
1049
1050    #[test]
1051    fn test_postgres_create_partition() {
1052        let partition = Partition::builder("orders")
1053            .range_partition()
1054            .column("created_at")
1055            .add_range(
1056                "orders_2024",
1057                RangeBound::date("2024-01-01"),
1058                RangeBound::date("2025-01-01"),
1059            )
1060            .build()
1061            .unwrap();
1062
1063        let sqls = partition.to_postgres_create_all_partitions();
1064        assert_eq!(sqls.len(), 1);
1065        assert!(sqls[0].contains("CREATE TABLE orders_2024 PARTITION OF orders"));
1066        assert!(sqls[0].contains("FOR VALUES FROM ('2024-01-01') TO ('2025-01-01')"));
1067    }
1068
1069    #[test]
1070    fn test_list_partition() {
1071        let partition = Partition::builder("users")
1072            .list_partition()
1073            .column("country")
1074            .add_list("users_us", ["US", "USA"])
1075            .add_list("users_eu", ["DE", "FR", "GB", "IT"])
1076            .build()
1077            .unwrap();
1078
1079        assert_eq!(partition.partition_type, PartitionType::List);
1080
1081        let sqls = partition.to_postgres_create_all_partitions();
1082        assert_eq!(sqls.len(), 2);
1083        assert!(sqls[0].contains("FOR VALUES IN"));
1084    }
1085
1086    #[test]
1087    fn test_hash_partition() {
1088        let partition = Partition::builder("events")
1089            .hash_partition()
1090            .column("user_id")
1091            .add_hash_partitions(4, "events")
1092            .build()
1093            .unwrap();
1094
1095        assert_eq!(partition.partition_type, PartitionType::Hash);
1096
1097        let sqls = partition.to_postgres_create_all_partitions();
1098        assert_eq!(sqls.len(), 4);
1099        assert!(sqls[0].contains("MODULUS 4"));
1100        assert!(sqls[0].contains("REMAINDER 0"));
1101    }
1102
1103    #[test]
1104    fn test_mysql_partition_clause() {
1105        let partition = Partition::builder("orders")
1106            .range_partition()
1107            .column("created_at")
1108            .add_range(
1109                "p2024",
1110                RangeBound::MinValue,
1111                RangeBound::date("2025-01-01"),
1112            )
1113            .add_range(
1114                "p_future",
1115                RangeBound::date("2025-01-01"),
1116                RangeBound::MaxValue,
1117            )
1118            .build()
1119            .unwrap();
1120
1121        let clause = partition.to_mysql_partition_clause();
1122        assert!(clause.contains("PARTITION BY RANGE"));
1123        assert!(clause.contains("PARTITION p2024"));
1124        assert!(clause.contains("PARTITION p_future"));
1125    }
1126
1127    #[test]
1128    fn test_detach_partition() {
1129        let partition = Partition::builder("orders")
1130            .range_partition()
1131            .column("created_at")
1132            .add_range(
1133                "orders_2024",
1134                RangeBound::date("2024-01-01"),
1135                RangeBound::date("2025-01-01"),
1136            )
1137            .build()
1138            .unwrap();
1139
1140        let sql = partition
1141            .detach_partition_sql("orders_2024", DatabaseType::PostgreSQL)
1142            .unwrap();
1143        assert_eq!(sql, "ALTER TABLE orders DETACH PARTITION orders_2024;");
1144    }
1145
1146    #[test]
1147    fn test_drop_partition() {
1148        let partition = Partition::builder("orders")
1149            .range_partition()
1150            .column("created_at")
1151            .add_range(
1152                "orders_2024",
1153                RangeBound::date("2024-01-01"),
1154                RangeBound::date("2025-01-01"),
1155            )
1156            .build()
1157            .unwrap();
1158
1159        let pg_sql = partition
1160            .drop_partition_sql("orders_2024", DatabaseType::PostgreSQL)
1161            .unwrap();
1162        assert_eq!(pg_sql, "DROP TABLE IF EXISTS orders_2024;");
1163
1164        let mysql_sql = partition
1165            .drop_partition_sql("orders_2024", DatabaseType::MySQL)
1166            .unwrap();
1167        assert!(mysql_sql.contains("DROP PARTITION"));
1168    }
1169
1170    #[test]
1171    fn test_missing_partition_type() {
1172        let result = Partition::builder("orders")
1173            .column("created_at")
1174            .add_range("p1", RangeBound::MinValue, RangeBound::MaxValue)
1175            .build();
1176
1177        assert!(result.is_err());
1178    }
1179
1180    #[test]
1181    fn test_missing_columns() {
1182        let result = Partition::builder("orders")
1183            .range_partition()
1184            .add_range("p1", RangeBound::MinValue, RangeBound::MaxValue)
1185            .build();
1186
1187        assert!(result.is_err());
1188    }
1189
1190    mod time_partition_tests {
1191        use super::super::time_partitions;
1192
1193        #[test]
1194        fn test_monthly_partitions() {
1195            let builder = time_partitions::monthly_partitions("orders", "created_at", 2024, 1, 3);
1196            let partition = builder.build().unwrap();
1197
1198            if let super::PartitionDef::Range(ranges) = &partition.partitions {
1199                assert_eq!(ranges.len(), 3);
1200                assert_eq!(ranges[0].name, "orders_2024_01");
1201                assert_eq!(ranges[1].name, "orders_2024_02");
1202                assert_eq!(ranges[2].name, "orders_2024_03");
1203            } else {
1204                panic!("Expected range partitions");
1205            }
1206        }
1207
1208        #[test]
1209        fn test_quarterly_partitions() {
1210            let builder = time_partitions::quarterly_partitions("sales", "order_date", 2024, 4);
1211            let partition = builder.build().unwrap();
1212
1213            if let super::PartitionDef::Range(ranges) = &partition.partitions {
1214                assert_eq!(ranges.len(), 4);
1215                assert_eq!(ranges[0].name, "sales_2024q1");
1216                assert_eq!(ranges[3].name, "sales_2024q4");
1217            } else {
1218                panic!("Expected range partitions");
1219            }
1220        }
1221
1222        #[test]
1223        fn test_yearly_partitions() {
1224            let builder = time_partitions::yearly_partitions("logs", "timestamp", 2020, 5);
1225            let partition = builder.build().unwrap();
1226
1227            if let super::PartitionDef::Range(ranges) = &partition.partitions {
1228                assert_eq!(ranges.len(), 5);
1229                assert_eq!(ranges[0].name, "logs_2020");
1230                assert_eq!(ranges[4].name, "logs_2024");
1231            } else {
1232                panic!("Expected range partitions");
1233            }
1234        }
1235    }
1236
1237    mod mongodb_tests {
1238        use super::super::mongodb::*;
1239
1240        #[test]
1241        fn test_shard_key_builder() {
1242            let key = shard_key()
1243                .hashed_field("tenant_id")
1244                .range_field("created_at")
1245                .unique(false)
1246                .build();
1247
1248            assert_eq!(key.fields.len(), 2);
1249            assert_eq!(
1250                key.fields[0],
1251                ("tenant_id".to_string(), ShardKeyType::Hashed)
1252            );
1253            assert_eq!(
1254                key.fields[1],
1255                ("created_at".to_string(), ShardKeyType::Range)
1256            );
1257        }
1258
1259        #[test]
1260        fn test_shard_collection_command() {
1261            let key = shard_key().hashed_field("user_id").build();
1262
1263            let cmd = key.shard_collection_command("mydb", "users");
1264            assert_eq!(cmd["shardCollection"], "mydb.users");
1265            assert_eq!(cmd["key"]["user_id"], "hashed");
1266        }
1267
1268        #[test]
1269        fn test_zone_sharding() {
1270            let builder = zone_sharding()
1271                .add_zone(
1272                    "US",
1273                    serde_json::json!({"region": "US"}),
1274                    serde_json::json!({"region": "US~"}),
1275                )
1276                .add_zone(
1277                    "EU",
1278                    serde_json::json!({"region": "EU"}),
1279                    serde_json::json!({"region": "EU~"}),
1280                )
1281                .assign_shard("shard0", "US")
1282                .assign_shard("shard1", "EU");
1283
1284            let commands = builder.build_commands("mydb.users");
1285            assert_eq!(commands.len(), 4); // 2 shard assignments + 2 zone ranges
1286        }
1287
1288        #[test]
1289        fn test_shard_key_index_spec() {
1290            let key = shard_key()
1291                .range_field("tenant_id")
1292                .range_field("created_at")
1293                .build();
1294
1295            let spec = key.index_spec();
1296            assert_eq!(spec["tenant_id"], 1);
1297            assert_eq!(spec["created_at"], 1);
1298        }
1299    }
1300}
prax_query/partition.rs

prax_query/
partition.rs
