diff --git a/pageindex/backend/cloud.py b/pageindex/backend/cloud.py index 48006f2..ddf6b93 100644 --- a/pageindex/backend/cloud.py +++ b/pageindex/backend/cloud.py @@ -6,6 +6,7 @@ API reference: https://github.com/VectifyAI/pageindex_sdk from __future__ import annotations import json import logging +import os import time import urllib.parse import requests @@ -30,6 +31,11 @@ def _as_int(value): return None +def _doc_type_from_name(name: str) -> str: + ext = os.path.splitext(name)[1].lstrip(".").lower() + return ext or "pdf" + + class CloudBackend: def __init__(self, api_key: str | Callable[[], str], base_url: str | Callable[[], str] | None = None): @@ -292,11 +298,12 @@ class CloudBackend: params={"type": "tree", "summary": "true"}) raw_tree = tree_resp.get("result", []) page_num = _as_int(resp.get("pageNum")) + doc_name = resp.get("name", "") result = { "doc_id": resp.get("id", doc_id), - "doc_name": resp.get("name", ""), + "doc_name": doc_name, "doc_description": resp.get("description", ""), - "doc_type": "pdf", + "doc_type": _doc_type_from_name(doc_name), "status": resp.get("status", ""), "structure": self._normalize_tree(raw_tree, max_page=page_num), } @@ -385,9 +392,8 @@ class CloudBackend: def list_documents(self, collection: str) -> list[dict]: folder_id = self._get_folder_id(collection) - # The API caps `limit` at 100; paginate with `offset` until a short - # page comes back so collections with >100 docs aren't silently - # truncated (queries over the whole collection rely on this list). + # Paginate with `offset` until a short page comes back so large + # collections aren't silently truncated. page_size = 100 offset = 0 docs: list[dict] = [] @@ -397,15 +403,14 @@ class CloudBackend: params["folder_id"] = folder_id data = self._request("GET", "/docs/", params=params) batch = data.get("documents", []) or [] - docs.extend( - { + for d in batch: + name = d.get("name", "") + docs.append({ "doc_id": d.get("id", ""), - "doc_name": d.get("name", ""), + "doc_name": name, "doc_description": d.get("description", ""), - "doc_type": "pdf", - } - for d in batch - ) + "doc_type": _doc_type_from_name(name), + }) if len(batch) < page_size: return docs offset += page_size diff --git a/pageindex/backend/local.py b/pageindex/backend/local.py index 3eb3f17..b8a9746 100644 --- a/pageindex/backend/local.py +++ b/pageindex/backend/local.py @@ -15,11 +15,13 @@ from ..index.pipeline import build_index from ..index.utils import parse_pages, get_pdf_page_content, remove_fields from ..backend.protocol import AgentTools from ..errors import (FileTypeError, DocumentNotFoundError, CollectionNotFoundError, - IndexingError, PageIndexError) + IndexingError) +from .._validation import validate_collection_name -# Matched with .fullmatch() (not .match()): a $-anchored .match() would accept a -# trailing newline ("papers\n") because $ matches just before a final \n. -_COLLECTION_NAME_RE = re.compile(r'[a-zA-Z0-9_-]{1,128}') +# Collections created by older SDK versions used their names directly as safe +# directory names. Preserve that layout for backward compatibility; names newly +# allowed by the cloud-compatible contract use an opaque directory instead. +_LEGACY_COLLECTION_DIR_RE = re.compile(r'[a-zA-Z0-9_-]{1,128}') class LocalBackend: @@ -47,8 +49,19 @@ class LocalBackend: # Collection management def _validate_collection_name(self, name: str) -> None: - if not _COLLECTION_NAME_RE.fullmatch(name): - raise PageIndexError(f"Invalid collection name: {name!r}. Must be 1-128 chars of [a-zA-Z0-9_-].") + validate_collection_name(name) + + def _collection_dir(self, name: str) -> Path: + """Return a stable, contained directory for a logical collection name. + + Legacy-safe names retain their historical ``files/{name}`` layout. Any + other cloud-valid name is hashed so spaces, Unicode, slashes, ``..``, or + platform-specific path characters can never escape ``files_dir``. + """ + if _LEGACY_COLLECTION_DIR_RE.fullmatch(name): + return self._files_dir / name + digest = hashlib.sha256(name.encode("utf-8")).hexdigest() + return self._files_dir / ".collections" / digest def create_collection(self, name: str) -> None: self._validate_collection_name(name) @@ -62,11 +75,9 @@ class LocalBackend: return self._storage.list_collections() def delete_collection(self, name: str) -> None: - # Validate before touching the filesystem — an unvalidated name like - # "../.." would make the rmtree below escape files_dir entirely. self._validate_collection_name(name) self._storage.delete_collection(name) - col_dir = self._files_dir / name + col_dir = self._collection_dir(name) if col_dir.exists(): shutil.rmtree(col_dir) @@ -107,13 +118,13 @@ class LocalBackend: # Copy file to managed directory ext = os.path.splitext(file_path)[1] - col_dir = self._files_dir / collection + col_dir = self._collection_dir(collection) col_dir.mkdir(parents=True, exist_ok=True) managed_path = col_dir / f"{doc_id}{ext}" shutil.copy2(file_path, managed_path) try: - # Store images alongside the document: files/{collection}/{doc_id}/images/ + # Store images alongside the managed document directory. images_dir = str(col_dir / doc_id / "images") parsed = parser.parse(file_path, model=self._model, images_dir=images_dir) result = build_index(parsed, model=self._model, opt=self._index_config) @@ -134,6 +145,7 @@ class LocalBackend: "file_path": str(managed_path), "file_hash": file_hash, "doc_type": ext.lstrip("."), + "status": "completed", **(parsed.metadata or {}), # parser-reported, e.g. page_count / line_count "structure": result["structure"], "pages": pages, @@ -185,7 +197,6 @@ class LocalBackend: use in agent/LLM contexts as it can exhaust the context window. """ doc = self._require_document(collection, doc_id) - doc["status"] = "completed" # local indexing is synchronous doc["structure"] = self._storage.get_document_structure(collection, doc_id) if include_text: pages = self._storage.get_pages(collection, doc_id) or [] @@ -248,7 +259,7 @@ class LocalBackend: self._storage.delete_document(collection, doc_id) if doc.get("file_path"): Path(doc["file_path"]).unlink(missing_ok=True) - doc_dir = self._files_dir / collection / doc_id + doc_dir = self._collection_dir(collection) / doc_id if doc_dir.exists(): shutil.rmtree(doc_dir) diff --git a/pageindex/storage/sqlite.py b/pageindex/storage/sqlite.py index 9526833..c65d683 100644 --- a/pageindex/storage/sqlite.py +++ b/pageindex/storage/sqlite.py @@ -1,23 +1,16 @@ import json -import re import sqlite3 import threading from pathlib import Path -from ..errors import CollectionAlreadyExistsError, PageIndexError - -# Mirrors LocalBackend's own collection-name rule. SQLiteStorage enforces this -# itself (not just relying on LocalBackend's pre-check or the schema's CHECK -# constraint below) because it's a public StorageEngine that can be used -# directly, bypassing LocalBackend entirely. -# Matched with .fullmatch() (not .match()): a $-anchored .match() would accept a -# trailing newline ("papers\n") because $ matches just before a final \n. -_COLLECTION_NAME_RE = re.compile(r'[a-zA-Z0-9_-]{1,128}') +from ..errors import CollectionAlreadyExistsError +from .._validation import validate_collection_name def _validate_collection_name(name: str) -> None: - if not _COLLECTION_NAME_RE.fullmatch(name): - raise PageIndexError(f"Invalid collection name: {name!r}. Must be 1-128 chars of [a-zA-Z0-9_-].") + # SQLiteStorage is a public StorageEngine and may be used without + # LocalBackend, so enforce the shared contract at this boundary too. + validate_collection_name(name) class SQLiteStorage: @@ -75,46 +68,98 @@ class SQLiteStorage: def _init_schema(self): conn = self._get_conn() - conn.execute("PRAGMA user_version = 1") - conn.executescript(""" - CREATE TABLE IF NOT EXISTS collections ( - -- GLOB '*' is "any characters", not a regex quantifier over the - -- preceding class — '[a-zA-Z0-9_-]*' alone only constrains the - -- FIRST character. The second GLOB (NOT ... '*[^...]*') checks - -- every remaining character too, so this is real defense-in-depth - -- for direct SQLiteStorage use (bypassing _validate_collection_name - -- above), not just a first-character gate. - name TEXT PRIMARY KEY CHECK( - length(name) BETWEEN 1 AND 128 - AND name GLOB '[a-zA-Z0-9_-]*' - AND name NOT GLOB '*[^a-zA-Z0-9_-]*' - ), - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP - ); - CREATE TABLE IF NOT EXISTS documents ( - doc_id TEXT PRIMARY KEY, - collection_name TEXT NOT NULL REFERENCES collections(name) ON DELETE CASCADE, - doc_name TEXT, - doc_description TEXT, - file_path TEXT, - file_hash TEXT, - doc_type TEXT NOT NULL, - page_count INTEGER, - line_count INTEGER, - structure JSON, - pages JSON, - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - UNIQUE(collection_name, file_hash) - ); - CREATE INDEX IF NOT EXISTS idx_docs_collection ON documents(collection_name); - CREATE INDEX IF NOT EXISTS idx_docs_hash ON documents(collection_name, file_hash); - """) - # DBs created before the count columns existed: add them in place. + # SQLite cannot ALTER a CHECK constraint. Rebuild the small parent table + # transactionally when opening a v1 database; documents keep referring to + # the same table name and are verified after foreign keys are re-enabled. + conn.execute("PRAGMA foreign_keys=OFF") + try: + conn.execute("BEGIN IMMEDIATE") + schema_version = conn.execute("PRAGMA user_version").fetchone()[0] + conn.execute(""" + CREATE TABLE IF NOT EXISTS collections ( + name TEXT PRIMARY KEY NOT NULL + CHECK( + length(name) BETWEEN 1 AND 255 + -- SQLite length(TEXT) stops at the first NUL, while + -- Python and the cloud API count it as a character. + -- The Python boundary still enforces the 255 limit. + OR (instr(name, char(0)) > 0 AND name <> '') + ), + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + ) + """) + schema_row = conn.execute( + "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'collections'" + ).fetchone() + schema_sql = schema_row[0] if schema_row else "" + schema_upper = schema_sql.upper() + if "BETWEEN 1 AND 128" in schema_upper or "NAME GLOB" in schema_upper: + conn.execute(""" + CREATE TABLE _pageindex_collections_v2 ( + name TEXT PRIMARY KEY NOT NULL + CHECK( + length(name) BETWEEN 1 AND 255 + OR (instr(name, char(0)) > 0 AND name <> '') + ), + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + ) + """) + conn.execute(""" + INSERT INTO _pageindex_collections_v2 (name, created_at) + SELECT name, created_at FROM collections + """) + conn.execute("DROP TABLE collections") + conn.execute("ALTER TABLE _pageindex_collections_v2 RENAME TO collections") + + conn.execute(""" + CREATE TABLE IF NOT EXISTS documents ( + doc_id TEXT PRIMARY KEY, + collection_name TEXT NOT NULL REFERENCES collections(name) ON DELETE CASCADE, + doc_name TEXT, + doc_description TEXT, + file_path TEXT, + file_hash TEXT, + doc_type TEXT NOT NULL, + status TEXT NOT NULL, + page_count INTEGER, + line_count INTEGER, + structure JSON, + pages JSON, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + UNIQUE(collection_name, file_hash) + ) + """) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_docs_collection ON documents(collection_name)" + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_docs_hash ON documents(collection_name, file_hash)" + ) + if schema_version < 2: + conn.execute("PRAGMA user_version = 2") + conn.commit() + except Exception: + conn.rollback() + raise + finally: + conn.execute("PRAGMA foreign_keys=ON") + + violations = conn.execute("PRAGMA foreign_key_check").fetchall() + if violations: + raise sqlite3.IntegrityError( + f"Foreign-key violations after SQLite schema migration: {violations!r}" + ) + + # DBs created before these columns existed: add them in place. cols = {r[1] for r in conn.execute("PRAGMA table_info(documents)")} - for name in ("page_count", "line_count"): - if name not in cols: + for col_name, col_def in ( + ("page_count", "INTEGER"), + ("line_count", "INTEGER"), + ("status", "TEXT NOT NULL DEFAULT 'completed'"), + ): + if col_name not in cols: try: - conn.execute(f"ALTER TABLE documents ADD COLUMN {name} INTEGER") + conn.execute(f"ALTER TABLE documents ADD COLUMN {col_name} {col_def}") except sqlite3.OperationalError: pass # concurrent open of the same legacy DB already added it conn.commit() @@ -133,7 +178,16 @@ class SQLiteStorage: _validate_collection_name(name) with self._write_lock: conn = self._get_conn() + # INSERT OR IGNORE works with older SQLite versions, but it can also + # suppress CHECK failures. Verify the row so ignored non-duplicate + # constraints never falsely report success. conn.execute("INSERT OR IGNORE INTO collections (name) VALUES (?)", (name,)) + if conn.execute( + "SELECT 1 FROM collections WHERE name = ?", (name,) + ).fetchone() is None: + raise sqlite3.IntegrityError( + f"Collection {name!r} was rejected by the SQLite schema" + ) conn.commit() def list_collections(self) -> list[str]: @@ -155,10 +209,11 @@ class SQLiteStorage: conn = self._get_conn() conn.execute( """INSERT INTO documents - (doc_id, collection_name, doc_name, doc_description, file_path, file_hash, doc_type, page_count, line_count, structure, pages) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + (doc_id, collection_name, doc_name, doc_description, file_path, file_hash, doc_type, status, page_count, line_count, structure, pages) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", (doc_id, collection, doc.get("doc_name"), doc.get("doc_description"), doc.get("file_path"), doc.get("file_hash"), doc["doc_type"], + doc["status"], doc.get("page_count"), doc.get("line_count"), json.dumps(doc.get("structure", [])), json.dumps(doc.get("pages")) if doc.get("pages") else None), @@ -176,14 +231,14 @@ class SQLiteStorage: def get_document(self, collection: str, doc_id: str) -> dict: conn = self._get_conn() row = conn.execute( - "SELECT doc_id, doc_name, doc_description, file_path, doc_type, page_count, line_count FROM documents WHERE doc_id = ? AND collection_name = ?", + "SELECT doc_id, doc_name, doc_description, file_path, doc_type, status, page_count, line_count FROM documents WHERE doc_id = ? AND collection_name = ?", (doc_id, collection), ).fetchone() if not row: return {} doc = {"doc_id": row[0], "doc_name": row[1], "doc_description": row[2], - "file_path": row[3], "doc_type": row[4]} - doc.update({k: v for k, v in (("page_count", row[5]), ("line_count", row[6])) if v is not None}) + "file_path": row[3], "doc_type": row[4], "status": row[5]} + doc.update({k: v for k, v in (("page_count", row[6]), ("line_count", row[7])) if v is not None}) return doc def get_document_structure(self, collection: str, doc_id: str) -> list: