(feat) Add SchemaValidatorUtil to Enable Safe Decoupling from DB Schema Evolution (#537)

* Add SchemaValidatorUtil to Enable Safe Decoupling from DB Schema Evolution

* Checkstyle

* Checkstyle

* Fix index

* Fix index

* Fix index

* Fix index

* Added TTL

* Adding check column for filter/index

* Adding unit test to SQLStatementUtilsTest

* Adding unit test to EbeanLocalAccessTest

* Checkstyle

* Checkstyle

* Fix unit test

* Fix unit test

* Checkstyle

* Unit tests

* Unit tests

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

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

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

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Addressed comment

---------

Co-authored-by: Rakhi Agrawal <rakagrawal@linkedin.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

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

@@ -12,6 +12,7 @@
 import com.linkedin.metadata.dao.utils.RecordUtils;
 import com.linkedin.metadata.dao.utils.SQLSchemaUtils;
 import com.linkedin.metadata.dao.utils.SQLStatementUtils;
+import com.linkedin.metadata.dao.utils.SchemaValidatorUtil;
 import com.linkedin.metadata.events.IngestionTrackingContext;
 import com.linkedin.metadata.query.ExtraInfo;
 import com.linkedin.metadata.query.ExtraInfoArray;
@@ -74,6 +75,7 @@ public class EbeanLocalAccess<URN extends Urn> implements IEbeanLocalAccess<URN>
   // key: table_name,
   // value: Set(column1, column2, column3 ...)
   private final Map<String, Set<String>> tableColumns = new ConcurrentHashMap<>();
+  private final SchemaValidatorUtil validator;
 
   public EbeanLocalAccess(EbeanServer server, ServerConfig serverConfig, @Nonnull Class<URN> urnClass,
       UrnPathExtractor<URN> urnPathExtractor, boolean nonDollarVirtualColumnsEnabled) {
@@ -83,6 +85,7 @@ public EbeanLocalAccess(EbeanServer server, ServerConfig serverConfig, @Nonnull
     _entityType = ModelUtils.getEntityTypeFromUrnClass(_urnClass);
     _schemaEvolutionManager = createSchemaEvolutionManager(serverConfig);
     _nonDollarVirtualColumnsEnabled = nonDollarVirtualColumnsEnabled;
+    validator = new SchemaValidatorUtil(server);
   }
 
   public void setUrnPathExtractor(@Nonnull UrnPathExtractor<URN> urnPathExtractor) {
@@ -276,7 +279,7 @@ public <ASPECT extends RecordTemplate> List<EbeanMetadataAspect> batchGetUnion(
     for (int index = position; index < end; index++) {
       final Urn entityUrn = aspectKeys.get(index).getUrn();
       final Class<ASPECT> aspectClass = (Class<ASPECT>) aspectKeys.get(index).getAspectClass();
-      if (checkColumnExists(isTestMode ? getTestTableName(entityUrn) : getTableName(entityUrn),
+      if (validator.columnExists(isTestMode ? getTestTableName(entityUrn) : getTableName(entityUrn),
           getAspectColumnName(entityUrn.getEntityType(), aspectClass))) {
         keysToQueryMap.computeIfAbsent(aspectClass, unused -> new HashSet<>()).add(entityUrn);
       }
@@ -326,7 +329,7 @@ public ListResult<URN> listUrns(@Nullable IndexFilter indexFilter, @Nullable Ind
       int start, int pageSize) {
     final SqlQuery sqlQuery = createFilterSqlQuery(indexFilter, indexSortCriterion, start, pageSize);
     final List<SqlRow> sqlRows = sqlQuery.findList();
-    if (sqlRows.size() == 0) {
+    if (sqlRows.isEmpty()) {
       final List<SqlRow> totalCountResults = createFilterSqlQuery(indexFilter, indexSortCriterion, 0, DEFAULT_PAGE_SIZE).findList();
       final int actualTotalCount = totalCountResults.isEmpty() ? 0 : totalCountResults.get(0).getInteger("_total_count");
       return toListResult(actualTotalCount, start, pageSize);
@@ -428,13 +431,13 @@ public Map<String, Long> countAggregate(@Nullable IndexFilter indexFilter,
         getGeneratedColumnName(_entityType, indexGroupByCriterion.getAspect(), indexGroupByCriterion.getPath(),
             _nonDollarVirtualColumnsEnabled);
     // first, check for existence of the column we want to GROUP BY
-    if (!checkColumnExists(tableName, groupByColumn)) {
+    if (!validator.columnExists(tableName, groupByColumn)) {
       // if we are trying to GROUP BY the results on a column that does not exist, just return an empty map
       return Collections.emptyMap();
     }
 
     // now run the actual GROUP BY query
-    final String groupBySql = SQLStatementUtils.createGroupBySql(_entityType, indexFilter, indexGroupByCriterion, _nonDollarVirtualColumnsEnabled);
+    final String groupBySql = SQLStatementUtils.createGroupBySql(_entityType, indexFilter, indexGroupByCriterion, _nonDollarVirtualColumnsEnabled, validator);
     final SqlQuery sqlQuery = _server.createSqlQuery(groupBySql);
     final List<SqlRow> sqlRows = sqlQuery.findList();
     Map<String, Long> resultMap = new HashMap<>();
@@ -461,7 +464,7 @@ public Map<String, Long> countAggregate(@Nullable IndexFilter indexFilter,
   private SqlQuery createFilterSqlQuery(@Nullable IndexFilter indexFilter,
       @Nullable IndexSortCriterion indexSortCriterion, int offset, int pageSize) {
     StringBuilder filterSql = new StringBuilder();
-    filterSql.append(SQLStatementUtils.createFilterSql(_entityType, indexFilter, true, _nonDollarVirtualColumnsEnabled));
+    filterSql.append(SQLStatementUtils.createFilterSql(_entityType, indexFilter, true, _nonDollarVirtualColumnsEnabled, validator));
     filterSql.append("\n");
     filterSql.append(parseSortCriteria(_entityType, indexSortCriterion, _nonDollarVirtualColumnsEnabled));
     filterSql.append(String.format(" LIMIT %d", Math.max(pageSize, 0)));
@@ -475,7 +478,7 @@ private SqlQuery createFilterSqlQuery(@Nullable IndexFilter indexFilter,
   private SqlQuery createFilterSqlQuery(@Nullable IndexFilter indexFilter,
       @Nullable IndexSortCriterion indexSortCriterion, @Nullable URN lastUrn, int pageSize) {
     StringBuilder filterSql = new StringBuilder();
-    filterSql.append(SQLStatementUtils.createFilterSql(_entityType, indexFilter, false, _nonDollarVirtualColumnsEnabled));
+    filterSql.append(SQLStatementUtils.createFilterSql(_entityType, indexFilter, false, _nonDollarVirtualColumnsEnabled, validator));
 
     if (lastUrn != null) {
       // because createFilterSql will only include a WHERE clause if there are non-urn filters, we need to make sure
@@ -621,23 +624,6 @@ private SchemaEvolutionManager createSchemaEvolutionManager(@Nonnull ServerConfi
     return new FlywaySchemaEvolutionManager(config);
   }
 
-  /**
-   * Check column exists in table.
-   */
-  public boolean checkColumnExists(@Nonnull String tableName, @Nonnull String columnName) {
-    // Fetch table columns on very first read and cache it in tableColumns
-    if (!tableColumns.containsKey(tableName)) {
-      final List<SqlRow> rows = _server.createSqlQuery(SQLStatementUtils.getAllColumnForTable(tableName)).findList();
-      Set<String> columns = new HashSet<>();
-      for (SqlRow row : rows) {
-        columns.add(row.getString("COLUMN_NAME").toLowerCase());
-      }
-      tableColumns.put(tableName, columns);
-    }
-
-    return tableColumns.get(tableName).contains(columnName.toLowerCase());
-  }
-
   /**
    * SQL implementation of find the latest {@link EbeanMetadataAspect}.
    * @param connection {@link Connection} get from the current transaction, it should not be closed manually

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

@@ -12,6 +12,7 @@
 import com.linkedin.metadata.dao.utils.RelationshipLookUpContext;
 import com.linkedin.metadata.dao.utils.SQLSchemaUtils;
 import com.linkedin.metadata.dao.utils.SQLStatementUtils;
+import com.linkedin.metadata.dao.utils.SchemaValidatorUtil;
 import com.linkedin.metadata.query.Condition;
 import com.linkedin.metadata.query.LocalRelationshipCriterion;
 import com.linkedin.metadata.query.LocalRelationshipFilter;
@@ -46,6 +47,7 @@ public class EbeanLocalRelationshipQueryDAO {
   public static final String RELATIONSHIP_RETURN_TYPE = "relationship.return.type";
   public static final String MG_INTERNAL_ASSET_RELATIONSHIP_TYPE = "AssetRelationship.proto";
   private static final int FILTER_BATCH_SIZE = 200;
+  private static final String IDX_DESTINATION_DELETED_TS = "idx_destination_deleted_ts";
   private static final String FORCE_IDX_ON_DESTINATION = " FORCE INDEX (idx_destination_deleted_ts) ";
   private static final String DESTINATION_FIELD =  "destination";
   private final EbeanServer _server;
@@ -55,17 +57,20 @@ public class EbeanLocalRelationshipQueryDAO {
 
   private Set<String> _mgEntityTypeNameSet;
   private EbeanLocalDAO.SchemaConfig _schemaConfig = EbeanLocalDAO.SchemaConfig.NEW_SCHEMA_ONLY;
+  private SchemaValidatorUtil _schemaValidatorUtil;
 
   public EbeanLocalRelationshipQueryDAO(EbeanServer server, EBeanDAOConfig eBeanDAOConfig) {
     _server = server;
     _eBeanDAOConfig = eBeanDAOConfig;
     _sqlGenerator = new MultiHopsTraversalSqlGenerator(SUPPORTED_CONDITIONS);
+    _schemaValidatorUtil = new SchemaValidatorUtil(server);
   }
 
   public EbeanLocalRelationshipQueryDAO(EbeanServer server) {
     _server = server;
     _eBeanDAOConfig = new EBeanDAOConfig();
     _sqlGenerator = new MultiHopsTraversalSqlGenerator(SUPPORTED_CONDITIONS);
+    _schemaValidatorUtil = new SchemaValidatorUtil(server);
   }
 
   static final Map<Condition, String> SUPPORTED_CONDITIONS =
@@ -552,12 +557,14 @@ private String buildFindRelationshipSQL(
         // non-mg entity case, applying dest filter on relationship table
         filters.add(new Pair<>(destinationEntityFilter, "rt"));
       } else if (!relationshipFilter.getCriteria().isEmpty()) {
-        //TODO: Add a safeguard to check if the FORCE_IDX_ON_DESTINATION is present in the table, we can check once on bootup and then caching the result
-        // Check if any relationship-level filter is on "destination"
+        // Apply FORCE INDEX if destination field is being filtered, and the index exists
         for (LocalRelationshipCriterion criterion : relationshipFilter.getCriteria()) {
           LocalRelationshipCriterion.Field field = criterion.getField();
           if (field.getUrnField() != null && DESTINATION_FIELD.equals(field.getUrnField().getName())) {
-            sqlBuilder.append(FORCE_IDX_ON_DESTINATION);
+            // Check if index exists on 'destination' before applying FORCE INDEX
+            if (_schemaValidatorUtil.indexExists(relationshipTableName, IDX_DESTINATION_DELETED_TS)) {
+              sqlBuilder.append(FORCE_IDX_ON_DESTINATION);
+            }
             break;
           }
         }

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

@@ -15,6 +15,7 @@
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang.StringEscapeUtils;
 
 import static com.linkedin.metadata.dao.utils.SQLSchemaUtils.*;
@@ -24,6 +25,7 @@
 /**
  * Condition SQL script utils based on {@link IndexFilter}.
  */
+@Slf4j
 public class SQLIndexFilterUtils {
 
   private SQLIndexFilterUtils() {
@@ -90,11 +92,12 @@ public static String parseSortCriteria(@Nonnull String entityType, @Nullable Ind
   /**
    * Parse {@link IndexFilter} into MySQL syntax.
    * @param entityType entity type from the Urn
-   * @param indexFilter                    index filter
+   * @param indexFilter index filter
    * @param nonDollarVirtualColumnsEnabled whether to enable non-dollar virtual columns
    * @return translated SQL condition expression, e.g. WHERE ...
    */
-  public static String parseIndexFilter(@Nonnull String entityType, @Nullable IndexFilter indexFilter, boolean nonDollarVirtualColumnsEnabled) {
+  public static String parseIndexFilter(@Nonnull String entityType, @Nullable IndexFilter indexFilter,
+      boolean nonDollarVirtualColumnsEnabled, @Nonnull SchemaValidatorUtil schemaValidator) {
     List<String> sqlFilters = new ArrayList<>();
 
     if (indexFilter == null || !indexFilter.hasCriteria()) {
@@ -115,6 +118,12 @@ public static String parseIndexFilter(@Nonnull String entityType, @Nullable Inde
         validateConditionAndValue(indexCriterion);
         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
+        if (!schemaValidator.columnExists(tableName, indexColumn)) {
+          log.warn("Skipping filter: virtual column '{}' not found in table '{}'", indexColumn, tableName);
+          continue;
+        }
         sqlFilters.add(parseSqlFilter(indexColumn, condition, pathParams.getValue()));
       }
     }

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

@@ -21,6 +21,7 @@
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import javax.annotation.ParametersAreNonnullByDefault;
+import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang.StringEscapeUtils;
 import org.javatuples.Pair;
 
@@ -31,6 +32,7 @@
 /**
  * SQL statement util class to generate executable SQL query / execution statements.
  */
+@Slf4j
 public class SQLStatementUtils {
   private static final Escaper URN_ESCAPER = Escapers.builder()
       .addEscape('\'', "''")
@@ -112,8 +114,6 @@ public class SQLStatementUtils {
 
   private static final String INDEX_GROUP_BY_CRITERION = "SELECT count(*) as COUNT, %s FROM %s";
 
-  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_URN_EXIST_TEMPLATE = "SELECT urn FROM %s WHERE urn = '%s'";
 
   private static final String INSERT_LOCAL_RELATIONSHIP = "INSERT INTO %s (metadata, source, destination, source_type, "
@@ -302,9 +302,10 @@ public static <ASPECT extends RecordTemplate> String createAspectUpdateWithOptim
    * @param nonDollarVirtualColumnsEnabled  true if virtual column does not contain $, false otherwise
    * @return translated SQL where statement
    */
-  public static String createFilterSql(String entityType, @Nullable IndexFilter indexFilter, boolean hasTotalCount, boolean nonDollarVirtualColumnsEnabled) {
+  public static String createFilterSql(String entityType, @Nullable IndexFilter indexFilter, boolean hasTotalCount, boolean nonDollarVirtualColumnsEnabled,
+      @Nonnull SchemaValidatorUtil schemaValidator) {
     final String tableName = getTableName(entityType);
-    String whereClause = parseIndexFilter(entityType, indexFilter, nonDollarVirtualColumnsEnabled);
+    String whereClause = parseIndexFilter(entityType, indexFilter, nonDollarVirtualColumnsEnabled, schemaValidator);
     String totalCountSql = String.format("SELECT COUNT(urn) FROM %s %s", tableName, whereClause);
     StringBuilder sb = new StringBuilder();
 
@@ -328,23 +329,25 @@ public static String createFilterSql(String entityType, @Nullable IndexFilter in
    * @return translated group by SQL
    */
   public static String createGroupBySql(String entityType, @Nullable IndexFilter indexFilter,
-      @Nonnull IndexGroupByCriterion indexGroupByCriterion, boolean nonDollarVirtualColumnsEnabled) {
+      @Nonnull IndexGroupByCriterion indexGroupByCriterion, boolean nonDollarVirtualColumnsEnabled, @Nonnull SchemaValidatorUtil schemaValidator) {
     final String tableName = getTableName(entityType);
     final String columnName =
         getGeneratedColumnName(entityType, indexGroupByCriterion.getAspect(), indexGroupByCriterion.getPath(),
             nonDollarVirtualColumnsEnabled);
+    // Check if the column exists in the schema
+    if (!schemaValidator.columnExists(tableName, columnName)) {
+      log.warn("Skipping group-by: column '{}' not found in table '{}'", columnName, tableName);
+      return ""; // skip query generation
+    }
     StringBuilder sb = new StringBuilder();
     sb.append(String.format(INDEX_GROUP_BY_CRITERION, columnName, tableName));
     sb.append("\n");
-    sb.append(parseIndexFilter(entityType, indexFilter, nonDollarVirtualColumnsEnabled));
+    sb.append(parseIndexFilter(entityType, indexFilter, nonDollarVirtualColumnsEnabled, schemaValidator));
     sb.append("\nGROUP BY ");
     sb.append(columnName);
     return sb.toString();
   }
 
-  public static String getAllColumnForTable(String tableName) {
-    return String.format(SQL_GET_ALL_COLUMNS, tableName);
-  }
 
   /**
    * Create aspect browse SQL statement.

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

@@ -0,0 +1,130 @@
+package com.linkedin.metadata.dao.utils;
+
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.google.common.annotations.VisibleForTesting;
+import io.ebean.EbeanServer;
+import io.ebean.SqlRow;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nonnull;
+import lombok.extern.slf4j.Slf4j;
+
+
+/**
+ * Utility class for validating presence of columns and indexes in MySQL tables
+ * by querying information_schema. Uses Caffeine caches to reduce DB load and
+ * support eventual consistency of schema evolution.
+ */
+@Slf4j
+public class SchemaValidatorUtil {
+  private final EbeanServer server;
+
+  // Cache: tableName → Set of index names
+  // 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.
+  private final Cache<String, Set<String>> indexCache = Caffeine.newBuilder()
+      .expireAfterWrite(10, TimeUnit.MINUTES)
+      .maximumSize(1000)
+      .build();
+
+  // Cache: tableName → Set of column names
+  private final Cache<String, Set<String>> columnCache = 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'";
+
+  public SchemaValidatorUtil(EbeanServer server) {
+    this.server = server;
+  }
+
+  /**
+   * Clears all caches, including indexCache, columnCache, missingColumnCache, and missingIndexCache.
+   * Useful for testing.
+   */
+  @VisibleForTesting
+  void clearCaches() {
+    indexCache.invalidateAll();
+    columnCache.invalidateAll();
+  }
+
+
+  /**
+   * Checks whether the given column exists in the specified table.
+   *
+   * @param tableName  Table name
+   * @param columnName Column name
+   * @return true if column exists, false otherwise
+   */
+  public boolean columnExists(@Nonnull String tableName, @Nonnull String columnName) {
+    String lowerTable = tableName.toLowerCase();
+    String lowerColumn = columnName.toLowerCase();
+
+    Set<String> columns = columnCache.get(lowerTable, tbl -> {
+      log.info("Refreshing column cache for table '{}'", tbl);
+      return loadColumns(tbl);
+    });
+
+    return columns.contains(lowerColumn);
+  }
+
+  /**
+   * Checks whether the given index exists in the specified table.
+   *
+   * @param tableName Table name
+   * @param indexName Index name
+   * @return true if index exists, false otherwise
+   */
+  public boolean indexExists(@Nonnull String tableName, @Nonnull String indexName) {
+    String lowerTable = tableName.toLowerCase();
+    String lowerIndex = indexName.toLowerCase();
+
+    Set<String> indexes = indexCache.get(lowerTable, tbl -> {
+      log.info("Refreshing index cache for table '{}'", tbl);
+      return loadIndexes(tbl);
+    });
+
+    return indexes.contains(lowerIndex);
+  }
+
+  /**
+   * Loads all columns for the given table from information_schema.
+   *
+   * @param tableName Table to query
+   * @return Set of lowercase column names
+   */
+  private Set<String> loadColumns(String tableName) {
+    List<SqlRow> rows = server.createSqlQuery(String.format(SQL_GET_ALL_COLUMNS, tableName)).findList();
+    Set<String> columns = new HashSet<>();
+    for (SqlRow row : rows) {
+      columns.add(row.getString("COLUMN_NAME").toLowerCase());
+    }
+    return columns;
+  }
+
+  /**
+   * Loads all index names for the given table from information_schema.
+   *
+   * @param tableName Table to query
+   * @return Set of lowercase index names
+   */
+  private Set<String> loadIndexes(String tableName) {
+    List<SqlRow> rows = server.createSqlQuery(String.format(SQL_GET_ALL_INDEXES, tableName)).findList();
+    Set<String> indexes = new HashSet<>();
+    for (SqlRow row : rows) {
+      indexes.add(row.getString("INDEX_NAME").toLowerCase());
+    }
+    return indexes;
+  }
+
+}