implement cancel function (#1367)

* implement cancel function * addressing comments * addressing comments * Future implementation * Future implementation * Future implementation * Future implementation * Future implementation * adding Future Task * adding Future Task * adding Future Task * fixing future implementation * fixing future implementation * fixing future implementation * fixing future implementation * fixing future implementation * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing issues * fixing bugs * fixing issues * fixing issues * fixing issues * fixing issues * fixing bugs * fixing bugs * fixing bugs * fixing bugs * fixing bugs * fixing bugs * fixing bugs * fixing bugs * ExecutionException * fixing issues * aggregation changes * aggregation changes * aggregation changes * fixing bugs * fixing bugs * fixing bugs * rebasing * adddressing comments * adddressing comments * addressing comments * addressing comments Co-authored-by: Ramsha Rao <[email protected]>
yahoo · Aug 7, 2020 · 4e62b2a · 4e62b2a
1 parent fa22727
commit 4e62b2a
Show file tree

Hide file tree

Showing 29 changed files with 161 additions and 44 deletions.
diff --git a/elide-core/src/main/java/com/yahoo/elide/core/DataStoreTransaction.java b/elide-core/src/main/java/com/yahoo/elide/core/DataStoreTransaction.java
@@ -268,4 +268,10 @@ default boolean supportsPagination(Class<?> entityClass, FilterExpression expres
      * @return UUID id
      */
     UUID getRequestId();
+
+    /**
+     * Cancel running transaction.
+     * Implementation must be thread-safe.
+     */
+    void cancel();
 }
diff --git a/...e-core/src/main/java/com/yahoo/elide/core/datastore/inmemory/HashMapStoreTransaction.java b/...e-core/src/main/java/com/yahoo/elide/core/datastore/inmemory/HashMapStoreTransaction.java
@@ -182,4 +182,9 @@ public boolean supportsPagination(Class<?> entityClass, FilterExpression express
     private boolean containsObject(Object obj) {
         return dataStore.get(obj.getClass()).containsValue(obj);
     }
+
+    @Override
+    public void cancel() {
+        //nothing to cancel in HashMap store transaction
+    }
 }
diff --git a/...-core/src/main/java/com/yahoo/elide/core/datastore/inmemory/InMemoryStoreTransaction.java b/...-core/src/main/java/com/yahoo/elide/core/datastore/inmemory/InMemoryStoreTransaction.java
@@ -468,4 +468,9 @@ private Pair<Optional<Pagination>, Optional<Pagination>> splitPagination(
             return Pair.of(pagination, Optional.empty());
         }
     }
+
+   @Override
+   public void cancel() {
+       tx.cancel();
+   }
 }
diff --git a/elide-core/src/main/java/com/yahoo/elide/core/datastore/wrapped/TransactionWrapper.java b/elide-core/src/main/java/com/yahoo/elide/core/datastore/wrapped/TransactionWrapper.java
@@ -125,4 +125,9 @@ public Iterable<Object> loadObjects(EntityProjection projection, RequestScope sc
     public void close() throws IOException {
         tx.close();
     }
+
+    @Override
+    public void cancel() {
+        tx.cancel();
+    }
 }
diff --git a/elide-core/src/test/java/com/yahoo/elide/core/DataStoreTransactionTest.java b/elide-core/src/test/java/com/yahoo/elide/core/DataStoreTransactionTest.java
@@ -160,4 +160,9 @@ public void commit(RequestScope scope) {
     public void createObject(Object entity, RequestScope scope) {
         // nothing
     }
+
+    @Override
+    public void cancel() {
+       //nothing
+    }
 }
diff --git a/elide-core/src/test/java/com/yahoo/elide/core/EntityDictionaryTest.java b/elide-core/src/test/java/com/yahoo/elide/core/EntityDictionaryTest.java
@@ -1005,7 +1005,6 @@ public void testGetBoundByVersion() {
         assertEquals(3, models.size());  //Also includes com.yahoo.elide inner classes from this file.
         assertTrue(models.contains(BookV2.class));
 
-
         models = getBoundClassesByVersion(NO_VERSION);
         assertEquals(14, models.size());
     }

diff --git a/...src/main/java/com/yahoo/elide/datastores/aggregation/AggregationDataStoreTransaction.java b/...src/main/java/com/yahoo/elide/datastores/aggregation/AggregationDataStoreTransaction.java
@@ -99,4 +99,9 @@ Query buildQuery(EntityProjection entityProjection, RequestScope scope) {
                 scope.getDictionary());
         return translator.getQuery();
     }
+
+    @Override
+    public void cancel() {
+        queryEngineTransaction.cancel();
+    }
 }
diff --git a/...tastore-aggregation/src/main/java/com/yahoo/elide/datastores/aggregation/QueryEngine.java b/...tastore-aggregation/src/main/java/com/yahoo/elide/datastores/aggregation/QueryEngine.java
@@ -27,7 +27,6 @@
 
 import java.util.Map;
 import java.util.stream.Collectors;
-
 /**
  * A {@link QueryEngine} is an abstraction that an AggregationDataStore leverages to run analytic queries (OLAP style)
  * against an underlying persistence layer.
@@ -157,6 +156,11 @@ private void populateMetaData(MetaDataStore metaDataStore) {
     public interface Transaction extends AutoCloseable {
         @Override
         void close();
+
+        /**
+         * Cancels running transaction
+         */
+        void cancel();
     }
 
     public abstract Transaction beginTransaction();

diff --git a/...src/main/java/com/yahoo/elide/datastores/aggregation/queryengines/sql/SQLQueryEngine.java b/...src/main/java/com/yahoo/elide/datastores/aggregation/queryengines/sql/SQLQueryEngine.java
@@ -49,6 +49,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import javax.persistence.EntityManager;
 import javax.persistence.EntityManagerFactory;
@@ -60,13 +61,14 @@
 @Slf4j
 public class SQLQueryEngine extends QueryEngine {
     private final EntityManagerFactory entityManagerFactory;
-
+    private final Consumer<EntityManager> transactionCancel;
     private final SQLReferenceTable referenceTable;
 
-    public SQLQueryEngine(MetaDataStore metaDataStore, EntityManagerFactory entityManagerFactory) {
+    public SQLQueryEngine(MetaDataStore metaDataStore, EntityManagerFactory eMFactory, Consumer<EntityManager> txC) {
         super(metaDataStore);
-        this.entityManagerFactory = entityManagerFactory;
+        this.entityManagerFactory = eMFactory;
         this.referenceTable = new SQLReferenceTable(metaDataStore);
+        this.transactionCancel = txC;
     }
 
     @Override
@@ -118,14 +120,17 @@ public MetricProjection constructMetricProjection(Metric metric,
     /**
      * State needed for SQLQueryEngine to execute queries.
      */
-    static class SqlTransaction implements QueryEngine.Transaction {
+    static class SqlTransaction implements QueryEngine.Transaction  {
 
         private final EntityManager entityManager;
         private final EntityTransaction transaction;
+        private final Consumer<EntityManager> transactionCancel;
+
+        SqlTransaction(EntityManagerFactory emf, Consumer<EntityManager> transactionCancel) {
 
-        SqlTransaction(EntityManagerFactory emf) {
             entityManager = emf.createEntityManager();
             transaction = entityManager.getTransaction();
+            this.transactionCancel = transactionCancel;
             if (!transaction.isActive()) {
                 transaction.begin();
             }
@@ -140,11 +145,17 @@ public void close() {
                 entityManager.close();
             }
         }
+
+        @Override
+        public void cancel() {
+            transactionCancel.accept(entityManager);
+        }
+
     }
 
     @Override
     public QueryEngine.Transaction beginTransaction() {
-        return new SqlTransaction(entityManagerFactory);
+        return new SqlTransaction(entityManagerFactory, transactionCancel);
     }
 
     @Override

diff --git a/.../yahoo/elide/datastores/aggregation/annotation/dimensionformula/DimensionFormulaTest.java b/.../yahoo/elide/datastores/aggregation/annotation/dimensionformula/DimensionFormulaTest.java
@@ -24,7 +24,7 @@ public void testReferenceLoop() {
 
         IllegalArgumentException exception = assertThrows(
                 IllegalArgumentException.class,
-                () -> new SQLQueryEngine(metaDataStore, null));
+                () -> new SQLQueryEngine(metaDataStore, null, null));
         assertEquals(
                 "Formula reference loop found: loop.playerLevel->loop.playerLevel",
                 exception.getMessage());
@@ -36,7 +36,7 @@ public void testJoinToLoop() {
 
         IllegalArgumentException exception = assertThrows(
                 IllegalArgumentException.class,
-                () -> new SQLQueryEngine(metaDataStore, null));
+                () -> new SQLQueryEngine(metaDataStore, null, null));
         assertEquals(
                 "Formula reference loop found: joinToLoop.playerLevel->joinToLoop.playerLevel",
                 exception.getMessage());
@@ -49,7 +49,7 @@ public void testCrossClassReferenceLoop() {
 
         IllegalArgumentException exception = assertThrows(
                 IllegalArgumentException.class,
-                () -> new SQLQueryEngine(metaDataStore, null));
+                () -> new SQLQueryEngine(metaDataStore, null, null));
 
         String exception1 = "Formula reference loop found: loopCountryA.inUsa->loopCountryB.inUsa->loopCountryA.inUsa";
 

diff --git a/...va/com/yahoo/elide/datastores/aggregation/annotation/metricformula/MetricFormulaTest.java b/...va/com/yahoo/elide/datastores/aggregation/annotation/metricformula/MetricFormulaTest.java
@@ -21,7 +21,7 @@ public void testReferenceLoop() {
 
         IllegalArgumentException exception = assertThrows(
                 IllegalArgumentException.class,
-                () -> new SQLQueryEngine(metaDataStore, null));
+                () -> new SQLQueryEngine(metaDataStore, null, null));
         assertEquals(
                 "Formula reference loop found: loop.highScore->loop.highScore",
                 exception.getMessage());

diff --git a/...ava/com/yahoo/elide/datastores/aggregation/framework/AggregationDataStoreTestHarness.java b/...ava/com/yahoo/elide/datastores/aggregation/framework/AggregationDataStoreTestHarness.java
@@ -14,6 +14,11 @@
 import com.yahoo.elide.datastores.jpa.transaction.NonJtaTransaction;
 import com.yahoo.elide.datastores.multiplex.MultiplexManager;
 
+import org.hibernate.Session;
+
+import java.util.function.Consumer;
+
+import javax.persistence.EntityManager;
 import javax.persistence.EntityManagerFactory;
 
 public class AggregationDataStoreTestHarness implements DataStoreTestHarness {
@@ -26,14 +31,15 @@ public AggregationDataStoreTestHarness(EntityManagerFactory entityManagerFactory
     @Override
     public DataStore getDataStore() {
         MetaDataStore metaDataStore = new MetaDataStore();
+        Consumer<EntityManager> txCancel = (em) -> { em.unwrap(Session.class).cancelQuery(); };
 
         AggregationDataStore aggregationDataStore = AggregationDataStore.builder()
-                .queryEngine(new SQLQueryEngine(metaDataStore, entityManagerFactory))
+                .queryEngine(new SQLQueryEngine(metaDataStore, entityManagerFactory, txCancel))
                 .build();
 
         DataStore jpaStore = new JpaDataStore(
                 () -> entityManagerFactory.createEntityManager(),
-                NonJtaTransaction::new
+                (em) -> { return new NonJtaTransaction(em, txCancel); }
         );
 
         return new MultiplexManager(jpaStore, metaDataStore, aggregationDataStore);

diff --git a/...gregation/src/test/java/com/yahoo/elide/datastores/aggregation/framework/SQLUnitTest.java b/...gregation/src/test/java/com/yahoo/elide/datastores/aggregation/framework/SQLUnitTest.java
@@ -30,12 +30,14 @@
 import com.yahoo.elide.request.Argument;
 import com.yahoo.elide.utils.ClassScanner;
 
+import org.hibernate.Session;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
 
@@ -83,8 +85,9 @@ public static void init() {
         filterParser = new RSQLFilterDialect(dictionary);
 
         metaDataStore.populateEntityDictionary(dictionary);
+        Consumer<EntityManager> txCancel = (entityManager) -> { entityManager.unwrap(Session.class).cancelQuery(); };
+        engine = new SQLQueryEngine(metaDataStore, emf, txCancel);
 
-        engine = new SQLQueryEngine(metaDataStore, emf);
         playerStatsTable = engine.getTable("playerStats");
 
         ASIA.setName("Asia");

diff --git a/...rc/test/java/com/yahoo/elide/datastores/aggregation/queryengines/sql/QueryEngineTest.java b/...rc/test/java/com/yahoo/elide/datastores/aggregation/queryengines/sql/QueryEngineTest.java
@@ -46,7 +46,7 @@ public static void init() {
      * Test loading all three records from the table.
      */
     @Test
-    public void testFullTableLoad() {
+    public void testFullTableLoad() throws Exception {
         Query query = Query.builder()
                 .table(playerStatsTable)
                 .metric(invoke(playerStatsTable.getMetric("lowScore")))
@@ -84,7 +84,7 @@ public void testFullTableLoad() {
      * Test loading records using {@link FromSubquery}
      */
     @Test
-    public void testFromSubQuery() {
+    public void testFromSubQuery() throws Exception {
         Query query = Query.builder()
                 .table(playerStatsViewTable)
                 .metric(invoke(playerStatsViewTable.getMetric("highScore")))
@@ -185,7 +185,7 @@ public void testNotProjectedFilter() throws Exception {
     }
 
     @Test
-    public void testSortAggregatedMetric() {
+    public void testSortAggregatedMetric() throws Exception {
         Map<String, Sorting.SortOrder> sortMap = new TreeMap<>();
         sortMap.put("lowScore", Sorting.SortOrder.desc);
 
@@ -217,7 +217,7 @@ public void testSortAggregatedMetric() {
      * Test sorting by dimension attribute which is not present in the query.
      */
     @Test
-    public void testSortJoin() {
+    public void testSortJoin() throws Exception {
         Map<String, Sorting.SortOrder> sortMap = new TreeMap<>();
         sortMap.put("playerName", Sorting.SortOrder.asc);
 
@@ -259,7 +259,7 @@ public void testSortJoin() {
      * Test pagination.
      */
     @Test
-    public void testPagination() {
+    public void testPagination() throws Exception {
         Query query = Query.builder()
                 .table(playerStatsTable)
                 .metric(invoke(playerStatsTable.getMetric("lowScore")))
@@ -349,7 +349,7 @@ public void testHavingClauseJoin() throws Exception {
      * Test sorting by two different columns-one metric and one dimension.
      */
     @Test
-    public void testSortByMultipleColumns() {
+    public void testSortByMultipleColumns() throws Exception {
         Map<String, Sorting.SortOrder> sortMap = new TreeMap<>();
         sortMap.put("lowScore", Sorting.SortOrder.desc);
         sortMap.put("playerName", Sorting.SortOrder.asc);
@@ -392,7 +392,7 @@ public void testSortByMultipleColumns() {
      * Test grouping by a dimension with a JoinTo annotation.
      */
     @Test
-    public void testJoinToGroupBy() {
+    public void testJoinToGroupBy() throws Exception {
         Query query = Query.builder()
                 .table(playerStatsTable)
                 .metric(invoke(playerStatsTable.getMetric("highScore")))
@@ -452,7 +452,7 @@ public void testJoinToFilter() throws Exception {
      * Test grouping by a dimension with a JoinTo annotation.
      */
     @Test
-    public void testJoinToSort() {
+    public void testJoinToSort() throws Exception {
         Map<String, Sorting.SortOrder> sortMap = new TreeMap<>();
         sortMap.put("countryIsoCode", Sorting.SortOrder.asc);
         sortMap.put("highScore", Sorting.SortOrder.asc);
@@ -495,7 +495,7 @@ public void testJoinToSort() {
      * Test month grain query.
      */
     @Test
-    public void testTotalScoreByMonth() {
+    public void testTotalScoreByMonth() throws Exception {
         Query query = Query.builder()
                 .table(playerStatsTable)
                 .metric(invoke(playerStatsTable.getMetric("highScore")))
@@ -517,7 +517,7 @@ public void testTotalScoreByMonth() {
      * Test filter by time dimension.
      */
     @Test
-    public void testFilterByTemporalDimension() {
+    public void testFilterByTemporalDimension() throws Exception {
         FilterPredicate predicate = new FilterPredicate(
                 new Path(PlayerStats.class, dictionary, "recordedDate"),
                 Operator.IN,
@@ -542,7 +542,7 @@ public void testFilterByTemporalDimension() {
     }
 
     @Test
-    public void testAmbiguousFields() {
+    public void testAmbiguousFields() throws Exception {
         Map<String, Sorting.SortOrder> sortMap = new TreeMap<>();
         sortMap.put("lowScore", Sorting.SortOrder.asc);
 
@@ -581,7 +581,7 @@ public void testAmbiguousFields() {
     }
 
     @Test
-    public void testNullJoinToStringValue() {
+    public void testNullJoinToStringValue() throws Exception {
         Query query = Query.builder()
                 .table(playerStatsTable)
                 .metric(invoke(playerStatsTable.getMetric("highScore")))
@@ -606,7 +606,7 @@ public void testNullJoinToStringValue() {
     }
 
     @Test
-    public void testNullJoinToIntValue() {
+    public void testNullJoinToIntValue() throws Exception {
         Query query = Query.builder()
                 .table(playerStatsTable)
                 .metric(invoke(playerStatsTable.getMetric("highScore")))

diff --git a/...ation/src/test/java/com/yahoo/elide/datastores/aggregation/queryengines/sql/ViewTest.java b/...ation/src/test/java/com/yahoo/elide/datastores/aggregation/queryengines/sql/ViewTest.java
@@ -33,7 +33,7 @@ public static void init() {
     }
 
     @Test
-    public void testViewAttribute() {
+    public void testViewAttribute() throws Exception {
         Map<String, Sorting.SortOrder> sortMap = new TreeMap<>();
         sortMap.put("countryViewIsoCode", Sorting.SortOrder.desc);
 

diff --git a/...-hibernate3/src/main/java/com/yahoo/elide/datastores/hibernate3/HibernateTransaction.java b/...-hibernate3/src/main/java/com/yahoo/elide/datastores/hibernate3/HibernateTransaction.java
@@ -305,4 +305,9 @@ public Integer getQueryLimit() {
         // no limit
         return null;
     }
+
+    @Override
+    public void cancel() {
+        session.cancelQuery();
+    }
 }
diff --git a/...-hibernate5/src/main/java/com/yahoo/elide/datastores/hibernate5/HibernateTransaction.java b/...-hibernate5/src/main/java/com/yahoo/elide/datastores/hibernate5/HibernateTransaction.java
@@ -301,4 +301,9 @@ public void close() throws IOException {
             throw new IOException("Transaction not closed");
         }
     }
+
+    @Override
+    public void cancel() {
+        session.cancelQuery();
+    }
 }