Added manual override in JpaDataStore to explicitly bind entities (#1114

)

* Added manual override in JpaDataStore to explicitly bind entities

* Fixed checkstyles

elide-datastore/elide-datastore-jpa/src/main/java/com/yahoo/elide/datastores/jpa/JpaDataStore.java

@@ -10,6 +10,9 @@
 import com.yahoo.elide.core.datastore.JPQLDataStore;
 import com.yahoo.elide.datastores.jpa.transaction.JpaTransaction;
 
+import java.util.HashSet;
+import java.util.Set;
+
 import javax.persistence.EntityManager;
 import javax.persistence.metamodel.EntityType;
 
@@ -19,15 +22,28 @@
 public class JpaDataStore implements JPQLDataStore {
     protected final EntityManagerSupplier entityManagerSupplier;
     protected final JpaTransactionSupplier transactionSupplier;
+    protected final Set<Class<?>> modelsToBind;
 
     public JpaDataStore(EntityManagerSupplier entityManagerSupplier,
-                        JpaTransactionSupplier transactionSupplier) {
+                        JpaTransactionSupplier transactionSupplier,
+                        Class<?> ... models) {
         this.entityManagerSupplier = entityManagerSupplier;
         this.transactionSupplier = transactionSupplier;
+        this.modelsToBind = new HashSet<>();
+        for (Class<?> model : models) {
+            modelsToBind.add(model);
+        }
     }
 
     @Override
     public void populateEntityDictionary(EntityDictionary dictionary) {
+        // If the user provided models, we'll manually add them and skip scanning for entities.
+        if (! modelsToBind.isEmpty()) {
+            modelsToBind.stream().forEach((model) -> bindEntityClass(model, dictionary));
+            return;
+        }
+
+        // Use the entities defined in the entity manager factory.
         for (EntityType type : entityManagerSupplier.get().getMetamodel().getEntities()) {
             try {
                 Class<?> mappedClass = type.getJavaType();

elide-datastore/elide-datastore-jpa/src/test/java/com/yahoo/elide/datastores/jpa/JpaDataStoreTest.java

@@ -30,7 +30,6 @@
 import javax.persistence.metamodel.EntityType;
 import javax.persistence.metamodel.Metamodel;
 
-
 public class JpaDataStoreTest {
     public static class TestGenerator implements JPQLPredicateGenerator {
         @Override
@@ -63,7 +62,6 @@ class Test {
 
         JpaDataStore store = new JpaDataStore(() -> { return managerMock; }, (unused) -> { return null; });
         EntityDictionary dictionary = new EntityDictionary(new HashMap<>());
-        dictionary.bindEntity(Test.class);
 
 
         try {
@@ -74,4 +72,29 @@ class Test {
             FilterTranslator.registerJPQLGenerator(Operator.IN, Test.class, "name", null);
         }
     }
+
+    @Test
+    public void verifyManualEntityBinding() {
+
+        @Include
+        @Entity
+        class Test {
+            @Id
+            private long id;
+
+            private String name;
+        }
+
+        Metamodel mockModel = mock(Metamodel.class);
+        when(mockModel.getEntities()).thenReturn(Sets.newHashSet());
+
+        EntityManager managerMock = mock(EntityManager.class);
+        when(managerMock.getMetamodel()).thenReturn(mockModel);
+
+        JpaDataStore store = new JpaDataStore(() -> { return managerMock; }, (unused) -> { return null; }, Test.class);
+        EntityDictionary dictionary = new EntityDictionary(new HashMap<>());
+        store.populateEntityDictionary(dictionary);
+
+        assertNotNull(dictionary.lookupBoundClass(Test.class));
+    }
 }