From b7f8dbb62c203bc7a4f11164286994459f13af06 Mon Sep 17 00:00:00 2001 From: Sean Klein Date: Tue, 28 Nov 2023 18:07:47 -0800 Subject: [PATCH 1/2] [nexus] Make 'update_and_check' CTE explicitly request columns --- nexus/db-queries/src/db/column_walker.rs | 114 ++++++++++++++++++++ nexus/db-queries/src/db/mod.rs | 1 + nexus/db-queries/src/db/update_and_check.rs | 48 +++++---- 3 files changed, 143 insertions(+), 20 deletions(-) create mode 100644 nexus/db-queries/src/db/column_walker.rs diff --git a/nexus/db-queries/src/db/column_walker.rs b/nexus/db-queries/src/db/column_walker.rs new file mode 100644 index 00000000000..75538c70433 --- /dev/null +++ b/nexus/db-queries/src/db/column_walker.rs @@ -0,0 +1,114 @@ +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! CTE utility for iterating over all columns in a table. + +use diesel::prelude::*; +use std::marker::PhantomData; + +/// Used to iterate over a tuple of columns ("T"). +/// +/// Diesel exposes "AllColumns" as a tuple, which is difficult to iterate over +/// -- after all, all the types are distinct. However, each of these types +/// implements "Column", so we can use a macro to provide a +/// "convertion-to-iterator" implemenation for our expected tuples. +pub(crate) struct ColumnWalker { + remaining: PhantomData, +} + +impl ColumnWalker { + pub fn new() -> Self { + Self { remaining: PhantomData } + } +} + +macro_rules! impl_column_walker { + ( $($column:ident)+ ) => ( + impl<$($column: Column),+> IntoIterator for ColumnWalker<($($column,)+)> { + type Item = &'static str; + type IntoIter = std::vec::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + // TODO: don't convert to vec? You'll need to figure + // out how to state the type of IntoIter. + [$($column::NAME,)+].to_vec().into_iter() + } + } + ); +} + +// implementations for 1 - 32 columns +impl_column_walker! { A } +impl_column_walker! { A B } +impl_column_walker! { A B C } +impl_column_walker! { A B C D } +impl_column_walker! { A B C D E } +impl_column_walker! { A B C D E F } +impl_column_walker! { A B C D E F G } +impl_column_walker! { A B C D E F G H } +impl_column_walker! { A B C D E F G H I } +impl_column_walker! { A B C D E F G H I J } +impl_column_walker! { A B C D E F G H I J K } +impl_column_walker! { A B C D E F G H I J K L } +impl_column_walker! { A B C D E F G H I J K L M } +impl_column_walker! { A B C D E F G H I J K L M N } +impl_column_walker! { A B C D E F G H I J K L M N O } +impl_column_walker! { A B C D E F G H I J K L M N O P } +impl_column_walker! { A B C D E F G H I J K L M N O P Q } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 C1 } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 C1 D1 } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 C1 D1 E1 } +impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 C1 D1 E1 F1 } + +#[cfg(test)] +mod test { + use super::*; + + table! { + test_schema.test_table (id) { + id -> Uuid, + value -> Int4, + time_deleted -> Nullable, + } + } + + // We can convert all a tables columns into an iteratable format. + #[test] + fn test_walk_table() { + let all_columns = + ColumnWalker::<::AllColumns>::new(); + + let mut iter = all_columns.into_iter(); + assert_eq!(iter.next(), Some("id")); + assert_eq!(iter.next(), Some("value")); + assert_eq!(iter.next(), Some("time_deleted")); + assert_eq!(iter.next(), None); + } + + // We can, if we want to, also make a ColumnWalker out of an arbitrary tuple + // of columns. + #[test] + fn test_walk_columns() { + let all_columns = ColumnWalker::<( + test_table::columns::id, + test_table::columns::value, + )>::new(); + + let mut iter = all_columns.into_iter(); + assert_eq!(iter.next(), Some("id")); + assert_eq!(iter.next(), Some("value")); + assert_eq!(iter.next(), None); + } +} diff --git a/nexus/db-queries/src/db/mod.rs b/nexus/db-queries/src/db/mod.rs index 8b7424a0565..b7c7079b541 100644 --- a/nexus/db-queries/src/db/mod.rs +++ b/nexus/db-queries/src/db/mod.rs @@ -12,6 +12,7 @@ pub mod collection_attach; pub mod collection_detach; pub mod collection_detach_many; pub mod collection_insert; +mod column_walker; mod config; mod cte_utils; // This is marked public for use by the integration tests diff --git a/nexus/db-queries/src/db/update_and_check.rs b/nexus/db-queries/src/db/update_and_check.rs index d6bf14c0834..fed79d5254a 100644 --- a/nexus/db-queries/src/db/update_and_check.rs +++ b/nexus/db-queries/src/db/update_and_check.rs @@ -4,6 +4,7 @@ //! CTE implementation for "UPDATE with extended return status". +use super::column_walker::ColumnWalker; use super::pool::DbConnection; use async_bb8_diesel::AsyncRunQueryDsl; use diesel::associations::HasTable; @@ -21,7 +22,7 @@ use std::marker::PhantomData; /// allows referencing generics with names (and extending usage /// without re-stating those generic parameters everywhere). pub trait UpdateStatementExt { - type Table: QuerySource; + type Table: Table + QuerySource; type WhereClause; type Changeset; @@ -32,7 +33,7 @@ pub trait UpdateStatementExt { impl UpdateStatementExt for UpdateStatement where - T: QuerySource, + T: Table + QuerySource, { type Table = T; type WhereClause = U; @@ -201,11 +202,11 @@ where /// /// ```text /// // WITH found AS (SELECT FROM T WHERE ) -/// // updated AS (UPDATE T SET RETURNING *) +/// // updated AS (UPDATE T SET RETURNING ) /// // SELECT /// // found. /// // updated. -/// // found.* +/// // found. /// // FROM /// // found /// // LEFT JOIN @@ -217,41 +218,48 @@ impl QueryFragment for UpdateAndQueryStatement where US: UpdateStatementExt, US::Table: HasTable + Table, + ColumnWalker<<::Table as Table>::AllColumns>: + IntoIterator, PrimaryKey: diesel::Column, UpdateStatement: QueryFragment, { fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> QueryResult<()> { + let primary_key = as Column>::NAME; + out.push_sql("WITH found AS ("); self.find_subquery.walk_ast(out.reborrow())?; out.push_sql("), updated AS ("); self.update_statement.walk_ast(out.reborrow())?; - // TODO: Only need primary? Or would we actually want - // to pass the returned rows back through the result? - out.push_sql(" RETURNING *) "); + out.push_sql(" RETURNING "); + out.push_identifier(primary_key)?; + out.push_sql(") "); out.push_sql("SELECT"); - let name = as Column>::NAME; out.push_sql(" found."); - out.push_identifier(name)?; + out.push_identifier(primary_key)?; out.push_sql(", updated."); - out.push_identifier(name)?; - // TODO: I'd prefer to list all columns explicitly. But how? - // The types exist within Table::AllColumns, and each one - // has a name as "::Name". - // But Table::AllColumns is a tuple, which makes iteration - // a pain. - // - // TODO: Technically, we're repeating the PK here. - out.push_sql(", found.*"); + out.push_identifier(primary_key)?; + + // List all the "found" columns explicitly. + // This admittedly repeats the primary key, but that keeps the query + // "simple" since it returns all columns in the same order as + // AllColumns. + let all_columns = ColumnWalker::< + <::Table as Table>::AllColumns, + >::new(); + for column in all_columns.into_iter() { + out.push_sql(", found."); + out.push_identifier(column)?; + } out.push_sql(" FROM found LEFT JOIN updated ON"); out.push_sql(" found."); - out.push_identifier(name)?; + out.push_identifier(primary_key)?; out.push_sql(" = "); out.push_sql("updated."); - out.push_identifier(name)?; + out.push_identifier(primary_key)?; Ok(()) } From b8169350450dd66de51679f7bae0c2ff2f5f8a31 Mon Sep 17 00:00:00 2001 From: Sean Klein Date: Wed, 29 Nov 2023 10:45:04 -0800 Subject: [PATCH 2/2] no more vec conversion! --- nexus/db-queries/src/db/column_walker.rs | 72 ++++++++++++------------ 1 file changed, 35 insertions(+), 37 deletions(-) diff --git a/nexus/db-queries/src/db/column_walker.rs b/nexus/db-queries/src/db/column_walker.rs index 75538c70433..64c3b450c88 100644 --- a/nexus/db-queries/src/db/column_walker.rs +++ b/nexus/db-queries/src/db/column_walker.rs @@ -24,53 +24,51 @@ impl ColumnWalker { } macro_rules! impl_column_walker { - ( $($column:ident)+ ) => ( + ( $len:literal $($column:ident)+ ) => ( impl<$($column: Column),+> IntoIterator for ColumnWalker<($($column,)+)> { type Item = &'static str; - type IntoIter = std::vec::IntoIter; + type IntoIter = std::array::IntoIter; fn into_iter(self) -> Self::IntoIter { - // TODO: don't convert to vec? You'll need to figure - // out how to state the type of IntoIter. - [$($column::NAME,)+].to_vec().into_iter() + [$($column::NAME,)+].into_iter() } } ); } // implementations for 1 - 32 columns -impl_column_walker! { A } -impl_column_walker! { A B } -impl_column_walker! { A B C } -impl_column_walker! { A B C D } -impl_column_walker! { A B C D E } -impl_column_walker! { A B C D E F } -impl_column_walker! { A B C D E F G } -impl_column_walker! { A B C D E F G H } -impl_column_walker! { A B C D E F G H I } -impl_column_walker! { A B C D E F G H I J } -impl_column_walker! { A B C D E F G H I J K } -impl_column_walker! { A B C D E F G H I J K L } -impl_column_walker! { A B C D E F G H I J K L M } -impl_column_walker! { A B C D E F G H I J K L M N } -impl_column_walker! { A B C D E F G H I J K L M N O } -impl_column_walker! { A B C D E F G H I J K L M N O P } -impl_column_walker! { A B C D E F G H I J K L M N O P Q } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 C1 } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 C1 D1 } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 C1 D1 E1 } -impl_column_walker! { A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 C1 D1 E1 F1 } +impl_column_walker! { 1 A } +impl_column_walker! { 2 A B } +impl_column_walker! { 3 A B C } +impl_column_walker! { 4 A B C D } +impl_column_walker! { 5 A B C D E } +impl_column_walker! { 6 A B C D E F } +impl_column_walker! { 7 A B C D E F G } +impl_column_walker! { 8 A B C D E F G H } +impl_column_walker! { 9 A B C D E F G H I } +impl_column_walker! { 10 A B C D E F G H I J } +impl_column_walker! { 11 A B C D E F G H I J K } +impl_column_walker! { 12 A B C D E F G H I J K L } +impl_column_walker! { 13 A B C D E F G H I J K L M } +impl_column_walker! { 14 A B C D E F G H I J K L M N } +impl_column_walker! { 15 A B C D E F G H I J K L M N O } +impl_column_walker! { 16 A B C D E F G H I J K L M N O P } +impl_column_walker! { 17 A B C D E F G H I J K L M N O P Q } +impl_column_walker! { 18 A B C D E F G H I J K L M N O P Q R } +impl_column_walker! { 19 A B C D E F G H I J K L M N O P Q R S } +impl_column_walker! { 20 A B C D E F G H I J K L M N O P Q R S T } +impl_column_walker! { 21 A B C D E F G H I J K L M N O P Q R S T U } +impl_column_walker! { 22 A B C D E F G H I J K L M N O P Q R S T U V } +impl_column_walker! { 23 A B C D E F G H I J K L M N O P Q R S T U V W } +impl_column_walker! { 24 A B C D E F G H I J K L M N O P Q R S T U V W X } +impl_column_walker! { 25 A B C D E F G H I J K L M N O P Q R S T U V W X Y } +impl_column_walker! { 26 A B C D E F G H I J K L M N O P Q R S T U V W X Y Z } +impl_column_walker! { 27 A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 } +impl_column_walker! { 28 A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 } +impl_column_walker! { 29 A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 C1 } +impl_column_walker! { 30 A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 C1 D1 } +impl_column_walker! { 31 A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 C1 D1 E1 } +impl_column_walker! { 32 A B C D E F G H I J K L M N O P Q R S T U V W X Y Z A1 B1 C1 D1 E1 F1 } #[cfg(test)] mod test {