Skip to content

Commit

Permalink
Fix missing db.operation for CREATE/DROP/ALTER SQL statement (#10020)
Browse files Browse the repository at this point in the history
  • Loading branch information
bjrara committed Mar 12, 2024
1 parent 0029332 commit df83eaa
Show file tree
Hide file tree
Showing 7 changed files with 204 additions and 70 deletions.
93 changes: 91 additions & 2 deletions instrumentation-api-incubator/src/main/jflex/SqlSanitizer.jflex
Original file line number Diff line number Diff line change
Expand Up @@ -128,13 +128,56 @@ WHITESPACE = [ \t\r\n]+
/** @return true if all statement info is gathered */
boolean handleNext() {
return false;
}
}

/** @return true if all statement info is gathered */
boolean handleOperationTarget(String target) {
return false;
}

boolean expectingOperationTarget() {
return false;
}

SqlStatementInfo getResult(String fullStatement) {
return SqlStatementInfo.create(fullStatement, getClass().getSimpleName().toUpperCase(java.util.Locale.ROOT), mainIdentifier);
}
}

private abstract class DdlOperation extends Operation {
private String operationTarget = "";
private boolean expectingOperationTarget = true;

boolean expectingOperationTarget() {
return expectingOperationTarget;
}

boolean handleOperationTarget(String target) {
operationTarget = target;
expectingOperationTarget = false;
return false;
}

boolean shouldHandleIdentifier() {
// Return true only if the provided value corresponds to a table, as it will be used to set the attribute `db.sql.table`.
return "TABLE".equals(operationTarget);
}

boolean handleIdentifier() {
if (shouldHandleIdentifier()) {
mainIdentifier = readIdentifierName();
}
return true;
}

SqlStatementInfo getResult(String fullStatement) {
if (!"".equals(operationTarget)) {
return SqlStatementInfo.create(fullStatement, getClass().getSimpleName().toUpperCase(java.util.Locale.ROOT) + " " + operationTarget, mainIdentifier);
}
return super.getResult(fullStatement);
}
}

private static class NoOp extends Operation {
static final Operation INSTANCE = new NoOp();

Expand Down Expand Up @@ -279,6 +322,15 @@ WHITESPACE = [ \t\r\n]+
}
}

private class Create extends DdlOperation {
}

private class Drop extends DdlOperation {
}

private class Alter extends DdlOperation {
}

private SqlStatementInfo getResult() {
if (builder.length() > LIMIT) {
builder.delete(LIMIT, builder.length());
Expand Down Expand Up @@ -339,6 +391,27 @@ WHITESPACE = [ \t\r\n]+
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
"CREATE" {
if (!insideComment) {
setOperation(new Create());
}
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
"DROP" {
if (!insideComment) {
setOperation(new Drop());
}
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
"ALTER" {
if (!insideComment) {
setOperation(new Alter());
}
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
"FROM" {
if (!insideComment && !extractionDone) {
if (operation == NoOp.INSTANCE) {
Expand Down Expand Up @@ -367,11 +440,27 @@ WHITESPACE = [ \t\r\n]+
}
"NEXT" {
if (!insideComment && !extractionDone) {
extractionDone = operation.handleNext();
extractionDone = operation.handleNext();
}
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
"IF" | "NOT" | "EXISTS" {
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
"TABLE" | "INDEX" | "DATABASE" | "PROCEDURE" | "VIEW" {
if (!insideComment && !extractionDone) {
if (operation.expectingOperationTarget()) {
extractionDone = operation.handleOperationTarget(yytext());
} else {
extractionDone = operation.handleIdentifier();
}
}
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}

{COMMA} {
if (!insideComment && !extractionDone) {
extractionDone = operation.handleComma();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,17 @@ void veryLongSelectStatementsAreOk() {
assertThat(result).isEqualTo(expected);
}

@ParameterizedTest
@ArgumentsSource(DdlArgs.class)
void checkDdlOperationStatementsAreOk(
String actual, Function<String, SqlStatementInfo> expectFunc) {
SqlStatementInfo result = SqlStatementSanitizer.create(true).sanitize(actual);
SqlStatementInfo expected = expectFunc.apply(actual);
assertThat(result.getFullStatement()).isEqualTo(expected.getFullStatement());
assertThat(result.getOperation()).isEqualTo(expected.getOperation());
assertThat(result.getMainIdentifier()).isEqualTo(expected.getMainIdentifier());
}

@Test
void lotsOfTicksDontCauseStackOverflowOrLongRuntimes() {
String s = "'";
Expand Down Expand Up @@ -357,4 +368,34 @@ public Stream<? extends Arguments> provideArguments(ExtensionContext context) th
Arguments.of(null, expect(null, null)));
}
}

static class DdlArgs implements ArgumentsProvider {

static Function<String, SqlStatementInfo> expect(String operation, String identifier) {
return sql -> SqlStatementInfo.create(sql, operation, identifier);
}

static Function<String, SqlStatementInfo> expect(
String sql, String operation, String identifier) {
return ignored -> SqlStatementInfo.create(sql, operation, identifier);
}

@Override
public Stream<? extends Arguments> provideArguments(ExtensionContext context) throws Exception {
return Stream.of(
Arguments.of("CREATE TABLE `table`", expect("CREATE TABLE", "table")),
Arguments.of("CREATE TABLE IF NOT EXISTS table", expect("CREATE TABLE", "table")),
Arguments.of("DROP TABLE `if`", expect("DROP TABLE", "if")),
Arguments.of(
"ALTER TABLE table ADD CONSTRAINT c FOREIGN KEY (foreign_id) REFERENCES ref (id)",
expect("ALTER TABLE", "table")),
Arguments.of("CREATE INDEX types_name ON types (name)", expect("CREATE INDEX", null)),
Arguments.of("DROP INDEX types_name ON types (name)", expect("DROP INDEX", null)),
Arguments.of(
"CREATE VIEW tmp AS SELECT type FROM table WHERE id = ?",
expect("CREATE VIEW", null)),
Arguments.of(
"CREATE PROCEDURE p AS SELECT * FROM table GO", expect("CREATE PROCEDURE", null)));
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -225,8 +225,8 @@ private static Stream<Arguments> provideSyncParameters() {
null,
"DROP KEYSPACE IF EXISTS sync_test",
"DROP KEYSPACE IF EXISTS sync_test",
"DB Query",
null,
"DROP",
"DROP",
null))),
Arguments.of(
named(
Expand All @@ -235,8 +235,8 @@ private static Stream<Arguments> provideSyncParameters() {
null,
"CREATE KEYSPACE sync_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}",
"CREATE KEYSPACE sync_test WITH REPLICATION = {?:?, ?:?}",
"DB Query",
null,
"CREATE",
"CREATE",
null))),
Arguments.of(
named(
Expand All @@ -245,9 +245,9 @@ private static Stream<Arguments> provideSyncParameters() {
"sync_test",
"CREATE TABLE sync_test.users ( id UUID PRIMARY KEY, name text )",
"CREATE TABLE sync_test.users ( id UUID PRIMARY KEY, name text )",
"sync_test",
null,
null))),
"CREATE TABLE sync_test.users",
"CREATE TABLE",
"sync_test.users"))),
Arguments.of(
named(
"Insert data",
Expand Down Expand Up @@ -279,8 +279,8 @@ private static Stream<Arguments> provideAsyncParameters() {
null,
"DROP KEYSPACE IF EXISTS async_test",
"DROP KEYSPACE IF EXISTS async_test",
"DB Query",
null,
"DROP",
"DROP",
null))),
Arguments.of(
named(
Expand All @@ -289,8 +289,8 @@ private static Stream<Arguments> provideAsyncParameters() {
null,
"CREATE KEYSPACE async_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}",
"CREATE KEYSPACE async_test WITH REPLICATION = {?:?, ?:?}",
"DB Query",
null,
"CREATE",
"CREATE",
null))),
Arguments.of(
named(
Expand All @@ -299,9 +299,9 @@ private static Stream<Arguments> provideAsyncParameters() {
"async_test",
"CREATE TABLE async_test.users ( id UUID PRIMARY KEY, name text )",
"CREATE TABLE async_test.users ( id UUID PRIMARY KEY, name text )",
"async_test",
null,
null))),
"CREATE TABLE async_test.users",
"CREATE TABLE",
"async_test.users"))),
Arguments.of(
named(
"Insert data",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -200,8 +200,8 @@ private static Stream<Arguments> provideSyncParameters() {
null,
"DROP KEYSPACE IF EXISTS sync_test",
"DROP KEYSPACE IF EXISTS sync_test",
"DB Query",
null,
"DROP",
"DROP",
null))),
Arguments.of(
named(
Expand All @@ -210,8 +210,8 @@ private static Stream<Arguments> provideSyncParameters() {
null,
"CREATE KEYSPACE sync_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}",
"CREATE KEYSPACE sync_test WITH REPLICATION = {?:?, ?:?}",
"DB Query",
null,
"CREATE",
"CREATE",
null))),
Arguments.of(
named(
Expand All @@ -220,9 +220,9 @@ private static Stream<Arguments> provideSyncParameters() {
"sync_test",
"CREATE TABLE sync_test.users ( id UUID PRIMARY KEY, name text )",
"CREATE TABLE sync_test.users ( id UUID PRIMARY KEY, name text )",
"sync_test",
null,
null))),
"CREATE TABLE sync_test.users",
"CREATE TABLE",
"sync_test.users"))),
Arguments.of(
named(
"Insert data",
Expand Down Expand Up @@ -254,8 +254,8 @@ private static Stream<Arguments> provideAsyncParameters() {
null,
"DROP KEYSPACE IF EXISTS async_test",
"DROP KEYSPACE IF EXISTS async_test",
"DB Query",
null,
"DROP",
"DROP",
null))),
Arguments.of(
named(
Expand All @@ -264,8 +264,8 @@ private static Stream<Arguments> provideAsyncParameters() {
null,
"CREATE KEYSPACE async_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}",
"CREATE KEYSPACE async_test WITH REPLICATION = {?:?, ?:?}",
"DB Query",
null,
"CREATE",
"CREATE",
null))),
Arguments.of(
named(
Expand All @@ -274,9 +274,9 @@ private static Stream<Arguments> provideAsyncParameters() {
"async_test",
"CREATE TABLE async_test.users ( id UUID PRIMARY KEY, name text )",
"CREATE TABLE async_test.users ( id UUID PRIMARY KEY, name text )",
"async_test",
null,
null))),
"CREATE TABLE async_test.users",
"CREATE TABLE",
"async_test.users"))),
Arguments.of(
named(
"Insert data",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -106,8 +106,8 @@ private static Stream<Arguments> provideReactiveParameters() {
null,
"DROP KEYSPACE IF EXISTS reactive_test",
"DROP KEYSPACE IF EXISTS reactive_test",
"DB Query",
null,
"DROP",
"DROP",
null))),
Arguments.of(
named(
Expand All @@ -116,8 +116,8 @@ private static Stream<Arguments> provideReactiveParameters() {
null,
"CREATE KEYSPACE reactive_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}",
"CREATE KEYSPACE reactive_test WITH REPLICATION = {?:?, ?:?}",
"DB Query",
null,
"CREATE",
"CREATE",
null))),
Arguments.of(
named(
Expand All @@ -126,9 +126,9 @@ private static Stream<Arguments> provideReactiveParameters() {
"reactive_test",
"CREATE TABLE reactive_test.users ( id UUID PRIMARY KEY, name text )",
"CREATE TABLE reactive_test.users ( id UUID PRIMARY KEY, name text )",
"reactive_test",
null,
null))),
"CREATE TABLE reactive_test.users",
"CREATE TABLE",
"reactive_test.users"))),
Arguments.of(
named(
"Insert data",
Expand Down
Loading

0 comments on commit df83eaa

Please sign in to comment.