WIP: Transactions

api/extensions/utils.py

@@ -1,7 +1,9 @@
 from pathlib import Path
-from api.app import app
+
 from fastapi.responses import FileResponse
 
+from api.app import app
+
 
 @app.get("/favicon.ico", response_class=FileResponse, include_in_schema=False)
 async def favicon():

classes/errors/__init__.py

@@ -1,2 +1,7 @@
 from .pycord_user import UserNotFoundError
-from .wallet import WalletNotFoundError
+from .wallet import (
+    WalletBalanceLimitExceeded,
+    WalletInsufficientFunds,
+    WalletNotFoundError,
+    WalletOverdraftLimitExceeded,
+)

classes/errors/wallet.py

@@ -24,3 +24,39 @@ class WalletInsufficientFunds(Exception):
         super().__init__(
             f"Wallet of a user with id {self.wallet.owner_id} for the guild with id {self.wallet.guild_id} does not have sufficient funds to perform the operation (balance: {self.wallet.balance}, requested: {self.amount})"
         )
+
+
+class WalletOverdraftLimitExceeded(Exception):
+    """Wallet's overdraft limit is not sufficient to perform the operation"""
+
+    def __init__(
+        self,
+        wallet: "Wallet",
+        amount: float,
+        overdraft_limit: float,
+    ) -> None:
+        self.wallet = wallet
+        self.amount = amount
+        self.overdraft_limit = overdraft_limit
+
+        super().__init__(
+            f"Wallet of a user with id {self.wallet.owner_id} for the guild with id {self.wallet.guild_id} does not have sufficient funds to perform the operation (balance: {self.wallet.balance}, requested: {self.amount}, overdraft limit: {self.overdraft_limit})"
+        )
+
+
+class WalletBalanceLimitExceeded(Exception):
+    """Wallet's balance limit is not high enough to perform the operation"""
+
+    def __init__(
+        self,
+        wallet: "Wallet",
+        amount: float,
+        balance_limit: float,
+    ) -> None:
+        self.wallet = wallet
+        self.amount = amount
+        self.balance_limit = balance_limit
+
+        super().__init__(
+            f"Wallet of a user with id {self.wallet.owner_id} for the guild with id {self.wallet.guild_id} would have too much funds after the operation (balance: {self.wallet.balance}, deposited: {self.amount}, balance limit: {self.balance_limit})"
+        )

classes/pycord_user.py

@@ -82,9 +82,7 @@ class PycordUser:
 
         setattr(self, key, value)
 
-        await col_users.update_one(
+        await col_users.update_one({"_id": self._id}, {"$set": {key: value}}, upsert=True)
-            {"_id": self._id}, {"$set": {key: value}}, upsert=True
-        )
 
         self._update_cache(cache)
 

classes/wallet.py

@@ -7,7 +7,11 @@ from typing import Any, Dict, Optional
 from bson import ObjectId
 from pymongo.results import InsertOneResult
 
-from classes.errors import WalletNotFoundError
+from classes.errors import (
+    WalletBalanceLimitExceeded,
+    WalletNotFoundError,
+    WalletOverdraftLimitExceeded,
+)
 from classes.errors.wallet import WalletInsufficientFunds
 from modules.database import col_wallets
 
@@ -28,9 +32,7 @@ class Wallet:
     async def from_id(
         cls, owner_id: int, guild_id: int, allow_creation: bool = True
     ) -> "Wallet":
-        db_entry = await col_wallets.find_one(
+        db_entry = await col_wallets.find_one({"owner_id": owner_id, "guild_id": guild_id})
-            {"owner_id": owner_id, "guild_id": guild_id}
-        )
 
         if db_entry is None:
             if not allow_creation:
@@ -60,13 +62,9 @@ class Wallet:
 
         setattr(self, key, value)
 
-        await col_wallets.update_one(
+        await col_wallets.update_one({"_id": self._id}, {"$set": {key: value}}, upsert=True)
-            {"_id": self._id}, {"$set": {key: value}}, upsert=True
-        )
 
-        logger.info(
+        logger.info("Set attribute '%s' of the wallet %s to '%s'", key, str(self._id), value)
-            "Set attribute '%s' of the wallet %s to '%s'", key, str(self._id), value
-        )
 
     @staticmethod
     def get_defaults(
@@ -96,8 +94,15 @@ class Wallet:
         await self._set("is_frozen", False)
 
     # TODO Write a docstring
-    async def deposit(self, amount: float) -> None:
+    async def deposit(self, amount: float, balance_limit: Optional[float] = None) -> float:
-        await self._set("balance", round(self.balance + amount, 2))
+        new_balance: float = round(self.balance + amount, 2)
+
+        if balance_limit is not None and new_balance > balance_limit:
+            raise WalletBalanceLimitExceeded(self, amount, balance_limit)
+
+        await self._set("balance", new_balance)
+
+        return new_balance
 
     # TODO Write a docstring
     async def withdraw(
@@ -105,10 +110,47 @@ class Wallet:
         amount: float,
         allow_overdraft: bool = False,
         overdraft_limit: Optional[float] = None,
-    ) -> None:
+    ) -> float:
-        if amount > self.balance and (
+        if amount > self.balance:
-            not allow_overdraft or (allow_overdraft and amount > overdraft_limit)
+            if not allow_overdraft or overdraft_limit is None:
-        ):
+                raise WalletInsufficientFunds(self, amount)
-            raise WalletInsufficientFunds(self, amount)
 
-        await self._set("balance", round(self.balance - amount, 2))
+            if allow_overdraft and amount > overdraft_limit:
+                raise WalletOverdraftLimitExceeded(self, amount, overdraft_limit)
+
+        new_balance: float = round(self.balance - amount, 2)
+
+        await self._set("balance", new_balance)
+
+        return new_balance
+
+    async def transfer(
+        self,
+        wallet_owner_id: int,
+        wallet_guild_id: int,
+        amount: float,
+        balance_limit: Optional[float] = None,
+        allow_overdraft: bool = False,
+        overdraft_limit: Optional[float] = None,
+    ) -> None:
+        # TODO Replace with a concrete exception
+        if amount < 0:
+            raise ValueError()
+
+        wallet: Wallet = await self.from_id(
+            wallet_owner_id, wallet_guild_id, allow_creation=False
+        )
+
+        if balance_limit is not None and amount + wallet.balance > balance_limit:
+            raise WalletBalanceLimitExceeded(wallet, amount, balance_limit)
+
+        if amount > self.balance:
+            if not allow_overdraft or overdraft_limit is None:
+                raise WalletInsufficientFunds(self, amount)
+
+            if allow_overdraft and amount > overdraft_limit:
+                raise WalletOverdraftLimitExceeded(self, amount, overdraft_limit)
+
+        # TODO Make a sanity check to revert the transaction if anything goes wrong
+        await self.withdraw(amount, allow_overdraft, overdraft_limit)
+        await wallet.deposit(amount, balance_limit)

migrate.py

@@ -1,4 +1,3 @@
 from modules.migrator import migrate_database
 
-
 migrate_database()

modules/database.py

@@ -35,6 +35,4 @@ col_wallets: AsyncCollection = db.get_collection("wallets")
 
 # Update indexes
 db.dispatch.get_collection("users").create_index("id", unique=True)
-db.dispatch.get_collection("wallets").create_index(
+db.dispatch.get_collection("wallets").create_index(["owner_id", "guild_id"], unique=False)
-    ["owner_id", "guild_id"], unique=False
-)

modules/extensions_loader.py

@@ -15,9 +15,7 @@ def get_py_files(src: str | Path) -> List[str]:
     py_files = []
 
     for root, dirs, files in walk(src):
-        py_files.extend(
+        py_files.extend(Path(f"{cwd}/{root}/{file}") for file in files if file.endswith(".py"))
-            Path(f"{cwd}/{root}/{file}") for file in files if file.endswith(".py")
-        )
 
     return py_files
 

pyproject.toml

@@ -0,0 +1,6 @@
+[tool.black]
+line-length = 96
+target-version = ["py311"]
+
+[tool.isort]
+profile = "black"