CORDA-3209 Fix vault query for participants specified in common criteria (#5508)

* Generalise participant parsing code & additional test cases.

* Use a common predicate to expand the participants query (when specified more than once - eg. in fungible and linear query criteria).

* Introduce some re-usable functions.

* Additional code clean-up and improvements.

* Fix detekt MaxLineLength errors.

node/src/main/kotlin/net/corda/node/services/vault/HibernateQueryCriteriaParser.kt

@@ -27,6 +27,7 @@ import net.corda.node.services.persistence.NodeAttachmentService
 import org.hibernate.query.criteria.internal.expression.LiteralExpression
 import org.hibernate.query.criteria.internal.path.SingularAttributePath
 import org.hibernate.query.criteria.internal.predicate.ComparisonPredicate
+import org.hibernate.query.criteria.internal.predicate.CompoundPredicate
 import org.hibernate.query.criteria.internal.predicate.InPredicate
 import java.security.PublicKey
 import java.time.Instant
@@ -409,63 +410,74 @@ class HibernateQueryCriteriaParser(val contractStateType: Class<out ContractStat
         }
     }
 
+    private fun getVaultFungibleStateRoot(): Root<out Any> {
+        val entityStateClass = VaultSchemaV1.VaultFungibleStates::class.java
+        return rootEntities.getOrElse(entityStateClass) {
+            val entityRoot = criteriaQuery.from(entityStateClass)
+            rootEntities[entityStateClass] = entityRoot
+            entityRoot
+        }
+    }
+
+    private fun getVaultLinearStatesRoot(): Root<out Any> {
+        val entityStateClass = VaultSchemaV1.VaultLinearStates::class.java
+        return rootEntities.getOrElse(entityStateClass) {
+            val entityRoot = criteriaQuery.from(entityStateClass)
+            rootEntities[entityStateClass] = entityRoot
+            entityRoot
+        }
+    }
+
+    private fun getPersistentPartyRoot(): Root<out Any> {
+        val persistentPartyEntity = VaultSchemaV1.PersistentParty::class.java
+        return rootEntities.getOrElse(persistentPartyEntity) {
+            val entityRoot = criteriaQuery.from(persistentPartyEntity)
+            rootEntities[persistentPartyEntity] = entityRoot
+            entityRoot
+        }
+    }
+
     override fun parseCriteria(criteria: QueryCriteria.FungibleAssetQueryCriteria): Collection<Predicate> {
         log.trace { "Parsing FungibleAssetQueryCriteria: $criteria" }
 
         val predicateSet = mutableSetOf<Predicate>()
 
         // ensure we re-use any existing instance of the same root entity
-        val entityStateClass = VaultSchemaV1.VaultFungibleStates::class.java
+        val vaultFungibleStatesRoot = getVaultFungibleStateRoot()
-        val vaultFungibleStates =
+        val joinPredicate = criteriaBuilder.equal(vaultStates.get<PersistentStateRef>("stateRef"),
-                rootEntities.getOrElse(entityStateClass) {
+                vaultFungibleStatesRoot.get<PersistentStateRef>("stateRef"))
-                    val entityRoot = criteriaQuery.from(entityStateClass)
-                    rootEntities[entityStateClass] = entityRoot
-                    entityRoot
-                }
-
-        val joinPredicate = criteriaBuilder.equal(vaultStates.get<PersistentStateRef>("stateRef"), vaultFungibleStates.get<PersistentStateRef>("stateRef"))
         predicateSet.add(joinPredicate)
 
         // owner
         criteria.owner?.let {
             val owners = criteria.owner as List<AbstractParty>
-            predicateSet.add(criteriaBuilder.and(vaultFungibleStates.get<AbstractParty>("owner").`in`(owners)))
+            predicateSet.add(criteriaBuilder.and(vaultFungibleStatesRoot.get<AbstractParty>("owner").`in`(owners)))
         }
 
         // quantity
         criteria.quantity?.let {
-            predicateSet.add(columnPredicateToPredicate(vaultFungibleStates.get<Long>("quantity"), it))
+            predicateSet.add(columnPredicateToPredicate(vaultFungibleStatesRoot.get<Long>("quantity"), it))
         }
 
         // issuer party
         criteria.issuer?.let {
             val issuerParties = criteria.issuer as List<AbstractParty>
-            predicateSet.add(criteriaBuilder.and(vaultFungibleStates.get<AbstractParty>("issuer").`in`(issuerParties)))
+            predicateSet.add(criteriaBuilder.and(vaultFungibleStatesRoot.get<AbstractParty>("issuer").`in`(issuerParties)))
         }
 
         // issuer reference
         criteria.issuerRef?.let {
             val issuerRefs = (criteria.issuerRef as List<OpaqueBytes>).map { it.bytes }
-            predicateSet.add(criteriaBuilder.and(vaultFungibleStates.get<ByteArray>("issuerRef").`in`(issuerRefs)))
+            predicateSet.add(criteriaBuilder.and(vaultFungibleStatesRoot.get<ByteArray>("issuerRef").`in`(issuerRefs)))
         }
 
         // Participants.
         criteria.participants?.let {
-            val participants = criteria.participants!!
+            // Join VaultFungibleState and PersistentParty tables (participant values are added to the common query criteria predicate)
-
+            val statePartyToFungibleStatesJoin = criteriaBuilder.and(
-            // Get the persistent party entity.
+                criteriaBuilder.equal(vaultFungibleStatesRoot.get<VaultSchemaV1.VaultFungibleStates>("stateRef"),
-            val persistentPartyEntity = VaultSchemaV1.PersistentParty::class.java
+                    getPersistentPartyRoot().get<VaultSchemaV1.PersistentParty>("compositeKey").get<PersistentStateRef>("stateRef")))
-            val entityRoot = rootEntities.getOrElse(persistentPartyEntity) {
+            predicateSet.add(statePartyToFungibleStatesJoin)
-                val entityRoot = criteriaQuery.from(persistentPartyEntity)
-                rootEntities[persistentPartyEntity] = entityRoot
-                entityRoot
-            }
-
-            // Add the join and participants predicates.
-            val statePartyJoin = criteriaBuilder.equal(vaultStates.get<VaultSchemaV1.VaultStates>("stateRef"), entityRoot.get<VaultSchemaV1.PersistentParty>("compositeKey").get<PersistentStateRef>("stateRef"))
-            val participantsPredicate = criteriaBuilder.and(entityRoot.get<VaultSchemaV1.PersistentParty>("x500Name").`in`(participants))
-            predicateSet.add(statePartyJoin)
-            predicateSet.add(participantsPredicate)
         }
 
         return predicateSet
@@ -477,47 +489,32 @@ class HibernateQueryCriteriaParser(val contractStateType: Class<out ContractStat
         val predicateSet = mutableSetOf<Predicate>()
 
         // ensure we re-use any existing instance of the same root entity
-        val entityStateClass = VaultSchemaV1.VaultLinearStates::class.java
+        val vaultLinearStatesRoot = getVaultLinearStatesRoot()
-        val vaultLinearStates =
-                rootEntities.getOrElse(entityStateClass) {
-                    val entityRoot = criteriaQuery.from(entityStateClass)
-                    rootEntities[entityStateClass] = entityRoot
-                    entityRoot
-                }
 
-        val joinPredicate = criteriaBuilder.equal(vaultStates.get<PersistentStateRef>("stateRef"), vaultLinearStates.get<PersistentStateRef>("stateRef"))
+        val joinPredicate = criteriaBuilder.equal(vaultStates.get<PersistentStateRef>("stateRef"),
+                vaultLinearStatesRoot.get<PersistentStateRef>("stateRef"))
         predicateSet.add(joinPredicate)
 
         // linear ids UUID
         criteria.uuid?.let {
             val uuids = criteria.uuid as List<UUID>
-            predicateSet.add(criteriaBuilder.and(vaultLinearStates.get<UUID>("uuid").`in`(uuids)))
+            predicateSet.add(criteriaBuilder.and(vaultLinearStatesRoot.get<UUID>("uuid").`in`(uuids)))
         }
 
         // linear ids externalId
         criteria.externalId?.let {
             val externalIds = criteria.externalId as List<String>
             if (externalIds.isNotEmpty())
-                predicateSet.add(criteriaBuilder.and(vaultLinearStates.get<String>("externalId").`in`(externalIds)))
+                predicateSet.add(criteriaBuilder.and(vaultLinearStatesRoot.get<String>("externalId").`in`(externalIds)))
         }
 
         // Participants.
         criteria.participants?.let {
-            val participants = criteria.participants!!
+            // Join VaultLinearState and PersistentParty tables (participant values are added to the common query criteria predicate)
-
+            val statePartyToLinearStatesJoin = criteriaBuilder.and(
-            // Get the persistent party entity.
+                criteriaBuilder.equal(vaultLinearStatesRoot.get<VaultSchemaV1.VaultLinearStates>("stateRef"),
-            val persistentPartyEntity = VaultSchemaV1.PersistentParty::class.java
+                        getPersistentPartyRoot().get<VaultSchemaV1.PersistentParty>("compositeKey").get<PersistentStateRef>("stateRef")))
-            val entityRoot = rootEntities.getOrElse(persistentPartyEntity) {
+            predicateSet.add(statePartyToLinearStatesJoin)
-                val entityRoot = criteriaQuery.from(persistentPartyEntity)
-                rootEntities[persistentPartyEntity] = entityRoot
-                entityRoot
-            }
-
-            // Add the join and participants predicates.
-            val statePartyJoin = criteriaBuilder.equal(vaultStates.get<VaultSchemaV1.VaultStates>("stateRef"), entityRoot.get<VaultSchemaV1.PersistentParty>("compositeKey").get<PersistentStateRef>("stateRef"))
-            val participantsPredicate = criteriaBuilder.and(entityRoot.get<VaultSchemaV1.PersistentParty>("x500Name").`in`(participants))
-            predicateSet.add(statePartyJoin)
-            predicateSet.add(participantsPredicate)
         }
 
         return predicateSet
@@ -678,6 +675,33 @@ class HibernateQueryCriteriaParser(val contractStateType: Class<out ContractStat
             constraintPredicates.add(externalIdPredicate)
         }
 
+        // Participants.
+        criteria.participants?.let {
+            val participants = criteria.participants!!
+
+            // use a single predicate for querying the persistent party table (incrementally add additional parties)
+            val predicateID = Pair(VaultSchemaV1.PersistentParty::x500Name.name, EQUAL)
+            if (commonPredicates.containsKey(predicateID)) {
+                val existingParticipants = ((((commonPredicates[predicateID]) as CompoundPredicate).expressions[0]) as InPredicate<*>)
+                        .values.map { participant -> (participant as LiteralExpression<*>).literal }
+                log.warn("Adding new participants: $participants to existing participants: $existingParticipants")
+                commonPredicates.replace(predicateID, criteriaBuilder.and(
+                        getPersistentPartyRoot().get<VaultSchemaV1.PersistentParty>("x500Name").`in`(existingParticipants + participants)))
+            }
+            else {
+                // Get the persistent party entity.
+                commonPredicates[predicateID] = criteriaBuilder.and(
+                        getPersistentPartyRoot().get<VaultSchemaV1.PersistentParty>("x500Name").`in`(participants))
+            }
+
+            // Add the join for vault states to persistent entities (if this is not a Fungible nor Linear criteria query)
+            if (criteria !is QueryCriteria.FungibleAssetQueryCriteria && criteria !is QueryCriteria.LinearStateQueryCriteria ) {
+                val statePartyJoin = criteriaBuilder.equal(vaultStates.get<VaultSchemaV1.VaultStates>("stateRef"),
+                        getPersistentPartyRoot().get<VaultSchemaV1.PersistentParty>("compositeKey").get<PersistentStateRef>("stateRef"))
+                constraintPredicates.add(statePartyJoin)
+            }
+        }
+
         return emptySet()
     }
 

node/src/test/kotlin/net/corda/node/services/vault/VaultQueryTests.kt

@@ -229,6 +229,34 @@ abstract class VaultQueryTestsBase : VaultQueryParties {
      * Query API tests
      */
 
+    /** Generic Query tests: using CommonQueryCriteria */
+
+    @Test
+    fun `unconsumed base contract states for single participant`() {
+        database.transaction {
+            identitySvc.verifyAndRegisterIdentity(BIG_CORP_IDENTITY)
+            vaultFiller.fillWithDummyState(participants = listOf(MEGA_CORP, MINI_CORP))
+            vaultFiller.fillWithDummyState(participants = listOf(MEGA_CORP))
+            vaultFiller.fillWithDummyState(participants = listOf(MEGA_CORP, BIG_CORP))  // true
+            val criteria = VaultQueryCriteria(participants = listOf(BIG_CORP))
+            val results = vaultService.queryBy<ContractState>(criteria)
+            assertThat(results.states).hasSize(1)
+        }
+    }
+
+    @Test
+    fun `unconsumed base contract states for two participants`() {
+        database.transaction {
+            identitySvc.verifyAndRegisterIdentity(BIG_CORP_IDENTITY)
+            vaultFiller.fillWithDummyState(participants = listOf(MEGA_CORP, MINI_CORP)) // true
+            vaultFiller.fillWithDummyState(participants = listOf(MEGA_CORP, BIG_CORP))  // true
+            vaultFiller.fillWithDummyState(participants = listOf(MEGA_CORP))
+            val criteria = VaultQueryCriteria(participants = listOf(MINI_CORP, BIG_CORP))
+            val results = vaultService.queryBy<ContractState>(criteria)
+            assertThat(results.states).hasSize(2)
+        }
+    }
+
     /** Generic Query tests
     (combining both FungibleState and LinearState contract types) */
 
@@ -2379,8 +2407,8 @@ abstract class VaultQueryTestsBase : VaultQueryParties {
             identitySvc.verifyAndRegisterIdentity(BOB_IDENTITY)
             identitySvc.verifyAndRegisterIdentity(CHARLIE_IDENTITY)
             vaultFiller.fillWithSomeTestLinearStates(1, "TEST1", listOf(ALICE))
-            vaultFiller.fillWithSomeTestLinearStates(1,  "TEST2", listOf(BOB))
+            vaultFiller.fillWithSomeTestLinearStates(1, "TEST2", listOf(BOB))
-            vaultFiller.fillWithSomeTestLinearStates(1,  "TEST3", listOf(CHARLIE))
+            vaultFiller.fillWithSomeTestLinearStates(1, "TEST3", listOf(CHARLIE))
             vaultFiller.fillWithSomeTestCash(100.DOLLARS, notaryServices, 1, DUMMY_CASH_ISSUER)
             vaultFiller.fillWithSomeTestCommodity(Amount(100, Commodity.getInstance("FCOJ")!!), notaryServices, DUMMY_OBLIGATION_ISSUER.ref(1))
             vaultFiller.fillWithDummyState()
@@ -2397,7 +2425,11 @@ abstract class VaultQueryTestsBase : VaultQueryParties {
             assertThat(resultsFASC.states).hasSize(2)
             // composite query for both linear and fungible asset states by participants
             val resultsComposite = vaultService.queryBy<ContractState>(linearStateCriteria.or(fungibleAssetStateCriteria))
-            assertThat(resultsComposite.states).hasSize(4)
+            assertThat(resultsComposite.states).hasSize(5)
+            // composite query both linear and fungible and dummy asset states by participants
+            val commonQueryCriteria = VaultQueryCriteria(participants = listOf(MEGA_CORP))
+            val resultsAll = vaultService.queryBy<ContractState>(linearStateCriteria.or(fungibleAssetStateCriteria).or(commonQueryCriteria))
+            assertThat(resultsAll.states).hasSize(6)
         }
     }
 

testing/test-utils/src/main/kotlin/net/corda/testing/internal/vault/VaultFiller.kt

@@ -224,15 +224,16 @@ class VaultFiller @JvmOverloads constructor(
     /**
      * Records a dummy state in the Vault (useful for creating random states when testing vault queries)
      */
-    fun fillWithDummyState() : Vault<DummyState> {
+    fun fillWithDummyState(participants: List<AbstractParty> = listOf(services.myInfo.singleIdentity())) : Vault<DummyState> {
         val outputState = TransactionState(
-                data = DummyState(Random().nextInt(), participants = listOf(services.myInfo.singleIdentity())),
+                data = DummyState(Random().nextInt(), participants = participants),
                 contract = DummyContract.PROGRAM_ID,
                 notary = defaultNotary.party
         )
+        val participantKeys : List<PublicKey> = participants.map { it.owningKey }
         val builder = TransactionBuilder()
                 .addOutputState(outputState)
-                .addCommand(DummyCommandData, defaultNotary.party.owningKey)
+                .addCommand(DummyCommandData, participantKeys)
         val stxn = services.signInitialTransaction(builder)
         services.recordTransactions(stxn)
         return Vault(setOf(stxn.tx.outRef(0)))