Merge pull request #2632 from piotr-roslaniec/refactor-queries#2153

Refactor queries into a separate file

nucypher/characters/lawful.py

@@ -91,7 +91,7 @@
 )
 from nucypher.crypto.signing import InvalidSignature
 from nucypher.datastore.datastore import DatastoreTransactionError, RecordNotFound
-from nucypher.datastore.models import PolicyArrangement, TreasureMap as DatastoreTreasureMap
+from nucypher.datastore.queries import find_expired_policies, find_expired_treasure_maps
 from nucypher.network.exceptions import NodeSeemsToBeDown
 from nucypher.network.middleware import RestMiddleware
 from nucypher.network.nodes import NodeSprout, Teacher
@@ -1221,10 +1221,7 @@ def __prune_datastore(self) -> None:
         """Deletes all expired arrangements, kfrags, and treasure maps in the datastore."""
         now = maya.MayaDT.from_datetime(datetime.fromtimestamp(self._datastore_pruning_task.clock.seconds()))
         try:
-            with self.datastore.query_by(PolicyArrangement,
-                                         filter_field='expiration',
-                                         filter_func=lambda expiration: expiration <= now,
-                                         writeable=True) as expired_policies:
+            with find_expired_policies(self.datastore, now) as expired_policies:
                 for policy in expired_policies:
                     policy.delete()
                 result = len(expired_policies)
@@ -1237,10 +1234,7 @@ def __prune_datastore(self) -> None:
                 self.log.debug(f"Pruned {result} policy arrangements.")
 
         try:
-            with self.datastore.query_by(DatastoreTreasureMap,
-                                         filter_field='expiration',
-                                         filter_func=lambda expiration: expiration <= now,
-                                         writeable=True) as expired_treasure_maps:
+            with find_expired_treasure_maps(self.datastore, now) as expired_treasure_maps:
                 for treasure_map in expired_treasure_maps:
                     treasure_map.delete()
                 result = len(expired_treasure_maps)

nucypher/cli/painting/nodes.py

@@ -17,8 +17,7 @@
 import maya
 
 from nucypher.config.constants import SEEDNODES
-from nucypher.datastore.datastore import RecordNotFound
-from nucypher.datastore.models import Workorder
+from nucypher.datastore.queries import get_work_orders
 
 
 def build_fleet_state_status(ursula) -> str:
@@ -42,11 +41,8 @@ def paint_node_status(emitter, ursula, start_time):
     # Build FleetState status line
     fleet_state = build_fleet_state_status(ursula=ursula)
 
-    try:
-        with ursula.datastore.query_by(Workorder) as work_orders:
-            num_work_orders = len(work_orders)
-    except RecordNotFound:
-        num_work_orders = 0
+    work_orders = get_work_orders(ursula.datastore)
+    num_work_orders = len(work_orders)
 
     stats = ['⇀URSULA {}↽'.format(ursula.nickname.icon),
              '{}'.format(ursula),

nucypher/datastore/datastore.py

@@ -15,16 +15,13 @@
 along with nucypher.  If not, see <https://www.gnu.org/licenses/>.
 """
 import lmdb
-import maya
 from contextlib import contextmanager, suppress
 from functools import partial
-from typing import Any, Callable, List, NamedTuple, Optional, Type, Union
+from typing import Any, Callable, Generator, List, NamedTuple, Optional, Type, Union
 
-from bytestring_splitter import BytestringSplitter
-from nucypher.crypto.signing import Signature
-from nucypher.datastore.base import DatastoreRecord, DBWriteError, RecordField
-from nucypher.datastore.models import PolicyArrangement, Workorder
+from nucypher.datastore.base import DatastoreRecord, DBWriteError
 
+DatastoreQueryResult = Generator[List[Type['DatastoreRecord']], None, None]
 
 class RecordNotFound(Exception):
     """
@@ -141,7 +138,7 @@ def query_by(self,
               filter_func: Optional[Callable[[Union[Any, Type['DatastoreRecord']]], bool]] = None,
               filter_field: str = "",
               writeable: bool = False,
-              ) -> List[Type['DatastoreRecord']]:
+              ) -> DatastoreQueryResult:
         """
         Performs a query on the datastore for the record by `record_type`.
 

nucypher/datastore/queries.py

@@ -0,0 +1,74 @@
+"""
+ This file is part of nucypher.
+
+ nucypher is free software: you can redistribute it and/or modify
+ it under the terms of the GNU Affero General Public License as published by
+ the Free Software Foundation, either version 3 of the License, or
+ (at your option) any later version.
+
+ nucypher is distributed in the hope that it will be useful,
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ GNU Affero General Public License for more details.
+
+ You should have received a copy of the GNU Affero General Public License
+ along with nucypher.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+import functools
+from typing import Callable, List, Type
+
+import maya
+
+from nucypher.datastore.base import DatastoreRecord
+from nucypher.datastore.datastore import Datastore, DatastoreQueryResult, RecordNotFound
+from nucypher.datastore.models import PolicyArrangement, TreasureMap, Workorder
+
+
+def unwrap_records(func: Callable[..., DatastoreQueryResult]) -> Callable[..., List[Type['DatastoreRecord']]]:
+    """
+    Used to safely unwrap results of a query.
+    Suitable only for reading `DatastoreRecord`s. Use `find_*` functions if you want to modify records.
+
+    Since results returned by `Datastore.query_by()` are lazy (wrapped in a `@contextmanager` generator)
+    we have to unwrap them and handle `RecordNotFound` error, if any. `DatastoreRecord`s are not writable
+    after unwrapping, because exiting `@contextmanager` is also closing `Datastore` transaction.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs) -> List[Type['DatastoreRecord']]:
+        try:
+            with func(*args, **kwargs) as results:
+                return results
+        except RecordNotFound:
+            return []
+
+    return wrapper
+
+
+def find_expired_policies(ds: Datastore, cutoff: maya.MayaDT) -> DatastoreQueryResult:
+    return ds.query_by(PolicyArrangement,
+                       filter_field='expiration',
+                       filter_func=lambda expiration: expiration <= cutoff,
+                       writeable=True)
+
+
+def find_expired_treasure_maps(ds: Datastore, cutoff: maya.MayaDT) -> DatastoreQueryResult:
+    return ds.query_by(TreasureMap,
+                       filter_field='expiration',
+                       filter_func=lambda expiration: expiration <= cutoff,
+                       writeable=True)
+
+
+@unwrap_records
+def get_work_orders(ds: Datastore) -> List[Workorder]:
+    return ds.query_by(Workorder)
+
+
+def find_policy_arrangements(ds: Datastore) -> DatastoreQueryResult:
+    return ds.query_by(PolicyArrangement, writeable=True)
+
+
+@unwrap_records
+def get_policy_arrangements(ds: Datastore) -> List[PolicyArrangement]:
+    return ds.query_by(PolicyArrangement)

nucypher/utilities/prometheus/collector.py

@@ -30,8 +30,7 @@
 from nucypher.blockchain.eth.agents import ContractAgency, PolicyManagerAgent, StakingEscrowAgent, WorkLockAgent
 from nucypher.blockchain.eth.interfaces import BlockchainInterfaceFactory
 from nucypher.blockchain.eth.registry import BaseContractRegistry
-from nucypher.datastore.datastore import RecordNotFound
-from nucypher.datastore.models import Workorder, PolicyArrangement
+from nucypher.datastore.queries import get_policy_arrangements, get_work_orders
 
 from prometheus_client.registry import CollectorRegistry
 
@@ -129,11 +128,9 @@ def _collect_internal(self) -> None:
             self.metrics["availability_score_gauge"].set(self.ursula._availability_tracker.score)
         else:
             self.metrics["availability_score_gauge"].set(-1)
-        try:
-            with self.ursula.datastore.query_by(Workorder) as work_orders:
-                self.metrics["work_orders_gauge"].set(len(work_orders))
-        except RecordNotFound:
-            self.metrics["work_orders_gauge"].set(0)
+
+        work_orders = get_work_orders(self.ursula.datastore)
+        self.metrics["work_orders_gauge"].set(len(work_orders))
 
         if not self.ursula.federated_only:
             staking_agent = ContractAgency.get_agent(StakingEscrowAgent, registry=self.ursula.registry)
@@ -145,11 +142,8 @@ def _collect_internal(self) -> None:
                                      'missing_commitments': str(missing_commitments)}
             base_payload.update(decentralized_payload)
 
-            try:
-                with self.ursula.datastore.query_by(PolicyArrangement) as policy_arrangements:
-                    self.metrics["policies_held_gauge"].set(len(policy_arrangements))
-            except RecordNotFound:
-                self.metrics["policies_held_gauge"].set(0)
+            policy_arrangements = get_policy_arrangements(self.ursula.datastore)
+            self.metrics["policies_held_gauge"].set(len(policy_arrangements))
 
         self.metrics["host_info"].info(base_payload)
 

tests/acceptance/characters/test_freerider_attacks.py

@@ -22,6 +22,7 @@
 from nucypher.characters.unlawful import Amonia
 from nucypher.datastore.models import PolicyArrangement, TreasureMap as DatastoreTreasureMap
 from nucypher.datastore.datastore import RecordNotFound
+from nucypher.datastore.queries import find_policy_arrangements
 from nucypher.network.middleware import RestMiddleware
 
 
@@ -47,8 +48,9 @@ def test_policy_simple_sinpa(blockchain_ursulas, blockchain_alice, blockchain_bo
         # Reset the Ursula for the next test.
         ursula.suspicious_activities_witnessed['freeriders'] = []
         try:
-            with ursula.datastore.query_by(PolicyArrangement, writeable=True) as arrangements:
-                [arrangement.delete() for arrangement in arrangements]
+            with find_policy_arrangements(ursula.datastore) as arrangements:
+                for arrangement in arrangements:
+                    arrangement.delete()
         except RecordNotFound:
             # No records were found; this Ursula didn't have the arrangement.
             continue
@@ -76,7 +78,7 @@ def test_try_to_post_free_arrangement_by_hacking_enact(blockchain_ursulas, block
     for ursula in blockchain_ursulas:
         # Even though the grant executed without error...
         try:
-            with ursula.datastore.query_by(PolicyArrangement, writeable=True) as all_arrangements:
+            with find_policy_arrangements(ursula.datastore) as all_arrangements:
                 arrangement = all_arrangements[0] # ...and Ursula did save the Arrangement after considering it...
                 with pytest.raises(AttributeError):
                     should_error = arrangement.kfrag # ...Ursula did *not* save a KFrag and will not service this Policy.
@@ -88,7 +90,8 @@ def test_try_to_post_free_arrangement_by_hacking_enact(blockchain_ursulas, block
 
                 # Reset the Ursula for the next test.
                 ursula.suspicious_activities_witnessed['freeriders'] = []
-                [arrangement.delete() for arrangement in all_arrangements]
+                for arrangement in all_arrangements:
+                    arrangement.delete()
         except RecordNotFound:
             # No records were found; this Ursula didn't have the arrangement.
             continue
@@ -120,7 +123,7 @@ def test_pay_a_flunky_instead_of_the_arranged_ursula(blockchain_alice, blockchai
     for ursula in blockchain_ursulas:
         # Even though the grant executed without error...
         try:
-            with ursula.datastore.query_by(PolicyArrangement, writeable=True) as all_arrangements:
+            with find_policy_arrangements(ursula.datastore) as all_arrangements:
                 arrangement = all_arrangements[0] # ...and Ursula did save the Arrangement after considering it...
                 with pytest.raises(AttributeError):
                     should_error = arrangement.kfrag # ...Ursula did *not* save a KFrag and will not service this Policy.
@@ -132,7 +135,8 @@ def test_pay_a_flunky_instead_of_the_arranged_ursula(blockchain_alice, blockchai
 
                 # Reset the Ursula for the next test.
                 ursula.suspicious_activities_witnessed['freeriders'] = []
-                [arrangement.delete() for arrangement in all_arrangements]
+                for arrangement in all_arrangements:
+                    arrangement.delete()
         except RecordNotFound:
             # No records were found; this Ursula didn't have the arrangement.
             continue