[Fix #2175]Adding count process intance and count user task queries

apache · Feb 24, 2025 · 1f2bad3 · 1f2bad3
1 parent 891401f
commit 1f2bad3
Show file tree

Hide file tree

Showing 10 changed files with 110 additions and 17 deletions.
diff --git a/...ndex-graphql/src/main/java/org/kie/kogito/index/graphql/AbstractGraphQLSchemaManager.java b/...ndex-graphql/src/main/java/org/kie/kogito/index/graphql/AbstractGraphQLSchemaManager.java
@@ -40,6 +40,7 @@
 import org.kie.kogito.index.model.UserTaskInstance;
 import org.kie.kogito.index.service.DataIndexServiceException;
 import org.kie.kogito.index.storage.DataIndexStorageService;
+import org.kie.kogito.index.storage.StorageServiceCapability;
 import org.kie.kogito.persistence.api.StorageFetcher;
 import org.kie.kogito.persistence.api.query.Query;
 import org.slf4j.Logger;
@@ -109,6 +110,23 @@ protected final void loadAdditionalMutations(TypeDefinitionRegistry typeRegistry
         mutations.stream().map(GraphQLMutationsProvider::registry).forEach(typeRegistry::merge);
     }
 
+    protected final void addCountQueries(TypeDefinitionRegistry typeRegistry) {
+        if (supportsCount()) {
+            typeRegistry.merge(loadSchemaDefinitionFile("count.schema.graphqls"));
+        }
+    }
+
+    protected final void addCountQueries(Builder builder) {
+        if (supportsCount()) {
+            builder.dataFetcher("CountProcessInstances", this::countProcessInstances);
+            builder.dataFetcher("CountUserTaskInstances", this::countUserTaskInstances);
+        }
+    }
+
+    private boolean supportsCount() {
+        return cacheService.capabilities().contains(StorageServiceCapability.COUNT);
+    }
+
     protected TypeDefinitionRegistry loadSchemaDefinitionFile(String fileName) {
         return CommonUtils.loadSchemaDefinitionFile(fileName);
     }
@@ -182,18 +200,17 @@ protected Collection<ProcessInstance> getProcessInstancesValues(DataFetchingEnvi
         return executeAdvancedQueryForCache(cacheService.getProcessInstanceStorage(), env);
     }
 
-    protected <K, T> List<T> executeAdvancedQueryForCache(StorageFetcher<K, T> cache, DataFetchingEnvironment env) {
-        Objects.requireNonNull(cache, "Cache not found");
-
-        String inputTypeName = ((GraphQLNamedType) env.getFieldDefinition().getArgument("where").getType()).getName();
-
-        Query<T> query = cache.query();
+    protected long countProcessInstances(DataFetchingEnvironment env) {
+        return executeCount(cacheService.getProcessInstanceStorage(), env);
+    }
 
-        Map<String, Object> where = env.getArgument("where");
-        query.filter(GraphQLQueryParserRegistry.get().getParser(inputTypeName).apply(where));
+    protected long countUserTaskInstances(DataFetchingEnvironment env) {
+        return executeCount(cacheService.getUserTaskInstanceStorage(), env);
+    }
 
+    protected <K, T> List<T> executeAdvancedQueryForCache(StorageFetcher<K, T> cache, DataFetchingEnvironment env) {
+        Query<T> query = setupQuery(cache, env);
         query.sort(new GraphQLQueryOrderByParser().apply(env));
-
         Map<String, Integer> pagination = env.getArgument("pagination");
         if (pagination != null) {
             Integer limit = pagination.get("limit");
@@ -205,10 +222,22 @@ protected <K, T> List<T> executeAdvancedQueryForCache(StorageFetcher<K, T> cache
                 query.offset(offset);
             }
         }
-
         return query.execute();
     }
 
+    protected <K, T> long executeCount(StorageFetcher<K, T> cache, DataFetchingEnvironment env) {
+        return setupQuery(cache, env).count();
+    }
+
+    private <K, T> Query<T> setupQuery(StorageFetcher<K, T> cache, DataFetchingEnvironment env) {
+        Objects.requireNonNull(cache, "Cache not found");
+        String inputTypeName = ((GraphQLNamedType) env.getFieldDefinition().getArgument("where").getType()).getName();
+        Query<T> query = cache.query();
+        Map<String, Object> where = env.getArgument("where");
+        query.filter(GraphQLQueryParserRegistry.get().getParser(inputTypeName).apply(where));
+        return query;
+    }
+
     protected Collection<UserTaskInstance> getUserTaskInstancesValues(DataFetchingEnvironment env) {
         return executeAdvancedQueryForCache(cacheService.getUserTaskInstanceStorage(), env);
     }

diff --git a/data-index/data-index-graphql/src/main/resources/count.schema.graphqls b/data-index/data-index-graphql/src/main/resources/count.schema.graphqls
@@ -0,0 +1,4 @@
+extend type Query {
+    CountProcessInstances(where: ProcessInstanceArgument): Int
+    CountUserTaskInstances(where: UserTaskInstanceArgument): Int
+}
diff --git a/...e-common/src/main/java/org/kie/kogito/index/service/graphql/GraphQLSchemaManagerImpl.java b/...e-common/src/main/java/org/kie/kogito/index/service/graphql/GraphQLSchemaManagerImpl.java
@@ -74,6 +74,7 @@ public GraphQLSchema createSchema() {
         TypeDefinitionRegistry typeDefinitionRegistry = new TypeDefinitionRegistry();
         typeDefinitionRegistry.merge(loadSchemaDefinitionFile("basic.schema.graphqls"));
         typeDefinitionRegistry.merge(loadSchemaDefinitionFile("domain.schema.graphqls"));
+        addCountQueries(typeDefinitionRegistry);
         loadAdditionalMutations(typeDefinitionRegistry);
 
         RuntimeWiring runtimeWiring = RuntimeWiring.newRuntimeWiring()
@@ -83,6 +84,7 @@ public GraphQLSchema createSchema() {
                     builder.dataFetcher("ProcessInstances", this::getProcessInstancesValues);
                     builder.dataFetcher("UserTaskInstances", this::getUserTaskInstancesValues);
                     builder.dataFetcher("Jobs", this::getJobsValues);
+                    addCountQueries(typeDefinitionRegistry);
                     return builder;
                 })
                 .type("Mutation", builder -> {

diff --git a/...index-storage-api/src/main/java/org/kie/kogito/index/storage/DataIndexStorageService.java b/...index-storage-api/src/main/java/org/kie/kogito/index/storage/DataIndexStorageService.java
@@ -18,6 +18,9 @@
  */
 package org.kie.kogito.index.storage;
 
+import java.util.EnumSet;
+import java.util.Set;
+
 import org.kie.kogito.index.model.Job;
 import org.kie.kogito.index.model.ProcessDefinition;
 import org.kie.kogito.index.model.ProcessDefinitionKey;
@@ -26,7 +29,6 @@
 import com.fasterxml.jackson.databind.node.ObjectNode;
 
 public interface DataIndexStorageService {
-
     Storage<ProcessDefinitionKey, ProcessDefinition> getProcessDefinitionStorage();
 
     ProcessInstanceStorage getProcessInstanceStorage();
@@ -40,4 +42,8 @@ public interface DataIndexStorageService {
     String getDomainModelCacheName(String processId);
 
     Storage<String, String> getProcessIdModelCache();
+
+    default Set<StorageServiceCapability> capabilities() {
+        return EnumSet.noneOf(StorageServiceCapability.class);
+    }
 }
diff --git a/...ndex-storage-api/src/main/java/org/kie/kogito/index/storage/StorageServiceCapability.java b/...ndex-storage-api/src/main/java/org/kie/kogito/index/storage/StorageServiceCapability.java
@@ -0,0 +1,5 @@
+package org.kie.kogito.index.storage;
+
+public enum StorageServiceCapability {
+    COUNT
+}
diff --git a/...jpa-common/src/main/java/org/kie/kogito/index/jpa/storage/JPADataIndexStorageService.java b/...jpa-common/src/main/java/org/kie/kogito/index/jpa/storage/JPADataIndexStorageService.java
@@ -18,11 +18,15 @@
  */
 package org.kie.kogito.index.jpa.storage;
 
+import java.util.EnumSet;
+import java.util.Set;
+
 import org.kie.kogito.index.model.Job;
 import org.kie.kogito.index.model.ProcessDefinition;
 import org.kie.kogito.index.model.ProcessDefinitionKey;
 import org.kie.kogito.index.storage.DataIndexStorageService;
 import org.kie.kogito.index.storage.ProcessInstanceStorage;
+import org.kie.kogito.index.storage.StorageServiceCapability;
 import org.kie.kogito.index.storage.UserTaskInstanceStorage;
 import org.kie.kogito.persistence.api.Storage;
 
@@ -80,4 +84,9 @@ public String getDomainModelCacheName(String processId) {
     public Storage<String, String> getProcessIdModelCache() {
         throw new UnsupportedOperationException("Generic String cache not available in JPA");
     }
+
+    @Override
+    public Set<StorageServiceCapability> capabilities() {
+        return EnumSet.of(StorageServiceCapability.COUNT);
+    }
 }
diff --git a/...ata-index-storage-jpa-common/src/main/java/org/kie/kogito/index/jpa/storage/JPAQuery.java b/...ata-index-storage-jpa-common/src/main/java/org/kie/kogito/index/jpa/storage/JPAQuery.java
@@ -86,10 +86,7 @@ public List<T> execute() {
         CriteriaBuilder builder = repository.getEntityManager().getCriteriaBuilder();
         CriteriaQuery<E> criteriaQuery = builder.createQuery(entityClass);
         Root<E> root = criteriaQuery.from(entityClass);
-        if (filters != null && !filters.isEmpty()) {
-            List<Predicate> predicates = getPredicates(builder, root);
-            criteriaQuery.where(predicates.toArray(new Predicate[] {}));
-        }
+        addWhere(builder, criteriaQuery, root);
         if (sortBy != null && !sortBy.isEmpty()) {
             List<Order> orderBy = sortBy.stream().map(f -> {
                 Path attributePath = getAttributePath(root, f.getAttribute());
@@ -99,7 +96,6 @@ public List<T> execute() {
         }
 
         jakarta.persistence.Query query = repository.getEntityManager().createQuery(criteriaQuery);
-
         if (limit != null) {
             query.setMaxResults(limit);
         }
@@ -195,4 +191,21 @@ private List<Predicate> getRecursivePredicate(AttributeFilter<?> filter, Root<E>
                 .collect(toList());
     }
 
+    @Override
+    public long count() {
+        CriteriaBuilder builder = repository.getEntityManager().getCriteriaBuilder();
+        CriteriaQuery<Long> criteriaQuery = builder.createQuery(Long.class);
+        Root<E> root = criteriaQuery.from(entityClass);
+        criteriaQuery.select(builder.count(root));
+        addWhere(builder, criteriaQuery, root);
+        return repository.getEntityManager().createQuery(criteriaQuery).getSingleResult();
+    }
+
+    private <V> void addWhere(CriteriaBuilder builder, CriteriaQuery<V> criteriaQuery, Root<E> root) {
+        if (filters != null && !filters.isEmpty()) {
+            List<Predicate> predicates = getPredicates(builder, root);
+            criteriaQuery.where(predicates.toArray(new Predicate[] {}));
+        }
+    }
+
 }
diff --git a/...on/src/test/java/org/kie/kogito/index/jpa/query/AbstractProcessInstanceEntityQueryIT.java b/...on/src/test/java/org/kie/kogito/index/jpa/query/AbstractProcessInstanceEntityQueryIT.java
@@ -18,11 +18,20 @@
  */
 package org.kie.kogito.index.jpa.query;
 
+import java.util.List;
+import java.util.UUID;
+
+import org.junit.jupiter.api.Test;
+import org.kie.kogito.event.process.ProcessInstanceVariableDataEvent;
 import org.kie.kogito.index.jpa.storage.ProcessInstanceEntityStorage;
+import org.kie.kogito.index.storage.ProcessInstanceStorage;
+import org.kie.kogito.index.test.TestUtils;
 import org.kie.kogito.index.test.query.AbstractProcessInstanceQueryIT;
 
 import jakarta.inject.Inject;
 
+import static org.assertj.core.api.Assertions.assertThat;
+
 public abstract class AbstractProcessInstanceEntityQueryIT extends AbstractProcessInstanceQueryIT {
 
     @Inject
@@ -32,4 +41,15 @@ public abstract class AbstractProcessInstanceEntityQueryIT extends AbstractProce
     public ProcessInstanceEntityStorage getStorage() {
         return storage;
     }
+
+    @Test
+    void testCount() {
+        String processId = "persons";
+        String processInstanceId = UUID.randomUUID().toString();
+        ProcessInstanceStorage storage = getStorage();
+        ProcessInstanceVariableDataEvent variableEvent = TestUtils.createProcessInstanceVariableEvent(processInstanceId, processId, "John", 28, false,
+                List.of("Super", "Astonishing", "TheRealThing"));
+        storage.indexVariable(variableEvent);
+        assertThat(storage.query().count()).isOne();
+    }
 }
diff --git a/...ntime/src/main/java/org/kie/kogito/index/addon/graphql/GraphQLAddonSchemaManagerImpl.java b/...ntime/src/main/java/org/kie/kogito/index/addon/graphql/GraphQLAddonSchemaManagerImpl.java
@@ -34,6 +34,7 @@ public class GraphQLAddonSchemaManagerImpl extends AbstractGraphQLSchemaManager
     public GraphQLSchema createSchema() {
         TypeDefinitionRegistry typeDefinitionRegistry = new TypeDefinitionRegistry();
         typeDefinitionRegistry.merge(loadSchemaDefinitionFile("basic.schema.graphqls"));
+        addCountQueries(typeDefinitionRegistry);
         loadAdditionalMutations(typeDefinitionRegistry);
 
         RuntimeWiring runtimeWiring = RuntimeWiring.newRuntimeWiring()
@@ -42,6 +43,7 @@ public GraphQLSchema createSchema() {
                     builder.dataFetcher("ProcessInstances", this::getProcessInstancesValues);
                     builder.dataFetcher("UserTaskInstances", this::getUserTaskInstancesValues);
                     builder.dataFetcher("Jobs", this::getJobsValues);
+                    addCountQueries(builder);
                     return builder;
                 })
                 .type("Mutation", builder -> {
@@ -86,5 +88,4 @@ public GraphQLSchema createSchema() {
         SchemaGenerator schemaGenerator = new SchemaGenerator();
         return schemaGenerator.makeExecutableSchema(typeDefinitionRegistry, runtimeWiring);
     }
-
 }
diff --git a/...ons/persistence-commons-api/src/main/java/org/kie/kogito/persistence/api/query/Query.java b/...ons/persistence-commons-api/src/main/java/org/kie/kogito/persistence/api/query/Query.java
@@ -31,4 +31,8 @@ public interface Query<T> {
     Query<T> sort(List<AttributeSort> sortBy);
 
     List<T> execute();
+
+    default long count() {
+        throw new UnsupportedOperationException();
+    }
 }