82 lines
3.2 KiB
Python
82 lines
3.2 KiB
Python
import aiosqlite
|
|
from pathlib import Path
|
|
|
|
from config import get_settings
|
|
|
|
|
|
async def get_db() -> aiosqlite.Connection:
|
|
settings = get_settings()
|
|
path = Path(settings.database_path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
db = await aiosqlite.connect(path)
|
|
db.row_factory = aiosqlite.Row
|
|
await db.execute("PRAGMA foreign_keys = ON")
|
|
return db
|
|
|
|
|
|
async def _ensure_column(db: aiosqlite.Connection, table: str, column: str, ddl: str) -> None:
|
|
cursor = await db.execute(f"PRAGMA table_info({table})")
|
|
cols = {row[1] for row in await cursor.fetchall()}
|
|
if column not in cols:
|
|
await db.execute(f"ALTER TABLE {table} ADD COLUMN {column} {ddl}")
|
|
|
|
|
|
async def init_db() -> None:
|
|
db = await get_db()
|
|
try:
|
|
await db.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS guest_links (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
token TEXT NOT NULL UNIQUE,
|
|
link_type TEXT NOT NULL CHECK(link_type IN ('extend', 'traffic')),
|
|
title TEXT NOT NULL DEFAULT '',
|
|
value INTEGER NOT NULL,
|
|
max_uses INTEGER NOT NULL DEFAULT 1,
|
|
used_count INTEGER NOT NULL DEFAULT 0,
|
|
is_active INTEGER NOT NULL DEFAULT 1,
|
|
expires_at TEXT,
|
|
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
|
note TEXT NOT NULL DEFAULT '',
|
|
duration_days INTEGER NOT NULL DEFAULT 0,
|
|
plan_traffic_gb INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS guest_claims (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
link_id INTEGER NOT NULL,
|
|
username TEXT NOT NULL,
|
|
user_uuid TEXT,
|
|
detail TEXT NOT NULL DEFAULT '',
|
|
claimed_at TEXT NOT NULL DEFAULT (datetime('now')),
|
|
FOREIGN KEY (link_id) REFERENCES guest_links(id) ON DELETE CASCADE
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS traffic_grants (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
claim_id INTEGER,
|
|
link_id INTEGER NOT NULL,
|
|
user_uuid TEXT NOT NULL,
|
|
username TEXT NOT NULL DEFAULT '',
|
|
bytes_added INTEGER NOT NULL,
|
|
expires_at TEXT NOT NULL,
|
|
reverted_at TEXT,
|
|
status TEXT NOT NULL DEFAULT 'active'
|
|
CHECK(status IN ('active', 'reverted', 'error')),
|
|
last_error TEXT NOT NULL DEFAULT '',
|
|
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
|
FOREIGN KEY (link_id) REFERENCES guest_links(id) ON DELETE CASCADE
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_guest_links_token ON guest_links(token);
|
|
CREATE INDEX IF NOT EXISTS idx_guest_claims_link ON guest_claims(link_id);
|
|
CREATE INDEX IF NOT EXISTS idx_traffic_grants_active
|
|
ON traffic_grants(status, expires_at);
|
|
"""
|
|
)
|
|
await _ensure_column(db, "guest_links", "duration_days", "INTEGER NOT NULL DEFAULT 0")
|
|
await _ensure_column(db, "guest_links", "plan_traffic_gb", "INTEGER NOT NULL DEFAULT 0")
|
|
await db.commit()
|
|
finally:
|
|
await db.close()
|