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

mz_mysql_util/
probe.rs
