fixup crud and models

This commit is contained in:
Arc 2024-11-19 13:18:17 +00:00
commit 9503a909eb
2 changed files with 81 additions and 56 deletions

110
crud.py
View file

@ -4,7 +4,16 @@ from typing import List, Optional
from lnbits.db import Database from lnbits.db import Database
from .execution_queue import enqueue from .execution_queue import enqueue
from .models import NWCBudget, NWCKey, CreateNWCKey from .models import (
NWCBudget,
NWCKey,
CreateNWCKey,
GetNWCKey,
GetWalletNWCKey,
GetBudgetsNWC,
TrackedSpendNWC,
DeleteNWC
)
db = Database("ext_nwcprovider") db = Database("ext_nwcprovider")
@ -30,46 +39,35 @@ async def create_nwc(data: CreateNWCKey) -> NWCKey:
await db.insert("nwcprovider.budgets", budget_entry) await db.insert("nwcprovider.budgets", budget_entry)
return NWCKey(**data.dict()) return NWCKey(**data.dict())
async def delete_nwc(data:DeleteNWC) -> None:
async def delete_nwc(pubkey: str, wallet_id: str):
nwc = await get_nwc(pubkey, wallet_id)
if not nwc:
raise Exception("Public key does not exist")
await db.execute( await db.execute(
""" "DELETE FROM nwcprovider.keys WHERE pubkey = :pubkey AND wallet = :wallet", {"pubkey": data.pubkey, "wallet": data.wallet_id}
DELETE FROM nwcprovider.keys WHERE pubkey = ? AND wallet = ?
""",
(pubkey, wallet_id),
) )
async def get_wallet_nwcs(data: GetWalletNWCKey) -> List[NWCKey]:
async def get_wallet_nwcs( return await db.fetchall(
wallet_id: str, include_expired: Optional[bool] = False
) -> List[NWCKey]:
rows = await db.fetchall(
""" """
SELECT * FROM nwcprovider.keys SELECT * FROM nwcprovider.keys
WHERE wallet = ? AND (expires_at = 0 OR expires_at > ?) WHERE wallet = :wallet AND (expires_at = 0 OR expires_at > :expires)
""", """,
(wallet_id, int(time.time()) if not include_expired else -1), {"wallet": data.wallet_id, "expires": int(time.time()) if not data.include_expired else -1},
model=NWCKey,
) )
return [NWCKey(**row) for row in rows]
async def get_nwc( async def get_nwc(
pubkey: str, data: GetNWCKey
wallet_id: Optional[str] = None,
include_expired: Optional[bool] = False,
refresh_last_used: Optional[bool] = False,
) -> Optional[NWCKey]: ) -> Optional[NWCKey]:
# expires_at = 0 means it never expires # expires_at = 0 means it never expires
if wallet_id: if data.wallet_id:
row = await db.fetchone( row = await db.fetchone(
""" """
SELECT * FROM nwcprovider.keys SELECT * FROM nwcprovider.keys
WHERE pubkey = ? AND wallet = ? AND (expires_at = 0 OR expires_at > ?) WHERE pubkey = :pubkey AND wallet = :wallet
AND (expires_at = 0 OR expires_at > :expires)
""", """,
(pubkey, wallet_id, int(time.time()) if not include_expired else -1), {"pubkey": data.pubkey, "wallet": data.wallet_id,
"expires": int(time.time()) if not data.include_expired else -1},
NWCKey,
) )
else: else:
row = await db.fetchone( row = await db.fetchone(
@ -77,44 +75,52 @@ async def get_nwc(
SELECT * FROM nwcprovider.keys SELECT * FROM nwcprovider.keys
WHERE pubkey = ? AND (expires_at = 0 OR expires_at > ?) WHERE pubkey = ? AND (expires_at = 0 OR expires_at > ?)
""", """,
(pubkey, int(time.time()) if not include_expired else -1), (data.pubkey, int(time.time()) if not data.include_expired else -1),
)
row = await db.fetchone(
"""
SELECT * FROM nwcprovider.keys
WHERE pubkey = :pubkey AND (expires_at = 0 OR expires_at > :expires)
""",
{"pubkey": data.pubkey, "expires": int(time.time()) if not data.include_expired else -1},
NWCKey,
) )
if not row: if not row:
return None return None
if refresh_last_used: if data.refresh_last_used:
await db.execute( await db.execute(
""" """
UPDATE nwcprovider.keys SET last_used = ? WHERE pubkey = ? UPDATE nwcprovider.keys SET last_used = :last_used WHERE pubkey = :pubkey
""", """,
(int(time.time()), pubkey), {"last_used":int(time.time()), "pubkey":data.pubkey},
) )
return NWCKey(**row) return NWCKey(**row)
async def get_budgets_nwc(pubkey, calculate_spent=False): async def get_budgets_nwc(data: GetBudgetsNWC) -> Optional[NWCBudget]:
rows = await db.fetchall( rows = await db.fetchall(
"SELECT * FROM nwcprovider.budgets WHERE pubkey = ?", (pubkey) "SELECT * FROM nwcprovider.budgets WHERE pubkey = :pubkey", {"pubkey":data.pubkey}
) )
budgets = [NWCBudget(**row) for row in rows] budgets = [NWCBudget(**row) for row in rows]
if calculate_spent: if data.calculate_spent:
for budget in budgets: for budget in budgets:
last_cycle, next_cycle = budget.get_timestamp_range() last_cycle, next_cycle = budget.get_timestamp_range()
tot_spent_in_range_msats = await db.fetchone( tot_spent_in_range_msats = await db.fetchone(
""" """
SELECT SUM(amount_msats) FROM nwcprovider.spent SELECT SUM(amount_msats) FROM nwcprovider.spent
WHERE pubkey = ? AND created_at >= ? AND created_at < ? WHERE pubkey = :pubkey AND created_at >= :last_cycle AND created_at < next_cycle
""", """,
(pubkey, last_cycle, next_cycle), {"pubkey":data.pubkey, "last_cycle":last_cycle, "next_cycle":next_cycle},
) )
tot_spent_in_range_msats = tot_spent_in_range_msats[0] or 0 tot_spent_in_range_msats = tot_spent_in_range_msats[0] or 0
budget.used_budget_msats = tot_spent_in_range_msats budget.used_budget_msats = tot_spent_in_range_msats
return budgets return budgets
async def tracked_spend_nwc(pubkey: str, amount_msats: int, action): async def tracked_spend_nwc(data:TrackedSpendNWC, action):
async def r(): async def r():
created_at = int(time.time()) created_at = int(time.time())
budgets = await get_budgets_nwc(pubkey) budgets = await get_budgets_nwc(data.pubkey)
in_budget = True in_budget = True
for budget in budgets: for budget in budgets:
last_cycle, next_cycle = budget.get_timestamp_range() last_cycle, next_cycle = budget.get_timestamp_range()
@ -123,14 +129,14 @@ async def tracked_spend_nwc(pubkey: str, amount_msats: int, action):
await db.fetchone( await db.fetchone(
""" """
SELECT SUM(amount_msats) FROM nwcprovider.spent SELECT SUM(amount_msats) FROM nwcprovider.spent
WHERE pubkey = ? AND created_at >= ? AND created_at < ? WHERE pubkey = :pubkey AND created_at >= :last_cycle AND created_at < :next_cycle
""", """,
(pubkey, last_cycle, next_cycle), {"pubkey":data.pubkey, "last_cycle":last_cycle, "next_cycle":next_cycle},
) )
)[0] )[0]
or 0 or 0
) )
if tot_spent_in_range_msats + amount_msats > budget.budget_msats: if tot_spent_in_range_msats + data.amount_msats > budget.budget_msats:
in_budget = False in_budget = False
break break
if not in_budget: if not in_budget:
@ -139,9 +145,9 @@ async def tracked_spend_nwc(pubkey: str, amount_msats: int, action):
await db.execute( await db.execute(
""" """
INSERT INTO nwcprovider.spent (pubkey, amount_msats, created_at) INSERT INTO nwcprovider.spent (pubkey, amount_msats, created_at)
VALUES (?, ?, ?) VALUES (:pubkey, :amount_msats, :created_at)
""", """,
(pubkey, amount_msats, created_at), {"pubkey":data.pubkey, "amount_msats":data.amount_msats, "created_at":created_at},
) )
return True, out return True, out
@ -149,29 +155,27 @@ async def tracked_spend_nwc(pubkey: str, amount_msats: int, action):
async def get_config_nwc(key: str): async def get_config_nwc(key: str):
row = await db.fetchone("SELECT * FROM nwcprovider.config WHERE key = ?", (key,)) row = await db.fetchone("SELECT * FROM nwcprovider.config WHERE key = :key", {"key":key})
if not row: if not row:
return None return None
return row["value"] return row["value"]
async def get_all_config_nwc():
rows = await db.fetchall("SELECT * FROM nwcprovider.config")
return {row["key"]: row["value"] for row in rows}
async def set_config_nwc(key: str, value: str): async def set_config_nwc(key: str, value: str):
await db.execute( await db.execute(
""" """
DELETE FROM nwcprovider.config DELETE FROM nwcprovider.config
WHERE key = ? WHERE key = :key
""", """,
(key,), {"key":key},
) )
await db.execute( await db.execute(
""" """
INSERT INTO nwcprovider.config (key, value) INSERT INTO nwcprovider.config (key, value)
VALUES (?, ?) VALUES (:key, :value)
""", """,
(key, value), {"key":key, "value":value},
) )
async def get_all_config_nwc():
rows = await db.fetchall("SELECT * FROM nwcprovider.config")
return {row["key"]: row["value"] for row in rows}

View file

@ -26,6 +26,27 @@ class NWCKey(BaseModel):
def from_row(cls, row: Dict[str, Any]) -> "NWCKey": def from_row(cls, row: Dict[str, Any]) -> "NWCKey":
return cls(**row) return cls(**row)
class GetNWCKey(BaseModel):
pubkey: str
wallet_id: Optional[str] = None
include_expired: Optional[bool] = False
refresh_last_used: Optional[bool] = False
class GetWalletNWCKey(BaseModel):
wallet_id: Optional[str] = None
include_expired: Optional[bool] = False
class GetBudgetsNWC(BaseModel):
pubkey: str
calculate_spent: Optional[bool] = False
class TrackedSpendNWC(BaseModel):
pubkey: str
amount_msats: int
class DeleteNWC(BaseModel):
pubkey: str
wallet_id: str
class NWCBudget(BaseModel): class NWCBudget(BaseModel):
id: int id: int