Skip to content

Commit

Permalink
nom-sql: Add support for index hints
Browse files Browse the repository at this point in the history
This commit adds support for index hints to the SQL parser. Index hints
are a MySQL-specific feature that allows the user to specify which index
to use or ignore for a query.
Readyset doesn't require those, but we should be able to parse and
treat any query if a variant of index hint as the same query as if
it didn't have any index hint.

Refs: REA-4320

Change-Id: I25938a3a549a28fc768a8b09f3d68a0327e1eb11
Reviewed-on: https://gerrit.readyset.name/c/readyset/+/7274
Tested-by: Buildkite CI
Reviewed-by: Luke Osborne <luke@readyset.io>
  • Loading branch information
altmannmarcelo committed Apr 11, 2024
1 parent 201c029 commit c89c902
Show file tree
Hide file tree
Showing 11 changed files with 402 additions and 9 deletions.
1 change: 1 addition & 0 deletions benchmarks/src/write_benchmark.rs
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,7 @@ fn query_indexed_by_columns(
tables: vec![TableExpr {
inner: TableExprInner::Table(table.name.clone().into()),
alias: None,
index_hint: None,
}],
fields: vec![FieldDefinitionExpr::All],
where_clause: cols
Expand Down
248 changes: 248 additions & 0 deletions nom-sql/src/index_hint.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,248 @@
use std::fmt::Display;

use itertools::Itertools;
use nom::branch::alt;
use nom::bytes::complete::{tag, tag_no_case};
use nom::combinator::{map, opt, value};
use nom::multi::separated_list1;
use nom::sequence::preceded;
use nom_locate::LocatedSpan;
use readyset_util::fmt::fmt_with;
use serde::{Deserialize, Serialize};
use test_strategy::Arbitrary;

use crate::common::ws_sep_comma;
use crate::whitespace::{whitespace0, whitespace1};
use crate::{Dialect, DialectDisplay, NomSqlResult, SqlIdentifier};

/// Type of index hint.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Arbitrary)]
pub enum IndexHintType {
/// Use index hint.
Use,

/// Ignore index hint.
Ignore,

/// Force index hint.
Force,
}

impl From<&str> for IndexHintType {
fn from(s: &str) -> Self {
match s {
"USE" => IndexHintType::Use,
"IGNORE" => IndexHintType::Ignore,
"FORCE" => IndexHintType::Force,
_ => panic!("Invalid index hint type: {}", s),
}
}
}

impl From<&&str> for IndexHintType {
fn from(s: &&str) -> Self {
IndexHintType::from(*s)
}
}

impl From<IndexHintType> for &str {
fn from(t: IndexHintType) -> &'static str {
match t {
IndexHintType::Use => "USE",
IndexHintType::Ignore => "IGNORE",
IndexHintType::Force => "FORCE",
}
}
}

impl<'a> From<LocatedSpan<&'a [u8]>> for IndexHintType {
fn from(span: LocatedSpan<&'a [u8]>) -> Self {
let s = span.fragment();
let str_slice = std::str::from_utf8(s).expect("Invalid UTF-8 string");
IndexHintType::from(str_slice)
}
}

/// Type of index or key.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Arbitrary)]
pub enum IndexOrKeyType {
/// Index.
Index,

/// Key.
Key,
}

impl From<&str> for IndexOrKeyType {
fn from(s: &str) -> Self {
match s {
"INDEX" => IndexOrKeyType::Index,
"KEY" => IndexOrKeyType::Key,
_ => panic!("Invalid index or key type: {}", s),
}
}
}

impl From<&&str> for IndexOrKeyType {
fn from(s: &&str) -> Self {
IndexOrKeyType::from(*s)
}
}

impl From<IndexOrKeyType> for &str {
fn from(t: IndexOrKeyType) -> &'static str {
match t {
IndexOrKeyType::Index => "INDEX",
IndexOrKeyType::Key => "KEY",
}
}
}

/// Index usage type.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Arbitrary)]
pub enum IndexUsageType {
/// FOR JOIN
Join,

/// FOR ORDER BY
OrderBy,

/// FOR GROUP BY
GroupBy,
}

impl From<&str> for IndexUsageType {
fn from(s: &str) -> Self {
match s {
"FOR JOIN" => IndexUsageType::Join,
"FOR ORDER BY" => IndexUsageType::OrderBy,
"FOR GROUP BY" => IndexUsageType::GroupBy,
_ => panic!("Invalid index usage type: {}", s),
}
}
}

impl From<&&str> for IndexUsageType {
fn from(s: &&str) -> Self {
IndexUsageType::from(*s)
}
}

impl From<IndexUsageType> for &str {
fn from(t: IndexUsageType) -> &'static str {
match t {
IndexUsageType::Join => " FOR JOIN",
IndexUsageType::OrderBy => " FOR ORDER BY",
IndexUsageType::GroupBy => " FOR GROUP BY",
}
}
}
/// Index hints for a query.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Arbitrary)]
pub struct IndexHint {
/// Type of index hint.
pub hint_type: IndexHintType,

/// Type of index or key.
pub index_or_key: IndexOrKeyType,

/// Index usage type.
pub index_usage_type: Option<IndexUsageType>,

/// List of indexes.
pub index_list: Vec<SqlIdentifier>,
}

impl IndexHint {
/// Create a new index hint.
pub fn default() -> Self {
Self {
hint_type: IndexHintType::Use,
index_or_key: IndexOrKeyType::Index,
index_usage_type: None,
index_list: vec![],
}
}
}

impl DialectDisplay for IndexHint {
fn display(&self, _dialect: Dialect) -> impl Display + '_ {
fmt_with(move |f| {
let hint_type: &str = self.hint_type.clone().into();
let index_or_key: &str = self.index_or_key.clone().into();
let index_usage: &str = self
.index_usage_type
.as_ref()
.map(|t| t.clone().into())
.unwrap_or("");
let index_list = self.index_list.iter().map(|t| t.to_owned()).join(", ");
write!(
f,
"{} {}{} ({})",
hint_type, index_or_key, index_usage, index_list
)
})
}
}

/// Parse an identifier or the PRIMARY keyword.
/// This is used to parse the index list.
/// PRIMARY is a reserved KEYWORD, and identifier() does not accept reserved keywords.
fn identifier_or_primary(
dialect: Dialect,
) -> impl Fn(LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], SqlIdentifier> {
move |i| {
let (i, _) = whitespace0(i)?;
let (i, identifier) = alt((
map(dialect.identifier(), |ident| ident),
value(SqlIdentifier::from("PRIMARY"), tag_no_case("PRIMARY")),
))(i)?;
let (i, _) = whitespace0(i)?;

Ok((i, identifier))
}
}

/// Parse an index hint.
pub fn index_hint_list(
dialect: Dialect,
) -> impl Fn(LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], IndexHint> {
move |i| {
let (i, _) = whitespace1(i)?;
let (i, hint_type) = alt((
value(IndexHintType::Use, tag_no_case("USE")),
value(IndexHintType::Ignore, tag_no_case("IGNORE")),
value(IndexHintType::Force, tag_no_case("FORCE")),
))(i)?;
let (i, _) = whitespace1(i)?;
let (i, index_or_key) = alt((
value(IndexOrKeyType::Index, tag_no_case("INDEX")),
value(IndexOrKeyType::Key, tag_no_case("KEY")),
))(i)?;
let (i, index_usage_type) = opt(preceded(
whitespace1,
alt((
value(IndexUsageType::Join, tag_no_case("FOR JOIN")),
value(IndexUsageType::OrderBy, tag_no_case("FOR ORDER BY")),
value(IndexUsageType::GroupBy, tag_no_case("FOR GROUP BY")),
)),
))(i)?;
let (i, _) = whitespace0(i)?;
let (i, _) = tag("(")(i)?;
let (i, _) = whitespace0(i)?;
let (i, index_list) = separated_list1(ws_sep_comma, identifier_or_primary(dialect))(i)?;

let (i, _) = whitespace0(i)?;
let (i, _) = tag(")")(i)?;

Ok((
i,
IndexHint {
hint_type,
index_or_key,
index_usage_type,
index_list,
},
))
}
}
8 changes: 5 additions & 3 deletions nom-sql/src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,7 @@ fn keyword_e_to_i(i: LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], &[u8]> {
terminated(tag_no_case("EXPLAIN"), keyword_follow_char),
terminated(tag_no_case("FAIL"), keyword_follow_char),
terminated(tag_no_case("FOR"), keyword_follow_char),
terminated(tag_no_case("FORCE"), keyword_follow_char),
terminated(tag_no_case("FOREIGN"), keyword_follow_char),
terminated(tag_no_case("FROM"), keyword_follow_char),
terminated(tag_no_case("FULL"), keyword_follow_char),
Expand All @@ -112,7 +113,6 @@ fn keyword_e_to_i(i: LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], &[u8]> {
terminated(tag_no_case("INITIALLY"), keyword_follow_char),
terminated(tag_no_case("INNER"), keyword_follow_char),
terminated(tag_no_case("INSTEAD"), keyword_follow_char),
terminated(tag_no_case("INTERSECT"), keyword_follow_char),
)),
|i| *i,
)(i)
Expand All @@ -121,6 +121,7 @@ fn keyword_e_to_i(i: LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], &[u8]> {
fn keyword_i_to_p(i: LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], &[u8]> {
map(
alt((
terminated(tag_no_case("INTERSECT"), keyword_follow_char),
terminated(tag_no_case("INTO"), keyword_follow_char),
terminated(tag_no_case("IS"), keyword_follow_char),
terminated(tag_no_case("JOIN"), keyword_follow_char),
Expand All @@ -141,7 +142,6 @@ fn keyword_i_to_p(i: LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], &[u8]> {
terminated(tag_no_case("ORDER"), keyword_follow_char),
terminated(tag_no_case("OUTER"), keyword_follow_char),
terminated(tag_no_case("PLAN"), keyword_follow_char),
terminated(tag_no_case("PRAGMA"), keyword_follow_char),
)),
|i| *i,
)(i)
Expand All @@ -150,6 +150,7 @@ fn keyword_i_to_p(i: LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], &[u8]> {
fn keyword_p_to_t(i: LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], &[u8]> {
map(
alt((
terminated(tag_no_case("PRAGMA"), keyword_follow_char),
terminated(tag_no_case("PRIMARY"), keyword_follow_char),
terminated(tag_no_case("QUERY"), keyword_follow_char),
terminated(tag_no_case("RAISE"), keyword_follow_char),
Expand All @@ -170,7 +171,6 @@ fn keyword_p_to_t(i: LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], &[u8]> {
terminated(tag_no_case("TEMP"), keyword_follow_char),
terminated(tag_no_case("TEMPORARY"), keyword_follow_char),
terminated(tag_no_case("THEN"), keyword_follow_char),
terminated(tag_no_case("TO"), keyword_follow_char),
)),
|i| *i,
)(i)
Expand All @@ -179,11 +179,13 @@ fn keyword_p_to_t(i: LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], &[u8]> {
fn keyword_t_to_z(i: LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], &[u8]> {
map(
alt((
terminated(tag_no_case("TO"), keyword_follow_char),
terminated(tag_no_case("TRANSACTION"), keyword_follow_char),
terminated(tag_no_case("TRIGGER"), keyword_follow_char),
terminated(tag_no_case("UNION"), keyword_follow_char),
terminated(tag_no_case("UNIQUE"), keyword_follow_char),
terminated(tag_no_case("UPDATE"), keyword_follow_char),
terminated(tag_no_case("USE"), keyword_follow_char),
terminated(tag_no_case("USING"), keyword_follow_char),
terminated(tag_no_case("VACUUM"), keyword_follow_char),
terminated(tag_no_case("VIEW"), keyword_follow_char),
Expand Down
1 change: 1 addition & 0 deletions nom-sql/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,7 @@ mod dialect;
mod drop;
mod explain;
mod expression;
mod index_hint;
mod insert;
mod join;
mod keywords;
Expand Down
Loading

0 comments on commit c89c902

Please sign in to comment.