import sqlite3 import struct import json from pathlib import Path from contextlib import contextmanager from typing import Optional DB_PATH = Path(__file__).parent.parent / "data" / "bookmarks.db" SCHEMA = """ CREATE TABLE IF NOT EXISTS bookmarks ( id INTEGER PRIMARY KEY AUTOINCREMENT, url TEXT NOT NULL UNIQUE, title TEXT NOT NULL, description TEXT DEFAULT '', context TEXT DEFAULT '', folder_path TEXT DEFAULT '', domain TEXT DEFAULT '', favicon_url TEXT DEFAULT '', auto_tags TEXT DEFAULT '[]', user_tags TEXT DEFAULT '[]', content_snippet TEXT DEFAULT '', embedding_text TEXT DEFAULT '', created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, last_visited_at TIMESTAMP, visit_count INTEGER DEFAULT 0, is_dead BOOLEAN DEFAULT FALSE ); CREATE INDEX IF NOT EXISTS idx_bookmarks_domain ON bookmarks(domain); CREATE INDEX IF NOT EXISTS idx_bookmarks_created ON bookmarks(created_at); """ FTS_SCHEMA = """ CREATE VIRTUAL TABLE IF NOT EXISTS bookmarks_fts USING fts5( title, description, context, auto_tags, user_tags, content_snippet, content='bookmarks', content_rowid='id' ); """ VEC_SCHEMA = """ CREATE VIRTUAL TABLE IF NOT EXISTS vec_bookmarks USING vec0( id INTEGER PRIMARY KEY, embedding FLOAT[384] ); """ FTS_TRIGGERS = """ CREATE TRIGGER IF NOT EXISTS bookmarks_ai AFTER INSERT ON bookmarks BEGIN INSERT INTO bookmarks_fts(rowid, title, description, context, auto_tags, user_tags, content_snippet) VALUES (new.id, new.title, new.description, new.context, new.auto_tags, new.user_tags, new.content_snippet); END; CREATE TRIGGER IF NOT EXISTS bookmarks_ad AFTER DELETE ON bookmarks BEGIN INSERT INTO bookmarks_fts(bookmarks_fts, rowid, title, description, context, auto_tags, user_tags, content_snippet) VALUES ('delete', old.id, old.title, old.description, old.context, old.auto_tags, old.user_tags, old.content_snippet); END; CREATE TRIGGER IF NOT EXISTS bookmarks_au AFTER UPDATE ON bookmarks BEGIN INSERT INTO bookmarks_fts(bookmarks_fts, rowid, title, description, context, auto_tags, user_tags, content_snippet) VALUES ('delete', old.id, old.title, old.description, old.context, old.auto_tags, old.user_tags, old.content_snippet); INSERT INTO bookmarks_fts(rowid, title, description, context, auto_tags, user_tags, content_snippet) VALUES (new.id, new.title, new.description, new.context, new.auto_tags, new.user_tags, new.content_snippet); END; """ def _row_to_dict(row: sqlite3.Row, columns: Optional[list[str]] = None) -> dict: """Convert a sqlite3.Row to a plain dict.""" return dict(row) @contextmanager def get_db(db_path: Optional[str] = None): """Context manager yielding a sqlite3 Connection with sqlite_vec loaded and schema ensured.""" path = Path(db_path) if db_path else DB_PATH path.parent.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect(str(path)) conn.row_factory = sqlite3.Row conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA foreign_keys=ON") conn.enable_load_extension(True) try: import sqlite_vec sqlite_vec.load(conn) except ImportError: raise RuntimeError("sqlite_vec is required for vector operations. Install it with 'pip install sqlite-vec'") conn.enable_load_extension(False) _ensure_schema(conn) try: yield conn finally: conn.close() def _ensure_schema(conn: sqlite3.Connection) -> None: """Create tables and triggers if they do not exist.""" for statement in SCHEMA.strip().split(";"): stmt = statement.strip() if stmt: conn.execute(stmt) for statement in FTS_SCHEMA.strip().split(";"): stmt = statement.strip() if stmt: conn.execute(stmt) for statement in VEC_SCHEMA.strip().split(";"): stmt = statement.strip() if stmt: conn.execute(stmt) for statement in FTS_TRIGGERS.strip().split(";"): stmt = statement.strip() if stmt: conn.execute(stmt) conn.commit() def init_db(db_path: Optional[str] = None) -> None: """Explicitly initialise the database (creates all tables/triggers).""" with get_db(db_path) as conn: pass # schema is ensured inside get_db def add_bookmark( conn: sqlite3.Connection, url: str, title: str, description: str, context: str, folder_path: str, domain: str, auto_tags: list[str], user_tags: list[str], content_snippet: str, embedding_text: str, embedding_blob: bytes, ) -> int: """Insert a bookmark and its vector. Returns the new row id.""" cursor = conn.execute( """ INSERT INTO bookmarks (url, title, description, context, folder_path, domain, auto_tags, user_tags, content_snippet, embedding_text) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( url, title, description, context, folder_path, domain, json.dumps(auto_tags, ensure_ascii=False), json.dumps(user_tags, ensure_ascii=False), content_snippet, embedding_text, ), ) new_id = cursor.lastrowid conn.execute( "INSERT INTO vec_bookmarks(id, embedding) VALUES (?, ?)", (new_id, embedding_blob), ) return new_id def get_bookmark(conn: sqlite3.Connection, id: int) -> Optional[dict]: """Fetch a single bookmark by id, or None if not found.""" row = conn.execute("SELECT * FROM bookmarks WHERE id = ?", (id,)).fetchone() if row is None: return None bm = dict(row) bm["auto_tags"] = json.loads(bm.get("auto_tags") or "[]") bm["user_tags"] = json.loads(bm.get("user_tags") or "[]") return bm def update_bookmark(conn: sqlite3.Connection, id: int, **kwargs) -> None: """Update specified fields on a bookmark. updated_at is set automatically.""" if not kwargs: return # Define allowed fields for update ALLOWED_FIELDS = { "url", "title", "description", "context", "folder_path", "domain", "auto_tags", "user_tags", "content_snippet", "embedding_text", "created_at", "updated_at", "last_visited_at", "visit_count", "is_dead" } # Validate that all provided fields are allowed for key in kwargs: if key not in ALLOWED_FIELDS: raise ValueError(f"Unknown field '{key}' for bookmark update. Allowed fields: {ALLOWED_FIELDS}") # Serialise tag lists if provided for key in ("auto_tags", "user_tags"): if key in kwargs and isinstance(kwargs[key], list): kwargs[key] = json.dumps(kwargs[key], ensure_ascii=False) kwargs["updated_at"] = "CURRENT_TIMESTAMP" set_clauses = [] values = [] for key, value in kwargs.items(): if key == "updated_at": set_clauses.append("updated_at = CURRENT_TIMESTAMP") else: set_clauses.append(f"{key} = ?") values.append(value) values.append(id) conn.execute( f"UPDATE bookmarks SET {', '.join(set_clauses)} WHERE id = ?", values, ) def delete_bookmark(conn: sqlite3.Connection, id: int) -> None: """Delete a bookmark and its vector entry.""" conn.execute("DELETE FROM bookmarks WHERE id = ?", (id,)) conn.execute("DELETE FROM vec_bookmarks WHERE id = ?", (id,)) def list_bookmarks( conn: sqlite3.Connection, page: int = 1, per_page: int = 20 ) -> tuple[list[dict], int]: """Return a page of bookmarks and the total count.""" total = conn.execute("SELECT COUNT(*) FROM bookmarks").fetchone()[0] offset = (page - 1) * per_page rows = conn.execute( "SELECT * FROM bookmarks ORDER BY created_at DESC LIMIT ? OFFSET ?", (per_page, offset), ).fetchall() bookmarks = [] for row in rows: bm = dict(row) bm["auto_tags"] = json.loads(bm.get("auto_tags") or "[]") bm["user_tags"] = json.loads(bm.get("user_tags") or "[]") bookmarks.append(bm) return bookmarks, total def get_all_tags(conn: sqlite3.Connection) -> list[tuple[str, int]]: """Return all unique tags with their occurrence counts, sorted by count desc.""" rows = conn.execute("SELECT auto_tags, user_tags FROM bookmarks").fetchall() counts: dict[str, int] = {} for row in rows: for col in ("auto_tags", "user_tags"): tags = json.loads(row[col] or "[]") for tag in tags: if tag: counts[tag] = counts.get(tag, 0) + 1 return sorted(counts.items(), key=lambda x: -x[1]) def get_stats(conn: sqlite3.Connection) -> dict: """Return aggregate statistics.""" total = conn.execute("SELECT COUNT(*) FROM bookmarks").fetchone()[0] dead_links = conn.execute( "SELECT COUNT(*) FROM bookmarks WHERE is_dead = TRUE" ).fetchone()[0] top_domains_rows = conn.execute( """ SELECT domain, COUNT(*) as cnt FROM bookmarks WHERE domain != '' GROUP BY domain ORDER BY cnt DESC LIMIT 10 """ ).fetchall() top_domains = [{"domain": r["domain"], "count": r["cnt"]} for r in top_domains_rows] return { "total_bookmarks": total, "unique_domains": len(top_domains), "dead_links": dead_links, "top_domains": top_domains, } def search_fts( conn: sqlite3.Connection, query: str, limit: int = 20 ) -> list[tuple[int, float]]: """FTS5 full-text search. Returns list of (rowid, rank).""" try: rows = conn.execute( """ SELECT rowid, rank FROM bookmarks_fts WHERE bookmarks_fts MATCH ? ORDER BY rank LIMIT ? """, (query, limit), ).fetchall() return [(row[0], row[1]) for row in rows] except sqlite3.OperationalError: # Invalid FTS query (e.g. bare special chars) return [] def search_vec( conn: sqlite3.Connection, query_blob: bytes, limit: int = 20 ) -> list[tuple[int, float]]: """sqlite-vec KNN search. Returns list of (id, distance).""" rows = conn.execute( """ SELECT id, distance FROM vec_bookmarks WHERE embedding MATCH ? ORDER BY distance LIMIT ? """, (query_blob, limit), ).fetchall() return [(row[0], row[1]) for row in rows]