qoherent/ria-ran--sim-drop
23
0
Fork 0
Code Issues Pull Requests Workflows Packages Projects Releases Wiki Activity

Phase 2: pre-registered batch + atomic claim (DB layer)

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
A ashkan@beigi.net 2026-07-02 03:13:34 +00:00
parent 2dc1520b34
commit 90490e649c
3 changed files with 138 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

71
server/database.py
View File

@ -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

1
server/models.py
View File

@ -13,3 +13,4 @@ class Subscriber:
mac_address: Optional[str]
created_at: datetime
synced_to_core: bool
status: str = "claimed"

68
tests/test_database.py
View File

@ -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())