Sql sanitizer: sanitize double quoted strings only in couchbase queries

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/db/SqlDialect.java

@@ -0,0 +1,15 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.db;
+
+/**
+ * Enumeration of sql dialects that are handled differently by {@link SqlStatementSanitizer}.
+ */
+public enum SqlDialect {
+  DEFAULT,
+  // couchbase uses double quotes for string literals
+  COUCHBASE
+}

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/db/SqlStatementSanitizer.java

@@ -10,6 +10,7 @@
 
 import io.opentelemetry.instrumentation.api.caching.Cache;
 import io.opentelemetry.instrumentation.api.internal.SupportabilityMetrics;
+import java.util.Objects;
 import javax.annotation.Nullable;
 
 /**
@@ -19,20 +20,56 @@
 public final class SqlStatementSanitizer {
   private static final SupportabilityMetrics supportability = SupportabilityMetrics.instance();
 
-  private static final Cache<String, SqlStatementInfo> sqlToStatementInfoCache =
+  private static final Cache<StatementKey, SqlStatementInfo> sqlToStatementInfoCache =
       Cache.builder().setMaximumSize(1000).build();
 
   public static SqlStatementInfo sanitize(@Nullable String statement) {
+    return sanitize(statement, SqlDialect.DEFAULT);
+  }
+
+  public static SqlStatementInfo sanitize(@Nullable String statement, SqlDialect dialect) {
     if (!isStatementSanitizationEnabled() || statement == null) {
       return SqlStatementInfo.create(statement, null, null);
     }
     return sqlToStatementInfoCache.computeIfAbsent(
-        statement,
+        new StatementKey(statement, dialect),
         k -> {
           supportability.incrementCounter(SQL_STATEMENT_SANITIZER_CACHE_MISS);
-          return AutoSqlSanitizer.sanitize(statement);
+          return AutoSqlSanitizer.sanitize(statement, dialect);
         });
   }
 
+  private static final class StatementKey {
+    private final String statement;
+    private final SqlDialect dialect;
+
+    StatementKey(String statement, SqlDialect dialect) {
+      this.statement = statement;
+      this.dialect = dialect;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (o == null || getClass() != o.getClass()) {
+        return false;
+      }
+      StatementKey that = (StatementKey) o;
+      return statement.equals(that.statement) && dialect == that.dialect;
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(statement, dialect);
+    }
+
+    @Override
+    public String toString() {
+      return "StatementKey(statement=" + statement + ", dialect=" + dialect + ")";
+    }
+  }
+
   private SqlStatementSanitizer() {}
 }

instrumentation-api/src/main/jflex/SqlSanitizer.jflex

@@ -30,8 +30,9 @@ DOLLAR_QUOTED_STR = "$$" [^$]* "$$"
 WHITESPACE        = [ \t\r\n]+
 
 %{
-  static SqlStatementInfo sanitize(String statement) {
+  static SqlStatementInfo sanitize(String statement, SqlDialect dialect) {
     AutoSqlSanitizer sanitizer = new AutoSqlSanitizer(new java.io.StringReader(statement));
+    sanitizer.dialect = dialect;
     try {
       while (!sanitizer.yyatEOF()) {
         int token = sanitizer.yylex();
@@ -71,6 +72,7 @@ WHITESPACE        = [ \t\r\n]+
   private boolean insideComment = false;
   private Operation operation = NoOp.INSTANCE;
   private boolean extractionDone = false;
+  private SqlDialect dialect;
 
   private void setOperation(Operation operation) {
     if (this.operation == NoOp.INSTANCE) {
@@ -361,11 +363,20 @@ WHITESPACE        = [ \t\r\n]+
       }
 
   // here is where the actual sanitization happens
-  {BASIC_NUM} | {HEX_NUM} | {QUOTED_STR} | {DOUBLE_QUOTED_STR} | {DOLLAR_QUOTED_STR} {
+  {BASIC_NUM} | {HEX_NUM} | {QUOTED_STR} | {DOLLAR_QUOTED_STR} {
           builder.append('?');
           if (isOverLimit()) return YYEOF;
       }
 
+  {DOUBLE_QUOTED_STR} {
+          if (dialect == SqlDialect.COUCHBASE) {
+            builder.append('?');
+          } else {
+            appendCurrentFragment();
+          }
+          if (isOverLimit()) return YYEOF;
+      }
+
   {WHITESPACE} {
           builder.append(' ');
           if (isOverLimit()) return YYEOF;

instrumentation-api/src/test/groovy/io/opentelemetry/instrumentation/api/db/SqlStatementSanitizerTest.groovy

@@ -37,6 +37,7 @@ class SqlStatementSanitizerTest extends Specification {
     "+7"                                                                       | "?"
     "SELECT 0x0af764"                                                          | "SELECT ?"
     "SELECT 0xdeadBEEF"                                                        | "SELECT ?"
+    "SELECT * FROM \"TABLE\""                                                  | "SELECT * FROM \"TABLE\""
 
     // Not numbers but could be confused as such
     "SELECT A + B"                                                             | "SELECT A + B"
@@ -59,15 +60,6 @@ class SqlStatementSanitizerTest extends Specification {
     "SELECT * FROM TABLE WHERE FIELD = '\"\$\$\$\$\"'"                         | "SELECT * FROM TABLE WHERE FIELD = ?"
     "SELECT * FROM TABLE WHERE FIELD = 'a single \" doublequote inside'"       | "SELECT * FROM TABLE WHERE FIELD = ?"
 
-    // Some databases support/encourage " instead of ' with same escape rules
-    "SELECT * FROM TABLE WHERE FIELD = \"\""                                   | "SELECT * FROM TABLE WHERE FIELD = ?"
-    "SELECT * FROM TABLE WHERE FIELD = \"words and spaces'\""                  | "SELECT * FROM TABLE WHERE FIELD = ?"
-    "SELECT * FROM TABLE WHERE FIELD = \" an escaped \"\" quote mark inside\"" | "SELECT * FROM TABLE WHERE FIELD = ?"
-    "SELECT * FROM TABLE WHERE FIELD = \"\\\\\""                               | "SELECT * FROM TABLE WHERE FIELD = ?"
-    "SELECT * FROM TABLE WHERE FIELD = \"'inside singles'\""                   | "SELECT * FROM TABLE WHERE FIELD = ?"
-    "SELECT * FROM TABLE WHERE FIELD = \"'\$\$\$\$'\""                         | "SELECT * FROM TABLE WHERE FIELD = ?"
-    "SELECT * FROM TABLE WHERE FIELD = \"a single ' singlequote inside\""      | "SELECT * FROM TABLE WHERE FIELD = ?"
-
     // Some databases allow using dollar-quoted strings
     "SELECT * FROM TABLE WHERE FIELD = \$\$\$\$"                               | "SELECT * FROM TABLE WHERE FIELD = ?"
     "SELECT * FROM TABLE WHERE FIELD = \$\$words and spaces\$\$"               | "SELECT * FROM TABLE WHERE FIELD = ?"
@@ -85,6 +77,25 @@ class SqlStatementSanitizerTest extends Specification {
     "FROM TABLE WHERE FIELD=1234"                                              | "FROM TABLE WHERE FIELD=?"
   }
 
+  def "normalize couchbase #originalSql"() {
+    setup:
+    def actualSanitized = SqlStatementSanitizer.sanitize(originalSql, SqlDialect.COUCHBASE)
+
+    expect:
+    actualSanitized.getFullStatement() == sanitizedSql
+
+    where:
+    originalSql                                                                | sanitizedSql
+    // Some databases support/encourage " instead of ' with same escape rules
+    "SELECT * FROM TABLE WHERE FIELD = \"\""                                   | "SELECT * FROM TABLE WHERE FIELD = ?"
+    "SELECT * FROM TABLE WHERE FIELD = \"words and spaces'\""                  | "SELECT * FROM TABLE WHERE FIELD = ?"
+    "SELECT * FROM TABLE WHERE FIELD = \" an escaped \"\" quote mark inside\"" | "SELECT * FROM TABLE WHERE FIELD = ?"
+    "SELECT * FROM TABLE WHERE FIELD = \"\\\\\""                               | "SELECT * FROM TABLE WHERE FIELD = ?"
+    "SELECT * FROM TABLE WHERE FIELD = \"'inside singles'\""                   | "SELECT * FROM TABLE WHERE FIELD = ?"
+    "SELECT * FROM TABLE WHERE FIELD = \"'\$\$\$\$'\""                         | "SELECT * FROM TABLE WHERE FIELD = ?"
+    "SELECT * FROM TABLE WHERE FIELD = \"a single ' singlequote inside\""      | "SELECT * FROM TABLE WHERE FIELD = ?"
+  }
+
   @Unroll
   def "should simplify #sql"() {
     expect:

instrumentation/couchbase/couchbase-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/couchbase/v2_0/CouchbaseQuerySanitizer.java

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.couchbase.v2_0;
 
+import io.opentelemetry.instrumentation.api.db.SqlDialect;
 import io.opentelemetry.instrumentation.api.db.SqlStatementInfo;
 import io.opentelemetry.instrumentation.api.db.SqlStatementSanitizer;
 import java.lang.invoke.MethodHandle;
@@ -115,7 +116,7 @@ private static String getStatementString(MethodHandle handle, Object query) {
   }
 
   private static SqlStatementInfo sanitizeString(String query) {
-    return SqlStatementSanitizer.sanitize(query);
+    return SqlStatementSanitizer.sanitize(query, SqlDialect.COUCHBASE);
   }
 
   private CouchbaseQuerySanitizer() {}