Revert "Splitting filter queries (#579)" (#580)

This reverts commit 5583423.

Author: rakhiagr

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

@@ -53,6 +53,7 @@
 import static com.linkedin.metadata.dao.utils.SQLSchemaUtils.*;
 import static com.linkedin.metadata.dao.utils.SQLStatementUtils.*;
 
+
 /**
  * EBeanLocalAccess provides model agnostic data access (read / write) to MySQL database.
  */
@@ -335,77 +336,23 @@ public int softDeleteAsset(@Nonnull URN urn, boolean isTestMode) {
   @Override
   public List<URN> listUrns(@Nullable IndexFilter indexFilter, @Nullable IndexSortCriterion indexSortCriterion,
       @Nullable URN lastUrn, int pageSize) {
-    SqlQuery sqlQuery = createSelectFilterSqlQuery(indexFilter, indexSortCriterion, lastUrn, pageSize);
+    SqlQuery sqlQuery = createFilterSqlQuery(indexFilter, indexSortCriterion, lastUrn, pageSize);
     final List<SqlRow> sqlRows = sqlQuery.setFirstRow(0).findList();
     return sqlRows.stream().map(sqlRow -> getUrn(sqlRow.getString("urn"), _urnClass)).collect(Collectors.toList());
   }
 
-  /**
-   * Retrieves a paginated list of URNs matching the provided filter and sort criteria.
-   * This method performs two queries:
-   * <ul>
-   *   <li>A count query to determine the total number of matching URNs.</li>
-   *   <li>A paginated query to fetch the URNs for the requested page.</li>
-   * </ul>
-   * The results are returned as a {@link ListResult} containing the page of URNs, total count, and pagination metadata.
-   *
-   * @param indexFilter The filter criteria to apply to the URNs (nullable).
-   * @param indexSortCriterion The sort criteria for ordering the URNs (nullable).
-   * @param start The starting offset (zero-based) of the page.
-   * @param pageSize The maximum number of URNs to return in the page.
-   * @return A {@link ListResult} containing the page of URNs, total count, and pagination metadata.
-   */
   @Override
   public ListResult<URN> listUrns(@Nullable IndexFilter indexFilter, @Nullable IndexSortCriterion indexSortCriterion,
       int start, int pageSize) {
-    // 1. Run the count query to get total count of matching rows
-    final SqlQuery countQuery = createCountSqlQuery(indexFilter);
-    final SqlRow countRow = countQuery.findOne();
-    final int totalCount = (countRow == null) ? 0 : countRow.getInteger("total_count");
-
-    // If there are no matching rows, return an empty result immediately
-    if (totalCount == 0) {
-      return toListResult(0, start, pageSize);
-    }
-
-    // 2. Run the paginated URN query for the current page
-    final SqlQuery pageQuery = createSelectFilterSqlQuery(indexFilter, indexSortCriterion, start, pageSize);
-    final List<SqlRow> sqlRows = pageQuery.findList();
-
-    // Map the SQL rows to URN objects
-    final List<URN> values =
-        sqlRows.stream().map(sqlRow -> getUrn(sqlRow.getString("urn"), _urnClass)).collect(Collectors.toList());
-
-    // 3. Build and return the ListResult with values, total count, and pagination metadata
-    int adjustedCount = resolveTotalCount(values.size(), totalCount, start, pageSize);
-    return toListResult(values, adjustedCount, start, pageSize);
-  }
-
-  /**
-   * Resolve totalCount to handle race conditions between count and data queries.
-   * 
-   * @param valuesSize Number of results returned from the data query
-   * @param totalCount Total count from the count query (potentially stale)
-   * @param start Starting offset for pagination
-   * @param pageSize Requested page size
-   * @return Resolved totalCount that accounts for potential concurrent insertions/deletions
-   */
-  protected int resolveTotalCount(int valuesSize, int totalCount, int start, int pageSize) {
-    // If requesting beyond available data with no results, preserve original totalCount
-    if (start >= totalCount && valuesSize == 0) {
-      return totalCount;
-    }
-    
-    if (valuesSize < pageSize && start + valuesSize < totalCount) {
-      // Deletion race condition detected: hit end early due to records being deleted between queries
-      // Adjust totalCount to reflect the actual end position
-      return start + valuesSize;
-    } else {
-      // Normal pagination OR insertion race condition:
-      // - Normal case: use original totalCount
-      // - Insertion case: ensure totalCount >= actual results returned to avoid undercount
-      return Math.max(totalCount, start + valuesSize);
+    final SqlQuery sqlQuery = createFilterSqlQuery(indexFilter, indexSortCriterion, start, pageSize);
+    final List<SqlRow> sqlRows = sqlQuery.findList();
+    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);
     }
+    final List<URN> values = sqlRows.stream().map(sqlRow -> getUrn(sqlRow.getString("urn"), _urnClass)).collect(Collectors.toList());
+    return toListResult(values, sqlRows, null, start, pageSize);
   }
 
   @Override
@@ -526,48 +473,29 @@ public Map<String, Long> countAggregate(@Nullable IndexFilter indexFilter,
   }
 
   /**
-   * Builds a {@link SqlQuery} for fetching a paginated list of URNs matching the provided filter and sort criteria.
-   * Generates a query like:
-   *   SELECT urn FROM table WHERE filters ORDER BY sort LIMIT pageSize OFFSET offset
-   * Uses the provided IndexFilter and IndexSortCriterion to construct the WHERE and ORDER BY clauses.
-   *
-   * @param indexFilter The filter criteria to apply (nullable).
-   * @param indexSortCriterion The sort criteria to apply (nullable).
-   * @param offset The starting offset for pagination.
-   * @param pageSize The maximum number of results to return.
-   * @return A parametrized SqlQuery for fetching the paginated URNs.
+   * Produce {@link SqlQuery} for list urn by offset (start) and limit (pageSize).
+   * @param indexFilter index filter conditions
+   * @param indexSortCriterion sorting criterion, default ACS
+   * @return SqlQuery a SQL query which can be executed by ebean server.
    */
-  private SqlQuery createSelectFilterSqlQuery(@Nullable IndexFilter indexFilter,
+  private SqlQuery createFilterSqlQuery(@Nullable IndexFilter indexFilter,
       @Nullable IndexSortCriterion indexSortCriterion, int offset, int pageSize) {
-    String filterSql =
-        SQLStatementUtils.createSelectFilterSql(_entityType, indexFilter, _nonDollarVirtualColumnsEnabled, validator)
-            + "\n" + parseSortCriteria(_entityType, indexSortCriterion, _nonDollarVirtualColumnsEnabled)
-            + String.format(" LIMIT %d", Math.max(pageSize, 0)) + String.format(" OFFSET %d", Math.max(offset, 0));
-    return _server.createSqlQuery(filterSql);
-  }
-
-  /**
-   * Builds a {@link SqlQuery} for counting the number of rows matching the provided filter.
-   * Generates a query like:
-   *   SELECT COUNT(urn) AS total_count FROM table WHERE filters ...
-   * Uses the provided IndexFilter and schema validator to construct the WHERE clause.
-   *
-   * @param indexFilter The filter criteria to apply (nullable).
-   * @return A parametrized SqlQuery for counting matching rows.
-   */
-  private SqlQuery createCountSqlQuery(@Nullable IndexFilter indexFilter) {
-    String countSql = SQLStatementUtils.createCountFilterSql(
-        _entityType, indexFilter, _nonDollarVirtualColumnsEnabled, validator);
-    return _server.createSqlQuery(countSql);
+    StringBuilder filterSql = new StringBuilder();
+    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)));
+    filterSql.append(String.format(" OFFSET %d", Math.max(offset, 0)));
+    return _server.createSqlQuery(filterSql.toString());
   }
 
   /**
    * Produce {@link SqlQuery} for list urns by last urn.
    */
-  private SqlQuery createSelectFilterSqlQuery(@Nullable IndexFilter indexFilter,
+  private SqlQuery createFilterSqlQuery(@Nullable IndexFilter indexFilter,
       @Nullable IndexSortCriterion indexSortCriterion, @Nullable URN lastUrn, int pageSize) {
     StringBuilder filterSql = new StringBuilder();
-    filterSql.append(SQLStatementUtils.createSelectFilterSql(_entityType, indexFilter, _nonDollarVirtualColumnsEnabled, validator));
+    filterSql.append(SQLStatementUtils.createFilterSql(_entityType, indexFilter, false, _nonDollarVirtualColumnsEnabled, validator));
 
     if (lastUrn != null) {
       // because createFilterSql will always include a WHERE clause to filter by deleted_ts is NULL
@@ -667,51 +595,6 @@ protected <T> ListResult<T> toListResult(@Nonnull List<T> values, @Nonnull List<
         .build();
   }
 
-  /**
-   * Convert values + totalCount into {@link ListResult}.
-   * This is used when totalCount comes from a separate COUNT query,
-   * not from an embedded _total_count in the SqlRow.
-   *
-   * @param values the page of results
-   * @param totalCount the total number of results matching the filter
-   * @param start starting offset
-   * @param pageSize number of rows in this page
-   * @param <T> type of query response
-   * @return {@link ListResult} containing results and pagination metadata
-   */
-  @Nonnull
-  protected <T> ListResult<T> toListResult(@Nonnull List<T> values,
-      int totalCount,
-      int start,
-      int pageSize) {
-    if (pageSize == 0) {
-      pageSize = DEFAULT_PAGE_SIZE;
-    }
-    final int totalPageCount = ceilDiv(totalCount, pageSize);
-
-    boolean hasNext;
-    int nextStart;
-
-    if (start + values.size() < totalCount) {
-      hasNext = true;
-      nextStart = start + values.size();
-    } else {
-      hasNext = false;
-      nextStart = ListResult.INVALID_NEXT_START;
-    }
-
-    return ListResult.<T>builder()
-        .values(values)
-        .metadata(null) // or construct ListResultMetadata if needed
-        .nextStart(nextStart)
-        .havingMore(hasNext)
-        .totalCount(totalCount)
-        .totalPageCount(totalPageCount)
-        .pageSize(pageSize)
-        .build();
-  }
-
-
   /**
    * Given an AuditedAspect object, serialize it into a json string in a format that will be saved in DB.
    * @param auditedAspect AuditedAspect object to be serialized

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

@@ -33,6 +33,7 @@
 import static com.linkedin.metadata.dao.utils.SQLIndexFilterUtils.*;
 import static com.linkedin.metadata.dao.utils.SQLSchemaUtils.*;
 
+
 /**
  * SQL statement util class to generate executable SQL query / execution statements.
  */
@@ -151,12 +152,15 @@ public class SQLStatementUtils {
   private static final String DELETE_BY_SOURCE_AND_ASPECT = "UPDATE %s SET deleted_ts=NOW() "
       + "WHERE source = :source AND (aspect = :aspect OR aspect = :pegasus_aspect) AND deleted_ts IS NULL";
 
-  private static final String SQL_COUNT_TEMPLATE =
-      "SELECT COUNT(urn) AS total_count FROM %s %s";
-
-  private static final String SQL_SELECT_URN_WHERE_TEMPLATE =
-      "SELECT urn FROM %s %s";
-
+  /**
+   *  Filter query has pagination params in the existing APIs. To accommodate this, we use subquery to include total result counts in the query response.
+   *  For example, we will build the following filter query statement:
+   *
+   *  <p>SELECT *, (SELECT COUNT(urn) FROM metadata_entity_foo WHERE i_aspectfoo$value >= 25\n"
+   *  AND i_aspectfoo$value < 50 AND JSON_EXTRACT(a_aspectfoo, '$.gma_deleted') IS NULL) as _total_count FROM metadata_entity_foo\n"
+   *  WHERE i_aspectfoo$value >= 25 AND i_aspectfoo$value < 50 AND JSON_EXTRACT(a_aspectfoo, '$.gma_deleted') IS NULL;
+   */
+  private static final String SQL_FILTER_TEMPLATE = "SELECT *, (%s) as _total_count FROM %s";
   private static final String SQL_BROWSE_ASPECT_TEMPLATE =
       String.format("SELECT urn, %%s, lastmodifiedon, lastmodifiedby, (SELECT COUNT(urn) FROM %%s) as _total_count "
           + "FROM %%s WHERE %s LIMIT %%d OFFSET %%d", SOFT_DELETED_CHECK);
@@ -321,6 +325,33 @@ public static <ASPECT extends RecordTemplate> String createAspectUpdateWithOptim
           columnName, columnName, columnName);
     }
   }
+
+  /**
+   * Create filter SQL statement.
+   * @param entityType entity type from urn
+   * @param indexFilter index filter
+   * @param hasTotalCount whether to calculate total count in SQL.
+   * @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,
+      @Nonnull SchemaValidatorUtil schemaValidator) {
+    final String tableName = getTableName(entityType);
+    String whereClause = parseIndexFilter(entityType, indexFilter, nonDollarVirtualColumnsEnabled, schemaValidator);
+    String totalCountSql = String.format("SELECT COUNT(urn) FROM %s %s", tableName, whereClause);
+    StringBuilder sb = new StringBuilder();
+
+    if (hasTotalCount) {
+      sb.append(String.format(SQL_FILTER_TEMPLATE, totalCountSql, tableName));
+    } else {
+      sb.append("SELECT urn FROM ").append(tableName);
+    }
+
+    sb.append("\n");
+    sb.append(whereClause);
+    return sb.toString();
+  }
+
   /**
    * Create index group by SQL statement.
    * @param entityType entity type
@@ -365,28 +396,6 @@ public static <ASPECT extends RecordTemplate> String createAspectBrowseSql(Strin
         Math.max(pageSize, 0), Math.max(offset, 0));
   }
 
-  /**
-   * Creates an SQL statement for fetching URNs from an entity table, applying the provided filter.
-   */
-  public static String createSelectFilterSql(String entityType, @Nullable IndexFilter indexFilter, boolean nonDollarVirtualColumnsEnabled,
-      @Nonnull SchemaValidatorUtil schemaValidator) {
-    final String tableName = getTableName(entityType);
-    String whereClause = parseIndexFilter(entityType, indexFilter, nonDollarVirtualColumnsEnabled, schemaValidator);
-    // Build select query
-    return String.format(SQL_SELECT_URN_WHERE_TEMPLATE, tableName, whereClause);
-  }
-
-  /**
-   * Creates an SQL statement for counting rows from an entity table, applying the provided filter.
-   */
-  public static String createCountFilterSql(String entityType, @Nullable IndexFilter indexFilter, boolean nonDollarVirtualColumnsEnabled,
-      @Nonnull SchemaValidatorUtil schemaValidator) {
-    final String tableName = getTableName(entityType);
-    String whereClause = parseIndexFilter(entityType, indexFilter, nonDollarVirtualColumnsEnabled, schemaValidator);
-    // Build count query
-    return String.format(SQL_COUNT_TEMPLATE, tableName, whereClause);
-  }
-
   /**
    * Generate the create SQL statement for inserting local relationships. There can be multiple relationships added in
    * a single statement. The SQL generated should look like the following, where N is the number of relationships to insert: