diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java index 2639047f8..ba9a99f9f 100644 --- a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java +++ b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java @@ -79,7 +79,9 @@ public class EbeanLocalDAO protected final EbeanServer _server; protected final Class _urnClass; - private int _queryKeysCount = 0; // 0 means no pagination on keys + + private final static int DEFAULT_BATCH_SIZE = 50; + private int _queryKeysCount = DEFAULT_BATCH_SIZE; private IEbeanLocalAccess _localAccess; private UrnPathExtractor _urnPathExtractor; private SchemaConfig _schemaConfig = SchemaConfig.OLD_SCHEMA_ONLY; @@ -899,9 +901,12 @@ protected void applyTimeBasedRetention(@Nonnull if (keys.isEmpty()) { return Collections.emptyMap(); } - - final List records = batchGet(keys, keys.size()); - + final List records; + if (_queryKeysCount == 0) { + records = batchGet(keys, keys.size()); + } else { + records = batchGet(keys, _queryKeysCount); + } final Map, AspectWithExtraInfo> result = new HashMap<>(); keys.forEach(key -> records.stream() @@ -941,13 +946,18 @@ public boolean exists(@Nonnull URN urn) { } /** - * Sets the max keys allowed for each single query. + * Sets the max keys allowed for each single query, not allowed more than the default batch size. */ public void setQueryKeysCount(int keysCount) { if (keysCount < 0) { throw new IllegalArgumentException("Query keys count must be non-negative: " + keysCount); + } else if (keysCount > DEFAULT_BATCH_SIZE) { + log.warn("Setting query keys count greater than " + DEFAULT_BATCH_SIZE + + " may cause performance issues. Defaulting to " + DEFAULT_BATCH_SIZE + "."); + _queryKeysCount = DEFAULT_BATCH_SIZE; + } else { + _queryKeysCount = keysCount; } - _queryKeysCount = keysCount; } /**