Druid Avatica - Handle escaping of search characters correctly (#10040)

Fix Avatica based metadata queries by appending ESCAPE '\' clause to the LIKE expressions

sql/src/main/java/org/apache/druid/sql/avatica/DruidMeta.java

@@ -364,7 +364,7 @@ public class DruidMeta extends MetaImpl
     }
 
     if (schemaPattern.s != null) {
-      whereBuilder.add("SCHEMATA.SCHEMA_NAME LIKE " + Calcites.escapeStringLiteral(schemaPattern.s));
+      whereBuilder.add("SCHEMATA.SCHEMA_NAME LIKE " + withEscapeClause(schemaPattern.s));
     }
 
     final String where = whereBuilder.isEmpty() ? "" : "WHERE " + Joiner.on(" AND ").join(whereBuilder);
@@ -395,11 +395,11 @@ public class DruidMeta extends MetaImpl
     }
 
     if (schemaPattern.s != null) {
-      whereBuilder.add("TABLES.TABLE_SCHEMA LIKE " + Calcites.escapeStringLiteral(schemaPattern.s));
+      whereBuilder.add("TABLES.TABLE_SCHEMA LIKE " + withEscapeClause(schemaPattern.s));
     }
 
     if (tableNamePattern.s != null) {
-      whereBuilder.add("TABLES.TABLE_NAME LIKE " + Calcites.escapeStringLiteral(tableNamePattern.s));
+      whereBuilder.add("TABLES.TABLE_NAME LIKE " + withEscapeClause(tableNamePattern.s));
     }
 
     if (typeList != null) {
@@ -446,15 +446,16 @@ public class DruidMeta extends MetaImpl
     }
 
     if (schemaPattern.s != null) {
-      whereBuilder.add("COLUMNS.TABLE_SCHEMA LIKE " + Calcites.escapeStringLiteral(schemaPattern.s));
+      whereBuilder.add("COLUMNS.TABLE_SCHEMA LIKE " + withEscapeClause(schemaPattern.s));
     }
 
     if (tableNamePattern.s != null) {
-      whereBuilder.add("COLUMNS.TABLE_NAME LIKE " + Calcites.escapeStringLiteral(tableNamePattern.s));
+      whereBuilder.add("COLUMNS.TABLE_NAME LIKE " + withEscapeClause(tableNamePattern.s));
     }
 
     if (columnNamePattern.s != null) {
-      whereBuilder.add("COLUMNS.COLUMN_NAME LIKE " + Calcites.escapeStringLiteral(columnNamePattern.s));
+      whereBuilder.add("COLUMNS.COLUMN_NAME LIKE "
+                       + withEscapeClause(columnNamePattern.s));
     }
 
     final String where = whereBuilder.isEmpty() ? "" : "WHERE " + Joiner.on(" AND ").join(whereBuilder);
@@ -638,4 +639,9 @@ public class DruidMeta extends MetaImpl
     }
     return Math.min(clientMaxRowsPerFrame, config.getMaxRowsPerFrame());
   }
+
+  private static String withEscapeClause(String toEscape)
+  {
+    return Calcites.escapeStringLiteral(toEscape) + " ESCAPE '\\'";
+  }
 }

sql/src/test/java/org/apache/druid/sql/avatica/DruidAvaticaHandlerTest.java

@@ -414,8 +414,19 @@ public class DruidAvaticaHandlerTest extends CalciteTestBase
                 Pair.of("TABLE_NAME", CalciteTests.DATASOURCE3),
                 Pair.of("TABLE_SCHEM", "druid"),
                 Pair.of("TABLE_TYPE", "TABLE")
+            ),
+            row(
+                Pair.of("TABLE_CAT", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE),
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_TYPE", "TABLE")
+            ),
+            row(
+                Pair.of("TABLE_CAT", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOMEXDATASOURCE),
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_TYPE", "TABLE")
             )
-
         ),
         getRows(
             metaData.getTables(null, "druid", "%", null),
@@ -465,8 +476,19 @@ public class DruidAvaticaHandlerTest extends CalciteTestBase
                 Pair.of("TABLE_NAME", CalciteTests.DATASOURCE3),
                 Pair.of("TABLE_SCHEM", "druid"),
                 Pair.of("TABLE_TYPE", "TABLE")
+            ),
+            row(
+                Pair.of("TABLE_CAT", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE),
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_TYPE", "TABLE")
+            ),
+            row(
+                Pair.of("TABLE_CAT", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOMEXDATASOURCE),
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_TYPE", "TABLE")
             )
-
         ),
         getRows(
             metaData.getTables(null, "druid", "%", null),
@@ -957,6 +979,210 @@ public class DruidAvaticaHandlerTest extends CalciteTestBase
     );
   }
 
+  @Test
+  public void testEscapingForGetColumns() throws Exception
+  {
+    final DatabaseMetaData metaData = client.getMetaData();
+
+    ImmutableList<Map<String, Object>> someDatasourceColumns = ImmutableList.of(
+        row(
+            Pair.of("TABLE_SCHEM", "druid"),
+            Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE),
+            Pair.of("COLUMN_NAME", "__time")
+        ),
+        row(
+            Pair.of("TABLE_SCHEM", "druid"),
+            Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE),
+            Pair.of("COLUMN_NAME", "cnt")
+        ),
+        row(
+            Pair.of("TABLE_SCHEM", "druid"),
+            Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE),
+            Pair.of("COLUMN_NAME", "dim1")
+        ),
+        row(
+            Pair.of("TABLE_SCHEM", "druid"),
+            Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE),
+            Pair.of("COLUMN_NAME", "dim2")
+        ),
+        row(
+            Pair.of("TABLE_SCHEM", "druid"),
+            Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE),
+            Pair.of("COLUMN_NAME", "dim3")
+        ),
+        row(
+            Pair.of("TABLE_SCHEM", "druid"),
+            Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE),
+            Pair.of("COLUMN_NAME", "m1")
+        ),
+        row(
+            Pair.of("TABLE_SCHEM", "druid"),
+            Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE),
+            Pair.of("COLUMN_NAME", "m2")
+        ),
+        row(
+            Pair.of("TABLE_SCHEM", "druid"),
+            Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE),
+            Pair.of("COLUMN_NAME", "unique_dim1")
+        )
+    );
+    // If the escape clause wasn't correctly set, rows for potentially none or more than
+    // one datasource (some_datasource and somexdatasource) would have been returned
+    Assert.assertEquals(
+        someDatasourceColumns,
+        getRows(
+            metaData.getColumns(null, "dr_id", CalciteTests.SOME_DATSOURCE_ESCAPED, null),
+            ImmutableSet.of("TABLE_NAME", "TABLE_SCHEM", "COLUMN_NAME")
+        )
+    );
+
+    ImmutableList<Map<String, Object>> someXDatasourceColumns = ImmutableList.of(
+        row(
+            Pair.of("TABLE_SCHEM", "druid"),
+            Pair.of("TABLE_NAME", CalciteTests.SOMEXDATASOURCE),
+            Pair.of("COLUMN_NAME", "__time")
+        ),
+        row(
+            Pair.of("TABLE_SCHEM", "druid"),
+            Pair.of("TABLE_NAME", CalciteTests.SOMEXDATASOURCE),
+            Pair.of("COLUMN_NAME", "cnt_x")
+        ),
+        row(
+            Pair.of("TABLE_SCHEM", "druid"),
+            Pair.of("TABLE_NAME", CalciteTests.SOMEXDATASOURCE),
+            Pair.of("COLUMN_NAME", "m1_x")
+        ),
+        row(
+            Pair.of("TABLE_SCHEM", "druid"),
+            Pair.of("TABLE_NAME", CalciteTests.SOMEXDATASOURCE),
+            Pair.of("COLUMN_NAME", "m2_x")
+        ),
+        row(
+            Pair.of("TABLE_SCHEM", "druid"),
+            Pair.of("TABLE_NAME", CalciteTests.SOMEXDATASOURCE),
+            Pair.of("COLUMN_NAME", "unique_dim1_x")
+        )
+    );
+    Assert.assertEquals(
+        someXDatasourceColumns,
+        getRows(
+            metaData.getColumns(null, "dr_id", "somexdatasource", null),
+            ImmutableSet.of("TABLE_NAME", "TABLE_SCHEM", "COLUMN_NAME")
+        )
+    );
+
+    List<Map<String, Object>> columnsOfBothTables = new ArrayList<>(someDatasourceColumns);
+    columnsOfBothTables.addAll(someXDatasourceColumns);
+    // Assert that the pattern matching still works when no escape string is provided
+    Assert.assertEquals(
+        columnsOfBothTables,
+        getRows(
+            metaData.getColumns(null, "dr_id", "some_datasource", null),
+            ImmutableSet.of("TABLE_NAME", "TABLE_SCHEM", "COLUMN_NAME")
+        )
+    );
+
+    // Assert column name pattern works correctly when _ is in the column names
+    Assert.assertEquals(
+        ImmutableList.of(
+            row(
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOMEXDATASOURCE),
+                Pair.of("COLUMN_NAME", "m1_x")
+            ),
+            row(
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOMEXDATASOURCE),
+                Pair.of("COLUMN_NAME", "m2_x")
+            )
+        ),
+        getRows(
+            metaData.getColumns("druid", "dr_id", CalciteTests.SOMEXDATASOURCE, "m_\\_x"),
+            ImmutableSet.of("TABLE_NAME", "TABLE_SCHEM", "COLUMN_NAME")
+        )
+    );
+
+    // Assert column name pattern with % works correctly for column names starting with m
+    Assert.assertEquals(
+        ImmutableList.of(
+            row(
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE),
+                Pair.of("COLUMN_NAME", "m1")
+            ),
+            row(
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE),
+                Pair.of("COLUMN_NAME", "m2")
+            ),
+            row(
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOMEXDATASOURCE),
+                Pair.of("COLUMN_NAME", "m1_x")
+            ),
+            row(
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOMEXDATASOURCE),
+                Pair.of("COLUMN_NAME", "m2_x")
+            )
+        ),
+        getRows(
+            metaData.getColumns("druid", "dr_id", CalciteTests.SOME_DATASOURCE, "m%"),
+            ImmutableSet.of("TABLE_NAME", "TABLE_SCHEM", "COLUMN_NAME")
+        )
+    );
+  }
+
+  @Test
+  public void testEscapingForGetTables() throws Exception
+  {
+    final DatabaseMetaData metaData = client.getMetaData();
+
+    Assert.assertEquals(
+        ImmutableList.of(
+            row(
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE)
+            )
+        ),
+        getRows(
+            metaData.getTables("druid", "dr_id", CalciteTests.SOME_DATSOURCE_ESCAPED, null),
+            ImmutableSet.of("TABLE_SCHEM", "TABLE_NAME")
+        )
+    );
+
+    Assert.assertEquals(
+        ImmutableList.of(
+            row(
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOMEXDATASOURCE)
+            )
+        ),
+        getRows(
+            metaData.getTables("druid", "dr_id", CalciteTests.SOMEXDATASOURCE, null),
+            ImmutableSet.of("TABLE_SCHEM", "TABLE_NAME")
+        )
+    );
+
+    // Assert that some_datasource is treated as a pattern that matches some_datasource and somexdatasource
+    Assert.assertEquals(
+        ImmutableList.of(
+            row(
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOME_DATASOURCE)
+            ),
+            row(
+                Pair.of("TABLE_SCHEM", "druid"),
+                Pair.of("TABLE_NAME", CalciteTests.SOMEXDATASOURCE)
+            )
+        ),
+        getRows(
+            metaData.getTables("druid", "dr_id", CalciteTests.SOME_DATASOURCE, null),
+            ImmutableSet.of("TABLE_SCHEM", "TABLE_NAME")
+        )
+    );
+  }
+
   private static List<Map<String, Object>> getRows(final ResultSet resultSet) throws SQLException
   {
     return getRows(resultSet, null);

sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java

@@ -718,6 +718,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
             .add(new Object[]{"druid", CalciteTests.DATASOURCE4, "TABLE"})
             .add(new Object[]{"druid", CalciteTests.DATASOURCE5, "TABLE"})
             .add(new Object[]{"druid", CalciteTests.DATASOURCE3, "TABLE"})
+            .add(new Object[]{"druid", CalciteTests.SOME_DATASOURCE, "TABLE"})
+            .add(new Object[]{"druid", CalciteTests.SOMEXDATASOURCE, "TABLE"})
             .add(new Object[]{"druid", "aview", "VIEW"})
             .add(new Object[]{"druid", "bview", "VIEW"})
             .add(new Object[]{"INFORMATION_SCHEMA", "COLUMNS", "SYSTEM_TABLE"})
@@ -746,6 +748,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
             .add(new Object[]{"druid", CalciteTests.FORBIDDEN_DATASOURCE, "TABLE"})
             .add(new Object[]{"druid", CalciteTests.DATASOURCE5, "TABLE"})
             .add(new Object[]{"druid", CalciteTests.DATASOURCE3, "TABLE"})
+            .add(new Object[]{"druid", CalciteTests.SOME_DATASOURCE, "TABLE"})
+            .add(new Object[]{"druid", CalciteTests.SOMEXDATASOURCE, "TABLE"})
             .add(new Object[]{"druid", "aview", "VIEW"})
             .add(new Object[]{"druid", "bview", "VIEW"})
             .add(new Object[]{"INFORMATION_SCHEMA", "COLUMNS", "SYSTEM_TABLE"})

sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java

@@ -140,6 +140,9 @@ public class CalciteTests
   public static final String DATASOURCE4 = "foo4";
   public static final String DATASOURCE5 = "lotsocolumns";
   public static final String FORBIDDEN_DATASOURCE = "forbiddenDatasource";
+  public static final String SOME_DATASOURCE = "some_datasource";
+  public static final String SOME_DATSOURCE_ESCAPED = "some\\_datasource";
+  public static final String SOMEXDATASOURCE = "somexdatasource";
   public static final String DRUID_SCHEMA_NAME = "druid";
   public static final String INFORMATION_SCHEMA_NAME = "INFORMATION_SCHEMA";
   public static final String SYSTEM_SCHEMA_NAME = "sys";
@@ -294,6 +297,16 @@ public class CalciteTests
       .withRollup(false)
       .build();
 
+  private static final IncrementalIndexSchema INDEX_SCHEMA_WITH_X_COLUMNS = new IncrementalIndexSchema.Builder()
+      .withMetrics(
+          new CountAggregatorFactory("cnt_x"),
+          new FloatSumAggregatorFactory("m1_x", "m1_x"),
+          new DoubleSumAggregatorFactory("m2_x", "m2_x"),
+          new HyperUniquesAggregatorFactory("unique_dim1_x", "dim1_x")
+      )
+      .withRollup(false)
+      .build();
+
   private static final IncrementalIndexSchema INDEX_SCHEMA_NUMERIC_DIMS = new IncrementalIndexSchema.Builder()
       .withMetrics(
           new CountAggregatorFactory("cnt"),
@@ -361,6 +374,68 @@ public class CalciteTests
           .put("dim1", "abc")
           .build()
   );
+
+  public static final List<InputRow> RAW_ROWS1_X = ImmutableList.of(
+      createRow(
+          ImmutableMap.<String, Object>builder()
+              .put("t", "2000-01-01")
+              .put("m1_x", "1.0")
+              .put("m2_x", "1.0")
+              .put("dim1_x", "")
+              .put("dim2_x", ImmutableList.of("a"))
+              .put("dim3_x", ImmutableList.of("a", "b"))
+              .build()
+      ),
+      createRow(
+          ImmutableMap.<String, Object>builder()
+              .put("t", "2000-01-02")
+              .put("m1_x", "2.0")
+              .put("m2_x", "2.0")
+              .put("dim1_x", "10.1")
+              .put("dim2_x", ImmutableList.of())
+              .put("dim3_x", ImmutableList.of("b", "c"))
+              .build()
+      ),
+      createRow(
+          ImmutableMap.<String, Object>builder()
+              .put("t", "2000-01-03")
+              .put("m1_x", "3.0")
+              .put("m2_x", "3.0")
+              .put("dim1_x", "2")
+              .put("dim2_x", ImmutableList.of(""))
+              .put("dim3_x", ImmutableList.of("d"))
+              .build()
+      ),
+      createRow(
+          ImmutableMap.<String, Object>builder()
+              .put("t", "2001-01-01")
+              .put("m1_x", "4.0")
+              .put("m2_x", "4.0")
+              .put("dim1_x", "1")
+              .put("dim2_x", ImmutableList.of("a"))
+              .put("dim3_x", ImmutableList.of(""))
+              .build()
+      ),
+      createRow(
+          ImmutableMap.<String, Object>builder()
+              .put("t", "2001-01-02")
+              .put("m1_x", "5.0")
+              .put("m2_x", "5.0")
+              .put("dim1_x", "def")
+              .put("dim2_x", ImmutableList.of("abc"))
+              .put("dim3_x", ImmutableList.of())
+              .build()
+      ),
+      createRow(
+          ImmutableMap.<String, Object>builder()
+              .put("t", "2001-01-03")
+              .put("m1_x", "6.0")
+              .put("m2_x", "6.0")
+              .put("dim1_x", "abc")
+              .build()
+      )
+  );
+
   public static final List<InputRow> ROWS1 =
       RAW_ROWS1.stream().map(CalciteTests::createRow).collect(Collectors.toList());
 
@@ -635,6 +710,22 @@ public class CalciteTests
         .rows(ROWS_LOTS_OF_COLUMNS)
         .buildMMappedIndex();
 
+    final QueryableIndex someDatasourceIndex = IndexBuilder
+        .create()
+        .tmpDir(new File(tmpDir, "1"))
+        .segmentWriteOutMediumFactory(OffHeapMemorySegmentWriteOutMediumFactory.instance())
+        .schema(INDEX_SCHEMA)
+        .rows(ROWS1)
+        .buildMMappedIndex();
+
+    final QueryableIndex someXDatasourceIndex = IndexBuilder
+        .create()
+        .tmpDir(new File(tmpDir, "1"))
+        .segmentWriteOutMediumFactory(OffHeapMemorySegmentWriteOutMediumFactory.instance())
+        .schema(INDEX_SCHEMA_WITH_X_COLUMNS)
+        .rows(RAW_ROWS1_X)
+        .buildMMappedIndex();
+
 
     return new SpecificSegmentsQuerySegmentWalker(
         conglomerate,
@@ -695,6 +786,24 @@ public class CalciteTests
                    .size(0)
                    .build(),
         indexLotsOfColumns
+    ).add(
+        DataSegment.builder()
+                   .dataSource(SOME_DATASOURCE)
+                   .interval(indexLotsOfColumns.getDataInterval())
+                   .version("1")
+                   .shardSpec(new LinearShardSpec(0))
+                   .size(0)
+                   .build(),
+        someDatasourceIndex
+    ).add(
+        DataSegment.builder()
+                   .dataSource(SOMEXDATASOURCE)
+                   .interval(indexLotsOfColumns.getDataInterval())
+                   .version("1")
+                   .shardSpec(new LinearShardSpec(0))
+                   .size(0)
+                   .build(),
+        someXDatasourceIndex
     );
   }