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 Dec 8, 2023
1 parent 4d7fc64 commit 8f87387
Show file tree
Hide file tree
Showing 7 changed files with 219 additions and 69 deletions.
118 changes: 117 additions & 1 deletion instrumentation-api-incubator/src/main/jflex/SqlSanitizer.jflex
Original file line number Diff line number Diff line change
Expand Up @@ -122,13 +122,48 @@ 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 waitingOperationTarget() {
return false;
}

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

private static abstract class CompositeOp extends Operation {
private String operationTarget = "";
private boolean waitingOperationTarget = true;

boolean waitingOperationTarget() {
return waitingOperationTarget;
}

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

boolean shouldHandleIdentifier() {
return "TABLE".equals(operationTarget);
}

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 +308,33 @@ WHITESPACE = [ \t\r\n]+
}
}

private class Create extends CompositeOp {
boolean handleIdentifier() {
if (shouldHandleIdentifier()) {
mainIdentifier = readIdentifierName();
}
return true;
}
}

private class Drop extends CompositeOp {
boolean handleIdentifier() {
if (shouldHandleIdentifier()) {
mainIdentifier = readIdentifierName();
}
return true;
}
}

private class Alter extends CompositeOp {
boolean handleIdentifier() {
if (shouldHandleIdentifier()) {
mainIdentifier = readIdentifierName();
}
return true;
}
}

private SqlStatementInfo getResult() {
if (builder.length() > LIMIT) {
builder.delete(LIMIT, builder.length());
Expand Down Expand Up @@ -329,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 @@ -358,7 +441,40 @@ WHITESPACE = [ \t\r\n]+
"NEXT" {
if (!insideComment && !extractionDone) {
extractionDone = operation.handleNext();
}
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
"TABLE" {
if (!insideComment && !extractionDone) {
if (operation.waitingOperationTarget()) {
extractionDone = operation.handleOperationTarget("TABLE");
} else {
extractionDone = operation.handleIdentifier();
}
}
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
"INDEX" {
if (!insideComment && !extractionDone) {
if (operation.waitingOperationTarget()) {
extractionDone = operation.handleOperationTarget("INDEX");
} else {
extractionDone = operation.handleIdentifier();
}
}
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
"DATABASE" {
if (!insideComment && !extractionDone) {
if (operation.waitingOperationTarget()) {
extractionDone = operation.handleOperationTarget("DATABASE");
} else {
extractionDone = operation.handleIdentifier();
}
}
appendCurrentFragment();
if (isOverLimit()) return YYEOF;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,36 @@ void veryLongSelectStatementsAreOk() {
assertThat(result).isEqualTo(expected);
}

@Test
void changeSchemaStatementsAreOk() {
String createTableStatement = "CREATE TABLE `table`";
SqlStatementInfo expected =
SqlStatementInfo.create(createTableStatement, "CREATE TABLE", "table");
SqlStatementInfo result = SqlStatementSanitizer.create(true).sanitize(createTableStatement);
assertThat(result).isEqualTo(expected);

String dropTableStatement = "DROP TABLE anotherTable";
expected = SqlStatementInfo.create(dropTableStatement, "DROP TABLE", "anotherTable");
result = SqlStatementSanitizer.create(true).sanitize(dropTableStatement);
assertThat(result).isEqualTo(expected);

String alterStatement =
"ALTER TABLE table ADD CONSTRAINT c FOREIGN KEY (foreign_id) REFERENCES ref (id)";
expected = SqlStatementInfo.create(alterStatement, "ALTER TABLE", "table");
result = SqlStatementSanitizer.create(true).sanitize(alterStatement);
assertThat(result).isEqualTo(expected);

String createStatement = "CREATE INDEX types_name ON types (name)";
expected = SqlStatementInfo.create(createStatement, "CREATE INDEX", null);
result = SqlStatementSanitizer.create(true).sanitize(createStatement);
assertThat(result).isEqualTo(expected);

String dropStatement = "DROP INDEX types_name ON types (name)";
expected = SqlStatementInfo.create(dropStatement, "DROP INDEX", null);
result = SqlStatementSanitizer.create(true).sanitize(dropStatement);
assertThat(result).isEqualTo(expected);
}

@Test
void lotsOfTicksDontCauseStackOverflowOrLongRuntimes() {
String s = "'";
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -213,8 +213,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 @@ -223,8 +223,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 @@ -233,9 +233,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 @@ -267,8 +267,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 @@ -277,8 +277,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 @@ -287,9 +287,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 @@ -172,8 +172,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 @@ -182,8 +182,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 @@ -192,9 +192,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 @@ -226,8 +226,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 @@ -236,8 +236,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 @@ -246,9 +246,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 @@ -91,8 +91,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 @@ -101,8 +101,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 @@ -111,9 +111,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 8f87387

Please sign in to comment.