mz_mysql_util/
probe.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use mysql_async::prelude::Queryable;
11use mysql_async::{Params, Transaction, Value};
12use mz_ore::str::redact;
13
14use crate::{MySqlError, QualifiedTableRef, quote_identifier};
15
16/// The escape character for `LIKE` patterns built by [`like_prefix_pattern`].
17const LIKE_ESCAPE: char = '|';
18
19/// The longest key the probe bounds cover, in characters.
20/// <https://dev.mysql.com/doc/refman/8.4/en/innodb-limits.html> caps an index
21/// key at 3072 bytes, or 768 utf8mb4 characters. Longer keys (possible
22/// through prefix indexes or narrower charsets) are not supported.
23pub const MAX_KEY_LENGTH: u32 = 768;
24
25/// Probes a string primary key column. Only supports `utf8mb4_bin` against CHAR/VARCHAR
26/// columns up to 768 characters. Enforcement is deferred to the caller. There may be
27/// other collations we can support, but we should do more validation.
28///
29/// NOTE: The partitioning walk's local bookkeeping scales quadratically with the
30/// number of probes. Local benchmarking of 5k probes over 5k entries was fast
31/// (under a second of CPU), so exceed that budget at your own risk.
32pub struct KeyProber<'a, 't> {
33    tx: &'a mut Transaction<'t>,
34    /// Quoted `` `schema`.`table` `` for SQL interpolation.
35    table: String,
36    /// Quoted key column for SQL interpolation.
37    col: String,
38    /// Unquoted `schema.table` for error reporting.
39    table_name: String,
40    /// Unquoted key column for error reporting.
41    col_name: String,
42}
43
44impl<'a, 't> KeyProber<'a, 't> {
45    /// NOTE: `tx` is assumed to use a utf8mb4 connection character set (the
46    /// driver's handshake default), so key values arrive converted to UTF-8.
47    /// `tx` should be `REPEATABLE READ` so sequential probes see one
48    /// snapshot of the table.
49    pub fn new(tx: &'a mut Transaction<'t>, table: QualifiedTableRef<'_>, key_col: &str) -> Self {
50        Self {
51            tx,
52            table: format!(
53                "{}.{}",
54                quote_identifier(table.schema_name),
55                quote_identifier(table.table_name)
56            ),
57            col: quote_identifier(key_col),
58            table_name: format!("{}.{}", table.schema_name, table.table_name),
59            col_name: key_col.to_string(),
60        }
61    }
62
63    /// Estimates the row count for the given range. Estimates vary widely. On a static table
64    /// with 2.2B rows we observed estimates that should be near 2B report exactly half the
65    /// `TABLE_ROWS` reported by `information_schema.tables`. The sum of the row estimates
66    /// from this function were around 4B for the same test case, or about a 2x overcount relative
67    /// to the 2.05B reported by `TABLE_ROWS` reported by `information_schema.tables` and the 2.2B
68    /// actual rows. The underlying estimates are computed by sampling a small number of pages
69    /// after traversing the index (assuming this is a primary key being filtered on), so extrapolated
70    /// row counts can be inaccurate but appear to eventually converge towards more accurate estimates
71    /// as the sampled range shrinks on a static table.
72    pub async fn estimate_range_rows(
73        &mut self,
74        lower_bound_exclusive: &str,
75        upper_bound_exclusive: Option<&str>,
76    ) -> Result<u64, MySqlError> {
77        let (clause, params) =
78            self.range_filter(Some(lower_bound_exclusive), upper_bound_exclusive);
79        let select = format!(
80            "SELECT {col} FROM {table} WHERE {clause}",
81            col = self.col,
82            table = self.table,
83        );
84        explain_row_estimate(&mut *self.tx, &select, Params::Positional(params))
85            .await?
86            .ok_or_else(|| MySqlError::MissingRowEstimate {
87                qualified_table_name: self.table_name.clone(),
88                // The bounds are column values, redact them so the error
89                // stays loggable outside of CI.
90                lower_bound: format!("{:?}", redact(&lower_bound_exclusive)),
91                upper_bound: format!("{:?}", redact(&upper_bound_exclusive)),
92            })
93    }
94
95    /// Grabs a prefix of up to `max_prefix_length` characters for the first
96    /// key in the given range. If the key is shorter than `max_prefix_length`,
97    /// it returns that shorter value.
98    ///
99    /// The query will generally look something like:
100    ///
101    /// ```sql
102    /// SELECT LEFT(pk_col, 3) FROM table
103    /// WHERE pk_col > 'ab' AND pk_col < RPAD('ac', 768, CHAR(0))
104    /// ORDER BY pk_col
105    /// LIMIT 1
106    /// ```
107    pub async fn prefix_of_first_key_in_range(
108        &mut self,
109        lower_bound_exclusive: &str,
110        upper_bound_exclusive: Option<&str>,
111        max_prefix_length: usize,
112    ) -> Result<Option<String>, MySqlError> {
113        let (clause, params) =
114            self.range_filter(Some(lower_bound_exclusive), upper_bound_exclusive);
115        let sql = format!(
116            "SELECT LEFT({col}, {max_prefix_length}) FROM {table} WHERE {clause} ORDER BY {col} LIMIT 1",
117            col = self.col,
118            table = self.table,
119        );
120        self.query_string(sql, params).await
121    }
122
123    /// Returns the prefix of up to `max_prefix_length` characters of the first key after `prefix`,
124    /// but below `upper_bound_exclusive`. Returns None if no key matching these conditions exists.
125    ///
126    /// Note: this should be run inside a REPEATABLE READ transaction because
127    /// it issues two queries sequentially.
128    pub async fn prefix_of_first_row_not_matching_prefix(
129        &mut self,
130        prefix: &str,
131        upper_bound_exclusive: Option<&str>,
132        max_prefix_length: usize,
133    ) -> Result<Option<String>, MySqlError> {
134        let Some(max_key) = self
135            .max_key_with_prefix(prefix, upper_bound_exclusive)
136            .await?
137        else {
138            return Ok(None);
139        };
140        self.prefix_of_first_key_in_range(&max_key, upper_bound_exclusive, max_prefix_length)
141            .await
142    }
143
144    /// Quick way to grab the maximum key matching the prefix below the
145    /// exclusive upper bound.
146    ///
147    /// The query will generally look something like:
148    ///
149    /// ```sql
150    ///     SELECT pk_col FROM table
151    ///     WHERE pk_col LIKE /* prefix% */ 'abc%' AND pk_col < /* upper_bound_exclusive */ RPAD('ac', 768, CHAR(0))
152    ///     ORDER BY pk_col DESC
153    ///     LIMIT 1
154    /// ```
155    async fn max_key_with_prefix(
156        &mut self,
157        prefix: &str,
158        upper_bound_exclusive: Option<&str>,
159    ) -> Result<Option<String>, MySqlError> {
160        let (range_clause, range_params) = self.range_filter(None, upper_bound_exclusive);
161        let sql = format!(
162            "SELECT {col} FROM {table} WHERE {col} LIKE ? ESCAPE '{LIKE_ESCAPE}' \
163             AND {range_clause} ORDER BY {col} DESC LIMIT 1",
164            col = self.col,
165            table = self.table,
166        );
167        let mut params: Vec<Value> = vec![like_prefix_pattern(prefix).into()];
168        params.extend(range_params);
169        self.query_string(sql, params).await
170    }
171
172    /// Returns clause with upper and lower bounds enforced if present.
173    /// If both are None returns TRUE so this can plug in cleanly after a
174    /// leading "WHERE" or "AND".
175    ///
176    /// The upper bound is padded with NUL characters so that no key it
177    /// prefixes falls inside the range. Under PAD SPACE collations like
178    /// `utf8mb4_bin` "ab" is ordered as equivalent to "ab        " (however
179    /// many spaces are needed to fill remaining char/varchar length), so "ab\0"
180    /// sorts before either of those (because NUL is below all other characters
181    /// in `utf8mb4_bin`). Padding to [`MAX_KEY_LENGTH`] bounds every key an
182    /// utf8mb4 primary key column can hold.
183    fn range_filter(
184        &self,
185        lower_bound_exclusive: Option<&str>,
186        upper_bound_exclusive: Option<&str>,
187    ) -> (String, Vec<Value>) {
188        let col = &self.col;
189        let mut conditions = Vec::new();
190        let mut params: Vec<Value> = Vec::new();
191        if let Some(lower) = lower_bound_exclusive {
192            conditions.push(format!("{col} > ?"));
193            params.push(lower.into());
194        }
195        if let Some(upper) = upper_bound_exclusive {
196            conditions.push(format!(
197                "{col} < RPAD(?, {MAX_KEY_LENGTH}, CHAR(0 USING utf8mb4))"
198            ));
199            params.push(upper.into());
200        }
201        if conditions.is_empty() {
202            ("TRUE".to_string(), params)
203        } else {
204            (conditions.join(" AND "), params)
205        }
206    }
207
208    async fn query_string(
209        &mut self,
210        sql: String,
211        params: Vec<Value>,
212    ) -> Result<Option<String>, MySqlError> {
213        let row: Option<mysql_async::Row> =
214            self.tx.exec_first(sql, Params::Positional(params)).await?;
215        match row.and_then(|mut row| row.take_opt::<String, _>(0)) {
216            None => Ok(None),
217            Some(Ok(value)) => Ok(Some(value)),
218            Some(Err(err)) => Err(MySqlError::NonUtf8KeyValue {
219                qualified_table_name: self.table_name.clone(),
220                column_name: self.col_name.clone(),
221                error: err.to_string(),
222            }),
223        }
224    }
225}
226
227/// LIKE uses % and _ as wildcard characters. By default MySQL uses a backslash as an escape character
228/// but that can be disabled via config, so we instead specify a specific escape character.
229/// LIKE operator: <https://dev.mysql.com/doc/refman/8.0/en/string-comparison-functions.html#operator_like>
230/// Backslash escapes: <https://dev.mysql.com/doc/refman/8.0/en/sql-mode.html#sqlmode_no_backslash_escapes>
231fn like_prefix_pattern(prefix: &str) -> String {
232    let mut pattern = String::with_capacity(prefix.len() + 1);
233    for c in prefix.chars() {
234        if c == LIKE_ESCAPE || matches!(c, '%' | '_') {
235            pattern.push(LIKE_ESCAPE);
236        }
237        pattern.push(c);
238    }
239    pattern.push('%');
240    pattern
241}
242
243async fn explain_row_estimate<P>(
244    tx: &mut Transaction<'_>,
245    select: &str,
246    params: P,
247) -> Result<Option<u64>, MySqlError>
248where
249    P: Into<Params> + Send,
250{
251    let plan: Option<mysql_async::Row> = tx
252        .exec_first(format!("EXPLAIN FORMAT=TRADITIONAL {select}"), params)
253        .await?;
254    let estimate = plan.and_then(|row| {
255        row.get_opt::<Option<u64>, _>("rows")
256            .and_then(Result::ok)
257            .flatten()
258    });
259    Ok(estimate)
260}
261
262/// The live MySQL harness here is shared with [`crate::partition`]'s tests.
263#[cfg(test)]
264pub(crate) mod tests {
265    use std::collections::BTreeSet;
266
267    use mz_ore::cast::CastFrom;
268
269    use super::*;
270
271    #[mz_ore::test]
272    fn test_like_prefix_pattern() {
273        assert_eq!(like_prefix_pattern("abc"), "abc%");
274        assert_eq!(like_prefix_pattern(""), "%");
275        assert_eq!(like_prefix_pattern("a_b"), "a|_b%");
276        assert_eq!(like_prefix_pattern("50%"), "50|%%");
277        assert_eq!(like_prefix_pattern("a|b"), "a||b%");
278        // Backslash has no special meaning under an explicit ESCAPE '|'.
279        assert_eq!(like_prefix_pattern("a\\b"), "a\\b%");
280        assert_eq!(like_prefix_pattern("héllo"), "héllo%");
281    }
282
283    #[mz_ore::test(tokio::test)]
284    #[cfg_attr(miri, ignore)]
285    async fn test_basic_prefix_traversal() -> Result<(), anyhow::Error> {
286        let Some(mut conn) = connect().await? else {
287            return Ok(());
288        };
289        const DB: &str = "mz_probe_basic";
290        let keys = ["aa", "ab", "b", "bb", "bbb", "c"];
291        let table = setup_table(&mut conn, DB, "utf8mb4_bin", &keys).await?;
292
293        let mut tx = start_tx(&mut conn).await?;
294        let p = &mut KeyProber::new(&mut tx, table, "id");
295        assert_eq!(
296            prefix_of_first_key_in_range(p, "", None, 1).await,
297            some("a")
298        );
299        assert_eq!(
300            prefix_of_first_row_not_matching_prefix(p, "a", None, 1).await,
301            some("b")
302        );
303        assert_eq!(
304            prefix_of_first_row_not_matching_prefix(p, "b", None, 1).await,
305            some("c")
306        );
307        assert_eq!(
308            prefix_of_first_row_not_matching_prefix(p, "c", None, 1).await,
309            None
310        );
311
312        assert_eq!(
313            prefix_of_first_key_in_range(p, "a", Some("b"), 2).await,
314            some("aa")
315        );
316        assert_eq!(
317            prefix_of_first_row_not_matching_prefix(p, "aa", Some("b"), 2).await,
318            some("ab")
319        );
320        assert_eq!(
321            prefix_of_first_row_not_matching_prefix(p, "ab", Some("b"), 2).await,
322            None
323        );
324
325        // Bounds are exclusive: the exact key "b" is skipped as a split
326        // point, and its extensions surface as their own prefixes.
327        assert_eq!(
328            prefix_of_first_key_in_range(p, "b", Some("c"), 2).await,
329            some("bb")
330        );
331        assert_eq!(
332            prefix_of_first_row_not_matching_prefix(p, "bb", Some("c"), 2).await,
333            None
334        );
335        assert_eq!(prefix_of_first_key_in_range(p, "c", None, 2).await, None);
336
337        tx.rollback().await?;
338        drop_db(&mut conn, DB).await?;
339        conn.disconnect().await?;
340        Ok(())
341    }
342
343    #[mz_ore::test(tokio::test)]
344    #[cfg_attr(miri, ignore)]
345    async fn test_explain_row_estimate_sizing() -> Result<(), anyhow::Error> {
346        let Some(mut conn) = connect().await? else {
347            return Ok(());
348        };
349        const DB: &str = "mz_probe_explain_test";
350        let ids: Vec<String> = (0..1000).map(|i| format!("a{i:05}")).collect();
351        let table = setup_table(&mut conn, DB, "utf8mb4_bin", &ids).await?;
352        let mut tx = start_tx(&mut conn).await?;
353        let mut p = KeyProber::new(&mut tx, table, "id");
354
355        // Estimates are index dives, near reality but never exact by
356        // contract, so the bounds are deliberately loose.
357        let all = p.estimate_range_rows("", None).await?;
358        assert!((500..=2000).contains(&all), "all={all}");
359        let half = p.estimate_range_rows("a00500", None).await?;
360        assert!((250..=1000).contains(&half), "half={half}");
361        let none = p.estimate_range_rows("zzz", None).await?;
362        assert!(none <= 5, "none={none}");
363
364        tx.rollback().await?;
365        drop_db(&mut conn, DB).await?;
366        conn.disconnect().await?;
367        Ok(())
368    }
369
370    #[mz_ore::test(tokio::test)]
371    #[cfg_attr(miri, ignore)]
372    async fn test_case_insensitive_prefix_traversal() -> Result<(), anyhow::Error> {
373        let Some(mut conn) = connect().await? else {
374            return Ok(());
375        };
376        const DB: &str = "mz_probe_case_insensitive";
377        let keys = ["Aa", "ab", "b", "Bb", "bbb", "C"];
378        let table = setup_table(&mut conn, DB, "utf8mb4_general_ci", &keys).await?;
379
380        let mut tx = start_tx(&mut conn).await?;
381        let p = &mut KeyProber::new(&mut tx, table, "id");
382        // Sorting is case-insensitive but returned prefixes are the stored
383        // bytes: "A", "b", "C".
384        // Grab the initial prefix.
385        assert_eq!(
386            prefix_of_first_key_in_range(p, "", None, 1).await,
387            some("A")
388        );
389        // Traverse through sibling prefixes at depth 1.
390        assert_eq!(
391            prefix_of_first_row_not_matching_prefix(p, "A", None, 1).await,
392            some("b")
393        );
394        assert_eq!(
395            prefix_of_first_row_not_matching_prefix(p, "b", None, 1).await,
396            some("C")
397        );
398        assert_eq!(
399            prefix_of_first_row_not_matching_prefix(p, "C", None, 1).await,
400            None
401        );
402
403        // Same traversal at depth 2, bounded by the depth-1 prefixes.
404        assert_eq!(
405            prefix_of_first_key_in_range(p, "A", Some("b"), 2).await,
406            some("Aa")
407        );
408        assert_eq!(
409            prefix_of_first_row_not_matching_prefix(p, "Aa", Some("b"), 2).await,
410            some("ab")
411        );
412        assert_eq!(
413            prefix_of_first_row_not_matching_prefix(p, "ab", Some("b"), 2).await,
414            None
415        );
416
417        // The exclusive bound skips the exact key "b", and its extensions
418        // surface as their own prefixes under this case-insensitive
419        // collation.
420        assert_eq!(
421            prefix_of_first_key_in_range(p, "b", Some("C"), 2).await,
422            some("Bb")
423        );
424        assert_eq!(
425            prefix_of_first_row_not_matching_prefix(p, "Bb", Some("C"), 2).await,
426            None
427        );
428        // Every key matching 'b%' is covered by the prefix match.
429        assert_eq!(
430            prefix_of_first_row_not_matching_prefix(p, "b", Some("C"), 2).await,
431            None
432        );
433
434        assert_eq!(prefix_of_first_key_in_range(p, "C", None, 2).await, None);
435
436        tx.rollback().await?;
437        drop_db(&mut conn, DB).await?;
438        conn.disconnect().await?;
439        Ok(())
440    }
441
442    #[mz_ore::test(tokio::test)]
443    #[cfg_attr(miri, ignore)]
444    async fn test_wildcard_char_in_data() -> Result<(), anyhow::Error> {
445        let Some(mut conn) = connect().await? else {
446            return Ok(());
447        };
448        const DB: &str = "mz_probe_wildcard_test";
449        // Keys are a_1, a\2, a\3, a%4, a|5, covering the LIKE wildcards
450        // and the escape character itself. utf8mb4_bin orders them by byte:
451        // a%4 < a\2 < a\3 < a_1 < a|5.
452        let keys = ["a_1", "a\\2", "a\\3", "a%4", "a|5"];
453        let table = setup_table(&mut conn, DB, "utf8mb4_bin", &keys).await?;
454
455        let mut tx = start_tx(&mut conn).await?;
456        let p = &mut KeyProber::new(&mut tx, table, "id");
457        assert_eq!(
458            prefix_of_first_key_in_range(p, "", None, 1).await,
459            some("a")
460        );
461        assert_eq!(
462            prefix_of_first_row_not_matching_prefix(p, "a", None, 1).await,
463            None
464        );
465        assert_eq!(
466            prefix_of_first_key_in_range(p, "a", None, 2).await,
467            some("a%")
468        );
469        assert_eq!(
470            prefix_of_first_row_not_matching_prefix(p, "a%", None, 2).await,
471            some("a\\")
472        );
473        assert_eq!(
474            prefix_of_first_row_not_matching_prefix(p, "a\\", None, 2).await,
475            some("a_")
476        );
477        assert_eq!(
478            prefix_of_first_row_not_matching_prefix(p, "a_", None, 2).await,
479            some("a|")
480        );
481        assert_eq!(
482            prefix_of_first_row_not_matching_prefix(p, "a|", None, 2).await,
483            None
484        );
485
486        // Range bounds that are themselves wildcard characters.
487        assert_eq!(
488            prefix_of_first_key_in_range(p, "a%", Some("a\\"), 3).await,
489            some("a%4")
490        );
491        assert_eq!(
492            prefix_of_first_row_not_matching_prefix(p, "a%4", Some("a\\"), 3).await,
493            None
494        );
495        assert_eq!(
496            prefix_of_first_key_in_range(p, "a\\", Some("a_"), 3).await,
497            some("a\\2")
498        );
499        assert_eq!(
500            prefix_of_first_row_not_matching_prefix(p, "a\\2", Some("a_"), 3).await,
501            some("a\\3")
502        );
503        assert_eq!(
504            prefix_of_first_row_not_matching_prefix(p, "a\\3", Some("a_"), 3).await,
505            None
506        );
507        assert_eq!(
508            prefix_of_first_key_in_range(p, "a_", Some("a|"), 3).await,
509            some("a_1")
510        );
511        assert_eq!(
512            prefix_of_first_row_not_matching_prefix(p, "a_1", None, 3).await,
513            some("a|5")
514        );
515        assert_eq!(
516            prefix_of_first_key_in_range(p, "a|", None, 3).await,
517            some("a|5")
518        );
519        assert_eq!(
520            prefix_of_first_row_not_matching_prefix(p, "a|5", None, 3).await,
521            None
522        );
523
524        tx.rollback().await?;
525        drop_db(&mut conn, DB).await?;
526        conn.disconnect().await?;
527        Ok(())
528    }
529
530    #[mz_ore::test(tokio::test)]
531    #[cfg_attr(miri, ignore)]
532    async fn test_multibyte_chars_in_data() -> Result<(), anyhow::Error> {
533        let Some(mut conn) = connect().await? else {
534            return Ok(());
535        };
536        const DB: &str = "mz_probe_multibyte_test";
537        // utf8mb4_general_ci gives every supplementary character one shared
538        // weight, so emoji sort last: a < a😀 < 日本 < 日本語 < 😀 < 😀a < 😀😀.
539        let keys = ["a", "a😀", "日本", "日本語", "😀", "😀a", "😀😀"];
540        let table = setup_table(&mut conn, DB, "utf8mb4_general_ci", &keys).await?;
541
542        let mut tx = start_tx(&mut conn).await?;
543        let p = &mut KeyProber::new(&mut tx, table, "id");
544        // Prefix lengths count characters, not bytes: a one-char prefix of a
545        // four-byte emoji is the whole emoji, never a broken fragment.
546        assert_eq!(
547            prefix_of_first_key_in_range(p, "", None, 1).await,
548            some("a")
549        );
550        assert_eq!(
551            prefix_of_first_row_not_matching_prefix(p, "a", None, 1).await,
552            some("日")
553        );
554        assert_eq!(
555            prefix_of_first_row_not_matching_prefix(p, "日", None, 1).await,
556            some("😀")
557        );
558        assert_eq!(
559            prefix_of_first_row_not_matching_prefix(p, "😀", None, 1).await,
560            None
561        );
562
563        // Depth 2 walk for each prefix from depth 1.
564        assert_eq!(
565            prefix_of_first_key_in_range(p, "a", Some("日"), 2).await,
566            some("a😀")
567        );
568        assert_eq!(
569            prefix_of_first_row_not_matching_prefix(p, "a😀", Some("日"), 2).await,
570            None
571        );
572        assert_eq!(
573            prefix_of_first_key_in_range(p, "日", Some("😀"), 2).await,
574            some("日本")
575        );
576        assert_eq!(
577            prefix_of_first_row_not_matching_prefix(p, "日本", Some("😀"), 2).await,
578            None
579        );
580        assert_eq!(
581            prefix_of_first_key_in_range(p, "😀", None, 2).await,
582            some("😀a")
583        );
584        assert_eq!(
585            prefix_of_first_row_not_matching_prefix(p, "😀a", None, 2).await,
586            some("😀😀")
587        );
588        assert_eq!(
589            prefix_of_first_row_not_matching_prefix(p, "😀😀", None, 2).await,
590            None
591        );
592
593        tx.rollback().await?;
594        drop_db(&mut conn, DB).await?;
595        conn.disconnect().await?;
596        Ok(())
597    }
598
599    #[mz_ore::test(tokio::test)]
600    #[cfg_attr(miri, ignore)]
601    async fn test_live_mysql_uuid_pk() -> Result<(), anyhow::Error> {
602        let Some(mut conn) = connect().await? else {
603            return Ok(());
604        };
605        const DB: &str = "mz_probe_uuid_test";
606        // Hyphenated lowercase v4-shaped UUIDs, unique via the last group,
607        // with the leading group scattered like random UUIDs.
608        let ids: Vec<String> = (0..1000u64)
609            .map(|i| {
610                let h = i.wrapping_mul(2654435761) % 0x1_0000_0000;
611                format!("{h:08x}-0000-4000-8000-{i:012x}")
612            })
613            .collect();
614        let table = setup_table(&mut conn, DB, "utf8mb4_bin", &ids).await?;
615        let mut tx = start_tx(&mut conn).await?;
616        let mut p = KeyProber::new(&mut tx, table, "id");
617
618        // Lowercase hex order matches byte order under this collation.
619        assert_eq!(
620            prefix_of_first_key_in_range(&mut p, "", None, 36).await,
621            ids.iter().min().cloned()
622        );
623
624        let walked = walk_prefixes(&mut p, 1).await?;
625        let expected: Vec<String> = ids
626            .iter()
627            .map(|id| id[..1].to_string())
628            .collect::<BTreeSet<_>>()
629            .into_iter()
630            .collect();
631        assert_eq!(walked, expected);
632
633        tx.rollback().await?;
634        drop_db(&mut conn, DB).await?;
635        conn.disconnect().await?;
636        Ok(())
637    }
638
639    #[mz_ore::test(tokio::test)]
640    #[cfg_attr(miri, ignore)]
641    async fn test_live_mysql_like_metacharacters() -> Result<(), anyhow::Error> {
642        let Some(mut conn) = connect().await? else {
643            return Ok(());
644        };
645        const DB: &str = "mz_probe_like_test";
646        // Every walk step matches on `LIKE '<prefix>%'`, so keys whose
647        // prefixes are LIKE metacharacters exercise the escaping.
648        let ids = [
649            "%a",
650            "%%",
651            "_a",
652            "__",
653            "\\a",
654            "\\\\",
655            "a%",
656            "a%b",
657            "a_",
658            "a_b",
659            "a\\",
660            "a\\b",
661            "ab",
662            "a b",
663            "|a",
664            "||",
665            "a|",
666            "a|b",
667            "100%",
668            "50%off",
669            "under_score",
670            "back\\slash",
671        ];
672        let table = setup_table(&mut conn, DB, "utf8mb4_bin", &ids).await?;
673
674        // Assert that walking the prefixes gives range boundaries that
675        // partition the table and every key falls into exactly one range.
676        for len in [1, 2] {
677            let mut tx = start_tx(&mut conn).await?;
678            let walked =
679                walk_prefixes(&mut KeyProber::new(&mut tx, table.clone(), "id"), len).await?;
680            tx.rollback().await?;
681            let mut total = 0;
682            for (i, lo) in walked.iter().enumerate() {
683                let (n, prefixed) = count_range(&mut conn, DB, lo, walked.get(i + 1)).await?;
684                assert!(
685                    n > 0,
686                    "empty interval: len={len} lo={lo:?} walked={walked:?}"
687                );
688                assert_eq!(
689                    prefixed, n,
690                    "keys outside prefix: len={len} lo={lo:?} walked={walked:?}"
691                );
692                total += n;
693            }
694            assert_eq!(
695                total,
696                u64::cast_from(ids.len()),
697                "len={len} walked={walked:?}"
698            );
699        }
700
701        drop_db(&mut conn, DB).await?;
702        conn.disconnect().await?;
703        Ok(())
704    }
705
706    #[mz_ore::test(tokio::test)]
707    #[cfg_attr(miri, ignore)]
708    async fn test_live_mysql_collations() -> Result<(), anyhow::Error> {
709        let Some(mut conn) = connect().await? else {
710            return Ok(());
711        };
712        const CI_DB: &str = "mz_probe_collation_ci_test";
713        const BIN_DB: &str = "mz_probe_collation_bin_test";
714
715        // Case-insensitive collation: case variants of one key collide, so
716        // keys differ by letter, in mixed case.
717        let ci_keys = ["Apple", "apricot", "banana", "Cherry"];
718        let t_ci = setup_table(&mut conn, CI_DB, "utf8mb4_general_ci", &ci_keys).await?;
719        let mut tx = start_tx(&mut conn).await?;
720        let mut prober = KeyProber::new(&mut tx, t_ci, "id");
721        // 'A' covers 'apricot' too: LIKE is case-insensitive here, so a
722        // returned prefix covers every case variant of it.
723        assert_eq!(walk_prefixes(&mut prober, 1).await?, ["A", "b", "C"]);
724        tx.rollback().await?;
725
726        // Binary collation: case variants coexist and order by byte value.
727        let bin_keys = ["ABC", "ABD", "abc", "abd"];
728        let t_bin = setup_table(&mut conn, BIN_DB, "utf8mb4_bin", &bin_keys).await?;
729        let mut tx = start_tx(&mut conn).await?;
730        let mut prober = KeyProber::new(&mut tx, t_bin, "id");
731        // Uppercase sorts before lowercase in byte order, and case variants
732        // are distinct prefixes.
733        assert_eq!(walk_prefixes(&mut prober, 1).await?, ["A", "a"]);
734        assert_eq!(walk_prefixes(&mut prober, 3).await?, bin_keys);
735        tx.rollback().await?;
736
737        drop_db(&mut conn, CI_DB).await?;
738        drop_db(&mut conn, BIN_DB).await?;
739        conn.disconnect().await?;
740        Ok(())
741    }
742
743    #[mz_ore::test(tokio::test)]
744    #[cfg_attr(miri, ignore)]
745    async fn test_live_mysql_invalid_utf8_keys() -> Result<(), anyhow::Error> {
746        let Some(mut conn) = connect().await? else {
747            return Ok(());
748        };
749        const DB: &str = "mz_probe_binary_test";
750        recreate_db(&mut conn, DB).await?;
751        // A binary key column passes bytes through unconverted, so this is
752        // the one way invalid UTF-8 can reach the client. The snapshot
753        // operator will only support char and varchar columns, so this
754        // shouldn't happen in practice.
755        #[allow(clippy::disallowed_methods)]
756        conn.query_drop(format!(
757            "CREATE TABLE {DB}.t (id VARBINARY(36) PRIMARY KEY NOT NULL)"
758        ))
759        .await?;
760        let keys: Vec<Vec<u8>> = vec![b"a1".to_vec(), b"a2".to_vec(), vec![0xff, 0xfe, 0x31]];
761        conn.exec_batch(
762            format!("INSERT INTO {DB}.t VALUES (?)"),
763            keys.iter().map(|k| (Value::Bytes(k.clone()),)),
764        )
765        .await?;
766        #[allow(clippy::disallowed_methods)]
767        conn.query_drop(format!("ANALYZE TABLE {DB}.t")).await?;
768        let table = QualifiedTableRef {
769            schema_name: DB,
770            table_name: "t",
771        };
772        let mut tx = start_tx(&mut conn).await?;
773        let mut p = KeyProber::new(&mut tx, table, "id");
774
775        // Estimates never decode key values, they keep working.
776        assert!(p.estimate_range_rows("", None).await.is_ok());
777
778        // ASCII keys order before the 0xff key and decode fine.
779        assert_eq!(
780            prefix_of_first_key_in_range(&mut p, "", None, 2).await,
781            some("a1")
782        );
783        assert_eq!(
784            prefix_of_first_row_not_matching_prefix(&mut p, "a1", None, 2).await,
785            some("a2")
786        );
787        // The next key is invalid UTF-8. The probe reports it as a named
788        // error so callers can log it and fall back.
789        let err = p
790            .prefix_of_first_row_not_matching_prefix("a2", None, 2)
791            .await
792            .unwrap_err();
793        assert!(matches!(err, MySqlError::NonUtf8KeyValue { .. }), "{err:?}");
794
795        tx.rollback().await?;
796        drop_db(&mut conn, DB).await?;
797        conn.disconnect().await?;
798        Ok(())
799    }
800
801    #[mz_ore::test(tokio::test)]
802    #[cfg_attr(miri, ignore)]
803    async fn test_live_mysql_stale_statistics() -> Result<(), anyhow::Error> {
804        let Some(mut conn) = connect().await? else {
805            return Ok(());
806        };
807        const DB: &str = "mz_probe_stale_test";
808        recreate_db(&mut conn, DB).await?;
809        // This setup stays bespoke: STATS_AUTO_RECALC=0 plus an ANALYZE while
810        // empty pins the persisted statistics at zero rows, no matter what is
811        // inserted afterwards.
812        #[allow(clippy::disallowed_methods)]
813        {
814            conn.query_drop(format!(
815                "CREATE TABLE {DB}.t (id VARCHAR(36) CHARACTER SET utf8mb4 \
816                 COLLATE utf8mb4_bin PRIMARY KEY NOT NULL) \
817                 STATS_AUTO_RECALC=0, STATS_PERSISTENT=1"
818            ))
819            .await?;
820            conn.query_drop(format!("ANALYZE TABLE {DB}.t")).await?;
821        }
822        let ids: Vec<String> = (0..1000).map(|i| format!("a{i:05}")).collect();
823        conn.exec_batch(
824            format!("INSERT INTO {DB}.t VALUES (?)"),
825            ids.iter().map(|id| (id.as_str(),)),
826        )
827        .await?;
828
829        // The staleness this test is about: table_rows reports 0.
830        let table_rows: Option<u64> = conn
831            .exec_first(
832                "SELECT table_rows FROM information_schema.tables \
833                 WHERE table_schema = ? AND table_name = 't'",
834                (DB,),
835            )
836            .await?;
837        assert_eq!(table_rows, Some(0));
838
839        let table = QualifiedTableRef {
840            schema_name: DB,
841            table_name: "t",
842        };
843        let mut tx = start_tx(&mut conn).await?;
844        let mut prober = KeyProber::new(&mut tx, table, "id");
845
846        // Range estimates come from index dives on the real B-tree, not the
847        // stale table statistics, so they still reflect the actual data.
848        let all = prober.estimate_range_rows("", None).await?;
849        assert!((500..=2000).contains(&all), "all={all}");
850        let range = prober.estimate_range_rows("a00100", Some("a00200")).await?;
851        assert!((50..=200).contains(&range), "range={range}");
852
853        tx.rollback().await?;
854        drop_db(&mut conn, DB).await?;
855        conn.disconnect().await?;
856        Ok(())
857    }
858
859    #[mz_ore::test(tokio::test)]
860    #[cfg_attr(miri, ignore)]
861    async fn test_probe_sargability() -> Result<(), anyhow::Error> {
862        let Some(mut conn) = connect().await? else {
863            return Ok(());
864        };
865        const DB: &str = "mz_probe_sargable_test";
866        let ids: Vec<String> = (0..1000).map(|i| format!("a{i:05}")).collect();
867        let table = setup_table(&mut conn, DB, "utf8mb4_bin", &ids).await?;
868
869        let mut tx = start_tx(&mut conn).await?;
870
871        // Prove the methodology first: a deliberately non-sargable predicate
872        // reads every row, and the session handler counters see it.
873        let before = handler_reads(&mut tx).await?;
874        let _: Option<u64> = tx
875            .exec_first(
876                format!("SELECT COUNT(*) FROM {DB}.t WHERE LEFT(id, 2) = 'a0'"),
877                (),
878            )
879            .await?;
880        let scan_reads = handler_reads(&mut tx).await? - before;
881        assert!(scan_reads >= 1000, "scan_reads={scan_reads}");
882
883        // Every probe must stay a handful of index operations. A regression
884        // to a scan costs >= 1000 reads, far past the generous bound.
885        let before = handler_reads(&mut tx).await?;
886        let got = prefix_of_first_key_in_range(
887            &mut KeyProber::new(&mut tx, table.clone(), "id"),
888            "a00500",
889            None,
890            6,
891        )
892        .await;
893        let reads = handler_reads(&mut tx).await? - before;
894        // The exclusive bound skips the exact key a00500.
895        assert_eq!(got, some("a00501"));
896        assert!(reads < 50, "prefix_of_first_key_in_range reads={reads}");
897
898        let before = handler_reads(&mut tx).await?;
899        let got = prefix_of_first_row_not_matching_prefix(
900            &mut KeyProber::new(&mut tx, table.clone(), "id"),
901            "a00500",
902            None,
903            6,
904        )
905        .await;
906        let reads = handler_reads(&mut tx).await? - before;
907        assert_eq!(got, some("a00501"));
908        assert!(reads < 50, "max_key probe reads={reads}");
909
910        let before = handler_reads(&mut tx).await?;
911        let got = prefix_of_first_row_not_matching_prefix(
912            &mut KeyProber::new(&mut tx, table.clone(), "id"),
913            "a0",
914            None,
915            6,
916        )
917        .await;
918        let reads = handler_reads(&mut tx).await? - before;
919        // Every key matches 'a0%', so the prefix match covers the whole table and
920        // there is no next prefix, at the cost of two dives rather than a
921        // scan.
922        assert_eq!(got, None);
923        assert!(reads < 50, "whole-table match reads={reads}");
924
925        let before = handler_reads(&mut tx).await?;
926        let got = prefix_of_first_key_in_range(
927            &mut KeyProber::new(&mut tx, table.clone(), "id"),
928            "a00500",
929            Some("a00501"),
930            6,
931        )
932        .await;
933        let reads = handler_reads(&mut tx).await? - before;
934        assert_eq!(got, None);
935        assert!(reads < 50, "empty bounded range reads={reads}");
936
937        let bounded = KeyProber::new(&mut tx, table.clone(), "id")
938            .estimate_range_rows("a00100", Some("a00200"))
939            .await?;
940        assert!((50..=300).contains(&bounded), "bounded estimate={bounded}");
941
942        tx.rollback().await?;
943        drop_db(&mut conn, DB).await?;
944        conn.disconnect().await?;
945        Ok(())
946    }
947
948    /// `utf8mb4_bin` has no contractions or expansions: Czech `ch` stays an
949    /// ordinary `c` extension and `ß` an ordinary character, so the walk
950    /// visits every prefix. This would not work with the standard default collation.
951    #[mz_ore::test(tokio::test)]
952    #[cfg_attr(miri, ignore)]
953    async fn test_live_mysql_bin_no_contraction_or_expansion() -> Result<(), anyhow::Error> {
954        let Some(mut conn) = connect().await? else {
955            return Ok(());
956        };
957        const DB: &str = "mz_probe_bin_no_hazards";
958        let keys = [
959            "aaa", "asz", "aßx", "cesta", "chleba", "duha", "hora", "ibis",
960        ];
961        let table = setup_table(&mut conn, DB, "utf8mb4_bin", &keys).await?;
962
963        let mut tx = start_tx(&mut conn).await?;
964        let p = &mut KeyProber::new(&mut tx, table, "id");
965        assert_eq!(walk_prefixes(p, 1).await?, ["a", "c", "d", "h", "i"]);
966        assert_eq!(
967            walk_prefixes(p, 2).await?,
968            ["aa", "as", "aß", "ce", "ch", "du", "ho", "ib"]
969        );
970
971        tx.rollback().await?;
972        drop_db(&mut conn, DB).await?;
973        conn.disconnect().await?;
974        Ok(())
975    }
976
977    /// `utf8mb4_bin` compares character by character but is PAD SPACE, so
978    /// keys starting below space sort below the empty string. A walk seeded
979    /// with the empty string drops them, they land in the snapshot range
980    /// left of the first boundary. This means keys starting below space
981    /// will just be included in the first open range, which will be fine
982    /// for our partitioning, just a little unbalanced.
983    #[mz_ore::test(tokio::test)]
984    #[cfg_attr(miri, ignore)]
985    async fn test_live_mysql_keys_below_empty_string() -> Result<(), anyhow::Error> {
986        let Some(mut conn) = connect().await? else {
987            return Ok(());
988        };
989        const DB: &str = "mz_probe_below_empty_test";
990        let keys = ["\0a", "\u{1}a", "\u{9}b", "a1", "a1\u{1}x", "b1"];
991        let table = setup_table(&mut conn, DB, "utf8mb4_bin", &keys).await?;
992
993        let mut tx = start_tx(&mut conn).await?;
994        let p = &mut KeyProber::new(&mut tx, table, "id");
995        assert_eq!(
996            prefix_of_first_key_in_range(p, "", None, 2).await,
997            some("a1")
998        );
999        assert_eq!(
1000            prefix_of_first_row_not_matching_prefix(p, "a1", None, 2).await,
1001            some("b1")
1002        );
1003        assert_eq!(
1004            prefix_of_first_row_not_matching_prefix(p, "b1", None, 2).await,
1005            None
1006        );
1007
1008        assert_eq!(
1009            prefix_of_first_key_in_range(p, "", Some("a1"), 2).await,
1010            None
1011        );
1012        assert_eq!(
1013            prefix_of_first_row_not_matching_prefix(p, "\u{9}", Some("a1"), 2).await,
1014            None
1015        );
1016        assert_eq!(
1017            prefix_of_first_key_in_range(p, "", Some("b1"), 2).await,
1018            some("a1")
1019        );
1020
1021        tx.rollback().await?;
1022        drop_db(&mut conn, DB).await?;
1023        conn.disconnect().await?;
1024        Ok(())
1025    }
1026
1027    // Test helpers.
1028
1029    /// Connects to the server named by `MZ_TEST_MYSQL_URL`, or `None` to skip
1030    /// the test when it is unset. Skipping is a local-only convenience, CI
1031    /// must always provide the URL.
1032    pub(crate) async fn connect() -> Result<Option<mysql_async::Conn>, anyhow::Error> {
1033        let Ok(url) = std::env::var("MZ_TEST_MYSQL_URL") else {
1034            if mz_ore::env::is_var_truthy("CI") {
1035                panic!("CI is supposed to run this test but something has gone wrong!");
1036            }
1037            tracing::info!("MZ_TEST_MYSQL_URL not set: skipping live MySQL test");
1038            return Ok(None);
1039        };
1040        Ok(Some(
1041            mysql_async::Conn::new(mysql_async::Opts::from_url(&url)?).await?,
1042        ))
1043    }
1044
1045    /// Opens the transaction shape [`KeyProber`]'s contract expects:
1046    /// `REPEATABLE READ` and read-only.
1047    pub(crate) async fn start_tx(
1048        conn: &mut mysql_async::Conn,
1049    ) -> Result<Transaction<'_>, anyhow::Error> {
1050        let mut tx_opts = mysql_async::TxOpts::default();
1051        tx_opts
1052            .with_isolation_level(mysql_async::IsolationLevel::RepeatableRead)
1053            .with_readonly(true);
1054        Ok(conn.start_transaction(tx_opts).await?)
1055    }
1056
1057    /// Drops and recreates the scratch database `db`. Each test must use its
1058    /// own database name, tests on the shared server run concurrently.
1059    async fn recreate_db(conn: &mut mysql_async::Conn, db: &str) -> Result<(), anyhow::Error> {
1060        #[allow(clippy::disallowed_methods)]
1061        {
1062            conn.query_drop(format!("DROP DATABASE IF EXISTS {db}"))
1063                .await?;
1064            conn.query_drop(format!("CREATE DATABASE {db}")).await?;
1065        }
1066        Ok(())
1067    }
1068
1069    /// Recreates scratch database `db` holding one table `t` whose string
1070    /// primary key `id` is pinned to the given `collation`, containing
1071    /// `keys`, with fresh statistics. Returns a ref for [`KeyProber::new`].
1072    pub(crate) async fn setup_table<'a>(
1073        conn: &mut mysql_async::Conn,
1074        db: &'a str,
1075        collation: &str,
1076        keys: &[impl AsRef<str> + Sync],
1077    ) -> Result<QualifiedTableRef<'a>, anyhow::Error> {
1078        recreate_db(conn, db).await?;
1079        // MySQL collation names start with their character set's name, so
1080        // the charset is pinned explicitly without a second parameter.
1081        let charset = collation.split('_').next().expect("nonempty collation");
1082        #[allow(clippy::disallowed_methods)]
1083        conn.query_drop(format!(
1084            "CREATE TABLE {db}.t (id VARCHAR(36) CHARACTER SET {charset} \
1085             COLLATE {collation} PRIMARY KEY NOT NULL)"
1086        ))
1087        .await?;
1088
1089        for chunk in keys.chunks(1000) {
1090            conn.exec_drop(
1091                format!(
1092                    "INSERT INTO {db}.t VALUES {}",
1093                    vec!["(?)"; chunk.len()].join(",")
1094                ),
1095                chunk
1096                    .iter()
1097                    .map(|id| id.as_ref().into())
1098                    .collect::<Vec<mysql_async::Value>>(),
1099            )
1100            .await?;
1101        }
1102        #[allow(clippy::disallowed_methods)]
1103        conn.query_drop(format!("ANALYZE TABLE {db}.t")).await?;
1104        Ok(QualifiedTableRef {
1105            schema_name: db,
1106            table_name: "t",
1107        })
1108    }
1109
1110    /// Drops the scratch database `db`.
1111    pub(crate) async fn drop_db(
1112        conn: &mut mysql_async::Conn,
1113        db: &str,
1114    ) -> Result<(), anyhow::Error> {
1115        #[allow(clippy::disallowed_methods)]
1116        conn.query_drop(format!("DROP DATABASE {db}")).await?;
1117        Ok(())
1118    }
1119
1120    /// Keys in `[lo, hi)` of `db`'s table: the total, and how many have `lo`
1121    /// as a prefix, counted by the server so the comparisons happen under the
1122    /// column's collation.
1123    async fn count_range(
1124        conn: &mut mysql_async::Conn,
1125        db: &str,
1126        lo: &str,
1127        hi: Option<&String>,
1128    ) -> Result<(u64, u64), anyhow::Error> {
1129        let mut clause = "id >= ?".to_string();
1130        let mut params: Vec<Value> = vec![lo.into(), lo.into(), lo.into()];
1131        if let Some(hi) = hi {
1132            clause.push_str(" AND id < ?");
1133            params.push(hi.as_str().into());
1134        }
1135        let row: Option<(u64, Option<u64>)> = conn
1136            .exec_first(
1137                format!(
1138                    "SELECT COUNT(*), SUM(LEFT(id, CHAR_LENGTH(?)) = ?) FROM {db}.t WHERE {clause}"
1139                ),
1140                Params::Positional(params),
1141            )
1142            .await?;
1143        let (total, prefixed) = row.expect("COUNT returns a row");
1144        Ok((total, prefixed.unwrap_or(0)))
1145    }
1146
1147    /// Sum of this session's `Handler_read_*` counters: how many index or row
1148    /// read operations the connection has performed so far.
1149    async fn handler_reads(tx: &mut Transaction<'_>) -> Result<u64, anyhow::Error> {
1150        let rows: Vec<(String, String)> = tx
1151            .exec("SHOW SESSION STATUS LIKE 'Handler_read%'", ())
1152            .await?;
1153        Ok(rows.into_iter().map(|(_, v)| v.parse().unwrap_or(0)).sum())
1154    }
1155
1156    // Wrapped to limit boilerplate
1157    async fn prefix_of_first_key_in_range(
1158        prober: &mut KeyProber<'_, '_>,
1159        lower_bound_exclusive: &str,
1160        upper_bound_exclusive: Option<&str>,
1161        max_prefix_length: usize,
1162    ) -> Option<String> {
1163        prober
1164            .prefix_of_first_key_in_range(
1165                lower_bound_exclusive,
1166                upper_bound_exclusive,
1167                max_prefix_length,
1168            )
1169            .await
1170            .expect("prefix_of_first_key_in_range failed")
1171    }
1172
1173    // Wrapped to limit boilerplate
1174    async fn prefix_of_first_row_not_matching_prefix(
1175        prober: &mut KeyProber<'_, '_>,
1176        prefix: &str,
1177        upper_bound_exclusive: Option<&str>,
1178        max_prefix_length: usize,
1179    ) -> Option<String> {
1180        prober
1181            .prefix_of_first_row_not_matching_prefix(
1182                prefix,
1183                upper_bound_exclusive,
1184                max_prefix_length,
1185            )
1186            .await
1187            .expect("prefix_of_first_row_not_matching_prefix failed")
1188    }
1189
1190    /// `Some` for comparing against [`prefix_of_first_key_in_range`] and
1191    /// [`prefix_of_first_row_not_matching_prefix`] results without `as_deref` noise at
1192    /// every assertion.
1193    fn some(s: &str) -> Option<String> {
1194        Some(s.into())
1195    }
1196
1197    /// Test helper to walk prefixes at a consistent depth. Only works when
1198    /// all keys have length >= len.
1199    async fn walk_prefixes(
1200        prober: &mut KeyProber<'_, '_>,
1201        len: usize,
1202    ) -> Result<Vec<String>, anyhow::Error> {
1203        let mut walked = Vec::new();
1204        let Some(mut cur) = prober.prefix_of_first_key_in_range("", None, len).await? else {
1205            return Ok(walked);
1206        };
1207        loop {
1208            assert!(
1209                !walked.contains(&cur),
1210                "prefix repeated: {cur:?} (walked: {walked:?})"
1211            );
1212            walked.push(cur.clone());
1213            match prober
1214                .prefix_of_first_row_not_matching_prefix(&cur, None, len)
1215                .await?
1216            {
1217                Some(next) => cur = next,
1218                None => break,
1219            }
1220        }
1221        Ok(walked)
1222    }
1223}
mz_mysql_util/probe.rs

mz_mysql_util/
probe.rs
