From 90490e649ceb453725a2a00e2adc45945c2f3a31 Mon Sep 17 00:00:00 2001 From: "ashkan@beigi.net" Date: Thu, 2 Jul 2026 03:13:34 +0000 Subject: [PATCH] Phase 2: pre-registered batch + atomic claim (DB layer) Add a subscribers.status column ('available'|'claimed') with migration for legacy DBs, plus add_available_subscriber, count_available, and claim_available (atomic BEGIN IMMEDIATE select+update so concurrent claims can't double-hand-out a profile). On-demand add_subscriber stays 'claimed'. Tests: +7; 115 pass. Co-Authored-By: Claude Opus 4.8 --- server/database.py | 71 ++++++++++++++++++++++++++++++++++++++++-- server/models.py | 1 + tests/test_database.py | 68 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 138 insertions(+), 2 deletions(-) diff --git a/server/database.py b/server/database.py index c8528d6..34e9c45 100644 --- a/server/database.py +++ b/server/database.py @@ -7,6 +7,8 @@ from server.models import Subscriber logger = logging.getLogger(__name__) +# status: 'claimed' (on-demand / already handed out) or 'available' (pre-registered +# batch, waiting to be claimed by a device). _SCHEMA = """ CREATE TABLE IF NOT EXISTS subscribers ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -16,7 +18,8 @@ CREATE TABLE IF NOT EXISTS subscribers ( opc TEXT NOT NULL, mac_address TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - synced_to_core BOOLEAN DEFAULT 0 + synced_to_core BOOLEAN DEFAULT 0, + status TEXT NOT NULL DEFAULT 'claimed' ); CREATE INDEX IF NOT EXISTS idx_imsi ON subscribers(imsi); CREATE INDEX IF NOT EXISTS idx_mac_address ON subscribers(mac_address); @@ -24,10 +27,18 @@ CREATE INDEX IF NOT EXISTS idx_mac_address ON subscribers(mac_address); def init_db(db_path: str) -> None: - """Create the database schema if it does not already exist.""" + """Create the database schema if it does not already exist (and migrate).""" conn = sqlite3.connect(db_path) try: conn.executescript(_SCHEMA) + # Migrate pre-existing databases that lack the status column. This must + # happen before creating the status index (a legacy table has neither). + cols = {row[1] for row in conn.execute("PRAGMA table_info(subscribers)")} + if "status" not in cols: + conn.execute( + "ALTER TABLE subscribers ADD COLUMN status TEXT NOT NULL DEFAULT 'claimed'" + ) + conn.execute("CREATE INDEX IF NOT EXISTS idx_status ON subscribers(status)") conn.commit() logger.info(f"Database initialized at {db_path}") finally: @@ -59,6 +70,7 @@ def _row_to_subscriber(row: sqlite3.Row) -> Subscriber: mac_address=row["mac_address"], created_at=created, synced_to_core=bool(row["synced_to_core"]), + status=row["status"], ) @@ -88,6 +100,60 @@ def add_subscriber( raise ValueError(f"Duplicate subscriber (imsi or iccid already exists): {exc}") from exc +def add_available_subscriber( + conn: sqlite3.Connection, imsi: str, iccid: str, ki: str, opc: str, +) -> int: + """Insert a pre-registered ('available') subscriber with no MAC assigned yet. + + Raises ValueError if imsi or iccid already exists. + """ + try: + cur = conn.execute( + """ + INSERT INTO subscribers (imsi, iccid, ki, opc, mac_address, status) + VALUES (?, ?, ?, ?, NULL, 'available') + """, + (imsi, iccid, ki, opc), + ) + conn.commit() + return cur.lastrowid + except sqlite3.IntegrityError as exc: + raise ValueError(f"Duplicate subscriber (imsi or iccid already exists): {exc}") from exc + + +def count_available(conn: sqlite3.Connection) -> int: + """Return the number of unclaimed pre-registered subscribers.""" + return conn.execute( + "SELECT COUNT(*) FROM subscribers WHERE status = 'available'" + ).fetchone()[0] + + +def claim_available(conn: sqlite3.Connection, mac_address: str) -> Optional[Subscriber]: + """Atomically claim the oldest available subscriber for ``mac_address``. + + Returns the claimed Subscriber, or None if none are available. Uses a write + transaction so concurrent claims cannot hand out the same profile twice. + """ + try: + conn.execute("BEGIN IMMEDIATE") + row = conn.execute( + "SELECT id FROM subscribers WHERE status = 'available' ORDER BY id LIMIT 1" + ).fetchone() + if row is None: + conn.commit() + return None + sub_id = row["id"] + conn.execute( + "UPDATE subscribers SET status = 'claimed', mac_address = ? WHERE id = ?", + (mac_address, sub_id), + ) + conn.commit() + except Exception: + conn.rollback() + raise + return get_subscriber_by_id(conn, sub_id) + + def get_subscriber_by_imsi(conn: sqlite3.Connection, imsi: str) -> Optional[Subscriber]: row = conn.execute("SELECT * FROM subscribers WHERE imsi = ?", (imsi,)).fetchone() return _row_to_subscriber(row) if row else None @@ -145,5 +211,6 @@ def export_all(conn: sqlite3.Connection) -> list: "mac_address": row["mac_address"], "created_at": str(row["created_at"]), "synced_to_core": bool(row["synced_to_core"]), + "status": row["status"], }) return result diff --git a/server/models.py b/server/models.py index 1cac631..0c0dd76 100644 --- a/server/models.py +++ b/server/models.py @@ -13,3 +13,4 @@ class Subscriber: mac_address: Optional[str] created_at: datetime synced_to_core: bool + status: str = "claimed" diff --git a/tests/test_database.py b/tests/test_database.py index c92c9a5..3a9a4e3 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -2,8 +2,11 @@ import pytest from datetime import datetime, timedelta from server.database import ( + add_available_subscriber, add_subscriber, check_rate_limit, + claim_available, + count_available, export_all, get_connection, get_subscriber_by_id, @@ -108,6 +111,71 @@ class TestSync: assert sub.synced_to_core is True +class TestBatchClaim: + def test_available_defaults_and_count(self, db_conn): + assert count_available(db_conn) == 0 + add_available_subscriber(db_conn, "302720000000010", "8910302720000000010", + "aa" * 16, "bb" * 16) + assert count_available(db_conn) == 1 + sub = get_subscriber_by_imsi(db_conn, "302720000000010") + assert sub.status == "available" + assert sub.mac_address is None + + def test_on_demand_add_is_claimed(self, db_conn): + add_subscriber(db_conn, **_sub()) + sub = get_subscriber_by_imsi(db_conn, _sub()["imsi"]) + assert sub.status == "claimed" + + def test_claim_assigns_and_decrements(self, db_conn): + add_available_subscriber(db_conn, "302720000000010", "8910302720000000010", + "aa" * 16, "bb" * 16) + claimed = claim_available(db_conn, "aa:bb:cc:dd:ee:ff") + assert claimed is not None + assert claimed.status == "claimed" + assert claimed.mac_address == "aa:bb:cc:dd:ee:ff" + assert count_available(db_conn) == 0 + + def test_claim_returns_none_when_empty(self, db_conn): + assert claim_available(db_conn, "aa:bb:cc:dd:ee:ff") is None + + def test_claim_is_fifo_and_unique(self, db_conn): + for i in range(3): + add_available_subscriber(db_conn, f"30272000000002{i}", + f"891030272000000002{i}", "aa" * 16, "bb" * 16) + first = claim_available(db_conn, "11:11:11:11:11:11") + second = claim_available(db_conn, "22:22:22:22:22:22") + assert first.imsi != second.imsi + assert first.id < second.id # oldest first + assert count_available(db_conn) == 1 + + +class TestMigration: + def test_adds_status_column_to_legacy_db(self, tmp_path): + import sqlite3 + db_path = str(tmp_path / "legacy.db") + legacy = sqlite3.connect(db_path) + legacy.executescript( + "CREATE TABLE subscribers (id INTEGER PRIMARY KEY AUTOINCREMENT, " + "imsi TEXT UNIQUE NOT NULL, iccid TEXT UNIQUE NOT NULL, ki TEXT NOT NULL, " + "opc TEXT NOT NULL, mac_address TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, " + "synced_to_core BOOLEAN DEFAULT 0);" + ) + legacy.execute( + "INSERT INTO subscribers (imsi, iccid, ki, opc) VALUES " + "('302720000000099', '8910302720000000099', 'aa', 'bb')" + ) + legacy.commit() + legacy.close() + + init_db(db_path) # should add the status column, defaulting old rows to 'claimed' + conn = get_connection(db_path) + try: + sub = get_subscriber_by_imsi(conn, "302720000000099") + assert sub.status == "claimed" + finally: + conn.close() + + class TestExport: def test_export_returns_list_with_correct_keys(self, db_conn): add_subscriber(db_conn, **_sub())