Refactor ARRAY_CONTAINS query syntax (#592)

Author: jphui

dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLIndexFilterUtils.java

@@ -1,5 +1,6 @@
 package com.linkedin.metadata.dao.utils;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.linkedin.data.template.GetMode;
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.data.template.StringArray;
@@ -162,6 +163,28 @@ public static String parseIndexFilter(@Nonnull String entityType, @Nullable Inde
 
   }
 
+  /**
+   * Strip CAST() statement from the input string. The intention for this is to be used for array-based functional
+   * indexes in which the CAST() statement must be removed before consumption from a JSON_CONTAINS() on the query side.
+   * Note that the CAST() statement is expected to be the **outermost** statement in the input string.
+   * Specifics:
+   * - leading '(' is optional
+   * - trailing ')' is optional
+   * - whitespace is lenient -- before and after CAST, AS, between closing parentheses
+   * - case-insensitive
+   * - has to account for closing parens from the "AS" statement -- ie. as char(128)
+   * - has to account for array casting: ... as char(128) array
+   * @param inputString input string
+   *                    ex. (CAST(JSON_EXTRACT(a_asset_labels, '$.aspect.derived_labels') AS CHAR(128) ARRAY))
+   * @return stripped string, enclosed in parentheses
+   *                    ex. (JSON_EXTRACT(a_asset_labels, '$.aspect.derived_labels'))
+   */
+  @VisibleForTesting
+  @Nonnull
+  protected static String stripCastStatement(@Nonnull String inputString) {
+    return inputString.replaceFirst("^\\(?\\s*(?i)CAST\\s*\\(", "(").replaceFirst("\\s+(?i)AS\\s+[^)]*(?:\\([^)]*\\))?[^)]*\\s*\\){1,2}$", ")");
+  }
+
   /**
    * Parse condition expression.
    * @param index the name of the virtual generated column OR the actual expression of a functional index
@@ -174,7 +197,7 @@ private static String parseSqlFilter(String index, Condition condition, IndexVal
     switch (condition) {
       // TODO: add validation to check that the index column value is an array type
       case ARRAY_CONTAINS:
-        return String.format("'%s' MEMBER OF(%s)", parseIndexValue(indexValue), index);  // JSON Array
+        return String.format("JSON_CONTAINS(%s, '%s')", stripCastStatement(index), parseIndexValue(indexValue));  // JSON Array
       case CONTAIN:
         return String.format("JSON_SEARCH(%s, 'one', '%s') IS NOT NULL", index, parseIndexValue(indexValue));  // JSON String, Array, Struct
       case IN:

dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/utils/SQLIndexFilterUtilsTest.java

@@ -123,7 +123,7 @@ public void testParseIndexFilterWithFunctionalIndex() {
         SQLIndexFilterUtils.createIndexCriterion(AspectBar.class, "value_array", Condition.ARRAY_CONTAINS,
             IndexValue.create(12L)));
     final String expectedSql1 =
-        "WHERE a_aspectbar IS NOT NULL\nAND JSON_EXTRACT(a_aspectbar, '$.gma_deleted') IS NULL\nAND '12' MEMBER OF((cast(json_extract(`a_aspectbar`, '$.aspect.value_array') as char(128) array)))\nAND deleted_ts IS NULL";
+        "WHERE a_aspectbar IS NOT NULL\nAND JSON_EXTRACT(a_aspectbar, '$.gma_deleted') IS NULL\nAND JSON_CONTAINS((json_extract(`a_aspectbar`, '$.aspect.value_array')), '12')\nAND deleted_ts IS NULL";
     assertValidSql(expectedSql1);  // assert that the expected SQL is valid to begin with
     assertEquals(SQLIndexFilterUtils.parseIndexFilter(FooUrn.ENTITY_TYPE, indexFilter, false, mockValidator),
         expectedSql1);
@@ -240,4 +240,39 @@ public void testGetIndexedExpressionOrColumn() {
         false, mockValidator),
         "(cast(json_extract(`a_aspectbar`, '$.aspect.value') as char(1024)))");
   }
+
+  @Test
+  public void testStripCastStatement() {
+    // lowercase
+    assertEquals(SQLIndexFilterUtils.stripCastStatement("(cast(json_extract(`a_aspectbar`, '$.aspect.value') as char(1024)))"),
+        "(json_extract(`a_aspectbar`, '$.aspect.value'))");
+
+    // all caps
+    assertEquals(SQLIndexFilterUtils.stripCastStatement("(CAST(json_extract(`a_aspectbar`, '$.aspect.value') AS char(1024)))"),
+        "(json_extract(`a_aspectbar`, '$.aspect.value'))");
+
+    // mixed case
+    assertEquals(SQLIndexFilterUtils.stripCastStatement("(CaSt(json_extract(`a_aspectbar`, '$.aspect.value') As char(1024)))"),
+        "(json_extract(`a_aspectbar`, '$.aspect.value'))");
+
+    // extra spaces (before 'cast', before 'as')
+    assertEquals(SQLIndexFilterUtils.stripCastStatement("(      cast(json_extract(`a_aspectbar`, '$.aspect.value')     as char(1024)))"),
+        "(json_extract(`a_aspectbar`, '$.aspect.value'))");
+
+    // extra spaces (after 'cast', after 'as')
+    assertEquals(SQLIndexFilterUtils.stripCastStatement("(cast    (json_extract(`a_aspectbar`, '$.aspect.value') as      char(1024)))"),
+        "(json_extract(`a_aspectbar`, '$.aspect.value'))");
+
+    // casting as an array
+    assertEquals(SQLIndexFilterUtils.stripCastStatement("(cast(json_extract(`a_aspectbar`, '$.aspect.value') as char(1024) array))"),
+        "(json_extract(`a_aspectbar`, '$.aspect.value'))");
+
+    // NOT the outermost statement
+    assertEquals(SQLIndexFilterUtils.stripCastStatement("(foo(cast(json_extract(`a_aspectbar`, '$.aspect.value') as char(1024))))"),
+        "(foo(cast(json_extract(`a_aspectbar`, '$.aspect.value') as char(1024))))");
+
+    // no enclosing parents in original statement
+    assertEquals(SQLIndexFilterUtils.stripCastStatement("cast(json_extract(`a_aspectbar`, '$.aspect.value') as char(1024))"),
+        "(json_extract(`a_aspectbar`, '$.aspect.value'))");
+  }
 }

dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/utils/SQLStatementUtilsTest.java

@@ -197,10 +197,10 @@ public void testCreateFilterSqlWithArrayContainsCondition() {
 
     String sql1 = SQLStatementUtils.createFilterSql("foo", indexFilter, true, false, mockValidator);
     String expectedSql1 = "SELECT urn, (SELECT COUNT(urn) FROM metadata_entity_foo WHERE a_aspectfoobar IS NOT NULL\n"
-        + "AND JSON_EXTRACT(a_aspectfoobar, '$.gma_deleted') IS NULL\n" + "AND 'bar1' MEMBER OF(i_aspectfoobar$bars)\n"
+        + "AND JSON_EXTRACT(a_aspectfoobar, '$.gma_deleted') IS NULL\n" + "AND JSON_CONTAINS(i_aspectfoobar$bars, 'bar1')\n"
         + "AND deleted_ts IS NULL)" + " as _total_count FROM metadata_entity_foo\n"
         + "WHERE a_aspectfoobar IS NOT NULL\n" + "AND JSON_EXTRACT(a_aspectfoobar, '$.gma_deleted') IS NULL\n"
-        + "AND 'bar1' MEMBER OF(i_aspectfoobar$bars)\n" + "AND deleted_ts IS NULL";
+        + "AND JSON_CONTAINS(i_aspectfoobar$bars, 'bar1')\n" + "AND deleted_ts IS NULL";
 
     assertEquals(sql1, expectedSql1);
   }