From ffc6c751b3c144160652421d79e03f754c464767 Mon Sep 17 00:00:00 2001 From: Erdem Date: Sat, 10 Jan 2026 21:23:10 +0300 Subject: [PATCH 1/2] feat: add support for table partitioning --- src/backend/mysql/table.rs | 83 +++++++++++++++++++ src/backend/postgres/table.rs | 71 ++++++++++++++++ src/backend/table_builder.rs | 117 ++++++++++++++++++-------- src/table/create.rs | 150 ++++++++++++++++++++++++++++++++-- tests/mysql/table.rs | 64 +++++++++++++++ tests/postgres/table.rs | 73 +++++++++++++++++ 6 files changed, 518 insertions(+), 40 deletions(-) diff --git a/src/backend/mysql/table.rs b/src/backend/mysql/table.rs index 5b8ef1288..c54af4363 100644 --- a/src/backend/mysql/table.rs +++ b/src/backend/mysql/table.rs @@ -237,6 +237,65 @@ impl TableBuilder for MysqlQueryBuilder { } } + fn prepare_partition_by(&self, partition_by: &PartitionBy, sql: &mut impl SqlWriter) { + match partition_by { + PartitionBy::Range(cols) => { + sql.write_str("RANGE (").unwrap(); + self.prepare_partition_cols(cols, sql); + sql.write_char(')').unwrap(); + } + PartitionBy::List(cols) => { + sql.write_str("LIST (").unwrap(); + self.prepare_partition_cols(cols, sql); + sql.write_char(')').unwrap(); + } + PartitionBy::Hash(cols) => { + sql.write_str("HASH (").unwrap(); + self.prepare_partition_cols(cols, sql); + sql.write_char(')').unwrap(); + } + PartitionBy::Key(cols) => { + sql.write_str("KEY (").unwrap(); + self.prepare_partition_cols(cols, sql); + sql.write_char(')').unwrap(); + } + } + } + + fn prepare_partition_definition( + &self, + partition_definition: &PartitionDefinition, + sql: &mut impl SqlWriter, + ) { + sql.write_str("PARTITION ").unwrap(); + self.prepare_iden(&partition_definition.name, sql); + if let Some(values) = &partition_definition.values { + sql.write_str(" ").unwrap(); + self.prepare_partition_values(values, sql); + } + } + + fn prepare_partition_values( + &self, + partition_values: &PartitionValues, + sql: &mut impl SqlWriter, + ) { + match partition_values { + PartitionValues::In(values) => { + sql.write_str("VALUES IN (").unwrap(); + self.prepare_partition_exprs(values, sql); + sql.write_char(')').unwrap(); + } + PartitionValues::FromTo(_, _) => panic!("MySQL does not support VALUES FROM ... TO"), + PartitionValues::LessThan(values) => { + sql.write_str("VALUES LESS THAN (").unwrap(); + self.prepare_partition_exprs(values, sql); + sql.write_char(')').unwrap(); + } + PartitionValues::With(_, _) => panic!("MySQL does not support VALUES WITH"), + } + } + /// column comment fn column_comment(&self, comment: &str, sql: &mut impl SqlWriter) { sql.write_str(" COMMENT '").unwrap(); @@ -244,3 +303,27 @@ impl TableBuilder for MysqlQueryBuilder { sql.write_str("'").unwrap(); } } + +impl MysqlQueryBuilder { + fn prepare_partition_cols(&self, cols: &[DynIden], sql: &mut impl SqlWriter) { + let mut first = true; + for col in cols { + if !first { + sql.write_str(", ").unwrap(); + } + self.prepare_iden(col, sql); + first = false; + } + } + + fn prepare_partition_exprs(&self, exprs: &[SimpleExpr], sql: &mut impl SqlWriter) { + let mut first = true; + for expr in exprs { + if !first { + sql.write_str(", ").unwrap(); + } + self.prepare_expr(expr, sql); + first = false; + } + } +} diff --git a/src/backend/postgres/table.rs b/src/backend/postgres/table.rs index c8b072486..21ebc4c85 100644 --- a/src/backend/postgres/table.rs +++ b/src/backend/postgres/table.rs @@ -222,6 +222,77 @@ impl TableBuilder for PostgresQueryBuilder { self.prepare_table_ref_table_stmt(to_name, sql); } } + + fn prepare_partition_by(&self, partition_by: &PartitionBy, sql: &mut impl SqlWriter) { + match partition_by { + PartitionBy::Range(cols) => { + sql.write_str("RANGE (").unwrap(); + self.prepare_partition_cols(cols, sql); + sql.write_char(')').unwrap(); + } + PartitionBy::List(cols) => { + sql.write_str("LIST (").unwrap(); + self.prepare_partition_cols(cols, sql); + sql.write_char(')').unwrap(); + } + PartitionBy::Hash(cols) => { + sql.write_str("HASH (").unwrap(); + self.prepare_partition_cols(cols, sql); + sql.write_char(')').unwrap(); + } + PartitionBy::Key(_) => panic!("Postgres does not support PARTITION BY KEY"), + } + } + + fn prepare_partition_values( + &self, + partition_values: &PartitionValues, + sql: &mut impl SqlWriter, + ) { + sql.write_str("FOR VALUES ").unwrap(); + match partition_values { + PartitionValues::In(values) => { + sql.write_str("IN (").unwrap(); + self.prepare_partition_exprs(values, sql); + sql.write_char(')').unwrap(); + } + PartitionValues::FromTo(from, to) => { + sql.write_str("FROM (").unwrap(); + self.prepare_partition_exprs(from, sql); + sql.write_str(") TO (").unwrap(); + self.prepare_partition_exprs(to, sql); + sql.write_char(')').unwrap(); + } + PartitionValues::LessThan(_) => panic!("Postgres does not support VALUES LESS THAN"), + PartitionValues::With(modulus, remainder) => { + write!(sql, "WITH (MODULUS {modulus}, REMAINDER {remainder})").unwrap(); + } + } + } +} + +impl PostgresQueryBuilder { + fn prepare_partition_cols(&self, cols: &[DynIden], sql: &mut impl SqlWriter) { + let mut first = true; + for col in cols { + if !first { + sql.write_str(", ").unwrap(); + } + self.prepare_iden(col, sql); + first = false; + } + } + + fn prepare_partition_exprs(&self, exprs: &[SimpleExpr], sql: &mut impl SqlWriter) { + let mut first = true; + for expr in exprs { + if !first { + sql.write_str(", ").unwrap(); + } + self.prepare_expr(expr, sql); + first = false; + } + } } impl PostgresQueryBuilder { diff --git a/src/backend/table_builder.rs b/src/backend/table_builder.rs index 3f266ead1..92f3dffbf 100644 --- a/src/backend/table_builder.rs +++ b/src/backend/table_builder.rs @@ -21,42 +21,80 @@ pub trait TableBuilder: self.prepare_table_ref_table_stmt(table_ref, sql); } - sql.write_str(" ( ").unwrap(); - let mut first = true; + if let Some(partition_of) = &create.partition_of { + sql.write_str(" PARTITION OF ").unwrap(); + self.prepare_table_ref_table_stmt(partition_of, sql); + } - create.columns.iter().for_each(|column_def| { - if !first { - sql.write_str(", ").unwrap(); - } - self.prepare_column_def(column_def, sql); - first = false; - }); + if !create.columns.is_empty() + || !create.indexes.is_empty() + || !create.foreign_keys.is_empty() + || !create.check.is_empty() + { + sql.write_str(" ( ").unwrap(); + let mut first = true; + + create.columns.iter().for_each(|column_def| { + if !first { + sql.write_str(", ").unwrap(); + } + self.prepare_column_def(column_def, sql); + first = false; + }); - create.indexes.iter().for_each(|index| { - if !first { - sql.write_str(", ").unwrap(); - } - self.prepare_table_index_expression(index, sql); - first = false; - }); + create.indexes.iter().for_each(|index| { + if !first { + sql.write_str(", ").unwrap(); + } + self.prepare_table_index_expression(index, sql); + first = false; + }); - create.foreign_keys.iter().for_each(|foreign_key| { - if !first { - sql.write_str(", ").unwrap(); - } - self.prepare_foreign_key_create_statement_internal(foreign_key, sql, Mode::Creation); - first = false; - }); + create.foreign_keys.iter().for_each(|foreign_key| { + if !first { + sql.write_str(", ").unwrap(); + } + self.prepare_foreign_key_create_statement_internal( + foreign_key, + sql, + Mode::Creation, + ); + first = false; + }); + + create.check.iter().for_each(|check| { + if !first { + sql.write_str(", ").unwrap(); + } + self.prepare_check_constraint(check, sql); + first = false; + }); - create.check.iter().for_each(|check| { - if !first { - sql.write_str(", ").unwrap(); - } - self.prepare_check_constraint(check, sql); - first = false; - }); + sql.write_str(" )").unwrap(); + } - sql.write_str(" )").unwrap(); + if let Some(partition_values) = &create.partition_values { + sql.write_str(" ").unwrap(); + self.prepare_partition_values(partition_values, sql); + } + + if let Some(partition_by) = &create.partition_by { + sql.write_str(" PARTITION BY ").unwrap(); + self.prepare_partition_by(partition_by, sql); + } + + if !create.partitions.is_empty() { + sql.write_str(" ( ").unwrap(); + let mut first = true; + for partition in create.partitions.iter() { + if !first { + sql.write_str(", ").unwrap(); + } + self.prepare_partition_definition(partition, sql); + first = false; + } + sql.write_str(" )").unwrap(); + } self.prepare_table_opt(create, sql); @@ -197,10 +235,21 @@ pub trait TableBuilder: } } - /// Translate [`TablePartition`] into SQL statement. - fn prepare_table_partition( + /// Translate [`PartitionBy`] into SQL statement. + fn prepare_partition_by(&self, _partition_by: &PartitionBy, _sql: &mut impl SqlWriter) {} + + /// Translate [`PartitionValues`] into SQL statement. + fn prepare_partition_values( + &self, + _partition_values: &PartitionValues, + _sql: &mut impl SqlWriter, + ) { + } + + /// Translate [`PartitionDefinition`] into SQL statement. + fn prepare_partition_definition( &self, - _table_partition: &TablePartition, + _partition_definition: &PartitionDefinition, _sql: &mut impl SqlWriter, ) { } diff --git a/src/table/create.rs b/src/table/create.rs index 8bbac1834..036a3a33e 100644 --- a/src/table/create.rs +++ b/src/table/create.rs @@ -1,8 +1,8 @@ use inherent::inherent; use crate::{ - ColumnDef, IntoColumnDef, SchemaStatementBuilder, backend::SchemaBuilder, foreign_key::*, - index::*, table::constraint::Check, types::*, + ColumnDef, IntoColumnDef, SchemaStatementBuilder, SimpleExpr, backend::SchemaBuilder, + foreign_key::*, index::*, table::constraint::Check, types::*, }; /// Create a table @@ -84,7 +84,10 @@ pub struct TableCreateStatement { pub(crate) table: Option, pub(crate) columns: Vec, pub(crate) options: Vec, - pub(crate) partitions: Vec, + pub(crate) partition_by: Option, + pub(crate) partition_of: Option, + pub(crate) partition_values: Option, + pub(crate) partitions: Vec, pub(crate) indexes: Vec, pub(crate) foreign_keys: Vec, pub(crate) if_not_exists: bool, @@ -105,8 +108,26 @@ pub enum TableOpt { /// All available table partition options #[derive(Debug, Clone)] -#[non_exhaustive] -pub enum TablePartition {} +pub enum PartitionBy { + Range(Vec), + List(Vec), + Hash(Vec), + Key(Vec), +} + +#[derive(Debug, Clone)] +pub enum PartitionValues { + In(Vec), + FromTo(Vec, Vec), + LessThan(Vec), + With(u32, u32), +} + +#[derive(Debug, Clone)] +pub struct PartitionDefinition { + pub(crate) name: DynIden, + pub(crate) values: Option, +} impl TableCreateStatement { /// Construct create table statement @@ -274,11 +295,125 @@ impl TableCreateStatement { } #[allow(dead_code)] - fn partition(&mut self, partition: TablePartition) -> &mut Self { + fn partition(&mut self, partition: PartitionDefinition) -> &mut Self { self.partitions.push(partition); self } + /// Set partition by range + pub fn partition_by_range(&mut self, cols: I) -> &mut Self + where + I: IntoIterator, + T: IntoIden, + { + self.partition_by = Some(PartitionBy::Range( + cols.into_iter().map(|c| c.into_iden()).collect(), + )); + self + } + + /// Set partition by list + pub fn partition_by_list(&mut self, cols: I) -> &mut Self + where + I: IntoIterator, + T: IntoIden, + { + self.partition_by = Some(PartitionBy::List( + cols.into_iter().map(|c| c.into_iden()).collect(), + )); + self + } + + /// Set partition by hash + pub fn partition_by_hash(&mut self, cols: I) -> &mut Self + where + I: IntoIterator, + T: IntoIden, + { + self.partition_by = Some(PartitionBy::Hash( + cols.into_iter().map(|c| c.into_iden()).collect(), + )); + self + } + + /// Set partition by key. MySQL only. + pub fn partition_by_key(&mut self, cols: I) -> &mut Self + where + I: IntoIterator, + T: IntoIden, + { + self.partition_by = Some(PartitionBy::Key( + cols.into_iter().map(|c| c.into_iden()).collect(), + )); + self + } + + /// Set partition of table. Postgres only. + pub fn partition_of(&mut self, table: T) -> &mut Self + where + T: IntoTableRef, + { + self.partition_of = Some(table.into_table_ref()); + self + } + + /// Set partition values IN. Postgres and MySQL. + pub fn values_in(&mut self, values: I) -> &mut Self + where + I: IntoIterator, + T: Into, + { + self.partition_values = Some(PartitionValues::In( + values.into_iter().map(|v| v.into()).collect(), + )); + self + } + + /// Set partition values FROM ... TO. Postgres only. + pub fn values_from_to(&mut self, from: I, to: J) -> &mut Self + where + I: IntoIterator, + T: Into, + J: IntoIterator, + U: Into, + { + self.partition_values = Some(PartitionValues::FromTo( + from.into_iter().map(|v| v.into()).collect(), + to.into_iter().map(|v| v.into()).collect(), + )); + self + } + + /// Set partition values LESS THAN. MySQL only. + pub fn values_less_than(&mut self, values: I) -> &mut Self + where + I: IntoIterator, + T: Into, + { + self.partition_values = Some(PartitionValues::LessThan( + values.into_iter().map(|v| v.into()).collect(), + )); + self + } + + /// Set partition values WITH (modulus, remainder). Postgres only. + pub fn values_with(&mut self, modulus: u32, remainder: u32) -> &mut Self { + self.partition_values = Some(PartitionValues::With(modulus, remainder)); + self + } + + /// Add a partition definition. MySQL only. + pub fn add_partition(&mut self, name: T, values: Option) -> &mut Self + where + T: IntoIden, + { + self.partitions.push(PartitionDefinition { + name: name.into_iden(), + values, + }); + self + } + pub fn get_table_name(&self) -> Option<&TableRef> { self.table.as_ref() } @@ -412,6 +547,9 @@ impl TableCreateStatement { table: self.table.take(), columns: std::mem::take(&mut self.columns), options: std::mem::take(&mut self.options), + partition_by: self.partition_by.take(), + partition_of: self.partition_of.take(), + partition_values: self.partition_values.take(), partitions: std::mem::take(&mut self.partitions), indexes: std::mem::take(&mut self.indexes), foreign_keys: std::mem::take(&mut self.foreign_keys), diff --git a/tests/mysql/table.rs b/tests/mysql/table.rs index f44148326..15abf86e3 100644 --- a/tests/mysql/table.rs +++ b/tests/mysql/table.rs @@ -481,3 +481,67 @@ fn alter_with_named_check_constraint() { r#"ALTER TABLE `glyph` ADD COLUMN `aspect` int NOT NULL DEFAULT 101 CONSTRAINT `positive_aspect` CHECK (`aspect` > 100)"#, ); } + +#[test] +fn create_partition_range() { + assert_eq!( + Table::create() + .table(Glyph::Table) + .col(ColumnDef::new(Glyph::Id).integer().not_null()) + .partition_by_range([Glyph::Id]) + .add_partition( + Alias::new("p0"), + Some(PartitionValues::LessThan(vec![6.into()])) + ) + .add_partition( + Alias::new("p1"), + Some(PartitionValues::LessThan(vec![11.into()])) + ) + .to_string(MysqlQueryBuilder), + "CREATE TABLE `glyph` ( `id` int NOT NULL ) PARTITION BY RANGE (`id`) ( PARTITION `p0` VALUES LESS THAN (6), PARTITION `p1` VALUES LESS THAN (11) )" + ); +} + +#[test] +fn create_partition_list() { + assert_eq!( + Table::create() + .table(Glyph::Table) + .col(ColumnDef::new(Glyph::Id).integer().not_null()) + .partition_by_list([Glyph::Id]) + .add_partition( + Alias::new("p0"), + Some(PartitionValues::In(vec![1.into(), 2.into()])) + ) + .add_partition( + Alias::new("p1"), + Some(PartitionValues::In(vec![3.into(), 4.into()])) + ) + .to_string(MysqlQueryBuilder), + "CREATE TABLE `glyph` ( `id` int NOT NULL ) PARTITION BY LIST (`id`) ( PARTITION `p0` VALUES IN (1, 2), PARTITION `p1` VALUES IN (3, 4) )" + ); +} + +#[test] +fn create_partition_hash() { + assert_eq!( + Table::create() + .table(Glyph::Table) + .col(ColumnDef::new(Glyph::Id).integer().not_null()) + .partition_by_hash([Glyph::Id]) + .to_string(MysqlQueryBuilder), + "CREATE TABLE `glyph` ( `id` int NOT NULL ) PARTITION BY HASH (`id`)" + ); +} + +#[test] +fn create_partition_key() { + assert_eq!( + Table::create() + .table(Glyph::Table) + .col(ColumnDef::new(Glyph::Id).integer().not_null()) + .partition_by_key([Glyph::Id]) + .to_string(MysqlQueryBuilder), + "CREATE TABLE `glyph` ( `id` int NOT NULL ) PARTITION BY KEY (`id`)" + ); +} diff --git a/tests/postgres/table.rs b/tests/postgres/table.rs index 15e6fa9f1..86108ce8c 100644 --- a/tests/postgres/table.rs +++ b/tests/postgres/table.rs @@ -740,3 +740,76 @@ fn create_19() { .join(" "), ); } + +#[test] +fn create_partition_master_range() { + assert_eq!( + Table::create() + .table(Glyph::Table) + .col(ColumnDef::new(Glyph::Id).integer().not_null()) + .col(ColumnDef::new(Glyph::Aspect).integer().not_null()) + .partition_by_range([Glyph::Aspect]) + .to_string(PostgresQueryBuilder), + r#"CREATE TABLE "glyph" ( "id" integer NOT NULL, "aspect" integer NOT NULL ) PARTITION BY RANGE ("aspect")"# + ); +} + +#[test] +fn create_partition_child_range() { + assert_eq!( + Table::create() + .table(Alias::new("glyph_1")) + .partition_of(Glyph::Table) + .values_from_to([1], [10]) + .to_string(PostgresQueryBuilder), + r#"CREATE TABLE "glyph_1" PARTITION OF "glyph" FOR VALUES FROM (1) TO (10)"# + ); +} + +#[test] +fn create_partition_master_list() { + assert_eq!( + Table::create() + .table(Glyph::Table) + .col(ColumnDef::new(Glyph::Id).integer().not_null()) + .partition_by_list([Glyph::Id]) + .to_string(PostgresQueryBuilder), + r#"CREATE TABLE "glyph" ( "id" integer NOT NULL ) PARTITION BY LIST ("id")"# + ); +} + +#[test] +fn create_partition_child_list() { + assert_eq!( + Table::create() + .table(Alias::new("glyph_p1")) + .partition_of(Glyph::Table) + .values_in([1, 2, 3]) + .to_string(PostgresQueryBuilder), + r#"CREATE TABLE "glyph_p1" PARTITION OF "glyph" FOR VALUES IN (1, 2, 3)"# + ); +} + +#[test] +fn create_partition_master_hash() { + assert_eq!( + Table::create() + .table(Glyph::Table) + .col(ColumnDef::new(Glyph::Id).integer().not_null()) + .partition_by_hash([Glyph::Id]) + .to_string(PostgresQueryBuilder), + r#"CREATE TABLE "glyph" ( "id" integer NOT NULL ) PARTITION BY HASH ("id")"# + ); +} + +#[test] +fn create_partition_child_hash() { + assert_eq!( + Table::create() + .table(Alias::new("glyph_p1")) + .partition_of(Glyph::Table) + .values_with(4, 0) + .to_string(PostgresQueryBuilder), + r#"CREATE TABLE "glyph_p1" PARTITION OF "glyph" FOR VALUES WITH (MODULUS 4, REMAINDER 0)"# + ); +} From d5f62f9ac559015d936f3e73caebc14fa533965e Mon Sep 17 00:00:00 2001 From: Erdem Date: Sun, 11 Jan 2026 17:15:32 +0300 Subject: [PATCH 2/2] refactor: replace SimpleExpr with Expr in table partitioning --- src/backend/mysql/table.rs | 2 +- src/backend/postgres/table.rs | 2 +- src/table/create.rs | 20 ++++++++++---------- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/src/backend/mysql/table.rs b/src/backend/mysql/table.rs index c54af4363..a176ff54f 100644 --- a/src/backend/mysql/table.rs +++ b/src/backend/mysql/table.rs @@ -316,7 +316,7 @@ impl MysqlQueryBuilder { } } - fn prepare_partition_exprs(&self, exprs: &[SimpleExpr], sql: &mut impl SqlWriter) { + fn prepare_partition_exprs(&self, exprs: &[Expr], sql: &mut impl SqlWriter) { let mut first = true; for expr in exprs { if !first { diff --git a/src/backend/postgres/table.rs b/src/backend/postgres/table.rs index 21ebc4c85..246ba7d2f 100644 --- a/src/backend/postgres/table.rs +++ b/src/backend/postgres/table.rs @@ -283,7 +283,7 @@ impl PostgresQueryBuilder { } } - fn prepare_partition_exprs(&self, exprs: &[SimpleExpr], sql: &mut impl SqlWriter) { + fn prepare_partition_exprs(&self, exprs: &[Expr], sql: &mut impl SqlWriter) { let mut first = true; for expr in exprs { if !first { diff --git a/src/table/create.rs b/src/table/create.rs index 036a3a33e..4ed3ab8ec 100644 --- a/src/table/create.rs +++ b/src/table/create.rs @@ -1,8 +1,8 @@ use inherent::inherent; use crate::{ - ColumnDef, IntoColumnDef, SchemaStatementBuilder, SimpleExpr, backend::SchemaBuilder, - foreign_key::*, index::*, table::constraint::Check, types::*, + ColumnDef, Expr, IntoColumnDef, SchemaStatementBuilder, backend::SchemaBuilder, foreign_key::*, + index::*, table::constraint::Check, types::*, }; /// Create a table @@ -117,9 +117,9 @@ pub enum PartitionBy { #[derive(Debug, Clone)] pub enum PartitionValues { - In(Vec), - FromTo(Vec, Vec), - LessThan(Vec), + In(Vec), + FromTo(Vec, Vec), + LessThan(Vec), With(u32, u32), } @@ -361,7 +361,7 @@ impl TableCreateStatement { pub fn values_in(&mut self, values: I) -> &mut Self where I: IntoIterator, - T: Into, + T: Into, { self.partition_values = Some(PartitionValues::In( values.into_iter().map(|v| v.into()).collect(), @@ -369,13 +369,13 @@ impl TableCreateStatement { self } - /// Set partition values FROM ... TO. Postgres only. + /// Set partition values FROM ... TO .... Postgres only. pub fn values_from_to(&mut self, from: I, to: J) -> &mut Self where I: IntoIterator, - T: Into, + T: Into, J: IntoIterator, - U: Into, + U: Into, { self.partition_values = Some(PartitionValues::FromTo( from.into_iter().map(|v| v.into()).collect(), @@ -388,7 +388,7 @@ impl TableCreateStatement { pub fn values_less_than(&mut self, values: I) -> &mut Self where I: IntoIterator, - T: Into, + T: Into, { self.partition_values = Some(PartitionValues::LessThan( values.into_iter().map(|v| v.into()).collect(),