cache credential exchange by thread ID & initiator

Signed-off-by: Andrew Whitehead <cywolf@gmail.com>

Author: andrewwhitehead

aries_cloudagent/holder/tests/test_indy.py

@@ -13,7 +13,7 @@
 import pytest
 
 from ...messaging.issue_credential.v1_0.messages.inner.credential_preview import (
-    CredentialPreview
+    CredentialPreview,
 )
 
 
@@ -73,7 +73,7 @@ async def test_get_credential_attrs_mime_types(self, mock_nonsec_get_wallet_reco
             "type": IndyHolder.RECORD_TYPE_MIME_TYPES,
             "id": cred_id,
             "value": "value",
-            "tags": dummy_tags
+            "tags": dummy_tags,
         }
         mock_nonsec_get_wallet_record.return_value = json.dumps(dummy_rec)
 
@@ -88,12 +88,8 @@ async def test_get_credential_attrs_mime_types(self, mock_nonsec_get_wallet_reco
             dummy_rec["type"],
             f"{IndyHolder.RECORD_TYPE_MIME_TYPES}::{dummy_rec['id']}",
             json.dumps(
-                {
-                    "retrieveType": True,
-                    "retrieveValue": True,
-                    "retrieveTags": True
-                }
-            )
+                {"retrieveType": False, "retrieveValue": True, "retrieveTags": True}
+            ),
         )
 
         assert mime_types == dummy_tags
@@ -106,7 +102,7 @@ async def test_get_credential_attr_mime_type(self, mock_nonsec_get_wallet_record
             "type": IndyHolder.RECORD_TYPE_MIME_TYPES,
             "id": cred_id,
             "value": "value",
-            "tags": dummy_tags
+            "tags": dummy_tags,
         }
         mock_nonsec_get_wallet_record.return_value = json.dumps(dummy_rec)
 
@@ -121,12 +117,8 @@ async def test_get_credential_attr_mime_type(self, mock_nonsec_get_wallet_record
             dummy_rec["type"],
             f"{IndyHolder.RECORD_TYPE_MIME_TYPES}::{dummy_rec['id']}",
             json.dumps(
-                {
-                    "retrieveType": True,
-                    "retrieveValue": True,
-                    "retrieveTags": True
-                }
-            )
+                {"retrieveType": False, "retrieveValue": True, "retrieveTags": True}
+            ),
         )
 
         assert a_mime_type == dummy_tags["a"]
@@ -233,7 +225,7 @@ async def test_delete_credential(
         self,
         mock_nonsec_del_wallet_record,
         mock_nonsec_get_wallet_record,
-        mock_prover_del_cred
+        mock_prover_del_cred,
     ):
         mock_wallet = async_mock.MagicMock()
         holder = IndyHolder(mock_wallet)
@@ -242,18 +234,14 @@ async def test_delete_credential(
                 "type": "typ",
                 "id": "ident",
                 "value": "value",
-                "tags": {
-                    "a": json.dumps("1"),
-                    "b": json.dumps("2")
-                }
+                "tags": {"a": json.dumps("1"), "b": json.dumps("2")},
             }
         )
 
         credential = await holder.delete_credential("credential_id")
 
         mock_prover_del_cred.assert_called_once_with(
-            mock_wallet.handle,
-            "credential_id"
+            mock_wallet.handle, "credential_id"
         )
 
     @async_mock.patch("indy.anoncreds.prover_create_proof")

aries_cloudagent/messaging/credentials/manager.py

@@ -340,14 +340,10 @@ async def receive_request(self, credential_request_message: CredentialRequest):
 
         credential_request = json.loads(credential_request_message.request)
 
-        credential_exchange_record = await CredentialExchange.retrieve_by_tag_filter(
-            self.context,
-            {
-                "thread_id": credential_request_message._thread_id,
-            },
-            {
-                "initiator": "self",
-            },
+        (
+            credential_exchange_record
+        ) = await CredentialExchange.retrieve_by_thread_and_initiator(
+            self.context, credential_request_message._thread_id, "self"
         )
         credential_exchange_record.credential_request = credential_request
         credential_exchange_record.state = CredentialExchange.STATE_REQUEST_RECEIVED
@@ -443,14 +439,8 @@ async def receive_credential(self, credential_message: CredentialIssue):
         try:
             (
                 credential_exchange_record
-            ) = await CredentialExchange.retrieve_by_tag_filter(
-                self.context,
-                {
-                    "thread_id": credential_message._thread_id,
-                },
-                {
-                    "initiator": "external",
-                },
+            ) = await CredentialExchange.retrieve_by_thread_and_initiator(
+                self.context, credential_message._thread_id, "external"
             )
         except StorageNotFoundError:
 
@@ -464,14 +454,8 @@ async def receive_credential(self, credential_message: CredentialIssue):
             # credential_request_metadata
             (
                 credential_exchange_record
-            ) = await CredentialExchange.retrieve_by_tag_filter(
-                self.context,
-                {
-                    "thread_id": credential_message._thread.pthid,
-                },
-                {
-                    "initiator": "external",
-                },
+            ) = await CredentialExchange.retrieve_by_thread_and_initiator(
+                self.context, credential_message._thread.pthid, "external"
             )
 
             # Copy values from parent but create new record on save (no id)
@@ -587,14 +571,10 @@ async def credential_stored(self, credential_stored_message: CredentialStored):
         """
 
         # Get current exchange record by thread id
-        credential_exchange_record = await CredentialExchange.retrieve_by_tag_filter(
-            self.context,
-            {
-                "thread_id": credential_stored_message._thread_id,
-            },
-            {
-                "initiator": "self",
-            },
+        (
+            credential_exchange_record
+        ) = await CredentialExchange.retrieve_by_thread_and_initiator(
+            self.context, credential_stored_message._thread_id, "self"
         )
 
         credential_exchange_record.state = CredentialExchange.STATE_STORED

aries_cloudagent/messaging/credentials/models/credential_exchange.py

@@ -2,6 +2,8 @@
 
 from marshmallow import fields
 
+from ....config.injection_context import InjectionContext
+
 from ...models.base_record import BaseRecord, BaseRecordSchema
 
 
@@ -100,6 +102,44 @@ def record_value(self) -> dict:
             )
         }
 
+    @classmethod
+    async def retrieve_by_thread_and_initiator(
+        cls, context: InjectionContext, thread_id: str, initiator: str
+    ) -> "CredentialExchange":
+        """Retrieve a credential exchange record by thread ID and inititator."""
+        cache_key = f"credential_exchange_tidx::{thread_id}::{initiator}"
+        record_id = await cls.get_cached_key(context, cache_key)
+        if record_id:
+            record = await cls.retrieve_by_id(context, record_id)
+        else:
+            record = await cls.retrieve_by_tag_filter(
+                context, {"thread_id": thread_id}, {"initiator": initiator}
+            )
+            await cls.set_cached_key(context, cache_key, record.credential_exchange_id)
+        return record
+
+    async def post_save(
+        self,
+        context: InjectionContext,
+        new_record: bool,
+        last_state: str,
+        webhook: bool = None,
+    ):
+        """Perform post-save actions.
+
+        Args:
+            context: The injection context to use
+            new_record: Flag indicating if the record was just created
+            last_state: The previous state value
+            webhook: Adjust whether the webhook is called
+        """
+        await super(CredentialExchange, self).post_save(
+            context, new_record, last_state, webhook
+        )
+        if self.thread_id and self.initiator:
+            cache_key = f"credential_exchange_tidx::{self.thread_id}::{self.initiator}"
+            await self.set_cached_key(context, cache_key, self.credential_exchange_id)
+
 
 class CredentialExchangeSchema(BaseRecordSchema):
     """Schema to allow serialization/deserialization of credential exchange records."""

aries_cloudagent/messaging/models/base_record.py

@@ -222,14 +222,14 @@ async def retrieve_by_tag_filter(
             post_filter: Additional value filters to apply after retrieval
         """
         storage: BaseStorage = await context.inject(BaseStorage)
-        result = storage.search_records(
+        query = storage.search_records(
             cls.RECORD_TYPE,
             cls.prefix_tag_filter(tag_filter),
             None,
             {"retrieveTags": False},
         )
         found = None
-        async for record in result:
+        async for record in query:
             vals = json.loads(record.value)
             if not post_filter or match_post_filter(vals, post_filter):
                 if found:
@@ -254,14 +254,14 @@ async def query(
             post_filter: Additional value filters to apply
         """
         storage: BaseStorage = await context.inject(BaseStorage)
-        result = storage.search_records(
+        query = storage.search_records(
             cls.RECORD_TYPE,
             cls.prefix_tag_filter(tag_filter),
             None,
             {"retrieveTags": False},
         )
         result = []
-        async for record in result:
+        async for record in query:
             vals = json.loads(record.value)
             if not post_filter or match_post_filter(vals, post_filter):
                 result.append(cls.from_storage(record.id, vals))

aries_cloudagent/messaging/models/tests/test_base_record.py

@@ -126,7 +126,7 @@ async def test_retrieve_uncached_id(self):
             mock_storage.get_record.return_value = stored
             result = await BaseRecordImpl.retrieve_by_id(context, record_id, False)
             mock_storage.get_record.assert_awaited_once_with(
-                BaseRecordImpl.RECORD_TYPE, record_id
+                BaseRecordImpl.RECORD_TYPE, record_id, {"retrieveTags": False}
             )
             set_cached_key.assert_awaited_once_with(
                 context, cache_key.return_value, record_value
@@ -146,12 +146,10 @@ async def test_query(self):
             BaseRecordImpl.RECORD_TYPE, json.dumps(record_value), {}, record_id
         )
 
-        mock_storage.search_records.return_value.fetch_all = async_mock.CoroutineMock(
-            return_value=[stored]
-        )
+        mock_storage.search_records.return_value.__aiter__.return_value = [stored]
         result = await BaseRecordImpl.query(context, tag_filter)
         mock_storage.search_records.assert_called_once_with(
-            BaseRecordImpl.RECORD_TYPE, tag_filter
+            BaseRecordImpl.RECORD_TYPE, tag_filter, None, {"retrieveTags": False}
         )
         assert result and isinstance(result[0], BaseRecordImpl)
         assert result[0]._id == record_id