diff --git a/lnbits/core/crud.py b/lnbits/core/crud.py index 5573f44a9..3c837a515 100644 --- a/lnbits/core/crud.py +++ b/lnbits/core/crud.py @@ -31,6 +31,7 @@ from .models import ( TinyURL, User, Wallet, + WalletBalance, WebPushSubscription, ) @@ -419,7 +420,7 @@ async def delete_unused_wallets( async def get_wallet( wallet_id: str, deleted: Optional[bool] = None, conn: Optional[Connection] = None -) -> Optional[Wallet]: +) -> Optional[WalletBalance]: where = "AND deleted = :deleted" if deleted is not None else "" return await (conn or db).fetchone( f""" @@ -429,13 +430,13 @@ async def get_wallet( WHERE id = :wallet {where} """, {"wallet": wallet_id, "deleted": deleted}, - Wallet, + WalletBalance, ) async def get_wallets( user_id: str, deleted: Optional[bool] = None, conn: Optional[Connection] = None -) -> list[Wallet]: +) -> list[WalletBalance]: where = "AND deleted = :deleted" if deleted is not None else "" return await (conn or db).fetchall( f""" @@ -445,14 +446,14 @@ async def get_wallets( WHERE "user" = :user {where} """, {"user": user_id, "deleted": deleted}, - Wallet, + WalletBalance, ) async def get_wallet_for_key( key: str, conn: Optional[Connection] = None, -) -> Optional[Wallet]: +) -> Optional[WalletBalance]: return await (conn or db).fetchone( """ SELECT *, COALESCE(( @@ -462,7 +463,7 @@ async def get_wallet_for_key( WHERE (adminkey = :key OR inkey = :key) AND deleted = false """, {"key": key}, - Wallet, + WalletBalance, ) diff --git a/lnbits/core/models.py b/lnbits/core/models.py index 2d1cc0e0e..2093e6df0 100644 --- a/lnbits/core/models.py +++ b/lnbits/core/models.py @@ -42,25 +42,10 @@ class Wallet(BaseModel): adminkey: str inkey: str deleted: bool = False - balance_msat: int = 0 created_at: datetime = datetime.now(timezone.utc) updated_at: datetime = datetime.now(timezone.utc) currency: Optional[str] = None - # @property - # def balance_msat(self) -> int: - # return self.balance_msat // 1000 - - @property - def balance(self) -> int: - return self.balance_msat // 1000 - - @property - def withdrawable_balance(self) -> int: - from .services import fee_reserve - - return self.balance_msat - fee_reserve(self.balance_msat) - @property def lnurlwithdraw_full(self) -> str: url = url_for("/withdraw", external=True, usr=self.user, wal=self.id) @@ -78,6 +63,22 @@ class Wallet(BaseModel): ) +class WalletBalance(Wallet): + """Wallet with balance properties""" + + balance_msat: int = 0 + + @property + def balance(self) -> int: + return self.balance_msat // 1000 + + @property + def withdrawable_balance(self) -> int: + from .services import fee_reserve + + return self.balance_msat - fee_reserve(self.balance_msat) + + class KeyType(Enum): admin = 0 invoice = 1 @@ -91,7 +92,7 @@ class KeyType(Enum): @dataclass class WalletTypeInfo: key_type: KeyType - wallet: Wallet + wallet: WalletBalance class UserExtra(BaseModel): @@ -173,7 +174,7 @@ class User(BaseModel): username: Optional[str] = None pubkey: Optional[str] = None extensions: list[str] = [] - wallets: list[Wallet] = [] + wallets: list[WalletBalance] = [] admin: bool = False super_user: bool = False has_password: bool = False @@ -275,13 +276,6 @@ class CreatePayment(BaseModel): fee: int = 0 -# TODO: thanks about paymeny extra -# class PaymentExtra(Extra): -# tag: Optional[str] = None -# def __getitem__(self, key): -# return self[key] or self._raw_json[key] - - class Payment(BaseModel): status: str checking_id: str diff --git a/lnbits/core/services.py b/lnbits/core/services.py index 4c847a0e3..2e47063a1 100644 --- a/lnbits/core/services.py +++ b/lnbits/core/services.py @@ -78,7 +78,7 @@ from .models import ( PaymentState, User, UserExtra, - Wallet, + WalletBalance, ) @@ -452,7 +452,7 @@ async def _create_external_payment( def _check_wallet_balance( - wallet: Wallet, + wallet: WalletBalance, fee_reserve_total_msat: int, internal_checking_id: Optional[str] = None, ): @@ -701,7 +701,7 @@ def fee_reserve_total(amount_msat: int, internal: bool = False) -> int: return fee_reserve(amount_msat, internal) + service_fee(amount_msat, internal) -async def send_payment_notification(wallet: Wallet, payment: Payment): +async def send_payment_notification(wallet: WalletBalance, payment: Payment): await websocket_updater( wallet.inkey, json.dumps( diff --git a/lnbits/core/views/generic.py b/lnbits/core/views/generic.py index 85694cac3..7db37af89 100644 --- a/lnbits/core/views/generic.py +++ b/lnbits/core/views/generic.py @@ -13,7 +13,7 @@ from pydantic.types import UUID4 from lnbits.core.extensions.models import Extension, ExtensionMeta, InstallableExtension from lnbits.core.helpers import to_valid_user_id -from lnbits.core.models import User +from lnbits.core.models import User, WalletBalance from lnbits.core.services import create_invoice, create_user_account from lnbits.decorators import check_admin, check_user_exists from lnbits.helpers import template_renderer @@ -167,7 +167,8 @@ async def wallet( if wal: wallet = await get_wallet(wal.hex) elif len(user.wallets) == 0: - wallet = await create_wallet(user_id=user.id) + _wallet = await create_wallet(user_id=user.id) + wallet = WalletBalance(**_wallet.dict()) user.wallets.append(wallet) elif lnbits_last_active_wallet and user.get_wallet(lnbits_last_active_wallet): wallet = await get_wallet(lnbits_last_active_wallet) diff --git a/lnbits/core/views/user_api.py b/lnbits/core/views/user_api.py index b4780ee48..ccd527146 100644 --- a/lnbits/core/views/user_api.py +++ b/lnbits/core/views/user_api.py @@ -21,7 +21,7 @@ from lnbits.core.models import ( AccountOverview, CreateTopup, User, - Wallet, + WalletBalance, ) from lnbits.core.services import update_wallet_balance from lnbits.db import Filters, Page @@ -103,7 +103,7 @@ async def api_users_toggle_admin(user_id: str) -> None: @users_router.get("/user/{user_id}/wallet") -async def api_users_get_user_wallet(user_id: str) -> List[Wallet]: +async def api_users_get_user_wallet(user_id: str) -> List[WalletBalance]: return await get_wallets(user_id) diff --git a/lnbits/core/views/wallet_api.py b/lnbits/core/views/wallet_api.py index 428d2a28d..4ce1ee9b1 100644 --- a/lnbits/core/views/wallet_api.py +++ b/lnbits/core/views/wallet_api.py @@ -12,6 +12,7 @@ from lnbits.core.models import ( CreateWallet, KeyType, Wallet, + WalletBalance, ) from lnbits.decorators import ( WalletTypeInfo, @@ -61,7 +62,7 @@ async def api_update_wallet( name: Optional[str] = Body(None), currency: Optional[str] = Body(None), key_info: WalletTypeInfo = Depends(require_admin_key), -) -> Wallet: +) -> WalletBalance: wallet = await get_wallet(key_info.wallet.id) if not wallet: raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Wallet not found") @@ -81,7 +82,7 @@ async def api_delete_wallet( ) -@wallet_router.post("", response_model=Wallet) +@wallet_router.post("") async def api_create_wallet( data: CreateWallet, wallet: WalletTypeInfo = Depends(require_admin_key),