linkedin · rakhiagr · Feb 16, 2024
diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java
@@ -67,7 +67,6 @@ public class EbeanLocalAccess<URN extends Urn> implements IEbeanLocalAccess<URN>
   private final EbeanLocalRelationshipWriterDAO _localRelationshipWriterDAO;
   private LocalRelationshipBuilderRegistry _localRelationshipBuilderRegistry;
   private final SchemaEvolutionManager _schemaEvolutionManager;
-  private final boolean _nonDollarVirtualColumnsEnabled;
 
   // TODO confirm if the default page size is 1000 in other code context.
   private static final int DEFAULT_PAGE_SIZE = 1000;
@@ -86,7 +85,7 @@ public EbeanLocalAccess(EbeanServer server, ServerConfig serverConfig, @Nonnull
     _entityType = ModelUtils.getEntityTypeFromUrnClass(_urnClass);
     _localRelationshipWriterDAO = new EbeanLocalRelationshipWriterDAO(_server);
     _schemaEvolutionManager = createSchemaEvolutionManager(serverConfig);
-    _nonDollarVirtualColumnsEnabled = nonDollarVirtualColumnsEnabled;
+    SQLSchemaUtils.setNonDollarVirtualColumnsEnabled(nonDollarVirtualColumnsEnabled);
   }
 
   public void setUrnPathExtractor(@Nonnull UrnPathExtractor<URN> urnPathExtractor) {
@@ -339,15 +338,15 @@ public <ASPECT extends RecordTemplate> ListResult<ASPECT> list(@Nonnull Class<AS
   public Map<String, Long> countAggregate(@Nullable IndexFilter indexFilter,
       @Nonnull IndexGroupByCriterion indexGroupByCriterion) {
     final String tableName = SQLSchemaUtils.getTableName(_entityType);
-    final String groupByColumn = getGeneratedColumnName(indexGroupByCriterion.getAspect(), indexGroupByCriterion.getPath(), _nonDollarVirtualColumnsEnabled);
+    final String groupByColumn = getGeneratedColumnName(indexGroupByCriterion.getAspect(), indexGroupByCriterion.getPath());
     // first, check for existence of the column we want to GROUP BY
     if (!checkColumnExists(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(tableName, indexFilter, indexGroupByCriterion, _nonDollarVirtualColumnsEnabled);
+    final String groupBySql = SQLStatementUtils.createGroupBySql(tableName, indexFilter, indexGroupByCriterion);
     final SqlQuery sqlQuery = _server.createSqlQuery(groupBySql);
     final List<SqlRow> sqlRows = sqlQuery.findList();
     Map<String, Long> resultMap = new HashMap<>();
@@ -376,9 +375,9 @@ private SqlQuery createFilterSqlQuery(@Nullable IndexFilter indexFilter,
 
     final String tableName = SQLSchemaUtils.getTableName(_entityType);
     StringBuilder filterSql = new StringBuilder();
-    filterSql.append(SQLStatementUtils.createFilterSql(tableName, indexFilter, true, _nonDollarVirtualColumnsEnabled));
+    filterSql.append(SQLStatementUtils.createFilterSql(tableName, indexFilter, true));
     filterSql.append("\n");
-    filterSql.append(parseSortCriteria(indexSortCriterion, _nonDollarVirtualColumnsEnabled));
+    filterSql.append(parseSortCriteria(indexSortCriterion));
     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());
@@ -391,7 +390,7 @@ private SqlQuery createFilterSqlQuery(@Nullable IndexFilter indexFilter,
       @Nullable IndexSortCriterion indexSortCriterion, @Nullable URN lastUrn, int pageSize) {
     StringBuilder filterSql = new StringBuilder();
     final String tableName = SQLSchemaUtils.getTableName(_entityType);
-    filterSql.append(SQLStatementUtils.createFilterSql(tableName, indexFilter, false, _nonDollarVirtualColumnsEnabled));
+    filterSql.append(SQLStatementUtils.createFilterSql(tableName, indexFilter, false));
 
     if (lastUrn != null) {
       // because createFilterSql will only include a WHERE clause if there are non-urn filters, we need to make sure
@@ -408,7 +407,7 @@ private SqlQuery createFilterSqlQuery(@Nullable IndexFilter indexFilter,
     }
 
     filterSql.append("\n");
-    filterSql.append(parseSortCriteria(indexSortCriterion, _nonDollarVirtualColumnsEnabled));
+    filterSql.append(parseSortCriteria(indexSortCriterion));
     filterSql.append(String.format(" LIMIT %d", Math.max(pageSize, 0)));
     return _server.createSqlQuery(filterSql.toString());
   }

diff --git a/...mpl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalRelationshipQueryDAO.java b/...mpl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalRelationshipQueryDAO.java
@@ -34,17 +34,14 @@ public class EbeanLocalRelationshipQueryDAO {
   private final EbeanServer _server;
   private final MultiHopsTraversalSqlGenerator _sqlGenerator;
 
-  private final EBeanDAOConfig _eBeanDAOConfig;
-
   public EbeanLocalRelationshipQueryDAO(EbeanServer server, EBeanDAOConfig eBeanDAOConfig) {
     _server = server;
-    _eBeanDAOConfig = eBeanDAOConfig;
+    SQLSchemaUtils.setNonDollarVirtualColumnsEnabled(eBeanDAOConfig.isNonDollarVirtualColumnsEnabled());
     _sqlGenerator = new MultiHopsTraversalSqlGenerator(SUPPORTED_CONDITIONS);
   }
 
   public EbeanLocalRelationshipQueryDAO(EbeanServer server) {
     _server = server;
-    _eBeanDAOConfig = new EBeanDAOConfig();
     _sqlGenerator = new MultiHopsTraversalSqlGenerator(SUPPORTED_CONDITIONS);
   }
 
@@ -79,8 +76,7 @@ public <SNAPSHOT extends RecordTemplate> List<SNAPSHOT> findEntities(@Nonnull Cl
     final StringBuilder sqlBuilder = new StringBuilder();
     sqlBuilder.append("SELECT * FROM ").append(tableName);
     if (filter.hasCriteria() && filter.getCriteria().size() > 0) {
-      sqlBuilder.append(" WHERE ").append(SQLStatementUtils.whereClause(filter, SUPPORTED_CONDITIONS, null,
-          _eBeanDAOConfig.isNonDollarVirtualColumnsEnabled()));
+      sqlBuilder.append(" WHERE ").append(SQLStatementUtils.whereClause(filter, SUPPORTED_CONDITIONS, null));
     }
     sqlBuilder.append(" ORDER BY urn LIMIT ").append(Math.max(1, count)).append(" OFFSET ").append(Math.max(0, offset));
 
@@ -105,8 +101,7 @@ public <SRC_SNAPSHOT extends RecordTemplate, DEST_SNAPSHOT extends RecordTemplat
     final String srcEntityTable = SQLSchemaUtils.getTableName(ModelUtils.getUrnTypeFromSnapshot(sourceEntityClass));
     final String destEntityTable = SQLSchemaUtils.getTableName(ModelUtils.getUrnTypeFromSnapshot(destinationEntityClass));
     final String sql = _sqlGenerator.multiHopTraversalSql(minHops, maxHops, Math.max(1, count), Math.max(0, offset), relationshipTable,
-        srcEntityTable, destEntityTable, relationshipFilter, sourceEntityFilter, destinationEntityFilter,
-        _eBeanDAOConfig.isNonDollarVirtualColumnsEnabled());
+        srcEntityTable, destEntityTable, relationshipFilter, sourceEntityFilter, destinationEntityFilter);
 
     final Class snapshotClass = relationshipFilter.getDirection() == RelationshipDirection.INCOMING ? sourceEntityClass : destinationEntityClass;
 
@@ -258,7 +253,6 @@ private String buildFindRelationshipSQL(@Nullable final String destTableName, @N
 
     sqlBuilder.append("WHERE deleted_ts is NULL");
     String whereClause = SQLStatementUtils.whereClause(SUPPORTED_CONDITIONS,
-        _eBeanDAOConfig.isNonDollarVirtualColumnsEnabled(),
         new Pair<>(sourceEntityFilter, "st"),
         new Pair<>(destinationEntityFilter, "dt"),
         new Pair<>(relationshipFilter, "rt"));

diff --git a/...ean-dao/src/main/java/com/linkedin/metadata/dao/utils/MultiHopsTraversalSqlGenerator.java b/...ean-dao/src/main/java/com/linkedin/metadata/dao/utils/MultiHopsTraversalSqlGenerator.java
@@ -27,7 +27,7 @@ public MultiHopsTraversalSqlGenerator(Map<Condition, String> supportedConditions
   @ParametersAreNonnullByDefault
   public String multiHopTraversalSql(int minHop, int maxHop, int count, int offset, String relationshipTable,
       String srcEntityTable, String destEntityTable, LocalRelationshipFilter relationshipFilter, LocalRelationshipFilter srcFilter,
-      LocalRelationshipFilter destFilter, boolean nonDollarVirtualColumnsEnabled) {
+      LocalRelationshipFilter destFilter) {
 
     /*
      * For now, only one-hop traversal is supported because multi-hops traversal using SQL is expensive
@@ -48,13 +48,13 @@ public String multiHopTraversalSql(int minHop, int maxHop, int count, int offset
     if (relationshipFilter.getDirection() == RelationshipDirection.INCOMING
         || relationshipFilter.getDirection() == RelationshipDirection.OUTGOING) {
       String urnSql = firstHopUrnsDirected(relationshipTable, srcEntityTable, destEntityTable, relationshipFilter, srcFilter,
-          destFilter, relationshipFilter.getDirection(), nonDollarVirtualColumnsEnabled);
+          destFilter, relationshipFilter.getDirection());
       return String.format("SELECT * FROM %s WHERE urn IN (%s) ORDER BY urn LIMIT %d OFFSET %d", entityTable, urnSql, count, offset);
     }
 
     // Relationship is undirected.
-    String urnSql = firstHopUrnsUndirected(relationshipTable, entityTable, relationshipFilter, srcFilter, nonDollarVirtualColumnsEnabled);
-    return findEntitiesUndirected(entityTable, relationshipTable, urnSql, destFilter, nonDollarVirtualColumnsEnabled);
+    String urnSql = firstHopUrnsUndirected(relationshipTable, entityTable, relationshipFilter, srcFilter);
+    return findEntitiesUndirected(entityTable, relationshipTable, urnSql, destFilter);
   }
 
   /**
@@ -64,8 +64,7 @@ public String multiHopTraversalSql(int minHop, int maxHop, int count, int offset
   @Nonnull
   @ParametersAreNonnullByDefault
   private String firstHopUrnsDirected(String relationshipTable, String srcEntityTable, String destEntityTable,
-      LocalRelationshipFilter relationshipFilter, LocalRelationshipFilter srcFilter, LocalRelationshipFilter destFilter, RelationshipDirection direction,
-      boolean nonDollarVirtualColumnsEnabled) {
+      LocalRelationshipFilter relationshipFilter, LocalRelationshipFilter srcFilter, LocalRelationshipFilter destFilter, RelationshipDirection direction) {
 
     String urnColumn = "destination";
     if (direction == RelationshipDirection.INCOMING) {
@@ -76,7 +75,7 @@ private String firstHopUrnsDirected(String relationshipTable, String srcEntityTa
         String.format("SELECT rt.%s FROM %s rt INNER JOIN %s dt ON rt.destination=dt.urn INNER JOIN %s st ON rt.source=st.urn WHERE rt.deleted_ts IS NULL",
             urnColumn, relationshipTable, destEntityTable, srcEntityTable));
 
-    String whereClause = SQLStatementUtils.whereClause(_supportedConditions,  nonDollarVirtualColumnsEnabled,
+    String whereClause = SQLStatementUtils.whereClause(_supportedConditions,
         new Pair<>(relationshipFilter, "rt"),
         new Pair<>(destFilter, "dt"),
         new Pair<>(srcFilter, "st"));
@@ -94,7 +93,7 @@ private String firstHopUrnsDirected(String relationshipTable, String srcEntityTa
   @Nonnull
   @ParametersAreNonnullByDefault
   private String firstHopUrnsUndirected(String relationshipTable, String entityTable, LocalRelationshipFilter relationshipFilter,
-      LocalRelationshipFilter srcFilter, boolean nonDollarVirtualColumnsEnabled) {
+      LocalRelationshipFilter srcFilter) {
 
     StringBuilder sourceUrnsSql = new StringBuilder(
         String.format("SELECT rt.source FROM %s rt INNER JOIN %s et ON rt.source=et.urn WHERE rt.deleted_ts IS NULL",
@@ -104,7 +103,7 @@ private String firstHopUrnsUndirected(String relationshipTable, String entityTab
         String.format("SELECT rt.destination FROM %s rt INNER JOIN %s et ON rt.destination=et.urn WHERE rt.deleted_ts IS NULL",
             relationshipTable, entityTable));
 
-    String whereClause = SQLStatementUtils.whereClause(_supportedConditions, nonDollarVirtualColumnsEnabled,
+    String whereClause = SQLStatementUtils.whereClause(_supportedConditions,
         new Pair<>(relationshipFilter, "rt"),
         new Pair<>(srcFilter, "et"));
 
@@ -121,9 +120,8 @@ private String firstHopUrnsUndirected(String relationshipTable, String entityTab
    */
   @Nonnull
   @ParametersAreNonnullByDefault
-  private String findEntitiesUndirected(String entityTable, String relationshipTable, String firstHopUrnSql, LocalRelationshipFilter destFilter,
-  boolean nonDollarVirtualColumnsEnabled) {
-    String whereClause = SQLStatementUtils.whereClause(_supportedConditions, nonDollarVirtualColumnsEnabled, new Pair<>(destFilter, "et"));
+  private String findEntitiesUndirected(String entityTable, String relationshipTable, String firstHopUrnSql, LocalRelationshipFilter destFilter) {
+    String whereClause = SQLStatementUtils.whereClause(_supportedConditions, 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)",

diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLIndexFilterUtils.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLIndexFilterUtils.java
@@ -66,17 +66,15 @@ private static String parseIndexValue(@Nullable IndexValue indexValue) {
   /**
    * Parse {@link IndexSortCriterion} into SQL syntax.
    * @param indexSortCriterion filter sorting criterion
-   * @param nonDollarVirtualColumnsEnabled  true if virtual column does not contain $, false otherwise
    * @return SQL statement of sorting, e.g. ORDER BY ... DESC ..etc.
    */
-  public static String parseSortCriteria(@Nullable IndexSortCriterion indexSortCriterion, boolean nonDollarVirtualColumnsEnabled) {
+  public static String parseSortCriteria(@Nullable IndexSortCriterion indexSortCriterion) {
     if (indexSortCriterion == null) {
       // Default to order by urn if user does not provide sort criterion.
       return "ORDER BY URN";
     }
     final String indexColumn =
-        SQLSchemaUtils.getGeneratedColumnName(indexSortCriterion.getAspect(), indexSortCriterion.getPath(),
-            nonDollarVirtualColumnsEnabled);
+        SQLSchemaUtils.getGeneratedColumnName(indexSortCriterion.getAspect(), indexSortCriterion.getPath());
 
     if (!indexSortCriterion.hasOrder()) {
       return "ORDER BY " + indexColumn;
@@ -89,10 +87,9 @@ public static String parseSortCriteria(@Nullable IndexSortCriterion indexSortCri
    * Parse {@link IndexFilter} into MySQL syntax.
    *
    * @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(@Nullable IndexFilter indexFilter, boolean nonDollarVirtualColumnsEnabled) {
+  public static String parseIndexFilter(@Nullable IndexFilter indexFilter) {
     List<String> sqlFilters = new ArrayList<>();
 
     if (indexFilter == null || !indexFilter.hasCriteria()) {
@@ -112,7 +109,7 @@ public static String parseIndexFilter(@Nullable IndexFilter indexFilter, boolean
       if (pathParams != null) {
         validateConditionAndValue(indexCriterion);
         final Condition condition = pathParams.getCondition();
-        final String indexColumn = getGeneratedColumnName(aspect, pathParams.getPath(), nonDollarVirtualColumnsEnabled);
+        final String indexColumn = getGeneratedColumnName(aspect, pathParams.getPath());
         sqlFilters.add(parseSqlFilter(indexColumn, condition, pathParams.getValue()));
       }
     }

diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLSchemaUtils.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLSchemaUtils.java
@@ -8,6 +8,8 @@
 import com.linkedin.metadata.dao.exception.MissingAnnotationException;
 import java.util.Map;
 import javax.annotation.Nonnull;
+import lombok.Getter;
+import lombok.Setter;
 
 
 /**
@@ -23,6 +25,10 @@ public class SQLSchemaUtils {
 
   private static final int MYSQL_MAX_COLUMN_NAME_LENGTH = 64 - ASPECT_PREFIX.length();
 
+  @Getter
+  @Setter
+  private static boolean nonDollarVirtualColumnsEnabled = false; // default value
+
   private SQLSchemaUtils() {
   }
 
@@ -88,13 +94,13 @@ public static <ASPECT extends RecordTemplate> String getAspectColumnName(@Nonnul
    * Get generated column name from aspect and path.
    */
   @Nonnull
-  public static String getGeneratedColumnName(@Nonnull String aspect, @Nonnull String path, boolean nonDollarVirtualColumnsEnabled) {
+  public static String getGeneratedColumnName(@Nonnull String aspect, @Nonnull String path) {
     char delimiter = nonDollarVirtualColumnsEnabled ? '0' : '$';
     if (isUrn(aspect)) {
-      return INDEX_PREFIX + "urn" + processPath(path, delimiter);
+      return INDEX_PREFIX + "urn" + processPath(path);
     }
 
-    return INDEX_PREFIX + getColumnNameFromAnnotation(aspect) + processPath(path, delimiter);
+    return INDEX_PREFIX + getColumnNameFromAnnotation(aspect) + processPath(path);
   }
 
   /**
@@ -107,11 +113,11 @@ public static boolean isUrn(@Nonnull String className) {
   /**
    * process 'path' into mysql column name convention.
    * @param path path in string e.g. /name/value, /name
-   * @param delimiter delimiter i.e '$' or '0'
    * @return $name$value or $name or 0name$value or 0name
    */
   @Nonnull
-  public static String processPath(@Nonnull String path, char delimiter) {
+  public static String processPath(@Nonnull String path) {
+    char delimiter = nonDollarVirtualColumnsEnabled ? '0' : '$';
     path = path.replace("/", String.valueOf(delimiter));
     if (!path.startsWith(String.valueOf(delimiter))) {
       path = delimiter + path;