Skip to content

Commit 08ce99b

Browse files
hawkwhawkw
committed
rewrite CTE to use raw query builder
as @smklein suggested in #9335 (comment)
1 parent cac60f3 commit 08ce99b

File tree

3 files changed

+171
-159
lines changed

3 files changed

+171
-159
lines changed

‎nexus/db-queries/src/db/datastore/fm.rs‎

Lines changed: 112 additions & 159 deletions
Original file line numberDiff line numberDiff line change
@@ -15,20 +15,21 @@ use crate::db::datastore::RunnableQuery;
1515
use crate::db::model;
1616
use crate::db::model::SqlU32;
1717
use crate::db::pagination::paginated;
18+
use crate::db::raw_query_builder::QueryBuilder;
19+
use crate::db::raw_query_builder::TypedSqlQuery;
1820
use async_bb8_diesel::AsyncRunQueryDsl;
1921
use diesel::pg::Pg;
2022
use diesel::prelude::*;
2123
use diesel::query_builder::AstPass;
22-
use diesel::query_builder::Query;
2324
use diesel::query_builder::QueryFragment;
2425
use diesel::query_builder::QueryId;
2526
use diesel::result::DatabaseErrorKind;
2627
use diesel::result::Error as DieselError;
2728
use diesel::sql_types;
29+
use dropshot::PaginationOrder;
2830
use nexus_db_errors::ErrorHandler;
2931
use nexus_db_errors::public_error_from_diesel;
3032
use nexus_db_lookup::DbConnection;
31-
use nexus_db_schema::schema;
3233
use nexus_db_schema::schema::fm_sitrep::dsl as sitrep_dsl;
3334
use nexus_db_schema::schema::fm_sitrep_history::dsl as history_dsl;
3435
use nexus_types::fm;
@@ -224,7 +225,7 @@ impl DataStore {
224225
// TODO(eliza): there should probably be an authz object for the fm sitrep?
225226
opctx.authorize(authz::Action::Modify, &authz::FLEET).await?;
226227

227-
let list = ListOrphanedQuery::new(&pagparams)
228+
let list = Self::list_orphaned_query(&pagparams)
228229
.load_async::<Uuid>(&*conn)
229230
.await
230231
.map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))?
@@ -234,6 +235,71 @@ impl DataStore {
234235
Ok(list)
235236
}
236237

238+
fn list_orphaned_query(
239+
pagparams: &DataPageParams<'_, SitrepUuid>,
240+
) -> TypedSqlQuery<sql_types::Uuid> {
241+
let mut builder = QueryBuilder::new();
242+
builder.sql(
243+
"WITH current_sitrep_id AS ( \
244+
SELECT sitrep_id \
245+
FROM omicron.public.fm_sitrep_history \
246+
ORDER BY version DESC \
247+
LIMIT 1 \
248+
),",
249+
);
250+
251+
// batch AS (
252+
// SELECT s.id, s.parent_sitrep_id
253+
// FROM omicron.public.fm_sitrep s
254+
// WHERE s.id > $1
255+
// ORDER BY s.id
256+
// LIMIT $2
257+
// )
258+
let (dir, cmp) = match pagparams.direction {
259+
PaginationOrder::Ascending => ("ASC", "> "),
260+
PaginationOrder::Descending => ("DESC", "< "),
261+
};
262+
builder.sql(
263+
"batch AS ( \
264+
SELECT s.id, s.parent_sitrep_id \
265+
FROM omicron.public.fm_sitrep s ",
266+
);
267+
if let Some(marker) = pagparams.marker {
268+
builder
269+
.sql("WHERE s.id ")
270+
.sql(cmp)
271+
.param()
272+
.bind::<sql_types::Uuid, _>(marker.into_untyped_uuid());
273+
}
274+
builder
275+
.sql(" ORDER BY s.id ")
276+
.sql(dir)
277+
.sql(" LIMIT ")
278+
.param()
279+
.bind::<sql_types::Int8, _>(pagparams.limit.get() as i64)
280+
.sql(") ");
281+
282+
builder.sql(
283+
"SELECT id \
284+
FROM omicron.public.fm_sitrep \
285+
WHERE id IN ( \
286+
SELECT b.id \
287+
FROM batch b \
288+
LEFT JOIN omicron.public.fm_sitrep_history h \
289+
ON h.sitrep_id = b.id \
290+
WHERE
291+
h.sitrep_id IS NULL \
292+
AND (
293+
b.parent_sitrep_id IS NULL \
294+
OR b.parent_sitrep_id != ( \
295+
SELECT sitrep_id FROM current_sitrep_id \
296+
) \
297+
) \
298+
);",
299+
);
300+
builder.query()
301+
}
302+
237303
/// Deletes all sitreps with the provided IDs.
238304
pub async fn fm_sitrep_delete_all(
239305
&self,
@@ -469,19 +535,21 @@ impl QueryId for InsertSitrepVersionQuery {
469535
type QueryId = ();
470536
const HAS_STATIC_QUERY_ID: bool = false;
471537
}
472-
473-
type FromClause<T> =
474-
diesel::internal::table_macro::StaticQueryFragmentInstance<T>;
475-
const SITREP_FROM_CLAUSE: FromClause<schema::fm_sitrep::table> =
476-
FromClause::new();
477-
const SITREP_HISTORY_FROM_CLAUSE: FromClause<schema::fm_sitrep_history::table> =
478-
FromClause::new();
479-
480538
impl QueryFragment<Pg> for InsertSitrepVersionQuery {
481539
fn walk_ast<'a>(
482540
&'a self,
483541
mut out: AstPass<'_, 'a, Pg>,
484542
) -> diesel::QueryResult<()> {
543+
use nexus_db_schema::schema;
544+
545+
type FromClause<T> =
546+
diesel::internal::table_macro::StaticQueryFragmentInstance<T>;
547+
const SITREP_FROM_CLAUSE: FromClause<schema::fm_sitrep::table> =
548+
FromClause::new();
549+
const SITREP_HISTORY_FROM_CLAUSE: FromClause<
550+
schema::fm_sitrep_history::table,
551+
> = FromClause::new();
552+
485553
const CURRENT_SITREP: &'static str = "current_sitrep";
486554

487555
out.push_sql("WITH ");
@@ -605,150 +673,13 @@ impl QueryFragment<Pg> for InsertSitrepVersionQuery {
605673

606674
impl RunQueryDsl<DbConnection> for InsertSitrepVersionQuery {}
607675

608-
/// CTE for listing orphaned sitreps (paginated).
609-
#[derive(Debug, Clone, Copy)]
610-
struct ListOrphanedQuery {
611-
marker: Option<SitrepUuid>,
612-
order: dropshot::PaginationOrder,
613-
limit: i32,
614-
}
615-
616-
impl ListOrphanedQuery {
617-
fn new(pagparams: &DataPageParams<'_, SitrepUuid>) -> Self {
618-
Self {
619-
marker: pagparams.marker.copied(),
620-
order: pagparams.direction,
621-
limit: pagparams.limit.get() as i32,
622-
}
623-
}
624-
}
625-
626-
impl QueryFragment<Pg> for ListOrphanedQuery {
627-
fn walk_ast<'a>(
628-
&'a self,
629-
mut out: AstPass<'_, 'a, Pg>,
630-
) -> diesel::QueryResult<()> {
631-
// WITH current_sitrep_id AS (
632-
// SELECT sitrep_id
633-
// FROM omicron.public.fm_sitrep_history
634-
// ORDER BY version DESC
635-
// LIMIT 1
636-
// ),
637-
const CURRENT_SITREP_ID: &'static str = "current_sitrep_id";
638-
out.push_sql("WITH ");
639-
out.push_identifier(CURRENT_SITREP_ID)?;
640-
out.push_sql(" AS (SELECT ");
641-
out.push_identifier(history_dsl::sitrep_id::NAME)?;
642-
out.push_sql(" FROM ");
643-
SITREP_HISTORY_FROM_CLAUSE.walk_ast(out.reborrow())?;
644-
out.push_sql(" ORDER BY ");
645-
out.push_identifier(history_dsl::version::NAME)?;
646-
out.push_sql(" DESC LIMIT 1),");
647-
648-
// batch AS (
649-
// SELECT s.id, s.parent_sitrep_id
650-
// FROM omicron.public.fm_sitrep s
651-
// WHERE s.id > $1
652-
// ORDER BY s.id
653-
// LIMIT $2
654-
// )
655-
const BATCH: &'static str = "batch";
656-
out.push_identifier(BATCH)?;
657-
out.push_sql(" AS (SELECT ");
658-
out.push_identifier(sitrep_dsl::id::NAME)?;
659-
out.push_sql(", ");
660-
out.push_identifier(sitrep_dsl::parent_sitrep_id::NAME)?;
661-
out.push_sql(" FROM ");
662-
SITREP_FROM_CLAUSE.walk_ast(out.reborrow())?;
663-
if let Some(ref marker) = self.marker {
664-
out.push_sql(" WHERE ");
665-
out.push_identifier(sitrep_dsl::id::NAME)?;
666-
out.push_sql(match self.order {
667-
dropshot::PaginationOrder::Ascending => " > ",
668-
dropshot::PaginationOrder::Descending => " < ",
669-
});
670-
out.push_bind_param::<sql_types::Uuid, Uuid>(
671-
marker.as_untyped_uuid(),
672-
)?;
673-
}
674-
out.push_sql(" ORDER BY ");
675-
out.push_identifier(sitrep_dsl::id::NAME)?;
676-
match self.order {
677-
dropshot::PaginationOrder::Ascending => out.push_sql(" ASC"),
678-
dropshot::PaginationOrder::Descending => out.push_sql(" DESC"),
679-
}
680-
out.push_sql(" LIMIT ");
681-
out.push_bind_param::<sql_types::Integer, i32>(&self.limit)?;
682-
out.push_sql(") ");
683-
684-
// SELECT id
685-
// FROM omicron.public.fm_sitrep
686-
// WHERE id IN (
687-
// SELECT b.id
688-
// FROM batch b
689-
// -- Lookup the sitrep in the history
690-
// LEFT JOIN omicron.public.fm_sitrep_history h ON h.sitrep_id = b.id
691-
// -- Find sitreps missing from history
692-
// WHERE h.sitrep_id IS NULL
693-
// -- ... where the sitrep cannot be made active
694-
// AND (b.parent_sitrep_id IS NULL
695-
// OR b.parent_sitrep_id != (SELECT sitrep_id FROM current_sitrep_id))
696-
// );
697-
out.push_sql("SELECT ");
698-
out.push_identifier(sitrep_dsl::id::NAME)?;
699-
out.push_sql(" FROM ");
700-
SITREP_FROM_CLAUSE.walk_ast(out.reborrow())?;
701-
out.push_sql(" WHERE ");
702-
out.push_identifier(sitrep_dsl::id::NAME)?;
703-
out.push_sql(" IN (");
704-
out.push_sql("SELECT ");
705-
out.push_identifier(BATCH)?;
706-
out.push_sql(".id");
707-
out.push_sql(" FROM ");
708-
out.push_identifier(BATCH)?;
709-
out.push_sql(" LEFT JOIN ");
710-
SITREP_HISTORY_FROM_CLAUSE.walk_ast(out.reborrow())?;
711-
out.push_sql(" h ON h.");
712-
out.push_identifier(history_dsl::sitrep_id::NAME)?;
713-
out.push_sql(" = ");
714-
out.push_identifier(BATCH)?;
715-
out.push_sql(".");
716-
out.push_identifier(sitrep_dsl::id::NAME)?;
717-
out.push_sql(" WHERE h.");
718-
out.push_identifier(history_dsl::sitrep_id::NAME)?;
719-
out.push_sql(" IS NULL AND (");
720-
out.push_identifier(BATCH)?;
721-
out.push_sql(".");
722-
out.push_identifier(sitrep_dsl::parent_sitrep_id::NAME)?;
723-
out.push_sql(" IS NULL OR ");
724-
out.push_identifier(BATCH)?;
725-
out.push_sql(".");
726-
out.push_identifier(sitrep_dsl::parent_sitrep_id::NAME)?;
727-
out.push_sql(" != (SELECT sitrep_id FROM ");
728-
out.push_identifier(CURRENT_SITREP_ID)?;
729-
out.push_sql(")));");
730-
731-
Ok(())
732-
}
733-
}
734-
735-
impl QueryId for ListOrphanedQuery {
736-
type QueryId = ();
737-
const HAS_STATIC_QUERY_ID: bool = false;
738-
}
739-
740-
impl RunQueryDsl<DbConnection> for ListOrphanedQuery {}
741-
742-
impl Query for ListOrphanedQuery {
743-
type SqlType = sql_types::Uuid;
744-
}
745-
746676
#[cfg(test)]
747677
mod tests {
748678
use super::*;
749679
use crate::db::explain::ExplainableAsync;
750680
use crate::db::pagination::Paginator;
751681
use crate::db::pub_test_utils::TestDatabase;
682+
use crate::db::raw_query_builder::expectorate_query_contents;
752683
use chrono::Utc;
753684
use nexus_types::fm;
754685
use omicron_test_utils::dev;
@@ -788,6 +719,37 @@ mod tests {
788719
logctx.cleanup_successful();
789720
}
790721

722+
// Ensure we have the right query contents.
723+
#[tokio::test]
724+
async fn expectorate_sitrep_list_orphans_no_marker() {
725+
let pagparams = DataPageParams {
726+
marker: None,
727+
limit: std::num::NonZeroU32::new(420).unwrap(),
728+
direction: dropshot::PaginationOrder::Descending,
729+
};
730+
let query = DataStore::list_orphaned_query(&pagparams);
731+
expectorate_query_contents(
732+
&query,
733+
"tests/output/sitrep_list_orphans_no_marker.sql",
734+
)
735+
.await;
736+
}
737+
738+
#[tokio::test]
739+
async fn expectorate_sitrep_list_orphans_with_marker() {
740+
let pagparams = DataPageParams {
741+
marker: Some(&SitrepUuid::nil()),
742+
limit: std::num::NonZeroU32::new(420).unwrap(),
743+
direction: dropshot::PaginationOrder::Descending,
744+
};
745+
let query = DataStore::list_orphaned_query(&pagparams);
746+
expectorate_query_contents(
747+
&query,
748+
"tests/output/sitrep_list_orphans_with_marker.sql",
749+
)
750+
.await;
751+
}
752+
791753
#[tokio::test]
792754
async fn explain_sitrep_list_orphaned_query() {
793755
let logctx = dev::test_setup_log("explain_sitrep_list_orphaned_query");
@@ -800,16 +762,7 @@ mod tests {
800762
limit: std::num::NonZeroU32::new(420).unwrap(),
801763
direction: dropshot::PaginationOrder::Descending,
802764
};
803-
let query = ListOrphanedQuery::new(&pagparams);
804-
805-
// Before trying to explain the query, elt's start by making sure it's
806-
// valid SQL...
807-
let q = diesel::debug_query::<Pg, _>(&query).to_string();
808-
match dev::db::format_sql(&q).await {
809-
Ok(q) => eprintln!("query: {q}"),
810-
Err(e) => panic!("query is malformed: {e}\n{q}"),
811-
}
812-
765+
let query = DataStore::list_orphaned_query(&pagparams);
813766
let explanation = query
814767
.explain_async(&conn)
815768
.await

‎nexus/db-queries/tests/output/sitrep_list_orphans_no_marker.sql‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
WITH
2+
current_sitrep_id
3+
AS (SELECT sitrep_id FROM omicron.public.fm_sitrep_history ORDER BY version DESC LIMIT 1),
4+
batch
5+
AS (
6+
SELECT s.id, s.parent_sitrep_id FROM omicron.public.fm_sitrep AS s ORDER BY s.id DESC LIMIT $1
7+
)
8+
SELECT
9+
id
10+
FROM
11+
omicron.public.fm_sitrep
12+
WHERE
13+
id
14+
IN (
15+
SELECT
16+
b.id
17+
FROM
18+
batch AS b LEFT JOIN omicron.public.fm_sitrep_history AS h ON h.sitrep_id = b.id
19+
WHERE
20+
h.sitrep_id IS NULL
21+
AND (
22+
b.parent_sitrep_id IS NULL
23+
OR b.parent_sitrep_id != (SELECT sitrep_id FROM current_sitrep_id)
24+
)
25+
)

0 commit comments

Comments
 (0)