Refactor queries into a separate file

pull/2632/head
Piotr Roslaniec 2021-04-09 20:08:28 +02:00
parent 555f7c631c
commit 269c7c2fc2
5 changed files with 100 additions and 71 deletions

View File

@ -90,8 +90,8 @@ from nucypher.crypto.powers import (
TransactingPower
)
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.datastore import DatastoreTransactionError
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,35 +1221,29 @@ class Ursula(Teacher, Character, Worker):
"""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)
except RecordNotFound:
self.log.debug("No expired policy arrangements found.")
except DatastoreTransactionError:
self.log.warn(f"Failed to prune policy arrangements; DB session rolled back.")
else:
if result > 0:
if result == 0:
self.log.debug("No expired policy arrangements found.")
elif result > 0:
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)
except RecordNotFound:
self.log.debug("No expired treasure maps found.")
except DatastoreTransactionError:
self.log.warn(f"Failed to prune expired treasure maps; DB session rolled back.")
else:
if result > 0:
if result == 0:
self.log.debug("No expired treasure maps found.")
elif result > 0:
self.log.debug(f"Pruned {result} treasure maps.")
def __preflight(self) -> None:

View File

@ -17,8 +17,7 @@ along with nucypher. If not, see <https://www.gnu.org/licenses/>.
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 find_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
with find_work_orders(ursula.datastore) as work_orders:
num_work_orders = len(work_orders)
stats = ['⇀URSULA {}'.format(ursula.nickname.icon),
'{}'.format(ursula),

View File

@ -0,0 +1,56 @@
"""
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/>.
"""
from typing import List, Type
from nucypher.datastore.base import DatastoreRecord
from nucypher.datastore.datastore import Datastore, RecordNotFound
from nucypher.datastore.models import PolicyArrangement, TreasureMap, Workorder
def find_expired_policies(ds: Datastore, now) -> List[Type['DatastoreRecord']]:
try:
return ds.query_by(PolicyArrangement,
filter_field='expiration',
filter_func=lambda expiration: expiration <= now,
writeable=True)
except RecordNotFound:
return []
def find_expired_treasure_maps(ds: Datastore, now) -> List[Type['DatastoreRecord']]:
try:
return ds.query_by(TreasureMap,
filter_field='expiration',
filter_func=lambda expiration: expiration <= now,
writeable=True)
except RecordNotFound:
return []
def find_work_orders(ds: Datastore) -> List[Type['DatastoreRecord']]:
try:
return ds.query_by(Workorder)
except RecordNotFound:
return []
def find_policy_arrangements(ds: Datastore, writeable=False) -> List[Type['DatastoreRecord']]:
try:
return ds.query_by(PolicyArrangement, writeable=writeable)
except RecordNotFound:
return []

View File

@ -30,8 +30,7 @@ from nucypher.blockchain.eth.actors import NucypherTokenActor
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 find_policy_arrangements, find_work_orders
from prometheus_client.registry import CollectorRegistry
@ -129,11 +128,9 @@ class UrsulaInfoMetricsCollector(BaseMetricsCollector):
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)
with find_work_orders(self.ursula.datastore) as work_orders:
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 @@ class UrsulaInfoMetricsCollector(BaseMetricsCollector):
'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)
with find_policy_arrangements(self.ursula.datastore) as policy_arrangements:
self.metrics["policies_held_gauge"].set(len(policy_arrangements))
self.metrics["host_info"].info(base_payload)

View File

@ -22,6 +22,7 @@ import pytest
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
@ -46,12 +47,8 @@ def test_policy_simple_sinpa(blockchain_ursulas, blockchain_alice, blockchain_bo
for ursula in blockchain_ursulas:
# 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]
except RecordNotFound:
# No records were found; this Ursula didn't have the arrangement.
continue
with find_policy_arrangements(ursula.datastore, writeable=True) as arrangements:
[arrangement.delete() for arrangement in arrangements]
def test_try_to_post_free_arrangement_by_hacking_enact(blockchain_ursulas, blockchain_alice, blockchain_bob, agency,
@ -75,23 +72,19 @@ 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:
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.
with find_policy_arrangements(ursula.datastore, writeable=True) 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.
# Additionally, Ursula logged Amonia as a freerider:
freeriders = ursula.suspicious_activities_witnessed['freeriders']
assert len(freeriders) == 1
assert freeriders[0][0] == amonia
# Additionally, Ursula logged Amonia as a freerider:
freeriders = ursula.suspicious_activities_witnessed['freeriders']
assert len(freeriders) == 1
assert freeriders[0][0] == amonia
# Reset the Ursula for the next test.
ursula.suspicious_activities_witnessed['freeriders'] = []
[arrangement.delete() for arrangement in all_arrangements]
except RecordNotFound:
# No records were found; this Ursula didn't have the arrangement.
continue
# Reset the Ursula for the next test.
ursula.suspicious_activities_witnessed['freeriders'] = []
[arrangement.delete() for arrangement in all_arrangements]
def test_pay_a_flunky_instead_of_the_arranged_ursula(blockchain_alice, blockchain_bob, blockchain_ursulas,
@ -119,23 +112,19 @@ def test_pay_a_flunky_instead_of_the_arranged_ursula(blockchain_alice, blockchai
# Same exact set of assertions as the last test:
for ursula in blockchain_ursulas:
# Even though the grant executed without error...
try:
with ursula.datastore.query_by(PolicyArrangement, writeable=True) 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.
with find_policy_arrangements(ursula.datastore, writeable=True) 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.
# Additionally, Ursula logged Amonia as a freerider:
freeriders = ursula.suspicious_activities_witnessed['freeriders']
assert len(freeriders) == 1
assert freeriders[0][0] == amonia
# Additionally, Ursula logged Amonia as a freerider:
freeriders = ursula.suspicious_activities_witnessed['freeriders']
assert len(freeriders) == 1
assert freeriders[0][0] == amonia
# Reset the Ursula for the next test.
ursula.suspicious_activities_witnessed['freeriders'] = []
[arrangement.delete() for arrangement in all_arrangements]
except RecordNotFound:
# No records were found; this Ursula didn't have the arrangement.
continue
# Reset the Ursula for the next test.
ursula.suspicious_activities_witnessed['freeriders'] = []
[arrangement.delete() for arrangement in all_arrangements]
def test_put_additional_treasure_map_on_network(blockchain_ursulas, blockchain_alice, blockchain_bob, agency, testerchain):