Skip to content

Commit 2e2ea45

Browse files
AmitGoenkasvc-squareup-copybara
AmitGoenka
authored and
svc-squareup-copybara
committed
Update HibernateModule to allow skipping Hibernate and JDBC
HealthChecks. GitOrigin-RevId: 622685271e477b7a9dceecaf9644d8aee622c051
1 parent d22212a commit 2e2ea45

File tree

4 files changed

+98
-17
lines changed

4 files changed

+98
-17
lines changed

‎misk-hibernate/api/misk-hibernate.api‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,11 +124,14 @@ public final class misk/hibernate/HibernateExceptionLogLevelConfig : wisp/config
124124
public final class misk/hibernate/HibernateModule : misk/inject/KAbstractModule {
125125
public fun <init> (Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceClusterConfig;Lmisk/jdbc/DatabasePool;)V
126126
public synthetic fun <init> (Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceClusterConfig;Lmisk/jdbc/DatabasePool;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
127+
public fun <init> (Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceClusterConfig;Lmisk/jdbc/DatabasePool;Z)V
128+
public synthetic fun <init> (Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceClusterConfig;Lmisk/jdbc/DatabasePool;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V
127129
public fun <init> (Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;)V
128130
public fun <init> (Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lmisk/jdbc/DatabasePool;)V
129131
public fun <init> (Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lmisk/jdbc/DatabasePool;Lmisk/hibernate/HibernateExceptionLogLevelConfig;)V
130132
public fun <init> (Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lmisk/jdbc/DatabasePool;Lmisk/hibernate/HibernateExceptionLogLevelConfig;Z)V
131-
public synthetic fun <init> (Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lmisk/jdbc/DatabasePool;Lmisk/hibernate/HibernateExceptionLogLevelConfig;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V
133+
public fun <init> (Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lmisk/jdbc/DatabasePool;Lmisk/hibernate/HibernateExceptionLogLevelConfig;ZZ)V
134+
public synthetic fun <init> (Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lmisk/jdbc/DatabasePool;Lmisk/hibernate/HibernateExceptionLogLevelConfig;ZZILkotlin/jvm/internal/DefaultConstructorMarker;)V
132135
public fun <init> (Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lmisk/jdbc/DatabasePool;)V
133136
public synthetic fun <init> (Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lmisk/jdbc/DatabasePool;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
134137
public fun <init> (Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceConfig;Lmisk/jdbc/DatabasePool;Lmisk/hibernate/HibernateExceptionLogLevelConfig;)V

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

Lines changed: 40 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ class HibernateModule @JvmOverloads constructor(
5858
val databasePool: DatabasePool = RealDatabasePool,
5959
private val logLevelConfig: HibernateExceptionLogLevelConfig = HibernateExceptionLogLevelConfig(),
6060
private val jdbcModuleAlreadySetup: Boolean = false,
61+
private val installHealthChecks: Boolean = true,
6162
) : KAbstractModule() {
6263

6364
// Make sure Hibernate logs use slf4j. Otherwise, it will base its decision on the classpath and
@@ -89,14 +90,29 @@ class HibernateModule @JvmOverloads constructor(
8990
databasePool: DatabasePool = RealDatabasePool,
9091
) : this(qualifier, cluster.writer, readerQualifier, cluster.reader, databasePool)
9192

93+
constructor(
94+
qualifier: KClass<out Annotation>,
95+
readerQualifier: KClass<out Annotation>,
96+
cluster: DataSourceClusterConfig,
97+
databasePool: DatabasePool = RealDatabasePool,
98+
installHealthChecks: Boolean = true,
99+
) : this(
100+
qualifier = qualifier,
101+
config = cluster.writer,
102+
readerQualifier = readerQualifier,
103+
readerConfig = cluster.reader,
104+
databasePool = databasePool,
105+
installHealthChecks = installHealthChecks
106+
)
107+
92108
constructor(
93109
qualifier: KClass<out Annotation>,
94110
config: DataSourceConfig,
95111
databasePool: DatabasePool = RealDatabasePool,
96112
jdbcModuleAlreadySetup: Boolean,
97113
) : this(
98114
qualifier = qualifier,
99-
config = config,
115+
config = config,
100116
readerQualifier = null,
101117
readerConfig = null,
102118
databasePool = databasePool,
@@ -111,17 +127,26 @@ class HibernateModule @JvmOverloads constructor(
111127
}
112128

113129
if (!jdbcModuleAlreadySetup) {
114-
install(JdbcModule(qualifier, config, readerQualifier, readerConfig, databasePool))
130+
install(
131+
JdbcModule(
132+
qualifier = qualifier,
133+
config = config,
134+
readerQualifier = readerQualifier,
135+
readerConfig = readerConfig,
136+
databasePool = databasePool,
137+
installHealthCheck = installHealthChecks
138+
)
139+
)
115140
}
116141

117142
bind<Query.Factory>().to<ReflectionQuery.Factory>()
118143
bind<QueryLimitsConfig>()
119144
.toInstance(QueryLimitsConfig(MAX_MAX_ROWS, ROW_COUNT_ERROR_LIMIT, ROW_COUNT_WARNING_LIMIT))
120145
bind<HibernateExceptionLogLevelConfig>().toInstance(logLevelConfig)
121146

122-
bindDataSource(qualifier, config, true)
123-
if (readerQualifier != null && readerConfig != null) {
124-
bindDataSource(readerQualifier, readerConfig, false)
147+
bindDataSource(qualifier, true)
148+
if (readerQualifier != null) {
149+
bindDataSource(readerQualifier, false)
125150
}
126151

127152
newMultibinder<DataSourceDecorator>(qualifier)
@@ -210,7 +235,6 @@ class HibernateModule @JvmOverloads constructor(
210235

211236
private fun bindDataSource(
212237
qualifier: KClass<out Annotation>,
213-
config: DataSourceConfig,
214238
isWriter: Boolean,
215239
) {
216240
// These items are configured on the writer qualifier only
@@ -251,14 +275,16 @@ class HibernateModule @JvmOverloads constructor(
251275
)
252276
}
253277

254-
val healthCheckKey = keyOf<HealthCheck>(qualifier)
255-
bind(healthCheckKey)
256-
.toProvider(object : Provider<HibernateHealthCheck> {
257-
@Inject lateinit var clock: Clock
278+
if (this.installHealthChecks) {
279+
val healthCheckKey = keyOf<HealthCheck>(qualifier)
280+
bind(healthCheckKey)
281+
.toProvider(object : Provider<HibernateHealthCheck> {
282+
@Inject lateinit var clock: Clock
258283

259-
override fun get() = HibernateHealthCheck(qualifier, sessionFactoryServiceProvider, clock)
260-
})
261-
.asSingleton()
262-
multibind<HealthCheck>().to(healthCheckKey)
284+
override fun get() = HibernateHealthCheck(qualifier, sessionFactoryServiceProvider, clock)
285+
})
286+
.asSingleton()
287+
multibind<HealthCheck>().to(healthCheckKey)
288+
}
263289
}
264290
}

‎misk-hibernate/src/test/kotlin/misk/hibernate/DisabledHealthCheckTest.kt‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
package misk.hibernate
2+
3+
import com.google.inject.Provider
4+
import jakarta.inject.Inject
5+
import misk.healthchecks.HealthCheck
6+
import misk.jdbc.SchemaMigratorService
7+
import misk.testing.MiskTest
8+
import misk.testing.MiskTestModule
9+
import org.assertj.core.api.Assertions.assertThat
10+
import org.junit.jupiter.api.Test
11+
import wisp.logging.getLogger
12+
import wisp.time.FakeClock
13+
import java.sql.Timestamp
14+
15+
@MiskTest(startService = true)
16+
class DisabledHealthCheckTest {
17+
@MiskTestModule
18+
val module = MoviesTestModule(installHealthChecks = false)
19+
20+
@Inject @Movies private lateinit var sessionFactoryService: Provider<SessionFactoryService>
21+
@Inject private lateinit var fakeClock: FakeClock
22+
@Inject private lateinit var healthChecks: List<HealthCheck>
23+
24+
@Test
25+
fun isNotInjected() {
26+
assertThat(healthChecks).noneMatch { it is HibernateHealthCheck }
27+
assertThat(healthChecks).noneMatch { it is SchemaMigratorService }
28+
}
29+
30+
@Test
31+
fun isStillConnectedToDb() {
32+
val databaseInstant = try {
33+
val sessionFactory = sessionFactoryService.get().sessionFactory
34+
sessionFactory.openSession().use { session ->
35+
session.createNativeQuery("SELECT NOW()").uniqueResult() as Timestamp
36+
}.toInstant()
37+
} catch (e: Exception) {
38+
logger.error(e) { "error performing hibernate health check" }
39+
null
40+
}
41+
42+
assertThat(databaseInstant).isNotNull()
43+
assertThat(databaseInstant).isAfterOrEqualTo(fakeClock.instant())
44+
}
45+
46+
companion object {
47+
val logger = getLogger<DisabledHealthCheckTest>()
48+
}
49+
}

‎misk-hibernate/src/test/kotlin/misk/hibernate/MoviesTestModule.kt‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ class MoviesTestModule(
2323
addEntities(DbMovie::class, DbActor::class, DbCharacter::class)
2424
}
2525
},
26+
private val installHealthChecks: Boolean = true,
2627
) : KAbstractModule() {
2728
override fun configure() {
2829
install(LogCollectorModule())
@@ -45,8 +46,10 @@ class MoviesTestModule(
4546
)
4647
install(
4748
HibernateModule(
48-
Movies::class, MoviesReader::class,
49-
DataSourceClusterConfig(writer = dataSourceConfig, reader = dataSourceConfig)
49+
qualifier = Movies::class,
50+
readerQualifier = MoviesReader::class,
51+
cluster = DataSourceClusterConfig(writer = dataSourceConfig, reader = dataSourceConfig),
52+
installHealthChecks = installHealthChecks
5053
)
5154
)
5255
install(entitiesModule)

0 commit comments

Comments
 (0)