Spaces:
Paused
Paused
| """ | |
| Database access for the whole app: SQLite (default, local file) or MySQL (permanent, external server). | |
| MySQL is switched on by the MYSQL_HOST variable. The rest of the code keeps its SQLite-style SQL | |
| (`?` placeholders, `datetime('now','-7 days')`, `INSERT OR REPLACE`, ...); in MySQL mode `translate()` | |
| rewrites it, so both back ends run the same queries. | |
| Variables (Space secrets): MYSQL_HOST, MYSQL_PORT (3306; Hugging Face only allows 80/443/8080 outbound, so the | |
| server may have to listen on 8080), MYSQL_USER, MYSQL_PASSWORD, MYSQL_DB, MYSQL_SSL_CA (path, optional), | |
| MYSQL_SSL=1 (TLS without a CA file), MYSQL_POOL (default 8). | |
| """ | |
| import os | |
| import queue | |
| import re | |
| import sqlite3 | |
| import threading | |
| import time | |
| from pathlib import Path | |
| MYSQL_HOST = os.getenv("MYSQL_HOST", "").strip() | |
| IS_MYSQL = bool(MYSQL_HOST) | |
| _SQLITE_PATH = None # set by auth.py (it decides between /data and the app folder) | |
| # --------------------------------------------------------------------------- | |
| # SQLite -> MySQL SQL translation | |
| # --------------------------------------------------------------------------- | |
| _cache = {} | |
| def _scalar_minmax(sql): | |
| """MAX(a, b) / MIN(a, b) with a top-level comma are SQLite scalar functions -> GREATEST / LEAST.""" | |
| out, i = [], 0 | |
| for m in re.finditer(r"\b(MAX|MIN)\(", sql, flags=re.I): | |
| if m.start() < i: | |
| continue | |
| depth, j, comma = 1, m.end(), False | |
| while j < len(sql) and depth: | |
| c = sql[j] | |
| depth += c == "(" | |
| depth -= c == ")" | |
| if c == "," and depth == 1: | |
| comma = True | |
| j += 1 | |
| if comma: | |
| out.append(sql[i:m.start()]) | |
| out.append("GREATEST(" if m.group(1).upper() == "MAX" else "LEAST(") | |
| i = m.end() | |
| out.append(sql[i:]) | |
| return "".join(out) | |
| def _placeholders(sql, escape_percent): | |
| """`?` -> `%s` outside string literals; a literal % becomes %% when parameters are passed.""" | |
| out, quote = [], None | |
| for c in sql: | |
| if quote: | |
| out.append("%%" if (c == "%" and escape_percent) else c) | |
| if c == quote: | |
| quote = None | |
| elif c in ("'", '"', "`"): | |
| quote = c | |
| out.append(c) | |
| elif c == "?": | |
| out.append("%s") | |
| elif c == "%" and escape_percent: | |
| out.append("%%") | |
| else: | |
| out.append(c) | |
| return "".join(out) | |
| def translate(sql, has_params=True): | |
| key = (sql, has_params) | |
| if key in _cache: | |
| return _cache[key] | |
| s = sql | |
| s = re.sub(r"\bINSERT\s+OR\s+REPLACE\s+INTO\b", "REPLACE INTO", s, flags=re.I) | |
| s = re.sub(r"\bINSERT\s+OR\s+IGNORE\s+INTO\b", "INSERT IGNORE INTO", s, flags=re.I) | |
| s = re.sub(r"\bON\s+CONFLICT\s*\([^)]*\)\s*DO\s+UPDATE\s+SET\b", "ON DUPLICATE KEY UPDATE", s, flags=re.I) | |
| s = re.sub(r"\bexcluded\.(\w+)", r"VALUES(\1)", s, flags=re.I) | |
| s = re.sub(r"datetime\('now'\s*,\s*'-(\d+)\s+days?'\)", r"(UTC_TIMESTAMP() - INTERVAL \1 DAY)", s, flags=re.I) | |
| s = re.sub(r"date\('now'\s*,\s*'-(\d+)\s+days?'\)", r"DATE_SUB(UTC_DATE(), INTERVAL \1 DAY)", s, flags=re.I) | |
| s = re.sub(r"datetime\('now'\)", "UTC_TIMESTAMP()", s, flags=re.I) | |
| s = re.sub(r"date\('now'\)", "UTC_DATE()", s, flags=re.I) | |
| s = re.sub(r"\bdate\(", "DATE(", s) | |
| s = re.sub(r"\browid\b", "id", s, flags=re.I) | |
| s = re.sub(r"\b(FROM|INTO|UPDATE|TABLE)\s+usage\b", r"\1 `usage`", s, flags=re.I) | |
| s = _scalar_minmax(s) | |
| s = _placeholders(s, has_params) | |
| _cache[key] = s | |
| return s | |
| # --------------------------------------------------------------------------- | |
| # MySQL schema (created once; every statement is idempotent) | |
| # --------------------------------------------------------------------------- | |
| _T = "ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci" | |
| SCHEMA = [ | |
| f"""CREATE TABLE IF NOT EXISTS users ( | |
| email VARCHAR(190) PRIMARY KEY, name VARCHAR(255), salt VARCHAR(64), pw VARCHAR(255), | |
| created DATETIME DEFAULT CURRENT_TIMESTAMP, last_login DATETIME NULL, runs INT DEFAULT 0, | |
| active TINYINT DEFAULT 1, text_runs INT DEFAULT 0, doc_limit INT NULL, text_limit INT NULL, | |
| recovery_hash VARCHAR(255) NULL, plan VARCHAR(20) DEFAULT 'free', plan_until DATETIME NULL, | |
| period_start DATETIME NULL, period_docs INT DEFAULT 0, doc_credits INT DEFAULT 0, | |
| trial_used INT DEFAULT 0, plan_started DATETIME NULL, plan_status VARCHAR(20) DEFAULT 'free', | |
| credits_until DATETIME NULL, email_results TINYINT DEFAULT 1) {_T}""", | |
| f"CREATE TABLE IF NOT EXISTS meta (k VARCHAR(64) PRIMARY KEY, v TEXT) {_T}", | |
| f"""CREATE TABLE IF NOT EXISTS `usage` (who VARCHAR(190), day VARCHAR(10), kind VARCHAR(10), | |
| n INT DEFAULT 0, PRIMARY KEY (who, day, kind)) {_T}""", | |
| f"CREATE TABLE IF NOT EXISTS reset_fails (email VARCHAR(190), ts DOUBLE, KEY ix_email (email)) {_T}", | |
| f"""CREATE TABLE IF NOT EXISTS reset_requests (id BIGINT AUTO_INCREMENT PRIMARY KEY, | |
| ts DATETIME DEFAULT CURRENT_TIMESTAMP, email VARCHAR(190), contact VARCHAR(255), | |
| status VARCHAR(10) DEFAULT 'open') {_T}""", | |
| f"""CREATE TABLE IF NOT EXISTS reset_codes (id BIGINT AUTO_INCREMENT PRIMARY KEY, email VARCHAR(190), | |
| code_hash VARCHAR(255), expires DOUBLE, attempts INT DEFAULT 0, created DOUBLE, | |
| KEY ix_email (email)) {_T}""", | |
| f"""CREATE TABLE IF NOT EXISTS payments (id BIGINT AUTO_INCREMENT PRIMARY KEY, | |
| ts DATETIME DEFAULT CURRENT_TIMESTAMP, email VARCHAR(190), item VARCHAR(80), amount VARCHAR(40), | |
| source VARCHAR(30), note TEXT, KEY ix_email (email)) {_T}""", | |
| f"""CREATE TABLE IF NOT EXISTS upgrade_requests (id BIGINT AUTO_INCREMENT PRIMARY KEY, | |
| ts DATETIME DEFAULT CURRENT_TIMESTAMP, email VARCHAR(190), name VARCHAR(255), item VARCHAR(80), | |
| status VARCHAR(10) DEFAULT 'open') {_T}""", | |
| f"""CREATE TABLE IF NOT EXISTS plan_history (id BIGINT AUTO_INCREMENT PRIMARY KEY, | |
| ts DATETIME DEFAULT CURRENT_TIMESTAMP, email VARCHAR(190), plan VARCHAR(20), started DATETIME NULL, | |
| until DATETIME NULL, event VARCHAR(60), note TEXT, KEY ix_email (email)) {_T}""", | |
| f"""CREATE TABLE IF NOT EXISTS reruns (email VARCHAR(190), fkey VARCHAR(190), first_ts DATETIME, | |
| runs INT DEFAULT 0, PRIMARY KEY (email, fkey)) {_T}""", | |
| f"""CREATE TABLE IF NOT EXISTS daily_usage (email VARCHAR(190), day VARCHAR(10), tool VARCHAR(40), | |
| n INT DEFAULT 0, PRIMARY KEY (email, day, tool)) {_T}""", | |
| f"""CREATE TABLE IF NOT EXISTS runs (id BIGINT AUTO_INCREMENT PRIMARY KEY, | |
| ts DATETIME DEFAULT CURRENT_TIMESTAMP, who VARCHAR(190), email VARCHAR(190), tool VARCHAR(60), | |
| kind VARCHAR(10), name VARCHAR(255), words INT, ai_before INT, ai_after INT, secs DOUBLE, | |
| status VARCHAR(255), output TEXT, engine VARCHAR(20), | |
| KEY runs_email (email, ts), KEY runs_ts (ts)) {_T}""", | |
| f"""CREATE TABLE IF NOT EXISTS push_subs (endpoint VARCHAR(512) PRIMARY KEY, cid VARCHAR(64), | |
| email VARCHAR(190), sub TEXT, ts DOUBLE) {_T}""", | |
| f"""CREATE TABLE IF NOT EXISTS mail_log (id BIGINT AUTO_INCREMENT PRIMARY KEY, | |
| ts DATETIME DEFAULT CURRENT_TIMESTAMP, email VARCHAR(190), subject VARCHAR(255), | |
| status VARCHAR(400)) {_T}""", | |
| # --- billing (new): every credit movement and every online payment | |
| f"""CREATE TABLE IF NOT EXISTS credit_ledger (id BIGINT AUTO_INCREMENT PRIMARY KEY, | |
| ts DATETIME DEFAULT CURRENT_TIMESTAMP, email VARCHAR(190) NOT NULL, delta INT NOT NULL, | |
| bucket VARCHAR(10) NOT NULL, reason VARCHAR(30) NOT NULL, ref VARCHAR(120) NOT NULL, | |
| note VARCHAR(255) NULL, UNIQUE KEY uq_ref (reason, ref, bucket), KEY ix_email (email, ts)) {_T}""", | |
| f"""CREATE TABLE IF NOT EXISTS reservations (rid VARCHAR(40) PRIMARY KEY, | |
| ts DATETIME DEFAULT CURRENT_TIMESTAMP, email VARCHAR(190), month_n INT DEFAULT 0, bought_n INT DEFAULT 0, | |
| trial_n INT DEFAULT 0, rerun_fkey VARCHAR(190) NULL, rerun_new TINYINT DEFAULT 0, | |
| state VARCHAR(10) DEFAULT 'held', KEY ix_state (state, ts)) {_T}""", | |
| f"""CREATE TABLE IF NOT EXISTS orders (id BIGINT AUTO_INCREMENT PRIMARY KEY, | |
| ts DATETIME DEFAULT CURRENT_TIMESTAMP, merchant_order_id VARCHAR(64) NOT NULL UNIQUE, | |
| gateway_order_id VARCHAR(80) NULL, email VARCHAR(190) NOT NULL, item VARCHAR(40) NOT NULL, | |
| months INT DEFAULT 1, amount_paise INT NOT NULL, status VARCHAR(12) DEFAULT 'created', | |
| paid_at DATETIME NULL, gateway_txn_id VARCHAR(80) NULL, KEY ix_email (email, ts)) {_T}""", | |
| ] | |
| _schema_lock = threading.Lock() | |
| _schema_done = False | |
| def _ensure_schema(raw): | |
| global _schema_done | |
| with _schema_lock: | |
| if _schema_done: | |
| return | |
| cur = raw.cursor() | |
| for stmt in SCHEMA: | |
| cur.execute(stmt) | |
| raw.commit() | |
| _schema_done = True | |
| # --------------------------------------------------------------------------- | |
| # MySQL connection pool + sqlite3-like wrapper | |
| # --------------------------------------------------------------------------- | |
| _POOL_SIZE = int(os.getenv("MYSQL_POOL", "8")) | |
| _pool = queue.LifoQueue() | |
| _created = 0 | |
| _pool_lock = threading.Lock() | |
| def _new_raw(): | |
| import pymysql | |
| from pymysql.constants import FIELD_TYPE | |
| from pymysql.converters import conversions | |
| conv = conversions.copy() | |
| for t in (FIELD_TYPE.DATETIME, FIELD_TYPE.TIMESTAMP, FIELD_TYPE.DATE): # the app handles times as text | |
| conv[t] = lambda v: v | |
| conv[FIELD_TYPE.NEWDECIMAL] = lambda v: int(v) if re.fullmatch(r"-?\d+", v) else float(v) # SUM() / AVG() | |
| ssl = None | |
| if os.getenv("MYSQL_SSL_CA"): | |
| ssl = {"ca": os.getenv("MYSQL_SSL_CA")} | |
| elif os.getenv("MYSQL_SSL", "") in ("1", "true", "yes"): | |
| ssl = {"ssl": True} | |
| raw = pymysql.connect( | |
| host=MYSQL_HOST, port=int(os.getenv("MYSQL_PORT", "3306")), user=os.getenv("MYSQL_USER", ""), | |
| password=os.getenv("MYSQL_PASSWORD", ""), database=os.getenv("MYSQL_DB", ""), charset="utf8mb4", | |
| autocommit=False, conv=conv, ssl=ssl, connect_timeout=10, read_timeout=30, write_timeout=30) | |
| with raw.cursor() as c: | |
| c.execute("SET time_zone = '+00:00'") | |
| _ensure_schema(raw) | |
| return raw | |
| def _acquire(): | |
| global _created | |
| while True: | |
| try: | |
| raw = _pool.get_nowait() | |
| except queue.Empty: | |
| with _pool_lock: | |
| if _created < _POOL_SIZE: | |
| _created += 1 | |
| make = True | |
| else: | |
| make = False | |
| if make: | |
| try: | |
| return _new_raw() | |
| except Exception: | |
| with _pool_lock: | |
| _created -= 1 | |
| raise | |
| raw = _pool.get(timeout=30) | |
| try: | |
| raw.ping(reconnect=True) # server closed an idle connection -> reconnect | |
| return raw | |
| except Exception: | |
| with _pool_lock: | |
| _created -= 1 | |
| def _release(raw): | |
| try: | |
| raw.rollback() | |
| _pool.put(raw) | |
| except Exception: | |
| global _created | |
| with _pool_lock: | |
| _created -= 1 | |
| class _MyConn: | |
| """sqlite3.Connection look-alike: execute(), `with con:` (commit / rollback), rowcount, lastrowid.""" | |
| def __init__(self): | |
| self._raw = _acquire() | |
| def execute(self, sql, params=None): | |
| cur = self._raw.cursor() | |
| if params is not None and not isinstance(params, (tuple, list, dict)): | |
| params = (params,) | |
| cur.execute(translate(sql, params is not None and len(params) > 0), tuple(params) if params else None) | |
| return cur | |
| def commit(self): | |
| self._raw.commit() | |
| def rollback(self): | |
| self._raw.rollback() | |
| def close(self): | |
| raw, self._raw = self._raw, None | |
| if raw is not None: | |
| _release(raw) | |
| def __enter__(self): | |
| return self | |
| def __exit__(self, et, ev, tb): | |
| try: | |
| if et is None: | |
| self._raw.commit() | |
| else: | |
| self._raw.rollback() | |
| finally: | |
| self.close() # back to the pool | |
| return False | |
| def __del__(self): | |
| try: | |
| self.close() | |
| except Exception: | |
| pass | |
| # --------------------------------------------------------------------------- | |
| def connect(sqlite_path=None): | |
| """A connection: MySQL from the pool, or the SQLite file (sqlite_path / USERS_DB).""" | |
| if IS_MYSQL: | |
| return _MyConn() | |
| return sqlite3.connect(str(sqlite_path or _SQLITE_PATH), timeout=10) | |
| def ping(): | |
| """(ok, message) for the admin page / start-up log.""" | |
| try: | |
| con = connect() | |
| con.execute("SELECT 1").fetchone() | |
| con.close() | |
| return True, ("MySQL " + MYSQL_HOST) if IS_MYSQL else "SQLite" | |
| except Exception as err: | |
| return False, str(err)[:200] | |
| def table_exists_cols(con, table): | |
| """Column names of a table (SQLite PRAGMA; MySQL information_schema) -- used by the SQLite auto-migrations.""" | |
| if IS_MYSQL: | |
| return {r[0] for r in con.execute( | |
| "SELECT column_name FROM information_schema.columns WHERE table_schema=DATABASE() " | |
| "AND table_name=?", (table,)).fetchall()} | |
| return {r[1] for r in con.execute(f"PRAGMA table_info({table})")} | |