Coverage for polar/models/wallet.py: 90%

1from enum import StrEnum 

2from typing import TYPE_CHECKING 

3from uuid import UUID 

4 

5from sqlalchemy import ForeignKey, String, UniqueConstraint, Uuid, func, select 

6from sqlalchemy.ext.associationproxy import AssociationProxy, association_proxy 

7from sqlalchemy.orm import ( 

8 Mapped, 

9 column_property, 

10 declared_attr, 

11 mapped_column, 

12 relationship, 

13) 

14 

15from polar.kit.db.models.base import RecordModel 

16from polar.kit.extensions.sqlalchemy.types import StringEnum 

17 

18if TYPE_CHECKING: 18 ↛ 19line 18 didn't jump to line 19 because the condition on line 18 was never true

19 from .customer import Customer 

20 from .organization import Organization 

21 

22 

23class WalletType(StrEnum): 

24 usage = "usage" 

25 billing = "billing" 

26 

27 

28class Wallet(RecordModel): 

29 __tablename__ = "wallets" 

30 __table_args__ = (UniqueConstraint("type", "currency", "customer_id"),) 

31 

32 type: Mapped[WalletType] = mapped_column(StringEnum(WalletType), nullable=False) 

33 currency: Mapped[str] = mapped_column(String(3)) 

34 customer_id: Mapped[UUID] = mapped_column( 

35 Uuid, 

36 ForeignKey("customers.id", ondelete="cascade"), 

37 unique=True, 

38 ) 

39 

40 @declared_attr 

41 def customer(cls) -> Mapped["Customer"]: 

42 return relationship("Customer", lazy="raise_on_sql") 

43 

44 organization: AssociationProxy["Organization"] = association_proxy( 

45 "customer", "organization" 

46 ) 

47 

48 @declared_attr 

49 def balance(cls) -> Mapped[int]: 

50 from .wallet_transaction import WalletTransaction 

51 

52 return column_property( 

53 select(func.coalesce(func.sum(WalletTransaction.amount), 0)) 

54 .where(WalletTransaction.wallet_id == cls.id) 

55 .correlate_except(WalletTransaction) 

56 .scalar_subquery() 

57 )