From 5395a1c4857dab4af0e6529553e9abf4344fe99d Mon Sep 17 00:00:00 2001 From: superintendent2521 <67196752+superintendent2521@users.noreply.github.com> Date: Sat, 13 Dec 2025 14:34:38 -0500 Subject: [PATCH 1/9] Implement migration from MongoDB to Postgres with backup and export features --- requirements.txt | 1 + scripts/mongo_to_postgres_migration.py | 211 ++++++++ scripts/mongo_to_postgres_migration.sh | 210 ++++++++ src/database.py | 698 ++++++++++++++++--------- 4 files changed, 868 insertions(+), 252 deletions(-) create mode 100644 scripts/mongo_to_postgres_migration.py create mode 100644 scripts/mongo_to_postgres_migration.sh diff --git a/requirements.txt b/requirements.txt index c0de04f..6d4707e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -19,3 +19,4 @@ aiozipstream markdown-pdf boto3 pydantic-settings +asyncpg diff --git a/scripts/mongo_to_postgres_migration.py b/scripts/mongo_to_postgres_migration.py new file mode 100644 index 0000000..6bee524 --- /dev/null +++ b/scripts/mongo_to_postgres_migration.py @@ -0,0 +1,211 @@ +#!/usr/bin/env python3 +""" +MongoDB -> Postgres migration helper for WikiWare. + +Steps: +1) Creates a mongodump backup (gzip archive) before touching data. +2) Reads MongoDB collections and writes them into Postgres JSONB storage + (table: wikiware_documents). +3) Upserts documents so the script can be run multiple times. + +Usage: + python scripts/mongo_to_postgres_migration.py [--skip-backup] [--truncate] + +Flags: + --skip-backup Do not run mongodump first (not recommended). + --truncate Delete existing rows for the migrated collections before insert. + +Environment: + MONGODB_URL (default: mongodb://localhost:27017) + MONGODB_DB_NAME (default: wikiware) + POSTGRES_DSN (default: postgresql://postgres:postgres@localhost:5432/wikiware) +""" + +from __future__ import annotations + +import argparse +import asyncio +import datetime as dt +import os +import subprocess +from pathlib import Path +from typing import Dict, Iterable, List + +import asyncpg +from dotenv import load_dotenv +from loguru import logger +from pymongo import MongoClient + +TABLE_NAME = "wikiware_documents" +COLLECTIONS = [ + "pages", + "history", + "branches", + "users", + "sessions", + "image_hashes", + "analytics_events", + "settings", + "system_logs", +] + + +def _timestamp() -> str: + return dt.datetime.utcnow().strftime("%Y-%m-%dT%H%M%SZ") + + +def run_backup(uri_with_db: str, backup_dir: Path) -> Path: + backup_dir.mkdir(parents=True, exist_ok=True) + archive = backup_dir / f"mongodb-{_timestamp()}.archive.gz" + logger.info("Running mongodump to {}", archive) + cmd = [ + "mongodump", + f"--uri={uri_with_db}", + f"--archive={archive}", + "--gzip", + ] + subprocess.run(cmd, check=True) + logger.info("Backup complete: {}", archive) + return archive + + +def _uri_with_db(uri: str, db: str) -> str: + if uri.rstrip("/").endswith(f"/{db}"): + return uri + if "?" in uri: + base, query = uri.split("?", 1) + return f"{base.rstrip('/')}/{db}?{query}" + return f"{uri.rstrip('/')}/{db}" + + +def load_env_defaults() -> Dict[str, str]: + load_dotenv() + return { + "MONGODB_URL": os.getenv("MONGODB_URL", "mongodb://localhost:27017"), + "MONGODB_DB_NAME": os.getenv("MONGODB_DB_NAME", "wikiware"), + "POSTGRES_DSN": os.getenv( + "POSTGRES_DSN", "postgresql://postgres:postgres@localhost:5432/wikiware" + ), + } + + +async def ensure_table(pool: asyncpg.Pool) -> None: + async with pool.acquire() as conn: + await conn.execute( + f""" + CREATE TABLE IF NOT EXISTS {TABLE_NAME} ( + collection TEXT NOT NULL, + id TEXT NOT NULL, + doc JSONB NOT NULL, + PRIMARY KEY (collection, id) + ); + """ + ) + + +def fetch_collection(mongo_db, name: str) -> Iterable[Dict]: + logger.info("Fetching Mongo collection '{}'", name) + return mongo_db[name].find() + + +async def upsert_documents( + pool: asyncpg.Pool, collection: str, documents: Iterable[Dict] +) -> int: + inserted = 0 + async with pool.acquire() as conn: + async with conn.transaction(): + stmt = await conn.prepare( + f""" + INSERT INTO {TABLE_NAME} (collection, id, doc) + VALUES ($1, $2, $3) + ON CONFLICT (collection, id) DO UPDATE SET doc = EXCLUDED.doc + """ + ) + for doc in documents: + doc = dict(doc) + doc_id = str(doc.get("_id") or doc.get("id") or doc.get("uuid") or "") + if not doc_id: + continue + await stmt.fetch(collection, doc_id, doc) + inserted += 1 + return inserted + + +async def migrate_collections( + mongo_uri: str, + mongo_db_name: str, + postgres_dsn: str, + *, + backup_first: bool, + truncate: bool, +) -> None: + mongo_uri_with_db = _uri_with_db(mongo_uri, mongo_db_name) + backup_dir = Path("backups") + + if backup_first: + run_backup(mongo_uri_with_db, backup_dir) + else: + logger.warning("Skipping mongodump backup as requested.") + + mongo_client = MongoClient(mongo_uri_with_db) + mongo_db = mongo_client.get_database() + + pool = await asyncpg.create_pool(postgres_dsn, min_size=1, max_size=10) + await ensure_table(pool) + + if truncate: + async with pool.acquire() as conn: + logger.info("Truncating existing rows for selected collections") + await conn.execute( + f"DELETE FROM {TABLE_NAME} WHERE collection = ANY($1::text[])", + COLLECTIONS, + ) + + for name in COLLECTIONS: + documents = fetch_collection(mongo_db, name) + count = await upsert_documents(pool, name, documents) + logger.info("Migrated {} documents into '{}'", count, name) + + await pool.close() + mongo_client.close() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Migrate MongoDB data into Postgres JSONB storage.") + parser.add_argument("--skip-backup", action="store_true", help="Do not run mongodump before migrating.") + parser.add_argument("--truncate", action="store_true", help="Delete existing rows for migrated collections first.") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + env = load_env_defaults() + + logger.info("Starting Mongo -> Postgres migration") + logger.info("Mongo URI: {}", env["MONGODB_URL"]) + logger.info("Mongo DB: {}", env["MONGODB_DB_NAME"]) + logger.info("Postgres DSN: {}", env["POSTGRES_DSN"]) + + try: + asyncio.run( + migrate_collections( + env["MONGODB_URL"], + env["MONGODB_DB_NAME"], + env["POSTGRES_DSN"], + backup_first=not args.skip_backup, + truncate=args.truncate, + ) + ) + except subprocess.CalledProcessError as exc: + logger.error("Backup failed: {}", exc) + return 1 + except Exception as exc: + logger.exception("Migration failed: {}", exc) + return 1 + + logger.info("Migration completed successfully.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/mongo_to_postgres_migration.sh b/scripts/mongo_to_postgres_migration.sh new file mode 100644 index 0000000..87ba802 --- /dev/null +++ b/scripts/mongo_to_postgres_migration.sh @@ -0,0 +1,210 @@ +#!/usr/bin/env bash +# Helper to migrate WikiWare data from MongoDB to Postgres. +# - Always takes a mongodump backup first. +# - Exports MongoDB collections to NDJSON for inspection. +# - Optionally loads the NDJSON into Postgres staging tables as JSONB. + +set -euo pipefail + +usage() { + cat <<'USAGE' +Usage: ./scripts/mongo_to_postgres_migration.sh [--load-postgres] [--yes] + + --load-postgres Load exported Mongo documents into Postgres staging tables (jsonb). + --yes, -y Skip confirmation prompts. + --help, -h Show this help message. + +The script reads MONGODB_URL, MONGODB_DB_NAME, and POSTGRES_DSN from .env if present. +Defaults: + MONGODB_URL=mongodb://localhost:27017 + MONGODB_DB_NAME=wikiware + POSTGRES_DSN=postgresql://postgres:postgres@localhost:5432/wikiware +USAGE +} + +log() { + printf '[mongo->postgres] %s\n' "$*" +} + +fail() { + printf '[mongo->postgres] ERROR: %s\n' "$*" >&2 + exit 1 +} + +confirm() { + local prompt="${1:-Proceed?} [y/N] " + if [[ ${AUTO_CONFIRM:-0} -eq 1 ]]; then + return 0 + fi + read -r -p "$prompt" reply + [[ "$reply" =~ ^[Yy]$ ]] +} + +require_cmd() { + local name="$1" + command -v "$name" >/dev/null 2>&1 || fail "Missing required command: $name" +} + +uri_with_db() { + local uri="$1" + local db="$2" + if [[ "$uri" =~ /${db}($|[\?#/]) ]]; then + echo "$uri" + return + fi + if [[ "$uri" == *"?"* ]]; then + local base="${uri%%\?*}" + local query="${uri#"$base"}" + echo "${base%/}/$db${query}" + else + echo "${uri%/}/$db" + fi +} + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" +ENV_FILE="$ROOT_DIR/.env" + +if [[ -f "$ENV_FILE" ]]; then + set -o allexport + # shellcheck disable=SC1090 + source "$ENV_FILE" + set +o allexport +fi + +MONGODB_URL="${MONGODB_URL:-mongodb://localhost:27017}" +MONGODB_DB_NAME="${MONGODB_DB_NAME:-wikiware}" +POSTGRES_DSN="${POSTGRES_DSN:-postgresql://postgres:postgres@localhost:5432/wikiware}" +BACKUP_DIR="${BACKUP_DIR:-$ROOT_DIR/backups}" +EXPORT_DIR="${EXPORT_DIR:-$ROOT_DIR/migration_workdir}" +MONGO_EXPORT_DIR="$EXPORT_DIR/mongo_exports" +CSV_EXPORT_DIR="$EXPORT_DIR/postgres_csv" +TIMESTAMP="$(date -u +"%Y-%m-%dT%H%M%SZ")" +AUTO_CONFIRM=0 +LOAD_POSTGRES=0 + +for arg in "$@"; do + case "$arg" in + --load-postgres) LOAD_POSTGRES=1 ;; + --yes|-y) AUTO_CONFIRM=1 ;; + --help|-h) usage; exit 0 ;; + *) fail "Unknown argument: $arg" ;; + esac +done + +MONGODB_URI_WITH_DB="$(uri_with_db "$MONGODB_URL" "$MONGODB_DB_NAME")" + +backup_mongo() { + require_cmd mongodump + mkdir -p "$BACKUP_DIR" + local archive="$BACKUP_DIR/mongodb-${MONGODB_DB_NAME}-${TIMESTAMP}.archive.gz" + + log "Creating MongoDB backup at $archive" + mongodump \ + --uri="$MONGODB_URI_WITH_DB" \ + --archive="$archive" \ + --gzip + log "Backup complete: $archive" +} + +export_collections() { + require_cmd mongoexport + mkdir -p "$MONGO_EXPORT_DIR" + local collections=( + pages + history + branches + users + sessions + image_hashes + analytics_events + ) + + log "Exporting MongoDB collections to $MONGO_EXPORT_DIR" + for collection in "${collections[@]}"; do + local outfile="$MONGO_EXPORT_DIR/${collection}.ndjson" + log " - $collection -> $outfile" + mongoexport \ + --uri="$MONGODB_URI_WITH_DB" \ + --collection="$collection" \ + --type=json \ + --out="$outfile" + done + log "Exports complete." +} + +load_postgres_staging() { + require_cmd perl + require_cmd psql + mkdir -p "$CSV_EXPORT_DIR" + local collections=( + pages + history + branches + users + sessions + image_hashes + analytics_events + ) + + log "Preparing Postgres staging tables in schema wikiware_migration" + psql "$POSTGRES_DSN" -v ON_ERROR_STOP=1 <<'SQL' +CREATE SCHEMA IF NOT EXISTS wikiware_migration; +SQL + + for collection in "${collections[@]}"; do + local ndjson="$MONGO_EXPORT_DIR/${collection}.ndjson" + if [[ ! -s "$ndjson" ]]; then + log " - Skipping $collection (no export found)" + continue + fi + + local csv="$CSV_EXPORT_DIR/${collection}.csv" + perl -pe 's/"/""/g; $_="\"$_\""' "$ndjson" >"$csv" + + log " - Loading $collection into Postgres staging table" + psql "$POSTGRES_DSN" -v ON_ERROR_STOP=1 <= DB_OPERATION_LOG_THRESHOLD_MS: - logger.info(f"DB {method_name} on {collection_name} took {duration:.4f}s") - return result - except Exception as e: - duration = time.monotonic() - start_time - logger.error(f"DB {method_name} on {collection_name} failed after {duration:.4f}s: {e}") - raise - - return timed_method - - def timed_method(*args, **kwargs): - start_time = time.monotonic() - try: - result = original_method(*args, **kwargs) - duration = time.monotonic() - start_time - if config.DB_QUERY_LOGGING_ENABLED and duration >= DB_OPERATION_LOG_THRESHOLD_MS: - logger.info(f"DB {method_name} on {collection_name} took {duration:.4f}s") - return result - except Exception as e: - duration = time.monotonic() - start_time - logger.error(f"DB {method_name} on {collection_name} failed after {duration:.4f}s: {e}") - raise - - return timed_method +POSTGRES_DSN = os.getenv( + "POSTGRES_DSN", "postgresql://postgres:postgres@localhost:5432/wikiware" +) +DB_OPERATION_LOG_THRESHOLD_MS = ( + float(os.getenv("DB_OPERATION_LOG_THRESHOLD_MS", "100")) / 1000 +) +TABLE_NAME = "wikiware_documents" + + +# ---------------------- Result helpers ---------------------- + + +@dataclass +class InsertOneResult: + inserted_id: Optional[str] + + +@dataclass +class UpdateResult: + matched_count: int + modified_count: int + upserted_id: Optional[str] = None + + +@dataclass +class DeleteResult: + deleted_count: int + + +# ---------------------- Utility helpers ---------------------- + + +def _ensure_loop() -> Optional[asyncio.AbstractEventLoop]: + try: + return asyncio.get_running_loop() + except RuntimeError: + return None + + +def _get_by_path(doc: Dict[str, Any], path: str) -> Any: + parts = path.split(".") + current: Any = doc + for part in parts: + if not isinstance(current, dict) or part not in current: + return None + current = current[part] + return current + + +def _set_by_path(doc: Dict[str, Any], path: str, value: Any) -> None: + parts = path.split(".") + current = doc + for part in parts[:-1]: + if part not in current or not isinstance(current[part], dict): + current[part] = {} + current = current[part] + current[parts[-1]] = value + + +def _apply_projection(doc: Dict[str, Any], projection: Optional[Dict[str, int]]) -> Dict[str, Any]: + if projection is None: + return doc + include_keys = {k for k, v in projection.items() if v} + exclude_keys = {k for k, v in projection.items() if not v} + + if include_keys: + projected = {k: doc[k] for k in include_keys if k in doc} + if "_id" in doc and "_id" not in projected and projection.get("_id", 1): + projected["_id"] = doc["_id"] + return projected + + if exclude_keys: + return {k: v for k, v in doc.items() if k not in exclude_keys} + + return doc + + +def _matches_filter(doc: Dict[str, Any], filt: Dict[str, Any]) -> bool: + if not filt: + return True + + def compare(val: Any, condition: Any) -> bool: + if isinstance(condition, dict): + for op, expected in condition.items(): + if op == "$gte" and not (val is not None and val >= expected): + return False + if op == "$gt" and not (val is not None and val > expected): + return False + if op == "$lte" and not (val is not None and val <= expected): + return False + if op == "$lt" and not (val is not None and val < expected): + return False + if op == "$in" and val not in expected: + return False + if op == "$nin" and val in expected: + return False + return True + return val == condition + + for key, expected in filt.items(): + value = _get_by_path(doc, key) if "." in key else doc.get(key) + if not compare(value, expected): + return False + return True + + +def _apply_update(doc: Dict[str, Any], update: Dict[str, Any]) -> Dict[str, Any]: + updated = dict(doc) + for op, changes in update.items(): + if op == "$set": + for path, value in changes.items(): + _set_by_path(updated, path, value) + elif op == "$inc": + for path, delta in changes.items(): + current = _get_by_path(updated, path) + if current is None: + _set_by_path(updated, path, delta) + else: + _set_by_path(updated, path, current + delta) + else: + logger.warning("Unsupported update operator {}", op) + return updated + + +def _project_sort_key(doc: Dict[str, Any], key: str) -> Any: + if "." in key: + return _get_by_path(doc, key) + return doc.get(key) + + +# ---------------------- Cursor ---------------------- + + +class PostgresCursor: + def __init__( + self, + loader: Callable[[], asyncio.Future], + *, + limit: Optional[int] = None, + ): + self._loader = loader + self._docs: Optional[List[Dict[str, Any]]] = None + self._sorts: List[tuple[str, int]] = [] + self._limit = limit + + async def _ensure_loaded(self) -> None: + if self._docs is None: + self._docs = await self._loader() + self._apply_sorts_and_limit() + + def _apply_sorts_and_limit(self) -> None: + if self._docs is None: + return + for key, direction in reversed(self._sorts): + self._docs.sort(key=lambda d: _project_sort_key(d, key) or 0, reverse=direction < 0) + if self._limit is not None: + self._docs = self._docs[: self._limit] + + def sort(self, key: str, direction: int = 1) -> "PostgresCursor": + self._sorts.append((key, direction)) + return self + + async def to_list(self, length: Optional[int]) -> List[Dict[str, Any]]: + await self._ensure_loaded() + if self._docs is None: + return [] + if length is None: + return list(self._docs) + return list(self._docs[:length]) + + def __aiter__(self) -> AsyncIterator[Dict[str, Any]]: + async def iterator(): + await self._ensure_loaded() + for doc in self._docs or []: + yield doc + + return iterator() + + def __iter__(self) -> Iterable[Dict[str, Any]]: + loop = _ensure_loop() + if self._docs is None: + if loop and loop.is_running(): + raise RuntimeError("Synchronous iteration is not supported while loop is running.") + asyncio.run(self._ensure_loaded()) + return iter(self._docs or []) + + +# ---------------------- Collection ---------------------- + + +class PostgresCollection: + def __init__(self, name: str, db: "Database"): + self.name = name + self._db = db + + async def _fetch_docs(self) -> List[Dict[str, Any]]: + rows = await self._db.fetch( + f"SELECT doc FROM {TABLE_NAME} WHERE collection = $1", + self.name, + ) + return [row["doc"] for row in rows] + + def find( + self, + filt: Optional[Dict[str, Any]] = None, + projection: Optional[Dict[str, int]] = None, + limit: Optional[int] = None, + ) -> PostgresCursor: + async def loader(): + docs = await self._fetch_docs() + matched = [doc for doc in docs if _matches_filter(doc, filt or {})] + if projection: + matched = [_apply_projection(doc, projection) for doc in matched] + return matched + + return PostgresCursor(loader, limit=limit) + + async def find_one(self, filt: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]: + cursor = self.find(filt, limit=1) + results = await cursor.to_list(1) + return results[0] if results else None + + async def insert_one(self, document: Dict[str, Any]) -> InsertOneResult: + doc = dict(document) + if "_id" not in doc: + doc["_id"] = str(uuid.uuid4()) + await self._db.execute( + f""" + INSERT INTO {TABLE_NAME} (collection, id, doc) + VALUES ($1, $2, $3) + ON CONFLICT (collection, id) DO UPDATE SET doc = EXCLUDED.doc + """, + self.name, + str(doc["_id"]), + doc, + ) + return InsertOneResult(inserted_id=str(doc["_id"])) + + async def update_one( + self, + filt: Dict[str, Any], + update: Dict[str, Any], + *, + upsert: bool = False, + ) -> UpdateResult: + existing = await self.find_one(filt) + if existing is None: + if not upsert: + return UpdateResult(matched_count=0, modified_count=0, upserted_id=None) + base = {k: v for k, v in filt.items() if not isinstance(v, dict)} + new_doc = _apply_update(base, update) + result = await self.insert_one(new_doc) + return UpdateResult(matched_count=0, modified_count=1, upserted_id=result.inserted_id) + + updated_doc = _apply_update(existing, update) + await self._db.execute( + f"UPDATE {TABLE_NAME} SET doc = $3 WHERE collection = $1 AND id = $2", + self.name, + str(existing.get("_id")), + updated_doc, + ) + return UpdateResult(matched_count=1, modified_count=1, upserted_id=None) + + async def delete_one(self, filt: Dict[str, Any]) -> DeleteResult: + existing = await self.find_one(filt) + if existing is None: + return DeleteResult(deleted_count=0) + await self._db.execute( + f"DELETE FROM {TABLE_NAME} WHERE collection = $1 AND id = $2", + self.name, + str(existing.get("_id")), + ) + return DeleteResult(deleted_count=1) + + async def count_documents(self, filt: Optional[Dict[str, Any]] = None) -> int: + docs = await self._fetch_docs() + return sum(1 for doc in docs if _matches_filter(doc, filt or {})) + + async def distinct(self, key: str, filt: Optional[Dict[str, Any]] = None) -> List[Any]: + docs = await self._fetch_docs() + values = [] + for doc in docs: + if filt and not _matches_filter(doc, filt): + continue + value = _get_by_path(doc, key) if "." in key else doc.get(key) + if value is not None: + values.append(value) + return list({v for v in values}) + + def aggregate(self, pipeline: List[Dict[str, Any]]) -> PostgresCursor: + async def loader(): + docs = await self._fetch_docs() + for stage in pipeline: + if "$match" in stage: + docs = [doc for doc in docs if _matches_filter(doc, stage["$match"])] + elif "$project" in stage: + projected = [] + for doc in docs: + new_doc: Dict[str, Any] = {} + for key, expr in stage["$project"].items(): + if expr == 1: + new_doc[key] = doc.get(key) + elif isinstance(expr, str) and expr.startswith("$"): + new_doc[key] = _get_by_path(doc, expr[1:]) if "." in expr[1:] else doc.get(expr[1:]) + elif isinstance(expr, dict) and "$dateToString" in expr: + fmt = expr["$dateToString"]["format"] + date_val = _get_by_path(doc, expr["$dateToString"]["date"][1:]) + new_doc[key] = date_val.strftime(fmt) if hasattr(date_val, "strftime") else None + else: + new_doc[key] = expr + projected.append(new_doc) + docs = projected or docs + elif "$group" in stage: + grouped: Dict[Any, Dict[str, Any]] = {} + group_spec = stage["$group"] + for doc in docs: + group_id_expr = group_spec.get("_id") + if isinstance(group_id_expr, dict): + group_id = {} + for k, v in group_id_expr.items(): + if isinstance(v, dict) and "$dateToString" in v: + date_val = _get_by_path(doc, v["$dateToString"]["date"][1:]) + fmt = v["$dateToString"]["format"] + group_id[k] = date_val.strftime(fmt) if hasattr(date_val, "strftime") else None + elif isinstance(v, str) and v.startswith("$"): + group_id[k] = _get_by_path(doc, v[1:]) if "." in v[1:] else doc.get(v[1:]) + else: + group_id[k] = v + elif isinstance(group_id_expr, str) and group_id_expr.startswith("$"): + group_id = _get_by_path(doc, group_id_expr[1:]) if "." in group_id_expr[1:] else doc.get(group_id_expr[1:]) + else: + group_id = group_id_expr + + if group_id not in grouped: + grouped[group_id] = {"_id": group_id} + + for field, expr in group_spec.items(): + if field == "_id": + continue + if isinstance(expr, dict) and "$sum" in expr: + val = expr["$sum"] + addend = 0 + if isinstance(val, (int, float)): + addend = val + elif isinstance(val, str) and val.startswith("$"): + addend = _get_by_path(doc, val[1:]) if "." in val[1:] else doc.get(val[1:]) or 0 + grouped[group_id][field] = grouped[group_id].get(field, 0) + (addend or 0) + elif isinstance(expr, dict) and "$max" in expr: + candidate = expr["$max"] + if isinstance(candidate, str) and candidate.startswith("$"): + candidate = _get_by_path(doc, candidate[1:]) if "." in candidate[1:] else doc.get(candidate[1:]) + current = grouped[group_id].get(field) + if current is None or (candidate is not None and candidate > current): + grouped[group_id][field] = candidate + elif isinstance(expr, dict) and "$first" in expr: + if field not in grouped[group_id]: + val = expr["$first"] + if isinstance(val, str) and val.startswith("$"): + val = _get_by_path(doc, val[1:]) if "." in val[1:] else doc.get(val[1:]) + grouped[group_id][field] = val + docs = list(grouped.values()) + elif "$sort" in stage: + for key, direction in reversed(list(stage["$sort"].items())): + docs.sort(key=lambda d: _project_sort_key(d, key) or 0, reverse=direction < 0) + elif "$limit" in stage: + docs = docs[: stage["$limit"]] + return docs + + return PostgresCursor(loader) + + +# ---------------------- Database ---------------------- class Database: - """Manages MongoDB connection and provides access to collections.""" + """Manages Postgres connection and collection-style access.""" - def __init__(self, mongo_url: str = MONGODB_URL, db_name: str = MONGODB_DB_NAME): - self._mongo_url = mongo_url - self._db_name = db_name - self.client: Optional[AsyncIOMotorClient] = None # type: ignore - self.db: Optional[AsyncIOMotorDatabase] = None # type: ignore + def __init__(self, dsn: str = POSTGRES_DSN): + self._dsn = dsn + self.pool: Optional[asyncpg.Pool] = None self.is_connected = False self._connection_lock = asyncio.Lock() - self._max_pool_size = MONGODB_MAX_CONNECTIONS - self._min_pool_size = MONGODB_MIN_CONNECTIONS - self._wrapped_collections: Dict[str, AsyncIOMotorCollection] = {} - # List of methods to wrap with timing - self._methods_to_wrap = [ - 'find_one', 'insert_one', 'update_one', 'delete_one', - 'find', 'count_documents', 'aggregate', 'distinct' - ] + self._wrapped_collections: Dict[str, PostgresCollection] = {} def _reset_state(self) -> None: - """Clear cached client/db handles and connection flags.""" - if self.client is not None: - self.client.close() - self.client = None - self.db = None + self.pool = None self.is_connected = False self._wrapped_collections.clear() - async def connect(self, max_retries: Optional[int] = 10) -> None: - """Establish connection to MongoDB with connection pooling and retry logic.""" - if max_retries is None: - max_retries = float("inf") - delay = 10 - else: - delay = 5 - - retry_count = 0 - + async def connect(self) -> None: async with self._connection_lock: - while retry_count < max_retries: - try: - logger.info( - "Attempting to connect to MongoDB at {}... (attempt {}) with pool_size={}, min_pool_size={}", - self._mongo_url, - retry_count + 1, - self._max_pool_size, - self._min_pool_size, - ) - - self.client = AsyncIOMotorClient( - self._mongo_url, - maxPoolSize=self._max_pool_size, - minPoolSize=self._min_pool_size, - maxIdleTimeMS=MONGODB_MAX_IDLE_TIME_MS, - serverSelectionTimeoutMS=MONGODB_SERVER_SELECTION_TIMEOUT_MS, - socketTimeoutMS=MONGODB_SOCKET_TIMEOUT_MS, - retryWrites=True, - retryReads=True, - connectTimeoutMS=MONGODB_SERVER_SELECTION_TIMEOUT_MS, - ) - - # Test the connection - await self.client.admin.command("ping") - self.db = self.client[self._db_name] - self.is_connected = True - logger.info( - "Connected to MongoDB database '{}' with connection pool configured", - self._db_name - ) - return - except ServerSelectionTimeoutError: - retry_count += 1 - self._reset_state() - if max_retries != float("inf"): - logger.warning( - "MongoDB server not available. Attempt {}/{}. Retrying in {} seconds...", - retry_count, - max_retries, - delay, - ) - else: - logger.warning( - "MongoDB connection lost. Retrying in {} seconds...", delay - ) - await asyncio.sleep(delay) - except Exception: - logger.exception("Database connection error") - self._reset_state() - if max_retries != float("inf"): - return - await asyncio.sleep(delay) - - if max_retries != float("inf"): - logger.error( - "Failed to connect to MongoDB after multiple attempts. Running in offline mode." - ) + if self.is_connected: + return + try: + logger.info("Connecting to Postgres at {}", self._dsn) + self.pool = await asyncpg.create_pool(self._dsn, min_size=1, max_size=10) + await self._ensure_table() + self.is_connected = True + logger.info("Connected to Postgres and storage table ensured") + except Exception: + logger.exception("Failed to connect to Postgres") self._reset_state() - async def monitor_connection(self) -> None: - """Background task to monitor and retry connection if lost.""" - logger.info("Starting MongoDB connection monitor (retries every 10s)") - while True: - if not self.is_connected: - await self.connect(max_retries=None) - await asyncio.sleep(10) + async def _ensure_table(self) -> None: + if self.pool is None: + return + await self.execute( + f""" + CREATE TABLE IF NOT EXISTS {TABLE_NAME} ( + collection TEXT NOT NULL, + id TEXT NOT NULL, + doc JSONB NOT NULL, + PRIMARY KEY (collection, id) + ); + """ + ) async def disconnect(self) -> None: - """Close the MongoDB connection.""" async with self._connection_lock: + if self.pool: + await self.pool.close() self._reset_state() - def get_collection(self, name: str) -> Optional[AsyncIOMotorCollection]: # pyright: ignore[reportInvalidTypeForm] - """Get a collection by name if database is connected.""" - if self.is_connected and self.db is not None: - if name in self._wrapped_collections: - return self._wrapped_collections[name] - - collection = self.db[name] - # Wrap methods for timing once per collection - for method_name in self._methods_to_wrap: - if hasattr(collection, method_name): - original_method = getattr(collection, method_name) - if getattr(original_method, "_wikiware_timed", False): - continue - wrapped_method = _timed_wrapper(original_method, method_name, name) - setattr(wrapped_method, "_wikiware_timed", True) - setattr(collection, method_name, wrapped_method) - self._wrapped_collections[name] = collection - return collection - return None + def get_collection(self, name: str) -> Optional[PostgresCollection]: + if not self.is_connected: + return None + if name in self._wrapped_collections: + return self._wrapped_collections[name] + collection = PostgresCollection(name, self) + self._wrapped_collections[name] = collection + return collection + + async def execute(self, query: str, *args: Any) -> None: + if self.pool is None: + raise RuntimeError("Database not connected") + async with self.pool.acquire() as conn: + await conn.execute(query, *args) + + async def fetch(self, query: str, *args: Any) -> List[asyncpg.Record]: + if self.pool is None: + raise RuntimeError("Database not connected") + async with self.pool.acquire() as conn: + return await conn.fetch(query, *args) async def get_pool_stats(self) -> Dict[str, Any]: - """Get connection pool statistics.""" - if self.client is None: + if self.pool is None: return {"status": "not_connected"} - - try: - server_info = await self.client.server_info() - return { - "status": "connected", - "max_pool_size": self._max_pool_size, - "min_pool_size": self._min_pool_size, - "server_info": server_info, - } - except Exception: - logger.error("Error getting pool stats") - return {"status": "error", "error": "Failed to get pool stats"} + return { + "status": "connected", + "pool_size": self.pool.max_size, + "free": self.pool.free_size, + } # Global database instance @@ -238,85 +482,35 @@ async def get_pool_stats(self) -> Dict[str, Any]: # Collections def get_pages_collection(): - """Get the pages collection.""" return db_instance.get_collection("pages") def get_history_collection(): - """Get the history collection.""" return db_instance.get_collection("history") def get_branches_collection(): - """Get the branches collection.""" return db_instance.get_collection("branches") def get_users_collection(): - """Get the users collection.""" return db_instance.get_collection("users") def get_image_hashes_collection(): - """Get the image_hashes collection.""" return db_instance.get_collection("image_hashes") -# Helper functions -async def _drop_legacy_page_index(pages: AsyncIOMotorCollection) -> None: # type: ignore - """Remove deprecated single-field title index if present.""" - try: - existing_indexes = await pages.index_information() - if "title_1" in existing_indexes: - await pages.drop_index("title_1") - logger.info("Dropped legacy unique index on pages.title") - except Exception: - logger.warning("Failed to drop legacy title index") - -async def _ensure_pages_indexes(pages: AsyncIOMotorCollection) -> None: # type: ignore - """Ensure compound and text indexes exist for the pages collection.""" - await _drop_legacy_page_index(pages) - await pages.create_index([("title", 1), ("branch", 1)], unique=True) - await pages.create_index("updated_at") - await pages.create_index( - [("title", "text"), ("content", "text")], name="page_text_search" - ) - logger.info("Pages collection indexes created") - -async def _create_collection_indexes( - collection_name: str, collection: AsyncIOMotorCollection # type: ignore -) -> None: - """Create indexes declared in INDEX_CONFIGS for a given collection.""" - for keys, options in INDEX_CONFIGS.get(collection_name, []): - await collection.create_index(keys, **options) - logger.info("{} collection indexes created", collection_name.capitalize()) - async def create_indexes() -> None: - """Create required indexes on MongoDB collections.""" - if not db_instance.is_connected: - logger.warning("Skipping index creation because database is disconnected") - return - - pages = get_pages_collection() - if pages is not None: - await _ensure_pages_indexes(pages) - - for collection_name in ("users", "sessions", "image_hashes", "analytics_events"): - collection = db_instance.get_collection(collection_name) - if collection is None: - logger.warning( - "Collection '{}' unavailable while creating indexes", collection_name - ) - continue - await _create_collection_indexes(collection_name, collection) + # JSONB storage uses application-level filtering; explicit DB indexes can be added later. + logger.info("Index creation skipped (JSONB storage)") + async def init_database() -> None: - """Initialize database connection and create indexes.""" try: await db_instance.connect() if db_instance.is_connected: await create_indexes() - # Log connection pool stats pool_stats = await db_instance.get_pool_stats() logger.info("Database pool stats: {}", pool_stats) except Exception: From d596612eabbf4cba6e4695e7654e65e87e3244e0 Mon Sep 17 00:00:00 2001 From: superintendent2521 <67196752+superintendent2521@users.noreply.github.com> Date: Sat, 13 Dec 2025 14:36:11 -0500 Subject: [PATCH 2/9] add import awaitable --- src/database.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/database.py b/src/database.py index 8b89fbd..8377705 100644 --- a/src/database.py +++ b/src/database.py @@ -12,7 +12,7 @@ import os import uuid from dataclasses import dataclass -from typing import Any, AsyncIterator, Callable, Dict, Iterable, List, Optional +from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Iterable, List, Optional import asyncpg from dotenv import load_dotenv @@ -158,7 +158,7 @@ def _project_sort_key(doc: Dict[str, Any], key: str) -> Any: class PostgresCursor: def __init__( self, - loader: Callable[[], asyncio.Future], + loader: Callable[[], Awaitable[List[Dict[str, Any]]]], *, limit: Optional[int] = None, ): @@ -471,8 +471,8 @@ async def get_pool_stats(self) -> Dict[str, Any]: return {"status": "not_connected"} return { "status": "connected", - "pool_size": self.pool.max_size, - "free": self.pool.free_size, + "pool_size": getattr(self.pool, "get_max_size", lambda: None)(), + "free": getattr(self.pool, "get_idle_size", lambda: None)(), } From 243ef918d5e816610cd65988496f0a29bbbd0d7e Mon Sep 17 00:00:00 2001 From: superintendent2521 <67196752+superintendent2521@users.noreply.github.com> Date: Sat, 13 Dec 2025 14:55:48 -0500 Subject: [PATCH 3/9] Refactor migration scripts: add JSON sanitization, remove deprecated bash script --- .env | 3 +- .gitignore | 4 +- scripts/mongo_to_postgres_migration.py | 38 ++++- scripts/mongo_to_postgres_migration.sh | 210 ------------------------- 4 files changed, 40 insertions(+), 215 deletions(-) delete mode 100644 scripts/mongo_to_postgres_migration.sh diff --git a/.env b/.env index 3ed278f..cd2a619 100644 --- a/.env +++ b/.env @@ -1,3 +1,4 @@ MONGODB_URL=mongodb://localhost:27017 PORT=8000 -DEV=True \ No newline at end of file +DEV=True +POSTGRES_DSN=postgresql://postgres:postgres@localhost:5432/wikiware \ No newline at end of file diff --git a/.gitignore b/.gitignore index 5d3c3fa..294eefd 100644 --- a/.gitignore +++ b/.gitignore @@ -9,4 +9,6 @@ example-page.pdf # for local dev shennanigans involving piping output from cloudflared cloudflared.log backups/* -src/env_config.py \ No newline at end of file +src/env_config.py +scripts/backups/* +.DS_Store \ No newline at end of file diff --git a/scripts/mongo_to_postgres_migration.py b/scripts/mongo_to_postgres_migration.py index 6bee524..c1cbfb6 100644 --- a/scripts/mongo_to_postgres_migration.py +++ b/scripts/mongo_to_postgres_migration.py @@ -26,6 +26,7 @@ import argparse import asyncio import datetime as dt +import json import os import subprocess from pathlib import Path @@ -35,6 +36,8 @@ from dotenv import load_dotenv from loguru import logger from pymongo import MongoClient +from bson import ObjectId +from decimal import Decimal TABLE_NAME = "wikiware_documents" COLLECTIONS = [ @@ -47,11 +50,12 @@ "analytics_events", "settings", "system_logs", + "edit_sessions", ] def _timestamp() -> str: - return dt.datetime.utcnow().strftime("%Y-%m-%dT%H%M%SZ") + return dt.datetime.now(dt.timezone.utc).strftime("%Y-%m-%dT%H%M%SZ") def run_backup(uri_with_db: str, backup_dir: Path) -> Path: @@ -108,6 +112,31 @@ def fetch_collection(mongo_db, name: str) -> Iterable[Dict]: return mongo_db[name].find() +def _jsonable(value): + if isinstance(value, ObjectId): + return str(value) + if isinstance(value, dt.datetime): + if value.tzinfo is None: + value = value.replace(tzinfo=dt.timezone.utc) + return value.isoformat() + if isinstance(value, Decimal): + return float(value) + if isinstance(value, bytes): + return value.decode("utf-8", errors="ignore") + if isinstance(value, dict): + return {k: _jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple, set)): + return [_jsonable(v) for v in value] + return value + + +def sanitize_document(doc: Dict) -> Dict: + cleaned = _jsonable(doc) + if isinstance(cleaned, dict): + return cleaned + return {} + + async def upsert_documents( pool: asyncpg.Pool, collection: str, documents: Iterable[Dict] ) -> int: @@ -117,7 +146,7 @@ async def upsert_documents( stmt = await conn.prepare( f""" INSERT INTO {TABLE_NAME} (collection, id, doc) - VALUES ($1, $2, $3) + VALUES ($1, $2, $3::jsonb) ON CONFLICT (collection, id) DO UPDATE SET doc = EXCLUDED.doc """ ) @@ -126,7 +155,10 @@ async def upsert_documents( doc_id = str(doc.get("_id") or doc.get("id") or doc.get("uuid") or "") if not doc_id: continue - await stmt.fetch(collection, doc_id, doc) + clean_doc = sanitize_document(doc) + clean_doc["_id"] = doc_id + json_payload = json.dumps(clean_doc, ensure_ascii=False) + await stmt.fetch(collection, doc_id, json_payload) inserted += 1 return inserted diff --git a/scripts/mongo_to_postgres_migration.sh b/scripts/mongo_to_postgres_migration.sh deleted file mode 100644 index 87ba802..0000000 --- a/scripts/mongo_to_postgres_migration.sh +++ /dev/null @@ -1,210 +0,0 @@ -#!/usr/bin/env bash -# Helper to migrate WikiWare data from MongoDB to Postgres. -# - Always takes a mongodump backup first. -# - Exports MongoDB collections to NDJSON for inspection. -# - Optionally loads the NDJSON into Postgres staging tables as JSONB. - -set -euo pipefail - -usage() { - cat <<'USAGE' -Usage: ./scripts/mongo_to_postgres_migration.sh [--load-postgres] [--yes] - - --load-postgres Load exported Mongo documents into Postgres staging tables (jsonb). - --yes, -y Skip confirmation prompts. - --help, -h Show this help message. - -The script reads MONGODB_URL, MONGODB_DB_NAME, and POSTGRES_DSN from .env if present. -Defaults: - MONGODB_URL=mongodb://localhost:27017 - MONGODB_DB_NAME=wikiware - POSTGRES_DSN=postgresql://postgres:postgres@localhost:5432/wikiware -USAGE -} - -log() { - printf '[mongo->postgres] %s\n' "$*" -} - -fail() { - printf '[mongo->postgres] ERROR: %s\n' "$*" >&2 - exit 1 -} - -confirm() { - local prompt="${1:-Proceed?} [y/N] " - if [[ ${AUTO_CONFIRM:-0} -eq 1 ]]; then - return 0 - fi - read -r -p "$prompt" reply - [[ "$reply" =~ ^[Yy]$ ]] -} - -require_cmd() { - local name="$1" - command -v "$name" >/dev/null 2>&1 || fail "Missing required command: $name" -} - -uri_with_db() { - local uri="$1" - local db="$2" - if [[ "$uri" =~ /${db}($|[\?#/]) ]]; then - echo "$uri" - return - fi - if [[ "$uri" == *"?"* ]]; then - local base="${uri%%\?*}" - local query="${uri#"$base"}" - echo "${base%/}/$db${query}" - else - echo "${uri%/}/$db" - fi -} - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" -ENV_FILE="$ROOT_DIR/.env" - -if [[ -f "$ENV_FILE" ]]; then - set -o allexport - # shellcheck disable=SC1090 - source "$ENV_FILE" - set +o allexport -fi - -MONGODB_URL="${MONGODB_URL:-mongodb://localhost:27017}" -MONGODB_DB_NAME="${MONGODB_DB_NAME:-wikiware}" -POSTGRES_DSN="${POSTGRES_DSN:-postgresql://postgres:postgres@localhost:5432/wikiware}" -BACKUP_DIR="${BACKUP_DIR:-$ROOT_DIR/backups}" -EXPORT_DIR="${EXPORT_DIR:-$ROOT_DIR/migration_workdir}" -MONGO_EXPORT_DIR="$EXPORT_DIR/mongo_exports" -CSV_EXPORT_DIR="$EXPORT_DIR/postgres_csv" -TIMESTAMP="$(date -u +"%Y-%m-%dT%H%M%SZ")" -AUTO_CONFIRM=0 -LOAD_POSTGRES=0 - -for arg in "$@"; do - case "$arg" in - --load-postgres) LOAD_POSTGRES=1 ;; - --yes|-y) AUTO_CONFIRM=1 ;; - --help|-h) usage; exit 0 ;; - *) fail "Unknown argument: $arg" ;; - esac -done - -MONGODB_URI_WITH_DB="$(uri_with_db "$MONGODB_URL" "$MONGODB_DB_NAME")" - -backup_mongo() { - require_cmd mongodump - mkdir -p "$BACKUP_DIR" - local archive="$BACKUP_DIR/mongodb-${MONGODB_DB_NAME}-${TIMESTAMP}.archive.gz" - - log "Creating MongoDB backup at $archive" - mongodump \ - --uri="$MONGODB_URI_WITH_DB" \ - --archive="$archive" \ - --gzip - log "Backup complete: $archive" -} - -export_collections() { - require_cmd mongoexport - mkdir -p "$MONGO_EXPORT_DIR" - local collections=( - pages - history - branches - users - sessions - image_hashes - analytics_events - ) - - log "Exporting MongoDB collections to $MONGO_EXPORT_DIR" - for collection in "${collections[@]}"; do - local outfile="$MONGO_EXPORT_DIR/${collection}.ndjson" - log " - $collection -> $outfile" - mongoexport \ - --uri="$MONGODB_URI_WITH_DB" \ - --collection="$collection" \ - --type=json \ - --out="$outfile" - done - log "Exports complete." -} - -load_postgres_staging() { - require_cmd perl - require_cmd psql - mkdir -p "$CSV_EXPORT_DIR" - local collections=( - pages - history - branches - users - sessions - image_hashes - analytics_events - ) - - log "Preparing Postgres staging tables in schema wikiware_migration" - psql "$POSTGRES_DSN" -v ON_ERROR_STOP=1 <<'SQL' -CREATE SCHEMA IF NOT EXISTS wikiware_migration; -SQL - - for collection in "${collections[@]}"; do - local ndjson="$MONGO_EXPORT_DIR/${collection}.ndjson" - if [[ ! -s "$ndjson" ]]; then - log " - Skipping $collection (no export found)" - continue - fi - - local csv="$CSV_EXPORT_DIR/${collection}.csv" - perl -pe 's/"/""/g; $_="\"$_\""' "$ndjson" >"$csv" - - log " - Loading $collection into Postgres staging table" - psql "$POSTGRES_DSN" -v ON_ERROR_STOP=1 < Date: Sun, 14 Dec 2025 14:21:23 -0500 Subject: [PATCH 4/9] the fuck --- scripts/mongo_to_postgres_migration.py | 37 +++--- src/database.py | 152 +++++++++++++++++++++---- src/routes/api/stats.py | 5 +- src/services/page_service.py | 96 ++++++---------- src/services/user_service.py | 12 +- src/stats.py | 24 ++-- 6 files changed, 202 insertions(+), 124 deletions(-) diff --git a/scripts/mongo_to_postgres_migration.py b/scripts/mongo_to_postgres_migration.py index c1cbfb6..d4c2dcb 100644 --- a/scripts/mongo_to_postgres_migration.py +++ b/scripts/mongo_to_postgres_migration.py @@ -5,7 +5,7 @@ Steps: 1) Creates a mongodump backup (gzip archive) before touching data. 2) Reads MongoDB collections and writes them into Postgres JSONB storage - (table: wikiware_documents). + using one table per collection (tables prefixed with wikiware_). 3) Upserts documents so the script can be run multiple times. Usage: @@ -39,7 +39,7 @@ from bson import ObjectId from decimal import Decimal -TABLE_NAME = "wikiware_documents" +TABLE_PREFIX = "wikiware_" COLLECTIONS = [ "pages", "history", @@ -93,15 +93,21 @@ def load_env_defaults() -> Dict[str, str]: } -async def ensure_table(pool: asyncpg.Pool) -> None: +def table_name_for(collection: str) -> str: + if not collection or not collection.replace("_", "").isalnum(): + raise ValueError(f"Invalid collection name '{collection}'") + return f"{TABLE_PREFIX}{collection}" + + +async def ensure_table(pool: asyncpg.Pool, collection: str) -> None: + table_name = table_name_for(collection) async with pool.acquire() as conn: await conn.execute( f""" - CREATE TABLE IF NOT EXISTS {TABLE_NAME} ( - collection TEXT NOT NULL, + CREATE TABLE IF NOT EXISTS {table_name} ( id TEXT NOT NULL, doc JSONB NOT NULL, - PRIMARY KEY (collection, id) + PRIMARY KEY (id) ); """ ) @@ -140,14 +146,15 @@ def sanitize_document(doc: Dict) -> Dict: async def upsert_documents( pool: asyncpg.Pool, collection: str, documents: Iterable[Dict] ) -> int: + table_name = table_name_for(collection) inserted = 0 async with pool.acquire() as conn: async with conn.transaction(): stmt = await conn.prepare( f""" - INSERT INTO {TABLE_NAME} (collection, id, doc) - VALUES ($1, $2, $3::jsonb) - ON CONFLICT (collection, id) DO UPDATE SET doc = EXCLUDED.doc + INSERT INTO {table_name} (id, doc) + VALUES ($1, $2::jsonb) + ON CONFLICT (id) DO UPDATE SET doc = EXCLUDED.doc """ ) for doc in documents: @@ -158,7 +165,7 @@ async def upsert_documents( clean_doc = sanitize_document(doc) clean_doc["_id"] = doc_id json_payload = json.dumps(clean_doc, ensure_ascii=False) - await stmt.fetch(collection, doc_id, json_payload) + await stmt.fetch(doc_id, json_payload) inserted += 1 return inserted @@ -183,15 +190,15 @@ async def migrate_collections( mongo_db = mongo_client.get_database() pool = await asyncpg.create_pool(postgres_dsn, min_size=1, max_size=10) - await ensure_table(pool) + for name in COLLECTIONS: + await ensure_table(pool, name) if truncate: async with pool.acquire() as conn: logger.info("Truncating existing rows for selected collections") - await conn.execute( - f"DELETE FROM {TABLE_NAME} WHERE collection = ANY($1::text[])", - COLLECTIONS, - ) + for name in COLLECTIONS: + table_name = table_name_for(name) + await conn.execute(f"DELETE FROM {table_name}") for name in COLLECTIONS: documents = fetch_collection(mongo_db, name) diff --git a/src/database.py b/src/database.py index 8377705..59f4c88 100644 --- a/src/database.py +++ b/src/database.py @@ -11,6 +11,9 @@ import asyncio import os import uuid +import json +import datetime as dt +from decimal import Decimal from dataclasses import dataclass from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Iterable, List, Optional @@ -28,7 +31,31 @@ DB_OPERATION_LOG_THRESHOLD_MS = ( float(os.getenv("DB_OPERATION_LOG_THRESHOLD_MS", "100")) / 1000 ) -TABLE_NAME = "wikiware_documents" +TABLE_PREFIX = "wikiware_" +DEFAULT_COLLECTIONS = ("pages", "history", "branches", "users", "image_hashes") +INDEX_SPECS = { + "pages": [ + ("title_branch", "((doc->>'title'), (doc->>'branch'))"), + ("branch", "((doc->>'branch'))"), + ("updated_at", "((doc->>'updated_at'))"), + ], + "history": [ + ("title_branch", "((doc->>'title'), (doc->>'branch'))"), + ("updated_at", "((doc->>'updated_at'))"), + ], + "branches": [ + ("page_branch", "((doc->>'page_title'), (doc->>'branch_name'))"), + ("created_at", "((doc->>'created_at'))"), + ], + "users": [ + ("username", "((doc->>'username'))"), + ("email", "((doc->>'email'))"), + ], + "image_hashes": [ + ("filename", "((doc->>'filename'))"), + ("sha256", "((doc->>'sha256'))"), + ], +} # ---------------------- Result helpers ---------------------- @@ -152,6 +179,24 @@ def _project_sort_key(doc: Dict[str, Any], key: str) -> Any: return doc.get(key) +def _jsonable(value: Any) -> Any: + if isinstance(value, dt.datetime): + if value.tzinfo is None: + value = value.replace(tzinfo=dt.timezone.utc) + return value.isoformat() + if isinstance(value, dt.date): + return value.isoformat() + if isinstance(value, Decimal): + return float(value) + if isinstance(value, bytes): + return value.decode("utf-8", errors="ignore") + if isinstance(value, dict): + return {k: _jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple, set)): + return [_jsonable(v) for v in value] + return value + + # ---------------------- Cursor ---------------------- @@ -213,16 +258,35 @@ def __iter__(self) -> Iterable[Dict[str, Any]]: class PostgresCollection: - def __init__(self, name: str, db: "Database"): + def __init__(self, name: str, table_name: str, db: "Database"): self.name = name + self._table_name = table_name self._db = db + async def _ensure_table(self) -> None: + await self._db.ensure_table(self._table_name) + async def _fetch_docs(self) -> List[Dict[str, Any]]: + await self._ensure_table() rows = await self._db.fetch( - f"SELECT doc FROM {TABLE_NAME} WHERE collection = $1", - self.name, + f"SELECT doc FROM {self._table_name}", ) - return [row["doc"] for row in rows] + docs: List[Dict[str, Any]] = [] + for row in rows: + doc = row["doc"] + if isinstance(doc, bytes): + doc = doc.decode("utf-8", errors="ignore") + if isinstance(doc, str): + try: + doc = json.loads(doc) + except json.JSONDecodeError: + logger.warning("Skipping non-JSON doc in {}", self._table_name) + continue + if isinstance(doc, dict): + docs.append(doc) + else: + logger.warning("Skipping non-dict doc in {}", self._table_name) + return docs def find( self, @@ -245,18 +309,19 @@ async def find_one(self, filt: Optional[Dict[str, Any]] = None) -> Optional[Dict return results[0] if results else None async def insert_one(self, document: Dict[str, Any]) -> InsertOneResult: + await self._ensure_table() doc = dict(document) if "_id" not in doc: doc["_id"] = str(uuid.uuid4()) + json_payload = json.dumps(_jsonable(doc), ensure_ascii=False) await self._db.execute( f""" - INSERT INTO {TABLE_NAME} (collection, id, doc) - VALUES ($1, $2, $3) - ON CONFLICT (collection, id) DO UPDATE SET doc = EXCLUDED.doc + INSERT INTO {self._table_name} (id, doc) + VALUES ($1, $2::jsonb) + ON CONFLICT (id) DO UPDATE SET doc = EXCLUDED.doc """, - self.name, str(doc["_id"]), - doc, + json_payload, ) return InsertOneResult(inserted_id=str(doc["_id"])) @@ -277,11 +342,11 @@ async def update_one( return UpdateResult(matched_count=0, modified_count=1, upserted_id=result.inserted_id) updated_doc = _apply_update(existing, update) + json_payload = json.dumps(_jsonable(updated_doc), ensure_ascii=False) await self._db.execute( - f"UPDATE {TABLE_NAME} SET doc = $3 WHERE collection = $1 AND id = $2", - self.name, + f"UPDATE {self._table_name} SET doc = $2::jsonb WHERE id = $1", str(existing.get("_id")), - updated_doc, + json_payload, ) return UpdateResult(matched_count=1, modified_count=1, upserted_id=None) @@ -290,8 +355,7 @@ async def delete_one(self, filt: Dict[str, Any]) -> DeleteResult: if existing is None: return DeleteResult(deleted_count=0) await self._db.execute( - f"DELETE FROM {TABLE_NAME} WHERE collection = $1 AND id = $2", - self.name, + f"DELETE FROM {self._table_name} WHERE id = $1", str(existing.get("_id")), ) return DeleteResult(deleted_count=1) @@ -405,11 +469,15 @@ def __init__(self, dsn: str = POSTGRES_DSN): self.is_connected = False self._connection_lock = asyncio.Lock() self._wrapped_collections: Dict[str, PostgresCollection] = {} + self._ensured_tables: set[str] = set() + self._ensured_indexes: set[str] = set() def _reset_state(self) -> None: self.pool = None self.is_connected = False self._wrapped_collections.clear() + self._ensured_tables.clear() + self._ensured_indexes.clear() async def connect(self) -> None: async with self._connection_lock: @@ -418,26 +486,61 @@ async def connect(self) -> None: try: logger.info("Connecting to Postgres at {}", self._dsn) self.pool = await asyncpg.create_pool(self._dsn, min_size=1, max_size=10) - await self._ensure_table() + await self._ensure_tables(DEFAULT_COLLECTIONS) + await self._ensure_indexes(DEFAULT_COLLECTIONS) self.is_connected = True logger.info("Connected to Postgres and storage table ensured") except Exception: logger.exception("Failed to connect to Postgres") self._reset_state() - async def _ensure_table(self) -> None: - if self.pool is None: + def _table_name_for_collection(self, collection: str) -> str: + if not collection or not collection.replace("_", "").isalnum(): + raise ValueError(f"Invalid collection name '{collection}'") + return f"{TABLE_PREFIX}{collection}" + + async def _ensure_tables(self, collections: Iterable[str]) -> None: + for collection in collections: + table_name = self._table_name_for_collection(collection) + await self.ensure_table(table_name) + + async def ensure_table(self, table_name: str) -> None: + if self.pool is None or table_name in self._ensured_tables: return await self.execute( f""" - CREATE TABLE IF NOT EXISTS {TABLE_NAME} ( - collection TEXT NOT NULL, + CREATE TABLE IF NOT EXISTS {table_name} ( id TEXT NOT NULL, doc JSONB NOT NULL, - PRIMARY KEY (collection, id) + PRIMARY KEY (id) ); """ ) + try: + await self.execute( + f"ALTER TABLE {table_name} ALTER COLUMN doc TYPE JSONB USING doc::jsonb;" + ) + except Exception: + logger.warning("Could not coerce doc column to JSONB for {}", table_name) + self._ensured_tables.add(table_name) + + async def _ensure_indexes(self, collections: Iterable[str]) -> None: + for collection in collections: + await self.ensure_indexes_for_collection(collection) + + async def ensure_indexes_for_collection(self, collection: str) -> None: + if self.pool is None: + return + table_name = self._table_name_for_collection(collection) + await self.ensure_table(table_name) + specs = INDEX_SPECS.get(collection, []) + for suffix, expression in specs: + index_name = f"{table_name}_{suffix}_idx" + cache_key = f"{table_name}:{index_name}" + if cache_key in self._ensured_indexes: + continue + await self.execute(f"CREATE INDEX IF NOT EXISTS {index_name} ON {table_name} {expression};") + self._ensured_indexes.add(cache_key) async def disconnect(self) -> None: async with self._connection_lock: @@ -450,7 +553,8 @@ def get_collection(self, name: str) -> Optional[PostgresCollection]: return None if name in self._wrapped_collections: return self._wrapped_collections[name] - collection = PostgresCollection(name, self) + table_name = self._table_name_for_collection(name) + collection = PostgresCollection(name, table_name, self) self._wrapped_collections[name] = collection return collection @@ -502,8 +606,8 @@ def get_image_hashes_collection(): async def create_indexes() -> None: - # JSONB storage uses application-level filtering; explicit DB indexes can be added later. - logger.info("Index creation skipped (JSONB storage)") + logger.info("Ensuring JSONB indexes for default collections") + await db_instance._ensure_indexes(DEFAULT_COLLECTIONS) async def init_database() -> None: diff --git a/src/routes/api/stats.py b/src/routes/api/stats.py index de547bb..61a1021 100644 --- a/src/routes/api/stats.py +++ b/src/routes/api/stats.py @@ -28,10 +28,7 @@ async def get_user_stats(username: str): # Get stats for the specific user users_collection = get_users_collection() - user_stats = await users_collection.find_one( - {"username": username}, - projection={"_id": 0, "total_edits": 1, "page_edits": 1}, - ) + user_stats = await users_collection.find_one({"username": username}) if not user_stats: raise HTTPException(status_code=404, detail="User not found") diff --git a/src/services/page_service.py b/src/services/page_service.py index 24b0c29..93d77e7 100644 --- a/src/services/page_service.py +++ b/src/services/page_service.py @@ -20,6 +20,17 @@ class PageService: """Service class for page-related operations.""" + @staticmethod + def _normalize_timestamp(page: Dict[str, Any]) -> Dict[str, Any]: + """Ensure updated_at is a datetime for template usage.""" + ts = page.get("updated_at") + if isinstance(ts, str): + try: + page["updated_at"] = datetime.fromisoformat(ts) + except ValueError: + page["updated_at"] = None + return page + @staticmethod def _normalize_summary(edit_summary: Optional[str]) -> str: summary = (edit_summary or "").strip() @@ -222,26 +233,24 @@ async def update_page( any_existing_page = await pages_collection.find_one({"title": title}) if not any_existing_page: # Create both main and talk branches for new pages - async with await db_instance.client.start_session() as s: - async with s.start_transaction(): - created_main = await PageService.create_page( - title, - content, - author, - "main", - edit_summary=summary, - edit_permission=edit_permission, - allowed_users=allowed_users or [], - ) - created_talk = await PageService.create_page( - title, - "", - author, - "talk", - edit_summary="wikibot: Auto-created talk page", - edit_permission=edit_permission, - allowed_users=allowed_users or [], - ) + created_main = await PageService.create_page( + title, + content, + author, + "main", + edit_summary=summary, + edit_permission=edit_permission, + allowed_users=allowed_users or [], + ) + created_talk = await PageService.create_page( + title, + "", + author, + "talk", + edit_summary="wikibot: Auto-created talk page", + edit_permission=edit_permission, + allowed_users=allowed_users or [], + ) if created_main and created_talk: await log_action( author, @@ -318,6 +327,7 @@ async def get_pages_by_branch( .sort("updated_at", -1) .to_list(limit) ) + pages = [PageService._normalize_timestamp(p) for p in pages] return pages except Exception as e: logger.error(f"Error getting pages for branch {branch}: {str(e)}") @@ -339,48 +349,12 @@ async def search_pages( logger.error("Pages collection not available") return [] - # Shared helpers - collation = {"locale": "en", "strength": 1} # case + diacritic insensitive - projection = {"score": {"$meta": "textScore"}, "title": 1, "updated_at": 1, "branch": 1} # trim big fields - - try: - cursor = ( - pages_collection - .find( - {"$and": [{"branch": branch}, {"$text": {"$search": query}}]}, - projection - ) - .collation(collation) - .sort([("score", {"$meta": "textScore"}), ("updated_at", -1)]) - ) - pages = await cursor.to_list(limit) - except OperationFailure as op_err: - logger.warning( - f"Text search unavailable, falling back to regex search: {op_err}" - ) - import re - safe = re.escape(query) - cursor = ( - pages_collection - .find( - { - "$and": [ - {"branch": branch}, - { - "$or": [ - {"title": {"$regex": safe, "$options": "i"}}, - {"content": {"$regex": safe, "$options": "i"}}, - ] - }, - ] - }, - {"title": 1, "updated_at": 1, "branch": 1} # no text score in regex mode - ) - .collation(collation) - .sort([("updated_at", -1)]) - ) - pages = await cursor.to_list(limit) + pages = await pages_collection.find( + {"branch": branch, "title": {"$regex": query, "$options": "i"}}, + projection={"title": 1, "updated_at": 1, "branch": 1}, + ).sort("updated_at", -1).to_list(limit) + pages = [PageService._normalize_timestamp(p) for p in pages] logger.info(f"Search performed: {query!r} on branch {branch!r} - found {len(pages)} results") return pages diff --git a/src/services/user_service.py b/src/services/user_service.py index aff6703..aef89fe 100644 --- a/src/services/user_service.py +++ b/src/services/user_service.py @@ -491,11 +491,13 @@ async def get_session(session_id: str) -> Optional[Dict[str, Any]]: # Normalize expires_at to timezone-aware UTC for safe comparison expires_at = session.get("expires_at") - if isinstance(expires_at, datetime): - if expires_at.tzinfo is None: - expires_at = expires_at.replace(tzinfo=timezone.utc) - else: - expires_at = None + if isinstance(expires_at, str): + try: + expires_at = datetime.fromisoformat(expires_at) + except ValueError: + expires_at = None + if isinstance(expires_at, datetime) and expires_at.tzinfo is None: + expires_at = expires_at.replace(tzinfo=timezone.utc) now_utc = datetime.now(timezone.utc) if expires_at and expires_at > now_utc: diff --git a/src/stats.py b/src/stats.py index 750b052..5a6d3aa 100644 --- a/src/stats.py +++ b/src/stats.py @@ -99,21 +99,15 @@ async def get_total_characters(): start = time.perf_counter() pages_collection = get_pages_collection() if pages_collection is not None: - # Aggregate to sum the length of all content - pipeline = [ - {"$project": {"content_length": {"$strLenCP": "$content"}}}, - { - "$group": { - "_id": None, - "total_characters": {"$sum": "$content_length"}, - } - }, - ] - - result = await pages_collection.aggregate(pipeline).to_list(1) - total_characters = result[0]["total_characters"] if result else 0 - - # Update cache — this sets last_character_count_time to a valid datetime + total_characters = 0 + pages = await pages_collection.find({}, projection={"content": 1}).to_list(None) + for page in pages: + content = page.get("content", "") + if isinstance(content, dict): + # Legacy or malformed; skip + continue + total_characters += len(content or "") + last_character_count = total_characters last_character_count_time = datetime.now() end = time.perf_counter() From b85b9cd73b71ae828373a6bcf1708b4a58c8a601 Mon Sep 17 00:00:00 2001 From: superintendent2521 <67196752+superintendent2521@users.noreply.github.com> Date: Sun, 14 Dec 2025 17:43:10 -0500 Subject: [PATCH 5/9] chud life --- scripts/mongo_to_postgres_migration.py | 29 +- src/database.py | 847 ++++++++++++++++++++----- src/services/analytics_service.py | 20 + src/services/branch_service.py | 27 +- src/services/page_service.py | 47 +- src/stats.py | 13 +- src/utils/logs.py | 237 +++---- 7 files changed, 861 insertions(+), 359 deletions(-) diff --git a/scripts/mongo_to_postgres_migration.py b/scripts/mongo_to_postgres_migration.py index d4c2dcb..95f587d 100644 --- a/scripts/mongo_to_postgres_migration.py +++ b/scripts/mongo_to_postgres_migration.py @@ -147,7 +147,20 @@ async def upsert_documents( pool: asyncpg.Pool, collection: str, documents: Iterable[Dict] ) -> int: table_name = table_name_for(collection) - inserted = 0 + records_to_insert = [] + for doc in documents: + doc = dict(doc) + doc_id = str(doc.get("_id") or doc.get("id") or doc.get("uuid") or "") + if not doc_id: + continue + clean_doc = sanitize_document(doc) + clean_doc["_id"] = doc_id + json_payload = json.dumps(clean_doc, ensure_ascii=False) + records_to_insert.append((doc_id, json_payload)) + + if not records_to_insert: + return 0 + async with pool.acquire() as conn: async with conn.transaction(): stmt = await conn.prepare( @@ -157,17 +170,9 @@ async def upsert_documents( ON CONFLICT (id) DO UPDATE SET doc = EXCLUDED.doc """ ) - for doc in documents: - doc = dict(doc) - doc_id = str(doc.get("_id") or doc.get("id") or doc.get("uuid") or "") - if not doc_id: - continue - clean_doc = sanitize_document(doc) - clean_doc["_id"] = doc_id - json_payload = json.dumps(clean_doc, ensure_ascii=False) - await stmt.fetch(doc_id, json_payload) - inserted += 1 - return inserted + await stmt.executemany(records_to_insert) + + return len(records_to_insert) async def migrate_collections( diff --git a/src/database.py b/src/database.py index 59f4c88..2f6366b 100644 --- a/src/database.py +++ b/src/database.py @@ -13,6 +13,8 @@ import uuid import json import datetime as dt +import re +from contextlib import asynccontextmanager from decimal import Decimal from dataclasses import dataclass from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Iterable, List, Optional @@ -32,12 +34,23 @@ float(os.getenv("DB_OPERATION_LOG_THRESHOLD_MS", "100")) / 1000 ) TABLE_PREFIX = "wikiware_" -DEFAULT_COLLECTIONS = ("pages", "history", "branches", "users", "image_hashes") +DEFAULT_COLLECTIONS = ( + "pages", + "history", + "branches", + "users", + "image_hashes", + "analytics_events", + "sessions", + "settings", + "system_logs", +) INDEX_SPECS = { "pages": [ ("title_branch", "((doc->>'title'), (doc->>'branch'))"), ("branch", "((doc->>'branch'))"), ("updated_at", "((doc->>'updated_at'))"), + ("title_trgm", "USING gin ((doc->>'title') gin_trgm_ops)"), ], "history": [ ("title_branch", "((doc->>'title'), (doc->>'branch'))"), @@ -46,6 +59,7 @@ "branches": [ ("page_branch", "((doc->>'page_title'), (doc->>'branch_name'))"), ("created_at", "((doc->>'created_at'))"), + ("branch_name", "((doc->>'branch_name'))"), ], "users": [ ("username", "((doc->>'username'))"), @@ -55,6 +69,23 @@ ("filename", "((doc->>'filename'))"), ("sha256", "((doc->>'sha256'))"), ], + "analytics_events": [ + ("event_type_timestamp", "((doc->>'event_type'), ((doc->>'timestamp')::timestamptz))"), + ("query_normalized_trgm", "USING gin ((doc->>'query_normalized') gin_trgm_ops)"), + ("timestamp_only", "(((doc->>'timestamp'))::timestamptz)"), + ], + "sessions": [ + ("session_id", "((doc->>'session_id'))"), + ("expires_at", "(((doc->>'expires_at'))::timestamptz)"), + ("user_id", "((doc->>'user_id'))"), + ], + "settings": [ + ("doc_id", "((doc->>'_id'))"), + ], + "system_logs": [ + ("action_timestamp", "((doc->>'action'), ((doc->>'timestamp')::timestamptz))"), + ("timestamp_only", "(((doc->>'timestamp'))::timestamptz)"), + ], } @@ -130,9 +161,30 @@ def _matches_filter(doc: Dict[str, Any], filt: Dict[str, Any]) -> bool: if not filt: return True + if "$and" in filt: + return all(_matches_filter(doc, sub) for sub in filt.get("$and", [])) + if "$or" in filt: + return any(_matches_filter(doc, sub) for sub in filt.get("$or", [])) + def compare(val: Any, condition: Any) -> bool: if isinstance(condition, dict): + if "$regex" in condition: + pattern = condition.get("$regex") or "" + options = condition.get("$options", "") + flags = re.IGNORECASE if "i" in options else 0 + try: + return re.search(pattern, str(val or ""), flags) is not None + except re.error: + return False for op, expected in condition.items(): + if op == "$options": + continue + if op == "$exists": + if expected and val is None: + return False + if not expected and val is not None: + return False + continue if op == "$gte" and not (val is not None and val >= expected): return False if op == "$gt" and not (val is not None and val > expected): @@ -203,27 +255,31 @@ def _jsonable(value: Any) -> Any: class PostgresCursor: def __init__( self, - loader: Callable[[], Awaitable[List[Dict[str, Any]]]], - *, + collection: "PostgresCollection", + filt: Optional[Dict[str, Any]], + projection: Optional[Dict[str, int]], limit: Optional[int] = None, + *, + pipeline: Optional[List[Dict[str, Any]]] = None, ): - self._loader = loader + self._collection = collection + self._filt = filt or {} + self._projection = projection self._docs: Optional[List[Dict[str, Any]]] = None self._sorts: List[tuple[str, int]] = [] self._limit = limit + self._pipeline = pipeline async def _ensure_loaded(self) -> None: if self._docs is None: - self._docs = await self._loader() - self._apply_sorts_and_limit() - - def _apply_sorts_and_limit(self) -> None: - if self._docs is None: - return - for key, direction in reversed(self._sorts): - self._docs.sort(key=lambda d: _project_sort_key(d, key) or 0, reverse=direction < 0) - if self._limit is not None: - self._docs = self._docs[: self._limit] + if self._pipeline is not None: + self._docs = await self._collection._run_aggregate_pipeline( + self._pipeline, self._sorts, self._limit + ) + else: + self._docs = await self._collection._find_docs( + self._filt, self._projection, self._sorts, self._limit + ) def sort(self, key: str, direction: int = 1) -> "PostgresCursor": self._sorts.append((key, direction)) @@ -246,12 +302,9 @@ async def iterator(): return iterator() def __iter__(self) -> Iterable[Dict[str, Any]]: - loop = _ensure_loop() - if self._docs is None: - if loop and loop.is_running(): - raise RuntimeError("Synchronous iteration is not supported while loop is running.") - asyncio.run(self._ensure_loaded()) - return iter(self._docs or []) + raise RuntimeError( + "Synchronous iteration is not supported for PostgresCursor. Use async iteration or to_list()." + ) # ---------------------- Collection ---------------------- @@ -263,30 +316,224 @@ def __init__(self, name: str, table_name: str, db: "Database"): self._table_name = table_name self._db = db + def _normalize_filter_value(self, value: Any) -> Any: + if isinstance(value, dt.datetime): + return value if value.tzinfo else value.replace(tzinfo=dt.timezone.utc) + if isinstance(value, (list, tuple)): + return [self._normalize_filter_value(v) for v in value] + return value + + def _value_type(self, value: Any) -> str: + if isinstance(value, (int, float, Decimal)): + return "numeric" + if isinstance(value, dt.datetime): + return "timestamptz" + if isinstance(value, dt.date): + return "date" + if isinstance(value, bool): + return "boolean" + return "text" + + def _value_type_for_iterable(self, values: Iterable[Any]) -> str: + for val in values: + return self._value_type(val) + return "text" + + def _cast_expr_for_type(self, expr: str, value_type: str) -> str: + casts = { + "numeric": "::numeric", + "timestamptz": "::timestamptz", + "date": "::date", + "boolean": "::boolean", + } + suffix = casts.get(value_type) + return f"({expr}){suffix}" if suffix else expr + + def _append_param(self, value: Any, value_type: str, params: List[Any], *, is_array: bool = False) -> str: + normalized = self._normalize_filter_value(value) + params.append(normalized) + placeholder = f"${len(params)}" + if is_array: + casts = { + "numeric": "::numeric[]", + "timestamptz": "::timestamptz[]", + "date": "::date[]", + "boolean": "::boolean[]", + } + return f"{placeholder}{casts.get(value_type, '::text[]')}" + casts = { + "numeric": "::numeric", + "timestamptz": "::timestamptz", + "date": "::date", + "boolean": "::boolean", + } + return f"{placeholder}{casts.get(value_type, '::text')}" + + def _add_path_param(self, path: str, params: List[Any]) -> str: + parts = [p for p in path.split(".") if p] + params.append(parts) + return f"${len(params)}::text[]" + async def _ensure_table(self) -> None: await self._db.ensure_table(self._table_name) - async def _fetch_docs(self) -> List[Dict[str, Any]]: - await self._ensure_table() - rows = await self._db.fetch( - f"SELECT doc FROM {self._table_name}", - ) - docs: List[Dict[str, Any]] = [] - for row in rows: - doc = row["doc"] - if isinstance(doc, bytes): - doc = doc.decode("utf-8", errors="ignore") - if isinstance(doc, str): - try: - doc = json.loads(doc) - except json.JSONDecodeError: - logger.warning("Skipping non-JSON doc in {}", self._table_name) + def _json_path_expr(self, path: str, *, as_text: bool = True) -> str: + if not path: + raise ValueError(f"Invalid field path '{path}'") + if path == "_id": + return "id" + safe_parts = [json.dumps(part)[1:-1] for part in path.split(".")] + path_literal = '","'.join(safe_parts) + operator = "#>>" if as_text else "#>" + return f'doc {operator} \'{{"{path_literal}"}}\'' + + def _build_where_clause(self, filt: Optional[Dict[str, Any]], params: List[Any]) -> str: + if not filt: + return "" + + def _param_value(value: Any) -> str: + normalized = _jsonable(value) + if isinstance(normalized, (dict, list)): + return json.dumps(normalized, ensure_ascii=False) + return normalized + + clauses: List[str] = [] + for key, condition in filt.items(): + if key in ("$and", "$or"): + if not isinstance(condition, list): continue - if isinstance(doc, dict): - docs.append(doc) + nested_clauses = [] + for sub_filter in condition: + nested = self._build_where_clause(sub_filter, params) + if nested: + nested_clauses.append(f"({nested})") + if nested_clauses: + joiner = " AND " if key == "$and" else " OR " + clauses.append(joiner.join(nested_clauses)) + continue + expr = self._json_path_expr(key) + if isinstance(condition, dict): + if "$regex" in condition: + pattern = condition.get("$regex") or "" + options = condition.get("$options", "") + params.append(pattern) + operator = "~*" if "i" in options else "~" + clauses.append(f"{expr} {operator} ${len(params)}") + continue + if "$like" in condition or "$ilike" in condition: + op_key = "$ilike" if "$ilike" in condition else "$like" + operator = "ILIKE" if op_key == "$ilike" else "LIKE" + params.append(_param_value(condition[op_key])) + clauses.append(f"{expr} {operator} ${len(params)}") + continue + for op, expected in condition.items(): + if op == "$options": + continue + raw_expected = expected + if raw_expected is None: + clauses.append(f"{expr} IS NULL") + continue + if op == "$exists": + clauses.append(f"{expr} IS NOT NULL" if raw_expected else f"{expr} IS NULL") + continue + value_type = self._value_type(raw_expected) + typed_expr = self._cast_expr_for_type(expr, value_type) + if op == "$gte": + placeholder = self._append_param(raw_expected, value_type, params) + clauses.append(f"{typed_expr} >= {placeholder}") + elif op == "$gt": + placeholder = self._append_param(raw_expected, value_type, params) + clauses.append(f"{typed_expr} > {placeholder}") + elif op == "$lte": + placeholder = self._append_param(raw_expected, value_type, params) + clauses.append(f"{typed_expr} <= {placeholder}") + elif op == "$lt": + placeholder = self._append_param(raw_expected, value_type, params) + clauses.append(f"{typed_expr} < {placeholder}") + elif op == "$in": + values = expected if isinstance(expected, (list, tuple, set)) else [expected] + value_type = self._value_type_for_iterable(values) + placeholder = self._append_param(list(values), value_type, params, is_array=True) + typed_expr = self._cast_expr_for_type(expr, value_type) + clauses.append(f"{typed_expr} = ANY({placeholder})") + elif op == "$nin": + values = expected if isinstance(expected, (list, tuple, set)) else [expected] + value_type = self._value_type_for_iterable(values) + placeholder = self._append_param(list(values), value_type, params, is_array=True) + typed_expr = self._cast_expr_for_type(expr, value_type) + clauses.append(f"{typed_expr} <> ALL({placeholder})") + else: + placeholder = self._append_param(raw_expected, value_type, params) + clauses.append(f"{typed_expr} = {placeholder}") else: - logger.warning("Skipping non-dict doc in {}", self._table_name) - return docs + if condition is None: + clauses.append(f"{expr} IS NULL") + else: + value_type = self._value_type(condition) + typed_expr = self._cast_expr_for_type(expr, value_type) + placeholder = self._append_param(condition, value_type, params) + clauses.append(f"{typed_expr} = {placeholder}") + + if not clauses: + return "" + return " AND ".join(clauses) + + def _build_projection_clause( + self, projection: Optional[Dict[str, int]] + ) -> tuple[str, Optional[Dict[str, int]]]: + if projection is None: + return "doc AS doc", None + + include_keys = [k for k, v in projection.items() if v] + if include_keys: + parts = [] + for key in include_keys: + parts.append(f"'{key}', {self._json_path_expr(key, as_text=False)}") + if "_id" not in include_keys and projection.get("_id", 1): + parts.append("'_id', id") + fields = ", ".join(parts) + return f"jsonb_strip_nulls(jsonb_build_object({fields})) AS doc", None + + # Defer exclusion projection to Python, as includes were explicitly not requested + return "doc AS doc", projection + + def _build_order_clause(self, sorts: List[tuple[str, int]]) -> str: + if not sorts: + return "" + parts = [] + for key, direction in sorts: + expr = self._json_path_expr(key) + order = "DESC" if direction < 0 else "ASC" + parts.append(f"{expr} {order}") + return ", ".join(parts) + + def _decode_row(self, row: asyncpg.Record, projection: Optional[Dict[str, int]]) -> Dict[str, Any]: + doc = row["doc"] + if projection: + doc = _apply_projection(doc, projection) + return doc + + def _build_update_expression(self, update: Dict[str, Any], params: List[Any]) -> str: + if not update: + raise ValueError("Empty update payload is not supported") + + expr = "doc" + + for path, value in update.get("$set", {}).items(): + path_placeholder = self._add_path_param(path, params) + params.append(_jsonable(self._normalize_filter_value(value))) + value_placeholder = f"${len(params)}::jsonb" + expr = f"jsonb_set({expr}, {path_placeholder}, {value_placeholder}, true)" + + for path, delta in update.get("$inc", {}).items(): + path_placeholder = self._add_path_param(path, params) + params.append(_jsonable(self._normalize_filter_value(delta))) + delta_placeholder = f"${len(params)}::numeric" + current_numeric = f"COALESCE(({expr} #>> {path_placeholder})::numeric, 0)" + increment_expr = f"to_jsonb({current_numeric} + {delta_placeholder})" + expr = f"jsonb_set({expr}, {path_placeholder}, {increment_expr}, true)" + + return expr def find( self, @@ -294,21 +541,42 @@ def find( projection: Optional[Dict[str, int]] = None, limit: Optional[int] = None, ) -> PostgresCursor: - async def loader(): - docs = await self._fetch_docs() - matched = [doc for doc in docs if _matches_filter(doc, filt or {})] - if projection: - matched = [_apply_projection(doc, projection) for doc in matched] - return matched + return PostgresCursor(self, filt, projection, limit) - return PostgresCursor(loader, limit=limit) + async def _find_docs( + self, + filt: Optional[Dict[str, Any]], + projection: Optional[Dict[str, int]], + sorts: List[tuple[str, int]], + limit: Optional[int], + *, + connection: Optional[asyncpg.Connection] = None, + ) -> List[Dict[str, Any]]: + await self._ensure_table() + params: List[Any] = [] + where_clause = self._build_where_clause(filt or {}, params) + projection_clause, projection_for_python = self._build_projection_clause(projection) + order_clause = self._build_order_clause(sorts) + + query = f"SELECT {projection_clause} FROM {self._table_name}" + if where_clause: + query += f" WHERE {where_clause}" + if order_clause: + query += f" ORDER BY {order_clause}" + if limit is not None: + params.append(limit) + query += f" LIMIT ${len(params)}" + + rows = await self._db.fetch(query, *params, conn=connection) + return [self._decode_row(row, projection_for_python) for row in rows] async def find_one(self, filt: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]: - cursor = self.find(filt, limit=1) - results = await cursor.to_list(1) + results = await self._find_docs(filt, None, [], 1) return results[0] if results else None - async def insert_one(self, document: Dict[str, Any]) -> InsertOneResult: + async def insert_one( + self, document: Dict[str, Any], *, connection: Optional[asyncpg.Connection] = None + ) -> InsertOneResult: await self._ensure_table() doc = dict(document) if "_id" not in doc: @@ -322,6 +590,7 @@ async def insert_one(self, document: Dict[str, Any]) -> InsertOneResult: """, str(doc["_id"]), json_payload, + conn=connection, ) return InsertOneResult(inserted_id=str(doc["_id"])) @@ -331,130 +600,286 @@ async def update_one( update: Dict[str, Any], *, upsert: bool = False, + connection: Optional[asyncpg.Connection] = None, ) -> UpdateResult: - existing = await self.find_one(filt) - if existing is None: - if not upsert: - return UpdateResult(matched_count=0, modified_count=0, upserted_id=None) + await self._ensure_table() + params: List[Any] = [] + where_clause = self._build_where_clause(filt or {}, params) + update_expr = self._build_update_expression(update, params) + + target_cte = f"SELECT id FROM {self._table_name}" + if where_clause: + target_cte += f" WHERE {where_clause}" + target_cte += " LIMIT 1" + + query = ( + f"WITH target AS ({target_cte}) " + f"UPDATE {self._table_name} SET doc = {update_expr} " + f"FROM target WHERE target.id = {self._table_name}.id " + "RETURNING 1" + ) + + rows = await self._db.fetch(query, *params, conn=connection) + if rows: + return UpdateResult(matched_count=len(rows), modified_count=len(rows), upserted_id=None) + + if upsert: base = {k: v for k, v in filt.items() if not isinstance(v, dict)} new_doc = _apply_update(base, update) - result = await self.insert_one(new_doc) + result = await self.insert_one(new_doc, connection=connection) return UpdateResult(matched_count=0, modified_count=1, upserted_id=result.inserted_id) - updated_doc = _apply_update(existing, update) - json_payload = json.dumps(_jsonable(updated_doc), ensure_ascii=False) - await self._db.execute( - f"UPDATE {self._table_name} SET doc = $2::jsonb WHERE id = $1", - str(existing.get("_id")), - json_payload, - ) - return UpdateResult(matched_count=1, modified_count=1, upserted_id=None) + return UpdateResult(matched_count=0, modified_count=0, upserted_id=None) - async def delete_one(self, filt: Dict[str, Any]) -> DeleteResult: - existing = await self.find_one(filt) + async def delete_one( + self, filt: Dict[str, Any], *, connection: Optional[asyncpg.Connection] = None + ) -> DeleteResult: + existing_docs = await self._find_docs(filt, None, [], 1, connection=connection) + existing = existing_docs[0] if existing_docs else None if existing is None: return DeleteResult(deleted_count=0) await self._db.execute( f"DELETE FROM {self._table_name} WHERE id = $1", str(existing.get("_id")), + conn=connection, ) return DeleteResult(deleted_count=1) async def count_documents(self, filt: Optional[Dict[str, Any]] = None) -> int: - docs = await self._fetch_docs() - return sum(1 for doc in docs if _matches_filter(doc, filt or {})) + await self._ensure_table() + params: List[Any] = [] + where_clause = self._build_where_clause(filt or {}, params) + query = f"SELECT COUNT(*) AS count FROM {self._table_name}" + if where_clause: + query += f" WHERE {where_clause}" + rows = await self._db.fetch(query, *params) + if not rows: + return 0 + return int(rows[0]["count"]) async def distinct(self, key: str, filt: Optional[Dict[str, Any]] = None) -> List[Any]: - docs = await self._fetch_docs() - values = [] + await self._ensure_table() + params: List[Any] = [] + where_clause = self._build_where_clause(filt or {}, params) + expr = self._json_path_expr(key) + query = f"SELECT DISTINCT {expr} AS value FROM {self._table_name}" + if where_clause: + query += f" WHERE {where_clause}" + rows = await self._db.fetch(query, *params) + return [row["value"] for row in rows if row["value"] is not None] + + async def delete_many( + self, filt: Optional[Dict[str, Any]], *, connection: Optional[asyncpg.Connection] = None + ) -> DeleteResult: + await self._ensure_table() + params: List[Any] = [] + where_clause = self._build_where_clause(filt or {}, params) + query = f"DELETE FROM {self._table_name}" + if where_clause: + query += f" WHERE {where_clause}" + query += " RETURNING 1" + rows = await self._db.fetch(query, *params, conn=connection) + return DeleteResult(deleted_count=len(rows)) + + async def update_many( + self, + filt: Dict[str, Any], + update: Dict[str, Any], + *, + upsert: bool = False, + connection: Optional[asyncpg.Connection] = None, + ) -> UpdateResult: + await self._ensure_table() + params: List[Any] = [] + where_clause = self._build_where_clause(filt or {}, params) + update_expr = self._build_update_expression(update, params) + + query = f"UPDATE {self._table_name} SET doc = {update_expr}" + if where_clause: + query += f" WHERE {where_clause}" + query += " RETURNING 1" + + rows = await self._db.fetch(query, *params, conn=connection) + matched = len(rows) + modified = matched + upserted_id = None + + if upsert and matched == 0: + base = {k: v for k, v in filt.items() if not isinstance(v, dict)} + new_doc = _apply_update(base, update) + result = await self.insert_one(new_doc, connection=connection) + upserted_id = result.inserted_id + modified += 1 + + return UpdateResult(matched_count=matched, modified_count=modified, upserted_id=upserted_id) + + def aggregate(self, pipeline: List[Dict[str, Any]]) -> PostgresCursor: + return PostgresCursor(self, {}, None, None, pipeline=pipeline) + + def _translate_date_format(self, fmt: str) -> str: + replacements = { + "%Y": "YYYY", + "%m": "MM", + "%d": "DD", + "%H": "HH24", + "%M": "MI", + "%S": "SS", + } + for mongo_fmt, pg_fmt in replacements.items(): + fmt = fmt.replace(mongo_fmt, pg_fmt) + return fmt + + def _value_expression(self, value: Any, params: List[Any]) -> str: + if isinstance(value, dict) and "$dateToString" in value: + spec = value["$dateToString"] + fmt = self._translate_date_format(spec.get("format", "")) + date_val = spec.get("date") + if isinstance(date_val, str) and date_val.startswith("$"): + date_expr = self._json_path_expr(date_val[1:]) + return f"to_char({date_expr}::timestamptz, '{fmt}')" + if isinstance(value, str) and value.startswith("$"): + return self._json_path_expr(value[1:]) + params.append(_jsonable(value)) + return f"${len(params)}" + + def _build_group_id_expr(self, group_spec: Any, params: List[Any]) -> tuple[str, List[str]]: + if isinstance(group_spec, dict): + parts = [] + group_by_exprs: List[str] = [] + for key, val in group_spec.items(): + expr = self._value_expression(val, params) + parts.append(f"'{key}', {expr}") + group_by_exprs.append(expr) + return f"jsonb_build_object({', '.join(parts)})", group_by_exprs + if group_spec is None: + return "NULL", [] + expr = self._value_expression(group_spec, params) + return expr, [expr] + + def _build_agg_expression(self, expr: Any, params: List[Any]) -> str: + if isinstance(expr, dict): + if "$sum" in expr: + value_expr = self._value_expression(expr["$sum"], params) + return f"SUM(COALESCE({value_expr}::numeric, 0))" + if "$max" in expr: + value_expr = self._value_expression(expr["$max"], params) + return f"MAX({value_expr})" + if "$first" in expr: + value_expr = self._value_expression(expr["$first"], params) + return f"COALESCE((ARRAY_AGG({value_expr}))[1], NULL)" + return self._value_expression(expr, params) + + def _build_aggregate_order_clause(self, sorts: List[tuple[str, int]]) -> str: + if not sorts: + return "" + parts = [] + for key, direction in sorts: + if key.startswith("_id."): + subkey = key.split(".", 1)[1] + expr = f"_id ->> '{subkey}'" + elif key == "_id": + expr = "_id" + else: + expr = key + order = "DESC" if direction < 0 else "ASC" + parts.append(f"{expr} {order}") + return ", ".join(parts) + + def _apply_project_stage(self, docs: List[Dict[str, Any]], project: Dict[str, Any]) -> List[Dict[str, Any]]: + projected: List[Dict[str, Any]] = [] for doc in docs: - if filt and not _matches_filter(doc, filt): + new_doc: Dict[str, Any] = {} + for key, expr in project.items(): + if expr == 0: + continue + if expr == 1: + new_doc[key] = doc.get(key) + elif isinstance(expr, str) and expr.startswith("$"): + new_doc[key] = _get_by_path(doc, expr[1:]) if "." in expr[1:] else doc.get(expr[1:]) + elif isinstance(expr, dict) and "$dateToString" in expr: + fmt = expr["$dateToString"]["format"] + date_val = _get_by_path(doc, expr["$dateToString"]["date"][1:]) + new_doc[key] = date_val.strftime(fmt) if hasattr(date_val, "strftime") else None + else: + new_doc[key] = expr + projected.append(new_doc) + return projected + + async def _run_aggregate_pipeline( + self, + pipeline: List[Dict[str, Any]], + sorts: List[tuple[str, int]], + limit: Optional[int], + ) -> List[Dict[str, Any]]: + await self._ensure_table() + match_filter: Dict[str, Any] = {} + group_stage: Optional[Dict[str, Any]] = None + project_stage: Optional[Dict[str, Any]] = None + sort_stage: Optional[Dict[str, int]] = None + limit_stage: Optional[int] = None + + for stage in pipeline: + if "$match" in stage: + match_filter.update(stage["$match"]) + elif "$group" in stage: + group_stage = stage["$group"] + elif "$project" in stage: + project_stage = stage["$project"] + elif "$sort" in stage: + sort_stage = stage["$sort"] + elif "$limit" in stage: + limit_stage = stage["$limit"] + + params: List[Any] = [] + where_clause = self._build_where_clause(match_filter, params) + + if group_stage is None: + # Fall back to a find-like query if no grouping is requested + combined_sorts: List[tuple[str, int]] = [] + if sort_stage: + combined_sorts.extend(list(sort_stage.items())) + combined_sorts.extend(sorts) + effective_limit = limit_stage if limit_stage is not None else limit + results = await self._find_docs(match_filter, None, combined_sorts, effective_limit) + if project_stage: + results = self._apply_project_stage(results, project_stage) + return results + + group_id_expr, group_by_exprs = self._build_group_id_expr(group_stage.get("_id"), params) + select_parts = [f"{group_id_expr} AS _id"] + for field, expr in group_stage.items(): + if field == "_id": continue - value = _get_by_path(doc, key) if "." in key else doc.get(key) - if value is not None: - values.append(value) - return list({v for v in values}) + agg_expr = self._build_agg_expression(expr, params) + select_parts.append(f"{agg_expr} AS {field}") + + query = f"SELECT {', '.join(select_parts)} FROM {self._table_name}" + if where_clause: + query += f" WHERE {where_clause}" + if group_by_exprs: + query += " GROUP BY " + ", ".join(group_by_exprs) + + combined_sorts: List[tuple[str, int]] = [] + if sort_stage: + combined_sorts.extend(list(sort_stage.items())) + combined_sorts.extend(sorts) + order_clause = self._build_aggregate_order_clause(combined_sorts) + if order_clause: + query += f" ORDER BY {order_clause}" + + effective_limit = limit_stage if limit_stage is not None else limit + if effective_limit is not None: + params.append(effective_limit) + query += f" LIMIT ${len(params)}" + + rows = await self._db.fetch(query, *params) + docs: List[Dict[str, Any]] = [] + for row in rows: + docs.append({k: row[k] for k in row.keys()}) - def aggregate(self, pipeline: List[Dict[str, Any]]) -> PostgresCursor: - async def loader(): - docs = await self._fetch_docs() - for stage in pipeline: - if "$match" in stage: - docs = [doc for doc in docs if _matches_filter(doc, stage["$match"])] - elif "$project" in stage: - projected = [] - for doc in docs: - new_doc: Dict[str, Any] = {} - for key, expr in stage["$project"].items(): - if expr == 1: - new_doc[key] = doc.get(key) - elif isinstance(expr, str) and expr.startswith("$"): - new_doc[key] = _get_by_path(doc, expr[1:]) if "." in expr[1:] else doc.get(expr[1:]) - elif isinstance(expr, dict) and "$dateToString" in expr: - fmt = expr["$dateToString"]["format"] - date_val = _get_by_path(doc, expr["$dateToString"]["date"][1:]) - new_doc[key] = date_val.strftime(fmt) if hasattr(date_val, "strftime") else None - else: - new_doc[key] = expr - projected.append(new_doc) - docs = projected or docs - elif "$group" in stage: - grouped: Dict[Any, Dict[str, Any]] = {} - group_spec = stage["$group"] - for doc in docs: - group_id_expr = group_spec.get("_id") - if isinstance(group_id_expr, dict): - group_id = {} - for k, v in group_id_expr.items(): - if isinstance(v, dict) and "$dateToString" in v: - date_val = _get_by_path(doc, v["$dateToString"]["date"][1:]) - fmt = v["$dateToString"]["format"] - group_id[k] = date_val.strftime(fmt) if hasattr(date_val, "strftime") else None - elif isinstance(v, str) and v.startswith("$"): - group_id[k] = _get_by_path(doc, v[1:]) if "." in v[1:] else doc.get(v[1:]) - else: - group_id[k] = v - elif isinstance(group_id_expr, str) and group_id_expr.startswith("$"): - group_id = _get_by_path(doc, group_id_expr[1:]) if "." in group_id_expr[1:] else doc.get(group_id_expr[1:]) - else: - group_id = group_id_expr - - if group_id not in grouped: - grouped[group_id] = {"_id": group_id} - - for field, expr in group_spec.items(): - if field == "_id": - continue - if isinstance(expr, dict) and "$sum" in expr: - val = expr["$sum"] - addend = 0 - if isinstance(val, (int, float)): - addend = val - elif isinstance(val, str) and val.startswith("$"): - addend = _get_by_path(doc, val[1:]) if "." in val[1:] else doc.get(val[1:]) or 0 - grouped[group_id][field] = grouped[group_id].get(field, 0) + (addend or 0) - elif isinstance(expr, dict) and "$max" in expr: - candidate = expr["$max"] - if isinstance(candidate, str) and candidate.startswith("$"): - candidate = _get_by_path(doc, candidate[1:]) if "." in candidate[1:] else doc.get(candidate[1:]) - current = grouped[group_id].get(field) - if current is None or (candidate is not None and candidate > current): - grouped[group_id][field] = candidate - elif isinstance(expr, dict) and "$first" in expr: - if field not in grouped[group_id]: - val = expr["$first"] - if isinstance(val, str) and val.startswith("$"): - val = _get_by_path(doc, val[1:]) if "." in val[1:] else doc.get(val[1:]) - grouped[group_id][field] = val - docs = list(grouped.values()) - elif "$sort" in stage: - for key, direction in reversed(list(stage["$sort"].items())): - docs.sort(key=lambda d: _project_sort_key(d, key) or 0, reverse=direction < 0) - elif "$limit" in stage: - docs = docs[: stage["$limit"]] - return docs - - return PostgresCursor(loader) + if project_stage: + docs = self._apply_project_stage(docs, project_stage) + return docs # ---------------------- Database ---------------------- @@ -471,34 +896,90 @@ def __init__(self, dsn: str = POSTGRES_DSN): self._wrapped_collections: Dict[str, PostgresCollection] = {} self._ensured_tables: set[str] = set() self._ensured_indexes: set[str] = set() + self._ensured_extensions: set[str] = set() + self._monitor_task: Optional[asyncio.Task] = None - def _reset_state(self) -> None: + def _reset_state(self, *, preserve_monitor: bool = False) -> None: self.pool = None self.is_connected = False self._wrapped_collections.clear() self._ensured_tables.clear() self._ensured_indexes.clear() + self._ensured_extensions.clear() + if self._monitor_task and not preserve_monitor: + self._monitor_task.cancel() + self._monitor_task = None - async def connect(self) -> None: + async def connect(self, retries: int = 3, initial_delay: float = 1.0) -> None: async with self._connection_lock: if self.is_connected: return - try: - logger.info("Connecting to Postgres at {}", self._dsn) - self.pool = await asyncpg.create_pool(self._dsn, min_size=1, max_size=10) - await self._ensure_tables(DEFAULT_COLLECTIONS) - await self._ensure_indexes(DEFAULT_COLLECTIONS) - self.is_connected = True - logger.info("Connected to Postgres and storage table ensured") - except Exception: - logger.exception("Failed to connect to Postgres") - self._reset_state() + + attempt = 0 + delay = initial_delay + while attempt <= retries and not self.is_connected: + try: + logger.info("Connecting to Postgres at {}", self._dsn) + self.pool = await asyncpg.create_pool(self._dsn, min_size=1, max_size=10) + await self._ensure_extensions() + await self._ensure_tables(DEFAULT_COLLECTIONS) + await self._ensure_indexes(DEFAULT_COLLECTIONS) + self.is_connected = True + logger.info("Connected to Postgres and storage table ensured") + self._start_monitor() + return + except Exception: + attempt += 1 + logger.exception("Failed to connect to Postgres (attempt {})", attempt) + self._reset_state() + if attempt > retries: + break + await asyncio.sleep(delay) + delay = min(delay * 2, 30) + + def _start_monitor(self) -> None: + if self._monitor_task and not self._monitor_task.done(): + return + self._monitor_task = asyncio.create_task(self._monitor_connection()) + + async def _monitor_connection(self) -> None: + try: + while True: + await asyncio.sleep(5) + if self.pool is None: + await self.connect() + continue + try: + async with self.pool.acquire() as conn: + await conn.execute("SELECT 1") + except Exception: + logger.warning("Lost connection to Postgres, attempting to reconnect") + try: + if self.pool: + await self.pool.close() + except Exception: + logger.debug("Error while closing pool during reconnect", exc_info=True) + self._reset_state(preserve_monitor=True) + await self.connect() + except asyncio.CancelledError: + return def _table_name_for_collection(self, collection: str) -> str: if not collection or not collection.replace("_", "").isalnum(): raise ValueError(f"Invalid collection name '{collection}'") return f"{TABLE_PREFIX}{collection}" + async def _ensure_extensions(self) -> None: + if self.pool is None: + return + if "pg_trgm" in self._ensured_extensions: + return + try: + await self.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm;") + self._ensured_extensions.add("pg_trgm") + except Exception: + logger.warning("Could not ensure pg_trgm extension; trigram index may not be available") + async def _ensure_tables(self, collections: Iterable[str]) -> None: for collection in collections: table_name = self._table_name_for_collection(collection) @@ -558,17 +1039,45 @@ def get_collection(self, name: str) -> Optional[PostgresCollection]: self._wrapped_collections[name] = collection return collection - async def execute(self, query: str, *args: Any) -> None: + @asynccontextmanager + async def transaction(self) -> AsyncIterator[asyncpg.Connection]: if self.pool is None: raise RuntimeError("Database not connected") async with self.pool.acquire() as conn: + tx = conn.transaction() + await tx.start() + try: + yield conn + except Exception: + await tx.rollback() + raise + else: + await tx.commit() + + async def execute(self, query: str, *args: Any, conn: Optional[asyncpg.Connection] = None) -> None: + if self.pool is None and conn is None: + raise RuntimeError("Database not connected") + if conn is not None: await conn.execute(query, *args) + return + pool = self.pool + if pool is None: + raise RuntimeError("Database not connected") + async with pool.acquire() as pooled_conn: + await pooled_conn.execute(query, *args) - async def fetch(self, query: str, *args: Any) -> List[asyncpg.Record]: - if self.pool is None: + async def fetch( + self, query: str, *args: Any, conn: Optional[asyncpg.Connection] = None + ) -> List[asyncpg.Record]: + if self.pool is None and conn is None: raise RuntimeError("Database not connected") - async with self.pool.acquire() as conn: + if conn is not None: return await conn.fetch(query, *args) + pool = self.pool + if pool is None: + raise RuntimeError("Database not connected") + async with pool.acquire() as pooled_conn: + return await pooled_conn.fetch(query, *args) async def get_pool_stats(self) -> Dict[str, Any]: if self.pool is None: diff --git a/src/services/analytics_service.py b/src/services/analytics_service.py index a387765..9f73b61 100644 --- a/src/services/analytics_service.py +++ b/src/services/analytics_service.py @@ -24,6 +24,9 @@ class AnalyticsService: """Service responsible for recording and aggregating analytics events.""" _COLLECTION_NAME = "analytics_events" + _RETENTION_DAYS = 90 + _RETENTION_CHECK_INTERVAL = timedelta(hours=1) + _last_retention_check: datetime | None = None @staticmethod def _get_collection(): @@ -32,6 +35,19 @@ def _get_collection(): return None return db_instance.get_collection(AnalyticsService._COLLECTION_NAME) + @classmethod + async def _maybe_enforce_retention(cls, collection) -> None: + """Delete analytics events older than the retention window on a fixed cadence.""" + now = _utcnow() + if cls._last_retention_check and now - cls._last_retention_check < cls._RETENTION_CHECK_INTERVAL: + return + cutoff = now - timedelta(days=cls._RETENTION_DAYS) + try: + await collection.delete_many({"timestamp": {"$lt": cutoff}}) + except Exception as exc: # IGNORE W0718 + logger.warning(f"Failed to enforce analytics retention: {exc}") + cls._last_retention_check = now + @staticmethod async def record_page_view( request: Request, @@ -53,6 +69,7 @@ async def record_page_view( "referrer": request.headers.get("referer"), } await collection.insert_one(event) + await AnalyticsService._maybe_enforce_retention(collection) except Exception as exc: # IGNORE W0718 logger.warning(f"Failed to record page view for {page_title}: {exc}") @@ -79,6 +96,7 @@ async def record_search( "results": result_count, } await collection.insert_one(event) + await AnalyticsService._maybe_enforce_retention(collection) except Exception as exc: # IGNORE W0718 logger.warning(f"Failed to record search query {query!r}: {exc}") @@ -99,6 +117,7 @@ async def record_favorite_added( "branch": branch, } await collection.insert_one(event) + await AnalyticsService._maybe_enforce_retention(collection) except Exception as exc: # IGNORE W0718 logger.warning( f"Failed to record favorite for {page_title!r} on {branch}: {exc}" @@ -121,6 +140,7 @@ async def record_favorite_removed( "branch": branch, } await collection.insert_one(event) + await AnalyticsService._maybe_enforce_retention(collection) except Exception as exc: # IGNORE W0718 logger.warning( f"Failed to record favorite removal for {page_title!r} on {branch}: {exc}" diff --git a/src/services/branch_service.py b/src/services/branch_service.py index 93d7240..de8cb67 100644 --- a/src/services/branch_service.py +++ b/src/services/branch_service.py @@ -35,9 +35,10 @@ async def get_available_branches() -> List[str]: logger.warning("Branches collection not available") return ["main"] - branch_docs = await branches_collection.find().to_list(100) - branches = list(set(["main"] + [doc["branch_name"] for doc in branch_docs])) - return branches + branch_names = await branches_collection.distinct("branch_name") + branches = set(branch_names or []) + branches.add("main") + return list(branches) except Exception as e: logger.error(f"Error getting available branches: {str(e)}") return ["main"] @@ -63,23 +64,19 @@ async def get_branches_for_page(title: str) -> List[str]: branches_collection = get_branches_collection() pages_collection = get_pages_collection() - branches = ["main"] + branches = {"main"} if branches_collection is not None: - # Get branches from branches collection - branch_docs = await branches_collection.find( - {"page_title": title} - ).to_list(100) - branches.extend([doc["branch_name"] for doc in branch_docs]) + branch_docs = await branches_collection.distinct( + "branch_name", {"page_title": title} + ) + branches.update(branch_docs or []) - # Also check pages collection for any branches if pages_collection is not None: - page_docs = await pages_collection.find({"title": title}).to_list(100) - page_branches = [doc["branch"] for doc in page_docs if "branch" in doc] - branches.extend(page_branches) + page_branches = await pages_collection.distinct("branch", {"title": title}) + branches.update([b for b in page_branches or [] if b]) - # Remove duplicates and return - return list(set(branches)) + return list(branches) except Exception as e: logger.error(f"Error getting branches for page {title}: {str(e)}") return ["main"] diff --git a/src/services/page_service.py b/src/services/page_service.py index 93d77e7..098bd39 100644 --- a/src/services/page_service.py +++ b/src/services/page_service.py @@ -79,6 +79,7 @@ async def create_page( edit_summary: Optional[str] = None, edit_permission: str = "everybody", allowed_users: Optional[List[str]] = None, + connection: Any = None, ) -> bool: """ Create a new page. @@ -129,7 +130,7 @@ async def create_page( "allowed_users": normalized_allowed_users, } - await pages_collection.insert_one(page_data) + await pages_collection.insert_one(page_data, connection=connection) logger.info(f"Page created: {title} on branch: {branch} by {author}") return True except Exception as e: @@ -232,25 +233,28 @@ async def update_page( any_existing_page = await pages_collection.find_one({"title": title}) if not any_existing_page: - # Create both main and talk branches for new pages - created_main = await PageService.create_page( - title, - content, - author, - "main", - edit_summary=summary, - edit_permission=edit_permission, - allowed_users=allowed_users or [], - ) - created_talk = await PageService.create_page( - title, - "", - author, - "talk", - edit_summary="wikibot: Auto-created talk page", - edit_permission=edit_permission, - allowed_users=allowed_users or [], - ) + # Create both main and talk branches for new pages atomically + async with db_instance.transaction() as conn: + created_main = await PageService.create_page( + title, + content, + author, + "main", + edit_summary=summary, + edit_permission=edit_permission, + allowed_users=allowed_users or [], + connection=conn, + ) + created_talk = await PageService.create_page( + title, + "", + author, + "talk", + edit_summary="wikibot: Auto-created talk page", + edit_permission=edit_permission, + allowed_users=allowed_users or [], + connection=conn, + ) if created_main and created_talk: await log_action( author, @@ -349,8 +353,9 @@ async def search_pages( logger.error("Pages collection not available") return [] + pattern = f"%{query}%" pages = await pages_collection.find( - {"branch": branch, "title": {"$regex": query, "$options": "i"}}, + {"branch": branch, "title": {"$ilike": pattern}}, projection={"title": 1, "updated_at": 1, "branch": 1}, ).sort("updated_at", -1).to_list(limit) diff --git a/src/stats.py b/src/stats.py index 5a6d3aa..dd716ee 100644 --- a/src/stats.py +++ b/src/stats.py @@ -99,14 +99,11 @@ async def get_total_characters(): start = time.perf_counter() pages_collection = get_pages_collection() if pages_collection is not None: - total_characters = 0 - pages = await pages_collection.find({}, projection={"content": 1}).to_list(None) - for page in pages: - content = page.get("content", "") - if isinstance(content, dict): - # Legacy or malformed; skip - continue - total_characters += len(content or "") + await pages_collection._ensure_table() + rows = await db_instance.fetch( + f"SELECT COALESCE(SUM(LENGTH(doc->>'content')), 0) AS total_chars FROM {pages_collection._table_name}" + ) + total_characters = int(rows[0]["total_chars"]) if rows else 0 last_character_count = total_characters last_character_count_time = datetime.now() diff --git a/src/utils/logs.py b/src/utils/logs.py index 18a35fe..38355e8 100644 --- a/src/utils/logs.py +++ b/src/utils/logs.py @@ -4,9 +4,9 @@ """ from datetime import datetime, timezone -from typing import Dict, Any, Optional +from typing import Dict, Any, Optional, List from loguru import logger -from ..database import get_history_collection, get_branches_collection, db_instance +from ..database import TABLE_PREFIX, db_instance async def get_paginated_logs( @@ -50,16 +50,9 @@ async def get_paginated_logs( } if bypass: logger.warning("Bypass flag enabled - pagination limits are ignored") - - history_collection = get_history_collection() - branches_collection = get_branches_collection() - system_logs_collection = db_instance.get_collection("system_logs") - - if ( - (action_type in (None, "edit") and history_collection is None) - or (action_type in (None, "branch_create") and branches_collection is None) - ): - logger.error("Required collections not available for requested logs") + selects = _build_log_selects(action_type) + if not selects: + logger.warning("No log sources included for requested action type") return { "items": [], "total": 0, @@ -68,28 +61,10 @@ async def get_paginated_logs( "limit": 0 if bypass else sanitized_limit, } - history_count = 0 - if history_collection is not None and ( - action_type == "edit" or action_type is None - ): - history_count = await history_collection.count_documents({}) - - branch_count = 0 - if branches_collection is not None and ( - action_type == "branch_create" or action_type is None - ): - branch_count = await branches_collection.count_documents({}) - - page_create_filter = {"action": "page_create"} - create_count = 0 - if system_logs_collection is not None and ( - action_type == "page_create" or action_type is None - ): - create_count = await system_logs_collection.count_documents( - page_create_filter - ) - - total_items = history_count + branch_count + create_count + union_sql = " UNION ALL ".join(selects) + count_query = f"SELECT COUNT(*) AS count FROM ({union_sql}) AS combined" + count_rows = await db_instance.fetch(count_query) + total_items = int(count_rows[0]["count"]) if count_rows else 0 if total_items == 0: total_pages = 1 @@ -123,107 +98,14 @@ async def get_paginated_logs( offset = (page - 1) * effective_limit current_page = page - items = [] - fetch_limit = total_items if bypass else offset + effective_limit - - if history_collection is not None and ( - action_type == "edit" or action_type is None - ): - history_fetch_limit = ( - history_count if bypass else min(history_count, fetch_limit) - ) - if history_fetch_limit > 0: - history_cursor = history_collection.find().sort("updated_at", -1) - history_items = await history_cursor.limit(history_fetch_limit).to_list( - history_fetch_limit - ) - for item in history_items: - log_author = item.get("edited_by") or item.get("author", "Anonymous") - items.append( - { - "type": "edit", - "title": item["title"], - "author": log_author, - "branch": item["branch"], - "timestamp": item["updated_at"], - "action": "page_edit", - "details": { - "edited_by": log_author, - "previous_author": item.get("author"), - "content_length": ( - len(item.get("content", "")) - if "content" in item - else 0 - ), - }, - } - ) - - if branches_collection is not None and ( - action_type == "branch_create" or action_type is None - ): - branch_fetch_limit = ( - branch_count if bypass else min(branch_count, fetch_limit) - ) - if branch_fetch_limit > 0: - branches_cursor = branches_collection.find().sort("created_at", -1) - branches_items = await branches_cursor.limit(branch_fetch_limit).to_list( - branch_fetch_limit - ) - for item in branches_items: - items.append( - { - "type": "branch_create", - "title": item["page_title"], - "author": "System", - "branch": item["branch_name"], - "timestamp": item["created_at"], - "action": "branch_create", - "details": {"source_branch": item["created_from"]}, - } - ) - - if system_logs_collection is not None and ( - action_type == "page_create" or action_type is None - ): - creation_fetch_limit = ( - create_count if bypass else min(create_count, fetch_limit) - ) - if creation_fetch_limit > 0: - creations_cursor = system_logs_collection.find(page_create_filter).sort( - "timestamp", -1 - ) - creation_logs = await creations_cursor.limit( - creation_fetch_limit - ).to_list(creation_fetch_limit) - for item in creation_logs: - metadata = item.get("metadata") or {} - timestamp = item.get("timestamp") or datetime.now(timezone.utc) - items.append( - { - "type": "page_create", - "title": metadata.get("title") - or metadata.get("page_title") - or item.get("message", ""), - "author": metadata.get("author") - or item.get("username", "Unknown"), - "branch": metadata.get("branch", "main"), - "timestamp": timestamp, - "action": "page_create", - "details": { - "created_by": metadata.get("author") - or item.get("username"), - "branch": metadata.get("branch", "main"), - }, - } - ) - - items.sort(key=lambda x: x["timestamp"], reverse=True) - + data_params: List[Any] = [] + data_query = f"SELECT * FROM ({union_sql}) AS combined ORDER BY timestamp DESC" if not bypass: - start_index = offset - end_index = offset + effective_limit - items = items[start_index:end_index] + data_query += f" OFFSET ${len(data_params) + 1} LIMIT ${len(data_params) + 2}" + data_params.extend([offset, effective_limit]) + + rows = await db_instance.fetch(data_query, *data_params) + items = _rows_to_items(rows) response_limit = len(items) @@ -240,6 +122,93 @@ async def get_paginated_logs( raise e +def _build_log_selects(action_type: Optional[str]) -> List[str]: + history_table = f"{TABLE_PREFIX}history" + branches_table = f"{TABLE_PREFIX}branches" + system_logs_table = f"{TABLE_PREFIX}system_logs" + + selects: List[str] = [] + + if action_type in (None, "edit"): + selects.append( + f""" + SELECT + 'edit' AS type, + doc #>> '{{title}}' AS title, + COALESCE(NULLIF(doc #>> '{{edited_by}}', ''), doc #>> '{{author}}', 'Anonymous') AS author, + COALESCE(doc #>> '{{branch}}', 'main') AS branch, + (doc #>> '{{updated_at}}')::timestamptz AS timestamp, + 'page_edit' AS action, + jsonb_build_object( + 'edited_by', COALESCE(NULLIF(doc #>> '{{edited_by}}', ''), doc #>> '{{author}}', 'Anonymous'), + 'previous_author', doc #>> '{{author}}', + 'content_length', COALESCE(length(doc #>> '{{content}}'), 0) + ) AS details + FROM {history_table} + WHERE doc ? 'updated_at' + """ + ) + + if action_type in (None, "branch_create"): + selects.append( + f""" + SELECT + 'branch_create' AS type, + doc #>> '{{page_title}}' AS title, + 'System' AS author, + doc #>> '{{branch_name}}' AS branch, + (doc #>> '{{created_at}}')::timestamptz AS timestamp, + 'branch_create' AS action, + jsonb_build_object('source_branch', doc #>> '{{created_from}}') AS details + FROM {branches_table} + WHERE doc ? 'created_at' + """ + ) + + if action_type in (None, "page_create"): + selects.append( + f""" + SELECT + 'page_create' AS type, + COALESCE( + NULLIF(doc #>> '{{metadata,title}}', ''), + NULLIF(doc #>> '{{metadata,page_title}}', ''), + doc #>> '{{message}}', + '' + ) AS title, + COALESCE(doc #>> '{{metadata,author}}', doc #>> '{{username}}', 'Unknown') AS author, + COALESCE(doc #>> '{{metadata,branch}}', 'main') AS branch, + (doc #>> '{{timestamp}}')::timestamptz AS timestamp, + 'page_create' AS action, + jsonb_build_object( + 'created_by', COALESCE(doc #>> '{{metadata,author}}', doc #>> '{{username}}'), + 'branch', COALESCE(doc #>> '{{metadata,branch}}', 'main') + ) AS details + FROM {system_logs_table} + WHERE (doc #>> '{{action}}') = 'page_create' AND doc ? 'timestamp' + """ + ) + + return selects + + +def _rows_to_items(rows: List[Any]) -> List[Dict[str, Any]]: + items: List[Dict[str, Any]] = [] + for row in rows: + items.append( + { + "type": row.get("type"), + "title": row.get("title"), + "author": row.get("author"), + "branch": row.get("branch"), + "timestamp": row.get("timestamp"), + "action": row.get("action"), + "details": row.get("details") or {}, + } + ) + return items + + async def log_action( username: str, action: str, From 575610a901662603ccaf2d009a52f14e359f8d27 Mon Sep 17 00:00:00 2001 From: Superintendent <67196752+superintendent2521@users.noreply.github.com> Date: Sun, 14 Dec 2025 17:50:11 -0500 Subject: [PATCH 6/9] Update src/database.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/database.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/database.py b/src/database.py index 2f6366b..df777eb 100644 --- a/src/database.py +++ b/src/database.py @@ -570,8 +570,8 @@ async def _find_docs( rows = await self._db.fetch(query, *params, conn=connection) return [self._decode_row(row, projection_for_python) for row in rows] - async def find_one(self, filt: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]: - results = await self._find_docs(filt, None, [], 1) + async def find_one(self, filt: Optional[Dict[str, Any]] = None, projection: Optional[Dict[str, int]] = None) -> Optional[Dict[str, Any]]: + results = await self._find_docs(filt, projection, [], 1) return results[0] if results else None async def insert_one( From f1a5557f1d229a0b67dca203bad467d8d4838e17 Mon Sep 17 00:00:00 2001 From: superintendent2521 <67196752+superintendent2521@users.noreply.github.com> Date: Sun, 14 Dec 2025 17:52:21 -0500 Subject: [PATCH 7/9] fixes --- src/database.py | 95 +++++++++++-------------------------------------- 1 file changed, 21 insertions(+), 74 deletions(-) diff --git a/src/database.py b/src/database.py index 2f6366b..20137c1 100644 --- a/src/database.py +++ b/src/database.py @@ -112,13 +112,6 @@ class DeleteResult: # ---------------------- Utility helpers ---------------------- -def _ensure_loop() -> Optional[asyncio.AbstractEventLoop]: - try: - return asyncio.get_running_loop() - except RuntimeError: - return None - - def _get_by_path(doc: Dict[str, Any], path: str) -> Any: parts = path.split(".") current: Any = doc @@ -157,56 +150,6 @@ def _apply_projection(doc: Dict[str, Any], projection: Optional[Dict[str, int]]) return doc -def _matches_filter(doc: Dict[str, Any], filt: Dict[str, Any]) -> bool: - if not filt: - return True - - if "$and" in filt: - return all(_matches_filter(doc, sub) for sub in filt.get("$and", [])) - if "$or" in filt: - return any(_matches_filter(doc, sub) for sub in filt.get("$or", [])) - - def compare(val: Any, condition: Any) -> bool: - if isinstance(condition, dict): - if "$regex" in condition: - pattern = condition.get("$regex") or "" - options = condition.get("$options", "") - flags = re.IGNORECASE if "i" in options else 0 - try: - return re.search(pattern, str(val or ""), flags) is not None - except re.error: - return False - for op, expected in condition.items(): - if op == "$options": - continue - if op == "$exists": - if expected and val is None: - return False - if not expected and val is not None: - return False - continue - if op == "$gte" and not (val is not None and val >= expected): - return False - if op == "$gt" and not (val is not None and val > expected): - return False - if op == "$lte" and not (val is not None and val <= expected): - return False - if op == "$lt" and not (val is not None and val < expected): - return False - if op == "$in" and val not in expected: - return False - if op == "$nin" and val in expected: - return False - return True - return val == condition - - for key, expected in filt.items(): - value = _get_by_path(doc, key) if "." in key else doc.get(key) - if not compare(value, expected): - return False - return True - - def _apply_update(doc: Dict[str, Any], update: Dict[str, Any]) -> Dict[str, Any]: updated = dict(doc) for op, changes in update.items(): @@ -224,13 +167,6 @@ def _apply_update(doc: Dict[str, Any], update: Dict[str, Any]) -> Dict[str, Any] logger.warning("Unsupported update operator {}", op) return updated - -def _project_sort_key(doc: Dict[str, Any], key: str) -> Any: - if "." in key: - return _get_by_path(doc, key) - return doc.get(key) - - def _jsonable(value: Any) -> Any: if isinstance(value, dt.datetime): if value.tzinfo is None: @@ -630,20 +566,31 @@ async def update_one( return UpdateResult(matched_count=0, modified_count=1, upserted_id=result.inserted_id) return UpdateResult(matched_count=0, modified_count=0, upserted_id=None) - async def delete_one( self, filt: Dict[str, Any], *, connection: Optional[asyncpg.Connection] = None ) -> DeleteResult: - existing_docs = await self._find_docs(filt, None, [], 1, connection=connection) - existing = existing_docs[0] if existing_docs else None - if existing is None: - return DeleteResult(deleted_count=0) - await self._db.execute( - f"DELETE FROM {self._table_name} WHERE id = $1", - str(existing.get("_id")), - conn=connection, + await self._ensure_table() + params: List[Any] = [] + where_clause = self._build_where_clause(filt or {}, params) + + target_cte = f"SELECT id FROM {self._table_name}" + if where_clause: + target_cte += f" WHERE {where_clause}" + target_cte += " LIMIT 1" + + query = ( + f"WITH target AS ({target_cte}), " + f"deleted AS (" + f"DELETE FROM {self._table_name} t USING target " + f"WHERE t.id = target.id RETURNING 1" + f") " + "SELECT COUNT(*) AS count FROM deleted" ) - return DeleteResult(deleted_count=1) + + rows = await self._db.fetch(query, *params, conn=connection) + if not rows: + return DeleteResult(deleted_count=0) + return DeleteResult(deleted_count=int(rows[0]["count"])) async def count_documents(self, filt: Optional[Dict[str, Any]] = None) -> int: await self._ensure_table() From 381a17ae7f62389bed899fa876f05b3a82e3c9c7 Mon Sep 17 00:00:00 2001 From: superintendent2521 <67196752+superintendent2521@users.noreply.github.com> Date: Sun, 14 Dec 2025 18:25:44 -0500 Subject: [PATCH 8/9] fix critical --- src/database.py | 35 +++++++++++++++++++++++- src/services/page_service.py | 53 ++++++++++++++++++++++++++++++++---- 2 files changed, 82 insertions(+), 6 deletions(-) diff --git a/src/database.py b/src/database.py index 3fce58c..ae2b3e8 100644 --- a/src/database.py +++ b/src/database.py @@ -51,6 +51,10 @@ ("branch", "((doc->>'branch'))"), ("updated_at", "((doc->>'updated_at'))"), ("title_trgm", "USING gin ((doc->>'title') gin_trgm_ops)"), + ( + "content_fts", + "USING gin (to_tsvector('english', coalesce(doc->>'title', '') || ' ' || coalesce(doc->>'content', '')))", + ), ], "history": [ ("title_branch", "((doc->>'title'), (doc->>'branch'))"), @@ -628,6 +632,35 @@ async def delete_many( rows = await self._db.fetch(query, *params, conn=connection) return DeleteResult(deleted_count=len(rows)) + def _extract_upsert_base(self, filt: Dict[str, Any]) -> Dict[str, Any]: + """Collect simple equality filters to seed an upserted document.""" + base: Dict[str, Any] = {} + + def set_if_consistent(key: str, value: Any) -> None: + if key in base and base[key] != value: + base.pop(key, None) + return + base[key] = value + + def visit(obj: Any) -> None: + if not isinstance(obj, dict): + return + for key, condition in obj.items(): + if key == "$and" and isinstance(condition, list): + for sub in condition: + visit(sub) + continue + if key.startswith("$"): + continue + if isinstance(condition, dict): + if "$eq" in condition: + set_if_consistent(key, condition["$eq"]) + continue + set_if_consistent(key, condition) + + visit(filt) + return base + async def update_many( self, filt: Dict[str, Any], @@ -652,7 +685,7 @@ async def update_many( upserted_id = None if upsert and matched == 0: - base = {k: v for k, v in filt.items() if not isinstance(v, dict)} + base = self._extract_upsert_base(filt) new_doc = _apply_update(base, update) result = await self.insert_one(new_doc, connection=connection) upserted_id = result.inserted_id diff --git a/src/services/page_service.py b/src/services/page_service.py index 098bd39..9889e3d 100644 --- a/src/services/page_service.py +++ b/src/services/page_service.py @@ -8,6 +8,7 @@ from loguru import logger from pymongo.errors import OperationFailure from ..database import ( + TABLE_PREFIX, get_pages_collection, get_history_collection, get_users_collection, @@ -353,11 +354,53 @@ async def search_pages( logger.error("Pages collection not available") return [] - pattern = f"%{query}%" - pages = await pages_collection.find( - {"branch": branch, "title": {"$ilike": pattern}}, - projection={"title": 1, "updated_at": 1, "branch": 1}, - ).sort("updated_at", -1).to_list(limit) + normalized_query = query.strip() + if not normalized_query: + logger.info("Search attempted with empty query") + return [] + + table_name = f"{TABLE_PREFIX}pages" + search_sql = f""" + WITH search AS ( + SELECT + doc, + to_tsvector( + 'english', + coalesce(doc->>'title', '') || ' ' || coalesce(doc->>'content', '') + ) AS search_vector, + websearch_to_tsquery('english', $1) AS search_query + FROM {table_name} + ) + SELECT + doc #>> '{{title}}' AS title, + doc #>> '{{content}}' AS content, + NULLIF(doc #>> '{{updated_at}}', '')::timestamptz AS updated_at, + coalesce(doc #>> '{{branch}}', 'main') AS branch, + coalesce(NULLIF(doc #>> '{{author}}', ''), 'Anonymous') AS author, + ts_rank(search_vector, search_query) AS rank + FROM search + WHERE coalesce(doc #>> '{{branch}}', 'main') = $2 + AND search_vector @@ search_query + ORDER BY rank DESC, updated_at DESC + LIMIT $3 + """ + + effective_limit = 100 if limit is None else limit + rows = await db_instance.fetch( + search_sql, normalized_query, branch or "main", effective_limit + ) + + pages = [] + for row in rows: + pages.append( + { + "title": row["title"], + "content": row["content"], + "updated_at": row["updated_at"], + "branch": row["branch"], + "author": row["author"], + } + ) pages = [PageService._normalize_timestamp(p) for p in pages] logger.info(f"Search performed: {query!r} on branch {branch!r} - found {len(pages)} results") From bcde4e7f8743f3fc13cbc710fffc1f967b53ed30 Mon Sep 17 00:00:00 2001 From: superintendent2521 <67196752+superintendent2521@users.noreply.github.com> Date: Tue, 20 Jan 2026 20:09:57 -0500 Subject: [PATCH 9/9] ruff/black --- migrate_uploads_to_s3.py | 29 ++- scripts/mongo_to_postgres_migration.py | 18 +- src/config.py | 4 +- src/database.py | 181 ++++++++++---- src/middleware/user_agent_middleware.py | 13 +- src/models/user.py | 1 + src/routes/api/__init__.py | 12 +- src/routes/api/admin.py | 2 +- src/routes/api/exports.py | 4 +- src/routes/api/favorites.py | 2 +- src/routes/api/images.py | 7 +- src/routes/api/logs.py | 2 - src/routes/api/page_markdown.py | 1 + src/routes/api/pdf.py | 4 +- src/routes/api/uploads.py | 5 +- src/routes/media.py | 1 - src/routes/web/admin.py | 2 +- src/routes/web/auth.py | 10 +- src/routes/web/branches.py | 6 +- src/routes/web/pages.py | 6 +- src/server.py | 5 +- src/services/analytics_service.py | 9 +- src/services/branch_service.py | 4 +- src/services/page_service.py | 10 +- src/services/settings_service.py | 6 +- src/services/storage_service.py | 52 ++-- src/services/user_service.py | 14 +- src/stats.py | 9 +- src/tests/test.py | 313 +++++++++++++++--------- src/utils/logs.py | 8 +- src/utils/navigation_history.py | 17 +- src/utils/validation.py | 2 - 32 files changed, 488 insertions(+), 271 deletions(-) diff --git a/migrate_uploads_to_s3.py b/migrate_uploads_to_s3.py index eeb16c6..f1fe9ea 100644 --- a/migrate_uploads_to_s3.py +++ b/migrate_uploads_to_s3.py @@ -46,12 +46,16 @@ async def _upload_file(path: Path) -> tuple[str, str, int, int]: async def migrate(delete_local: bool) -> int: if not is_s3_configured(): - logger.error("S3 is not configured. Update src/config.py before running this script.") + logger.error( + "S3 is not configured. Update src/config.py before running this script." + ) return 1 upload_path = Path(UPLOAD_DIR) if not upload_path.exists(): - logger.info("Upload directory %s does not exist; nothing to migrate.", upload_path) + logger.info( + "Upload directory %s does not exist; nothing to migrate.", upload_path + ) return 0 await db_instance.connect() @@ -70,10 +74,14 @@ async def migrate(delete_local: bool) -> int: existing_doc = None if collection is not None: existing_doc = await collection.find_one({"filename": entry.name}) - if existing_doc and existing_doc.get("url") and not existing_doc["url"].startswith( - "/static/uploads/" + if ( + existing_doc + and existing_doc.get("url") + and not existing_doc["url"].startswith("/static/uploads/") ): - logger.info("Skipping %s (already points to remote storage).", entry.name) + logger.info( + "Skipping %s (already points to remote storage).", entry.name + ) skipped += 1 continue @@ -115,7 +123,12 @@ async def migrate(delete_local: bool) -> int: await db_instance.disconnect() - logger.info("Migration complete: %s migrated, %s skipped, %s failures.", migrated, skipped, len(failures)) + logger.info( + "Migration complete: %s migrated, %s skipped, %s failures.", + migrated, + skipped, + len(failures), + ) if failures: logger.error("Failures: %s", ", ".join(failures)) return 2 @@ -123,7 +136,9 @@ async def migrate(delete_local: bool) -> int: def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Migrate local uploads into S3 storage.") + parser = argparse.ArgumentParser( + description="Migrate local uploads into S3 storage." + ) parser.add_argument( "--delete-local", action="store_true", diff --git a/scripts/mongo_to_postgres_migration.py b/scripts/mongo_to_postgres_migration.py index 95f587d..ba2538a 100644 --- a/scripts/mongo_to_postgres_migration.py +++ b/scripts/mongo_to_postgres_migration.py @@ -30,7 +30,7 @@ import os import subprocess from pathlib import Path -from typing import Dict, Iterable, List +from typing import Dict, Iterable import asyncpg from dotenv import load_dotenv @@ -215,9 +215,19 @@ async def migrate_collections( def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Migrate MongoDB data into Postgres JSONB storage.") - parser.add_argument("--skip-backup", action="store_true", help="Do not run mongodump before migrating.") - parser.add_argument("--truncate", action="store_true", help="Delete existing rows for migrated collections first.") + parser = argparse.ArgumentParser( + description="Migrate MongoDB data into Postgres JSONB storage." + ) + parser.add_argument( + "--skip-backup", + action="store_true", + help="Do not run mongodump before migrating.", + ) + parser.add_argument( + "--truncate", + action="store_true", + help="Delete existing rows for migrated collections first.", + ) return parser.parse_args() diff --git a/src/config.py b/src/config.py index 177e06b..cd2a005 100644 --- a/src/config.py +++ b/src/config.py @@ -5,7 +5,7 @@ import os from dotenv import load_dotenv -from .env_config import SECRET_KEY, ACCESS_KEY # Importing sensitive data +from .env_config import SECRET_KEY, ACCESS_KEY # Importing sensitive data # Load environment variables load_dotenv() @@ -36,7 +36,7 @@ S3_PUBLIC_URL = "" # Optional public CDN/base URL override # Logging settings -#UNUSED +# UNUSED LOG_DIR = "logs" LOG_RETENTION = "7 days" LOG_LEVEL = "INFO" diff --git a/src/database.py b/src/database.py index ae2b3e8..fd6a1d2 100644 --- a/src/database.py +++ b/src/database.py @@ -13,17 +13,22 @@ import uuid import json import datetime as dt -import re from contextlib import asynccontextmanager from decimal import Decimal from dataclasses import dataclass -from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Iterable, List, Optional +from typing import ( + Any, + AsyncIterator, + Dict, + Iterable, + List, + Optional, +) import asyncpg from dotenv import load_dotenv from loguru import logger -from . import config load_dotenv() @@ -74,8 +79,14 @@ ("sha256", "((doc->>'sha256'))"), ], "analytics_events": [ - ("event_type_timestamp", "((doc->>'event_type'), ((doc->>'timestamp')::timestamptz))"), - ("query_normalized_trgm", "USING gin ((doc->>'query_normalized') gin_trgm_ops)"), + ( + "event_type_timestamp", + "((doc->>'event_type'), ((doc->>'timestamp')::timestamptz))", + ), + ( + "query_normalized_trgm", + "USING gin ((doc->>'query_normalized') gin_trgm_ops)", + ), ("timestamp_only", "(((doc->>'timestamp'))::timestamptz)"), ], "sessions": [ @@ -136,7 +147,9 @@ def _set_by_path(doc: Dict[str, Any], path: str, value: Any) -> None: current[parts[-1]] = value -def _apply_projection(doc: Dict[str, Any], projection: Optional[Dict[str, int]]) -> Dict[str, Any]: +def _apply_projection( + doc: Dict[str, Any], projection: Optional[Dict[str, int]] +) -> Dict[str, Any]: if projection is None: return doc include_keys = {k for k, v in projection.items() if v} @@ -171,6 +184,7 @@ def _apply_update(doc: Dict[str, Any], update: Dict[str, Any]) -> Dict[str, Any] logger.warning("Unsupported update operator {}", op) return updated + def _jsonable(value: Any) -> Any: if isinstance(value, dt.datetime): if value.tzinfo is None: @@ -289,7 +303,9 @@ def _cast_expr_for_type(self, expr: str, value_type: str) -> str: suffix = casts.get(value_type) return f"({expr}){suffix}" if suffix else expr - def _append_param(self, value: Any, value_type: str, params: List[Any], *, is_array: bool = False) -> str: + def _append_param( + self, value: Any, value_type: str, params: List[Any], *, is_array: bool = False + ) -> str: normalized = self._normalize_filter_value(value) params.append(normalized) placeholder = f"${len(params)}" @@ -325,9 +341,11 @@ def _json_path_expr(self, path: str, *, as_text: bool = True) -> str: safe_parts = [json.dumps(part)[1:-1] for part in path.split(".")] path_literal = '","'.join(safe_parts) operator = "#>>" if as_text else "#>" - return f'doc {operator} \'{{"{path_literal}"}}\'' + return f"doc {operator} '{{\"{path_literal}\"}}'" - def _build_where_clause(self, filt: Optional[Dict[str, Any]], params: List[Any]) -> str: + def _build_where_clause( + self, filt: Optional[Dict[str, Any]], params: List[Any] + ) -> str: if not filt: return "" @@ -374,36 +392,60 @@ def _param_value(value: Any) -> str: clauses.append(f"{expr} IS NULL") continue if op == "$exists": - clauses.append(f"{expr} IS NOT NULL" if raw_expected else f"{expr} IS NULL") + clauses.append( + f"{expr} IS NOT NULL" if raw_expected else f"{expr} IS NULL" + ) continue value_type = self._value_type(raw_expected) typed_expr = self._cast_expr_for_type(expr, value_type) if op == "$gte": - placeholder = self._append_param(raw_expected, value_type, params) + placeholder = self._append_param( + raw_expected, value_type, params + ) clauses.append(f"{typed_expr} >= {placeholder}") elif op == "$gt": - placeholder = self._append_param(raw_expected, value_type, params) + placeholder = self._append_param( + raw_expected, value_type, params + ) clauses.append(f"{typed_expr} > {placeholder}") elif op == "$lte": - placeholder = self._append_param(raw_expected, value_type, params) + placeholder = self._append_param( + raw_expected, value_type, params + ) clauses.append(f"{typed_expr} <= {placeholder}") elif op == "$lt": - placeholder = self._append_param(raw_expected, value_type, params) + placeholder = self._append_param( + raw_expected, value_type, params + ) clauses.append(f"{typed_expr} < {placeholder}") elif op == "$in": - values = expected if isinstance(expected, (list, tuple, set)) else [expected] + values = ( + expected + if isinstance(expected, (list, tuple, set)) + else [expected] + ) value_type = self._value_type_for_iterable(values) - placeholder = self._append_param(list(values), value_type, params, is_array=True) + placeholder = self._append_param( + list(values), value_type, params, is_array=True + ) typed_expr = self._cast_expr_for_type(expr, value_type) clauses.append(f"{typed_expr} = ANY({placeholder})") elif op == "$nin": - values = expected if isinstance(expected, (list, tuple, set)) else [expected] + values = ( + expected + if isinstance(expected, (list, tuple, set)) + else [expected] + ) value_type = self._value_type_for_iterable(values) - placeholder = self._append_param(list(values), value_type, params, is_array=True) + placeholder = self._append_param( + list(values), value_type, params, is_array=True + ) typed_expr = self._cast_expr_for_type(expr, value_type) clauses.append(f"{typed_expr} <> ALL({placeholder})") else: - placeholder = self._append_param(raw_expected, value_type, params) + placeholder = self._append_param( + raw_expected, value_type, params + ) clauses.append(f"{typed_expr} = {placeholder}") else: if condition is None: @@ -447,13 +489,17 @@ def _build_order_clause(self, sorts: List[tuple[str, int]]) -> str: parts.append(f"{expr} {order}") return ", ".join(parts) - def _decode_row(self, row: asyncpg.Record, projection: Optional[Dict[str, int]]) -> Dict[str, Any]: + def _decode_row( + self, row: asyncpg.Record, projection: Optional[Dict[str, int]] + ) -> Dict[str, Any]: doc = row["doc"] if projection: doc = _apply_projection(doc, projection) return doc - def _build_update_expression(self, update: Dict[str, Any], params: List[Any]) -> str: + def _build_update_expression( + self, update: Dict[str, Any], params: List[Any] + ) -> str: if not update: raise ValueError("Empty update payload is not supported") @@ -495,7 +541,9 @@ async def _find_docs( await self._ensure_table() params: List[Any] = [] where_clause = self._build_where_clause(filt or {}, params) - projection_clause, projection_for_python = self._build_projection_clause(projection) + projection_clause, projection_for_python = self._build_projection_clause( + projection + ) order_clause = self._build_order_clause(sorts) query = f"SELECT {projection_clause} FROM {self._table_name}" @@ -510,12 +558,19 @@ async def _find_docs( rows = await self._db.fetch(query, *params, conn=connection) return [self._decode_row(row, projection_for_python) for row in rows] - async def find_one(self, filt: Optional[Dict[str, Any]] = None, projection: Optional[Dict[str, int]] = None) -> Optional[Dict[str, Any]]: + async def find_one( + self, + filt: Optional[Dict[str, Any]] = None, + projection: Optional[Dict[str, int]] = None, + ) -> Optional[Dict[str, Any]]: results = await self._find_docs(filt, projection, [], 1) return results[0] if results else None async def insert_one( - self, document: Dict[str, Any], *, connection: Optional[asyncpg.Connection] = None + self, + document: Dict[str, Any], + *, + connection: Optional[asyncpg.Connection] = None, ) -> InsertOneResult: await self._ensure_table() doc = dict(document) @@ -561,15 +616,20 @@ async def update_one( rows = await self._db.fetch(query, *params, conn=connection) if rows: - return UpdateResult(matched_count=len(rows), modified_count=len(rows), upserted_id=None) + return UpdateResult( + matched_count=len(rows), modified_count=len(rows), upserted_id=None + ) if upsert: base = {k: v for k, v in filt.items() if not isinstance(v, dict)} new_doc = _apply_update(base, update) result = await self.insert_one(new_doc, connection=connection) - return UpdateResult(matched_count=0, modified_count=1, upserted_id=result.inserted_id) + return UpdateResult( + matched_count=0, modified_count=1, upserted_id=result.inserted_id + ) return UpdateResult(matched_count=0, modified_count=0, upserted_id=None) + async def delete_one( self, filt: Dict[str, Any], *, connection: Optional[asyncpg.Connection] = None ) -> DeleteResult: @@ -608,7 +668,9 @@ async def count_documents(self, filt: Optional[Dict[str, Any]] = None) -> int: return 0 return int(rows[0]["count"]) - async def distinct(self, key: str, filt: Optional[Dict[str, Any]] = None) -> List[Any]: + async def distinct( + self, key: str, filt: Optional[Dict[str, Any]] = None + ) -> List[Any]: await self._ensure_table() params: List[Any] = [] where_clause = self._build_where_clause(filt or {}, params) @@ -620,7 +682,10 @@ async def distinct(self, key: str, filt: Optional[Dict[str, Any]] = None) -> Lis return [row["value"] for row in rows if row["value"] is not None] async def delete_many( - self, filt: Optional[Dict[str, Any]], *, connection: Optional[asyncpg.Connection] = None + self, + filt: Optional[Dict[str, Any]], + *, + connection: Optional[asyncpg.Connection] = None, ) -> DeleteResult: await self._ensure_table() params: List[Any] = [] @@ -691,7 +756,9 @@ async def update_many( upserted_id = result.inserted_id modified += 1 - return UpdateResult(matched_count=matched, modified_count=modified, upserted_id=upserted_id) + return UpdateResult( + matched_count=matched, modified_count=modified, upserted_id=upserted_id + ) def aggregate(self, pipeline: List[Dict[str, Any]]) -> PostgresCursor: return PostgresCursor(self, {}, None, None, pipeline=pipeline) @@ -722,7 +789,9 @@ def _value_expression(self, value: Any, params: List[Any]) -> str: params.append(_jsonable(value)) return f"${len(params)}" - def _build_group_id_expr(self, group_spec: Any, params: List[Any]) -> tuple[str, List[str]]: + def _build_group_id_expr( + self, group_spec: Any, params: List[Any] + ) -> tuple[str, List[str]]: if isinstance(group_spec, dict): parts = [] group_by_exprs: List[str] = [] @@ -765,7 +834,9 @@ def _build_aggregate_order_clause(self, sorts: List[tuple[str, int]]) -> str: parts.append(f"{expr} {order}") return ", ".join(parts) - def _apply_project_stage(self, docs: List[Dict[str, Any]], project: Dict[str, Any]) -> List[Dict[str, Any]]: + def _apply_project_stage( + self, docs: List[Dict[str, Any]], project: Dict[str, Any] + ) -> List[Dict[str, Any]]: projected: List[Dict[str, Any]] = [] for doc in docs: new_doc: Dict[str, Any] = {} @@ -775,11 +846,19 @@ def _apply_project_stage(self, docs: List[Dict[str, Any]], project: Dict[str, An if expr == 1: new_doc[key] = doc.get(key) elif isinstance(expr, str) and expr.startswith("$"): - new_doc[key] = _get_by_path(doc, expr[1:]) if "." in expr[1:] else doc.get(expr[1:]) + new_doc[key] = ( + _get_by_path(doc, expr[1:]) + if "." in expr[1:] + else doc.get(expr[1:]) + ) elif isinstance(expr, dict) and "$dateToString" in expr: fmt = expr["$dateToString"]["format"] date_val = _get_by_path(doc, expr["$dateToString"]["date"][1:]) - new_doc[key] = date_val.strftime(fmt) if hasattr(date_val, "strftime") else None + new_doc[key] = ( + date_val.strftime(fmt) + if hasattr(date_val, "strftime") + else None + ) else: new_doc[key] = expr projected.append(new_doc) @@ -820,12 +899,16 @@ async def _run_aggregate_pipeline( combined_sorts.extend(list(sort_stage.items())) combined_sorts.extend(sorts) effective_limit = limit_stage if limit_stage is not None else limit - results = await self._find_docs(match_filter, None, combined_sorts, effective_limit) + results = await self._find_docs( + match_filter, None, combined_sorts, effective_limit + ) if project_stage: results = self._apply_project_stage(results, project_stage) return results - group_id_expr, group_by_exprs = self._build_group_id_expr(group_stage.get("_id"), params) + group_id_expr, group_by_exprs = self._build_group_id_expr( + group_stage.get("_id"), params + ) select_parts = [f"{group_id_expr} AS _id"] for field, expr in group_stage.items(): if field == "_id": @@ -900,7 +983,9 @@ async def connect(self, retries: int = 3, initial_delay: float = 1.0) -> None: while attempt <= retries and not self.is_connected: try: logger.info("Connecting to Postgres at {}", self._dsn) - self.pool = await asyncpg.create_pool(self._dsn, min_size=1, max_size=10) + self.pool = await asyncpg.create_pool( + self._dsn, min_size=1, max_size=10 + ) await self._ensure_extensions() await self._ensure_tables(DEFAULT_COLLECTIONS) await self._ensure_indexes(DEFAULT_COLLECTIONS) @@ -910,7 +995,9 @@ async def connect(self, retries: int = 3, initial_delay: float = 1.0) -> None: return except Exception: attempt += 1 - logger.exception("Failed to connect to Postgres (attempt {})", attempt) + logger.exception( + "Failed to connect to Postgres (attempt {})", attempt + ) self._reset_state() if attempt > retries: break @@ -933,12 +1020,16 @@ async def _monitor_connection(self) -> None: async with self.pool.acquire() as conn: await conn.execute("SELECT 1") except Exception: - logger.warning("Lost connection to Postgres, attempting to reconnect") + logger.warning( + "Lost connection to Postgres, attempting to reconnect" + ) try: if self.pool: await self.pool.close() except Exception: - logger.debug("Error while closing pool during reconnect", exc_info=True) + logger.debug( + "Error while closing pool during reconnect", exc_info=True + ) self._reset_state(preserve_monitor=True) await self.connect() except asyncio.CancelledError: @@ -958,7 +1049,9 @@ async def _ensure_extensions(self) -> None: await self.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm;") self._ensured_extensions.add("pg_trgm") except Exception: - logger.warning("Could not ensure pg_trgm extension; trigram index may not be available") + logger.warning( + "Could not ensure pg_trgm extension; trigram index may not be available" + ) async def _ensure_tables(self, collections: Iterable[str]) -> None: for collection in collections: @@ -1000,7 +1093,9 @@ async def ensure_indexes_for_collection(self, collection: str) -> None: cache_key = f"{table_name}:{index_name}" if cache_key in self._ensured_indexes: continue - await self.execute(f"CREATE INDEX IF NOT EXISTS {index_name} ON {table_name} {expression};") + await self.execute( + f"CREATE INDEX IF NOT EXISTS {index_name} ON {table_name} {expression};" + ) self._ensured_indexes.add(cache_key) async def disconnect(self) -> None: @@ -1034,7 +1129,9 @@ async def transaction(self) -> AsyncIterator[asyncpg.Connection]: else: await tx.commit() - async def execute(self, query: str, *args: Any, conn: Optional[asyncpg.Connection] = None) -> None: + async def execute( + self, query: str, *args: Any, conn: Optional[asyncpg.Connection] = None + ) -> None: if self.pool is None and conn is None: raise RuntimeError("Database not connected") if conn is not None: diff --git a/src/middleware/user_agent_middleware.py b/src/middleware/user_agent_middleware.py index de5e860..39876ae 100644 --- a/src/middleware/user_agent_middleware.py +++ b/src/middleware/user_agent_middleware.py @@ -5,7 +5,6 @@ from starlette.middleware.base import BaseHTTPMiddleware from starlette.requests import Request -from starlette.responses import Response from loguru import logger from .. import config @@ -18,16 +17,16 @@ async def dispatch(self, request: Request, call_next): """Log user agent and process the request.""" # Get user agent user_agent = request.headers.get("user-agent", "unknown") - + # Get client IP client_ip = "unknown" if request.client: client_ip = request.client.host - + # Get request method and path method = request.method path = request.url.path - + if config.REQUEST_LOGGING_ENABLED: # Log the request with user agent logger.info( @@ -35,10 +34,10 @@ async def dispatch(self, request: Request, call_next): f"Client: {client_ip} | " f"User-Agent: {user_agent}" ) - + # Process the request response = await call_next(request) - + # Log response status if it's an error if 400 <= response.status_code < 600: logger.warning( @@ -46,5 +45,5 @@ async def dispatch(self, request: Request, call_next): f"Client: {client_ip} | " f"User-Agent: {user_agent}" ) - + return response diff --git a/src/models/user.py b/src/models/user.py index d17b3fe..577f0f8 100644 --- a/src/models/user.py +++ b/src/models/user.py @@ -44,6 +44,7 @@ def validate_password_hash(cls, v): if not v or not v.strip(): raise ValueError("Password hash cannot be empty") return v + @validator("favorites") def validate_favorites(cls, v): if not isinstance(v, list): diff --git a/src/routes/api/__init__.py b/src/routes/api/__init__.py index 6e3dc40..7004fb7 100644 --- a/src/routes/api/__init__.py +++ b/src/routes/api/__init__.py @@ -5,4 +5,14 @@ from . import logs, stats, images, exports, uploads, pdf, history, favorites -__all__ = ["logs", "stats", "images", "exports", "uploads", "pdf", "history", "admin", "favorites"] +__all__ = [ + "logs", + "stats", + "images", + "exports", + "uploads", + "pdf", + "history", + "admin", + "favorites", +] diff --git a/src/routes/api/admin.py b/src/routes/api/admin.py index 73ca1ce..734b6b0 100644 --- a/src/routes/api/admin.py +++ b/src/routes/api/admin.py @@ -8,6 +8,7 @@ from ...middleware.auth_middleware import AuthMiddleware from ...services.settings_service import SettingsService + router = APIRouter() @@ -64,4 +65,3 @@ async def update_feature_flags(request: Request, csrf_protect: CsrfProtect = Dep url=f"{redirect_url}?status={status}", status_code=303, ) - diff --git a/src/routes/api/exports.py b/src/routes/api/exports.py index 3a7d00c..175a9d9 100644 --- a/src/routes/api/exports.py +++ b/src/routes/api/exports.py @@ -86,7 +86,9 @@ async def download_collections(request: Request): ) from unavailable_error except ValueError as missing_user: logger.warning("Export denied for {}: {}", username, missing_user) - raise HTTPException(status_code=404, detail=EXPORT_NOT_FOUND_MESSAGE) from missing_user + raise HTTPException( + status_code=404, detail=EXPORT_NOT_FOUND_MESSAGE + ) from missing_user headers = {"Content-Disposition": f'attachment; filename="{filename}"'} return StreamingResponse(stream, media_type=EXPORT_MEDIA_TYPE, headers=headers) diff --git a/src/routes/api/favorites.py b/src/routes/api/favorites.py index 8ad6407..4230cce 100644 --- a/src/routes/api/favorites.py +++ b/src/routes/api/favorites.py @@ -2,7 +2,7 @@ API routes for managing user favorites. """ -from typing import Any, Dict, List +from typing import Any, Dict from fastapi import APIRouter, HTTPException, Request from loguru import logger diff --git a/src/routes/api/images.py b/src/routes/api/images.py index d3cb631..a295fa0 100644 --- a/src/routes/api/images.py +++ b/src/routes/api/images.py @@ -5,7 +5,6 @@ from fastapi import APIRouter, Request, HTTPException from fastapi.responses import JSONResponse -import asyncio from ...middleware.auth_middleware import AuthMiddleware from ...services.storage_service import ( @@ -42,7 +41,7 @@ async def delete_image( status_code=403, detail="Access Denied: You must be an admin to perform this action.", ) - + # Security check - prevent directory traversal if "/" in filename or "\\" in filename: raise HTTPException(status_code=400, detail="Invalid filename") @@ -50,7 +49,9 @@ async def delete_image( try: exists = await storage_image_exists(filename) except StorageError as exc: - raise HTTPException(status_code=500, detail="Failed to access image storage") from exc + raise HTTPException( + status_code=500, detail="Failed to access image storage" + ) from exc if not exists: raise HTTPException(status_code=404, detail="Image not found") diff --git a/src/routes/api/logs.py b/src/routes/api/logs.py index 63eb8da..fdc3b26 100644 --- a/src/routes/api/logs.py +++ b/src/routes/api/logs.py @@ -16,8 +16,6 @@ router = APIRouter() - - @router.get("/logs", response_model=Dict[str, Any]) async def get_logs( request: Request, diff --git a/src/routes/api/page_markdown.py b/src/routes/api/page_markdown.py index e3df56d..d8cf3c0 100644 --- a/src/routes/api/page_markdown.py +++ b/src/routes/api/page_markdown.py @@ -4,6 +4,7 @@ router = APIRouter() + @router.get("/page/markdown/{title}") async def get_markdown_page(title: str): """Fetches the markdown content of a page by its title. useful for ai clients.""" diff --git a/src/routes/api/pdf.py b/src/routes/api/pdf.py index a6a9673..c6968df 100644 --- a/src/routes/api/pdf.py +++ b/src/routes/api/pdf.py @@ -192,9 +192,7 @@ async def _collect_linked_pages( "content": processed, } ) - logger.info( - f"Collected page {title} (branch: {page_branch}) for PDF export" - ) + logger.info(f"Collected page {title} (branch: {page_branch}) for PDF export") if len(collected) >= limit: break diff --git a/src/routes/api/uploads.py b/src/routes/api/uploads.py index fb662cd..f201b6a 100644 --- a/src/routes/api/uploads.py +++ b/src/routes/api/uploads.py @@ -3,7 +3,6 @@ Handles file upload operations. """ -import asyncio import hashlib import uuid @@ -199,7 +198,9 @@ def _matches_magic_signature(content_type: str, data: bytes) -> bool: # Store file in object storage try: - stored_image = await upload_image_bytes(file_content, unique_filename, file.content_type) + stored_image = await upload_image_bytes( + file_content, unique_filename, file.content_type + ) except StorageError as exc: logger.error(f"Image upload failed for '{unique_filename}': {exc}") return JSONResponse( diff --git a/src/routes/media.py b/src/routes/media.py index d9faca6..0189e73 100644 --- a/src/routes/media.py +++ b/src/routes/media.py @@ -2,7 +2,6 @@ Media proxy routes for serving uploaded images. """ -import asyncio import mimetypes from fastapi import APIRouter, HTTPException diff --git a/src/routes/web/admin.py b/src/routes/web/admin.py index 7e8d766..0c56940 100644 --- a/src/routes/web/admin.py +++ b/src/routes/web/admin.py @@ -4,7 +4,7 @@ """ from fastapi import APIRouter, Depends, Request, Response -from fastapi.responses import HTMLResponse, RedirectResponse +from fastapi.responses import HTMLResponse from fastapi_csrf_protect import CsrfProtect from ...database import db_instance, get_users_collection diff --git a/src/routes/web/auth.py b/src/routes/web/auth.py index 87799c0..142db1b 100644 --- a/src/routes/web/auth.py +++ b/src/routes/web/auth.py @@ -17,19 +17,20 @@ from ...utils.template_env import get_templates from ...utils.validation import sanitize_redirect_path + def get_client_ip(request: Request) -> str: """ Safely retrieve the client IP address from a FastAPI request. Handles various scenarios like proxies, VPNs, and privacy tools. - + Args: request: The FastAPI request object - + Returns: str: The client IP address or "unknown" if it cannot be determined """ xff = request.headers.get("x-forwarded-for") - + try: if xff: # x-forwarded-for can contain a list of IPs, take the first @@ -39,9 +40,10 @@ def get_client_ip(request: Request) -> str: except Exception: # In case accessing attributes raises for some ASGI servers, keep "unknown" pass - + return "unknown" + router = APIRouter() templates = get_templates() diff --git a/src/routes/web/branches.py b/src/routes/web/branches.py index d333d10..4b2aa26 100644 --- a/src/routes/web/branches.py +++ b/src/routes/web/branches.py @@ -149,7 +149,11 @@ async def set_branch( safe_branch = branch if is_safe_branch_parameter(branch) else "main" referer_header = request.headers.get("referer") # Simple validation - only allow relative URLs - safe_referer = referer_header if referer_header and not urlparse(referer_header).scheme else "/" + safe_referer = ( + referer_header + if referer_header and not urlparse(referer_header).scheme + else "/" + ) parsed = urlparse(safe_referer) if parsed.scheme or parsed.netloc: diff --git a/src/routes/web/pages.py b/src/routes/web/pages.py index cd293d6..e37f736 100644 --- a/src/routes/web/pages.py +++ b/src/routes/web/pages.py @@ -754,7 +754,9 @@ async def save_page( normalized_previous = sorted( {username for username in previous_allowed_users} ) - normalized_current = sorted({username for username in allowed_users_list}) + normalized_current = sorted( + {username for username in allowed_users_list} + ) permission_changed = edit_permission != previous_permission allowed_users_changed = normalized_previous != normalized_current became_protected = edit_permission != EDIT_PERMISSION_EVERYBODY @@ -958,7 +960,7 @@ async def delete_branch( if not parsed.netloc and not parsed.scheme: # relative path, safe to redirect return RedirectResponse(url=redirect_url, status_code=303) - # fallback to home page if unsafe + # fallback to home page if unsafe return RedirectResponse(url="/", status_code=303) else: logger.warning(f"Branch not found for deletion: {branch} from page {title}") diff --git a/src/server.py b/src/server.py index 7aec99e..dbbfa95 100644 --- a/src/server.py +++ b/src/server.py @@ -11,10 +11,11 @@ from fastapi_csrf_protect import CsrfProtect from pydantic_settings import BaseSettings from loguru import logger -from .config import NAME, APP_DESCRIPTION, STATIC_DIR, DEV, HELP_STATIC_DIR +from .config import NAME, APP_DESCRIPTION, STATIC_DIR, DEV from .database import init_database from .routes import media -#Why do we do it like this? Because otherwise we import a route that has both web and api routes, casuing circular imports + +# Why do we do it like this? Because otherwise we import a route that has both web and api routes, casuing circular imports from .routes.web import ( pages, search, diff --git a/src/services/analytics_service.py b/src/services/analytics_service.py index 9f73b61..d9d4e0b 100644 --- a/src/services/analytics_service.py +++ b/src/services/analytics_service.py @@ -39,7 +39,10 @@ def _get_collection(): async def _maybe_enforce_retention(cls, collection) -> None: """Delete analytics events older than the retention window on a fixed cadence.""" now = _utcnow() - if cls._last_retention_check and now - cls._last_retention_check < cls._RETENTION_CHECK_INTERVAL: + if ( + cls._last_retention_check + and now - cls._last_retention_check < cls._RETENTION_CHECK_INTERVAL + ): return cutoff = now - timedelta(days=cls._RETENTION_DAYS) try: @@ -159,9 +162,7 @@ async def get_admin_dashboard_metrics() -> Dict[str, Any]: return AnalyticsService._empty_metrics() now = _utcnow() - today_start = datetime( - now.year, now.month, now.day, tzinfo=timezone.utc - ) + today_start = datetime(now.year, now.month, now.day, tzinfo=timezone.utc) window_start = today_start - timedelta(days=6) metrics = AnalyticsService._empty_metrics() diff --git a/src/services/branch_service.py b/src/services/branch_service.py index de8cb67..6d5630b 100644 --- a/src/services/branch_service.py +++ b/src/services/branch_service.py @@ -73,7 +73,9 @@ async def get_branches_for_page(title: str) -> List[str]: branches.update(branch_docs or []) if pages_collection is not None: - page_branches = await pages_collection.distinct("branch", {"title": title}) + page_branches = await pages_collection.distinct( + "branch", {"title": title} + ) branches.update([b for b in page_branches or [] if b]) return list(branches) diff --git a/src/services/page_service.py b/src/services/page_service.py index 9889e3d..2dd55c3 100644 --- a/src/services/page_service.py +++ b/src/services/page_service.py @@ -6,7 +6,6 @@ from typing import Optional, List, Dict, Any from datetime import datetime, timezone from loguru import logger -from pymongo.errors import OperationFailure from ..database import ( TABLE_PREFIX, get_pages_collection, @@ -403,7 +402,9 @@ async def search_pages( ) pages = [PageService._normalize_timestamp(p) for p in pages] - logger.info(f"Search performed: {query!r} on branch {branch!r} - found {len(pages)} results") + logger.info( + f"Search performed: {query!r} on branch {branch!r} - found {len(pages)} results" + ) return pages except Exception as e: @@ -412,7 +413,6 @@ async def search_pages( ) return [] - @staticmethod async def delete_page(title: str) -> bool: """ @@ -616,7 +616,5 @@ async def rename_page(old_title: str, new_title: str) -> tuple[bool, str | None] logger.info(f"Page renamed from {old_title} to {new_title}") return True, None except Exception as e: - logger.error( - f"Error renaming page {old_title} to {new_title}: {str(e)}" - ) + logger.error(f"Error renaming page {old_title} to {new_title}: {str(e)}") return False, "error" diff --git a/src/services/settings_service.py b/src/services/settings_service.py index 947b289..99a00b1 100644 --- a/src/services/settings_service.py +++ b/src/services/settings_service.py @@ -99,7 +99,7 @@ def _parse_expires_at(value: Optional[object]) -> Optional[datetime]: except ValueError: logger.warning(f"Could not parse banner expiration '{value}'") return None - + if parsed_dt: if parsed_dt.tzinfo is None: return parsed_dt.replace(tzinfo=timezone.utc) @@ -247,7 +247,9 @@ async def get_feature_flags(cls, *, force_refresh: bool = False) -> FeatureFlags if not db_instance.is_connected: return cls._feature_flags_cache - if not force_refresh and cls._cache_is_fresh(cls._feature_flags_cache_fetched_at): + if not force_refresh and cls._cache_is_fresh( + cls._feature_flags_cache_fetched_at + ): return cls._feature_flags_cache settings_collection = db_instance.get_collection("settings") diff --git a/src/services/storage_service.py b/src/services/storage_service.py index f859389..67428bc 100644 --- a/src/services/storage_service.py +++ b/src/services/storage_service.py @@ -17,12 +17,9 @@ from botocore.config import Config from botocore.exceptions import BotoCoreError, ClientError from loguru import logger -from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase from ..config import ( ALLOWED_IMAGE_TYPES, - MONGODB_URL, - MONGODB_DB_NAME, S3_ACCESS_KEY, S3_BUCKET, S3_ENDPOINT, @@ -35,6 +32,7 @@ IMAGE_PREFIX = "uploads/" IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff"} + def _safe_local_image_path(filename: str) -> Path: """ Validate and return a safe local path for an image filename. @@ -101,9 +99,7 @@ async def _reset_s3_client() -> None: async def _handle_client_attribute_error(operation: str, exc: AttributeError) -> None: """Translate attribute errors into storage errors and reset the client.""" await _reset_s3_client() - logger.exception( - f"S3 client missing expected attribute during {operation}: {exc}" - ) + logger.exception(f"S3 client missing expected attribute during {operation}: {exc}") raise StorageError("Object storage client is not usable.") from exc @@ -115,33 +111,31 @@ async def _get_s3_client(): return _S3_CLIENT logger.warning("Cached S3 client missing methods; resetting cached client.") await _reset_s3_client() - + if not _s3_enabled(): raise StorageError("S3 client requested but storage is not configured.") - + async with _S3_LOCK: if _S3_CLIENT is not None: if hasattr(_S3_CLIENT, "get_object"): return _S3_CLIENT logger.warning("Cached S3 client missing methods after lock; resetting.") await _reset_s3_client() - + # Configure S3 client with connection pooling config_kwargs: Dict[str, Any] = { "signature_version": "s3v4", "max_pool_connections": 50, # Configure connection pool - "retries": { - "max_attempts": 3, - "mode": "adaptive" - } + "retries": {"max_attempts": 3, "mode": "adaptive"}, } - + if S3_FORCE_PATH_STYLE: config_kwargs["s3"] = {"addressing_style": "path"} - + from aiobotocore.session import get_session + session = get_session() - + stack = AsyncExitStack() try: client_cm = session.create_client( @@ -155,12 +149,12 @@ async def _get_s3_client(): except Exception: await stack.aclose() raise - + if not hasattr(client, "get_object"): await stack.aclose() logger.error("Created S3 client missing required methods; aborting.") raise StorageError("Failed to initialise object storage client.") - + _S3_CLIENT = client _S3_EXIT_STACK = stack return _S3_CLIENT @@ -198,7 +192,7 @@ async def upload_image_bytes( extra_args: Dict[str, Any] = {} if content_type and content_type in ALLOWED_IMAGE_TYPES: extra_args["ContentType"] = content_type - + try: await client.put_object( Bucket=S3_BUCKET, @@ -240,7 +234,7 @@ async def list_images() -> List[Dict[str, Any]]: if _s3_enabled(): client = await _get_s3_client() items: List[Dict[str, Any]] = [] - + try: paginator = client.get_paginator("list_objects_v2") async for page in paginator.paginate(Bucket=S3_BUCKET, Prefix=IMAGE_PREFIX): @@ -252,7 +246,9 @@ async def list_images() -> List[Dict[str, Any]]: filename = key.split("/")[-1] last_modified = obj.get("LastModified") modified_ts = ( - int(last_modified.timestamp()) if last_modified else int(time.time()) + int(last_modified.timestamp()) + if last_modified + else int(time.time()) ) items.append( { @@ -267,14 +263,14 @@ async def list_images() -> List[Dict[str, Any]]: except (BotoCoreError, ClientError) as exc: logger.exception(f"Failed to list images from S3: {exc}") raise StorageError("Listing images from object storage failed.") from exc - + items.sort(key=lambda item: item["modified"], reverse=True) return items upload_path = Path(UPLOAD_DIR) if not upload_path.exists(): return [] - + items: List[Dict[str, Any]] = [] for entry in upload_path.iterdir(): if entry.is_file() and entry.suffix.lower() in IMAGE_EXTENSIONS: @@ -301,7 +297,9 @@ async def download_image_bytes(filename: str) -> bytes: if _s3_enabled(): client = await _get_s3_client() try: - response = await client.get_object(Bucket=S3_BUCKET, Key=_object_key(filename)) + response = await client.get_object( + Bucket=S3_BUCKET, Key=_object_key(filename) + ) body = response.get("Body") if body is None: raise StorageError("Empty response body from object storage.") @@ -316,7 +314,9 @@ async def download_image_bytes(filename: str) -> bytes: f"Image '{filename}' missing in S3 (code={error_code or 'unknown'}); checking local storage." ) else: - logger.exception(f"Failed to download image '{filename}' from S3: {exc}") + logger.exception( + f"Failed to download image '{filename}' from S3: {exc}" + ) raise StorageError("Download from object storage failed.") from exc except BotoCoreError as exc: logger.exception(f"Failed to download image '{filename}' from S3: {exc}") @@ -334,7 +334,7 @@ async def download_image_bytes(filename: str) -> bytes: else: logger.warning(f"Image '{filename}' not found in local storage.") raise StorageError("Image not found.") - + try: async with aiofiles.open(file_path, "rb") as file_obj: return await file_obj.read() diff --git a/src/services/user_service.py b/src/services/user_service.py index aef89fe..53ccd91 100644 --- a/src/services/user_service.py +++ b/src/services/user_service.py @@ -148,9 +148,7 @@ async def list_favorites(username: str) -> Optional[List[Dict[str, str]]]: return None @staticmethod - async def add_favorite( - username: str, title: str, branch: str = "main" - ) -> bool: + async def add_favorite(username: str, title: str, branch: str = "main") -> bool: """ Add a favorite page for the specified user. @@ -168,8 +166,7 @@ async def add_favorite( return False if any( - fav["title"] == title and fav["branch"] == branch - for fav in favorites + fav["title"] == title and fav["branch"] == branch for fav in favorites ): return True @@ -207,9 +204,7 @@ async def add_favorite( return False @staticmethod - async def remove_favorite( - username: str, title: str, branch: str = "main" - ) -> bool: + async def remove_favorite(username: str, title: str, branch: str = "main") -> bool: """ Remove a favorite page for the specified user. @@ -227,8 +222,7 @@ async def remove_favorite( return False if not any( - fav["title"] == title and fav["branch"] == branch - for fav in favorites + fav["title"] == title and fav["branch"] == branch for fav in favorites ): return True diff --git a/src/stats.py b/src/stats.py index dd716ee..428aa92 100644 --- a/src/stats.py +++ b/src/stats.py @@ -1,6 +1,5 @@ import asyncio from datetime import datetime, timedelta -import os from loguru import logger from .database import ( get_pages_collection, @@ -20,6 +19,7 @@ last_image_count_time = None # Start as None to force first calculation image_count_cache_duration = timedelta(minutes=30) # Cache for 30 Minutes + async def get_total_edits(): """ Count total number of edits in the wiki (history collection entries). @@ -159,22 +159,23 @@ async def get_total_images(): logger.info( f"IMG: Cache is old or uninitialized, Updating! Time delta is {time_delta}" ) - try: + try: image_hashes_collection = get_image_hashes_collection() # only use the db since each image is added to the db when uploaded if db_instance.is_connected and image_hashes_collection is not None: total_images = await image_hashes_collection.count_documents({}) - + # CORRECTED CACHE UPDATE last_image_count = total_images last_image_count_time = datetime.now() - + return total_images return 0 except Exception as e: logger.error(f"Error counting total images: {str(e)}") return 0 + async def get_stats(): """ Get all statistics for the wiki. diff --git a/src/tests/test.py b/src/tests/test.py index 4ec85f5..6cb95fc 100644 --- a/src/tests/test.py +++ b/src/tests/test.py @@ -4,38 +4,39 @@ import time from urllib.parse import quote -BASE_URL = 'http://127.0.0.1:8000' -DEFAULT_PASSWORD = 'password123' +BASE_URL = "http://127.0.0.1:8000" +DEFAULT_PASSWORD = "password123" def fetch_csrf_token(session: requests.Session, path: str) -> str: """Fetch CSRF token from form pages and ensure cookie is stored on the session.""" - response = session.get(f'{BASE_URL}{path}', allow_redirects=True) + response = session.get(f"{BASE_URL}{path}", allow_redirects=True) response.raise_for_status() match = re.search(r'name="csrf_token" value="([^"]+)"', response.text) - assert match, f'CSRF token not found on {path}' + assert match, f"CSRF token not found on {path}" return match.group(1) def build_page_url(title: str, branch: str = "main") -> str: """Construct a properly encoded URL for viewing a page.""" encoded_title = quote(title, safe="") - base_url = f'{BASE_URL}/page/{encoded_title}' + base_url = f"{BASE_URL}/page/{encoded_title}" if branch and branch != "main": return f'{base_url}?branch={quote(branch, safe="")}' return base_url -@pytest.fixture(scope='module') +@pytest.fixture(scope="module") def user_credentials(): """Provide shared credentials for authentication tests.""" timestamp = str(int(time.time())) return { - 'username': f'test_user_{timestamp}', - 'password': DEFAULT_PASSWORD, + "username": f"test_user_{timestamp}", + "password": DEFAULT_PASSWORD, } -@pytest.fixture(scope='module') + +@pytest.fixture(scope="module") def test_client(): """ Test client fixture that connects to the running dev server. @@ -43,72 +44,92 @@ def test_client(): """ # Create a session that will persist cookies across requests session = requests.Session() - + # Set default timeout for requests session.timeout = 30 - + yield session + def test_server_available(test_client): """ Test to verify the server is available before running other tests. """ - # Wait for a moment to not hit the server real hard (race conditions?) + # Wait for a moment to not hit the server real hard (race conditions?) try: - response = test_client.get(f'{BASE_URL}/') - assert response.status_code in [200, 301, 302], f"Server not available: {response.status_code}" + response = test_client.get(f"{BASE_URL}/") + assert response.status_code in [ + 200, + 301, + 302, + ], f"Server not available: {response.status_code}" except requests.exceptions.ConnectionError: - pytest.fail("Server not available at http://127.0.0.1:8000. Make sure the dev server is running.") + pytest.fail( + "Server not available at http://127.0.0.1:8000. Make sure the dev server is running." + ) + def test_account_creation(test_client, user_credentials): """ Test account creation with a unique username """ # Fetch CSRF token before submitting registration form - # Wait for a moment to not hit the server real hard (race conditions?) + # Wait for a moment to not hit the server real hard (race conditions?) csrf_token = fetch_csrf_token(test_client, "/register") - + # Try registration (this is for web form, not API) - response = test_client.post(f'{BASE_URL}/register', data={ - 'username': user_credentials['username'], - 'password': user_credentials['password'], - 'confirm_password': user_credentials['password'], - 'csrf_token': csrf_token, - }, allow_redirects=True) - + response = test_client.post( + f"{BASE_URL}/register", + data={ + "username": user_credentials["username"], + "password": user_credentials["password"], + "confirm_password": user_credentials["password"], + "csrf_token": csrf_token, + }, + allow_redirects=True, + ) + # Check if registration was successful (redirect to home page) assert response.status_code == 200, f"Registration failed: {response.status_code}" - + # Check if we were redirected or if the page contains expected content content = response.text.lower() - assert "home" in content or "welcome" in content, "Expected home page after registration" + assert ( + "home" in content or "welcome" in content + ), "Expected home page after registration" + def test_login(test_client, user_credentials): """ Test login with the created user """ # Attempt login using the credentials created during registration - # Wait for a moment to not hit the server real hard (race conditions?) + # Wait for a moment to not hit the server real hard (race conditions?) test_client.cookies.clear() # Ensure a clean session for the login test csrf_token = fetch_csrf_token(test_client, "/login") # Try login - response = test_client.post(f'{BASE_URL}/login', data={ - 'username': user_credentials['username'], - 'password': user_credentials['password'], - 'csrf_token': csrf_token, - 'next': '/', - }, allow_redirects=True) - + response = test_client.post( + f"{BASE_URL}/login", + data={ + "username": user_credentials["username"], + "password": user_credentials["password"], + "csrf_token": csrf_token, + "next": "/", + }, + allow_redirects=True, + ) + # Check if login was successful assert response.status_code == 200, f"Login failed: {response.status_code}" - + # Check if we were redirected to home or see expected content content = response.text.lower() assert "home" in content, "Expected home page after login" - + # Check if session cookie was set - assert 'user_session' in test_client.cookies, "Session cookie not set after login" + assert "user_session" in test_client.cookies, "Session cookie not set after login" + def test_page_creation(test_client, user_credentials): """ @@ -119,104 +140,142 @@ def test_page_creation(test_client, user_credentials): # First, try to login test_client.cookies.clear() csrf_token = fetch_csrf_token(test_client, "/login") - login_response = test_client.post(f'{BASE_URL}/login', data={ - 'username': user_credentials['username'], - 'password': user_credentials['password'], - 'csrf_token': csrf_token, - 'next': '/', - }, allow_redirects=False) - + login_response = test_client.post( + f"{BASE_URL}/login", + data={ + "username": user_credentials["username"], + "password": user_credentials["password"], + "csrf_token": csrf_token, + "next": "/", + }, + allow_redirects=False, + ) + # If login gives a redirect, follow it if login_response.status_code in [301, 302, 303]: - redirect_location = login_response.headers.get('Location') + redirect_location = login_response.headers.get("Location") if redirect_location: - redirect_url = f'{BASE_URL}{redirect_location}' if redirect_location.startswith('/') else redirect_location + redirect_url = ( + f"{BASE_URL}{redirect_location}" + if redirect_location.startswith("/") + else redirect_location + ) test_client.get(redirect_url, allow_redirects=False) - + # Ensure we have a session cookie - assert 'user_session' in test_client.cookies, "Failed to login" - + assert "user_session" in test_client.cookies, "Failed to login" + # Now try to create a page page_title = f"Test Page {int(time.time())}" edit_csrf_token = fetch_csrf_token(test_client, f"/edit/{page_title}") - response = test_client.post(f'{BASE_URL}/edit/{page_title}', data={ - 'content': '# Test Page Content\n\nThis is a test page created by automated tests.', - 'edit_summary': 'Test page creation', - 'edit_permission': 'everybody', - 'allowed_users': '', - 'csrf_token': edit_csrf_token, - }, allow_redirects=False) - + response = test_client.post( + f"{BASE_URL}/edit/{page_title}", + data={ + "content": "# Test Page Content\n\nThis is a test page created by automated tests.", + "edit_summary": "Test page creation", + "edit_permission": "everybody", + "allowed_users": "", + "csrf_token": edit_csrf_token, + }, + allow_redirects=False, + ) + # Should redirect to the page after creation - assert response.status_code in [301, 302, 303], f"Page creation failed: {response.status_code}" - + assert response.status_code in [ + 301, + 302, + 303, + ], f"Page creation failed: {response.status_code}" + # Get the redirected URL - page_url = response.headers.get('Location') + page_url = response.headers.get("Location") assert page_url, "Page creation response missing redirect location" - full_page_url = f'{BASE_URL}{page_url}' if page_url.startswith('/') else page_url - - + full_page_url = f"{BASE_URL}{page_url}" if page_url.startswith("/") else page_url + # Verify the page was created by visiting it page_response = test_client.get(full_page_url) - assert page_response.status_code == 200, f"Page view failed: {page_response.status_code}" - + assert ( + page_response.status_code == 200 + ), f"Page view failed: {page_response.status_code}" + # Check if our content is in the page content = page_response.text.lower() - assert 'test page content' in content, "Page content not found" + assert "test page content" in content, "Page content not found" + def test_page_view(test_client, user_credentials): """ Test viewing an existing page """ - # Wait to avoid hammering the server + # Wait to avoid hammering the server page_title = f"View Test Page {int(time.time())}" # Ensure clean authentication state test_client.cookies.clear() csrf_token = fetch_csrf_token(test_client, "/login") - login_response = test_client.post(f'{BASE_URL}/login', data={ - 'username': user_credentials['username'], - 'password': user_credentials['password'], - 'csrf_token': csrf_token, - 'next': '/', - }, allow_redirects=False) + login_response = test_client.post( + f"{BASE_URL}/login", + data={ + "username": user_credentials["username"], + "password": user_credentials["password"], + "csrf_token": csrf_token, + "next": "/", + }, + allow_redirects=False, + ) if login_response.status_code in [301, 302, 303]: - redirect_location = login_response.headers.get('Location') + redirect_location = login_response.headers.get("Location") if redirect_location: - redirect_url = f'{BASE_URL}{redirect_location}' if redirect_location.startswith('/') else redirect_location + redirect_url = ( + f"{BASE_URL}{redirect_location}" + if redirect_location.startswith("/") + else redirect_location + ) test_client.get(redirect_url, allow_redirects=False) - if 'user_session' not in test_client.cookies: + if "user_session" not in test_client.cookies: pytest.fail("Could not login to create test page") # Create the page that we are going to view edit_csrf_token = fetch_csrf_token(test_client, f"/edit/{page_title}") - create_response = test_client.post(f'{BASE_URL}/edit/{page_title}', data={ - 'content': 'View Page Content\n\nThis page is for testing the view functionality.', - 'edit_summary': 'Create test view page', - 'edit_permission': 'everybody', - 'allowed_users': '', - 'csrf_token': edit_csrf_token, - }, allow_redirects=False) + create_response = test_client.post( + f"{BASE_URL}/edit/{page_title}", + data={ + "content": "View Page Content\n\nThis page is for testing the view functionality.", + "edit_summary": "Create test view page", + "edit_permission": "everybody", + "allowed_users": "", + "csrf_token": edit_csrf_token, + }, + allow_redirects=False, + ) - assert create_response.status_code in [301, 302, 303], f"Failed to create page for view test: {create_response.status_code}" + assert create_response.status_code in [ + 301, + 302, + 303, + ], f"Failed to create page for view test: {create_response.status_code}" - redirect_location = create_response.headers.get('Location') + redirect_location = create_response.headers.get("Location") if redirect_location: - page_url_to_check = f'{BASE_URL}{redirect_location}' if redirect_location.startswith('/') else redirect_location + page_url_to_check = ( + f"{BASE_URL}{redirect_location}" + if redirect_location.startswith("/") + else redirect_location + ) else: page_url_to_check = build_page_url(page_title) - - # Now test viewing the page page_response = test_client.get(page_url_to_check) - assert page_response.status_code == 200, f"Page view failed: {page_response.status_code}" + assert ( + page_response.status_code == 200 + ), f"Page view failed: {page_response.status_code}" # Check if page contains the expected content content = page_response.text.lower() - assert 'view page content' in content, "Page content not found in view" + assert "view page content" in content, "Page content not found in view" def test_favorites_flow(test_client, user_credentials): @@ -263,9 +322,11 @@ def test_favorites_flow(test_client, user_credentials): allow_redirects=False, ) - assert create_response.status_code in [301, 302, 303], ( - f"Failed to create page for favorites test: {create_response.status_code}" - ) + assert create_response.status_code in [ + 301, + 302, + 303, + ], f"Failed to create page for favorites test: {create_response.status_code}" redirect_location = create_response.headers.get("Location") if redirect_location: @@ -281,47 +342,63 @@ def test_favorites_flow(test_client, user_credentials): favorite_entry = {"title": page_title, "branch": "main"} favorites_response = test_client.get(favorites_url, headers=headers) - assert favorites_response.status_code == 200, ( - f"Initial favorites fetch failed: {favorites_response.status_code}" - ) + assert ( + favorites_response.status_code == 200 + ), f"Initial favorites fetch failed: {favorites_response.status_code}" favorites_data = favorites_response.json() - assert isinstance(favorites_data.get("favorites"), list), "Favorites response malformed" - assert favorite_entry not in favorites_data["favorites"], "Page already favorited unexpectedly" + assert isinstance( + favorites_data.get("favorites"), list + ), "Favorites response malformed" + assert ( + favorite_entry not in favorites_data["favorites"] + ), "Page already favorited unexpectedly" encoded_title = quote(page_title, safe="") add_response = test_client.post( f"{BASE_URL}/api/favorites/{encoded_title}", headers=headers ) - assert add_response.status_code == 200, f"Adding favorite failed: {add_response.status_code}" + assert ( + add_response.status_code == 200 + ), f"Adding favorite failed: {add_response.status_code}" add_data = add_response.json() - assert add_data.get("status") == "favorited", "Favorite status not returned as favorited" - assert favorite_entry in add_data.get("favorites", []), "Favorite not present after addition" + assert ( + add_data.get("status") == "favorited" + ), "Favorite status not returned as favorited" + assert favorite_entry in add_data.get( + "favorites", [] + ), "Favorite not present after addition" confirm_response = test_client.get(favorites_url, headers=headers) - assert confirm_response.status_code == 200, ( - f"Favorites check after add failed: {confirm_response.status_code}" - ) + assert ( + confirm_response.status_code == 200 + ), f"Favorites check after add failed: {confirm_response.status_code}" confirm_data = confirm_response.json() - assert favorite_entry in confirm_data.get("favorites", []), "Favorite missing in subsequent check" + assert favorite_entry in confirm_data.get( + "favorites", [] + ), "Favorite missing in subsequent check" delete_response = test_client.delete( f"{BASE_URL}/api/favorites/{encoded_title}", headers=headers ) - assert delete_response.status_code == 200, ( - f"Removing favorite failed: {delete_response.status_code}" - ) + assert ( + delete_response.status_code == 200 + ), f"Removing favorite failed: {delete_response.status_code}" delete_data = delete_response.json() - assert delete_data.get("status") == "unfavorited", "Favorite status not returned as unfavorited" - assert favorite_entry not in delete_data.get("favorites", []), "Favorite still present after removal" + assert ( + delete_data.get("status") == "unfavorited" + ), "Favorite status not returned as unfavorited" + assert favorite_entry not in delete_data.get( + "favorites", [] + ), "Favorite still present after removal" final_response = test_client.get(favorites_url, headers=headers) - assert final_response.status_code == 200, ( - f"Final favorites fetch failed: {final_response.status_code}" - ) + assert ( + final_response.status_code == 200 + ), f"Final favorites fetch failed: {final_response.status_code}" final_data = final_response.json() - assert favorite_entry not in final_data.get("favorites", []), ( - "Favorite entry still present after deletion" - ) + assert favorite_entry not in final_data.get( + "favorites", [] + ), "Favorite entry still present after deletion" # def test_login_rate_limit(test_client, user_credentials): @@ -338,10 +415,10 @@ def test_favorites_flow(test_client, user_credentials): # 'next': '/', # }) # assert "Too many login attempts" in response.text -# +# # # Wait for rate limit to reset (simulated) # time.sleep(60) -# +# # # Try again with correct password # csrf_token = fetch_csrf_token(test_client, "/login") # response = test_client.post(f'{BASE_URL}/login', data={ diff --git a/src/utils/logs.py b/src/utils/logs.py index 38355e8..d5a9893 100644 --- a/src/utils/logs.py +++ b/src/utils/logs.py @@ -84,9 +84,7 @@ async def get_paginated_logs( current_page = 1 else: effective_limit = sanitized_limit - total_pages = max( - 1, (total_items + effective_limit - 1) // effective_limit - ) + total_pages = max(1, (total_items + effective_limit - 1) // effective_limit) if page > total_pages: return { "items": [], @@ -101,7 +99,9 @@ async def get_paginated_logs( data_params: List[Any] = [] data_query = f"SELECT * FROM ({union_sql}) AS combined ORDER BY timestamp DESC" if not bypass: - data_query += f" OFFSET ${len(data_params) + 1} LIMIT ${len(data_params) + 2}" + data_query += ( + f" OFFSET ${len(data_params) + 1} LIMIT ${len(data_params) + 2}" + ) data_params.extend([offset, effective_limit]) rows = await db_instance.fetch(data_query, *data_params) diff --git a/src/utils/navigation_history.py b/src/utils/navigation_history.py index 6ea13eb..a549f6c 100644 --- a/src/utils/navigation_history.py +++ b/src/utils/navigation_history.py @@ -42,7 +42,9 @@ def load_history_cookie(request: Request) -> List[Dict[str, object]]: title = entry.get("title") if not isinstance(title, str) or not title: continue - branch_value = entry.get("branch") if isinstance(entry.get("branch"), str) else "main" + branch_value = ( + entry.get("branch") if isinstance(entry.get("branch"), str) else "main" + ) history.append( { "title": title, @@ -71,7 +73,7 @@ def apply_history_update( updated_history = [entry for entry in updated_history if entry != current_entry] updated_history.append(current_entry) if len(updated_history) > HISTORY_MAX_LENGTH: - updated_history = updated_history[-HISTORY_MAX_LENGTH :] + updated_history = updated_history[-HISTORY_MAX_LENGTH:] return updated_history @@ -84,10 +86,9 @@ def resolve_previous_entry( for index in range(len(history) - 2, -1, -1): entry = history[index] - if ( - entry.get("title") == current_entry.get("title") - and entry.get("branch") == current_entry.get("branch") - ): + if entry.get("title") == current_entry.get("title") and entry.get( + "branch" + ) == current_entry.get("branch"): continue return entry return None @@ -122,7 +123,9 @@ def prepare_navigation_context( """Return updated history data and previous page context for templates.""" history = load_history_cookie(request) current_entry = build_history_entry(title, branch, is_home) - is_back_navigation = request.query_params.get(HISTORY_QUERY_PARAM) == HISTORY_BACK_VALUE + is_back_navigation = ( + request.query_params.get(HISTORY_QUERY_PARAM) == HISTORY_BACK_VALUE + ) updated_history = apply_history_update(history, current_entry, is_back_navigation) previous_entry = resolve_previous_entry(updated_history, current_entry) diff --git a/src/utils/validation.py b/src/utils/validation.py index 43c6d54..5b16c17 100644 --- a/src/utils/validation.py +++ b/src/utils/validation.py @@ -72,8 +72,6 @@ def sanitize_redirect_path(target: Optional[str], default: str = "/") -> str: return f"{path}{query}" - - def sanitize_filename(filename: str) -> str: """Sanitize filename to prevent security issues.""" # Remove path separators and other dangerous characters