diff --git a/Cargo.lock b/Cargo.lock index d344ace..30e6cc6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2308,6 +2308,7 @@ dependencies = [ "serde-error", "serde_json", "snowflaked", + "sqlformat", "sqlx", "thiserror 2.0.18", "tokio", @@ -3494,6 +3495,16 @@ dependencies = [ "der", ] +[[package]] +name = "sqlformat" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0705994df478b895f05b8e290e0d46e53187b26f8d889d37b2a0881234922d94" +dependencies = [ + "unicode_categories", + "winnow", +] + [[package]] name = "sqlx" version = "0.8.6" @@ -4208,6 +4219,12 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + [[package]] name = "unindent" version = "0.2.4" @@ -4906,6 +4923,15 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" +[[package]] +name = "winnow" +version = "0.7.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" +dependencies = [ + "memchr", +] + [[package]] name = "wit-bindgen" version = "0.51.0" diff --git a/default.nix b/default.nix index 85e1d77..2e33771 100644 --- a/default.nix +++ b/default.nix @@ -46,6 +46,7 @@ let # Rust build tools pkgs.cargo-audit pkgs.cargo-edit + pkgs.cargo-insta pkgs.cargo-nextest pkgs.maturin diff --git a/opsqueue/Cargo.toml b/opsqueue/Cargo.toml index fd6b2e6..486ea65 100644 --- a/opsqueue/Cargo.toml +++ b/opsqueue/Cargo.toml @@ -75,6 +75,7 @@ humantime = "2.1.0" dashmap = "6.1.0" crossbeam-skiplist = "0.1.3" +sqlformat = "0.5.0" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html diff --git a/opsqueue/src/consumer/strategy.rs b/opsqueue/src/consumer/strategy.rs index 7aeffb2..68f0493 100644 --- a/opsqueue/src/consumer/strategy.rs +++ b/opsqueue/src/consumer/strategy.rs @@ -29,7 +29,6 @@ impl Strategy { metastate: &MetaState, ) -> &'a mut QueryBuilder<'a, Sqlite> { let qb = self.build_query_snippet(qb, metastate); - let qb = self.build_sort_order_query_snippet(qb); tracing::trace!("sql: {:?}", qb.sql()); qb } @@ -41,8 +40,8 @@ impl Strategy { ) -> &'a mut QueryBuilder<'a, Sqlite> { use Strategy::*; match self { - Oldest => qb.push("SELECT * FROM chunks"), - Newest => qb.push("SELECT * FROM chunks"), + Oldest => qb.push("SELECT * FROM chunks ORDER BY submission_id ASC"), + Newest => qb.push("SELECT * FROM chunks ORDER BY submission_id DESC"), Random => { let random_offset: u16 = rand::random(); qb.push("SELECT * FROM chunks WHERE random_order >= ") @@ -57,7 +56,6 @@ impl Strategy { } => { let qb = qb.push(format_args!("WITH inner_{meta_key} AS NOT MATERIALIZED (")); let qb = underlying.build_query_snippet(qb, metastate); - let qb = underlying.build_sort_order_query_snippet(qb); qb.push(format_args!( r#"), taken_{meta_key} AS ( @@ -76,7 +74,7 @@ impl Strategy { Some(field) => { let taken_values: Vec<_> = field.too_high_counts(1).collect(); let taken_values_string = - serde_json::to_string(&taken_values).expect("Always valid JSO"); + serde_json::to_string(&taken_values).expect("Always valid JSON"); tracing::trace!("Taken values that are left out of PreferDistinct: {taken_values_string:?}"); qb.push_bind(taken_values_string); } @@ -90,27 +88,6 @@ impl Strategy { } } } - - fn build_sort_order_query_snippet<'a>( - &'a self, - qb: &'a mut QueryBuilder<'a, Sqlite>, - ) -> &'a mut QueryBuilder<'a, Sqlite> { - use Strategy::*; - match self { - Oldest => qb.push("\nORDER BY submission_id ASC"), - Newest => qb.push("\nORDER BY submission_id DESC"), - Random => { - // It is **very** important that we do not apply extra sorting here. - // For the implementation of the 'cutting the deck' technique - // we rely on the order being 'whatever comes out of the UNION ALL' - qb - } - PreferDistinct { .. } => { - // **no** change in sort order. PreferDistinct passes the sort order on to the inner strategies that it unions. - qb - } - } - } } #[cfg(feature = "server-logic")] @@ -124,6 +101,7 @@ pub mod test { use super::*; use itertools::Itertools; + use sqlformat::{format, FormatOptions, QueryParams}; use sqlx::Row; use sqlx::{QueryBuilder, Sqlite, SqliteConnection}; @@ -131,10 +109,12 @@ pub mod test { qb: &mut sqlx::QueryBuilder<'_, Sqlite>, conn: &mut SqliteConnection, ) -> String { - sqlx::raw_sql(&format!("EXPLAIN QUERY PLAN {}", qb.sql())) + let formatted_query = format(qb.sql(), &QueryParams::None, &FormatOptions::default()); + + sqlx::raw_sql(&format!("EXPLAIN QUERY PLAN {}", formatted_query)) .fetch_all(&mut *conn) .await - .unwrap_or_else(|_| panic!("Invalid query: \n{}\n", qb.sql())) + .unwrap_or_else(|_| panic!("Invalid query: \n{}\n", formatted_query)) .into_iter() .map(|row| { let id = row.get::("id"); @@ -158,6 +138,16 @@ pub mod test { let metastate = MetaState::default(); let qb = Strategy::Oldest.build_query(&mut qb, &metastate); + let options = FormatOptions::default(); + let formatted_query = format(qb.sql(), &QueryParams::None, &options); + insta::assert_snapshot!(formatted_query, @" + SELECT + * + FROM + chunks + ORDER BY + submission_id ASC + "); let explained = explain(qb, &mut conn).await; assert_streaming_query(qb, &explained); @@ -171,6 +161,16 @@ pub mod test { let metastate = MetaState::default(); let qb = Strategy::Newest.build_query(&mut qb, &metastate); + let options = FormatOptions::default(); + let formatted_query = format(qb.sql(), &QueryParams::None, &options); + insta::assert_snapshot!(formatted_query, @" + SELECT + * + FROM + chunks + ORDER BY + submission_id DESC + "); let explained = explain(qb, &mut conn).await; assert_streaming_query(qb, &explained); @@ -184,8 +184,25 @@ pub mod test { let mut qb = QueryBuilder::new(""); let qb = Strategy::Random.build_query(&mut qb, &metastate); - let explained = explain(qb, &mut conn).await; + let formatted_query = format(qb.sql(), &QueryParams::None, &FormatOptions::default()); + insta::assert_snapshot!(formatted_query, @" + SELECT + * + FROM + chunks + WHERE + random_order >= ? + UNION ALL + SELECT + * + FROM + chunks + WHERE + random_order < ? + "); + + let explained = explain(qb, &mut conn).await; assert_streaming_query(qb, &explained); insta::assert_snapshot!(explained, @r" 1, 0, COMPOUND QUERY @@ -208,8 +225,62 @@ pub mod test { }; let mut qb = QueryBuilder::new(""); let qb = strategy.build_query(&mut qb, &metastate); - let explained = explain(qb, &mut conn).await; + let formatted_query = format(qb.sql(), &QueryParams::None, &FormatOptions::default()); + insta::assert_snapshot!(formatted_query, @" + WITH + inner_company_id AS NOT MATERIALIZED ( + SELECT + * + FROM + chunks + ORDER BY + submission_id ASC + ), + taken_company_id AS ( + SELECT + * + FROM + submissions_metadata + WHERE + submissions_metadata.metadata_key = ? + AND submissions_metadata.metadata_value IN ( + SELECT + value + FROM + json_each() + ) + ) + SELECT + * + FROM + inner_company_id + WHERE + NOT EXISTS ( + SELECT + 1 + FROM + taken_company_id + WHERE + inner_company_id.submission_id = taken_company_id.submission_id + ) + UNION ALL + SELECT + * + FROM + inner_company_id + WHERE + EXISTS ( + SELECT + 1 + FROM + taken_company_id + WHERE + inner_company_id.submission_id = taken_company_id.submission_id + ) + "); + + let explained = explain(qb, &mut conn).await; assert_streaming_query(qb, &explained); insta::assert_snapshot!(explained, @r" 1, 0, COMPOUND QUERY @@ -240,9 +311,62 @@ pub mod test { }; let mut qb = QueryBuilder::new(""); let qb = strategy.build_query(&mut qb, &metastate); - dbg!(qb.sql()); - let explained = explain(qb, &mut conn).await; + let formatted_query = format(qb.sql(), &QueryParams::None, &FormatOptions::default()); + insta::assert_snapshot!(formatted_query, @" + WITH + inner_company_id AS NOT MATERIALIZED ( + SELECT + * + FROM + chunks + ORDER BY + submission_id DESC + ), + taken_company_id AS ( + SELECT + * + FROM + submissions_metadata + WHERE + submissions_metadata.metadata_key = ? + AND submissions_metadata.metadata_value IN ( + SELECT + value + FROM + json_each() + ) + ) + SELECT + * + FROM + inner_company_id + WHERE + NOT EXISTS ( + SELECT + 1 + FROM + taken_company_id + WHERE + inner_company_id.submission_id = taken_company_id.submission_id + ) + UNION ALL + SELECT + * + FROM + inner_company_id + WHERE + EXISTS ( + SELECT + 1 + FROM + taken_company_id + WHERE + inner_company_id.submission_id = taken_company_id.submission_id + ) + "); + + let explained = explain(qb, &mut conn).await; assert_streaming_query(qb, &explained); insta::assert_snapshot!(explained, @r" 1, 0, COMPOUND QUERY @@ -273,8 +397,69 @@ pub mod test { }; let mut qb = QueryBuilder::new(""); let qb = strategy.build_query(&mut qb, &metastate); - let explained = explain(qb, &mut conn).await; + let formatted_query = format(qb.sql(), &QueryParams::None, &FormatOptions::default()); + insta::assert_snapshot!(formatted_query, @" + WITH + inner_company_id AS NOT MATERIALIZED ( + SELECT + * + FROM + chunks + WHERE + random_order >= ? + UNION ALL + SELECT + * + FROM + chunks + WHERE + random_order < ? + ), + taken_company_id AS ( + SELECT + * + FROM + submissions_metadata + WHERE + submissions_metadata.metadata_key = ? + AND submissions_metadata.metadata_value IN ( + SELECT + value + FROM + json_each() + ) + ) + SELECT + * + FROM + inner_company_id + WHERE + NOT EXISTS ( + SELECT + 1 + FROM + taken_company_id + WHERE + inner_company_id.submission_id = taken_company_id.submission_id + ) + UNION ALL + SELECT + * + FROM + inner_company_id + WHERE + EXISTS ( + SELECT + 1 + FROM + taken_company_id + WHERE + inner_company_id.submission_id = taken_company_id.submission_id + ) + "); + + let explained = explain(qb, &mut conn).await; assert_streaming_query(qb, &explained); insta::assert_snapshot!(explained, @r" 1, 0, COMPOUND QUERY @@ -325,11 +510,114 @@ pub mod test { let mut qb = QueryBuilder::new(""); let qb = strategy.build_query(&mut qb, &metastate); + + let formatted_query = format(qb.sql(), &QueryParams::None, &FormatOptions::default()); + insta::assert_snapshot!(formatted_query, @" + WITH + inner_company_id AS NOT MATERIALIZED ( + WITH + inner_priority AS NOT MATERIALIZED ( + SELECT + * + FROM + chunks + WHERE + random_order >= ? + UNION ALL + SELECT + * + FROM + chunks + WHERE + random_order < ? + ), + taken_priority AS ( + SELECT + * + FROM + submissions_metadata + WHERE + submissions_metadata.metadata_key = ? + AND submissions_metadata.metadata_value IN ( + SELECT + value + FROM + json_each() + ) + ) + SELECT + * + FROM + inner_priority + WHERE + NOT EXISTS ( + SELECT + 1 + FROM + taken_priority + WHERE + inner_priority.submission_id = taken_priority.submission_id + ) + UNION ALL + SELECT + * + FROM + inner_priority + WHERE + EXISTS ( + SELECT + 1 + FROM + taken_priority + WHERE + inner_priority.submission_id = taken_priority.submission_id + ) + ), + taken_company_id AS ( + SELECT + * + FROM + submissions_metadata + WHERE + submissions_metadata.metadata_key = ? + AND submissions_metadata.metadata_value IN ( + SELECT + value + FROM + json_each() + ) + ) + SELECT + * + FROM + inner_company_id + WHERE + NOT EXISTS ( + SELECT + 1 + FROM + taken_company_id + WHERE + inner_company_id.submission_id = taken_company_id.submission_id + ) + UNION ALL + SELECT + * + FROM + inner_company_id + WHERE + EXISTS ( + SELECT + 1 + FROM + taken_company_id + WHERE + inner_company_id.submission_id = taken_company_id.submission_id + ) + "); + let explained = explain(qb, &mut conn).await; - assert!( - !explained.contains("B-TREE"), - "Query should contain no materialization, but it did: {explained}" - ); + assert_streaming_query(qb, &explained); insta::assert_snapshot!(explained, @r" 1, 0, COMPOUND QUERY 2, 1, LEFT-MOST SUBQUERY