Refactor raw paginator for affinity specifically (#7851)

Refactor some of the "raw pagination" logic into `RawPaginator`, to help
de-duplicate pagination
logic in cases that too complex (e.g., the result of JOINs, UNIONs) for
the existing Diesel pagination
tools.

Author: smklein

nexus/db-queries/src/db/column_walker.rs

@@ -28,6 +28,10 @@ impl<T> ColumnWalker<T> {
 
 macro_rules! impl_column_walker {
     ( $len:literal $($column:ident)+ ) => (
+        /// Returns all columns with the "(prefix)." string prepended.
+        ///
+        /// For example, with the columns "id, name" and the prefix "foo:
+        /// The output string would be: "foo.id, foo.name"
         #[allow(dead_code)]
         impl<$($column: Column),+> ColumnWalker<($($column,)+)> {
             pub fn with_prefix(prefix: &'static str) -> TrustedStr {
@@ -42,6 +46,46 @@ macro_rules! impl_column_walker {
             }
         }
 
+        /// Identical to [Self::with_prefix], but also aliases each column.
+        ///
+        /// The aliased name is "(prefix)_(column name)".
+        ///
+        /// For example, with the columns "id, name" and the prefix "foo:
+        /// The output string would be: "foo.id as foo_id, foo.name as foo_name"
+        #[allow(dead_code)]
+        impl<$($column: Column),+> ColumnWalker<($($column,)+)> {
+            pub fn with_prefix_and_alias(prefix: &'static str) -> TrustedStr {
+                // This string is derived from:
+                // - The "table" type, with associated columns, which
+                // are not controlled by an arbitrary user, and
+                // - The "prefix" type, which is a "&'static str" (AKA,
+                // hopefully known at compile-time, and not leaked).
+                TrustedStr::i_take_responsibility_for_validating_this_string(
+                    [
+                        $(
+                            format!(
+                                "{prefix}.{column} as {prefix}_{column}",
+                                prefix = prefix, column = $column::NAME
+                            ),
+                        )+
+                    ].join(", ")
+                )
+            }
+        }
+
+        /// Returns all columns without any suffix.
+        #[allow(dead_code)]
+        impl<$($column: Column),+> ColumnWalker<($($column,)+)> {
+            pub fn as_str() -> TrustedStr {
+                // This string is derived from:
+                // - The "table" type, with associated columns, which
+                // are not controlled by an arbitrary user
+                TrustedStr::i_take_responsibility_for_validating_this_string(
+                    [$($column::NAME,)+].join(", ")
+                )
+            }
+        }
+
         impl<$($column: Column),+> IntoIterator for ColumnWalker<($($column,)+)> {
             type Item = &'static str;
             type IntoIter = std::array::IntoIter<Self::Item, $len>;

nexus/db-queries/src/db/datastore/affinity.rs

@@ -10,6 +10,7 @@ use crate::authz::ApiResource;
 use crate::db;
 use crate::db::collection_insert::AsyncInsertError;
 use crate::db::collection_insert::DatastoreCollection;
+use crate::db::column_walker::AllColumnsOf;
 use crate::db::datastore::InstanceStateComputer;
 use crate::db::datastore::OpContext;
 use crate::db::error::ErrorHandler;
@@ -27,15 +28,15 @@ use crate::db::model::Name;
 use crate::db::model::Project;
 use crate::db::model::VmmState;
 use crate::db::model::VmmStateEnum;
-use crate::db::pagination::paginated;
-use crate::db::raw_query_builder::QueryBuilder;
+use crate::db::pagination::RawPaginator;
 use crate::transaction_retry::OptionalError;
 use async_bb8_diesel::AsyncRunQueryDsl;
 use chrono::Utc;
+use diesel::helper_types::AsSelect;
+use diesel::pg::Pg;
 use diesel::prelude::*;
 use omicron_common::api::external;
 use omicron_common::api::external::CreateResult;
-use omicron_common::api::external::DataPageParams;
 use omicron_common::api::external::DeleteResult;
 use omicron_common::api::external::Error;
 use omicron_common::api::external::ListResultVec;
@@ -47,7 +48,6 @@ use omicron_uuid_kinds::AffinityGroupUuid;
 use omicron_uuid_kinds::AntiAffinityGroupUuid;
 use omicron_uuid_kinds::GenericUuid;
 use omicron_uuid_kinds::InstanceUuid;
-use ref_cast::RefCast;
 use uuid::Uuid;
 
 impl DataStore {
@@ -61,22 +61,21 @@ impl DataStore {
 
         opctx.authorize(authz::Action::ListChildren, authz_project).await?;
 
-        match pagparams {
-            PaginatedBy::Id(pagparams) => {
-                paginated(dsl::affinity_group, dsl::id, &pagparams)
-            }
-            PaginatedBy::Name(pagparams) => paginated(
-                dsl::affinity_group,
-                dsl::name,
-                &pagparams.map_name(|n| Name::ref_cast(n)),
-            ),
-        }
-        .filter(dsl::project_id.eq(authz_project.id()))
-        .filter(dsl::time_deleted.is_null())
-        .select(AffinityGroup::as_select())
-        .get_results_async(&*self.pool_connection_authorized(opctx).await?)
-        .await
-        .map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))
+        let mut paginator = RawPaginator::new();
+        paginator
+            .source()
+            .sql("SELECT ")
+            .sql(AllColumnsOf::<dsl::affinity_group>::as_str())
+            .sql(" FROM affinity_group WHERE project_id = ")
+            .param()
+            .bind::<diesel::sql_types::Uuid, _>(authz_project.id())
+            .sql(" AND time_deleted IS NULL");
+        paginator
+            .paginate_by_id_or_name(pagparams)
+            .query::<AsSelect<AffinityGroup, Pg>>()
+            .get_results_async(&*self.pool_connection_authorized(opctx).await?)
+            .await
+            .map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))
     }
 
     pub async fn anti_affinity_group_list(
@@ -89,22 +88,21 @@ impl DataStore {
 
         opctx.authorize(authz::Action::ListChildren, authz_project).await?;
 
-        match pagparams {
-            PaginatedBy::Id(pagparams) => {
-                paginated(dsl::anti_affinity_group, dsl::id, &pagparams)
-            }
-            PaginatedBy::Name(pagparams) => paginated(
-                dsl::anti_affinity_group,
-                dsl::name,
-                &pagparams.map_name(|n| Name::ref_cast(n)),
-            ),
-        }
-        .filter(dsl::project_id.eq(authz_project.id()))
-        .filter(dsl::time_deleted.is_null())
-        .select(AntiAffinityGroup::as_select())
-        .get_results_async(&*self.pool_connection_authorized(opctx).await?)
-        .await
-        .map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))
+        let mut paginator = RawPaginator::new();
+        paginator
+            .source()
+            .sql("SELECT ")
+            .sql(AllColumnsOf::<dsl::anti_affinity_group>::as_str())
+            .sql(" FROM anti_affinity_group WHERE project_id = ")
+            .param()
+            .bind::<diesel::sql_types::Uuid, _>(authz_project.id())
+            .sql(" AND time_deleted IS NULL");
+        paginator
+            .paginate_by_id_or_name(pagparams)
+            .query::<AsSelect<AntiAffinityGroup, Pg>>()
+            .get_results_async(&*self.pool_connection_authorized(opctx).await?)
+            .await
+            .map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))
     }
 
     pub async fn affinity_group_create(
@@ -352,11 +350,13 @@ impl DataStore {
     ) -> ListResultVec<external::AffinityGroupMember> {
         opctx.authorize(authz::Action::Read, authz_affinity_group).await?;
 
-        let mut query = QueryBuilder::new()
+        let mut paginator = RawPaginator::new();
+        paginator
+            .source()
             .sql(
                 "
-                SELECT * FROM (
-                SELECT instance.id as id,
+                SELECT
+                    instance.id as id,
                     instance.name as name,
                     instance.state,
                     instance.migration_id,
@@ -372,56 +372,8 @@ impl DataStore {
                     group_id = ",
             )
             .param()
-            .bind::<diesel::sql_types::Uuid, _>(authz_affinity_group.id())
-            .sql(") ");
-
-        let (direction, limit) = match pagparams {
-            PaginatedBy::Id(p) => (p.direction, p.limit),
-            PaginatedBy::Name(p) => (p.direction, p.limit),
-        };
-        let asc = match direction {
-            dropshot::PaginationOrder::Ascending => true,
-            dropshot::PaginationOrder::Descending => false,
-        };
-
-        match pagparams {
-            PaginatedBy::Id(DataPageParams { marker, .. }) => {
-                if let Some(id) = marker {
-                    query = query
-                        .sql("WHERE id ")
-                        .sql(if asc { ">" } else { "<" })
-                        .sql(" ")
-                        .param()
-                        .bind::<diesel::sql_types::Uuid, _>(**id);
-                };
-                query = query.sql(" ORDER BY id ");
-            }
-            PaginatedBy::Name(DataPageParams { marker, .. }) => {
-                if let Some(name) = marker {
-                    query = query
-                        .sql("WHERE name ")
-                        .sql(if asc { ">" } else { "<" })
-                        .sql(" ")
-                        .param()
-                        .bind::<diesel::sql_types::Text, _>(Name(
-                            (*name).clone(),
-                        ));
-                };
-                query = query.sql(" ORDER BY name ");
-            }
-        }
-        if asc {
-            query = query.sql("ASC ");
-        } else {
-            query = query.sql("DESC ");
-        }
-
-        query = query
-            .sql(" LIMIT ")
-            .param()
-            .bind::<diesel::sql_types::BigInt, _>(i64::from(limit.get()));
-
-        query
+            .bind::<diesel::sql_types::Uuid, _>(authz_affinity_group.id());
+        paginator.paginate_by_id_or_name(pagparams)
             .query::<(
                 diesel::sql_types::Uuid,
                 diesel::sql_types::Text,
@@ -457,77 +409,28 @@ impl DataStore {
     ) -> ListResultVec<external::AntiAffinityGroupMember> {
         opctx.authorize(authz::Action::Read, authz_anti_affinity_group).await?;
 
-        let (direction, limit) = match pagparams {
-            PaginatedBy::Id(p) => (p.direction, p.limit),
-            PaginatedBy::Name(p) => (p.direction, p.limit),
-        };
-        let asc = match direction {
-            dropshot::PaginationOrder::Ascending => true,
-            dropshot::PaginationOrder::Descending => false,
-        };
-
-        let mut query = QueryBuilder::new()
-            .sql(
-                "SELECT id,name,instance_state,migration_id,vmm_state
-                FROM (
-                    SELECT
-                        instance.id as id,
-                        instance.name as name,
-                        instance.state as instance_state,
-                        instance.migration_id as migration_id,
-                        vmm.state as vmm_state
-                    FROM anti_affinity_group_instance_membership
-                    INNER JOIN instance
-                    ON instance.id = anti_affinity_group_instance_membership.instance_id
-                    LEFT JOIN vmm
-                    ON instance.active_propolis_id = vmm.id
-                    WHERE
-                        instance.time_deleted IS NULL AND
-                        vmm.time_deleted IS NULL AND
-                        group_id = ",
+        let mut paginator = RawPaginator::new();
+        paginator.source()
+            .sql("
+                SELECT
+                    instance.id as id,
+                    instance.name as name,
+                    instance.state as instance_state,
+                    instance.migration_id as migration_id,
+                    vmm.state as vmm_state
+                FROM anti_affinity_group_instance_membership
+                INNER JOIN instance
+                ON instance.id = anti_affinity_group_instance_membership.instance_id
+                LEFT JOIN vmm
+                ON instance.active_propolis_id = vmm.id
+                WHERE
+                    instance.time_deleted IS NULL AND
+                    vmm.time_deleted IS NULL AND
+                    group_id = ",
             )
             .param()
-            .bind::<diesel::sql_types::Uuid, _>(authz_anti_affinity_group.id())
-            .sql(") ");
-
-        match pagparams {
-            PaginatedBy::Id(DataPageParams { marker, .. }) => {
-                if let Some(id) = marker {
-                    query = query
-                        .sql("WHERE id ")
-                        .sql(if asc { ">" } else { "<" })
-                        .sql(" ")
-                        .param()
-                        .bind::<diesel::sql_types::Uuid, _>(**id);
-                };
-                query = query.sql(" ORDER BY id ");
-            }
-            PaginatedBy::Name(DataPageParams { marker, .. }) => {
-                if let Some(name) = marker {
-                    query = query
-                        .sql("WHERE name ")
-                        .sql(if asc { ">" } else { "<" })
-                        .sql(" ")
-                        .param()
-                        .bind::<diesel::sql_types::Text, _>(Name(
-                            (*name).clone(),
-                        ));
-                };
-                query = query.sql(" ORDER BY name ");
-            }
-        }
-        if asc {
-            query = query.sql("ASC ");
-        } else {
-            query = query.sql("DESC ");
-        }
-
-        query = query
-            .sql(" LIMIT ")
-            .param()
-            .bind::<diesel::sql_types::BigInt, _>(i64::from(limit.get()));
-
-        query
+            .bind::<diesel::sql_types::Uuid, _>(authz_anti_affinity_group.id());
+        paginator.paginate_by_id_or_name(pagparams)
             .query::<(
                 diesel::sql_types::Uuid,
                 diesel::sql_types::Text,

nexus/db-queries/src/db/datastore/cockroachdb_settings.rs

@@ -66,18 +66,17 @@ impl DataStore {
         type QueryRow = (sql_types::Text, sql_types::Text, sql_types::Text);
 
         let conn = self.pool_connection_authorized(opctx).await?;
-        let output: QueryOutput = QueryBuilder::new()
+        let mut query = QueryBuilder::new();
+        query
             .sql("SELECT ")
             .sql(STATE_FINGERPRINT_SQL)
             .sql(", * FROM ")
             .sql("[SHOW CLUSTER SETTING version], ")
-            .sql("[SHOW CLUSTER SETTING cluster.preserve_downgrade_option]")
-            .query::<QueryRow>()
-            .get_result_async(&*conn)
-            .await
-            .map_err(|err| {
-                public_error_from_diesel(err, ErrorHandler::Server)
-            })?;
+            .sql("[SHOW CLUSTER SETTING cluster.preserve_downgrade_option]");
+        let output: QueryOutput =
+            query.query::<QueryRow>().get_result_async(&*conn).await.map_err(
+                |err| public_error_from_diesel(err, ErrorHandler::Server),
+            )?;
         Ok(CockroachDbSettings {
             state_fingerprint: output.state_fingerprint,
             version: output.version,
@@ -96,7 +95,8 @@ impl DataStore {
         value: String,
     ) -> Result<(), Error> {
         let conn = self.pool_connection_authorized(opctx).await?;
-        QueryBuilder::new()
+        let mut query = QueryBuilder::new();
+        query
             .sql("SET CLUSTER SETTING ")
             .sql(setting)
             // `CASE` is the one conditional statement we get out of the
@@ -114,13 +114,10 @@ impl DataStore {
             // below in `test_state_fingerprint`).
             .sql(" ELSE NULL END")
             .bind::<sql_types::Text, _>(state_fingerprint)
-            .bind::<sql_types::Text, _>(value.clone())
-            .query::<()>()
-            .execute_async(&*conn)
-            .await
-            .map_err(|err| {
-                public_error_from_diesel(err, ErrorHandler::Server)
-            })?;
+            .bind::<sql_types::Text, _>(value.clone());
+        query.query::<()>().execute_async(&*conn).await.map_err(|err| {
+            public_error_from_diesel(err, ErrorHandler::Server)
+        })?;
         info!(
             opctx.log,
             "set cockroachdb setting";