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/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/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..ba2538a --- /dev/null +++ b/scripts/mongo_to_postgres_migration.py @@ -0,0 +1,265 @@ +#!/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 + using one table per collection (tables prefixed with wikiware_). +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 json +import os +import subprocess +from pathlib import Path +from typing import Dict, Iterable + +import asyncpg +from dotenv import load_dotenv +from loguru import logger +from pymongo import MongoClient +from bson import ObjectId +from decimal import Decimal + +TABLE_PREFIX = "wikiware_" +COLLECTIONS = [ + "pages", + "history", + "branches", + "users", + "sessions", + "image_hashes", + "analytics_events", + "settings", + "system_logs", + "edit_sessions", +] + + +def _timestamp() -> str: + 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: + 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" + ), + } + + +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} ( + id TEXT NOT NULL, + doc JSONB NOT NULL, + PRIMARY KEY (id) + ); + """ + ) + + +def fetch_collection(mongo_db, name: str) -> Iterable[Dict]: + logger.info("Fetching Mongo collection '{}'", name) + 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: + table_name = table_name_for(collection) + 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( + f""" + INSERT INTO {table_name} (id, doc) + VALUES ($1, $2::jsonb) + ON CONFLICT (id) DO UPDATE SET doc = EXCLUDED.doc + """ + ) + await stmt.executemany(records_to_insert) + + return len(records_to_insert) + + +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) + 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") + 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) + 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/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 e20016f..fd6a1d2 100644 --- a/src/database.py +++ b/src/database.py @@ -1,235 +1,1169 @@ -"""Database connection and management for MongoDB.""" +""" +Postgres-backed database layer for WikiWare. + +This module replaces the previous MongoDB driver with a lightweight JSONB +storage model in Postgres while preserving the collection-style API used by +the rest of the codebase. +""" + +from __future__ import annotations import asyncio -import inspect import os -import time -from typing import Any, Dict, List, Optional +import uuid +import json +import datetime as dt +from contextlib import asynccontextmanager +from decimal import Decimal +from dataclasses import dataclass +from typing import ( + Any, + AsyncIterator, + Dict, + Iterable, + List, + Optional, +) +import asyncpg from dotenv import load_dotenv from loguru import logger -from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCollection, AsyncIOMotorDatabase -from pymongo.errors import ServerSelectionTimeoutError -from . import config load_dotenv() -MONGODB_URL = os.getenv("MONGODB_URL", "mongodb://localhost:27017") -MONGODB_DB_NAME = os.getenv("MONGODB_DB_NAME", "wikiware") -MONGODB_MAX_CONNECTIONS = int(os.getenv("MONGODB_MAX_CONNECTIONS", "100")) -MONGODB_MIN_CONNECTIONS = int(os.getenv("MONGODB_MIN_CONNECTIONS", "10")) -MONGODB_MAX_IDLE_TIME_MS = int(os.getenv("MONGODB_MAX_IDLE_TIME_MS", "30000")) -MONGODB_SERVER_SELECTION_TIMEOUT_MS = int(os.getenv("MONGODB_SERVER_SELECTION_TIMEOUT_MS", "5000")) -MONGODB_SOCKET_TIMEOUT_MS = int(os.getenv("MONGODB_SOCKET_TIMEOUT_MS", "30000")) -DB_OPERATION_LOG_THRESHOLD_MS = float(os.getenv("DB_OPERATION_LOG_THRESHOLD_MS", "100")) / 1000 # Convert to seconds - -IndexKey = str -IndexSpec = tuple[IndexKey, Dict[str, Any]] - -INDEX_CONFIGS: Dict[str, List[IndexSpec]] = { - "users": [ - ("username", {"unique": True}), - ("created_at", {}), +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_PREFIX = "wikiware_" +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)"), + ( + "content_fts", + "USING gin (to_tsvector('english', coalesce(doc->>'title', '') || ' ' || coalesce(doc->>'content', '')))", + ), ], - "sessions": [ - ("session_id", {"unique": True}), - ("user_id", {}), - ("expires_at", {"expireAfterSeconds": 0}), + "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'))"), + ("branch_name", "((doc->>'branch_name'))"), + ], + "users": [ + ("username", "((doc->>'username'))"), + ("email", "((doc->>'email'))"), ], "image_hashes": [ - ("filename", {"unique": True}), - ("sha256", {}), + ("filename", "((doc->>'filename'))"), + ("sha256", "((doc->>'sha256'))"), ], "analytics_events": [ - ([("event_type", 1), ("timestamp", -1)], {}), - ("timestamp", {}), - ("query_normalized", {}), + ( + "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)"), ], } -def _timed_wrapper(original_method, method_name, collection_name): - """Wrap collection methods with timing logs while preserving their sync/async nature.""" - if inspect.iscoroutinefunction(original_method): - async def timed_method(*args, **kwargs): - start_time = time.monotonic() - try: - result = await 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 +# ---------------------- Result helpers ---------------------- - 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 +@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 _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 _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 _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 ---------------------- + + +class PostgresCursor: + def __init__( + self, + collection: "PostgresCollection", + filt: Optional[Dict[str, Any]], + projection: Optional[Dict[str, int]], + limit: Optional[int] = None, + *, + pipeline: Optional[List[Dict[str, Any]]] = None, + ): + 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: + 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)) + 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]]: + raise RuntimeError( + "Synchronous iteration is not supported for PostgresCursor. Use async iteration or to_list()." + ) + + +# ---------------------- Collection ---------------------- + + +class PostgresCollection: + def __init__(self, name: str, table_name: str, db: "Database"): + self.name = name + 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) + + 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 + 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: + 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, + filt: Optional[Dict[str, Any]] = None, + projection: Optional[Dict[str, int]] = None, + limit: Optional[int] = None, + ) -> PostgresCursor: + return PostgresCursor(self, filt, projection, 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, + 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 - return timed_method + 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: + doc["_id"] = str(uuid.uuid4()) + json_payload = json.dumps(_jsonable(doc), ensure_ascii=False) + await self._db.execute( + f""" + INSERT INTO {self._table_name} (id, doc) + VALUES ($1, $2::jsonb) + ON CONFLICT (id) DO UPDATE SET doc = EXCLUDED.doc + """, + str(doc["_id"]), + json_payload, + conn=connection, + ) + return InsertOneResult(inserted_id=str(doc["_id"])) + + async def update_one( + 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) + + 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, connection=connection) + 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: + 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" + ) + + 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() + 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]: + 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)) + + 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], + 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 = 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 + 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: + 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 + 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()}) + + if project_stage: + docs = self._apply_project_stage(docs, project_stage) + return docs + + +# ---------------------- 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' - ] - - 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._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, *, 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, 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, retries: int = 3, initial_delay: float = 1.0) -> None: async with self._connection_lock: - while retry_count < max_retries: + if self.is_connected: + return + + attempt = 0 + delay = initial_delay + while attempt <= retries and not self.is_connected: 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, + logger.info("Connecting to Postgres at {}", self._dsn) + self.pool = await asyncpg.create_pool( + self._dsn, min_size=1, max_size=10 ) - - # Test the connection - await self.client.admin.command("ping") - self.db = self.client[self._db_name] + await self._ensure_extensions() + await self._ensure_tables(DEFAULT_COLLECTIONS) + await self._ensure_indexes(DEFAULT_COLLECTIONS) self.is_connected = True - logger.info( - "Connected to MongoDB database '{}' with connection pool configured", - self._db_name - ) + logger.info("Connected to Postgres and storage table ensured") + self._start_monitor() 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") + attempt += 1 + logger.exception( + "Failed to connect to Postgres (attempt {})", attempt + ) self._reset_state() - if max_retries != float("inf"): - return + if attempt > retries: + break await asyncio.sleep(delay) + delay = min(delay * 2, 30) - if max_retries != float("inf"): - logger.error( - "Failed to connect to MongoDB after multiple attempts. Running in offline mode." - ) - self._reset_state() + 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 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_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) + 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} ( + id TEXT NOT NULL, + doc JSONB NOT NULL, + 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: - """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] + table_name = self._table_name_for_collection(name) + collection = PostgresCollection(name, table_name, self) + self._wrapped_collections[name] = collection + return collection + + @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, conn: Optional[asyncpg.Connection] = None + ) -> List[asyncpg.Record]: + if self.pool is None and conn is None: + raise RuntimeError("Database not connected") + 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]: - """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": getattr(self.pool, "get_max_size", lambda: None)(), + "free": getattr(self.pool, "get_idle_size", lambda: None)(), + } # Global database instance @@ -238,85 +1172,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) + logger.info("Ensuring JSONB indexes for default collections") + await db_instance._ensure_indexes(DEFAULT_COLLECTIONS) + 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: 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/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/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 a387765..d9d4e0b 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,22 @@ 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 +72,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 +99,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 +120,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 +143,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}" @@ -139,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 93d7240..6d5630b 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,21 @@ 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 24b0c29..2dd55c3 100644 --- a/src/services/page_service.py +++ b/src/services/page_service.py @@ -6,8 +6,8 @@ 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, get_history_collection, get_users_collection, @@ -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() @@ -68,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. @@ -118,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: @@ -221,27 +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 - 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 [], - ) + # 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, @@ -318,6 +331,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,49 +353,58 @@ 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 + normalized_query = query.strip() + if not normalized_query: + logger.info("Search attempted with empty query") + return [] - 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}" + 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} ) - 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)]) + 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 = await cursor.to_list(limit) - logger.info(f"Search performed: {query!r} on branch {branch!r} - found {len(pages)} results") + 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 except Exception as e: @@ -390,7 +413,6 @@ async def search_pages( ) return [] - @staticmethod async def delete_page(title: str) -> bool: """ @@ -594,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 aff6703..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 @@ -491,11 +485,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..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). @@ -99,21 +99,12 @@ 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 + 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() end = time.perf_counter() @@ -168,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 18a35fe..d5a9893 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 @@ -109,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": [], @@ -123,107 +96,16 @@ 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) + 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}" ) - 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"), - }, - } - ) + data_params.extend([offset, effective_limit]) - items.sort(key=lambda x: x["timestamp"], reverse=True) - - if not bypass: - start_index = offset - end_index = offset + effective_limit - items = items[start_index:end_index] + 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, 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