DocsGPT/docsgpt/storage/db/repositories/sources.py

"""Repository for the ``sources`` table."""

from __future__ import annotations

import json
import uuid
from typing import Any, Optional

from sqlalchemy import case, Connection, func, or_, select, text

from docsgpt.storage.db.base_repository import looks_like_uuid, row_to_dict
from docsgpt.storage.db.models import ingest_chunk_progress_table, sources_table
from docsgpt.storage.db.source_config import SourceConfig


_SCALAR_COLUMNS = {
    "name", "type", "retriever", "sync_frequency", "tokens", "file_path",
    "language", "model", "date",
}
_JSONB_COLUMNS = {"metadata", "remote_data", "directory_structure", "file_name_map", "config"}
_ALLOWED_COLUMNS = _SCALAR_COLUMNS | _JSONB_COLUMNS

# Whitelist for sort columns exposed via ``list_for_user``. Anything not in
# this set falls back to ``date`` so user-supplied sort params can't be
# interpolated into SQL unchecked.
_SORTABLE_COLUMNS = {"date", "name", "tokens", "type", "created_at", "updated_at"}


def _coerce_uuid_ids(extra_ids: Optional[list]) -> list:
    """Coerce id strings to ``uuid.UUID`` for binding against UUID columns.

    Non-UUID-looking ids (e.g. legacy/synthetic) are dropped — they can never
    match a ``sources.id`` value. Used to OR team-shared source ids into the
    owner-scoped queries.
    """
    if not extra_ids:
        return []
    out: list = []
    for raw in extra_ids:
        s = str(raw)
        if looks_like_uuid(s):
            out.append(uuid.UUID(s))
    return out


def _owned_or_shared_scope(user_id: str, extra_ids: Optional[list]):
    """WHERE predicate matching the owner's rows plus any ``extra_ids``.

    Lets the paginated/count queries include team-shared sources (passed by id)
    alongside owned ones in a single query, so count/sort/search/pagination
    stay correct across the union.
    """
    scope = sources_table.c.user_id == user_id
    ids = _coerce_uuid_ids(extra_ids)
    if ids:
        scope = or_(scope, sources_table.c.id.in_(ids))
    return scope


def _escape_like(pattern: str) -> str:
    """Escape wildcards so a user-supplied substring is matched literally.

    We use ``LIKE ESCAPE '\\'`` on the query side so backslash, percent, and
    underscore in the input don't accidentally turn into regex-like wildcards.
    """
    return (
        pattern
        .replace("\\", "\\\\")
        .replace("%", "\\%")
        .replace("_", "\\_")
    )


def _coerce_jsonb(value: Any) -> Any:
    """Normalize incoming JSONB values for the Core ``Table.update()`` path.

    ``remote_data`` in particular arrives as either a dict or a JSON string
    (the legacy Mongo docs stored both shapes). Strings are parsed so the
    stored representation is always structured JSONB; dicts/lists pass
    through untouched for the SQLAlchemy JSONB type processor.
    """
    if value is None:
        return None
    if isinstance(value, (dict, list)):
        return value
    if isinstance(value, str):
        stripped = value.strip()
        if not stripped:
            return None
        try:
            return json.loads(stripped)
        except json.JSONDecodeError:
            return {"raw": value}
    return value


def _normalize_config(config: Any) -> dict:
    """Strict-validate the write-path ``config`` and return a plain dict.

    ``None`` becomes ``{}`` (classic defaults). A dict is validated through
    ``SourceConfig`` (raises on bad input, D7 strict-on-write) and dumped
    back to a normalized dict for JSONB storage.
    """
    if config is None:
        return {}
    if not isinstance(config, dict):
        raise ValueError("config must be a dict")
    return SourceConfig.model_validate(config).model_dump()


def _ingest_status_case():
    """Derive a user-facing ingest status from the joined progress row.

    ``failed`` — reconciler-escalated stall. ``processing`` — embed in
    flight. ``None`` — no progress row, or the embed completed.
    """
    icp = ingest_chunk_progress_table
    return case(
        (icp.c.source_id.is_(None), None),
        (icp.c.status == "stalled", "failed"),
        (icp.c.embedded_chunks < icp.c.total_chunks, "processing"),
        else_=None,
    ).label("ingest_status")


class SourcesRepository:
    def __init__(self, conn: Connection) -> None:
        self._conn = conn

    def create(
        self,
        name: str,
        *,
        source_id: Optional[str] = None,
        user_id: str,
        type: Optional[str] = None,
        metadata: Optional[dict] = None,
        config: Optional[dict] = None,
        retriever: Optional[str] = None,
        sync_frequency: Optional[str] = None,
        tokens: Optional[str] = None,
        file_path: Optional[str] = None,
        remote_data: Any = None,
        directory_structure: Any = None,
        file_name_map: Any = None,
        language: Optional[str] = None,
        model: Optional[str] = None,
        date: Any = None,
        legacy_mongo_id: Optional[str] = None,
    ) -> dict:
        result = self._conn.execute(
            text(
                """
                INSERT INTO sources (
                    id, user_id, name, type, metadata, config,
                    retriever, sync_frequency, tokens, file_path,
                    remote_data, directory_structure, file_name_map,
                    language, model, date, legacy_mongo_id
                )
                VALUES (
                    COALESCE(CAST(:source_id AS uuid), gen_random_uuid()),
                    :user_id, :name, :type, CAST(:metadata AS jsonb),
                    CAST(:config AS jsonb),
                    :retriever, :sync_frequency, :tokens, :file_path,
                    CAST(:remote_data AS jsonb),
                    CAST(:directory_structure AS jsonb),
                    CAST(:file_name_map AS jsonb),
                    :language, :model,
                    COALESCE(:date, now()),
                    :legacy_mongo_id
                )
                RETURNING *
                """
            ),
            {
                "source_id": source_id,
                "user_id": user_id,
                "name": name,
                "type": type,
                "metadata": json.dumps(metadata or {}),
                "config": json.dumps(_normalize_config(config)),
                "retriever": retriever,
                "sync_frequency": sync_frequency,
                "tokens": tokens,
                "file_path": file_path,
                "remote_data": (
                    None if remote_data is None
                    else json.dumps(_coerce_jsonb(remote_data))
                ),
                "directory_structure": (
                    None if directory_structure is None
                    else json.dumps(_coerce_jsonb(directory_structure))
                ),
                "file_name_map": (
                    None if file_name_map is None
                    else json.dumps(_coerce_jsonb(file_name_map))
                ),
                "language": language,
                "model": model,
                "date": date,
                "legacy_mongo_id": legacy_mongo_id,
            },
        )
        return row_to_dict(result.fetchone())

    def get(self, source_id: str, user_id: str) -> Optional[dict]:
        result = self._conn.execute(
            text("SELECT * FROM sources WHERE id = CAST(:id AS uuid) AND user_id = :user_id"),
            {"id": source_id, "user_id": user_id},
        )
        row = result.fetchone()
        return row_to_dict(row) if row is not None else None

    def get_any(self, source_id: str, user_id: str) -> Optional[dict]:
        """Resolve a source by either PG UUID or legacy Mongo ObjectId string.

        Cutover helper: URLs / bookmarks may still hold Mongo ObjectIds.
        Tries the UUID path first, then falls back to ``legacy_mongo_id``.
        Both paths are scoped by ``user_id``.
        """
        if looks_like_uuid(source_id):
            row = self.get(source_id, user_id)
            if row is not None:
                return row
        return self.get_by_legacy_id(source_id, user_id)

    def get_by_id(self, source_id: str) -> Optional[dict]:
        """Fetch a source by id with NO ownership scoping.

        Used ONLY after a team-grant authorization check (team sharing). Never
        call on a raw user-supplied id without that check.
        """
        if not looks_like_uuid(source_id):
            return None
        result = self._conn.execute(
            text("SELECT * FROM sources WHERE id = CAST(:id AS uuid)"),
            {"id": source_id},
        )
        row = result.fetchone()
        return row_to_dict(row) if row is not None else None

    def list_by_ids(self, source_ids) -> list[dict]:
        """Fetch sources whose id is in ``source_ids`` (team-shared listing path)."""
        ids = [str(s) for s in source_ids if looks_like_uuid(str(s))]
        if not ids:
            return []
        result = self._conn.execute(
            text("SELECT * FROM sources WHERE id = ANY(CAST(:ids AS uuid[])) ORDER BY created_at DESC"),
            {"ids": ids},
        )
        return [row_to_dict(r) for r in result.fetchall()]

    def find_by_name(self, user_id: str, name: str) -> Optional[dict]:
        """Return a user's source whose name matches ``name`` (case-insensitive).

        Used by agent YAML import to map a portable source name back to a
        concrete source id. Returns the oldest match when names collide.
        """
        if not name:
            return None
        result = self._conn.execute(
            text(
                "SELECT * FROM sources "
                "WHERE user_id = :user_id AND lower(name) = lower(:name) "
                "ORDER BY created_at LIMIT 1"
            ),
            {"user_id": user_id, "name": name},
        )
        row = result.fetchone()
        return row_to_dict(row) if row is not None else None

    def list_for_user(
        self,
        user_id: str,
        *,
        limit: Optional[int] = None,
        offset: int = 0,
        search_term: Optional[str] = None,
        sort_field: str = "created_at",
        sort_order: str = "desc",
        extra_ids: Optional[list] = None,
    ) -> list[dict]:
        """Return sources owned by ``user_id``, paginated and optionally filtered.

        All pagination, filtering, and sorting are pushed into SQL so large
        accounts don't materialize their full source list in Python for every
        page. See ``PaginatedSources`` in the sources routes for the matching
        call site.

        Args:
            user_id: Scope rows to this owner.
            limit: Page size. ``None`` returns every matching row (legacy
                full-list path used by ``CombinedJson``).
            offset: Rows to skip before collecting ``limit`` results.
            search_term: Case-insensitive substring filter on ``name``.
                ``%`` and ``_`` in the input are escaped so they match
                literally rather than as LIKE wildcards.
            sort_field: Column to sort by. Unknown values fall back to
                ``date``. Resolved against ``sources_table.c`` so the
                column identity is bound by SQLAlchemy — user input never
                reaches the emitted SQL as a string.
            sort_order: ``"asc"`` or ``"desc"``; anything else is treated
                as ``"desc"``.

        Returns:
            A list of source rows as plain dicts (via ``row_to_dict``),
            each carrying a derived ``ingest_status`` (``failed`` /
            ``processing`` / ``None``) from the joined progress row.
        """
        column_name = sort_field if sort_field in _SORTABLE_COLUMNS else "date"
        sort_column = sources_table.c[column_name]
        ascending = sort_order.lower() == "asc"

        stmt = (
            select(sources_table, _ingest_status_case())
            .select_from(
                sources_table.outerjoin(
                    ingest_chunk_progress_table,
                    ingest_chunk_progress_table.c.source_id
                    == sources_table.c.id,
                )
            )
            .where(_owned_or_shared_scope(user_id, extra_ids))
        )
        if search_term:
            stmt = stmt.where(
                sources_table.c.name.ilike(
                    f"%{_escape_like(search_term)}%",
                    escape="\\",
                )
            )

        # ``id`` is appended as a stable tiebreaker so paginated windows
        # are deterministic across equal sort keys.
        id_column = sources_table.c.id
        if ascending:
            stmt = stmt.order_by(sort_column.asc(), id_column.asc())
        else:
            stmt = stmt.order_by(sort_column.desc(), id_column.desc())

        if limit is not None:
            stmt = stmt.limit(limit).offset(offset)

        result = self._conn.execute(stmt)
        return [row_to_dict(r) for r in result.fetchall()]

    def count_for_user(
        self,
        user_id: str,
        *,
        search_term: Optional[str] = None,
        extra_ids: Optional[list] = None,
    ) -> int:
        """Return the count of rows that ``list_for_user`` would produce.

        The filter mirrors ``list_for_user`` exactly so ``total`` and the
        paginated window stay consistent page-to-page.

        Args:
            user_id: Scope rows to this owner.
            search_term: Same substring filter semantics as
                ``list_for_user``; ``None``/empty disables the filter.

        Returns:
            The total number of matching rows.
        """
        stmt = (
            select(func.count())
            .select_from(sources_table)
            .where(_owned_or_shared_scope(user_id, extra_ids))
        )
        if search_term:
            stmt = stmt.where(
                sources_table.c.name.ilike(
                    f"%{_escape_like(search_term)}%",
                    escape="\\",
                )
            )
        result = self._conn.execute(stmt)
        row = result.fetchone()
        return int(row[0]) if row is not None else 0

    def update(self, source_id: str, user_id: str, fields: dict) -> None:
        filtered = {k: v for k, v in fields.items() if k in _ALLOWED_COLUMNS}
        if not filtered:
            return

        values: dict = {}
        for col, val in filtered.items():
            values[col] = _coerce_jsonb(val) if col in _JSONB_COLUMNS else val
        values["updated_at"] = func.now()

        t = sources_table
        stmt = (
            t.update()
            .where(t.c.id == source_id)
            .where(t.c.user_id == user_id)
            .values(**values)
        )
        self._conn.execute(stmt)

    def set_wiki_outside_edits(self, source_id: str, user_id: str, allowed: bool) -> bool:
        """Record whether API-key and widget runs may edit this wiki.

        Kept out of :meth:`update`'s columns so no route that forwards a
        request body can change it; only the owner-checked wiki settings
        route calls this.

        Args:
            source_id: The source's UUID.
            user_id: The owner's id; the row is scoped to it.
            allowed: The new value.

        Returns:
            bool: Whether a row was updated.
        """
        t = sources_table
        result = self._conn.execute(
            t.update()
            .where(t.c.id == source_id)
            .where(t.c.user_id == user_id)
            .values(wiki_outside_edits=bool(allowed), updated_at=func.now())
        )
        return result.rowcount > 0

    def get_by_legacy_id(
        self, legacy_mongo_id: str, user_id: Optional[str] = None,
    ) -> Optional[dict]:
        legacy_mongo_id = str(legacy_mongo_id) if legacy_mongo_id is not None else None
        sql = "SELECT * FROM sources WHERE legacy_mongo_id = :legacy_id"
        params: dict[str, str] = {"legacy_id": legacy_mongo_id}
        if user_id is not None:
            sql += " AND user_id = :user_id"
            params["user_id"] = user_id
        result = self._conn.execute(text(sql), params)
        row = result.fetchone()
        return row_to_dict(row) if row is not None else None

    def update_by_legacy_id(
        self, legacy_mongo_id: str, user_id: str, fields: dict,
    ) -> bool:
        """Update a source addressed by the Mongo ObjectId string.

        Used by dual_write call sites that hold the Mongo ``_id`` but
        haven't resolved the PG UUID yet. Returns ``True`` if a row was
        updated (i.e. the legacy id was found).
        """
        legacy_mongo_id = str(legacy_mongo_id) if legacy_mongo_id is not None else None
        row = self.get_by_legacy_id(legacy_mongo_id, user_id)
        if row is None:
            return False
        self.update(str(row["id"]), user_id, fields)
        return True

    def delete_by_legacy_id(self, legacy_mongo_id: str, user_id: str) -> bool:
        """Delete by Mongo ObjectId. Used by dual_write in DeleteOldIndexes."""
        legacy_mongo_id = str(legacy_mongo_id) if legacy_mongo_id is not None else None
        result = self._conn.execute(
            text(
                "DELETE FROM sources "
                "WHERE legacy_mongo_id = :legacy_id AND user_id = :user_id"
            ),
            {"legacy_id": legacy_mongo_id, "user_id": user_id},
        )
        return result.rowcount > 0

    def delete(self, source_id: str, user_id: str) -> bool:
        result = self._conn.execute(
            text("DELETE FROM sources WHERE id = CAST(:id AS uuid) AND user_id = :user_id"),
            {"id": source_id, "user_id": user_id},
        )
        return result.rowcount > 0