Skip to content

Commit

Permalink
fix missing db.operation for create/drop/alter statement
Browse files Browse the repository at this point in the history
  • Loading branch information
bjrara committed Feb 27, 2024
1 parent b74ac2a commit 1500ced
Show file tree
Hide file tree
Showing 7 changed files with 249 additions and 70 deletions.
137 changes: 135 additions & 2 deletions instrumentation-api-incubator/src/main/jflex/SqlSanitizer.jflex
Original file line number Diff line number Diff line change
Expand Up @@ -122,13 +122,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 @@ -273,6 +316,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 @@ -329,6 +381,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 @@ -357,11 +430,71 @@ 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" {
if (!insideComment && !extractionDone) {
if (operation.expectingOperationTarget()) {
extractionDone = operation.handleOperationTarget("TABLE");
} else {
extractionDone = operation.handleIdentifier();
}
}
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
"INDEX" {
if (!insideComment && !extractionDone) {
if (operation.expectingOperationTarget()) {
extractionDone = operation.handleOperationTarget("INDEX");
} else {
extractionDone = operation.handleIdentifier();
}
}
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
"DATABASE" {
if (!insideComment && !extractionDone) {
if (operation.expectingOperationTarget()) {
extractionDone = operation.handleOperationTarget("DATABASE");
} else {
extractionDone = operation.handleIdentifier();
}
}
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
"PROCEDURE" {
if (!insideComment && !extractionDone) {
if (operation.expectingOperationTarget()) {
extractionDone = operation.handleOperationTarget("PROCEDURE");
} else {
extractionDone = operation.handleIdentifier();
}
}
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
"VIEW" {
if (!insideComment && !extractionDone) {
if (operation.expectingOperationTarget()) {
extractionDone = operation.handleOperationTarget("VIEW");
} 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 @@ -331,4 +342,35 @@ 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(
// Select
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
Loading

0 comments on commit 1500ced

Please sign in to comment.