reserved balance wip

This commit is contained in:
Shroominic
2025-08-19 16:54:17 -03:00
parent 0270bf2ca6
commit 4a338505cc
3 changed files with 55 additions and 7 deletions
@@ -0,0 +1,33 @@
"""introduce reserved balance
Revision ID: 042f6b77d69d
Revises: 898f00ea481e
Create Date: 2025-08-18 19:03:09.507368
"""
import sqlalchemy as sa
import sqlmodel
from alembic import op
# revision identifiers, used by Alembic.
revision = "042f6b77d69d"
down_revision = "898f00ea481e"
branch_labels = None
depends_on = None
def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.alter_column("api_keys", "balance", new_column_name="total_balance")
op.add_column(
"api_keys",
sa.Column("reserved_balance", sa.Integer(), nullable=False, server_default="0"),
)
# ### end Alembic commands ###
def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.drop_column("api_keys", "reserved_balance")
op.alter_column("api_keys", "total_balance", new_column_name="balance")
# ### end Alembic commands ###
+12 -6
View File
@@ -1,4 +1,5 @@
import hashlib
import math
from typing import Optional
from fastapi import HTTPException
@@ -323,10 +324,9 @@ async def pay_for_request(key: ApiKey, session: AsyncSession, body: dict) -> int
stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.balance) >= cost_per_request)
.where(col(ApiKey.total_balance) >= cost_per_request)
.values(
balance=col(ApiKey.balance) - cost_per_request,
total_spent=col(ApiKey.total_spent) + cost_per_request,
reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
total_requests=col(ApiKey.total_requests) + 1,
)
)
@@ -438,6 +438,7 @@ async def adjust_payment_for_tokens(
# If token-based pricing is enabled and base cost is 0, use token-based cost
# Otherwise, token cost is additional to the base cost
cost_difference = cost.total_msats - deducted_max_cost
total_cost_msats: int = math.ceil(cost.total_msats)
logger.info(
"Calculated token-based cost",
@@ -460,6 +461,7 @@ async def adjust_payment_for_tokens(
await session.commit()
return cost.dict()
# this should never happen why do we handle this???
if cost_difference > 0:
# Need to charge more
logger.info(
@@ -473,6 +475,7 @@ async def adjust_payment_for_tokens(
},
)
# this should never happen why do we handle this???
if key.balance < cost_difference:
logger.warning(
"Insufficient balance for token-based pricing adjustment",
@@ -486,6 +489,7 @@ async def adjust_payment_for_tokens(
)
await session.commit()
else:
# this should never happen why do we handle this???
charge_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
@@ -538,13 +542,15 @@ async def adjust_payment_for_tokens(
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.values(
balance=col(ApiKey.balance) + refund,
total_spent=col(ApiKey.total_spent) - refund,
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost,
total_balance=col(ApiKey.total_balance) - total_cost_msats,
total_spent=col(ApiKey.total_spent) + total_cost_msats,
)
)
await session.exec(refund_stmt) # type: ignore[call-overload]
await session.commit()
cost.total_msats = deducted_max_cost - refund
cost.total_msats = total_cost_msats
await session.refresh(key)
logger.info(
+10 -1
View File
@@ -22,7 +22,12 @@ class ApiKey(SQLModel, table=True): # type: ignore
__tablename__ = "api_keys"
hashed_key: str = Field(primary_key=True)
balance: int = Field(default=0, description="Balance in millisatoshis (msats)")
total_balance: int = Field(
default=0, description="Total balance in millisatoshis (msats)"
)
reserved_balance: int = Field(
default=0, description="Reserved balance in millisatoshis (msats)"
)
refund_address: str | None = Field(
default=None,
description="Lightning address to refund remaining balance after key expires",
@@ -44,6 +49,10 @@ class ApiKey(SQLModel, table=True): # type: ignore
description="Currency of the cashu-token",
)
@property
def balance(self) -> int:
return self.total_balance - self.reserved_balance
async def balances_for_mint_and_unit(
db_session: AsyncSession, mint_url: str, unit: str