diff --git a/docs/source/example-code/src/test/kotlin/net/corda/docs/WorkflowTransactionBuildTutorialTest.kt b/docs/source/example-code/src/test/kotlin/net/corda/docs/WorkflowTransactionBuildTutorialTest.kt
index b456cb8e02..17da495c0e 100644
--- a/docs/source/example-code/src/test/kotlin/net/corda/docs/WorkflowTransactionBuildTutorialTest.kt
+++ b/docs/source/example-code/src/test/kotlin/net/corda/docs/WorkflowTransactionBuildTutorialTest.kt
@@ -2,11 +2,14 @@ package net.corda.docs
 
 import net.corda.core.contracts.LinearState
 import net.corda.core.contracts.StateAndRef
-import net.corda.core.contracts.StateRef
+import net.corda.core.contracts.UniqueIdentifier
 import net.corda.core.getOrThrow
 import net.corda.core.node.ServiceHub
 import net.corda.core.node.services.ServiceInfo
-import net.corda.core.node.services.linearHeadsOfType
+import net.corda.core.node.services.Vault
+import net.corda.core.node.services.queryBy
+import net.corda.core.node.services.vault.QueryCriteria
+import net.corda.core.node.services.vault.and
 import net.corda.core.toFuture
 import net.corda.core.utilities.DUMMY_NOTARY
 import net.corda.core.utilities.DUMMY_NOTARY_KEY
@@ -26,10 +29,10 @@ class WorkflowTransactionBuildTutorialTest {
     lateinit var nodeB: MockNetwork.MockNode
 
     // Helper method to locate the latest Vault version of a LinearState from a possibly out of date StateRef
-    private inline fun <reified T : LinearState> ServiceHub.latest(ref: StateRef): StateAndRef<T> {
-        val linearHeads = vaultService.linearHeadsOfType<T>()
-        val original = storageService.validatedTransactions.getTransaction(ref.txhash)!!.tx.outRef<T>(ref.index)
-        return linearHeads[original.state.data.linearId]!!
+    private inline fun <reified T : LinearState> ServiceHub.latest(ref: UniqueIdentifier): StateAndRef<T> {
+        val linearHeads = vaultQueryService.queryBy<T>(QueryCriteria.LinearStateQueryCriteria(linearId = listOf(ref))
+                                                  .and(QueryCriteria.VaultQueryCriteria(status = Vault.StateStatus.UNCONSUMED)))
+        return linearHeads.states.single()
     }
 
     @Before
@@ -59,14 +62,15 @@ class WorkflowTransactionBuildTutorialTest {
         val flow1 = nodeA.services.startFlow(SubmitTradeApprovalFlow("1234", nodeB.info.legalIdentity))
         // Wait for the flow to finish
         val proposalRef = flow1.resultFuture.getOrThrow()
+        val proposalLinearId = proposalRef.state.data.linearId
         // Wait for NodeB to include it's copy in the vault
         nodeBVaultUpdate.get()
         // Fetch the latest copy of the state from both nodes
         val latestFromA = nodeA.database.transaction {
-            nodeA.services.latest<TradeApprovalContract.State>(proposalRef.ref)
+            nodeA.services.latest<TradeApprovalContract.State>(proposalLinearId)
         }
         val latestFromB = nodeB.database.transaction {
-            nodeB.services.latest<TradeApprovalContract.State>(proposalRef.ref)
+            nodeB.services.latest<TradeApprovalContract.State>(proposalLinearId)
         }
         // Confirm the state as as expected
         assertEquals(WorkflowState.NEW, proposalRef.state.data.state)
@@ -87,10 +91,10 @@ class WorkflowTransactionBuildTutorialTest {
         secondNodeBVaultUpdate.get()
         // Fetch the latest copies from the vault
         val finalFromA = nodeA.database.transaction {
-            nodeA.services.latest<TradeApprovalContract.State>(proposalRef.ref)
+            nodeA.services.latest<TradeApprovalContract.State>(proposalLinearId)
         }
         val finalFromB = nodeB.database.transaction {
-            nodeB.services.latest<TradeApprovalContract.State>(proposalRef.ref)
+            nodeB.services.latest<TradeApprovalContract.State>(proposalLinearId)
         }
         // Confirm the state is as expected
         assertEquals(WorkflowState.APPROVED, completedRef.state.data.state)
diff --git a/node/src/main/kotlin/net/corda/node/services/schema/HibernateObserver.kt b/node/src/main/kotlin/net/corda/node/services/schema/HibernateObserver.kt
index d53b29d5b1..b152f4c04b 100644
--- a/node/src/main/kotlin/net/corda/node/services/schema/HibernateObserver.kt
+++ b/node/src/main/kotlin/net/corda/node/services/schema/HibernateObserver.kt
@@ -6,7 +6,6 @@ import net.corda.core.contracts.StateRef
 import net.corda.core.node.services.Vault
 import net.corda.core.schemas.MappedSchema
 import net.corda.core.schemas.PersistentStateRef
-import net.corda.core.schemas.QueryableState
 import net.corda.core.utilities.debug
 import net.corda.core.utilities.loggerFor
 import net.corda.node.services.database.HibernateConfiguration
@@ -34,13 +33,11 @@ class HibernateObserver(vaultUpdates: Observable<Vault.Update>, val config: Hibe
 
     private fun persistState(stateAndRef: StateAndRef<ContractState>) {
         val state = stateAndRef.state.data
-        if (state is QueryableState) {
-            logger.debug { "Asked to persist state ${stateAndRef.ref}" }
-            config.schemaService.selectSchemas(state).forEach { persistStateWithSchema(state, stateAndRef.ref, it) }
-        }
+        logger.debug { "Asked to persist state ${stateAndRef.ref}" }
+        config.schemaService.selectSchemas(state).forEach { persistStateWithSchema(state, stateAndRef.ref, it) }
     }
 
-    fun persistStateWithSchema(state: QueryableState, stateRef: StateRef, schema: MappedSchema) {
+    fun persistStateWithSchema(state: ContractState, stateRef: StateRef, schema: MappedSchema) {
         val sessionFactory = config.sessionFactoryForSchema(schema)
         val session = sessionFactory.withOptions().
                 connection(TransactionManager.current().connection).