Refactor out hibernate query limits into dependency to speed up Refle…

…ctionQueryFactoryTest

misk-hibernate/src/main/kotlin/misk/hibernate/HibernateModule.kt

@@ -31,6 +31,11 @@ import kotlin.reflect.KClass
  * This also registers services to connect to the database ([SessionFactoryService]) and to verify
  * that the schema is up-to-date ([SchemaMigratorService]).
  */
+
+private const val MAX_MAX_ROWS = 10_000
+private const val ROW_COUNT_ERROR_LIMIT = 3000
+private const val ROW_COUNT_WARNING_LIMIT = 2000
+
 class HibernateModule(
   private val qualifier: KClass<out Annotation>,
   private val config: DataSourceConfig
@@ -131,6 +136,8 @@ class HibernateModule(
     }).asSingleton()
 
     bind<Query.Factory>().to<ReflectionQuery.Factory>()
+    bind<ReflectionQuery.QueryLimitsConfig>().toInstance(ReflectionQuery.QueryLimitsConfig(
+        MAX_MAX_ROWS, ROW_COUNT_ERROR_LIMIT, ROW_COUNT_WARNING_LIMIT))
 
     install(object : HibernateEntityModule(qualifier) {
       override fun configureHibernate() {

misk-hibernate/src/main/kotlin/misk/hibernate/ReflectionQueryFactory.kt

@@ -13,6 +13,7 @@ import java.lang.reflect.InvocationTargetException
 import java.lang.reflect.Method
 import java.lang.reflect.ParameterizedType
 import java.lang.reflect.Proxy
+import javax.inject.Inject
 import javax.inject.Singleton
 import javax.persistence.criteria.CriteriaBuilder
 import javax.persistence.criteria.Path
@@ -29,9 +30,6 @@ import kotlin.reflect.full.primaryConstructor
 import kotlin.reflect.jvm.javaMethod
 
 private val logger = getLogger<ReflectionQuery<*>>()
-private const val MAX_MAX_ROWS = 10_000
-private const val ROW_COUNT_ERROR_LIMIT = 3000
-private const val ROW_COUNT_WARNING_LIMIT = 2000
 
 /**
  * Implements the [Query] @[Constraint] methods using a dynamic proxy, and projections using
@@ -40,11 +38,12 @@ private const val ROW_COUNT_WARNING_LIMIT = 2000
 internal class ReflectionQuery<T : DbEntity<T>>(
   private val rootEntityType: KClass<T>,
   private val queryMethodHandlers: Map<Method, QueryMethodHandler>,
-  private val tracer: Tracer?
+  private val tracer: Tracer?,
+  private val queryLimitsConfig: QueryLimitsConfig
 ) : Query<T>, InvocationHandler {
   override var maxRows = -1
     set(value) {
-      require(value == -1 || (value in 1..MAX_MAX_ROWS)) { "out of range: $value" }
+      require(value == -1 || (value in 1..queryLimitsConfig.maxMaxRows)) { "out of range: $value" }
       field = value
     }
 
@@ -83,20 +82,20 @@ internal class ReflectionQuery<T : DbEntity<T>>(
     return when {
       !returnList -> 2 // Detect if the result wasn't actually unique.
       (maxRows != -1) -> maxRows
-      else -> MAX_MAX_ROWS + 1 // Detect if the result was truncated.
+      else -> queryLimitsConfig.maxMaxRows + 1 // Detect if the result was truncated.
     }
   }
 
   private fun checkRowCount(returnList: Boolean, rowCount: Int) {
     if (!returnList) {
       check(rowCount <= 1) { "query expected a unique result but was $rowCount" }
     } else if (maxRows == -1) {
-      if (rowCount > MAX_MAX_ROWS) {
+      if (rowCount > queryLimitsConfig.maxMaxRows) {
         throw IllegalStateException("query truncated at $rowCount rows")
-      } else if (rowCount > ROW_COUNT_ERROR_LIMIT) {
+      } else if (rowCount > queryLimitsConfig.rowCountErrorLimit) {
         logger.error("Unbounded query returned $rowCount rows. " +
             "(Specify maxRows to suppress this error)")
-      } else if (rowCount > ROW_COUNT_WARNING_LIMIT) {
+      } else if (rowCount > queryLimitsConfig.rowCountWarningLimit) {
         logger.warn("Unbounded query returned $rowCount rows. " +
             "(Specify maxRows to suppress this warning)")
       }
@@ -132,7 +131,8 @@ internal class ReflectionQuery<T : DbEntity<T>>(
   }
 
   @Singleton
-  internal class Factory : Query.Factory {
+  internal class Factory @Inject internal constructor(var queryLimitsConfig: QueryLimitsConfig)
+      : Query.Factory {
     @com.google.inject.Inject(optional = true) var tracer: Tracer? = null
 
     private val queryMethodHandlersCache = CacheBuilder.newBuilder()
@@ -153,7 +153,7 @@ internal class ReflectionQuery<T : DbEntity<T>>(
       return Proxy.newProxyInstance(
           classLoader,
           arrayOf<Class<*>>(queryClass.java),
-          ReflectionQuery(entityType.kotlin, queryMethodHandlers, tracer)
+          ReflectionQuery(entityType.kotlin, queryMethodHandlers, tracer, queryLimitsConfig)
       ) as T
     }
 
@@ -180,6 +180,12 @@ internal class ReflectionQuery<T : DbEntity<T>>(
     }
   }
 
+  data class QueryLimitsConfig(
+      var maxMaxRows: Int,
+      var rowCountErrorLimit: Int,
+      var rowCountWarningLimit: Int
+  )
+
   class SelectMethodHandler(
     private val returnList: Boolean,
     private val constructor: KFunction<*>?,

misk-hibernate/src/test/kotlin/misk/hibernate/ReflectionQueryFactoryTest.kt

@@ -16,8 +16,13 @@ class ReflectionQueryFactoryTest {
   @MiskTestModule
   val module = MoviesTestModule(disableCrossShardQueryDetector = true)
 
+  val maxMaxRows = 40
+  val rowCountErrorLimit = 30
+  val rowCountWarningLimit = 20
+  private val queryFactory = ReflectionQuery.Factory(ReflectionQuery.QueryLimitsConfig(
+      maxMaxRows, rowCountErrorLimit, rowCountWarningLimit))
+
   @Inject @Movies lateinit var transacter: Transacter
-  @Inject lateinit var queryFactory: Query.Factory
   @Inject lateinit var logCollector: LogCollector
 
   @Test
@@ -253,63 +258,63 @@ class ReflectionQueryFactoryTest {
    */
   @Test
   fun rowResultCountWarning() {
-    // 1900 rows is too small for the warning threshold (2000).
+    // 18 rows is too small for the warning threshold (20).
     transacter.transaction { session ->
-      for (i in 1..1900) {
+      for (i in 1..18) {
         session.save(DbMovie("Rocky $i", null))
       }
     }
     transacter.transaction { session ->
-      assertThat(queryFactory.newQuery<OperatorsMovieQuery>().list(session)).hasSize(1900)
+      assertThat(queryFactory.newQuery<OperatorsMovieQuery>().list(session)).hasSize(18)
     }
     assertThat(logCollector.takeMessages(loggerClass = ReflectionQuery::class)).isEmpty()
 
-    // 2100 rows logs a warning.
+    // 21 rows logs a warning.
     transacter.transaction { session ->
-      for (i in 1901..2100) {
+      for (i in 19..21) {
         session.save(DbMovie("Rocky $i", null))
       }
     }
     transacter.transaction { session ->
-      assertThat(queryFactory.newQuery<OperatorsMovieQuery>().list(session)).hasSize(2100)
+      assertThat(queryFactory.newQuery<OperatorsMovieQuery>().list(session)).hasSize(21)
     }
     assertThat(getOnlyElement(
         logCollector.takeMessages(loggerClass = ReflectionQuery::class, minLevel = Level.WARN))
-    ).startsWith("Unbounded query returned 2100 rows.")
+    ).startsWith("Unbounded query returned 21 rows.")
 
-    // 3100 rows logs an error.
+    // 31 rows logs an error.
     transacter.transaction { session ->
-      for (i in 2101..3100) {
+      for (i in 22..31) {
         session.save(DbMovie("Rocky $i", null))
       }
     }
     transacter.transaction { session ->
-      assertThat(queryFactory.newQuery<OperatorsMovieQuery>().list(session)).hasSize(3100)
+      assertThat(queryFactory.newQuery<OperatorsMovieQuery>().list(session)).hasSize(31)
     }
     assertThat(getOnlyElement(
         logCollector.takeMessages(loggerClass = ReflectionQuery::class, minLevel = Level.ERROR))
-    ).startsWith("Unbounded query returned 3100 rows.")
+    ).startsWith("Unbounded query returned 31 rows.")
 
     // An explicit max row count suppresses the warning.
     transacter.transaction { session ->
       assertThat(queryFactory.newQuery<OperatorsMovieQuery>()
-          .apply { maxRows = 3200 }
+          .apply { maxRows = 32 }
           .list(session))
-          .hasSize(3100)
+          .hasSize(31)
     }
     assertThat(logCollector.takeMessages(loggerClass = ReflectionQuery::class)).isEmpty()
 
-    // More than 10000 rows throws an exception.
+    // More than 40 rows throws an exception.
     transacter.transaction { session ->
-      for (i in 3101..10100) {
+      for (i in 32..41) {
         session.save(DbMovie("Rocky $i", null))
       }
     }
     assertThat(assertFailsWith<IllegalStateException> {
       transacter.transaction { session ->
         queryFactory.newQuery<OperatorsMovieQuery>().list(session)
       }
-    }).hasMessage("query truncated at 10001 rows")
+    }).hasMessage("query truncated at 41 rows")
   }
 
   @Test