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(
141        name: impl Into<String>,
142        from: RangeBound,
143        to: RangeBound,
144    ) -> Self {
145        Self {
146            name: name.into(),
147            from,
148            to,
149            tablespace: None,
150        }
151    }
152
153    /// Set the tablespace.
154    pub fn tablespace(mut self, tablespace: impl Into<String>) -> Self {
155        self.tablespace = Some(tablespace.into());
156        self
157    }
158}
159
160/// A list partition definition.
161#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
162pub struct ListPartitionDef {
163    /// Partition name.
164    pub name: String,
165    /// Values in this partition.
166    pub values: Vec<String>,
167    /// Tablespace (optional).
168    pub tablespace: Option<String>,
169}
170
171impl ListPartitionDef {
172    /// Create a new list partition definition.
173    pub fn new(name: impl Into<String>, values: impl IntoIterator<Item = impl Into<String>>) -> Self {
174        Self {
175            name: name.into(),
176            values: values.into_iter().map(Into::into).collect(),
177            tablespace: None,
178        }
179    }
180
181    /// Set the tablespace.
182    pub fn tablespace(mut self, tablespace: impl Into<String>) -> Self {
183        self.tablespace = Some(tablespace.into());
184        self
185    }
186}
187
188/// A hash partition definition.
189#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
190pub struct HashPartitionDef {
191    /// Partition name.
192    pub name: String,
193    /// Modulus for hash partitioning.
194    pub modulus: u32,
195    /// Remainder for this partition.
196    pub remainder: u32,
197    /// Tablespace (optional).
198    pub tablespace: Option<String>,
199}
200
201impl HashPartitionDef {
202    /// Create a new hash partition definition.
203    pub fn new(name: impl Into<String>, modulus: u32, remainder: u32) -> Self {
204        Self {
205            name: name.into(),
206            modulus,
207            remainder,
208            tablespace: None,
209        }
210    }
211
212    /// Set the tablespace.
213    pub fn tablespace(mut self, tablespace: impl Into<String>) -> Self {
214        self.tablespace = Some(tablespace.into());
215        self
216    }
217}
218
219/// Partition definitions based on partition type.
220#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
221pub enum PartitionDef {
222    /// Range partitions.
223    Range(Vec<RangePartitionDef>),
224    /// List partitions.
225    List(Vec<ListPartitionDef>),
226    /// Hash partitions.
227    Hash(Vec<HashPartitionDef>),
228}
229
230/// A table partition specification.
231#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
232pub struct Partition {
233    /// Table name.
234    pub table: String,
235    /// Schema name (optional).
236    pub schema: Option<String>,
237    /// Partition type.
238    pub partition_type: PartitionType,
239    /// Partition columns.
240    pub columns: Vec<String>,
241    /// Partition definitions.
242    pub partitions: PartitionDef,
243    /// Optional comment.
244    pub comment: Option<String>,
245}
246
247impl Partition {
248    /// Create a new partition builder.
249    pub fn builder(table: impl Into<String>) -> PartitionBuilder {
250        PartitionBuilder::new(table)
251    }
252
253    /// Get the fully qualified table name.
254    pub fn qualified_table(&self) -> Cow<'_, str> {
255        match &self.schema {
256            Some(schema) => Cow::Owned(format!("{}.{}", schema, self.table)),
257            None => Cow::Borrowed(&self.table),
258        }
259    }
260
261    /// Generate PostgreSQL CREATE TABLE with partitioning.
262    pub fn to_postgres_partition_clause(&self) -> String {
263        format!(
264            "PARTITION BY {} ({})",
265            self.partition_type.to_postgres_sql(),
266            self.columns.join(", ")
267        )
268    }
269
270    /// Generate PostgreSQL CREATE TABLE for a child partition.
271    pub fn to_postgres_create_partition(&self, def: &RangePartitionDef) -> String {
272        let mut sql = format!(
273            "CREATE TABLE {} PARTITION OF {}\n    FOR VALUES FROM ({}) TO ({})",
274            def.name,
275            self.qualified_table(),
276            def.from.to_sql(),
277            def.to.to_sql()
278        );
279
280        if let Some(ref ts) = def.tablespace {
281            sql.push_str(&format!("\n    TABLESPACE {}", ts));
282        }
283
284        sql.push(';');
285        sql
286    }
287
288    /// Generate PostgreSQL CREATE TABLE for a list partition.
289    pub fn to_postgres_create_list_partition(&self, def: &ListPartitionDef) -> String {
290        let values: Vec<String> = def.values.iter().map(|v| format!("'{}'", v)).collect();
291
292        let mut sql = format!(
293            "CREATE TABLE {} PARTITION OF {}\n    FOR VALUES IN ({})",
294            def.name,
295            self.qualified_table(),
296            values.join(", ")
297        );
298
299        if let Some(ref ts) = def.tablespace {
300            sql.push_str(&format!("\n    TABLESPACE {}", ts));
301        }
302
303        sql.push(';');
304        sql
305    }
306
307    /// Generate PostgreSQL CREATE TABLE for a hash partition.
308    pub fn to_postgres_create_hash_partition(&self, def: &HashPartitionDef) -> String {
309        let mut sql = format!(
310            "CREATE TABLE {} PARTITION OF {}\n    FOR VALUES WITH (MODULUS {}, REMAINDER {})",
311            def.name,
312            self.qualified_table(),
313            def.modulus,
314            def.remainder
315        );
316
317        if let Some(ref ts) = def.tablespace {
318            sql.push_str(&format!("\n    TABLESPACE {}", ts));
319        }
320
321        sql.push(';');
322        sql
323    }
324
325    /// Generate all PostgreSQL partition creation SQL.
326    pub fn to_postgres_create_all_partitions(&self) -> Vec<String> {
327        match &self.partitions {
328            PartitionDef::Range(ranges) => ranges
329                .iter()
330                .map(|r| self.to_postgres_create_partition(r))
331                .collect(),
332            PartitionDef::List(lists) => lists
333                .iter()
334                .map(|l| self.to_postgres_create_list_partition(l))
335                .collect(),
336            PartitionDef::Hash(hashes) => hashes
337                .iter()
338                .map(|h| self.to_postgres_create_hash_partition(h))
339                .collect(),
340        }
341    }
342
343    /// Generate MySQL PARTITION BY clause.
344    pub fn to_mysql_partition_clause(&self) -> String {
345        let columns_expr = if self.columns.len() == 1 {
346            self.columns[0].clone()
347        } else {
348            format!("({})", self.columns.join(", "))
349        };
350
351        let mut sql = format!(
352            "PARTITION BY {} ({})",
353            self.partition_type.to_mysql_sql(),
354            columns_expr
355        );
356
357        // Add partition definitions inline for MySQL
358        match &self.partitions {
359            PartitionDef::Range(ranges) => {
360                sql.push_str(" (\n");
361                for (i, r) in ranges.iter().enumerate() {
362                    if i > 0 {
363                        sql.push_str(",\n");
364                    }
365                    sql.push_str(&format!(
366                        "    PARTITION {} VALUES LESS THAN ({})",
367                        r.name,
368                        r.to.to_sql()
369                    ));
370                }
371                sql.push_str("\n)");
372            }
373            PartitionDef::List(lists) => {
374                sql.push_str(" (\n");
375                for (i, l) in lists.iter().enumerate() {
376                    if i > 0 {
377                        sql.push_str(",\n");
378                    }
379                    let values: Vec<String> = l.values.iter().map(|v| format!("'{}'", v)).collect();
380                    sql.push_str(&format!(
381                        "    PARTITION {} VALUES IN ({})",
382                        l.name,
383                        values.join(", ")
384                    ));
385                }
386                sql.push_str("\n)");
387            }
388            PartitionDef::Hash(hashes) => {
389                sql.push_str(&format!(" PARTITIONS {}", hashes.len()));
390            }
391        }
392
393        sql
394    }
395
396    /// Generate MSSQL partition function and scheme.
397    pub fn to_mssql_partition_sql(&self) -> QueryResult<Vec<String>> {
398        match &self.partitions {
399            PartitionDef::Range(ranges) => {
400                let mut sqls = Vec::new();
401
402                // Create partition function
403                let boundaries: Vec<String> = ranges
404                    .iter()
405                    .filter(|r| !matches!(r.to, RangeBound::MaxValue))
406                    .map(|r| r.to.to_sql().into_owned())
407                    .collect();
408
409                let func_name = format!("{}_pf", self.table);
410                sqls.push(format!(
411                    "CREATE PARTITION FUNCTION {}(datetime2)\nAS RANGE RIGHT FOR VALUES ({});",
412                    func_name,
413                    boundaries.join(", ")
414                ));
415
416                // Create partition scheme
417                let scheme_name = format!("{}_ps", self.table);
418                let filegroups: Vec<String> = 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(
576        mut self,
577        name: impl Into<String>,
578        from: RangeBound,
579        to: RangeBound,
580    ) -> Self {
581        self.range_partitions.push(RangePartitionDef::new(name, from, to));
582        self
583    }
584
585    /// Add a range partition with tablespace.
586    pub fn add_range_with_tablespace(
587        mut self,
588        name: impl Into<String>,
589        from: RangeBound,
590        to: RangeBound,
591        tablespace: impl Into<String>,
592    ) -> Self {
593        self.range_partitions.push(
594            RangePartitionDef::new(name, from, to).tablespace(tablespace),
595        );
596        self
597    }
598
599    /// Add a list partition.
600    pub fn add_list(
601        mut self,
602        name: impl Into<String>,
603        values: impl IntoIterator<Item = impl Into<String>>,
604    ) -> Self {
605        self.list_partitions.push(ListPartitionDef::new(name, values));
606        self
607    }
608
609    /// Add a hash partition.
610    pub fn add_hash(mut self, name: impl Into<String>, modulus: u32, remainder: u32) -> Self {
611        self.hash_partitions.push(HashPartitionDef::new(name, modulus, remainder));
612        self
613    }
614
615    /// Add multiple hash partitions automatically.
616    pub fn add_hash_partitions(mut self, count: u32, name_prefix: impl Into<String>) -> Self {
617        let prefix = name_prefix.into();
618        for i in 0..count {
619            self.hash_partitions.push(HashPartitionDef::new(
620                format!("{}_{}", prefix, i),
621                count,
622                i,
623            ));
624        }
625        self
626    }
627
628    /// Add a comment.
629    pub fn comment(mut self, comment: impl Into<String>) -> Self {
630        self.comment = Some(comment.into());
631        self
632    }
633
634    /// Build the partition specification.
635    pub fn build(self) -> QueryResult<Partition> {
636        let partition_type = self.partition_type.ok_or_else(|| {
637            QueryError::invalid_input(
638                "partition_type",
639                "Must specify partition type (range_partition, list_partition, or hash_partition)",
640            )
641        })?;
642
643        if self.columns.is_empty() {
644            return Err(QueryError::invalid_input(
645                "columns",
646                "Must specify at least one partition column",
647            ));
648        }
649
650        let partitions = match partition_type {
651            PartitionType::Range => {
652                if self.range_partitions.is_empty() {
653                    return Err(QueryError::invalid_input(
654                        "partitions",
655                        "Must define at least one range partition with add_range()",
656                    ));
657                }
658                PartitionDef::Range(self.range_partitions)
659            }
660            PartitionType::List => {
661                if self.list_partitions.is_empty() {
662                    return Err(QueryError::invalid_input(
663                        "partitions",
664                        "Must define at least one list partition with add_list()",
665                    ));
666                }
667                PartitionDef::List(self.list_partitions)
668            }
669            PartitionType::Hash => {
670                if self.hash_partitions.is_empty() {
671                    return Err(QueryError::invalid_input(
672                        "partitions",
673                        "Must define at least one hash partition with add_hash() or add_hash_partitions()",
674                    ));
675                }
676                PartitionDef::Hash(self.hash_partitions)
677            }
678        };
679
680        Ok(Partition {
681            table: self.table,
682            schema: self.schema,
683            partition_type,
684            columns: self.columns,
685            partitions,
686            comment: self.comment,
687        })
688    }
689}
690
691/// Time-based partition generation helpers.
692pub mod time_partitions {
693    use super::*;
694
695    /// Generate monthly partitions for a date range.
696    pub fn monthly_partitions(
697        table: &str,
698        column: &str,
699        start_year: i32,
700        start_month: u32,
701        count: u32,
702    ) -> PartitionBuilder {
703        let mut builder = Partition::builder(table)
704            .range_partition()
705            .column(column);
706
707        let mut year = start_year;
708        let mut month = start_month;
709
710        for _ in 0..count {
711            let from_date = format!("{:04}-{:02}-01", year, month);
712
713            // Calculate next month
714            let (next_year, next_month) = if month == 12 {
715                (year + 1, 1)
716            } else {
717                (year, month + 1)
718            };
719            let to_date = format!("{:04}-{:02}-01", next_year, next_month);
720
721            let partition_name = format!("{}_{:04}_{:02}", table, year, month);
722
723            builder = builder.add_range(
724                partition_name,
725                RangeBound::date(from_date),
726                RangeBound::date(to_date),
727            );
728
729            year = next_year;
730            month = next_month;
731        }
732
733        builder
734    }
735
736    /// Generate quarterly partitions for a date range.
737    pub fn quarterly_partitions(
738        table: &str,
739        column: &str,
740        start_year: i32,
741        count: u32,
742    ) -> PartitionBuilder {
743        let mut builder = Partition::builder(table)
744            .range_partition()
745            .column(column);
746
747        let mut year = start_year;
748        let mut quarter = 1;
749
750        for _ in 0..count {
751            let from_month = (quarter - 1) * 3 + 1;
752            let from_date = format!("{:04}-{:02}-01", year, from_month);
753
754            let (next_year, next_quarter) = if quarter == 4 {
755                (year + 1, 1)
756            } else {
757                (year, quarter + 1)
758            };
759            let to_month = (next_quarter - 1) * 3 + 1;
760            let to_date = format!("{:04}-{:02}-01", next_year, to_month);
761
762            let partition_name = format!("{}_{}q{}", table, year, quarter);
763
764            builder = builder.add_range(
765                partition_name,
766                RangeBound::date(from_date),
767                RangeBound::date(to_date),
768            );
769
770            year = next_year;
771            quarter = next_quarter;
772        }
773
774        builder
775    }
776
777    /// Generate yearly partitions.
778    pub fn yearly_partitions(
779        table: &str,
780        column: &str,
781        start_year: i32,
782        count: u32,
783    ) -> PartitionBuilder {
784        let mut builder = Partition::builder(table)
785            .range_partition()
786            .column(column);
787
788        for i in 0..count {
789            let year = start_year + i as i32;
790            let from_date = format!("{:04}-01-01", year);
791            let to_date = format!("{:04}-01-01", year + 1);
792            let partition_name = format!("{}_{}", table, year);
793
794            builder = builder.add_range(
795                partition_name,
796                RangeBound::date(from_date),
797                RangeBound::date(to_date),
798            );
799        }
800
801        builder
802    }
803}
804
805/// MongoDB sharding support.
806pub mod mongodb {
807    use serde::{Deserialize, Serialize};
808
809    /// Shard key type for MongoDB.
810    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
811    pub enum ShardKeyType {
812        /// Range-based sharding (good for range queries).
813        Range,
814        /// Hash-based sharding (better distribution).
815        Hashed,
816    }
817
818    impl ShardKeyType {
819        /// Get the MongoDB index specification value.
820        pub fn as_index_value(&self) -> serde_json::Value {
821            match self {
822                Self::Range => serde_json::json!(1),
823                Self::Hashed => serde_json::json!("hashed"),
824            }
825        }
826    }
827
828    /// A shard key definition.
829    #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
830    pub struct ShardKey {
831        /// Fields in the shard key.
832        pub fields: Vec<(String, ShardKeyType)>,
833        /// Whether the collection is unique on the shard key.
834        pub unique: bool,
835    }
836
837    impl ShardKey {
838        /// Create a new shard key builder.
839        pub fn builder() -> ShardKeyBuilder {
840            ShardKeyBuilder::new()
841        }
842
843        /// Get the shardCollection command.
844        pub fn shard_collection_command(
845            &self,
846            database: &str,
847            collection: &str,
848        ) -> serde_json::Value {
849            let mut key = serde_json::Map::new();
850            for (field, key_type) in &self.fields {
851                key.insert(field.clone(), key_type.as_index_value());
852            }
853
854            serde_json::json!({
855                "shardCollection": format!("{}.{}", database, collection),
856                "key": key,
857                "unique": self.unique
858            })
859        }
860
861        /// Get the index specification for the shard key.
862        pub fn index_spec(&self) -> serde_json::Value {
863            let mut spec = serde_json::Map::new();
864            for (field, key_type) in &self.fields {
865                spec.insert(field.clone(), key_type.as_index_value());
866            }
867            serde_json::Value::Object(spec)
868        }
869    }
870
871    /// Builder for shard keys.
872    #[derive(Debug, Clone, Default)]
873    pub struct ShardKeyBuilder {
874        fields: Vec<(String, ShardKeyType)>,
875        unique: bool,
876    }
877
878    impl ShardKeyBuilder {
879        /// Create a new shard key builder.
880        pub fn new() -> Self {
881            Self::default()
882        }
883
884        /// Add a range field to the shard key.
885        pub fn range_field(mut self, field: impl Into<String>) -> Self {
886            self.fields.push((field.into(), ShardKeyType::Range));
887            self
888        }
889
890        /// Add a hashed field to the shard key.
891        pub fn hashed_field(mut self, field: impl Into<String>) -> Self {
892            self.fields.push((field.into(), ShardKeyType::Hashed));
893            self
894        }
895
896        /// Set whether the shard key should enforce uniqueness.
897        pub fn unique(mut self, unique: bool) -> Self {
898            self.unique = unique;
899            self
900        }
901
902        /// Build the shard key.
903        pub fn build(self) -> ShardKey {
904            ShardKey {
905                fields: self.fields,
906                unique: self.unique,
907            }
908        }
909    }
910
911    /// Zone (tag-based) sharding configuration.
912    #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
913    pub struct ShardZone {
914        /// Zone name.
915        pub name: String,
916        /// Minimum shard key value.
917        pub min: serde_json::Value,
918        /// Maximum shard key value.
919        pub max: serde_json::Value,
920    }
921
922    impl ShardZone {
923        /// Create a new shard zone.
924        pub fn new(
925            name: impl Into<String>,
926            min: serde_json::Value,
927            max: serde_json::Value,
928        ) -> Self {
929            Self {
930                name: name.into(),
931                min,
932                max,
933            }
934        }
935
936        /// Get the updateZoneKeyRange command.
937        pub fn update_zone_key_range_command(
938            &self,
939            namespace: &str,
940        ) -> serde_json::Value {
941            serde_json::json!({
942                "updateZoneKeyRange": namespace,
943                "min": self.min,
944                "max": self.max,
945                "zone": self.name
946            })
947        }
948
949        /// Get the addShardToZone command.
950        pub fn add_shard_to_zone_command(&self, shard: &str) -> serde_json::Value {
951            serde_json::json!({
952                "addShardToZone": shard,
953                "zone": self.name
954            })
955        }
956    }
957
958    /// Builder for zone sharding configuration.
959    #[derive(Debug, Clone, Default)]
960    pub struct ZoneShardingBuilder {
961        zones: Vec<ShardZone>,
962        shard_assignments: Vec<(String, String)>, // (shard_name, zone_name)
963    }
964
965    impl ZoneShardingBuilder {
966        /// Create a new zone sharding builder.
967        pub fn new() -> Self {
968            Self::default()
969        }
970
971        /// Add a zone.
972        pub fn add_zone(
973            mut self,
974            name: impl Into<String>,
975            min: serde_json::Value,
976            max: serde_json::Value,
977        ) -> Self {
978            self.zones.push(ShardZone::new(name, min, max));
979            self
980        }
981
982        /// Assign a shard to a zone.
983        pub fn assign_shard(mut self, shard: impl Into<String>, zone: impl Into<String>) -> Self {
984            self.shard_assignments.push((shard.into(), zone.into()));
985            self
986        }
987
988        /// Get all configuration commands.
989        pub fn build_commands(&self, namespace: &str) -> Vec<serde_json::Value> {
990            let mut commands = Vec::new();
991
992            // Add shard to zone assignments
993            for (shard, zone) in &self.shard_assignments {
994                commands.push(serde_json::json!({
995                    "addShardToZone": shard,
996                    "zone": zone
997                }));
998            }
999
1000            // Add zone key ranges
1001            for zone in &self.zones {
1002                commands.push(zone.update_zone_key_range_command(namespace));
1003            }
1004
1005            commands
1006        }
1007    }
1008
1009    /// Helper to create a shard key.
1010    pub fn shard_key() -> ShardKeyBuilder {
1011        ShardKeyBuilder::new()
1012    }
1013
1014    /// Helper to create zone sharding configuration.
1015    pub fn zone_sharding() -> ZoneShardingBuilder {
1016        ZoneShardingBuilder::new()
1017    }
1018}
1019
1020#[cfg(test)]
1021mod tests {
1022    use super::*;
1023
1024    #[test]
1025    fn test_range_partition_builder() {
1026        let partition = Partition::builder("orders")
1027            .schema("sales")
1028            .range_partition()
1029            .column("created_at")
1030            .add_range("orders_2024_q1", RangeBound::date("2024-01-01"), RangeBound::date("2024-04-01"))
1031            .add_range("orders_2024_q2", RangeBound::date("2024-04-01"), RangeBound::date("2024-07-01"))
1032            .build()
1033            .unwrap();
1034
1035        assert_eq!(partition.table, "orders");
1036        assert_eq!(partition.partition_type, PartitionType::Range);
1037        assert_eq!(partition.columns, vec!["created_at"]);
1038    }
1039
1040    #[test]
1041    fn test_postgres_partition_clause() {
1042        let partition = Partition::builder("orders")
1043            .range_partition()
1044            .column("created_at")
1045            .add_range("orders_2024", RangeBound::date("2024-01-01"), RangeBound::date("2025-01-01"))
1046            .build()
1047            .unwrap();
1048
1049        let clause = partition.to_postgres_partition_clause();
1050        assert_eq!(clause, "PARTITION BY RANGE (created_at)");
1051    }
1052
1053    #[test]
1054    fn test_postgres_create_partition() {
1055        let partition = Partition::builder("orders")
1056            .range_partition()
1057            .column("created_at")
1058            .add_range("orders_2024", RangeBound::date("2024-01-01"), RangeBound::date("2025-01-01"))
1059            .build()
1060            .unwrap();
1061
1062        let sqls = partition.to_postgres_create_all_partitions();
1063        assert_eq!(sqls.len(), 1);
1064        assert!(sqls[0].contains("CREATE TABLE orders_2024 PARTITION OF orders"));
1065        assert!(sqls[0].contains("FOR VALUES FROM ('2024-01-01') TO ('2025-01-01')"));
1066    }
1067
1068    #[test]
1069    fn test_list_partition() {
1070        let partition = Partition::builder("users")
1071            .list_partition()
1072            .column("country")
1073            .add_list("users_us", ["US", "USA"])
1074            .add_list("users_eu", ["DE", "FR", "GB", "IT"])
1075            .build()
1076            .unwrap();
1077
1078        assert_eq!(partition.partition_type, PartitionType::List);
1079
1080        let sqls = partition.to_postgres_create_all_partitions();
1081        assert_eq!(sqls.len(), 2);
1082        assert!(sqls[0].contains("FOR VALUES IN"));
1083    }
1084
1085    #[test]
1086    fn test_hash_partition() {
1087        let partition = Partition::builder("events")
1088            .hash_partition()
1089            .column("user_id")
1090            .add_hash_partitions(4, "events")
1091            .build()
1092            .unwrap();
1093
1094        assert_eq!(partition.partition_type, PartitionType::Hash);
1095
1096        let sqls = partition.to_postgres_create_all_partitions();
1097        assert_eq!(sqls.len(), 4);
1098        assert!(sqls[0].contains("MODULUS 4"));
1099        assert!(sqls[0].contains("REMAINDER 0"));
1100    }
1101
1102    #[test]
1103    fn test_mysql_partition_clause() {
1104        let partition = Partition::builder("orders")
1105            .range_partition()
1106            .column("created_at")
1107            .add_range("p2024", RangeBound::MinValue, RangeBound::date("2025-01-01"))
1108            .add_range("p_future", RangeBound::date("2025-01-01"), RangeBound::MaxValue)
1109            .build()
1110            .unwrap();
1111
1112        let clause = partition.to_mysql_partition_clause();
1113        assert!(clause.contains("PARTITION BY RANGE"));
1114        assert!(clause.contains("PARTITION p2024"));
1115        assert!(clause.contains("PARTITION p_future"));
1116    }
1117
1118    #[test]
1119    fn test_detach_partition() {
1120        let partition = Partition::builder("orders")
1121            .range_partition()
1122            .column("created_at")
1123            .add_range("orders_2024", RangeBound::date("2024-01-01"), RangeBound::date("2025-01-01"))
1124            .build()
1125            .unwrap();
1126
1127        let sql = partition.detach_partition_sql("orders_2024", DatabaseType::PostgreSQL).unwrap();
1128        assert_eq!(sql, "ALTER TABLE orders DETACH PARTITION orders_2024;");
1129    }
1130
1131    #[test]
1132    fn test_drop_partition() {
1133        let partition = Partition::builder("orders")
1134            .range_partition()
1135            .column("created_at")
1136            .add_range("orders_2024", RangeBound::date("2024-01-01"), RangeBound::date("2025-01-01"))
1137            .build()
1138            .unwrap();
1139
1140        let pg_sql = partition.drop_partition_sql("orders_2024", DatabaseType::PostgreSQL).unwrap();
1141        assert_eq!(pg_sql, "DROP TABLE IF EXISTS orders_2024;");
1142
1143        let mysql_sql = partition.drop_partition_sql("orders_2024", DatabaseType::MySQL).unwrap();
1144        assert!(mysql_sql.contains("DROP PARTITION"));
1145    }
1146
1147    #[test]
1148    fn test_missing_partition_type() {
1149        let result = Partition::builder("orders")
1150            .column("created_at")
1151            .add_range("p1", RangeBound::MinValue, RangeBound::MaxValue)
1152            .build();
1153
1154        assert!(result.is_err());
1155    }
1156
1157    #[test]
1158    fn test_missing_columns() {
1159        let result = Partition::builder("orders")
1160            .range_partition()
1161            .add_range("p1", RangeBound::MinValue, RangeBound::MaxValue)
1162            .build();
1163
1164        assert!(result.is_err());
1165    }
1166
1167    mod time_partition_tests {
1168        use super::super::time_partitions;
1169
1170        #[test]
1171        fn test_monthly_partitions() {
1172            let builder = time_partitions::monthly_partitions("orders", "created_at", 2024, 1, 3);
1173            let partition = builder.build().unwrap();
1174
1175            if let super::PartitionDef::Range(ranges) = &partition.partitions {
1176                assert_eq!(ranges.len(), 3);
1177                assert_eq!(ranges[0].name, "orders_2024_01");
1178                assert_eq!(ranges[1].name, "orders_2024_02");
1179                assert_eq!(ranges[2].name, "orders_2024_03");
1180            } else {
1181                panic!("Expected range partitions");
1182            }
1183        }
1184
1185        #[test]
1186        fn test_quarterly_partitions() {
1187            let builder = time_partitions::quarterly_partitions("sales", "order_date", 2024, 4);
1188            let partition = builder.build().unwrap();
1189
1190            if let super::PartitionDef::Range(ranges) = &partition.partitions {
1191                assert_eq!(ranges.len(), 4);
1192                assert_eq!(ranges[0].name, "sales_2024q1");
1193                assert_eq!(ranges[3].name, "sales_2024q4");
1194            } else {
1195                panic!("Expected range partitions");
1196            }
1197        }
1198
1199        #[test]
1200        fn test_yearly_partitions() {
1201            let builder = time_partitions::yearly_partitions("logs", "timestamp", 2020, 5);
1202            let partition = builder.build().unwrap();
1203
1204            if let super::PartitionDef::Range(ranges) = &partition.partitions {
1205                assert_eq!(ranges.len(), 5);
1206                assert_eq!(ranges[0].name, "logs_2020");
1207                assert_eq!(ranges[4].name, "logs_2024");
1208            } else {
1209                panic!("Expected range partitions");
1210            }
1211        }
1212    }
1213
1214    mod mongodb_tests {
1215        use super::super::mongodb::*;
1216
1217        #[test]
1218        fn test_shard_key_builder() {
1219            let key = shard_key()
1220                .hashed_field("tenant_id")
1221                .range_field("created_at")
1222                .unique(false)
1223                .build();
1224
1225            assert_eq!(key.fields.len(), 2);
1226            assert_eq!(key.fields[0], ("tenant_id".to_string(), ShardKeyType::Hashed));
1227            assert_eq!(key.fields[1], ("created_at".to_string(), ShardKeyType::Range));
1228        }
1229
1230        #[test]
1231        fn test_shard_collection_command() {
1232            let key = shard_key()
1233                .hashed_field("user_id")
1234                .build();
1235
1236            let cmd = key.shard_collection_command("mydb", "users");
1237            assert_eq!(cmd["shardCollection"], "mydb.users");
1238            assert_eq!(cmd["key"]["user_id"], "hashed");
1239        }
1240
1241        #[test]
1242        fn test_zone_sharding() {
1243            let builder = zone_sharding()
1244                .add_zone("US", serde_json::json!({"region": "US"}), serde_json::json!({"region": "US~"}))
1245                .add_zone("EU", serde_json::json!({"region": "EU"}), serde_json::json!({"region": "EU~"}))
1246                .assign_shard("shard0", "US")
1247                .assign_shard("shard1", "EU");
1248
1249            let commands = builder.build_commands("mydb.users");
1250            assert_eq!(commands.len(), 4); // 2 shard assignments + 2 zone ranges
1251        }
1252
1253        #[test]
1254        fn test_shard_key_index_spec() {
1255            let key = shard_key()
1256                .range_field("tenant_id")
1257                .range_field("created_at")
1258                .build();
1259
1260            let spec = key.index_spec();
1261            assert_eq!(spec["tenant_id"], 1);
1262            assert_eq!(spec["created_at"], 1);
1263        }
1264    }
1265}
1266
1267
1268
1269
prax_query/partition.rs

prax_query/
partition.rs
