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

1from enum import StrEnum 2 ctx1ab

2from typing import TYPE_CHECKING 2 ctx1ab

3from uuid import UUID 2 ctx1ab

4 

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

6from sqlalchemy.ext.associationproxy import AssociationProxy, association_proxy 2 ctx1ab

7from sqlalchemy.orm import ( 2 ctx1ab

8 Mapped, 

9 column_property, 

10 declared_attr, 

11 mapped_column, 

12 relationship, 

13) 

14 

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

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

17 

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

19 from .customer import Customer 

20 from .organization import Organization 

21 

22 

23class WalletType(StrEnum): 2 ctx1ab

24 usage = "usage" 2 ctx1ab

25 billing = "billing" 2 ctx1ab

26 

27 

28class Wallet(RecordModel): 2 ctx1ab

29 __tablename__ = "wallets" 2 ctx1ab

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

31 

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

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

34 customer_id: Mapped[UUID] = mapped_column( 2 ctx1ab

35 Uuid, 

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

37 unique=True, 

38 ) 

39 

40 @declared_attr 2 ctx1ab

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

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

43 

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

45 "customer", "organization" 

46 ) 

47 

48 @declared_attr 2 ctx1ab

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

50 from .wallet_transaction import WalletTransaction 2 ctx1ab

51 

52 return column_property( 2 ctx1ab

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

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

55 .correlate_except(WalletTransaction) 

56 .scalar_subquery() 

57 )