Merge pull request #3921 from dvogel/pg-tablesample

Add Postgres TABLESAMPLE support.

Author: weiznich

diesel/src/internal/table_macro.rs

@@ -1,6 +1,9 @@
 #[doc(hidden)]
 pub use crate::expression::nullable::Nullable as NullableExpression;
 #[doc(hidden)]
+#[cfg(feature = "postgres_backend")]
+pub use crate::pg::query_builder::tablesample::TablesampleMethod;
+#[doc(hidden)]
 pub use crate::query_builder::from_clause::{FromClause, NoFromClause};
 #[doc(hidden)]
 pub use crate::query_builder::nodes::{

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 {

diesel/src/macros/mod.rs

@@ -243,6 +243,34 @@ macro_rules! __diesel_internal_backend_specific_allow_tables_to_appear_in_same_q
             for $left::table
         {
         }
+        impl<TSM> $crate::query_source::TableNotEqual<$left::table>
+            for $crate::query_builder::Tablesample<$right::table, TSM>
+        where
+            TSM: $crate::internal::table_macro::TablesampleMethod,
+        {
+        }
+        impl<TSM> $crate::query_source::TableNotEqual<$right::table>
+            for $crate::query_builder::Tablesample<$left::table, TSM>
+        where
+            TSM: $crate::internal::table_macro::TablesampleMethod,
+        {
+        }
+        impl<TSM>
+            $crate::query_source::TableNotEqual<
+                $crate::query_builder::Tablesample<$left::table, TSM>,
+            > for $right::table
+        where
+            TSM: $crate::internal::table_macro::TablesampleMethod,
+        {
+        }
+        impl<TSM>
+            $crate::query_source::TableNotEqual<
+                $crate::query_builder::Tablesample<$right::table, TSM>,
+            > for $left::table
+        where
+            TSM: $crate::internal::table_macro::TablesampleMethod,
+        {
+        }
     };
 }
 #[doc(hidden)]

diesel/src/pg/expression/extensions/mod.rs

@@ -3,6 +3,8 @@
 //! re-exported in `diesel::dsl`
 mod interval_dsl;
 mod only_dsl;
+mod tablesample_dsl;
 
 pub use self::interval_dsl::IntervalDsl;
 pub use self::only_dsl::OnlyDsl;
+pub use self::tablesample_dsl::TablesampleDsl;

diesel/src/pg/expression/extensions/tablesample_dsl.rs

@@ -0,0 +1,72 @@
+use crate::pg::query_builder::tablesample::{BernoulliMethod, SystemMethod};
+use crate::query_builder::Tablesample;
+use crate::Table;
+
+/// The `tablesample` method
+///
+/// The `TABLESAMPLE` clause is used to select a randomly sampled subset of rows from a table.
+///
+/// This is only implemented for the Postgres backend. While `TABLESAMPLE` is standardized in
+/// SQL:2003, in practice each RDBMS seems to implement a superset of the SQL:2003 syntax,
+/// supporting a wide variety of sampling methods.
+///
+/// Calling this function on a table (`mytable.tablesample(...)`) will result in the SQL
+/// `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.
+///
+/// To generate a `TABLESAMPLE ... REPEATABLE (f)` clause, you'll need to call
+/// [`.with_seed(f)`](Tablesample::with_seed).
+///
+/// Example:
+///
+/// ```rust
+/// # include!("../../../doctest_setup.rs");
+/// # use schema::{posts, users};
+/// # use diesel::dsl::*;
+/// # fn main() {
+/// # let connection = &mut establish_connection();
+/// let random_user_ids = users::table
+///     .tablesample_bernoulli(10)
+///     .select((users::id))
+///     .load::<i32>(connection);
+/// # }
+/// ```
+/// Selects the ids for a random 10 percent of users.
+///
+/// It can also be used in inner joins:
+///
+/// ```rust
+/// # include!("../../../doctest_setup.rs");
+/// # use schema::{posts, users};
+/// # use diesel::dsl::*;
+/// # fn main() {
+/// # let connection = &mut establish_connection();
+/// # let _ =
+/// users::table
+///     .tablesample_system(10).with_seed(42.0)
+///     .inner_join(posts::table)
+///     .select((users::name, posts::title))
+///     .load::<(String, String)>(connection);
+/// # }
+/// ```
+/// That query selects all of the posts for all of the users in a random 10 percent storage pages,
+/// 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_bernoulli(self, portion: i16) -> Tablesample<Self, BernoulliMethod> {
+        Tablesample::new(self, portion)
+    }
+
+    /// See the trait-level docs.
+    fn tablesample_system(self, portion: i16) -> Tablesample<Self, SystemMethod> {
+        Tablesample::new(self, portion)
+    }
+}
+
+impl<T: Table> TablesampleDsl for T {}

diesel/src/pg/query_builder/mod.rs

@@ -7,6 +7,7 @@ mod limit_offset;
 pub(crate) mod on_constraint;
 pub(crate) mod only;
 mod query_fragment_impls;
+pub(crate) mod tablesample;
 pub use self::distinct_on::DistinctOnClause;
 pub use self::distinct_on::OrderDecorator;
 

diesel/src/pg/query_builder/tablesample.rs

@@ -0,0 +1,214 @@
+use crate::expression::{Expression, ValidGrouping};
+use crate::pg::Pg;
+use crate::query_builder::{AsQuery, AstPass, FromClause, QueryFragment, QueryId, SelectStatement};
+use crate::query_source::QuerySource;
+use crate::result::QueryResult;
+use crate::sql_types::{Double, SmallInt};
+use crate::{JoinTo, SelectableExpression, Table};
+use std::marker::PhantomData;
+
+#[doc(hidden)]
+pub trait TablesampleMethod: Clone {
+    fn method_name_sql() -> &'static str;
+}
+
+#[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"
+    }
+}
+
+/// Represents a query with a `TABLESAMPLE` clause.
+#[derive(Debug, Clone, Copy)]
+pub struct Tablesample<S, TSM>
+where
+    TSM: TablesampleMethod,
+{
+    source: S,
+    method: PhantomData<TSM>,
+    portion: i16,
+    seed: Option<f64>,
+}
+
+impl<S, TSM> Tablesample<S, TSM>
+where
+    TSM: TablesampleMethod,
+{
+    pub(crate) fn new(source: S, portion: i16) -> Tablesample<S, TSM> {
+        Tablesample {
+            source,
+            method: PhantomData,
+            portion,
+            seed: None,
+        }
+    }
+
+    /// This method allows you to specify the random number generator seed to use in the sampling
+    /// method. This allows you to obtain repeatable results.
+    pub fn with_seed(self, seed: f64) -> Tablesample<S, TSM> {
+        Tablesample {
+            source: self.source,
+            method: self.method,
+            portion: self.portion,
+            seed: Some(seed),
+        }
+    }
+}
+
+impl<S, TSM> QueryId for Tablesample<S, TSM>
+where
+    S: QueryId,
+    TSM: TablesampleMethod,
+{
+    type QueryId = ();
+    const HAS_STATIC_QUERY_ID: bool = false;
+}
+
+impl<S, TSM> QuerySource for Tablesample<S, TSM>
+where
+    S: Table + Clone,
+    TSM: TablesampleMethod,
+    <S as QuerySource>::DefaultSelection:
+        ValidGrouping<()> + SelectableExpression<Tablesample<S, TSM>>,
+{
+    type FromClause = Self;
+    type DefaultSelection = <S as QuerySource>::DefaultSelection;
+
+    fn from_clause(&self) -> Self::FromClause {
+        self.clone()
+    }
+
+    fn default_selection(&self) -> Self::DefaultSelection {
+        self.source.default_selection()
+    }
+}
+
+impl<S, TSM> QueryFragment<Pg> for Tablesample<S, TSM>
+where
+    S: QueryFragment<Pg>,
+    TSM: TablesampleMethod,
+{
+    fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> QueryResult<()> {
+        self.source.walk_ast(out.reborrow())?;
+        out.push_sql(" TABLESAMPLE ");
+        out.push_sql(TSM::method_name_sql());
+        out.push_sql("(");
+        out.push_bind_param::<SmallInt, _>(&self.portion)?;
+        out.push_sql(")");
+        if let Some(f) = &self.seed {
+            out.push_sql(" REPEATABLE(");
+            out.push_bind_param::<Double, _>(f)?;
+            out.push_sql(")");
+        }
+        Ok(())
+    }
+}
+
+impl<S, TSM> AsQuery for Tablesample<S, TSM>
+where
+    S: Table + Clone,
+    TSM: TablesampleMethod,
+    <S as QuerySource>::DefaultSelection:
+        ValidGrouping<()> + SelectableExpression<Tablesample<S, TSM>>,
+{
+    type SqlType = <<Self as QuerySource>::DefaultSelection as Expression>::SqlType;
+    type Query = SelectStatement<FromClause<Self>>;
+    fn as_query(self) -> Self::Query {
+        SelectStatement::simple(self)
+    }
+}
+
+impl<S, T, TSM> JoinTo<T> for Tablesample<S, TSM>
+where
+    S: JoinTo<T>,
+    T: Table,
+    S: Table,
+    TSM: TablesampleMethod,
+{
+    type FromClause = <S as JoinTo<T>>::FromClause;
+    type OnClause = <S as JoinTo<T>>::OnClause;
+
+    fn join_target(rhs: T) -> (Self::FromClause, Self::OnClause) {
+        <S as JoinTo<T>>::join_target(rhs)
+    }
+}
+
+impl<S, TSM> Table for Tablesample<S, TSM>
+where
+    S: Table + Clone + AsQuery,
+    TSM: TablesampleMethod,
+
+    <S as Table>::PrimaryKey: SelectableExpression<Tablesample<S, TSM>>,
+    <S as Table>::AllColumns: SelectableExpression<Tablesample<S, TSM>>,
+    <S as QuerySource>::DefaultSelection:
+        ValidGrouping<()> + SelectableExpression<Tablesample<S, TSM>>,
+{
+    type PrimaryKey = <S as Table>::PrimaryKey;
+    type AllColumns = <S as Table>::AllColumns;
+
+    fn primary_key(&self) -> Self::PrimaryKey {
+        self.source.primary_key()
+    }
+
+    fn all_columns() -> Self::AllColumns {
+        S::all_columns()
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::backend::Backend;
+    use crate::pg::Pg;
+    use crate::query_builder::QueryBuilder;
+    use diesel::dsl::*;
+    use diesel::*;
+
+    macro_rules! assert_sql {
+        ($query:expr, $sql:expr) => {
+            let mut query_builder = <Pg as Backend>::QueryBuilder::default();
+            $query.to_sql(&mut query_builder, &Pg).unwrap();
+            let sql = query_builder.finish();
+            assert_eq!(sql, $sql);
+        };
+    }
+
+    table! {
+        users {
+            id -> Integer,
+            name -> VarChar,
+        }
+    }
+
+    #[test]
+    fn test_generated_tablesample_sql() {
+        assert_sql!(
+            users::table.tablesample_bernoulli(10),
+            "\"users\" TABLESAMPLE BERNOULLI($1)"
+        );
+
+        assert_sql!(
+            users::table.tablesample_system(10),
+            "\"users\" TABLESAMPLE SYSTEM($1)"
+        );
+
+        assert_sql!(
+            users::table.tablesample_system(10).with_seed(42.0),
+            "\"users\" TABLESAMPLE SYSTEM($1) REPEATABLE($2)"
+        );
+    }
+}

diesel/src/query_builder/mod.rs

@@ -120,6 +120,9 @@ pub(crate) use self::insert_statement::ColumnList;
 #[cfg(feature = "postgres_backend")]
 pub use crate::pg::query_builder::only::Only;
 
+#[cfg(feature = "postgres_backend")]
+pub use crate::pg::query_builder::tablesample::{Tablesample, TablesampleMethod};
+
 use crate::backend::Backend;
 use crate::result::QueryResult;
 use std::error::Error;