Skip to content

Commit

Permalink
nom-sql: Parse PostgreSQL-style comments
Browse files Browse the repository at this point in the history
Our benchmarking code relies upon table and column comments in order to
capture information about how to generate test data. Previously, we only
supported MySQL-style comments, so this commit adds parsing support for
PostgreSQL-style table and column comments.

Change-Id: I1d972d691ef13b205ec7a27636b11012a90d24f1
Reviewed-on: https://gerrit.readyset.name/c/readyset/+/5133
Tested-by: Buildkite CI
Reviewed-by: Griffin Smith <griffin@readyset.io>
Reviewed-by: Luke Osborne <luke.o@readyset.io>
  • Loading branch information
ethan-readyset committed Jun 23, 2023
1 parent cba739a commit a2e9302
Show file tree
Hide file tree
Showing 8 changed files with 293 additions and 21 deletions.
45 changes: 37 additions & 8 deletions nom-sql/src/analysis/visit.rs
Expand Up @@ -17,14 +17,14 @@ use crate::set::Variable;
use crate::transaction::{CommitStatement, RollbackStatement, StartTransactionStatement};
use crate::{
AlterColumnOperation, AlterTableDefinition, AlterTableStatement, CacheInner, CaseWhenBranch,
Column, ColumnConstraint, ColumnSpecification, CommonTableExpr, CompoundSelectStatement,
CreateCacheStatement, CreateTableStatement, CreateViewStatement, DeleteStatement,
DropAllCachesStatement, DropCacheStatement, DropTableStatement, DropViewStatement,
ExplainStatement, Expr, FieldDefinitionExpr, FieldReference, FunctionExpr, GroupByClause,
InValue, InsertStatement, JoinClause, JoinConstraint, JoinRightSide, Literal, OrderClause,
Relation, SelectSpecification, SelectStatement, SetNames, SetPostgresParameter, SetStatement,
SetVariables, ShowStatement, SqlIdentifier, SqlQuery, SqlType, TableExpr, TableExprInner,
TableKey, UpdateStatement, UseStatement,
Column, ColumnConstraint, ColumnSpecification, CommentStatement, CommonTableExpr,
CompoundSelectStatement, CreateCacheStatement, CreateTableStatement, CreateViewStatement,
DeleteStatement, DropAllCachesStatement, DropCacheStatement, DropTableStatement,
DropViewStatement, ExplainStatement, Expr, FieldDefinitionExpr, FieldReference, FunctionExpr,
GroupByClause, InValue, InsertStatement, JoinClause, JoinConstraint, JoinRightSide, Literal,
OrderClause, Relation, SelectSpecification, SelectStatement, SetNames, SetPostgresParameter,
SetStatement, SetVariables, ShowStatement, SqlIdentifier, SqlQuery, SqlType, TableExpr,
TableExprInner, TableKey, UpdateStatement, UseStatement,
};

/// Each method of the `Visitor` trait is a hook to be potentially overridden when recursively
Expand Down Expand Up @@ -396,6 +396,13 @@ pub trait Visitor<'ast>: Sized {
Ok(())
}

fn visit_comment_statement(
&mut self,
comment_statement: &'ast CommentStatement,
) -> Result<(), Self::Error> {
walk_comment_statement(self, comment_statement)
}

fn visit_sql_query(&mut self, sql_query: &'ast SqlQuery) -> Result<(), Self::Error> {
walk_sql_query(self, sql_query)
}
Expand Down Expand Up @@ -1111,7 +1118,29 @@ pub fn walk_sql_query<'a, V: Visitor<'a>>(
SqlQuery::Use(statement) => visitor.visit_use_statement(statement),
SqlQuery::Show(statement) => visitor.visit_show_statement(statement),
SqlQuery::Explain(statement) => visitor.visit_explain_statement(statement),
SqlQuery::Comment(statement) => visitor.visit_comment_statement(statement),
}
}

pub fn walk_comment_statement<'a, V: Visitor<'a>>(
visitor: &mut V,
comment_statement: &'a CommentStatement,
) -> Result<(), V::Error> {
match comment_statement {
CommentStatement::Column {
column_name,
table_name,
..
} => {
visitor.visit_sql_identifier(column_name)?;
visitor.visit_sql_identifier(table_name)?;
}
CommentStatement::Table { table_name, .. } => {
visitor.visit_sql_identifier(table_name)?;
}
}

Ok(())
}

#[cfg(test)]
Expand Down
45 changes: 37 additions & 8 deletions nom-sql/src/analysis/visit_mut.rs
Expand Up @@ -17,14 +17,14 @@ use crate::set::Variable;
use crate::transaction::{CommitStatement, RollbackStatement, StartTransactionStatement};
use crate::{
AlterColumnOperation, AlterTableDefinition, AlterTableStatement, CacheInner, CaseWhenBranch,
Column, ColumnConstraint, ColumnSpecification, CommonTableExpr, CompoundSelectStatement,
CreateCacheStatement, CreateTableStatement, CreateViewStatement, DeleteStatement,
DropAllCachesStatement, DropCacheStatement, DropTableStatement, DropViewStatement,
ExplainStatement, Expr, FieldDefinitionExpr, FieldReference, FunctionExpr, GroupByClause,
InValue, InsertStatement, JoinClause, JoinConstraint, JoinRightSide, Literal, OrderClause,
Relation, SelectSpecification, SelectStatement, SetNames, SetPostgresParameter, SetStatement,
SetVariables, ShowStatement, SqlIdentifier, SqlQuery, SqlType, TableExpr, TableExprInner,
TableKey, UpdateStatement, UseStatement,
Column, ColumnConstraint, ColumnSpecification, CommentStatement, CommonTableExpr,
CompoundSelectStatement, CreateCacheStatement, CreateTableStatement, CreateViewStatement,
DeleteStatement, DropAllCachesStatement, DropCacheStatement, DropTableStatement,
DropViewStatement, ExplainStatement, Expr, FieldDefinitionExpr, FieldReference, FunctionExpr,
GroupByClause, InValue, InsertStatement, JoinClause, JoinConstraint, JoinRightSide, Literal,
OrderClause, Relation, SelectSpecification, SelectStatement, SetNames, SetPostgresParameter,
SetStatement, SetVariables, ShowStatement, SqlIdentifier, SqlQuery, SqlType, TableExpr,
TableExprInner, TableKey, UpdateStatement, UseStatement,
};

/// Each method of the `VisitorMut` trait is a hook to be potentially overridden when recursively
Expand Down Expand Up @@ -411,6 +411,13 @@ pub trait VisitorMut<'ast>: Sized {
Ok(())
}

fn visit_comment_statement(
&mut self,
comment_statement: &'ast mut CommentStatement,
) -> Result<(), Self::Error> {
walk_comment_statement(self, comment_statement)
}

fn visit_sql_query(&mut self, sql_query: &'ast mut SqlQuery) -> Result<(), Self::Error> {
walk_sql_query(self, sql_query)
}
Expand Down Expand Up @@ -1128,7 +1135,29 @@ pub fn walk_sql_query<'a, V: VisitorMut<'a>>(
SqlQuery::Use(statement) => visitor.visit_use_statement(statement),
SqlQuery::Show(statement) => visitor.visit_show_statement(statement),
SqlQuery::Explain(statement) => visitor.visit_explain_statement(statement),
SqlQuery::Comment(statement) => visitor.visit_comment_statement(statement),
}
}

pub fn walk_comment_statement<'a, V: VisitorMut<'a>>(
visitor: &mut V,
comment_statement: &'a mut CommentStatement,
) -> Result<(), V::Error> {
match comment_statement {
CommentStatement::Column {
column_name,
table_name,
..
} => {
visitor.visit_sql_identifier(column_name)?;
visitor.visit_sql_identifier(table_name)?;
}
CommentStatement::Table { table_name, .. } => {
visitor.visit_sql_identifier(table_name)?;
}
}

Ok(())
}

#[cfg(test)]
Expand Down
187 changes: 187 additions & 0 deletions nom-sql/src/comment.rs
@@ -0,0 +1,187 @@
use std::{fmt, str};

use nom::branch::alt;
use nom::bytes::complete::{tag, tag_no_case};
use nom::combinator::map_res;
use nom::sequence::{delimited, terminated};
use nom_locate::LocatedSpan;
use readyset_util::fmt::fmt_with;
use serde::{Deserialize, Serialize};

use crate::common::statement_terminator;
use crate::whitespace::{whitespace0, whitespace1};
use crate::{literal, Dialect, NomSqlResult, SqlIdentifier};

#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub enum CommentStatement {
Column {
column_name: SqlIdentifier,
table_name: SqlIdentifier,
comment: String,
},
Table {
table_name: SqlIdentifier,
comment: String,
},
}

impl CommentStatement {
pub fn display(&self, dialect: Dialect) -> impl fmt::Display + Copy + '_ {
fmt_with(move |f| match self {
Self::Column {
column_name,
table_name,
comment,
} => {
write!(
f,
"COMMENT ON COLUMN {}.{} IS ",
dialect.quote_identifier(&table_name),
dialect.quote_identifier(&column_name),
)?;
literal::display_string_literal(f, comment)
}
Self::Table {
table_name,
comment,
} => {
write!(
f,
"COMMENT ON TABLE {} IS ",
dialect.quote_identifier(&table_name),
)?;
literal::display_string_literal(f, comment)
}
})
}
}

pub fn comment(
dialect: Dialect,
) -> impl Fn(LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], CommentStatement> {
move |i| alt((table_comment(dialect), column_comment(dialect)))(i)
}

fn table_comment(
dialect: Dialect,
) -> impl Fn(LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], CommentStatement> {
move |i| {
let (i, _) = tag_no_case("comment")(i)?;
let (i, _) = whitespace1(i)?;
let (i, _) = tag_no_case("on")(i)?;
let (i, _) = whitespace1(i)?;
let (i, _) = tag_no_case("table")(i)?;
let (i, _) = whitespace1(i)?;
let (i, table_name) = dialect.identifier()(i)?;
let (i, _) = whitespace1(i)?;
let (i, _) = tag_no_case("is")(i)?;
let (i, _) = whitespace1(i)?;
let (i, comment) = map_res(dialect.string_literal(), String::from_utf8)(i)?;
let (i, _) = statement_terminator(i)?;

Ok((
i,
CommentStatement::Table {
table_name,
comment,
},
))
}
}

fn column_comment(
dialect: Dialect,
) -> impl Fn(LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], CommentStatement> {
move |i| {
let (i, _) = tag_no_case("comment")(i)?;
let (i, _) = whitespace1(i)?;
let (i, _) = tag_no_case("on")(i)?;
let (i, _) = whitespace1(i)?;
let (i, _) = tag_no_case("column")(i)?;
let (i, _) = whitespace1(i)?;
let (i, table_name) = terminated(
dialect.identifier(),
delimited(whitespace0, tag("."), whitespace0),
)(i)?;
let (i, column_name) = dialect.identifier()(i)?;
let (i, _) = whitespace1(i)?;
let (i, _) = tag_no_case("is")(i)?;
let (i, _) = whitespace1(i)?;
let (i, comment) = map_res(dialect.string_literal(), String::from_utf8)(i)?;
let (i, _) = statement_terminator(i)?;

Ok((
i,
CommentStatement::Column {
table_name,
column_name,
comment,
},
))
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn table_comment() {
let res = test_parse!(
comment(Dialect::MySQL),
b"COMMENT ON TABLE test IS 'this is a comment'"
);

assert_eq!(
res,
CommentStatement::Table {
table_name: "test".into(),
comment: "this is a comment".into(),
}
)
}

#[test]
fn column_comment() {
let res = test_parse!(
comment(Dialect::MySQL),
b"COMMENT ON COLUMN test_table.test_column IS 'this is a comment'"
);

assert_eq!(
res,
CommentStatement::Column {
table_name: "test_table".into(),
column_name: "test_column".into(),
comment: "this is a comment".into(),
}
)
}

#[test]
fn display_column_comment() {
let comment = CommentStatement::Column {
table_name: "test_table".into(),
column_name: "test_column".into(),
comment: "this is a comment".into(),
};

assert_eq!(
comment.display(Dialect::PostgreSQL).to_string(),
"COMMENT ON COLUMN \"test_table\".\"test_column\" IS 'this is a comment'"
);
}

#[test]
fn display_table_comment() {
let comment = CommentStatement::Table {
table_name: "test_table".into(),
comment: "this is a comment".into(),
};

assert_eq!(
comment.display(Dialect::PostgreSQL).to_string(),
"COMMENT ON TABLE \"test_table\" IS 'this is a comment'"
);
}
}
2 changes: 2 additions & 0 deletions nom-sql/src/lib.rs
Expand Up @@ -16,6 +16,7 @@ pub use self::alter::{
AlterColumnOperation, AlterTableDefinition, AlterTableStatement, ReplicaIdentity,
};
pub use self::column::{Column, ColumnConstraint, ColumnSpecification};
pub use self::comment::CommentStatement;
pub use self::common::{FieldDefinitionExpr, FieldReference, IndexType, TableKey};
pub use self::compound_select::{CompoundSelectOperator, CompoundSelectStatement};
pub use self::create::{
Expand Down Expand Up @@ -63,6 +64,7 @@ mod macros;
mod alter;
pub mod analysis;
mod column;
mod comment;
mod common;
mod compound_select;
mod create;
Expand Down
8 changes: 5 additions & 3 deletions nom-sql/src/literal.rs
Expand Up @@ -195,9 +195,7 @@ impl Display for Literal {
Literal::Numeric(val, scale) => {
write!(f, "{}", Decimal::from_i128_with_scale(*val, *scale))
}
Literal::String(ref s) => {
write!(f, "'{}'", s.replace('\'', "''").replace('\\', "\\\\"))
}
Literal::String(ref s) => display_string_literal(f, s),
Literal::Blob(ref bv) => write!(
f,
"{}",
Expand All @@ -224,6 +222,10 @@ impl Display for Literal {
}
}

pub(crate) fn display_string_literal(f: &mut fmt::Formatter<'_>, s: &str) -> fmt::Result {
write!(f, "'{}'", s.replace('\'', "''").replace('\\', "\\\\"))
}

impl Literal {
pub fn arbitrary_with_type(sql_type: &SqlType) -> impl Strategy<Value = Self> + 'static {
use proptest::prelude::*;
Expand Down

0 comments on commit a2e9302

Please sign in to comment.