qmd_py.store¶

def add_line_numbers(text: str, start_line: int = 1) -> str:
    """Prefix each line with `N: `.

    Args:
        start_line: Number for the first line, so a slice taken from
            partway through a document still reports true line numbers.

    Returns:
        The text with every line prefixed. A trailing newline yields a
        numbered empty final line, since the split is unconditional.
    """
    lines = text.split("\n")
    return "\n".join(f"{start_line + i}: {line}" for i, line in enumerate(lines))

def extract_title(content: str, filename: str) -> str:
    """Port of the TS reference's `extractTitle()`: first `#`/`##` markdown
    heading (skipping a generic "Notes" heading in favor of the next one,
    a quirk carried over from the TS version's own note-taking-app
    interop), `#+TITLE:`/org heading for `.org` files, else the filename
    without its extension."""
    ext = filename[filename.rfind(".") :].lower() if "." in filename else ""

    if ext == ".md":
        match = _MD_HEADING.search(content)
        if match:
            title = match.group(1).strip()
            if title in ("\U0001f4dd Notes", "Notes"):
                next_match = re.search(r"^##\s+(.+)$", content, re.MULTILINE)
                if next_match:
                    return next_match.group(1).strip()
            return title
    elif ext == ".org":
        prop_match = _ORG_TITLE_PROP.search(content)
        if prop_match:
            return prop_match.group(1).strip()
        heading_match = _ORG_HEADING.search(content)
        if heading_match:
            return heading_match.group(1).strip()

    stem = re.sub(r"\.[^.]+$", "", filename)
    return stem.rsplit("/", 1)[-1] or filename

def hash_content(content: str) -> str:
    """SHA256 hex digest - matches the TS reference's `hashContent()`
    exactly (content-addressable hashes must agree byte-for-byte with the
    TS side for the Phase 3 parity check to mean anything).

    Sync: pure CPU, no I/O to await. It was `async` only because the TS
    original returns a Promise.
    """
    return hashlib.sha256(content.encode()).hexdigest()

def utcnow() -> datetime:
    """Timezone-aware current time.

    Returns:
        `datetime.now(UTC)` - aware, not naive, since every timestamp
        column is `TIMESTAMP WITH TIME ZONE` and psycopg rejects naive
        values for those.
    """
    return datetime.now(UTC)

async def cleanup_orphaned_content(session: AsyncSession) -> int:
    """A hash is only truly orphaned once NO document row - active or
    inactive - references it: `document.hash` is a NOT NULL FK with no
    `ON DELETE CASCADE`, so Postgres blocks deleting a content row an
    *inactive* document still points to, even though search/retrieval
    never reads inactive documents. Filtering this query to only active
    documents (as an earlier version did) raised a FK violation the first
    time an inactive document was the last remaining reference. Run
    `delete_inactive_documents` first (see the `cleanup` command) to
    actually reclaim content orphaned only by now-abandoned inactive
    rows.

    Returns:
        Number of content rows deleted.
    """
    result = await session.execute(
        delete(Content).where(~col(Content.hash).in_(select(col(Document.hash))))
    )
    await session.flush()
    return affected_rows(result)

async def delete_inactive_documents(session: AsyncSession, user: CurrentUser) -> int:
    """Hard-delete deactivated document rows across a user's collections.

    Removes the rows only; their content is `cleanup_orphaned_content`'s
    job, and has to run *after* this to reclaim bodies whose last
    remaining reference was one of these inactive rows.

    Deliberately irreversible: once these are gone, a file that reappears
    is indexed as a new document rather than reactivating its old row.

    Returns:
        Number of document rows deleted; 0 when the user has no
        accessible collections.
    """
    collection_ids = await resolve_collection_ids(session, user, None)
    if not collection_ids:
        return 0
    result = await session.execute(
        delete(Document).where(
            col(Document.collection_id).in_(collection_ids), ~col(Document.active)
        )
    )
    await session.flush()
    return affected_rows(result)

async def delete_llm_cache(session: AsyncSession) -> int:
    """Empty the cached-LLM-response table.

    Global, not per user: the cache is keyed by prompt, so there is
    nothing user-scoped to preserve.

    Returns:
        Number of cached responses deleted.
    """
    result = await session.execute(delete(LlmCache))
    await session.flush()
    return affected_rows(result)

async def add_collection(
    session: AsyncSession,
    user: CurrentUser,
    name: str,
    path: str,
    pattern: str = "**/*.md",
    ignore: list[str] | None = None,
) -> Collection:
    """Register a new collection. Does not index it.

    Indexing is a separate step: call `reindex_collection()` afterwards
    (that is what `marq collection add` does), so a caller can create and
    populate a collection in one transaction.

    Args:
        name: Unique per owner, not globally - two users may each have a
            collection called `notes`.
        path: Absolute filesystem path to walk. Not validated here; a
            path that doesn't exist simply yields no files at reindex.
        pattern: Glob relative to `path`. Brace groups are expanded
            (`{src,docs}/**/*.{md,py}`), which plain `glob` cannot do.
        ignore: Extra glob patterns to skip, on top of the always-excluded
            hidden files and `node_modules`/`.git`/`dist`-style directories.

    Returns:
        The persisted row, with `id` populated by the flush.

    Raises:
        sqlalchemy.exc.IntegrityError: This owner already has a collection
            of that name (`UNIQUE(owner_user_id, name)`). Raised on flush,
            so the caller must roll back before reusing the session.
    """
    collection = Collection(
        owner_user_id=user.id, name=name, path=path, pattern=pattern, ignore_patterns=ignore
    )
    session.add(collection)
    await session.flush()
    return collection

async def get_collection(session: AsyncSession, user: CurrentUser, name: str) -> Collection:
    """Fetch one collection by name, gated for read access.

    Returns:
        The ORM row, so callers can read `path`/`pattern`/`update_command`
        directly.

    Raises:
        CollectionNotFoundError: No collection of that name owned by `user`.
        PermissionDeniedError: `can_access()` refused `read` on it.
    """
    return await _resolve_owned_collection(session, user, name, "read")

async def list_collections(session: AsyncSession, user: CurrentUser) -> list[CollectionListRow]:
    """Every collection `user` can read, with per-collection statistics.

    Today that is all of them, since `can_access()` is mocked True - but
    the filter is real and applied per row, not skipped.

    Returns:
        Rows ordered by collection name.

    Note:
        Runs one statistics query per collection (an N+1). Fine at the
        scale this runs at - a handful of collections per user - but the
        counts would need folding into a single grouped query before that
        stopped being true.
    """
    result = await session.execute(select(Collection).order_by(col(Collection.name)))
    rows = []
    for collection in result.scalars():
        if not await can_read(user, collection):
            continue
        stats = await session.execute(
            select(
                func.count(col(Document.id)),
                func.max(col(Document.modified_at)),
            ).where(col(Document.collection_id) == collection.id, col(Document.active))
        )
        active_count, last_modified = stats.one()
        rows.append(
            CollectionListRow(
                name=collection.name,
                path=collection.path,
                pattern=collection.pattern,
                doc_count=active_count or 0,
                active_count=active_count or 0,
                last_modified=last_modified,
                include_by_default=collection.include_by_default,
            )
        )
    return rows

async def remove_collection(
    session: AsyncSession, user: CurrentUser, name: str
) -> RemoveCollectionResult:
    """Delete a collection, its documents, contexts and grants.

    Deletes the contexts and grants *before* the collection row -
    `collectioncontext_collection_id_fkey`/
    `collectiongrant_collection_id_fkey` have no `ON DELETE CASCADE`, so
    Postgres blocks the collection delete while either still references it
    (caught live: removing a collection with any per-path context still
    set raised a raw FK violation instead of succeeding).

    Returns:
        Counts of what was removed; see `RemoveCollectionResult`.

    Raises:
        CollectionNotFoundError: No collection of that name owned by
            `user`. Note this is also what a *non-owner* gets, because the
            lookup prefilters on ownership in SQL - see auth.py.
        PermissionDeniedError: `can_access()` refused `admin` on it.
    """
    collection = await _resolve_owned_collection(session, user, name, "admin")
    deleted = await session.execute(
        delete(Document).where(col(Document.collection_id) == collection.id)
    )
    await session.execute(
        delete(CollectionContext).where(col(CollectionContext.collection_id) == collection.id)
    )
    await session.execute(
        delete(CollectionGrant).where(col(CollectionGrant.collection_id) == collection.id)
    )
    await session.delete(collection)
    await session.flush()
    cleaned = await cleanup_orphaned_content(session)
    return RemoveCollectionResult(
        deleted_docs=affected_rows(deleted),
        cleaned_hashes=cleaned,
    )

async def rename_collection(
    session: AsyncSession, user: CurrentUser, old_name: str, new_name: str
) -> None:
    """Rename a collection in place, keeping its documents and contexts.

    Raises:
        CollectionNotFoundError: No collection named `old_name` owned by
            `user`.
        PermissionDeniedError: `can_access()` refused `admin` on it.
        sqlalchemy.exc.IntegrityError: `new_name` is already taken by
            another of this owner's collections. Raised on flush, not
            checked up front.
    """
    collection = await _resolve_owned_collection(session, user, old_name, "admin")
    collection.name = new_name
    session.add(collection)
    await session.flush()

async def set_include_by_default(
    session: AsyncSession, user: CurrentUser, name: str, include: bool
) -> None:
    """Include or exclude this collection from unscoped queries.

    Backs `marq collection include/exclude`. Excluding removes it from the
    default scope only: an explicit `-c <name>` still searches it.

    Note:
        Excluding *every* collection leaves the default scope empty, and an
        empty scope matches nothing rather than everything. That direction
        was once inverted - see `_filter_by_collections` in
        cli/commands/read.py.

    Args:
        include: True to put it back in the default scope, False to drop it.

    Raises:
        CollectionNotFoundError: No collection of that name owned by `user`.
        PermissionDeniedError: `can_access()` refused `admin` on it.
    """
    collection = await _resolve_owned_collection(session, user, name, "admin")
    collection.include_by_default = include
    session.add(collection)
    await session.flush()

async def set_update_command(
    session: AsyncSession, user: CurrentUser, name: str, command: str | None
) -> None:
    """Set the shell command run before this collection is re-indexed.

    Used for collections backed by something that has to be refreshed
    first - a `git pull`, an export script - so `marq update` can bring
    the source up to date before walking it.

    Args:
        command: Shell command to run from the collection's path, or None
            to clear it.

    Raises:
        CollectionNotFoundError: No collection of that name owned by `user`.
        PermissionDeniedError: `can_access()` refused `admin` on it.
    """
    collection = await _resolve_owned_collection(session, user, name, "admin")
    collection.update_command = command
    session.add(collection)
    await session.flush()

async def add_context(
    session: AsyncSession, user: CurrentUser, collection_name: str, path_prefix: str, text: str
) -> None:
    """Attach context prose to a collection, or to a path within it.

    Upserts: setting context for a prefix that already has one replaces
    the text rather than failing or accumulating.

    Contexts nest. A document inherits every context whose prefix it
    starts under, joined most-general-first, so a root context and a
    `journal/` context both reach `journal/entry.md`.

    Args:
        path_prefix: Path within the collection, or `""` for the whole
            collection. Matched as a path prefix, not a glob.
        text: The prose to attach. Search results carry it alongside the
            body so an agent knows what kind of document it is reading.

    Raises:
        CollectionNotFoundError: No collection of that name owned by `user`.
        PermissionDeniedError: `can_access()` refused `write` on it.
    """
    collection = await _resolve_owned_collection(session, user, collection_name, "write")
    stmt = (
        pg_insert(CollectionContext)
        .values(collection_id=collection.id, path_prefix=path_prefix, context=text)
        .on_conflict_do_update(
            index_elements=["collection_id", "path_prefix"], set_={"context": text}
        )
    )
    await session.execute(stmt)
    await session.flush()

async def context_check(
    session: AsyncSession, user: CurrentUser
) -> tuple[list[CollectionMissingContext], dict[str, list[str]]]:
    """Port of the TS reference's `getCollectionsWithoutContext` +
    `getTopLevelPathsWithoutContext`, combined into the one `context check`
    command they jointly back (which the TS CLI never actually wired up
    despite CLAUDE.md documenting it - see the qmd-py plan's list of fixed
    TS inconsistencies).

    Two different gaps are reported, which is why the return value is a
    pair: a collection with *no* context at all, and a collection that has
    some context but leaves whole top-level directories uncovered. A
    collection with a root (`""`) context is never reported for the
    second, since that covers everything beneath it.

    Returns:
        A `(collections, paths)` pair. `collections` lists collections
        with no context, sorted by name. `paths` maps a collection name to
        its uncovered top-level directories, sorted, and omits collections
        with nothing missing - so an empty dict means fully covered.

    Note:
        ACL (see auth.py): owner-prefiltered in SQL rather than gated
        through `can_access()` - same widening needed as `list_contexts`
        when grants go live.

        Runs two to three queries per collection (an N+1), which is fine
        for an interactive advisory command.
    """
    collections_result = await session.execute(
        select(Collection).where(col(Collection.owner_user_id) == user.id)
    )
    collections = list(collections_result.scalars())

    missing_context: list[CollectionMissingContext] = []
    missing_paths: dict[str, list[str]] = {}

    for collection in collections:
        contexts_result = await session.execute(
            select(CollectionContext).where(col(CollectionContext.collection_id) == collection.id)
        )
        contexts = list(contexts_result.scalars())

        if not contexts:
            count_result = await session.execute(
                select(func.count(col(Document.id))).where(
                    col(Document.collection_id) == collection.id, col(Document.active)
                )
            )
            missing_context.append(
                CollectionMissingContext(
                    name=collection.name,
                    path=collection.path,
                    doc_count=count_result.scalar_one(),
                )
            )
            continue

        prefixes = {c.path_prefix for c in contexts}
        if "" in prefixes:
            continue  # root context covers the whole collection

        paths_result = await session.execute(
            select(col(Document.path)).where(
                col(Document.collection_id) == collection.id, col(Document.active)
            )
        )
        top_level_dirs = {
            parts[0]
            for (path,) in paths_result.all()
            if len(parts := [p for p in path.split("/") if p]) > 1
        }
        missing = sorted(
            d
            for d in top_level_dirs
            if not any(p == d or d.startswith(p + "/") for p in prefixes)
        )
        if missing:
            missing_paths[collection.name] = missing

    missing_context.sort(key=lambda c: c.name)
    return missing_context, missing_paths

async def get_global_context(session: AsyncSession, user: CurrentUser) -> str | None:
    """Return this user's global context.

    Returns:
        The prose, or None when unset - which is also what a user row that
        somehow doesn't exist yields, since the caller has no use for the
        distinction.
    """
    result = await session.execute(select(col(User.global_context)).where(col(User.id) == user.id))
    return result.scalar_one_or_none()

async def list_contexts(session: AsyncSession, user: CurrentUser) -> list[ContextRow]:
    """Every context this user has set, across all their collections.

    Returns:
        Rows ordered by collection name, then path prefix - so a
        collection's root context sorts before its sub-path ones.

    Note:
        ACL (see auth.py): owner-prefiltered in SQL rather than gated
        through `can_access()`, so a granted-but-not-owned collection's
        contexts stay hidden until this query is widened.
    """
    result = await session.execute(
        select(
            col(Collection.name), col(CollectionContext.path_prefix), col(CollectionContext.context)
        )
        .join(Collection, col(CollectionContext.collection_id) == Collection.id)
        .where(col(Collection.owner_user_id) == user.id)
        .order_by(col(Collection.name), col(CollectionContext.path_prefix))
    )
    return [ContextRow(collection=r[0], path=r[1], context=r[2]) for r in result.all()]

async def remove_context(
    session: AsyncSession, user: CurrentUser, collection_name: str, path_prefix: str
) -> bool:
    """Delete the context set for one exact path prefix.

    Args:
        path_prefix: Must match what was stored exactly - this deletes one
            row, it does not clear a subtree.

    Returns:
        True if a row was deleted, False if there was nothing set for that
        prefix. Removing a context that isn't there is not an error.

    Raises:
        CollectionNotFoundError: No collection of that name owned by `user`.
        PermissionDeniedError: `can_access()` refused `write` on it.
    """
    collection = await _resolve_owned_collection(session, user, collection_name, "write")
    result = await session.execute(
        delete(CollectionContext).where(
            col(CollectionContext.collection_id) == collection.id,
            col(CollectionContext.path_prefix) == path_prefix,
        )
    )
    await session.flush()
    return affected_rows(result) > 0

async def set_global_context(session: AsyncSession, user: CurrentUser, text: str | None) -> None:
    """Set the context that applies across all of this user's collections.

    The TS reference's `context add /`. User-scoped here, not a single
    system-wide value (see `User.global_context` in db/models.py). It is
    prepended before any per-path context, so it reads as the most general
    statement about everything this user has indexed.

    Args:
        text: The prose, or None to clear it.
    """
    result = await session.execute(select(User).where(col(User.id) == user.id))
    user_row = result.scalar_one()
    user_row.global_context = text
    session.add(user_row)
    await session.flush()

async def deactivate_document(session: AsyncSession, collection_id: int, path: str) -> None:
    """Mark a document inactive, keeping the row.

    Soft delete: search and retrieval ignore inactive documents, but the
    row survives so a returning file can reactivate it rather than
    colliding with `UNIQUE(collection_id, path)`. `marq cleanup` is what
    hard-deletes them.

    A path with no row is a no-op, not an error.
    """
    await session.execute(
        update(Document)
        .where(col(Document.collection_id) == collection_id, col(Document.path) == path)
        .values(active=False)
    )
    await session.flush()

async def find_active_document(
    session: AsyncSession, collection_id: int, path: str
) -> Document | None:
    """Find the live document at a path, ignoring deactivated rows.

    Returns:
        The document, or None if there is none or it has been
        deactivated. Use `find_document_by_path()` when a deactivated row
        still matters.
    """
    result = await session.execute(
        select(Document).where(
            col(Document.collection_id) == collection_id,
            col(Document.path) == path,
            col(Document.active),
        )
    )
    return result.scalar_one_or_none()

async def find_document_by_path(
    session: AsyncSession, collection_id: int, path: str
) -> Document | None:
    """Find the document at a path whether or not it is active.

    `reindex_collection` needs this rather than `find_active_document`, to
    reactivate a deactivated document whose file reappears instead of
    violating `Document`'s table-wide `UNIQUE(collection_id, path)`
    constraint by inserting a second row.

    Returns:
        The document, active or not, or None if the path was never
        indexed.
    """
    result = await session.execute(
        select(Document).where(
            col(Document.collection_id) == collection_id, col(Document.path) == path
        )
    )
    return result.scalar_one_or_none()

async def get_active_document_paths(session: AsyncSession, collection_id: int) -> list[str]:
    """List the paths currently indexed and active in a collection.

    Returns:
        Paths relative to the collection root, unordered. `reindex_collection`
        diffs this against what it found on disk to decide what to
        deactivate.
    """
    result = await session.execute(
        select(col(Document.path)).where(
            col(Document.collection_id) == collection_id, col(Document.active)
        )
    )
    return [p for (p,) in result.all()]

async def insert_content(session: AsyncSession, hash_: str, doc: str) -> None:
    """Store a body under its content hash, if not already present.

    Idempotent by design: content is addressed by hash and shared across
    every collection, so re-indexing an unchanged file - or the same file
    in a second worktree - inserts nothing and costs nothing.

    Args:
        hash_: SHA-256 hex digest of `doc`, from `hash_content()`. Not
            recomputed or verified here.
        doc: The full body text.
    """
    stmt = (
        pg_insert(Content)
        .values(hash=hash_, doc=doc)
        .on_conflict_do_nothing(index_elements=["hash"])
    )
    await session.execute(stmt)

async def insert_document(
    session: AsyncSession,
    collection_id: int,
    path: str,
    title: str,
    hash_: str,
    created_at: datetime,
    modified_at: datetime,
) -> Document:
    """Create a document row and build its search vector.

    Assumes the content row already exists - call `insert_content()`
    first, since `document.hash` is a NOT NULL foreign key.

    Args:
        path: Path relative to the collection root.
        hash_: Content hash linking to the stored body.
        created_at: Usually the file's mtime rather than now, so the value
            survives a reindex from a fresh clone.
        modified_at: The file's mtime.

    Returns:
        The persisted row, with `id` populated by the flush.

    Raises:
        sqlalchemy.exc.IntegrityError: A row already exists for this
            `(collection_id, path)`, active or not, or `hash_` has no
            content row.
    """
    document = Document(
        collection_id=collection_id,
        path=path,
        title=title,
        hash=hash_,
        created_at=created_at,
        modified_at=modified_at,
    )
    session.add(document)
    await session.flush()
    await update_document_search_vector(session, document.id)
    return document