Skip to content

Commit

Permalink
fix(spans): Scrub more SQL identifiers (#2817)
Browse files Browse the repository at this point in the history 
Some numeric identifiers were not yet scrubbed, for example `ALTER TABLE
tmp_123`.
  • Loading branch information
@jjbayer
jjbayer authored Dec 5, 2023
1 parent 011cf4d commit 1d7672d
Show file tree
Hide file tree
Showing 2 changed files with 172 additions and 27 deletions.
30 changes: 30 additions & 0 deletions relay-event-normalization/src/normalize/span/description/sql/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -690,6 +690,30 @@ mod tests {
"SELECT * FROM temp_{%s}c"
);

scrub_sql_test!(
digits_in_alter_table,
r#"ALTER TABLE "foo"."1234_foo" ADD COLUMN "bar678" TEXT"#,
r#"ALTER TABLE {%s}_foo ADD COLUMN bar{%s} TEXT"#
);

scrub_sql_test!(
digits_in_create_table,
r#"CREATE TABLE "foo"."1234_foo" ()"#,
"CREATE TABLE {%s}_foo ()"
);

scrub_sql_test!(
digits_in_update,
r#"UPDATE "foo"."1234_foo" SET "bar678" = 0"#,
"UPDATE {%s}_foo SET bar{%s} = %s"
);

scrub_sql_test!(
digits_in_delete,
r#"DELETE FROM foo123 WHERE bar456 = 789"#,
"DELETE FROM foo{%s} WHERE bar{%s} = %s"
);

scrub_sql_test!(
uuid_in_table_name,
"SELECT * FROM prefix_0123456789abcdef0123456789ABCDEF_006_suffix",
Expand All @@ -702,6 +726,12 @@ mod tests {
"SELECT id FROM {%s} LIMIT %s OFFSET %s"
);

scrub_sql_test!(
rename_table,
r#"ALTER TABLE "foo"."tmp" RENAME TO "foo"."bar"#,
"ALTER TABLE tmp RENAME TO bar"
);

scrub_sql_test!(
clickhouse,
"SELECT (toStartOfHour(finish_ts, 'Universal') AS _snuba_time), (uniqIf((nullIf(user, '') AS _snuba_user), greater(multiIf(equals(tupleElement(('duration', 300), 1), 'lcp'), (if(has(measurements.key, 'lcp'), arrayElement(measurements.value, indexOf(measurements.key, 'lcp')), NULL) AS `_snuba_measurements[lcp]`), (duration AS _snuba_duration)), multiply(tupleElement(('duration', 300), 2), 4))) AS _snuba_count_miserable_user), (ifNull(divide(plus(_snuba_count_miserable_user, 4.56), plus(nullIf(uniqIf(_snuba_user, greater(multiIf(equals(tupleElement(('duration', 300), 1), 'lcp'), `_snuba_measurements[lcp]`, _snuba_duration), 0)), 0), 113.45)), 0) AS _snuba_user_misery), _snuba_count_miserable_user, (divide(countIf(notEquals(transaction_status, 0) AND notEquals(transaction_status, 1) AND notEquals(transaction_status, 2)), count()) AS _snuba_failure_rate), (divide(count(), divide(3600.0, 60)) AS _snuba_tpm_3600) FROM transactions_dist WHERE equals(('transaction' AS _snuba_type), 'transaction') AND greaterOrEquals((finish_ts AS _snuba_finish_ts), toDateTime('2023-06-13T09:08:51', 'Universal')) AND less(_snuba_finish_ts, toDateTime('2023-07-11T09:08:51', 'Universal')) AND in((project_id AS _snuba_project_id), [123, 456, 789]) AND equals((environment AS _snuba_environment), 'production') GROUP BY _snuba_time ORDER BY _snuba_time ASC LIMIT 10000 OFFSET 0",
Expand Down
169 changes: 142 additions & 27 deletions relay-event-normalization/src/normalize/span/description/sql/parser.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,9 @@ use itertools::Itertools;
use once_cell::sync::Lazy;
use regex::Regex;
use sqlparser::ast::{
Assignment, CloseCursor, Expr, Ident, Query, Select, SelectItem, SetExpr, Statement,
TableFactor, TableWithJoins, UnaryOperator, Value, VisitMut, VisitorMut,
AlterTableOperation, Assignment, CloseCursor, ColumnDef, Expr, Ident, ObjectName, Query,
Select, SelectItem, SetExpr, Statement, TableAlias, TableConstraint, TableFactor,
UnaryOperator, Value, VisitMut, VisitorMut,
};
use sqlparser::dialect::{Dialect, GenericDialect};

Expand Down Expand Up @@ -110,7 +111,6 @@ impl NormalizeVisitor {
fn transform_query(&mut self, query: &mut Query) {
if let SetExpr::Select(select) = &mut *query.body {
self.transform_select(&mut *select);
self.transform_from(&mut select.from);
}
}

Expand Down Expand Up @@ -139,28 +139,15 @@ impl NormalizeVisitor {
Self::collapse_items(&mut collapse, &mut select.projection);
}

fn transform_from(&mut self, from: &mut [TableWithJoins]) {
// Iterate "FROM".
for from in from {
self.transform_table_factor(&mut from.relation);
for join in &mut from.joins {
self.transform_table_factor(&mut join.relation);
fn simplify_table_alias(alias: &mut Option<TableAlias>) {
if let Some(TableAlias { name, columns }) = alias {
Self::scrub_name(name);
for column in columns {
Self::scrub_name(column);
}
}
}

fn transform_table_factor(&mut self, table_factor: &mut TableFactor) {
match table_factor {
// Recurse into subqueries.
TableFactor::Derived { subquery, .. } => {
self.transform_query(subquery);
}
// Strip quotes from identifiers in table names.
TableFactor::Table { name, .. } => Self::simplify_compound_identifier(&mut name.0),
_ => {}
}
}

fn simplify_compound_identifier(parts: &mut Vec<Ident>) {
if let Some(mut last) = parts.pop() {
Self::scrub_name(&mut last);
Expand All @@ -184,13 +171,59 @@ impl NormalizeVisitor {
impl VisitorMut for NormalizeVisitor {
type Break = ();

fn pre_visit_relation(&mut self, relation: &mut ObjectName) -> ControlFlow<Self::Break> {
Self::simplify_compound_identifier(&mut relation.0);
ControlFlow::Continue(())
}

fn pre_visit_table_factor(
&mut self,
table_factor: &mut TableFactor,
) -> ControlFlow<Self::Break> {
match table_factor {
TableFactor::Table { alias, .. } => {
Self::simplify_table_alias(alias);
}
TableFactor::Derived {
subquery, alias, ..
} => {
self.transform_query(subquery);
Self::simplify_table_alias(alias);
}
TableFactor::TableFunction { alias, .. } => {
Self::simplify_table_alias(alias);
}
TableFactor::UNNEST {
alias,
with_offset_alias,
..
} => {
Self::simplify_table_alias(alias);
if let Some(ident) = with_offset_alias {
Self::scrub_name(ident);
}
}
TableFactor::NestedJoin { alias, .. } => {
Self::simplify_table_alias(alias);
}
TableFactor::Pivot {
table_alias,
pivot_alias,
..
} => {
Self::simplify_table_alias(table_alias);
Self::simplify_table_alias(pivot_alias);
}
}
ControlFlow::Continue(())
}

fn pre_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow<Self::Break> {
if self.current_expr_depth > MAX_EXPRESSION_DEPTH {
*expr = Expr::Value(Value::Placeholder("..".to_owned()));
return ControlFlow::Continue(());
}
self.current_expr_depth += 1;

match expr {
// Simple values like numbers and strings are replaced by a placeholder:
Expr::Value(x) => *x = Self::placeholder(),
Expand Down Expand Up @@ -252,12 +285,8 @@ impl VisitorMut for NormalizeVisitor {
}
// `INSERT INTO col1, col2 VALUES (val1, val2)` becomes `INSERT INTO .. VALUES (%s)`.
Statement::Insert {
table_name,
columns,
source,
..
columns, source, ..
} => {
Self::simplify_compound_identifier(&mut table_name.0);
*columns = vec![Self::ellipsis()];
if let SetExpr::Values(v) = &mut *source.body {
v.rows = vec![vec![Expr::Value(Self::placeholder())]]
Expand Down Expand Up @@ -298,6 +327,92 @@ impl VisitorMut for NormalizeVisitor {
Statement::Close {
cursor: CloseCursor::Specific { name },
} => Self::erase_name(name),
Statement::AlterTable { name, operation } => {
Self::simplify_compound_identifier(&mut name.0);
match operation {
AlterTableOperation::AddConstraint(c) => match c {
TableConstraint::Unique { name, columns, .. } => {
if let Some(name) = name {
Self::scrub_name(name);
}
for column in columns {
Self::scrub_name(column);
}
}
TableConstraint::ForeignKey {
name,
columns,
referred_columns,
..
} => {
if let Some(name) = name {
Self::scrub_name(name);
}
for column in columns {
Self::scrub_name(column);
}
for column in referred_columns {
Self::scrub_name(column);
}
}
TableConstraint::Check { name, .. } => {
if let Some(name) = name {
Self::scrub_name(name);
}
}
TableConstraint::Index { name, columns, .. } => {
if let Some(name) = name {
Self::scrub_name(name);
}
for column in columns {
Self::scrub_name(column);
}
}
TableConstraint::FulltextOrSpatial {
opt_index_name,
columns,
..
} => {
if let Some(name) = opt_index_name {
Self::scrub_name(name);
}
for column in columns {
Self::scrub_name(column);
}
}
},
AlterTableOperation::AddColumn { column_def, .. } => {
let ColumnDef { name, .. } = column_def;
Self::scrub_name(name);
}
AlterTableOperation::DropConstraint { name, .. } => Self::scrub_name(name),
AlterTableOperation::DropColumn { column_name, .. } => {
Self::scrub_name(column_name)
}
AlterTableOperation::RenameColumn {
old_column_name,
new_column_name,
} => {
Self::scrub_name(old_column_name);
Self::scrub_name(new_column_name);
}
AlterTableOperation::ChangeColumn {
old_name, new_name, ..
} => {
Self::scrub_name(old_name);
Self::scrub_name(new_name);
}
AlterTableOperation::RenameConstraint { old_name, new_name } => {
Self::scrub_name(old_name);
Self::scrub_name(new_name);
}
AlterTableOperation::AlterColumn { column_name, .. } => {
Self::scrub_name(column_name);
}
_ => {}
}
}

_ => {}
}

Expand Down

0 comments on commit 1d7672d

Please sign in to comment.