Cleaned up the code and removed complexity for easier reading

backend/app/settings.py

@@ -117,7 +117,10 @@ class DevConfig(Config):
     CHAIN_ID = int(os.getenv("CHAIN_ID", 1337))
     # Put the db file in project root
     DB_PATH = os.path.join(Config.PROJECT_ROOT, DB_NAME)
-    SQLALCHEMY_DATABASE_URI = f"sqlite:///{DB_PATH}"
+    # SQLALCHEMY_DATABASE_URI = f"sqlite:///{DB_PATH}"
+    SQLALCHEMY_DATABASE_URI = (
+        f"postgresql://postgres:mysecretpassword@localhost:5433/postgres"
+    )
     SUBGRAPH_RETRY_TIMEOUT_SEC = 2
     X_REAL_IP_REQUIRED = parse_bool(os.getenv("X_REAL_IP_REQUIRED", "false"))
     CACHE_TYPE = "SimpleCache"

backend/v2/core/dependencies.py

@@ -1,4 +1,5 @@
 from functools import lru_cache
+import time
 from typing import Annotated, AsyncGenerator
 
 from fastapi import Depends
@@ -8,6 +9,8 @@
 from web3 import AsyncHTTPProvider, AsyncWeb3
 from web3.middleware import async_geth_poa_middleware
 
+from app.shared.blockchain_types import ChainTypes
+
 
 class OctantSettings(BaseSettings):
     model_config = SettingsConfigDict(env_file=".env", extra="ignore", frozen=True)
@@ -107,6 +110,10 @@ class ChainSettings(OctantSettings):
         description="The chain id to use for the signature verification.",
     )
 
+    @property
+    def is_mainnet(self) -> bool:
+        return self.chain_id == ChainTypes.MAINNET
+
 
 def get_chain_settings() -> ChainSettings:
     return ChainSettings()
@@ -127,7 +134,12 @@ def get_socketio_settings() -> SocketioSettings:
     return SocketioSettings()  # type: ignore[call-arg]
 
 
+def get_current_timestamp() -> int:
+    return int(time.time())
+
+
 GetSocketioSettings = Annotated[SocketioSettings, Depends(get_socketio_settings)]
 GetChainSettings = Annotated[ChainSettings, Depends(get_chain_settings)]
 Web3 = Annotated[AsyncWeb3, Depends(get_w3)]
 GetSession = Annotated[AsyncSession, Depends(get_db_session)]
+GetCurrentTimestamp = Annotated[int, Depends(get_current_timestamp)]

backend/v2/deposits/dependencies.py

@@ -1,10 +1,14 @@
 from typing import Annotated
 
 from fastapi import Depends
+from app.constants import (
+    SABLIER_UNLOCK_GRACE_PERIOD_24_HRS,
+    TEST_SABLIER_UNLOCK_GRACE_PERIOD_15_MIN,
+)
 from v2.deposits.repositories import DepositEventsRepository
 from v2.epochs.dependencies import GetEpochsSubgraph
 from v2.sablier.dependencies import GetSablierSubgraph
-from v2.core.dependencies import GetSession, OctantSettings, Web3
+from v2.core.dependencies import GetChainSettings, GetSession, OctantSettings, Web3
 from v2.deposits.contracts import DEPOSITS_ABI, DepositsContracts
 
 
@@ -22,15 +26,28 @@ def get_deposits_contracts(
     return DepositsContracts(w3, DEPOSITS_ABI, settings.deposits_contract_address)  # type: ignore[arg-type]
 
 
-
 def get_deposit_events_repository(
     session: GetSession,
     epochs_subgraph: GetEpochsSubgraph,
     sublier_subgraph: GetSablierSubgraph,
+    chain_settings: GetChainSettings,
 ) -> DepositEventsRepository:
-    return DepositEventsRepository(session, epochs_subgraph, sublier_subgraph)
+    sablier_unlock_grace_period = (
+        SABLIER_UNLOCK_GRACE_PERIOD_24_HRS
+        if chain_settings.is_mainnet
+        else TEST_SABLIER_UNLOCK_GRACE_PERIOD_15_MIN
+    )
+
+    return DepositEventsRepository(
+        session,
+        epochs_subgraph,
+        sublier_subgraph,
+        sablier_unlock_grace_period,
+    )
 
 
 # Annotated dependencies
 GetDepositsContracts = Annotated[DepositsContracts, Depends(get_deposits_contracts)]
-GetDepositEventsRepository = Annotated[DepositEventsRepository, Depends(get_deposit_events_repository)]
+GetDepositEventsRepository = Annotated[
+    DepositEventsRepository, Depends(get_deposit_events_repository)
+]

backend/v2/deposits/repositories.py

@@ -1,62 +1,53 @@
-
-
 from itertools import groupby
 from operator import attrgetter
-import os
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import selectinload
-from typing import Dict
 
-from app.engine.octant_rewards import OctantRewardsSettings
-from app.engine.user.budget.with_ppf import UserBudgetWithPPF
-from app.engine.user.effective_deposit import DepositEvent, EventType, UserDeposit, UserEffectiveDepositPayload
-from app.engine.user.effective_deposit.weighted_average.default_with_sablier_timebox import DefaultWeightedAverageWithSablierTimebox
+from app.engine.user.effective_deposit import (
+    DepositEvent,
+    EventType,
+)
 from app.infrastructure.database.models import Deposit
-from app.infrastructure.graphql.locks import get_locks_by_timestamp_range
-from app.infrastructure.graphql.unlocks import get_unlocks_by_timestamp_range
-from app.modules.common.sablier_events_mapper import FlattenStrategy, flatten_sablier_events, process_to_locks_and_unlocks
-from app.modules.dto import OctantRewardsDTO, PendingSnapshotDTO
-from app.modules.octant_rewards.core import calculate_rewards
-from app.modules.snapshots.pending.core import calculate_user_budgets
-from app.modules.user.budgets.core import get_upcoming_budget
+from app.modules.common.sablier_events_mapper import (
+    FlattenStrategy,
+    flatten_sablier_events,
+    process_to_locks_and_unlocks,
+)
 from app.modules.user.events_generator.core import unify_deposit_balances
-from app.modules.staking.proceeds.core import estimate_staking_proceeds
-from app.constants import SABLIER_SENDER_ADDRESS_SEPOLIA, SABLIER_TOKEN_ADDRESS_SEPOLIA, ZERO_ADDRESS
-from app.infrastructure import SubgraphEndpoints
+from v2.core.types import Address
+from v2.core.transformers import transform_to_checksum_address
 from v2.epochs.subgraphs import EpochsSubgraph
 from v2.sablier.subgraphs import SablierSubgraph
 
 
-
 async def get_all_deposit_events_for_epoch(
     session: AsyncSession,
     epoch_number: int,
-) -> dict[str, list[DepositEvent]]:
-
+) -> dict[Address, Deposit]:
     results = await session.scalars(
         select(Deposit)
         .options(selectinload(Deposit.user))
         .where(Deposit.epoch == epoch_number)
     )
 
     return {
-        result.user.address: result for result in results
+        transform_to_checksum_address(result.user.address): result for result in results
     }
 
 
-
-
 class DepositEventsRepository:
     def __init__(
-            self, 
-            session: AsyncSession, 
-            epochs_subgraph: EpochsSubgraph, 
-            sablier_subgraph: SablierSubgraph
-        ):
+        self,
+        session: AsyncSession,
+        epochs_subgraph: EpochsSubgraph,
+        sablier_subgraph: SablierSubgraph,
+        sablier_unlock_grace_period: int,
+    ):
         self.session = session
         self.epochs_subgraph = epochs_subgraph
         self.sablier_subgraph = sablier_subgraph
+        self.sablier_unlock_grace_period = sablier_unlock_grace_period
 
     async def get_all_users_events(
         self,
@@ -67,7 +58,7 @@ async def get_all_users_events(
         """
         Returns all user events (LOCK, UNLOCK) for a given epoch.
         """
-        
+
         # Get all locked amounts for the previous epoch
         epoch_start_locked_amounts = await get_all_deposit_events_for_epoch(
             self.session, epoch_number - 1
@@ -90,8 +81,12 @@ async def get_all_users_events(
         # print("Mapped streams", mapped_streams)
         epoch_events = []
         epoch_events += flatten_sablier_events(mapped_streams, FlattenStrategy.ALL)
-        epoch_events += await self.epochs_subgraph.fetch_locks_by_timestamp_range(start_sec, end_sec)
-        epoch_events += await self.epochs_subgraph.fetch_unlocks_by_timestamp_range(start_sec, end_sec)
+        epoch_events += await self.epochs_subgraph.fetch_locks_by_timestamp_range(
+            start_sec, end_sec
+        )
+        epoch_events += await self.epochs_subgraph.fetch_unlocks_by_timestamp_range(
+            start_sec, end_sec
+        )
 
         epoch_events = [DepositEvent.from_dict(event) for event in epoch_events]
         sorted_events = sorted(epoch_events, key=attrgetter("user", "timestamp"))
@@ -121,7 +116,8 @@ async def get_all_users_events(
                 )
 
             user_events[user_address] = unify_deposit_balances(
-                user_events[user_address]
+                user_events[user_address],
+                self.sablier_unlock_grace_period,
             )
 
         return user_events

backend/v2/epochs/contracts.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Dict, Optional
+from typing import Optional
 
 from v2.core.contracts import SmartContract
 from web3 import exceptions

backend/v2/epochs/dependencies.py

@@ -89,7 +89,6 @@ async def get_rewards_rate(epoch_number: int) -> float:
 async def get_epoch_state(
     session: GetSession, epochs_contracts: GetEpochsContracts, epoch_number: int
 ) -> EpochState:
-
     current_epoch_number = await epochs_contracts.get_current_epoch()
     if epoch_number > current_epoch_number:
         return EpochState.FUTURE

backend/v2/epochs/subgraphs.py

@@ -78,7 +78,9 @@ def __init__(
                 self.gql_client.execute_async
             )
 
-    async def fetch_locks_by_timestamp_range(self, from_ts: int, to_ts: int) -> list[LockEvent]:
+    async def fetch_locks_by_timestamp_range(
+        self, from_ts: int, to_ts: int
+    ) -> list[LockEvent]:
         """
         Get locks by timestamp range.
         """
@@ -108,7 +110,9 @@ async def fetch_locks_by_timestamp_range(self, from_ts: int, to_ts: int) -> list
         response = await self.gql_client.execute_async(query, variable_values=variables)
         return response["lockeds"]
 
-    async def fetch_unlocks_by_timestamp_range(self, from_ts: int, to_ts: int) -> list[UnlockEvent]:
+    async def fetch_unlocks_by_timestamp_range(
+        self, from_ts: int, to_ts: int
+    ) -> list[UnlockEvent]:
         """
         Get unlocks by timestamp range.
         """
@@ -154,9 +158,9 @@ async def fetch_epoch_by_number(self, epoch_number: int) -> EpochSubgraphItem:
             """
         )
         variables = {"epochNo": epoch_number}
-        
+
         response = await self.gql_client.execute_async(query, variable_values=variables)
-        
+
         data = response["epoches"]
         if not data:
             raise exceptions.EpochNotIndexed(epoch_number)

backend/v2/glms/dependencies.py

@@ -2,7 +2,6 @@
 from typing import Annotated
 
 from fastapi import Depends
-from v2.deposits.contracts import DepositsContracts
 from v2.deposits.dependencies import GetDepositsContracts
 from v2.core.dependencies import OctantSettings, Web3
 from v2.glms.contracts import ERC20_ABI, GLMContracts
@@ -30,4 +29,4 @@ async def get_glm_balance_of_deposits(
 
 
 GetGLMContracts = Annotated[GLMContracts, Depends(get_glm_contracts)]
-GetGLMBalanceOfDeposits = Annotated[int, Depends(get_glm_balance_of_deposits)]
+GetGLMBalanceOfDeposits = Annotated[int, Depends(get_glm_balance_of_deposits)]

backend/v2/project_rewards/dependencies.py

@@ -1,8 +1,7 @@
 from typing import Annotated
 
 from fastapi import Depends
-from v2.sablier.subgraphs import SablierSubgraph
-from v2.core.dependencies import GetSession, OctantSettings
+from v2.core.dependencies import GetSession
 from v2.epochs.dependencies import GetOpenAllocationWindowEpochNumber
 from v2.matched_rewards.dependencies import GetMatchedRewardsEstimator
 from v2.project_rewards.services import ProjectRewardsEstimator
@@ -22,6 +21,7 @@ async def get_project_rewards_estimator(
         epoch_number=epoch_number,
     )
 
+
 GetProjectRewardsEstimator = Annotated[
     ProjectRewardsEstimator,
     Depends(get_project_rewards_estimator),

backend/v2/project_rewards/repositories.py

@@ -1,7 +1,7 @@
-from email.headerregistry import Address
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 from app.infrastructure.database.models import Reward
+from v2.core.types import Address
 
 
 async def get_rewards_for_epoch(