CORDA-1888 - Fix Vault Query composite queries (#3775) (#3869)

* CORDA-1888 Fix Vault Query composite queries  (#3775)

* Reproduce composite query failures.

* Fixes to OR querying and composite queries that use the same QueryCriteria (Linear, Fungible, Custom) more than once.

* Revert debug logging for Hibernate SQL.

* Cleanup and remove redundant joinPredicates global var.

* Fix failing Java Unit test.

* Fix Java compilation error in example-code section of docs.

* Include copy() function for original constructor to maintain backwards API compatibility.

* Fixed compilation errors post cherry-pick.

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

@@ -33,7 +33,9 @@ abstract class AbstractQueryCriteriaParser<Q : GenericQueryCriteria<Q,P>, in P:
         val leftPredicates = parse(left)
         val rightPredicates = parse(right)
 
-        val orPredicate = criteriaBuilder.or(*leftPredicates.toTypedArray(), *rightPredicates.toTypedArray())
+        val leftAnd = criteriaBuilder.and(*leftPredicates.toTypedArray())
+        val rightAnd = criteriaBuilder.and(*rightPredicates.toTypedArray())
+        val orPredicate = criteriaBuilder.or(leftAnd,rightAnd)
         predicateSet.add(orPredicate)
 
         return predicateSet
@@ -173,8 +175,6 @@ class HibernateQueryCriteriaParser(val contractStateType: Class<out ContractStat
         private val log = contextLogger()
     }
 
-    // incrementally build list of join predicates
-    private val joinPredicates = mutableListOf<Predicate>()
     // incrementally build list of root entities (for later use in Sort parsing)
     private val rootEntities = mutableMapOf<Class<out PersistentState>, Root<*>>(Pair(VaultSchemaV1.VaultStates::class.java, vaultStates))
     private val aggregateExpressions = mutableListOf<Expression<*>>()
@@ -323,8 +323,14 @@ class HibernateQueryCriteriaParser(val contractStateType: Class<out ContractStat
 
         val predicateSet = mutableSetOf<Predicate>()
 
-        val vaultFungibleStates = criteriaQuery.from(VaultSchemaV1.VaultFungibleStates::class.java)
+        // ensure we re-use any existing instance of the same root entity
-        rootEntities.putIfAbsent(VaultSchemaV1.VaultFungibleStates::class.java, vaultFungibleStates)
+        val entityStateClass = VaultSchemaV1.VaultFungibleStates::class.java
+        val vaultFungibleStates =
+                rootEntities.getOrElse(entityStateClass) {
+                    val entityRoot = criteriaQuery.from(entityStateClass)
+                    rootEntities[entityStateClass] = entityRoot
+                    entityRoot
+                }
 
         val joinPredicate = criteriaBuilder.equal(vaultStates.get<PersistentStateRef>("stateRef"), vaultFungibleStates.get<PersistentStateRef>("stateRef"))
         predicateSet.add(joinPredicate)
@@ -355,7 +361,7 @@ class HibernateQueryCriteriaParser(val contractStateType: Class<out ContractStat
         // participants
         criteria.participants?.let {
             val participants = criteria.participants as List<AbstractParty>
-            val joinLinearStateToParty = vaultFungibleStates.joinSet<VaultSchemaV1.VaultLinearStates, AbstractParty>("participants")
+            val joinLinearStateToParty = vaultFungibleStates.joinSet<VaultSchemaV1.VaultFungibleStates, AbstractParty>("participants")
             predicateSet.add(criteriaBuilder.and(joinLinearStateToParty.`in`(participants)))
             criteriaQuery.distinct(true)
         }
@@ -367,11 +373,17 @@ class HibernateQueryCriteriaParser(val contractStateType: Class<out ContractStat
 
         val predicateSet = mutableSetOf<Predicate>()
 
-        val vaultLinearStates = criteriaQuery.from(VaultSchemaV1.VaultLinearStates::class.java)
+        // ensure we re-use any existing instance of the same root entity
-        rootEntities.putIfAbsent(VaultSchemaV1.VaultLinearStates::class.java, vaultLinearStates)
+        val entityStateClass = VaultSchemaV1.VaultLinearStates::class.java
+        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"))
-        joinPredicates.add(joinPredicate)
+        predicateSet.add(joinPredicate)
 
         // linear ids UUID
         criteria.uuid?.let {
@@ -400,22 +412,28 @@ class HibernateQueryCriteriaParser(val contractStateType: Class<out ContractStat
         log.trace { "Parsing VaultCustomQueryCriteria: $criteria" }
 
         val predicateSet = mutableSetOf<Predicate>()
-        val entityClass = resolveEnclosingObjectFromExpression(criteria.expression)
+        val entityStateClass = resolveEnclosingObjectFromExpression(criteria.expression)
 
         try {
-            val entityRoot = criteriaQuery.from(entityClass)
+            // ensure we re-use any existing instance of the same root entity
-            rootEntities.putIfAbsent(entityClass, entityRoot)
+            val entityRoot =
+                    rootEntities.getOrElse(entityStateClass) {
+                        val entityRoot = criteriaQuery.from(entityStateClass)
+                        rootEntities[entityStateClass] = entityRoot
+                        entityRoot
+                    }
 
             val joinPredicate = criteriaBuilder.equal(vaultStates.get<PersistentStateRef>("stateRef"), entityRoot.get<PersistentStateRef>("stateRef"))
-            joinPredicates.add(joinPredicate)
+            predicateSet.add(joinPredicate)
 
             // resolve general criteria expressions
-            parseExpression(entityRoot, criteria.expression, predicateSet)
+            @Suppress("UNCHECKED_CAST")
+            parseExpression(entityRoot as Root<L>, criteria.expression, predicateSet)
         } catch (e: Exception) {
             e.message?.let { message ->
                 if (message.contains("Not an entity"))
                     throw VaultQueryException("""
-                    Please register the entity '${entityClass.name}'
+                    Please register the entity '${entityStateClass.name}'
                     See https://docs.corda.net/api-persistence.html#custom-schema-registration for more information""")
             }
             throw VaultQueryException("Parsing error: ${e.message}")
@@ -437,7 +455,7 @@ class HibernateQueryCriteriaParser(val contractStateType: Class<out ContractStat
                 else
                     aggregateExpressions
         criteriaQuery.multiselect(selections)
-        val combinedPredicates = joinPredicates.plus(predicateSet).plus(commonPredicates.values)
+        val combinedPredicates = commonPredicates.values.plus(predicateSet)
         criteriaQuery.where(*combinedPredicates.toTypedArray())
 
         return predicateSet
@@ -494,9 +512,7 @@ class HibernateQueryCriteriaParser(val contractStateType: Class<out ContractStat
                     rootEntities.getOrElse(entityStateClass) {
                         // scenario where sorting on attributes not parsed as criteria
                         val entityRoot = criteriaQuery.from(entityStateClass)
-                        rootEntities.put(entityStateClass, entityRoot)
+                        rootEntities[entityStateClass] = entityRoot
-                        val joinPredicate = criteriaBuilder.equal(vaultStates.get<PersistentStateRef>("stateRef"), entityRoot.get<PersistentStateRef>("stateRef"))
-                        joinPredicates.add(joinPredicate)
                         entityRoot
                     }
             when (direction) {
@@ -515,7 +531,6 @@ class HibernateQueryCriteriaParser(val contractStateType: Class<out ContractStat
         }
         if (orderCriteria.isNotEmpty()) {
             criteriaQuery.orderBy(orderCriteria)
-            criteriaQuery.where(*joinPredicates.toTypedArray())
         }
     }
 

node/src/test/java/net/corda/node/services/vault/VaultQueryJavaTests.java

@@ -44,6 +44,7 @@ import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
 import static java.util.Arrays.asList;
+import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
 import static net.corda.core.node.services.vault.Builder.sum;
 import static net.corda.core.node.services.vault.QueryCriteriaUtils.*;
@@ -316,12 +317,12 @@ public class VaultQueryJavaTests {
     @Test
     public void trackDealStatesPagedSorted() {
         List<String> dealIds = asList("123", "456", "789");
-        UniqueIdentifier uid =
+        UniqueIdentifier uid = new UniqueIdentifier("999", UUID.randomUUID());
         database.transaction(tx -> {
-                    Vault<LinearState> states = vaultFiller.fillWithSomeTestLinearStates(10, null);
+            vaultFiller.fillWithSomeTestLinearStates(10, null);
-                    UniqueIdentifier _uid = states.getStates().iterator().next().component1().getData().getLinearId();
+            vaultFiller.fillWithSomeTestLinearStates(1, null, emptyList(), uid);
             vaultFiller.fillWithSomeTestDeals(dealIds);
-                    return _uid;
+            return tx;
         });
         database.transaction(tx -> {
             // DOCSTART VaultJavaQueryExample5
@@ -344,7 +345,7 @@ public class VaultQueryJavaTests {
             Vault.Page<ContractState> snapshot = results.getSnapshot();
             // DOCEND VaultJavaQueryExample5
 
-            assertThat(snapshot.getStates()).hasSize(13);
+            assertThat(snapshot.getStates()).hasSize(4);
 
             return tx;
         });

node/src/test/kotlin/net/corda/node/services/persistence/HibernateConfigurationTest.kt

@@ -5,6 +5,7 @@ import net.corda.core.contracts.Amount
 import net.corda.core.contracts.StateAndRef
 import net.corda.core.contracts.StateRef
 import net.corda.core.contracts.TransactionState
+import net.corda.core.contracts.UniqueIdentifier
 import net.corda.core.crypto.SecureHash
 import net.corda.core.crypto.generateKeyPair
 import net.corda.core.identity.AbstractParty
@@ -622,6 +623,52 @@ class HibernateConfigurationTest {
         assertThat(queryResults).hasSize(10)
     }
 
+    /**
+     *  Composite OR query
+     */
+    @Test
+    fun `composite or query across VaultStates, VaultLinearStates and DummyLinearStates`() {
+        val uniqueID456 = UniqueIdentifier("456")
+        database.transaction {
+            vaultFiller.fillWithSomeTestLinearStates(1, externalId = "123", linearString = "123", linearNumber = 123, linearBoolean = true)
+            vaultFiller.fillWithSomeTestLinearStates(1, uniqueIdentifier = uniqueID456)
+            vaultFiller.fillWithSomeTestLinearStates(1, externalId = "789")
+        }
+        val sessionFactory = sessionFactoryForSchemas(VaultSchemaV1, DummyLinearStateSchemaV1)
+        val criteriaBuilder = sessionFactory.criteriaBuilder
+        val entityManager = sessionFactory.createEntityManager()
+
+        // structure query
+        val criteriaQuery = criteriaBuilder.createQuery(VaultSchemaV1.VaultStates::class.java)
+        val vaultStates = criteriaQuery.from(VaultSchemaV1.VaultStates::class.java)
+        val vaultLinearStates = criteriaQuery.from(VaultSchemaV1.VaultLinearStates::class.java)
+        val dummyLinearStates = criteriaQuery.from(DummyLinearStateSchemaV1.PersistentDummyLinearState::class.java)
+
+        criteriaQuery.select(vaultStates)
+        val joinPredicate1 = criteriaBuilder.equal(vaultStates.get<PersistentStateRef>("stateRef"), vaultLinearStates.get<PersistentStateRef>("stateRef"))
+        val joinPredicate2 = criteriaBuilder.and(criteriaBuilder.equal(vaultStates.get<PersistentStateRef>("stateRef"), dummyLinearStates.get<PersistentStateRef>("stateRef")))
+
+        // and predicates on VaultLinearStates
+        val andLinearStatesPredicate1 = criteriaBuilder.and(criteriaBuilder.equal(vaultLinearStates.get<String>("externalId"), uniqueID456.externalId))
+        val andLinearStatesPredicate2 = criteriaBuilder.and(criteriaBuilder.equal(vaultLinearStates.get<UUID>("uuid"), uniqueID456.id))
+        val andLinearStatesPredicate = criteriaBuilder.and(andLinearStatesPredicate1, andLinearStatesPredicate2)
+
+        // and predicates on PersistentDummyLinearState
+        val andDummyLinearStatesPredicate1 = criteriaBuilder.and(criteriaBuilder.equal(dummyLinearStates.get<String>("linearString"), "123"))
+        val andDummyLinearStatesPredicate2 = criteriaBuilder.and(criteriaBuilder.equal(dummyLinearStates.get<Long>("linearNumber"), 123L))
+        val andDummyLinearStatesPredicate3 = criteriaBuilder.and(criteriaBuilder.equal(dummyLinearStates.get<Boolean>("linearBoolean"), true))
+        val andDummyLinearStatesPredicate = criteriaBuilder.and(andDummyLinearStatesPredicate1, criteriaBuilder.and(andDummyLinearStatesPredicate2, andDummyLinearStatesPredicate3))
+
+        // or predicates
+        val orPredicates = criteriaBuilder.or(andLinearStatesPredicate, andDummyLinearStatesPredicate)
+
+        criteriaQuery.where(joinPredicate1, joinPredicate2, orPredicates)
+
+        // execute query
+        val queryResults = entityManager.createQuery(criteriaQuery).resultList
+        assertThat(queryResults).hasSize(2)
+    }
+
     /**
      *  Test a OneToOne table mapping
      */

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

@@ -1276,11 +1276,12 @@ class VaultQueryTests {
     @Test
     fun `unconsumed deal states sorted`() {
         database.transaction {
-            val linearStates = vaultFiller.fillWithSomeTestLinearStates(10)
+            vaultFiller.fillWithSomeTestLinearStates(10)
+            val uid = UniqueIdentifier("999")
+            vaultFiller.fillWithSomeTestLinearStates(1, uniqueIdentifier = uid)
             vaultFiller.fillWithSomeTestDeals(listOf("123", "456", "789"))
-            val uid = linearStates.states.first().state.data.linearId.id
 
-            val linearStateCriteria = LinearStateQueryCriteria(uuid = listOf(uid))
+            val linearStateCriteria = LinearStateQueryCriteria(uuid = listOf(uid.id))
             val dealStateCriteria = LinearStateQueryCriteria(externalId = listOf("123", "456", "789"))
             val compositeCriteria = linearStateCriteria or dealStateCriteria
 
@@ -1810,6 +1811,63 @@ class VaultQueryTests {
         }
     }
 
+    @Test
+    fun `composite query for fungible and linear states`() {
+        database.transaction {
+            vaultFiller.fillWithSomeTestLinearStates(1, "TEST1")
+            vaultFiller.fillWithSomeTestDeals(listOf("123"))
+            vaultFiller.fillWithSomeTestCash(100.DOLLARS, notaryServices, 1, DUMMY_CASH_ISSUER, services.myInfo.singleIdentity())
+            vaultFiller.fillWithSomeTestCommodity(Amount(100, Commodity.getInstance("FCOJ")!!), notaryServices, DUMMY_OBLIGATION_ISSUER.ref(1))
+            vaultFiller.fillWithDummyState()
+            // all contract states query
+            val results = vaultService.queryBy<ContractState>()
+            assertThat(results.states).hasSize(5)
+            // linear states only query
+            val linearStateCriteria = LinearStateQueryCriteria()
+            val resultsLSC = vaultService.queryBy<ContractState>(linearStateCriteria)
+            assertThat(resultsLSC.states).hasSize(2)
+            // fungible asset states only query
+            val fungibleAssetStateCriteria = FungibleAssetQueryCriteria()
+            val resultsFASC = vaultService.queryBy<ContractState>(fungibleAssetStateCriteria)
+            assertThat(resultsFASC.states).hasSize(2)
+            // composite OR query for both linear and fungible asset states (eg. all states in either Fungible and Linear states tables)
+            val resultsCompositeOr = vaultService.queryBy<ContractState>(fungibleAssetStateCriteria.or(linearStateCriteria))
+            assertThat(resultsCompositeOr.states).hasSize(4)
+            // composite AND query for both linear and fungible asset states (eg. all states in both Fungible and Linear states tables)
+            val resultsCompositeAnd = vaultService.queryBy<ContractState>(fungibleAssetStateCriteria.and(linearStateCriteria))
+            assertThat(resultsCompositeAnd.states).hasSize(0)
+        }
+    }
+
+    @Test
+    fun `composite query for fungible and linear states for multiple participants`() {
+        database.transaction {
+            identitySvc.verifyAndRegisterIdentity(ALICE_IDENTITY)
+            identitySvc.verifyAndRegisterIdentity(BOB_IDENTITY)
+            identitySvc.verifyAndRegisterIdentity(CHARLIE_IDENTITY)
+            vaultFiller.fillWithSomeTestLinearStates(1, "TEST1", listOf(ALICE))
+            vaultFiller.fillWithSomeTestLinearStates(1,  "TEST2", listOf(BOB))
+            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()
+            // all contract states query
+            val results = vaultService.queryBy<ContractState>()
+            assertThat(results.states).hasSize(6)
+            // linear states by participants only query
+            val linearStateCriteria = LinearStateQueryCriteria(participants = listOf(ALICE,BOB))
+            val resultsLSC = vaultService.queryBy<ContractState>(linearStateCriteria)
+            assertThat(resultsLSC.states).hasSize(2)
+            // fungible asset states by participants only query
+            val fungibleAssetStateCriteria = FungibleAssetQueryCriteria(participants = listOf(services.myInfo.singleIdentity()))
+            val resultsFASC = vaultService.queryBy<ContractState>(fungibleAssetStateCriteria)
+            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)
+        }
+    }
+
     @Test
     fun `unconsumed linear heads where external id is null`() {
         database.transaction {

testing/test-utils/src/main/kotlin/net/corda/testing/contracts/DummyState.kt

@@ -6,8 +6,10 @@ import net.corda.core.identity.AbstractParty
 /**
  * Dummy state for use in testing. Not part of any contract, not even the [DummyContract].
  */
-data class DummyState(
+data class DummyState @JvmOverloads constructor (
         /** Some information that the state represents for test purposes. **/
-        val magicNumber: Int = 0) : ContractState {
+        val magicNumber: Int = 0,
-    override val participants: List<AbstractParty> get() = emptyList()
+        override val participants: List<AbstractParty> = listOf()) : ContractState {
+
+    fun copy(magicNumber: Int = this.magicNumber) = DummyState(magicNumber)
 }

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

@@ -17,6 +17,8 @@ import net.corda.finance.contracts.DealState
 import net.corda.finance.contracts.asset.Cash
 import net.corda.finance.contracts.asset.Obligation
 import net.corda.finance.contracts.asset.OnLedgerAsset
+import net.corda.testing.contracts.DummyContract
+import net.corda.testing.contracts.DummyState
 import net.corda.testing.core.*
 import net.corda.testing.internal.chooseIdentity
 import net.corda.testing.internal.chooseIdentityAndCert
@@ -96,6 +98,7 @@ class VaultFiller @JvmOverloads constructor(
     fun fillWithSomeTestLinearStates(numberToCreate: Int,
                                      externalId: String? = null,
                                      participants: List<AbstractParty> = emptyList(),
+                                     uniqueIdentifier: UniqueIdentifier? = null,
                                      linearString: String = "",
                                      linearNumber: Long = 0L,
                                      linearBoolean: Boolean = false,
@@ -108,7 +111,7 @@ class VaultFiller @JvmOverloads constructor(
             // Issue a Linear state
             val dummyIssue = TransactionBuilder(notary = defaultNotary.party).apply {
                 addOutputState(DummyLinearContract.State(
-                        linearId = UniqueIdentifier(externalId),
+                        linearId = uniqueIdentifier ?: UniqueIdentifier(externalId),
                         participants = participants.plus(me),
                         linearString = linearString,
                         linearNumber = linearNumber,
@@ -168,6 +171,23 @@ class VaultFiller @JvmOverloads constructor(
     }
 
 
+    /**
+     * Records a dummy state in the Vault (useful for creating random states when testing vault queries)
+     */
+    fun fillWithDummyState() : Vault<DummyState> {
+        val outputState = TransactionState(
+                data = DummyState(Random().nextInt(), participants = listOf(services.myInfo.singleIdentity())),
+                contract = DummyContract.PROGRAM_ID,
+                notary = defaultNotary.party
+        )
+        val builder = TransactionBuilder()
+                .addOutputState(outputState)
+                .addCommand(DummyCommandData, defaultNotary.party.owningKey)
+        val stxn = services.signInitialTransaction(builder)
+        services.recordTransactions(stxn)
+        return Vault(setOf(stxn.tx.outRef(0)))
+    }
+
     /**
      * Puts together an issuance transaction for the specified amount that starts out being owned by the given pubkey.
      */