diff --git a/README.md b/README.md index 936e0e9a..1dbc5d41 100644 --- a/README.md +++ b/README.md @@ -2,3 +2,16 @@ a reverse proxy that you can plug in front of any openai compatible api endpoint to handle payments using the cashu protocol (Bitcoin L3) + +## Database Migrations + +Alembic is used to manage the database schema for the `ApiKey` model defined in +`router/db.py`. Before running the application for the first time or after +pulling updates, apply the migrations with: + +```bash +alembic upgrade head +``` + +The configuration reads the `DATABASE_URL` environment variable to determine the +database connection. diff --git a/alembic.ini b/alembic.ini new file mode 100644 index 00000000..e5675df3 --- /dev/null +++ b/alembic.ini @@ -0,0 +1,35 @@ +[alembic] +script_location = migrations +sqlalchemy.url = + +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARN +handlers = console + +[logger_sqlalchemy] +level = WARN +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s diff --git a/migrations/env.py b/migrations/env.py new file mode 100644 index 00000000..10909f6b --- /dev/null +++ b/migrations/env.py @@ -0,0 +1,55 @@ +import asyncio +from logging.config import fileConfig +from sqlalchemy import pool +from sqlalchemy.ext.asyncio import create_async_engine +from alembic import context +from sqlmodel import SQLModel +import importlib.util +import pathlib + +db_path = pathlib.Path(__file__).resolve().parents[1] / "router" / "db.py" +spec = importlib.util.spec_from_file_location("db", db_path) +db = importlib.util.module_from_spec(spec) +spec.loader.exec_module(db) + +DATABASE_URL = getattr(db, "DATABASE_URL", "sqlite+aiosqlite:///keys.db") + +config = context.config +fileConfig(config.config_file_name) +config.set_main_option("sqlalchemy.url", DATABASE_URL) + +target_metadata = SQLModel.metadata + + +def run_migrations_offline() -> None: + context.configure( + url=DATABASE_URL, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def do_run_migrations(connection): + context.configure(connection=connection, target_metadata=target_metadata, compare_type=True) + + with context.begin_transaction(): + context.run_migrations() + + +async def run_migrations_online() -> None: + connectable = create_async_engine(DATABASE_URL, poolclass=pool.NullPool) + + async with connectable.connect() as connection: + await connection.run_sync(do_run_migrations) + + await connectable.dispose() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + asyncio.run(run_migrations_online()) diff --git a/migrations/script.py.mako b/migrations/script.py.mako new file mode 100644 index 00000000..7b2652ae --- /dev/null +++ b/migrations/script.py.mako @@ -0,0 +1,19 @@ +<%text># -*- coding: utf-8 -*- +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} +""" + +from alembic import op +import sqlalchemy as sa + +${imports if imports else ""} + +def upgrade(): + ${upgrades if upgrades else "pass"} + + +def downgrade(): + ${downgrades if downgrades else "pass"} diff --git a/migrations/versions/8eaf6006fa28_create_api_keys_table.py b/migrations/versions/8eaf6006fa28_create_api_keys_table.py new file mode 100644 index 00000000..2aff990b --- /dev/null +++ b/migrations/versions/8eaf6006fa28_create_api_keys_table.py @@ -0,0 +1,28 @@ +"""create api keys table + +Revision ID: 8eaf6006fa28 +Revises: +Create Date: 2025-06-06 13:47:00.000000 +""" +revision = "8eaf6006fa28" +down_revision = None +branch_labels = None +depends_on = None +from alembic import op +import sqlalchemy as sa + + +def upgrade() -> None: + op.create_table( + "api_keys", + sa.Column("hashed_key", sa.String(), primary_key=True, nullable=False), + sa.Column("balance", sa.Integer(), nullable=False, server_default="0"), + sa.Column("refund_address", sa.String(), nullable=True), + sa.Column("key_expiry_time", sa.Integer(), nullable=True), + sa.Column("total_spent", sa.Integer(), nullable=False, server_default="0"), + sa.Column("total_requests", sa.Integer(), nullable=False, server_default="0"), + ) + + +def downgrade() -> None: + op.drop_table("api_keys") diff --git a/migrations/versions/__init__.py b/migrations/versions/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/pyproject.toml b/pyproject.toml index abba2da8..deee9af1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "sqlmodel>=0.0.24", "httpx[socks]>=0.25.2", "greenlet>=3.2.1", + "alembic>=1.13", ] [dependency-groups] diff --git a/router/models.py b/router/models.py index 5bb50186..53493840 100644 --- a/router/models.py +++ b/router/models.py @@ -1,5 +1,6 @@ import asyncio import json +import os from pydantic.v1 import BaseModel from .price import sats_usd_ask_price @@ -42,7 +43,11 @@ class Model(BaseModel): MODELS: list[Model] = [] -with open("models.json", "r") as f: +models_file = "models.json" +if not os.path.exists(models_file): + models_file = "models.example.json" + +with open(models_file, "r") as f: MODELS = [Model(**model) for model in json.load(f)["models"]]