[PAY-1892] Implement weekly pool limit (#6289)

packages/discovery-provider/ddl/migrations/0036_add_challenge_disbursement_cap.sql

@@ -0,0 +1,12 @@
+begin;
+
+alter table challenge_disbursements 
+add column created_at timestamp with time zone default now();
+
+create index idx_challenge_disbursements_created_at 
+on challenge_disbursements (created_at);
+
+alter table challenges
+add column weekly_pool integer default null;
+
+commit;

packages/discovery-provider/integration_tests/queries/test_get_attestation.py

@@ -8,8 +8,11 @@
 from web3 import Web3
 
 from integration_tests.queries.test_get_challenges import setup_db
+from integration_tests.utils import populate_mock_db
+from src.models.rewards.challenge import ChallengeType
 from src.queries.get_attestation import (
     ADD_SENDER_MESSAGE_PREFIX,
+    POOL_EXHAUSTED,
     REWARDS_MANAGER_ACCOUNT,
     Attestation,
     AttestationError,
@@ -123,6 +126,149 @@ def test_get_attestation(app):
                 )
 
 
+def test_get_attestation_weekly_pool_exhausted(app):
+    with app.app_context():
+        db = get_db()
+
+    entities = {
+        "users": [
+            {"user_id": 1, "wallet": "0x38C68fF3926bf4E68289672F75ee1543117dDAAA"},
+            {"user_id": 2, "wallet": "0x38C68fF3926bf4E68289672F75ee1543117dD9B3"},
+            {"user_id": 3, "wallet": "0xDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEF"},
+        ],
+        "challenges": [
+            {
+                "id": "not-exhausted",
+                "type": ChallengeType.numeric,
+                "active": True,
+                "amount": 5,
+                "weekly_pool": 10,
+            },
+            {
+                "id": "exhausted",
+                "type": ChallengeType.numeric,
+                "active": True,
+                "amount": 10,
+                "weekly_pool": 20,
+            },
+        ],
+        "user_challenges": [
+            {
+                "challenge_id": "not-exhausted",
+                "user_id": 1,
+                "specifier": "1",
+                "is_complete": True,
+                "amount": 5,
+            },
+            {
+                "challenge_id": "not-exhausted",
+                "user_id": 2,
+                "specifier": "2",
+                "is_complete": True,
+                "amount": 5,
+            },
+            {
+                "challenge_id": "exhausted",
+                "user_id": 1,
+                "specifier": "1",
+                "is_complete": True,
+                "amount": 10,
+            },
+            {
+                "challenge_id": "exhausted",
+                "user_id": 2,
+                "specifier": "2",
+                "is_complete": True,
+                "amount": 10,
+            },
+            # User 3 has completed challenge `exhausted` but will not be able to disburse
+            {
+                "challenge_id": "exhausted",
+                "user_id": 3,
+                "specifier": "3",
+                "is_complete": True,
+                "amount": 10,
+            },
+        ],
+        "challenge_disbursements": [
+            {
+                "challenge_id": "not-exhausted",
+                "specifier": "1",
+                "user_id": 1,
+                "amount": "5",
+            },
+            {
+                "challenge_id": "exhausted",
+                "specifier": "1",
+                "user_id": 1,
+                "amount": "10",
+            },
+            {
+                "challenge_id": "exhausted",
+                "specifier": "2",
+                "user_id": 2,
+                "amount": "10",
+            },
+        ],
+    }
+    populate_mock_db(db, entities)
+    with db.scoped_session() as session:
+        oracle_address = "0x32a10e91820fd10366AC363eD0DEa40B2e598D22"
+        redis_handle.set(oracle_addresses_key, oracle_address)
+
+        delegate_owner_wallet, signature = get_attestation(
+            session,
+            user_id=2,
+            challenge_id="not-exhausted",
+            oracle_address=oracle_address,
+            specifier="2",
+        )
+
+        attestation = Attestation(
+            amount="5",
+            oracle_address=oracle_address,
+            user_address="0x38C68fF3926bf4E68289672F75ee1543117dD9B3",
+            challenge_id="not-exhausted",
+            challenge_specifier="2",
+        )
+
+        # User 2 should still be able to disburse the non-exhausted challenge
+
+        # confirm the attestation is what we think it should be
+        config_owner_wallet = shared_config["delegate"]["owner_wallet"]
+        config_private_key = shared_config["delegate"]["private_key"]
+
+        # Ensure we returned the correct owner wallet
+        assert delegate_owner_wallet == config_owner_wallet
+
+        # Ensure we can derive the owner wallet from the signed stringified attestation
+        attestation_bytes = attestation.get_attestation_bytes()
+        to_sign_hash = Web3.keccak(attestation_bytes)
+        private_key = keys.PrivateKey(HexBytes(config_private_key))
+        public_key = keys.PublicKey.from_private(private_key)
+        signture_bytes = to_bytes(hexstr=signature)
+        msg_signature = keys.Signature(signature_bytes=signture_bytes, vrs=None)
+
+        recovered_pubkey = public_key.recover_from_msg_hash(
+            message_hash=to_sign_hash, signature=msg_signature
+        )
+
+        assert (
+            Web3.to_checksum_address(recovered_pubkey.to_address())
+            == config_owner_wallet
+        )
+
+        with pytest.raises(AttestationError) as e:
+            get_attestation(
+                session,
+                user_id=3,
+                challenge_id="exhausted",
+                oracle_address=oracle_address,
+                specifier="3",
+            )
+        assert str(e.value) == POOL_EXHAUSTED
+
+
 @pytest.fixture
 def patch_get_all_other_nodes():
     with patch(

packages/discovery-provider/integration_tests/queries/test_get_disbursed_challenges_amount.py

@@ -0,0 +1,52 @@
+from datetime import datetime, timedelta
+
+from integration_tests.utils import populate_mock_db
+from src.queries.get_disbursed_challenges_amount import get_disbursed_challenges_amount
+from src.utils.db_session import get_db
+
+
+def test_get_disbursed_challenges_amount(app):
+    with app.app_context():
+        db = get_db()
+
+    entities = {
+        "challenge_disbursements": [
+            {
+                "challenge_id": "profile-completion",
+                "specifier": "1",
+                "user_id": 1,
+                "amount": "5",
+            },
+            {
+                "challenge_id": "profile-completion",
+                "specifier": "2",
+                "user_id": 2,
+                "amount": "10",
+            },
+            {
+                "challenge_id": "profile-completion",
+                "specifier": "3",
+                "user_id": 3,
+                "amount": "20",
+                "created_at": datetime.now() - timedelta(days=8),
+            },
+            {
+                "challenge_id": "listen-streak",
+                "specifier": "1",
+                "user_id": 1,
+                "amount": "4",
+            },
+            {
+                "challenge_id": "listen-streak",
+                "specifier": "2",
+                "user_id": 2,
+                "amount": "12",
+            },
+        ],
+    }
+    populate_mock_db(db, entities)
+    with db.scoped_session() as session:
+        amount_disbursed = get_disbursed_challenges_amount(
+            session, "profile-completion", datetime.now() - timedelta(days=7)
+        )
+        assert amount_disbursed == 15

packages/discovery-provider/integration_tests/utils.py

@@ -502,6 +502,7 @@ def populate_mock_db(db, entities, block_offset=None):
                 active=challenge_meta.get("active", True),
                 step_count=challenge_meta.get("step_count", None),
                 starting_block=challenge_meta.get("starting_block", None),
+                weekly_pool=challenge_meta.get("weekly_pool", None),
             )
             session.add(challenge)
         for i, user_challenge_meta in enumerate(user_challenges):
@@ -514,6 +515,7 @@ def populate_mock_db(db, entities, block_offset=None):
                     "completed_blocknumber", 1 + block_offset
                 ),
                 current_step_count=user_challenge_meta.get("current_step_count", None),
+                amount=user_challenge_meta.get("amount", None),
             )
             session.add(user_challenge)
         for i, user_bank_account in enumerate(user_bank_accounts):
@@ -596,6 +598,7 @@ def populate_mock_db(db, entities, block_offset=None):
                 signature=challenge_disbursement.get("signature", str(i)),
                 slot=challenge_disbursement.get("slot", i),
                 amount=challenge_disbursement.get("amount", i),
+                created_at=challenge_disbursement.get("created_at", datetime.now()),
             )
             session.add(cb)
         for i, playlist_seen in enumerate(playlist_seens):

packages/discovery-provider/src/challenges/challenges.dev.json

@@ -80,7 +80,8 @@
     "amount": 1,
     "active": true,
     "step_count": 2147483647,
-    "starting_block": 0
+    "starting_block": 0,
+    "weekly_pool": 25000
   },
   {
     "id": "s",
@@ -89,6 +90,7 @@
     "amount": 1,
     "active": true,
     "step_count": 2147483647,
-    "starting_block": 0
+    "starting_block": 0,
+    "weekly_pool": 25000
   }
 ]

packages/discovery-provider/src/challenges/challenges.json

@@ -102,7 +102,8 @@
     "amount": 1,
     "active": true,
     "step_count": 2147483647,
-    "starting_block": 220157041
+    "starting_block": 220157041,
+    "weekly_pool": 25000
   },
   {
     "id": "s",
@@ -111,6 +112,7 @@
     "amount": 1,
     "active": true,
     "step_count": 2147483647,
-    "starting_block": 220157041
+    "starting_block": 220157041,
+    "weekly_pool": 25000
   }
 ]

packages/discovery-provider/src/challenges/challenges.stage.json

@@ -80,7 +80,8 @@
     "amount": 1,
     "active": true,
     "step_count": 2147483647,
-    "starting_block": 0
+    "starting_block": 0,
+    "weekly_pool": 25000
   },
   {
     "id": "s",
@@ -89,6 +90,7 @@
     "amount": 1,
     "active": true,
     "step_count": 2147483647,
-    "starting_block": 0
+    "starting_block": 0,
+    "weekly_pool": 25000
   }
 ]

packages/discovery-provider/src/challenges/create_new_challenges.py

@@ -17,6 +17,7 @@ class ChallengeJSON(TypedDict):
     active: bool
     step_count: Optional[int]
     starting_block: Optional[int]
+    weekly_pool: Optional[int]
 
 
 class OverrideChallengeJson(TypedDict):
@@ -25,6 +26,7 @@ class OverrideChallengeJson(TypedDict):
     active: Optional[bool]
     step_count: Optional[int]
     starting_block: Optional[int]
+    weekly_pool: Optional[int]
 
 
 def get_challenges_dicts():
@@ -51,7 +53,14 @@ def get_challenges_dicts():
                 override = stage_challenges_map.get(challenge["id"])
                 if not override:
                     continue
-                for key in ["id", "amount", "active", "starting_block", "step_count"]:
+                for key in [
+                    "id",
+                    "amount",
+                    "active",
+                    "starting_block",
+                    "step_count",
+                    "weekly_pool",
+                ]:
                     if key in override:
                         challenge[key] = override[key]
 
@@ -88,6 +97,7 @@ def create_new_challenges(session, allowed_challenge_types=None):
             active=challenge_dict.get("active"),
             starting_block=challenge_dict.get("starting_block"),
             step_count=challenge_dict.get("step_count"),
+            weekly_pool=challenge_dict.get("weekly_pool"),
         )
         challenges.append(challenge)
     session.add_all(challenges)
@@ -103,3 +113,4 @@ def create_new_challenges(session, allowed_challenge_types=None):
             existing.amount = challenge_dict.get("amount")
             existing.step_count = challenge_dict.get("step_count")
             existing.starting_block = challenge_dict.get("starting_block")
+            existing.weekly_pool = challenge_dict.get("weekly_pool")

packages/discovery-provider/src/models/rewards/challenge.py

@@ -35,3 +35,5 @@ class Challenge(Base, RepresentableMixin):
     # if set, events emitted prior to the starting_block
     # will be ignord.
     starting_block = Column(Integer)
+    # Optional field to support challenges with a weekly pool
+    weekly_pool = Column(Integer)

packages/discovery-provider/src/models/rewards/challenge_disbursement.py

@@ -1,4 +1,4 @@
-from sqlalchemy import Column, Integer, String
+from sqlalchemy import Column, DateTime, Integer, String, text
 
 from src.models.base import Base
 from src.models.model_utils import RepresentableMixin
@@ -13,3 +13,6 @@ class ChallengeDisbursement(Base, RepresentableMixin):
     signature = Column(String, nullable=False)
     slot = Column(Integer, nullable=False, index=True)
     amount = Column(String, nullable=False)
+    created_at = Column(
+        DateTime, nullable=False, server_default=text("CURRENT_TIMESTAMP"), index=True
+    )