Update HibernateEntityModule to allow for binding the same set

of entities to multiple datasources. GitOrigin-RevId: 391e2194a72b5c6b7c79982550a9163853c529da
cashapp · Jan 14, 2025 · 300c730 · 300c730
1 parent aa11bb4
commit 300c730
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 0 deletions.
diff --git a/misk-hibernate/api/misk-hibernate.api b/misk-hibernate/api/misk-hibernate.api
@@ -99,6 +99,8 @@ public abstract class misk/hibernate/HibernateEntityModule : misk/inject/KAbstra
 	protected final fun addEntities ([Lkotlin/reflect/KClass;)V
 	protected final fun addEntity (Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;)V
 	public static synthetic fun addEntity$default (Lmisk/hibernate/HibernateEntityModule;Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;ILjava/lang/Object;)V
+	protected final fun addHibernateEntities ([Lkotlin/reflect/KClass;)V
+	protected final fun addHibernateEntity (Lkotlin/reflect/KClass;)V
 	protected final fun bindListener (Lorg/hibernate/event/spi/EventType;Lmisk/hibernate/BindPolicy;)Lcom/google/inject/binder/LinkedBindingBuilder;
 	public static synthetic fun bindListener$default (Lmisk/hibernate/HibernateEntityModule;Lorg/hibernate/event/spi/EventType;Lmisk/hibernate/BindPolicy;ILjava/lang/Object;)Lcom/google/inject/binder/LinkedBindingBuilder;
 	protected fun configure ()V

diff --git a/misk-hibernate/src/main/kotlin/misk/hibernate/HibernateEntityModule.kt b/misk-hibernate/src/main/kotlin/misk/hibernate/HibernateEntityModule.kt
@@ -84,6 +84,26 @@ abstract class HibernateEntityModule(
     addEntity(T::class)
   }
 
+  /** Bind just the HibernateEntities using the provided qualifier */
+  protected fun addHibernateEntities(vararg entities: KClass<out DbEntity<*>>) {
+    entities.forEach { entity ->
+      addHibernateEntity(dbEntityClass = entity)
+    }
+  }
+
+  /** Bind just the HibernateEntity using the provided qualifier */
+  protected fun <T : DbEntity<T>> addHibernateEntity(
+    dbEntityClass: KClass<T>
+  ) {
+    multibind<HibernateEntity>(qualifier)
+      .toInstance(HibernateEntity(dbEntityClass))
+  }
+
+  /** Bind just the HibernateEntity using the provided qualifier */
+  protected inline fun <reified T : DbEntity<T>> addHibernateEntity() {
+    addHibernateEntity(T::class)
+  }
+
   /** Adds a DbEntity to Database-Query Admin Dashboard Tab with a dynamic query
    * (not a static Misk.Query)
    */

diff --git a/misk-hibernate/src/test/kotlin/misk/hibernate/HibernateEntityModuleTest.kt b/misk-hibernate/src/test/kotlin/misk/hibernate/HibernateEntityModuleTest.kt
@@ -27,6 +27,26 @@ internal class HibernateEntityModuleTest {
       .containsExactly(Square::class, Circle::class)
   }
 
+  @Test fun multipleDataSourcesSameEntities() {
+    val injector = Guice.createInjector(
+      object : HibernateEntityModule(Dinosaurs::class) {
+        override fun configureHibernate() {
+          addHibernateEntities(Triceratops::class, Stegosaurus::class)
+        }
+      },
+      object : HibernateEntityModule(Shapes::class) {
+        override fun configureHibernate() {
+          addHibernateEntities(Triceratops::class, Stegosaurus::class)
+        }
+      }
+    )
+
+    assertThat(injector.getSetOf(HibernateEntity::class, Dinosaurs::class).unwrap())
+      .containsExactly(Triceratops::class, Stegosaurus::class)
+    assertThat(injector.getSetOf(HibernateEntity::class, Shapes::class).unwrap())
+      .containsExactly(Triceratops::class, Stegosaurus::class)
+  }
+
   @Test fun multipleModulesSameDataSource() {
     val injector = Guice.createInjector(
       object : HibernateEntityModule(Dinosaurs::class) {