Files
routstr-core/proxy/db.py
T
2025-04-25 13:28:11 +08:00

45 lines
1.4 KiB
Python

from contextlib import asynccontextmanager
import os
from typing import AsyncGenerator
from sqlmodel import Field, SQLModel
from sqlalchemy.ext.asyncio.engine import create_async_engine
from sqlmodel.ext.asyncio.session import AsyncSession
DATABASE_URL = os.environ.get("DATABASE_URL", "sqlite+aiosqlite:///keys.db")
engine = create_async_engine(DATABASE_URL, echo=False) # echo=True for debugging SQL
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)")
refund_address: str | None = Field(
default=None,
description="Lightning address to refund remaining balance after key expires",
)
total_spent: int = Field(
default=0, description="Total spent in millisatoshis (msats)"
)
total_requests: int = Field(default=0)
async def init_db() -> None:
"""Initializes the database and creates tables if they don't exist."""
async with engine.begin() as conn:
await conn.run_sync(SQLModel.metadata.create_all)
async def get_session() -> AsyncGenerator[AsyncSession, None]:
async with AsyncSession(engine, expire_on_commit=False) as session:
yield session
@asynccontextmanager
async def create_session() -> AsyncGenerator[AsyncSession, None]:
async with AsyncSession(engine, expire_on_commit=False) as session:
yield session