diff --git a/diesel/src/lib.rs b/diesel/src/lib.rs index ef8f22404528..96e0bf2817de 100644 --- a/diesel/src/lib.rs +++ b/diesel/src/lib.rs @@ -334,6 +334,14 @@ pub mod dsl { #[doc(inline)] pub use diesel_derives::auto_type; + + #[cfg(feature = "postgres_backend")] + #[doc(inline)] + pub use crate::pg::expression::extensions::OnlyDsl; + + #[cfg(feature = "postgres_backend")] + #[doc(inline)] + pub use crate::pg::expression::extensions::TablesampleDsl; } pub mod helper_types { diff --git a/diesel/src/pg/expression/extensions/mod.rs b/diesel/src/pg/expression/extensions/mod.rs index 1fc9be7aec09..98590b2ff9e6 100644 --- a/diesel/src/pg/expression/extensions/mod.rs +++ b/diesel/src/pg/expression/extensions/mod.rs @@ -7,4 +7,4 @@ mod tablesample_dsl; pub use self::interval_dsl::IntervalDsl; pub use self::only_dsl::OnlyDsl; -pub use self::tablesample_dsl::TablesampleDsl; +pub use self::tablesample_dsl::{BernoulliMethod, SystemMethod, TablesampleDsl}; diff --git a/diesel/src/pg/expression/extensions/tablesample_dsl.rs b/diesel/src/pg/expression/extensions/tablesample_dsl.rs index 1eb0b571308e..89f2a90b4009 100644 --- a/diesel/src/pg/expression/extensions/tablesample_dsl.rs +++ b/diesel/src/pg/expression/extensions/tablesample_dsl.rs @@ -1,6 +1,27 @@ use crate::query_builder::Tablesample; -pub(crate) use crate::query_builder::{TablesampleMethod, TablesampleSeed}; +pub(crate) use crate::query_builder::TablesampleMethod; use crate::Table; +use std::marker::PhantomData; + +#[derive(Clone, Copy, Debug)] +/// Used to specify the `BERNOULLI` sampling method. +pub struct BernoulliMethod; + +impl TablesampleMethod for BernoulliMethod { + fn method_name_sql() -> &'static str { + "BERNOULLI" + } +} + +#[derive(Clone, Copy, Debug)] +/// Used to specify the `SYSTEM` sampling method. +pub struct SystemMethod; + +impl TablesampleMethod for SystemMethod { + fn method_name_sql() -> &'static str { + "SYSTEM" + } +} /// The `tablesample` method /// @@ -11,23 +32,28 @@ use crate::Table; /// supporting a wide variety of sampling methods. /// /// Calling this function on a table (`mytable.tablesample(...)`) will result in the SQL -/// `FROM mytable TABLESAMPLE ...`. +/// `FROM mytable TABLESAMPLE ...` -- /// `mytable.tablesample(...)` can be used just like any table in diesel since it implements /// [Table](crate::Table). /// +/// The `BernoulliMethod` and `SystemMethod` types can be used to indicate the sampling method for +/// a `TABLESAMPLE method(p)` clause where p is specified by the portion argument. The provided +/// percentage should be an integer between 0 and 100. +/// +/// If the seed argument is is Some(f) then f becomes the seed in `TABLESAMPLE ... REPEATABLE (f)`. +/// /// Example: /// /// ```rust /// # include!("../../../doctest_setup.rs"); /// # use schema::{posts, users}; /// # use diesel::dsl::*; -/// # use crate::pg::query_builder::{TablesampleMethod, TablesampleSeed}; /// # fn main() { /// # let connection = &mut establish_connection(); /// let random_user_ids = users::table -/// .tablesample(TablesampleMethod::Bernoulli(10), TablesampleSeed::Auto) +/// .tablesample::(10, None) /// .select((users::id)) -/// .load::(connection); +/// .load::(connection); /// # } /// ``` /// Selects the ids for a random 10 percent of users. @@ -38,25 +64,30 @@ use crate::Table; /// # include!("../../../doctest_setup.rs"); /// # use schema::{posts, users}; /// # use diesel::dsl::*; -/// # use crate::query_builder::{TablesampleMethod, TablesampleSeed}; /// # fn main() { /// # let connection = &mut establish_connection(); /// # let _ = /// users::table -/// .tablesample(TablesampleMethod::Bernoulli(10), TablesampleSeed::Auto) -/// .inner_join(posts::table.only()) +/// .tablesample::(10, Some(42.0)) +/// .inner_join(posts::table) /// .select((users::name, posts::title)) /// .load::<(String, String)>(connection); /// # } /// ``` -/// That query selects all of the posts for a random 10 percent of users. +/// That query selects all of the posts for a random 10 percent of users, returning the same +/// results each time it is run due to the static seed of 42.0. /// pub trait TablesampleDsl: Table { /// See the trait-level docs. - fn tablesample(self, method: TablesampleMethod, seed: TablesampleSeed) -> Tablesample { + fn tablesample( + self, + portion: i16, + seed: Option, + ) -> Tablesample { Tablesample { source: self, - method, + method: PhantomData, + portion, seed, } } diff --git a/diesel/src/pg/query_builder/tablesample.rs b/diesel/src/pg/query_builder/tablesample.rs index bbd023e88a64..56d5a005edc3 100644 --- a/diesel/src/pg/query_builder/tablesample.rs +++ b/diesel/src/pg/query_builder/tablesample.rs @@ -7,49 +7,38 @@ use crate::{ sql_types::{Double, SmallInt}, JoinTo, SelectableExpression, Table, }; +use std::marker::PhantomData; -/// Indicates the sampling method for a `TABLESAMPLE method(n)` clause. The provided percentage -/// should be an integer between 0 and 100. -#[derive(Debug, Clone, Copy)] -pub enum TablesampleMethod { - /// Use the BERNOULLI sampline method. This is row-based, slower but more accurate. - Bernoulli(i16), - - /// Use the SYSTEM sampling method. This is page-based, faster but less accurate. - System(i16), -} - -/// Indicates the random number seed for a `TABLESAMPLE ... REPEATABLE(f)` clause. -#[derive(Debug, Clone, Copy)] -pub enum TablesampleSeed { - /// Have PostgreSQL generate an implied random number generator seed. - Auto, - - /// Provide your own random number generator seed. - Repeatable(f64), +#[doc(hidden)] +pub trait TablesampleMethod: Clone { + fn method_name_sql() -> &'static str; } /// Represents a query with a `TABLESAMPLE` clause. #[doc(hidden)] #[derive(Debug, Clone, Copy)] -pub struct Tablesample { +pub struct Tablesample { pub source: S, - pub method: TablesampleMethod, - pub seed: TablesampleSeed, + pub method: PhantomData, + pub portion: i16, + pub seed: Option, } -impl QueryId for Tablesample +impl QueryId for Tablesample where S: QueryId, + TSM: TablesampleMethod, { type QueryId = (); const HAS_STATIC_QUERY_ID: bool = false; } -impl QuerySource for Tablesample +impl QuerySource for Tablesample where S: Table + Clone, - ::DefaultSelection: ValidGrouping<()> + SelectableExpression>, + TSM: TablesampleMethod, + ::DefaultSelection: + ValidGrouping<()> + SelectableExpression>, { type FromClause = Self; type DefaultSelection = ::DefaultSelection; @@ -63,41 +52,33 @@ where } } -impl QueryFragment for Tablesample +impl QueryFragment for Tablesample where S: QueryFragment, + TSM: TablesampleMethod, { fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> QueryResult<()> { self.source.walk_ast(out.reborrow())?; out.push_sql(" TABLESAMPLE "); - match &self.method { - TablesampleMethod::Bernoulli(p) => { - out.push_sql("BERNOULLI("); - out.push_bind_param::(p)?; - out.push_sql(")"); - } - TablesampleMethod::System(p) => { - out.push_sql("SYSTEM("); - out.push_bind_param::(p)?; - out.push_sql(")"); - } - }; - match &self.seed { - TablesampleSeed::Auto => { /* no-op, this is the default */ } - TablesampleSeed::Repeatable(f) => { - out.push_sql(" REPEATABLE("); - out.push_bind_param::(f)?; - out.push_sql(")"); - } + out.push_sql(TSM::method_name_sql()); + out.push_sql("("); + out.push_bind_param::(&self.portion)?; + out.push_sql(")"); + if let Some(f) = &self.seed { + out.push_sql(" REPEATABLE("); + out.push_bind_param::(f)?; + out.push_sql(")"); } Ok(()) } } -impl AsQuery for Tablesample +impl AsQuery for Tablesample where S: Table + Clone, - ::DefaultSelection: ValidGrouping<()> + SelectableExpression>, + TSM: TablesampleMethod, + ::DefaultSelection: + ValidGrouping<()> + SelectableExpression>, { type SqlType = <::DefaultSelection as Expression>::SqlType; type Query = SelectStatement>; @@ -106,11 +87,12 @@ where } } -impl JoinTo for Tablesample +impl JoinTo for Tablesample where S: JoinTo, T: Table, S: Table, + TSM: TablesampleMethod, { type FromClause = >::FromClause; type OnClause = >::OnClause; @@ -120,13 +102,15 @@ where } } -impl Table for Tablesample +impl Table for Tablesample where S: Table + Clone + AsQuery, + TSM: TablesampleMethod, - ::PrimaryKey: SelectableExpression>, - ::AllColumns: SelectableExpression>, - ::DefaultSelection: ValidGrouping<()> + SelectableExpression>, + ::PrimaryKey: SelectableExpression>, + ::AllColumns: SelectableExpression>, + ::DefaultSelection: + ValidGrouping<()> + SelectableExpression>, { type PrimaryKey = ::PrimaryKey; type AllColumns = ::AllColumns; @@ -169,20 +153,17 @@ mod test { #[test] fn test_generated_tablesample_sql() { assert_sql!( - users::table.tablesample(TablesampleMethod::Bernoulli(10), TablesampleSeed::Auto), + users::table.tablesample::(10, None), "\"users\" TABLESAMPLE BERNOULLI($1)" ); assert_sql!( - users::table.tablesample(TablesampleMethod::System(10), TablesampleSeed::Auto), + users::table.tablesample::(10, None), "\"users\" TABLESAMPLE SYSTEM($1)" ); assert_sql!( - users::table.tablesample( - TablesampleMethod::System(10), - TablesampleSeed::Repeatable(42.0), - ), + users::table.tablesample::(10, Some(42.0),), "\"users\" TABLESAMPLE SYSTEM($1) REPEATABLE($2)" ); } diff --git a/diesel/src/query_builder/mod.rs b/diesel/src/query_builder/mod.rs index 6b45e04c0b71..6ce2750cbbf1 100644 --- a/diesel/src/query_builder/mod.rs +++ b/diesel/src/query_builder/mod.rs @@ -121,7 +121,7 @@ pub(crate) use self::insert_statement::ColumnList; pub use crate::pg::query_builder::only::Only; #[cfg(feature = "postgres_backend")] -pub use crate::pg::query_builder::tablesample::{Tablesample, TablesampleMethod, TablesampleSeed}; +pub use crate::pg::query_builder::tablesample::{Tablesample, TablesampleMethod}; use crate::backend::Backend; use crate::result::QueryResult; diff --git a/diesel_derives/src/table.rs b/diesel_derives/src/table.rs index 07ad02008e0b..875bfe1a69df 100644 --- a/diesel_derives/src/table.rs +++ b/diesel_derives/src/table.rs @@ -173,27 +173,32 @@ pub(crate) fn expand(input: TableDecl) -> TokenStream { type Count = diesel::query_source::Once; } - impl diesel::JoinTo> for table + impl diesel::JoinTo> for table where - diesel::query_builder::Tablesample: diesel::JoinTo, + diesel::query_builder::Tablesample: diesel::JoinTo
, + TSM: diesel::query_builder::TablesampleMethod { - type FromClause = diesel::query_builder::Tablesample; - type OnClause = as diesel::JoinTo
>::OnClause; + type FromClause = diesel::query_builder::Tablesample; + type OnClause = as diesel::JoinTo
>::OnClause; - fn join_target(__diesel_internal_rhs: diesel::query_builder::Tablesample) -> (Self::FromClause, Self::OnClause) { - let (_, __diesel_internal_on_clause) = diesel::query_builder::Tablesample::::join_target(table); + fn join_target(__diesel_internal_rhs: diesel::query_builder::Tablesample) -> (Self::FromClause, Self::OnClause) { + let (_, __diesel_internal_on_clause) = diesel::query_builder::Tablesample::::join_target(table); (__diesel_internal_rhs, __diesel_internal_on_clause) } } - impl diesel::query_source::AppearsInFromClause> + impl diesel::query_source::AppearsInFromClause> for table + where + TSM: diesel::query_builder::TablesampleMethod { type Count = diesel::query_source::Once; } - impl diesel::query_source::AppearsInFromClause
- for diesel::query_builder::Tablesample
+ impl diesel::query_source::AppearsInFromClause
+ for diesel::query_builder::Tablesample + where + TSM: diesel::query_builder::TablesampleMethod { type Count = diesel::query_source::Once; } @@ -693,12 +698,14 @@ fn expand_column_def(column_def: &ColumnDef) -> TokenStream { } impl diesel::SelectableExpression> for #column_name {} - impl diesel::query_source::AppearsInFromClause> + impl diesel::query_source::AppearsInFromClause> for #column_name + where TSM: diesel::query_builder::TablesampleMethod { type Count = diesel::query_source::Once; } - impl diesel::SelectableExpression> for #column_name {} + impl diesel::SelectableExpression> + for #column_name where TSM: diesel::query_builder::TablesampleMethod {} }) } else { None