Skip to content

Instantly share code, notes, and snippets.

@hgzimmerman
Last active February 20, 2024 18:16
Show Gist options
  • Save hgzimmerman/7440c270506dd90a758cbdcc1f26691e to your computer and use it in GitHub Desktop.
Save hgzimmerman/7440c270506dd90a758cbdcc1f26691e to your computer and use it in GitHub Desktop.
Diesel 2.x implementation of from_sql_function
//! # From SQL Function
//!
//! Probably 90% ready for general use, it should generate compatible output with the `diesel::table!()`
//! macro with the exception of the code commented out on the bottom of the main section.
//!
//! I ended up going in another direction and didn't end up using this in production, but for some time,
//! it replaced an ineffecient view with a marginally more performant function, and all of our tests passed,
//! so I think its usable.
//!
//!
//! ## What it does
//! It allows binding to a postgres SQL function that returns a table (ex: `CREATE FUNCTION function_name(...) RETURNS TABLE(...) ...`),
//! generating something that looks like a table to diesel. This "table" can then be bound to rust structs for easy integration with the rest of the code base.
//!
//! This macro also supports binding multiple functions to the same "table" representation.
//!
//! ## Example
//! ```rust
//! crate::from_sql_function!(
//! table_name = my_function,
//! my_function(
//! user_uuid: newtypes::UserUuid; Bound<diesel::sql_types::Uuid, uuid::Uuid>,
//! tenant_uuid: newtypes::TenantUuid; Bound<diesel::sql_types::Uuid, uuid::Uuid>
//! ) {
//! #[primary_key] user_uuid -> Uuid,
//! tenant_uuid -> Uuid,
//! count_one -> Int8,
//! count_two -> Int8,
//! count_three -> Int8,
//! status -> crate::schema::sql_types::Status,
//! }
//! )
//! ```
//!
//! ## Rules and Caveats
//! * The function names in the database _must_ match the name of the function here.
//! * Because this should work with diesel-async, the query itself needs to be both `Send` and `Sync`. Unfortunately,
//! the BoxableExpression used in the reference implementation cannot be used because it is not `Sync`.
//! This now requires passing in parameters like `; Bound<sql_type, external_type>` where `external_type`
//! must be the lowest level type that any wrapper in the actual arguments lowers to when binding.
//! * Multiple functions should be used in the query. While untested, it is easy to hypothesize
//! that diesel will have difficulty selecting the right columns if both are used at once.
//! * This does rely on a doc-private item of diesel, `Bound`. It may change in the future.
//! And most of diesel's machinery does rely on all of the traits in here being implemented.
//! If any of those change in a diesel 3.0 or a spicy 2.x release, we will have to look here and tweak this a little bit.
//!
//! ## Credits and references
//!
//! Based on this work: https://github.com/weiznich/wundergraph-workshop/blob/6785ce355d9e497909832bf9be79dede4c994c08/src/diesel_ext.rs
//! (By the maintainer of diesel)
//!
//! Extended to support diesel 2.0
//! This was tremendously helped by both cargo-expand to see what diesel already does and https://www.rustexplorer.com
//! to see what intermediate outputs of this macro were.
#[macro_export]
macro_rules! from_sql_function {
(
// Name of the table in rust, sets returned by functions called will be called by this name. No two functions that correspond to this pseudo-table should show up in the same query
table_name = $table_name: ident,
// Provide 1 or more function definitions, all of these functions must have identical returned rows.
$($fn_name: ident ($($arg: ident : $external_ty: ty; Bound<$arg_ty: ty, $bound_type: ty> ),*)),+
// The definition for the table, one field may have a `#[primary_key]` attribute to
{
$($(#[$($meta: tt)+])* $field_name: ident -> $field_ty: ty,)*
}
) => {
#[allow(dead_code)]
pub(crate) mod $table_name {
use diesel::query_source::*;
use diesel::expression::*;
use diesel::query_builder::*;
pub use self::columns::*;
#[derive(Clone)]
#[allow(non_camel_case_types, clippy::enum_variant_names)]
pub enum $table_name {
$($fn_name {$($arg: diesel::internal::derives::as_expression::Bound<$arg_ty, $bound_type>),*}),*
}
#[allow(non_camel_case_types)]
pub type table = $table_name;
impl Table for table {
type PrimaryKey = $crate::from_sql_function!(@collect_primary_key [] $($(#[$($meta)*])* $field_name,)*);
type AllColumns = ($(columns::$field_name,)*);
fn primary_key(&self) -> Self::PrimaryKey {
$crate::from_sql_function!(@collect_primary_key [] $($(#[$($meta)*])* $field_name,)*)
}
fn all_columns() -> Self::AllColumns {
($(columns::$field_name,)*)
}
}
impl QueryId for table {
type QueryId = ();
const HAS_STATIC_QUERY_ID: bool = false;
}
impl AppearsInFromClause<table> for table{
type Count = Once;
}
impl AppearsInFromClause<table> for () {
type Count = Never;
}
impl AsQuery for table {
type SqlType = <<Self as Table>::AllColumns as Expression>::SqlType;
type Query = SelectStatement<diesel::query_builder::FromClause<table>>;
fn as_query(self) -> Self::Query {
SelectStatement::simple(self)
}
}
impl<S> diesel::JoinTo<diesel::query_builder::Only<S>> for table
where
diesel::query_builder::Only<S>: diesel::JoinTo<table>,
{
type FromClause = diesel::query_builder::Only<S>;
type OnClause = <diesel::query_builder::Only<
S,
> as diesel::JoinTo<table>>::OnClause;
fn join_target(
__diesel_internal_rhs: diesel::query_builder::Only<S>,
) -> (Self::FromClause, Self::OnClause) {
let (_, __diesel_internal_on_clause) = diesel::query_builder::Only::<
S,
>::join_target(table::default());
(__diesel_internal_rhs, __diesel_internal_on_clause)
}
}
impl diesel::query_source::AppearsInFromClause<diesel::query_builder::Only<table>>
for table {
type Count = diesel::query_source::Once;
}
impl diesel::query_source::AppearsInFromClause<table>
for diesel::query_builder::Only<table> {
type Count = diesel::query_source::Once;
}
impl QuerySource for table {
type FromClause = Self;
type DefaultSelection = <Self as Table>::AllColumns;
fn from_clause(&self) -> Self::FromClause {
self.clone()
}
fn default_selection(&self) -> Self::DefaultSelection {
Self::all_columns()
}
}
impl diesel::associations::HasTable for table {
type Table = Self;
fn table() -> Self {
Self::default()
}
}
#[allow(unreachable_code)]
impl Default for table {
fn default() -> Self {
// expand all variants, but the first will return.
$(
return Self::$fn_name {
$($arg: diesel::expression::AsExpression::as_expression(<$external_ty>::default())),*
};
)*
}
}
// Need a module to allow destructuring of the args without colliding with
// columns with the same name due to the reexport.
#[allow(unused_mut, unused_variables)]
mod query_fragment {
use super::{table};
use diesel::query_builder::{QueryFragment, AstPass};
impl QueryFragment<diesel::pg::Pg> for table {
#[allow(dead_code, unused_assignments)]
fn walk_ast<'b>(&'b self, mut pass: AstPass<'_, 'b, diesel::pg::Pg>) -> diesel::result::QueryResult<()> {
match self {
// for different functions, we need to load different args into the fn call
$(Self::$fn_name{$($arg),*} => {
pass.push_sql(stringify!($fn_name));
pass.push_sql("(");
let mut first = true;
$(
if first {
first = false;
} else {
pass.push_sql(", ");
}
$arg.walk_ast(pass.reborrow())?;
)*
pass.push_sql(") AS ");
pass.push_sql(stringify!($table_name));
}),*
}
Ok(())
}
}
}
mod columns {
use diesel::sql_types::*;
use diesel::prelude::*;
use diesel::query_builder::{QueryFragment, AstPass};
$(
#[derive(Debug, Default, Copy, Clone)]
#[allow(non_camel_case_types)]
pub struct $field_name;
const _: () = {
use diesel;
use diesel::query_builder::QueryId;
#[allow(non_camel_case_types)]
impl QueryId for $field_name {
type QueryId = $field_name;
const HAS_STATIC_QUERY_ID: bool = true;
}
};
impl diesel::expression::Expression for $field_name {
type SqlType = $field_ty;
}
impl QueryFragment<diesel::pg::Pg> for $field_name {
fn walk_ast(&self, mut pass: AstPass<diesel::pg::Pg>) -> diesel::result::QueryResult<()> {
pass.push_identifier(stringify!($table_name))?; // must be fn name?
pass.push_sql(".");
pass.push_identifier(stringify!($field_name))?;
Ok(())
}
}
impl SelectableExpression<super::table> for $field_name {}
impl<QS> diesel::AppearsOnTable<QS> for $field_name
where
QS: diesel::query_source::AppearsInFromClause<
super::table,
Count = diesel::query_source::Once,
>,
{}
impl<
Left,
Right,
> diesel::SelectableExpression<
diesel::internal::table_macro::Join<
Left,
Right,
diesel::internal::table_macro::LeftOuter,
>,
> for $field_name
where
$field_name: diesel::AppearsOnTable<
diesel::internal::table_macro::Join<
Left,
Right,
diesel::internal::table_macro::LeftOuter,
>,
>,
Self: diesel::SelectableExpression<Left>,
Right: diesel::query_source::AppearsInFromClause<
super::table,
Count = diesel::query_source::Never,
> + diesel::query_source::QuerySource,
Left: diesel::query_source::QuerySource,
{}
impl<
Left,
Right,
> diesel::SelectableExpression<
diesel::internal::table_macro::Join<
Left,
Right,
diesel::internal::table_macro::Inner,
>,
> for $field_name
where
$field_name: diesel::AppearsOnTable<
diesel::internal::table_macro::Join<
Left,
Right,
diesel::internal::table_macro::Inner,
>,
>,
Left: diesel::query_source::AppearsInFromClause<super::table>
+ diesel::query_source::QuerySource,
Right: diesel::query_source::AppearsInFromClause<super::table>
+ diesel::query_source::QuerySource,
(
Left::Count,
Right::Count,
): diesel::internal::table_macro::Pick<Left, Right>,
Self: diesel::SelectableExpression<
<(
Left::Count,
Right::Count,
) as diesel::internal::table_macro::Pick<Left, Right>>::Selection,
>,
{}
impl<
Join,
On,
> diesel::SelectableExpression<diesel::internal::table_macro::JoinOn<Join, On>>
for $field_name
where
$field_name: diesel::SelectableExpression<Join>
+ diesel::AppearsOnTable<
diesel::internal::table_macro::JoinOn<Join, On>,
>,
{}
impl<
From,
> diesel::SelectableExpression<
diesel::internal::table_macro::SelectStatement<
diesel::internal::table_macro::FromClause<From>,
>,
> for $field_name
where
From: diesel::query_source::QuerySource,
$field_name: diesel::SelectableExpression<From>
+ diesel::AppearsOnTable<
diesel::internal::table_macro::SelectStatement<
diesel::internal::table_macro::FromClause<From>,
>,
>,
{}
impl<__GB> diesel::expression::ValidGrouping<__GB> for $field_name
where
__GB: diesel::expression::IsContainedInGroupBy<
$field_name,
Output = diesel::expression::is_contained_in_group_by::Yes,
>,
{
type IsAggregate = diesel::expression::is_aggregate::Yes;
}
impl diesel::expression::ValidGrouping<()> for $field_name {
type IsAggregate = diesel::expression::is_aggregate::No;
}
impl diesel::expression::IsContainedInGroupBy<$field_name> for $field_name {
type Output = diesel::expression::is_contained_in_group_by::Yes;
}
impl Column for $field_name {
type Table = super::table;
const NAME: &'static str = stringify!($field_name);
}
impl<T> diesel::EqAll<T> for $field_name where
T: diesel::expression::AsExpression<$field_ty>,
diesel::dsl::Eq<$field_name, T>: diesel::Expression<SqlType=diesel::sql_types::Bool>,
{
type Output = diesel::dsl::Eq<Self, T>;
fn eq_all(self, rhs: T) -> Self::Output {
self.eq(rhs)
}
}
impl diesel::query_source::AppearsInFromClause<
diesel::query_builder::Only<super::table>,
> for $field_name {
type Count = diesel::query_source::Once;
}
impl diesel::SelectableExpression<diesel::query_builder::Only<super::table>>
for $field_name {}
// There's a last little section I can't figure out how to implement in macro_rules! where
// this trait is implemented for this field against all other fields.
//
// It doesn't seem to be an impediment for the moment,
// I guess that we can't do group-by operations using these functions bound by this macro without this.
// impl diesel::expression::IsContainedInGroupBy<$other_field_name> for $field_name {
// type Output = diesel::expression::is_contained_in_group_by::Yes;
// }
//
// ...
// impl diesel::expression::IsContainedInGroupBy<$field_name> for $other_field_name {
// type Output = diesel::expression::is_contained_in_group_by::Yes;
// }
)*
}
#[allow(non_camel_case_types, unused_imports)]
pub(crate) mod function {
use diesel::sql_types::*;
$(
pub fn $fn_name($($arg: $external_ty,)*) -> super::$table_name
{
super::$table_name::$fn_name {
$($arg: diesel::expression::AsExpression::as_expression($arg),)*
}
}
)*
}
}
#[allow(dead_code,unused_imports)]
pub use self::$table_name::function::{$($fn_name),*};
};
(@collect_primary_key
[$($pk: ident,)*]
#[primary_key]
$field: ident, $($rest: tt)*
) => {
$crate::from_sql_function!(@collect_primary_key [$($pk,)* $field,] $($rest)*)
};
(@collect_primary_key
[$($pk: ident,)*]
$(#[$($meta: tt)*])*
$field: ident, $($rest: tt)*
) => {
$crate::from_sql_function!(@collect_primary_key [$($pk,)*] $($rest)*)
};
// One labeled as PK
(@collect_primary_key [$pk: ident,]) => {
$pk
};
// Many labeled as PK
(@collect_primary_key [$($pk: ident,)+]) => {
($($pk,)*)
};
// Nothing labeled as PK
(@collect_primary_key []) => {
columns::id
};
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment