| import sqlite3 |
| import os |
| import json |
| from typing import Optional |
|
|
| try: |
| import psycopg2 |
| from psycopg2 import pool |
| except ImportError: |
| psycopg2 = None |
|
|
| DB_FILE = os.getenv("MEP_SQLITE_PATH", "ledger.db") |
| DB_URL = os.getenv("MEP_DATABASE_URL") |
| PG_POOL_MIN = int(os.getenv("MEP_PG_POOL_MIN", "1")) |
| PG_POOL_MAX = int(os.getenv("MEP_PG_POOL_MAX", "5")) |
| _pg_pool: Optional["pool.SimpleConnectionPool"] = None |
|
|
| def _is_postgres() -> bool: |
| return bool(DB_URL) |
|
|
| def _get_pg_pool(): |
| global _pg_pool |
| if _pg_pool is None: |
| if psycopg2 is None: |
| raise RuntimeError("psycopg2 is required for Postgres") |
| _pg_pool = pool.SimpleConnectionPool(PG_POOL_MIN, PG_POOL_MAX, DB_URL) |
| return _pg_pool |
|
|
| def _get_conn(): |
| if _is_postgres(): |
| return _get_pg_pool().getconn() |
| return sqlite3.connect(DB_FILE, check_same_thread=False) |
|
|
| def _release_conn(conn): |
| if _is_postgres(): |
| _get_pg_pool().putconn(conn) |
| else: |
| conn.close() |
|
|
| def _row_to_dict(cursor, row): |
| if row is None: |
| return None |
| columns = [desc[0] for desc in cursor.description] |
| return dict(zip(columns, row)) |
|
|
| def init_db(): |
| conn = _get_conn() |
| cursor = conn.cursor() |
| cursor.execute(''' |
| CREATE TABLE IF NOT EXISTS ledger ( |
| node_id TEXT PRIMARY KEY, |
| pub_pem TEXT NOT NULL, |
| balance REAL NOT NULL |
| ) |
| ''') |
| cursor.execute(''' |
| CREATE TABLE IF NOT EXISTS tasks ( |
| task_id TEXT PRIMARY KEY, |
| consumer_id TEXT NOT NULL, |
| provider_id TEXT, |
| payload TEXT NOT NULL, |
| bounty REAL NOT NULL, |
| status TEXT NOT NULL, |
| target_node TEXT, |
| model_requirement TEXT, |
| result_payload TEXT, |
| created_at REAL NOT NULL, |
| updated_at REAL NOT NULL |
| ) |
| ''') |
| cursor.execute(''' |
| CREATE TABLE IF NOT EXISTS idempotency ( |
| node_id TEXT NOT NULL, |
| endpoint TEXT NOT NULL, |
| idem_key TEXT NOT NULL, |
| response TEXT NOT NULL, |
| status_code INTEGER NOT NULL, |
| created_at REAL NOT NULL, |
| PRIMARY KEY (node_id, endpoint, idem_key) |
| ) |
| ''') |
| conn.commit() |
| _release_conn(conn) |
|
|
| def register_node(node_id: str, pub_pem: str) -> float: |
| conn = _get_conn() |
| cursor = conn.cursor() |
| if _is_postgres(): |
| cursor.execute("SELECT balance FROM ledger WHERE node_id = %s", (node_id,)) |
| else: |
| cursor.execute("SELECT balance FROM ledger WHERE node_id = ?", (node_id,)) |
| row = cursor.fetchone() |
| if not row: |
| if _is_postgres(): |
| cursor.execute( |
| "INSERT INTO ledger (node_id, pub_pem, balance) VALUES (%s, %s, %s) ON CONFLICT (node_id) DO NOTHING", |
| (node_id, pub_pem, 10.0) |
| ) |
| else: |
| cursor.execute( |
| "INSERT OR IGNORE INTO ledger (node_id, pub_pem, balance) VALUES (?, ?, ?)", |
| (node_id, pub_pem, 10.0) |
| ) |
| conn.commit() |
| if _is_postgres(): |
| cursor.execute("SELECT balance FROM ledger WHERE node_id = %s", (node_id,)) |
| else: |
| cursor.execute("SELECT balance FROM ledger WHERE node_id = ?", (node_id,)) |
| row = cursor.fetchone() |
| _release_conn(conn) |
| return row[0] if row else 10.0 |
|
|
| def get_pub_pem(node_id: str) -> Optional[str]: |
| conn = _get_conn() |
| cursor = conn.cursor() |
| if _is_postgres(): |
| cursor.execute("SELECT pub_pem FROM ledger WHERE node_id = %s", (node_id,)) |
| else: |
| cursor.execute("SELECT pub_pem FROM ledger WHERE node_id = ?", (node_id,)) |
| row = cursor.fetchone() |
| _release_conn(conn) |
| return row[0] if row else None |
|
|
| def get_balance(node_id: str) -> Optional[float]: |
| conn = _get_conn() |
| cursor = conn.cursor() |
| if _is_postgres(): |
| cursor.execute("SELECT balance FROM ledger WHERE node_id = %s", (node_id,)) |
| else: |
| cursor.execute("SELECT balance FROM ledger WHERE node_id = ?", (node_id,)) |
| row = cursor.fetchone() |
| _release_conn(conn) |
| return row[0] if row else None |
|
|
| def set_balance(node_id: str, balance: float): |
| conn = _get_conn() |
| cursor = conn.cursor() |
| if _is_postgres(): |
| cursor.execute("UPDATE ledger SET balance = %s WHERE node_id = %s", (balance, node_id)) |
| else: |
| cursor.execute("UPDATE ledger SET balance = ? WHERE node_id = ?", (balance, node_id)) |
| conn.commit() |
| _release_conn(conn) |
|
|
| def add_balance(node_id: str, amount: float): |
| conn = _get_conn() |
| cursor = conn.cursor() |
| if _is_postgres(): |
| cursor.execute("UPDATE ledger SET balance = balance + %s WHERE node_id = %s", (amount, node_id)) |
| else: |
| cursor.execute("UPDATE ledger SET balance = balance + ? WHERE node_id = ?", (amount, node_id)) |
| conn.commit() |
| _release_conn(conn) |
|
|
| def deduct_balance(node_id: str, amount: float) -> bool: |
| conn = _get_conn() |
| cursor = conn.cursor() |
| if _is_postgres(): |
| cursor.execute( |
| "UPDATE ledger SET balance = balance - %s WHERE node_id = %s AND balance >= %s", |
| (amount, node_id, amount) |
| ) |
| else: |
| cursor.execute( |
| "UPDATE ledger SET balance = balance - ? WHERE node_id = ? AND balance >= ?", |
| (amount, node_id, amount) |
| ) |
| updated = cursor.rowcount |
| conn.commit() |
| _release_conn(conn) |
| return updated > 0 |
|
|
| def create_task(task_id: str, consumer_id: str, payload: str, bounty: float, status: str, target_node: Optional[str], model_requirement: Optional[str], created_at: float): |
| conn = _get_conn() |
| cursor = conn.cursor() |
| if _is_postgres(): |
| cursor.execute( |
| "INSERT INTO tasks (task_id, consumer_id, provider_id, payload, bounty, status, target_node, model_requirement, result_payload, created_at, updated_at) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)", |
| (task_id, consumer_id, None, payload, bounty, status, target_node, model_requirement, None, created_at, created_at) |
| ) |
| else: |
| cursor.execute( |
| "INSERT INTO tasks (task_id, consumer_id, provider_id, payload, bounty, status, target_node, model_requirement, result_payload, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", |
| (task_id, consumer_id, None, payload, bounty, status, target_node, model_requirement, None, created_at, created_at) |
| ) |
| conn.commit() |
| _release_conn(conn) |
|
|
| def update_task_assignment(task_id: str, provider_id: str, status: str, updated_at: float): |
| conn = _get_conn() |
| cursor = conn.cursor() |
| if _is_postgres(): |
| cursor.execute( |
| "UPDATE tasks SET provider_id = %s, status = %s, updated_at = %s WHERE task_id = %s", |
| (provider_id, status, updated_at, task_id) |
| ) |
| else: |
| cursor.execute( |
| "UPDATE tasks SET provider_id = ?, status = ?, updated_at = ? WHERE task_id = ?", |
| (provider_id, status, updated_at, task_id) |
| ) |
| conn.commit() |
| _release_conn(conn) |
|
|
| def update_task_result(task_id: str, provider_id: str, result_payload: str, status: str, updated_at: float): |
| conn = _get_conn() |
| cursor = conn.cursor() |
| if _is_postgres(): |
| cursor.execute( |
| "UPDATE tasks SET provider_id = %s, result_payload = %s, status = %s, updated_at = %s WHERE task_id = %s", |
| (provider_id, result_payload, status, updated_at, task_id) |
| ) |
| else: |
| cursor.execute( |
| "UPDATE tasks SET provider_id = ?, result_payload = ?, status = ?, updated_at = ? WHERE task_id = ?", |
| (provider_id, result_payload, status, updated_at, task_id) |
| ) |
| conn.commit() |
| _release_conn(conn) |
|
|
| def get_task(task_id: str) -> Optional[dict]: |
| conn = _get_conn() |
| if not _is_postgres(): |
| conn.row_factory = sqlite3.Row |
| cursor = conn.cursor() |
| if _is_postgres(): |
| cursor.execute("SELECT * FROM tasks WHERE task_id = %s", (task_id,)) |
| else: |
| cursor.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)) |
| row = cursor.fetchone() |
| if not row: |
| _release_conn(conn) |
| return None |
| if _is_postgres(): |
| result = _row_to_dict(cursor, row) |
| _release_conn(conn) |
| return result |
| result = dict(row) |
| _release_conn(conn) |
| return result |
|
|
| def get_active_tasks() -> list: |
| conn = _get_conn() |
| if not _is_postgres(): |
| conn.row_factory = sqlite3.Row |
| cursor = conn.cursor() |
| if _is_postgres(): |
| cursor.execute("SELECT * FROM tasks WHERE status IN ('bidding', 'assigned')") |
| else: |
| cursor.execute("SELECT * FROM tasks WHERE status IN ('bidding', 'assigned')") |
| rows = cursor.fetchall() |
| if _is_postgres(): |
| result = [_row_to_dict(cursor, row) for row in rows] |
| _release_conn(conn) |
| return result |
| result = [dict(row) for row in rows] |
| _release_conn(conn) |
| return result |
|
|
| def get_idempotency(node_id: str, endpoint: str, idem_key: str) -> Optional[dict]: |
| conn = _get_conn() |
| cursor = conn.cursor() |
| if _is_postgres(): |
| cursor.execute( |
| "SELECT response, status_code FROM idempotency WHERE node_id = %s AND endpoint = %s AND idem_key = %s", |
| (node_id, endpoint, idem_key) |
| ) |
| else: |
| cursor.execute( |
| "SELECT response, status_code FROM idempotency WHERE node_id = ? AND endpoint = ? AND idem_key = ?", |
| (node_id, endpoint, idem_key) |
| ) |
| row = cursor.fetchone() |
| if not row: |
| _release_conn(conn) |
| return None |
| response = json.loads(row[0]) |
| result = {"response": response, "status_code": row[1]} |
| _release_conn(conn) |
| return result |
|
|
| def set_idempotency(node_id: str, endpoint: str, idem_key: str, response: dict, status_code: int, created_at: float): |
| conn = _get_conn() |
| cursor = conn.cursor() |
| payload = json.dumps(response) |
| if _is_postgres(): |
| cursor.execute( |
| "INSERT INTO idempotency (node_id, endpoint, idem_key, response, status_code, created_at) VALUES (%s, %s, %s, %s, %s, %s) ON CONFLICT (node_id, endpoint, idem_key) DO NOTHING", |
| (node_id, endpoint, idem_key, payload, status_code, created_at) |
| ) |
| else: |
| cursor.execute( |
| "INSERT OR IGNORE INTO idempotency (node_id, endpoint, idem_key, response, status_code, created_at) VALUES (?, ?, ?, ?, ?, ?)", |
| (node_id, endpoint, idem_key, payload, status_code, created_at) |
| ) |
| conn.commit() |
| _release_conn(conn) |
|
|
| init_db() |
|
|
