Refactor relationship table operations to be functional index compatible

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

@@ -69,15 +69,15 @@ public class EbeanLocalRelationshipQueryDAO {
   public EbeanLocalRelationshipQueryDAO(EbeanServer server, EBeanDAOConfig eBeanDAOConfig) {
     _server = server;
     _eBeanDAOConfig = eBeanDAOConfig;
-    _sqlGenerator = new MultiHopsTraversalSqlGenerator(SUPPORTED_CONDITIONS);
     _schemaValidatorUtil = new SchemaValidatorUtil(server);
+    _sqlGenerator = new MultiHopsTraversalSqlGenerator(SUPPORTED_CONDITIONS, _schemaValidatorUtil);
   }
 
   public EbeanLocalRelationshipQueryDAO(EbeanServer server) {
     _server = server;
     _eBeanDAOConfig = new EBeanDAOConfig();
-    _sqlGenerator = new MultiHopsTraversalSqlGenerator(SUPPORTED_CONDITIONS);
     _schemaValidatorUtil = new SchemaValidatorUtil(server);
+    _sqlGenerator = new MultiHopsTraversalSqlGenerator(SUPPORTED_CONDITIONS, _schemaValidatorUtil);
   }
 
   static final Map<Condition, String> SUPPORTED_CONDITIONS =
@@ -125,8 +125,8 @@ private <SNAPSHOT extends RecordTemplate> List<SNAPSHOT> findEntitiesCore(@Nonnu
     final StringBuilder sqlBuilder = new StringBuilder();
     sqlBuilder.append("SELECT * FROM ").append(tableName);
     if (filterHasNonEmptyCriteria(filter)) {
-      sqlBuilder.append(" WHERE ").append(SQLStatementUtils.whereClause(filter, SUPPORTED_CONDITIONS, null,
-          _eBeanDAOConfig.isNonDollarVirtualColumnsEnabled()));
+      sqlBuilder.append(" WHERE ").append(SQLStatementUtils.whereClause(filter, SUPPORTED_CONDITIONS, null, tableName,
+          _schemaValidatorUtil, _eBeanDAOConfig.isNonDollarVirtualColumnsEnabled()));
     }
     sqlBuilder.append(" ORDER BY urn LIMIT ").append(Math.max(1, count)).append(" OFFSET ").append(Math.max(0, offset));
 
@@ -795,6 +795,7 @@ public String buildFindRelationshipSQL(@Nonnull final String relationshipTableNa
 
       String whereClause = SQLStatementUtils.whereClause(SUPPORTED_CONDITIONS,
           _eBeanDAOConfig.isNonDollarVirtualColumnsEnabled(),
+          relationshipTableName, _schemaValidatorUtil,
           filters.toArray(new Pair[filters.size()]));
 
       if (whereClause != null) {

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

@@ -15,9 +15,11 @@
  */
 public class MultiHopsTraversalSqlGenerator {
   private static Map<Condition, String> _supportedConditions;
+  private final SchemaValidatorUtil _schemaValidator;
 
-  public MultiHopsTraversalSqlGenerator(Map<Condition, String> supportedConditions) {
+  public MultiHopsTraversalSqlGenerator(Map<Condition, String> supportedConditions, SchemaValidatorUtil schemaValidator) {
     _supportedConditions = Collections.unmodifiableMap(supportedConditions);
+    _schemaValidator = schemaValidator;
   }
 
   /**
@@ -77,6 +79,7 @@ private String firstHopUrnsDirected(String relationshipTable, String srcEntityTa
             urnColumn, relationshipTable, destEntityTable, srcEntityTable));
 
     String whereClause = SQLStatementUtils.whereClause(_supportedConditions,  nonDollarVirtualColumnsEnabled,
+        relationshipTable, _schemaValidator,
         new Pair<>(relationshipFilter, "rt"),
         new Pair<>(destFilter, "dt"),
         new Pair<>(srcFilter, "st"));
@@ -105,6 +108,7 @@ private String firstHopUrnsUndirected(String relationshipTable, String entityTab
             relationshipTable, entityTable));
 
     String whereClause = SQLStatementUtils.whereClause(_supportedConditions, nonDollarVirtualColumnsEnabled,
+        relationshipTable, _schemaValidator,
         new Pair<>(relationshipFilter, "rt"),
         new Pair<>(srcFilter, "et"));
 
@@ -123,7 +127,8 @@ private String firstHopUrnsUndirected(String relationshipTable, String entityTab
   @ParametersAreNonnullByDefault
   private String findEntitiesUndirected(String entityTable, String relationshipTable, String firstHopUrnSql, LocalRelationshipFilter destFilter,
   boolean nonDollarVirtualColumnsEnabled) {
-    String whereClause = SQLStatementUtils.whereClause(_supportedConditions, nonDollarVirtualColumnsEnabled, new Pair<>(destFilter, "et"));
+    String whereClause = SQLStatementUtils.whereClause(_supportedConditions, nonDollarVirtualColumnsEnabled,
+        relationshipTable, _schemaValidator, new Pair<>(destFilter, "et"));
 
     StringBuilder sourceEntitySql = new StringBuilder(
         String.format("SELECT et.* FROM %s et INNER JOIN %s rt ON et.urn=rt.source WHERE rt.destination IN (%s)",

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

@@ -65,29 +65,47 @@ private static String parseIndexValue(@Nullable IndexValue indexValue) {
     }
   }
 
+  @Nullable
+  private static String getIndexedExpressionOrColumnGeneric(@Nonnull String expectedLegacyColumnName, @Nonnull String expectedExpressionIndexName,
+      @Nonnull String tableName, @Nonnull SchemaValidatorUtil schemaValidator) {
+    // Check if an expression-based index exists... if it does, use that
+    final String indexExpression = schemaValidator.getIndexExpression(tableName, expectedExpressionIndexName);
+    if (indexExpression != null) {
+      log.info("Using expression index '{}' in table '{}' with expression '{}'", expectedExpressionIndexName, tableName, indexExpression);
+      return indexExpression;
+    } else if (schemaValidator.columnExists(tableName, expectedLegacyColumnName)) {
+      // (Pre-functional-index logic) Check for existence of (virtual) column
+      return expectedLegacyColumnName;
+    } else {
+      return null;
+    }
+  }
+
   /**
    * Get the expression index "identifier", if it exists, otherwise retrieve the generated column name.
    * The idea behind this is that whatever is returned from this method can be used verbatim to query the database;
    * it's either the expression index itself (new approach) or the virtual column (old approach).
+   * Intended to be used with entity tables.
    */
   @Nullable
   public static String getIndexedExpressionOrColumn(@Nonnull String assetType, @Nonnull String aspect, @Nonnull String path,
       boolean nonDollarVirtualColumnsEnabled, @Nonnull SchemaValidatorUtil schemaValidator) {
     final String indexColumn = getGeneratedColumnName(assetType, aspect, path, nonDollarVirtualColumnsEnabled);
     final String tableName = getTableName(assetType);
+    return getIndexedExpressionOrColumnGeneric(indexColumn, getExpressionIndexName(assetType, aspect, path), tableName, schemaValidator);
+  }
 
-    // Check if an expression-based index exists... if it does, use that
-    final String expressionIndexName = getExpressionIndexName(assetType, aspect, path);
-    final String indexExpression = schemaValidator.getIndexExpression(tableName, expressionIndexName);
-    if (indexExpression != null) {
-      log.info("Using expression index '{}' in table '{}' with expression '{}'", expressionIndexName, tableName, indexExpression);
-      return indexExpression;
-    } else if (schemaValidator.columnExists(tableName, indexColumn)) {
-      // (Pre-functional-index logic) Check for existence of (virtual) column
-      return indexColumn;
-    } else {
-      return null;
-    }
+  /**
+   * Get the expression index "identifier", if it exists, otherwise retrieve the generated column name.
+   * The idea behind this is that whatever is returned from this method can be used verbatim to query the database;
+   * it's either the expression index itself (new approach) or the virtual column (old approach).
+   * Intended to be used with relationship tables.
+   */
+  @Nullable
+  public static String getIndexedExpressionOrColumnRelationship(@Nonnull String expectedLegacyColumnName, @Nonnull String path,
+      @Nonnull String tableName, @Nonnull SchemaValidatorUtil schemaValidator) {
+    final String expectedExpressionIndexName = SQLSchemaUtils.getExpressionIndexNameRelationship(path);
+    return getIndexedExpressionOrColumnGeneric(expectedLegacyColumnName, expectedExpressionIndexName, tableName, schemaValidator);
   }
 
   /**

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

@@ -25,6 +25,7 @@ public class SQLSchemaUtils {
   public static final String ASPECT_PREFIX = "a_";
   public static final String INDEX_PREFIX = "i_";
   public static final String EXPRESSION_INDEX_PREFIX = "e_";
+  public static final String RELATIONSHIP_TABLE_EXPRESSION_INDEX_INFIX = "metadata";
 
   private static final int MYSQL_MAX_COLUMN_NAME_LENGTH = 64 - ASPECT_PREFIX.length();
 
@@ -152,21 +153,31 @@ public static <ASPECT extends RecordTemplate> String getAspectColumnName(@Nonnul
     return getAspectColumnName(entityType, aspectClass.getCanonicalName());
   }
 
+  @Nonnull
+  private static String getExpectedNameFormatter(
+      @Nonnull String prefix,
+      @Nonnull String infix,
+      @Nonnull String path,
+      boolean nonDollarVirtualColumnsEnabled) {
+    char delimiter = nonDollarVirtualColumnsEnabled ? '0' : '$';
+    return prefix + infix + processPath(path, delimiter);
+  }
+
   @Nonnull
   private static String getExpectedNameHelper(
       @Nonnull String prefix,
       @Nonnull String assetType,
       @Nonnull String aspect,
       @Nonnull String path,
       boolean nonDollarVirtualColumnsEnabled) {
-    char delimiter = nonDollarVirtualColumnsEnabled ? '0' : '$';
     if (isUrn(aspect)) {
-      return prefix + "urn" + processPath(path, delimiter);
+      return getExpectedNameFormatter(prefix, "urn", path, nonDollarVirtualColumnsEnabled);
     }
     if (UNKNOWN_ASSET.equals(assetType)) {
-      log.warn("query with unknown asset type. aspect =  {}, path ={}, delimiter = {}", aspect, path, delimiter);
+      log.warn("query with unknown asset type. aspect =  {}, path ={}, nonDollarVirtualColumnsEnabled={}",
+          aspect, path, nonDollarVirtualColumnsEnabled);
     }
-    return prefix + getColumnName(assetType, aspect) + processPath(path, delimiter);
+    return getExpectedNameFormatter(prefix, getColumnName(assetType, aspect), path, nonDollarVirtualColumnsEnabled);
   }
 
   /**
@@ -191,6 +202,16 @@ public static String getExpressionIndexName(@Nonnull String assetType, @Nonnull
     return getExpectedNameHelper(EXPRESSION_INDEX_PREFIX, assetType, aspect, path, true);
   }
 
+  /**
+   * Get the expected expression index name for a relationship table given the path.
+   * With the expression index changes, we establish an expected naming as follows...
+   * ex. e_metadata0foo0bar
+   */
+  @Nonnull
+  public static String getExpressionIndexNameRelationship(@Nonnull String path) {
+    return getExpectedNameFormatter(EXPRESSION_INDEX_PREFIX, RELATIONSHIP_TABLE_EXPRESSION_INDEX_INFIX, path, true);
+  }
+
   /**
    * DEPRECATED, use getGeneratedColumnName(assetType, aspect, path, nonDollarVirtualColumnsEnabled) instead.
    */