| 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 |
|
|
|
|
| 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 |
|
|
| |
| 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" |
| } |
|
|
| |
| for key in kwargs: |
| if key not in ALLOWED_FIELDS: |
| raise ValueError(f"Unknown field '{key}' for bookmark update. Allowed fields: {ALLOWED_FIELDS}") |
|
|
| |
| 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: |
| |
| 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] |
|
|
