Skip to content

Re-write Blueprint Target Insertion using QueryBuilder #9381

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Re-write Blueprint Target Insertion using QueryBuilder #9381

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8d19141
Blueprint insertion query using QueryBuilder
smklein Nov 10, 2025
43c2dee
minor cleanup
smklein Nov 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
320 changes: 147 additions & 173 deletions nexus/db-queries/src/db/datastore/deployment.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ use crate::db::pagination::Paginator;
use crate::db::pagination::paginated;
use crate::db::queries::ALLOW_FULL_TABLE_SCAN_SQL;
use crate::db::raw_query_builder::QueryBuilder;
use crate::db::raw_query_builder::TypedSqlQuery;
use anyhow::Context;
use async_bb8_diesel::AsyncRunQueryDsl;
use async_bb8_diesel::AsyncSimpleConnection;
Expand All @@ -20,21 +21,15 @@ use clickhouse_admin_types::{KeeperId, ServerId};
use core::future::Future;
use core::pin::Pin;
use diesel::BoolExpressionMethods;
use diesel::Column;
use diesel::ExpressionMethods;
use diesel::Insertable;
use diesel::IntoSql;
use diesel::JoinOnDsl;
use diesel::NullableExpressionMethods;
use diesel::OptionalExtension;
use diesel::QueryDsl;
use diesel::RunQueryDsl;
use diesel::Table;
use diesel::expression::SelectableHelper;
use diesel::pg::Pg;
use diesel::query_builder::AstPass;
use diesel::query_builder::QueryFragment;
use diesel::query_builder::QueryId;
use diesel::result::DatabaseErrorKind;
use diesel::result::Error as DieselError;
use diesel::sql_types;
Expand Down Expand Up @@ -1997,6 +1992,7 @@ impl DataStore {
};

query
.to_query()
.execute_async(conn)
.await
.map_err(|e| Error::from(query.decode_error(e)))?;
Expand Down Expand Up @@ -2804,10 +2800,10 @@ impl From<InsertTargetError> for Error {
/// -- veresion in `bp_target`).
/// current_target AS (
/// SELECT
/// "version" AS version,
/// "blueprint_id" AS blueprint_id
/// FROM "bp_target"
/// ORDER BY "version" DESC
/// version,
/// blueprint_id
/// FROM bp_target
/// ORDER BY version DESC
/// LIMIT 1
/// ),
///
Expand Down Expand Up @@ -2939,172 +2935,108 @@ impl InsertTargetQuery {
}
}

impl QueryId for InsertTargetQuery {
type QueryId = ();
const HAS_STATIC_QUERY_ID: bool = false;
}

impl QueryFragment<Pg> for InsertTargetQuery {
fn walk_ast<'a>(
&'a self,
mut out: AstPass<'_, 'a, Pg>,
) -> diesel::QueryResult<()> {
use nexus_db_schema::schema::blueprint::dsl as bp_dsl;
use nexus_db_schema::schema::bp_target::dsl;

type FromClause<T> =
diesel::internal::table_macro::StaticQueryFragmentInstance<T>;
type BpTargetFromClause =
FromClause<nexus_db_schema::schema::bp_target::table>;
type BlueprintFromClause =
FromClause<nexus_db_schema::schema::blueprint::table>;
const BP_TARGET_FROM_CLAUSE: BpTargetFromClause =
BpTargetFromClause::new();
const BLUEPRINT_FROM_CLAUSE: BlueprintFromClause =
BlueprintFromClause::new();

out.push_sql("WITH ");

out.push_sql("current_target AS (SELECT ");
out.push_identifier(dsl::version::NAME)?;
out.push_sql(" AS version,");
out.push_identifier(dsl::blueprint_id::NAME)?;
out.push_sql(" AS blueprint_id FROM ");
BP_TARGET_FROM_CLAUSE.walk_ast(out.reborrow())?;
out.push_sql(" ORDER BY ");
out.push_identifier(dsl::version::NAME)?;
out.push_sql(" DESC LIMIT 1),");

out.push_sql(
"check_validity AS MATERIALIZED ( \
SELECT \
CAST( \
IF( \
(SELECT ",
);
out.push_identifier(bp_dsl::id::NAME)?;
out.push_sql(" FROM ");
BLUEPRINT_FROM_CLAUSE.walk_ast(out.reborrow())?;
out.push_sql(" WHERE ");
out.push_identifier(bp_dsl::id::NAME)?;
out.push_sql(" = ");
out.push_bind_param::<sql_types::Uuid, Uuid>(
self.target_id.as_untyped_uuid(),
)?;
out.push_sql(") IS NULL, ");
out.push_bind_param::<sql_types::Text, &'static str>(
&NO_SUCH_BLUEPRINT_SENTINEL,
)?;
out.push_sql(
", \
IF( \
(SELECT ",
);
out.push_identifier(bp_dsl::parent_blueprint_id::NAME)?;
out.push_sql(" FROM ");
BLUEPRINT_FROM_CLAUSE.walk_ast(out.reborrow())?;
out.push_sql(", current_target WHERE ");
out.push_identifier(bp_dsl::id::NAME)?;
out.push_sql(" = ");
out.push_bind_param::<sql_types::Uuid, Uuid>(
self.target_id.as_untyped_uuid(),
)?;
out.push_sql(" AND current_target.blueprint_id = ");
out.push_identifier(bp_dsl::parent_blueprint_id::NAME)?;
out.push_sql(
" ) IS NOT NULL \
OR \
(SELECT 1 FROM ",
);
BLUEPRINT_FROM_CLAUSE.walk_ast(out.reborrow())?;
out.push_sql(" WHERE ");
out.push_identifier(bp_dsl::id::NAME)?;
out.push_sql(" = ");
out.push_bind_param::<sql_types::Uuid, Uuid>(
self.target_id.as_untyped_uuid(),
)?;
out.push_sql(" AND ");
out.push_identifier(bp_dsl::parent_blueprint_id::NAME)?;
out.push_sql(
" IS NULL \
AND NOT EXISTS ( \
SELECT version FROM current_target) \
) = 1, ",
);
out.push_sql(" CAST(");
out.push_bind_param::<sql_types::Uuid, Uuid>(
self.target_id.as_untyped_uuid(),
)?;
out.push_sql(" AS text), ");
out.push_bind_param::<sql_types::Text, &'static str>(
&PARENT_NOT_TARGET_SENTINEL,
)?;
out.push_sql(
" ) \
) \
AS UUID) \
), ",
);

out.push_sql("new_target AS (SELECT 1 AS new_version FROM ");
BLUEPRINT_FROM_CLAUSE.walk_ast(out.reborrow())?;
out.push_sql(" WHERE ");
out.push_identifier(bp_dsl::id::NAME)?;
out.push_sql(" = ");
out.push_bind_param::<sql_types::Uuid, Uuid>(
self.target_id.as_untyped_uuid(),
)?;
out.push_sql(" AND ");
out.push_identifier(bp_dsl::parent_blueprint_id::NAME)?;
out.push_sql(
" IS NULL \
AND NOT EXISTS \
(SELECT version FROM current_target) \
UNION \
SELECT current_target.version + 1 FROM \
current_target, ",
);
BLUEPRINT_FROM_CLAUSE.walk_ast(out.reborrow())?;
out.push_sql(" WHERE ");
out.push_identifier(bp_dsl::id::NAME)?;
out.push_sql(" = ");
out.push_bind_param::<sql_types::Uuid, Uuid>(
self.target_id.as_untyped_uuid(),
)?;
out.push_sql(" AND ");
out.push_identifier(bp_dsl::parent_blueprint_id::NAME)?;
out.push_sql(" IS NOT NULL AND ");
out.push_identifier(bp_dsl::parent_blueprint_id::NAME)?;
out.push_sql(" = current_target.blueprint_id) ");

out.push_sql("INSERT INTO ");
BP_TARGET_FROM_CLAUSE.walk_ast(out.reborrow())?;
out.push_sql("(");
out.push_identifier(dsl::version::NAME)?;
out.push_sql(",");
out.push_identifier(dsl::blueprint_id::NAME)?;
out.push_sql(",");
out.push_identifier(dsl::enabled::NAME)?;
out.push_sql(",");
out.push_identifier(dsl::time_made_target::NAME)?;
out.push_sql(") SELECT new_target.new_version, ");
out.push_bind_param::<sql_types::Uuid, Uuid>(
self.target_id.as_untyped_uuid(),
)?;
out.push_sql(",");
out.push_bind_param::<sql_types::Bool, bool>(&self.enabled)?;
out.push_sql(",");
out.push_bind_param::<sql_types::Timestamptz, DateTime<Utc>>(
&self.time_made_target,
)?;
out.push_sql(" FROM new_target");

Ok(())
impl InsertTargetQuery {
fn to_query(self) -> TypedSqlQuery<()> {
Comment on lines +2938 to +2939
Copy link
Member

@hawkw hawkw Nov 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

take it or leave it: now that this is using QueryBuilder, it doesn't really need to be its own struct; we could just have a function that takes the query parameters and returns a TypedSqlQuery. i think it only had to be a struct in order to implement the QueryFragment trait, but now we don't need a trait imp.

it's up to you whether you think it's cleaner to just have a free function or not, though 🤷‍♀️

let mut builder = QueryBuilder::new();

let target_id = *self.target_id.as_untyped_uuid();
let enabled = self.enabled;
let time_made_target = self.time_made_target;

builder.sql(
"WITH \
current_target AS ( \
SELECT \
version, \
blueprint_id \
FROM bp_target \
ORDER BY version DESC \
LIMIT 1 \
), \
check_validity AS MATERIALIZED ( \
SELECT \
CAST( \
IF( \
(SELECT id FROM blueprint WHERE id = ",
)
.param()
.bind::<sql_types::Uuid, _>(target_id)
.sql(") IS NULL, '")
.sql(NO_SUCH_BLUEPRINT_SENTINEL)
.sql(
"', \
Copy link
Member

@hawkw hawkw Nov 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

turbo nit: it looks kinda weird to me how the comma is not indented to where the previous expression would be, here...

IF( \
(SELECT parent_blueprint_id FROM blueprint, current_target \
WHERE id = ",
)
.param()
.bind::<sql_types::Uuid, _>(target_id)
.sql(
" \
AND current_target.blueprint_id = parent_blueprint_id \
) IS NOT NULL \
OR \
(SELECT 1 FROM blueprint \
WHERE id = ",
)
.param()
.bind::<sql_types::Uuid, _>(target_id)
.sql(
" \
AND parent_blueprint_id IS NULL \
AND NOT EXISTS (SELECT version FROM current_target) \
) = 1, \
CAST(",
)
.param()
.bind::<sql_types::Uuid, _>(target_id)
.sql(" AS text), '")
.sql(PARENT_NOT_TARGET_SENTINEL)
.sql(
"' \
) \
) AS UUID \
) \
), \
new_target AS ( \
SELECT 1 AS new_version FROM blueprint \
WHERE id = ",
)
.param()
.bind::<sql_types::Uuid, _>(target_id)
.sql(
" \
AND parent_blueprint_id IS NULL \
AND NOT EXISTS (SELECT version FROM current_target) \
UNION \
SELECT current_target.version + 1 \
FROM current_target, blueprint \
WHERE id = ",
)
.param()
.bind::<sql_types::Uuid, _>(target_id)
.sql(
" \
AND parent_blueprint_id IS NOT NULL \
AND parent_blueprint_id = current_target.blueprint_id \
) \
INSERT INTO bp_target(version, blueprint_id, enabled, time_made_target) \
SELECT new_target.new_version, ",
)
.param()
.bind::<sql_types::Uuid, _>(target_id)
.sql(", ")
.param()
.bind::<sql_types::Bool, _>(enabled)
.sql(", ")
.param()
.bind::<sql_types::Timestamptz, _>(time_made_target)
.sql(" FROM new_target");

builder.query()
}
}

impl RunQueryDsl<DbConnection> for InsertTargetQuery {}

#[cfg(test)]
mod tests {
use super::*;
Expand Down Expand Up @@ -4862,4 +4794,46 @@ mod tests {
);
}
}

#[tokio::test]
async fn insert_target_query_sql_snapshot() {
Copy link
Member

@hawkw hawkw Nov 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit, take it or leave it: most of the other tests of this form i've seen are named expectorate_{query}, maybe we should make them consistent?

use crate::db::raw_query_builder::expectorate_query_contents;

let query = InsertTargetQuery {
target_id: BlueprintUuid::nil(),
enabled: true,
time_made_target: Utc::now(),
};

expectorate_query_contents(
&query.to_query(),
"tests/output/insert_target_blueprint_query.sql",
)
.await;
}

#[tokio::test]
async fn insert_target_query_is_valid() {
Copy link
Member

@hawkw hawkw Nov 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

similar nit: most of these tests are named explain_{query}; should we make this consistent? also, the expectorate test also checks for the query's validity...

use crate::db::explain::ExplainableAsync;

let logctx = dev::test_setup_log("insert_target_query_is_valid");
let db = TestDatabase::new_with_pool(&logctx.log).await;
let pool = db.pool();
let conn = pool.claim().await.unwrap();

let query = InsertTargetQuery {
target_id: BlueprintUuid::nil(),
enabled: false,
time_made_target: Utc::now(),
};

let _ = query
.to_query()
.explain_async(&conn)
.await
Comment on lines +4830 to +4833
Copy link
Member

@hawkw hawkw Nov 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit, take it or leave it: i find it useful to also eprintln! the EXPLAIN output in these, so that you can run cargo nextest run {name_of_test} --success-output final if you want to see the query plan.

also, the other EXPLAIN tests generally include assertions that there's no FULL SCAN in the explanation, is that worth having here?

.expect("Failed to explain query - is it valid SQL?");

db.terminate().await;
logctx.cleanup_successful();
}
}
Loading
Loading