(feat) Add Expression Index Support Utility (#585)

Author: jphui

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

@@ -118,8 +118,19 @@ public static String parseIndexFilter(@Nonnull String entityType, @Nullable Inde
           final Condition condition = pathParams.getCondition();
           final String indexColumn = getGeneratedColumnName(entityType, aspect, pathParams.getPath(), nonDollarVirtualColumnsEnabled);
           final String tableName = SQLSchemaUtils.getTableName(entityType);
-          // New: Skip filter if column doesn't exist
+
+          // NEW / TODO: Check if an expression-based index exists, if it does, use the new logic
+          final String indexExpression = schemaValidator.getIndexExpression(tableName, indexColumn);
+          if (indexExpression != null) {
+            log.debug("Using expression index '{}' in table '{}' with expression '{}'", indexColumn, tableName, indexExpression);
+            //// Commenting this out for now... to be extra safe, will not currently make this queryable yet
+            ////   and should verify that the above debug log is printed to properly acknoledge an expression.
+            // sqlFilters.add(parseSqlFilter(indexExpression, condition, pathParams.getValue()));
+          }
+
+          // FOR NOW: keep old logic to allow parallel usage of new indices and validation
           if (!schemaValidator.columnExists(tableName, indexColumn)) {
+            // Else: (old logic) Skip filter if column doesn't exist
             log.warn("Skipping filter: virtual column '{}' not found in table '{}'", indexColumn, tableName);
             continue;
           }

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

@@ -5,11 +5,14 @@
 import com.google.common.annotations.VisibleForTesting;
 import io.ebean.EbeanServer;
 import io.ebean.SqlRow;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import lombok.extern.slf4j.Slf4j;
 
 
@@ -39,10 +42,27 @@ public class SchemaValidatorUtil {
       .maximumSize(1000)
       .build();
 
+  // Cache: tableName → Set of index names -> expression that defines the index, used as a replacement for creating an index on virtual columns
+  // Configuration:
+  // - expireAfterWrite(10 minutes): Ensures that newly added indexes (e.g., via Pretzel) are picked up automatically
+  //   without requiring a service restart. After 10 minutes, the next request will trigger a DB refresh.
+  // - maximumSize(1000): Limits cache memory footprint by retaining entries for up to 1000 distinct tables.
+  //   Least recently used entries are evicted when the size limit is reached.
+  // ** THIS IS NEEDED ** because of local testing limitations by MariaDB: expression-based indexes are not supported,
+  //    so no existing logic should depend on anything introduced by the support of this. Otherwise, we'd need to mock
+  //    all indexing code in the test DB, which I want to avoid if possible.
+  // TODO: This can become the only cache needed for indexes once we are 100% migrated over to this logic.
+  private final Cache<String, Map<String, String>> indexExpressionCache = Caffeine.newBuilder()
+      .expireAfterWrite(10, TimeUnit.MINUTES)
+      .maximumSize(1000)
+      .build();
+
   private static final String SQL_GET_ALL_COLUMNS =
       "SELECT COLUMN_NAME FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = database() AND TABLE_NAME = '%s'";
   private static final String SQL_GET_ALL_INDEXES =
       "SELECT DISTINCT INDEX_NAME FROM information_schema.STATISTICS WHERE TABLE_SCHEMA = database() AND TABLE_NAME = '%s'";
+  private static final String SQL_GET_ALL_INDEXES_WITH_EXPRESSIONS =
+      "SELECT DISTINCT INDEX_NAME, EXPRESSION FROM information_schema.STATISTICS WHERE TABLE_SCHEMA = database() AND TABLE_NAME = '%s'";
 
   public SchemaValidatorUtil(EbeanServer server) {
     this.server = server;
@@ -97,6 +117,60 @@ public boolean indexExists(@Nonnull String tableName, @Nonnull String indexName)
     return indexes.contains(lowerIndex);
   }
 
+
+  /**
+   * Cleans SQL expression by removing MySQL-specific encoding artifacts.
+   * Removes _utf8mb4 charset prefix, unescapes quotes, and removes newlines.
+   * MySQL team is the POC for questions about this since there is preprocessing needed to transform the as-is
+   * index expression from the index table to a (string) expression that is usable directly in an indexed query.
+   *
+   * @param expression Raw SQL expression from database
+   * @return Cleaned expression string, with enclosing parentheses
+   */
+  @VisibleForTesting
+  protected String cleanIndexExpression(@Nullable String expression) {
+    if (expression == null) {
+      return null;
+    }
+
+    return "(" + expression
+        .replace("_utf8mb4\\'", "'")
+        .replace("\\'", "'")
+        .replace("\\\"", "\"")
+        .replace("\n", "") + ")";
+  }
+
+
+  /**
+   * Retrieves the expression associated with the given index.
+   *
+   * <p>NULL doesn't necessarily mean that an index doesn't exist, use {@link #indexExists(String, String)} to check for index existence.
+   *
+   * @param tableName Table name
+   * @param indexName Index name
+   * @return Expression string, or null if index does not exist OR is not created on an expression; will be enclosed in
+   *         parentheses '()'
+   */
+  @Nullable
+  public String getIndexExpression(@Nonnull String tableName, @Nonnull String indexName) {
+    String lowerTable = tableName.toLowerCase();
+    String lowerIndex = indexName.toLowerCase();
+
+    try {
+      Map<String, String> indexes = indexExpressionCache.get(lowerTable, tbl -> {
+        log.info("Refreshing index cache for table '{}' from expression retrieval call", tbl);
+        return loadIndexesAndExpressions(tbl);
+      });
+
+      // This will also return null if the Expression column is null itself
+      return cleanIndexExpression(indexes.getOrDefault(lowerIndex, null));
+    } catch (Exception e) {
+      // MariaDB for local testing doesn't support EXPRESSION column - gracefully degrade
+      log.debug("Unable to load index expressions for table '{}': {}", lowerTable, e.getMessage());
+      return null;  // same logic as "no expression exists", which is good (handled gracefully)
+    }
+  }
+
   /**
    * Loads all columns for the given table from information_schema.
    *
@@ -127,4 +201,21 @@ private Set<String> loadIndexes(String tableName) {
     return indexes;
   }
 
+  /**
+   * Loads all index names and  expressions for the given table from information_schema.
+   * See the comment for indexExpressionCache for more details.
+   *
+   * @param tableName Table to query
+   * @return Map of lowercase index names -> expressions
+   */
+  private Map<String, String> loadIndexesAndExpressions(String tableName) {
+    List<SqlRow> rows = server.createSqlQuery(String.format(SQL_GET_ALL_INDEXES_WITH_EXPRESSIONS, tableName)).findList();
+    Map<String, String> indexes = new HashMap<>();
+    for (SqlRow row : rows) {
+      // The Expression value will be null if the index is not created on an expression
+      indexes.put(row.getString("INDEX_NAME").toLowerCase(), row.getString("EXPRESSION"));
+    }
+    return indexes;
+  }
+
 }

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

@@ -4,8 +4,12 @@
 import com.linkedin.metadata.dao.EBeanDAOConfig;
 import io.ebean.Ebean;
 import io.ebean.EbeanServer;
+import io.ebean.SqlQuery;
+import io.ebean.SqlRow;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.DataProvider;
@@ -14,6 +18,7 @@
 
 import static com.linkedin.common.AuditStamps.*;
 import static com.linkedin.testing.TestUtils.*;
+import static org.mockito.Mockito.*;
 import static org.testng.AssertJUnit.*;
 
 
@@ -38,7 +43,9 @@ public static Object[][] inputList() {
 
   @BeforeClass
   public void init() {
-    server = EmbeddedMariaInstance.getServer(SchemaValidatorUtilTest.class.getSimpleName());
+    // need to mock this since we will be stubbing in for the EXPRESSION column retrieval for that test since
+    // MariaDB doesn't support functional indexes
+    server = spy(EmbeddedMariaInstance.getServer(SchemaValidatorUtilTest.class.getSimpleName()));
   }
 
   @BeforeMethod
@@ -77,4 +84,66 @@ public void testCheckIndexExists() {
     }
   }
 
+  @Test
+  public void testCleanExpression() {
+
+  }
+
+  /**
+   * These tests require mocking because MariaDB, our embedded test database, does not support functional indexes, which
+   * the code under test is trying to access.
+   */
+  @Test
+  public void testGetIndexExpression() {
+    // NEED to set up all mocks for DB access BEFORE running ANY tests because it will be cached
+    SqlQuery sqlQuery = mock(SqlQuery.class);
+    List<SqlRow> indexTable = new ArrayList<>();
+
+    when(sqlQuery.findList()).thenReturn(indexTable);
+    when(server.createSqlQuery(anyString())).thenReturn(sqlQuery);
+
+    // setup mock for the LEGACY index use case: no expression-based index, but the index still exists!
+    SqlRow row1 = mock(SqlRow.class);
+    indexTable.add(row1);
+    when(row1.getString("EXPRESSION")).thenReturn(null);
+
+    // setup mock for the EXPRESSION index use case
+    SqlRow row2 = mock(SqlRow.class);
+    indexTable.add(row2);
+    when(row2.getString("EXPRESSION")).thenReturn(
+        "cast(json_extract(`a_aspectfoo`, '$.aspect.value') as char(1024) charset utf8mb4)");
+
+    if (!ebeanConfig.isNonDollarVirtualColumnsEnabled()) {
+      when(row1.getString("INDEX_NAME")).thenReturn("i_aspectfoo$value");
+      when(row2.getString("INDEX_NAME")).thenReturn("idx_aspectfoo$value");
+    } else {
+      when(row1.getString("INDEX_NAME")).thenReturn("i_aspectfoo0value");
+      when(row2.getString("INDEX_NAME")).thenReturn("idx_aspectfoo0value");
+    }
+
+
+    // NONEXISTENT test
+    assertNull(validator.getIndexExpression("metadata_entity_burger", "idx_fake"));
+
+    if (!ebeanConfig.isNonDollarVirtualColumnsEnabled()) {
+      /// Verify!
+      assertNotNull(validator.getIndexExpression("metadata_entity_burger", "idx_aspectfoo$value"));
+      assertEquals("(cast(json_extract(`a_aspectfoo`, '$.aspect.value') as char(1024) charset utf8mb4))",
+          validator.getIndexExpression("metadata_entity_burger", "idx_aspectfoo$value"));
+
+      // Make sure that retrieving a "legacy" column-based index still returns true but returns null
+      assertTrue(validator.indexExists("metadata_entity_foo", "i_aspectfoo$value"));
+      assertNull(validator.getIndexExpression("metadata_entity_foo", "i_aspectfoo$value"));
+    } else {
+      /// Verify!
+      assertNotNull(validator.getIndexExpression("metadata_entity_burger", "idx_aspectfoo0value"));
+      assertEquals("(cast(json_extract(`a_aspectfoo`, '$.aspect.value') as char(1024) charset utf8mb4))",
+          validator.getIndexExpression("metadata_entity_burger", "idx_aspectfoo0value"));
+
+      // Make sure that retrieving a "legacy" column-based index still returns true but returns null
+      assertTrue(validator.indexExists("metadata_entity_foo", "i_aspectfoo0value"));
+      assertNull(validator.getIndexExpression("metadata_entity_foo", "i_aspectfoo0value"));
+    }
+  }
+
 }