vecshift 0.1.0__py3-none-any.whl

vecshift/cli_plan.py

@@ -0,0 +1,335 @@
+"""The ``vecshift init`` and ``vecshift plan`` commands."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import shutil
+import sys
+import textwrap
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Annotated
+
+import typer
+
+from vecshift.cli_style import finding_lines
+
+if TYPE_CHECKING:
+    from vecshift.connectors.pgvector import ConnectionSettings
+    from vecshift.connectors.pgvector.target import TargetState
+    from vecshift.doctor import IndexProfile
+    from vecshift.embeddings import ModelSpec
+    from vecshift.jobs import JobSpec
+    from vecshift.planning import Plan, ProbeResult
+
+DEFAULT_JOB = Path("vecshift.yaml")
+PROBE_DOCUMENTS = 16
+
+
+def _fail(message: str, hint: str | None = None) -> typer.Exit:
+    typer.secho(message, err=True, fg=typer.colors.RED)
+    if hint:
+        typer.echo(f"→ {hint}", err=True)
+    return typer.Exit(2)
+
+
+def init(
+    table: Annotated[
+        str, typer.Option(help="Table holding the vectors, as table or schema.table.")
+    ],
+    model: Annotated[
+        str, typer.Option(help="The new model, e.g. openai/text-embedding-3-large,dims=1024.")
+    ],
+    column: Annotated[str, typer.Option(help="Column for the new vectors.")] = "embedding_v2",
+    output: Annotated[Path, typer.Option("--output", "-o", help="Where to write the job.")] = (
+        DEFAULT_JOB
+    ),
+    force: Annotated[bool, typer.Option(help="Overwrite an existing file.")] = False,
+) -> None:
+    """Write a job file to start a migration from."""
+    import yaml
+    from pydantic import ValidationError
+
+    from vecshift.jobs.spec import JobSpec, format_errors, template
+
+    text = template(table, model, column)
+    try:
+        JobSpec.model_validate(yaml.safe_load(text))
+    except ValidationError as exc:
+        raise _fail(f"Those settings aren't valid:\n{format_errors(exc)}") from exc
+    if output.exists() and not force:
+        raise _fail(f"{output} already exists.", "Pass --force to overwrite it.")
+    try:
+        output.write_text(text, encoding="utf-8")
+    except OSError as exc:
+        raise _fail(f"Couldn't write {output}: {exc.strerror}") from exc
+    typer.echo(f"Wrote {output}. Review it, then run: vecshift plan {output}")
+
+
+def _size(n: int | None) -> str:
+    if n is None:
+        return "—"
+    for unit, size in (("TB", 1024**4), ("GB", 1024**3), ("MB", 1024**2), ("KB", 1024)):
+        if n >= size:
+            return f"{n / size:,.1f} {unit}"
+    return f"{n} B"
+
+
+def _duration(seconds: float | None) -> str:
+    if seconds is None:
+        return "unknown"
+    if seconds < 90:
+        return f"~{max(1, round(seconds))} s"
+    if seconds < 90 * 60:
+        return f"~{round(seconds / 60)} min"
+    if seconds < 48 * 3600:
+        return f"~{seconds / 3600:.1f} h"
+    return f"~{seconds / 86400:.1f} days"
+
+
+def _compact(n: int | None) -> str:
+    if n is None:
+        return "—"
+    for size, suffix in ((1_000_000_000, "B"), (1_000_000, "M"), (1_000, "K")):
+        if n >= size * 10:
+            return f"{n / size:,.1f}".rstrip("0").rstrip(".") + suffix
+    return f"{n:,}"
+
+
+def _note(text: str) -> str:
+    width = min(100, shutil.get_terminal_size((100, 24)).columns)
+    return textwrap.fill(text, width=width, initial_indent=" " * 6, subsequent_indent=" " * 6)
+
+
+def _money(value: float | None) -> str:
+    if value is None:
+        return "unknown"
+    if value == 0:
+        return "$0.00"
+    return "<$0.01" if value < 0.01 else f"~${value:,.2f}"
+
+
+def _render(plan: Plan) -> None:
+    e = plan.estimates
+    dims = f"{plan.dimensions:,}" if plan.dimensions else "?"
+    source_note = {
+        "spec": "set in the spec",
+        "known": "this model's size",
+        "probe": "measured",
+        None: "unknown",
+    }[plan.dimensions_source]
+    typer.secho(f"vecshift plan · {plan.job}", bold=True)
+    typer.echo(f"Source  {plan.source}")
+    typer.echo(f"Model   {plan.model} → {dims} dimensions ({source_note})")
+
+    typer.secho("\nChanges", bold=True)
+    marks = {"add_column": "+", "trigger": "+", "embed": "~", "index": "+", "cutover": "⇄"}
+    for change in plan.changes:
+        typer.echo(f"  {marks.get(change.kind, '·')} {change.summary}")
+        if change.sql:
+            typer.secho(f"      {change.sql}", fg=typer.colors.CYAN)
+        if change.note:
+            typer.secho(_note(change.note), dim=True)
+
+    typer.secho("\nEstimates", bold=True)
+    rows = "—" if e.rows is None else (f"{e.rows:,}" if e.rows_exact else f"~{e.rows:,}")
+    lines = [
+        ("Rows to embed", rows, "exact" if e.rows_exact else "from table statistics"),
+        ("Tokens", "~" + _compact(e.tokens) if e.tokens else "—", e.tokens_method or ""),
+        (
+            "Cost",
+            _money(e.cost_usd),
+            "",
+        ),
+        ("Requests", f"{e.requests:,}" if e.requests is not None else "—", ""),
+        ("Duration", _duration(e.seconds), e.seconds_method or "add a rate limit or --probe"),
+        (
+            "New column",
+            _size(e.new_bytes),
+            f"old column {_size(e.old_bytes)}" if e.old_bytes else "",
+        ),
+    ]
+    if e.index_memory_bytes:
+        lines.append(
+            (
+                "Index build",
+                f"~{_size(e.index_memory_bytes)}",
+                f"maintenance_work_mem is {_size(e.maintenance_work_mem)}",
+            )
+        )
+    width = max(len(k) for k, _, _ in lines)
+    for key, value, note in lines:
+        typer.echo(f"  {key:<{width}}  {value:<10}  " + typer.style(note, dim=True))
+
+    if plan.findings:
+        typer.secho("\nChecks", bold=True)
+        finding_lines(plan.findings)
+    typer.echo()
+    if plan.ok:
+        typer.secho("Ready to apply.", fg=typer.colors.GREEN, bold=True)
+    else:
+        noun = "error" if plan.errors == 1 else "errors"
+        typer.secho(f"Fix {plan.errors} {noun} before applying.", fg=typer.colors.RED, bold=True)
+
+
+async def _probe(spec: ModelSpec, texts: list[str]) -> ProbeResult:
+    from vecshift.embeddings import create_embedder
+    from vecshift.planning import ProbeResult
+
+    embedder = create_embedder(spec)
+    try:
+        start = time.perf_counter()
+        vectors = await embedder.embed(texts, "document")
+        elapsed = time.perf_counter() - start
+    finally:
+        await embedder.aclose()
+    chars = sum(len(spec.doc_prefix) + len(t) for t in texts)
+    return ProbeResult(
+        documents=len(texts),
+        dimensions=len(vectors[0]),
+        tokens_per_char=embedder.tokens / chars if chars else 0.25,
+        docs_per_second=len(texts) / elapsed if elapsed > 0 else 0.0,
+    )
+
+
+@dataclass(slots=True)
+class Prepared:
+    """Everything known about a job before anything changes."""
+
+    job: JobSpec
+    settings: ConnectionSettings
+    spec: ModelSpec
+    profile: IndexProfile
+    target: TargetState
+    plan: Plan
+
+
+def load(
+    job_file: Path, dsn_env: str | None = None
+) -> tuple[JobSpec, ConnectionSettings, ModelSpec]:
+    """The job, its connection settings (from the environment), and its model.
+
+    ``dsn_env`` reads the connection string from another variable, such as a replica's.
+    """
+    from vecshift.connectors import pgvector
+    from vecshift.embeddings import parse_spec
+    from vecshift.jobs import JobError, load_job
+
+    try:
+        job = load_job(job_file)
+    except JobError as exc:
+        hint = (
+            "Create one with: vecshift init --table ... --model ..."
+            if not job_file.exists()
+            else None
+        )
+        raise _fail(str(exc), hint) from exc
+    variable = dsn_env or job.source.dsn_env
+    dsn = os.environ.get(variable)
+    if not dsn:
+        raise _fail(
+            f"{variable} isn't set.",
+            f"Export the database connection string as {variable}.",
+        )
+    try:
+        settings = pgvector.prepare(dsn)
+    except pgvector.ConnectError as exc:
+        raise _fail(f"Couldn't connect: {exc}", exc.hint) from exc
+    return job, settings, parse_spec(job.model)
+
+
+def prepare(
+    job_file: Path, *, sample_size: int = 2000, probe: bool = False, yes: bool = False
+) -> Prepared:
+    """Load the job, inspect the database read-only, and build the plan."""
+    from vecshift.connectors import pgvector
+    from vecshift.connectors.pgvector.documents import sample_documents
+    from vecshift.connectors.pgvector.target import inspect_target, resolve_source_column
+    from vecshift.embeddings import EmbeddingError
+    from vecshift.planning import SampleStats, build_plan
+
+    job, settings, spec = load(job_file)
+    try:
+        conn = pgvector.connect(settings)
+    except pgvector.ConnectError as exc:
+        raise _fail(f"Couldn't connect: {exc}", exc.hint) from exc
+    try:
+        source_column = resolve_source_column(
+            conn, job.source.table, job.source.vector_column, job.target.column
+        )
+        profile = pgvector.inspect(
+            conn,
+            table=job.source.table,
+            column=source_column,
+            text_column=job.source.text_column,
+            sample_size=sample_size,
+        )
+        target = inspect_target(conn, job.source.table, source_column, job.target.column)
+        texts: list[str] = []
+        if profile.text_field:
+            rows, _ = sample_documents(
+                conn,
+                job.source.table,
+                source_column,
+                profile.text_field,
+                min(sample_size, 500),
+            )
+            texts = [t for _, t in rows]
+    except pgvector.TargetSelectionError as exc:
+        raise _fail(str(exc)) from exc
+    finally:
+        conn.rollback()
+        conn.close()
+
+    stats = SampleStats(
+        documents=len(texts), average_chars=sum(map(len, texts)) / len(texts) if texts else 0.0
+    )
+    measured = None
+    if probe and texts:
+        chosen = texts[:PROBE_DOCUMENTS]
+        if not spec.is_local and not yes:
+            typer.echo(
+                f"The probe sends {len(chosen)} sample rows to {spec.url} to measure {spec.name}.",
+                err=True,
+            )
+            if not sys.stdin.isatty():
+                raise _fail("Not sending data without confirmation.", "Pass --yes to proceed.")
+            if not typer.confirm("Continue?", err=True):
+                raise typer.Exit(1)
+        try:
+            measured = asyncio.run(_probe(spec, chosen))
+        except EmbeddingError as exc:
+            raise _fail(f"The probe failed: {exc}") from exc
+
+    result = build_plan(job, profile, target, stats, measured)
+    return Prepared(job, settings, spec, profile, target, result)
+
+
+def plan(
+    job_file: Annotated[
+        Path, typer.Argument(help="The job file.", show_default=True)
+    ] = DEFAULT_JOB,
+    probe: Annotated[
+        bool,
+        typer.Option(
+            help=f"Embed {PROBE_DOCUMENTS} sample rows with the real model to measure its "
+            "size, token counts, and speed."
+        ),
+    ] = False,
+    yes: Annotated[bool, typer.Option("--yes", "-y", help="Don't ask before probing.")] = False,
+    sample_size: Annotated[
+        int, typer.Option("--sample", min=50, max=100_000, help="Rows to sample.")
+    ] = 2000,
+    output_json: Annotated[bool, typer.Option("--json", help="Print the plan as JSON.")] = False,
+) -> None:
+    """Check a migration job and estimate it, without changing anything."""
+    result = prepare(job_file, sample_size=sample_size, probe=probe, yes=yes).plan
+    if output_json:
+        typer.echo(json.dumps(result.to_dict(), indent=2))
+    else:
+        _render(result)
+    if not result.ok:
+        raise typer.Exit(1)

vecshift/cli_style.py

@@ -0,0 +1,57 @@
+"""Terminal styling shared by the CLI commands."""
+
+from __future__ import annotations
+
+import shutil
+import textwrap
+from collections.abc import Sequence
+
+import typer
+
+from vecshift.doctor import Finding, Severity
+
+STYLE = {
+    Severity.ERROR: ("✖ ERROR  ", typer.colors.RED, "error"),
+    Severity.WARNING: ("⚠ WARNING", typer.colors.YELLOW, "warning"),
+    Severity.INFO: ("● INFO   ", typer.colors.BLUE, "info"),
+    Severity.OK: ("✔ OK     ", typer.colors.GREEN, "ok"),
+}
+
+
+INDENT = " " * 11
+
+
+def wrap(text: str) -> str:
+    width = min(100, shutil.get_terminal_size((100, 24)).columns)
+    return textwrap.fill(text, width=width, initial_indent=INDENT, subsequent_indent=INDENT)
+
+
+def finding_lines(findings: Sequence[Finding]) -> None:
+    """Print findings, most severe first, in the doctor style."""
+    for finding in sorted(findings, key=lambda f: -f.severity.rank):
+        label, color, _ = STYLE[finding.severity]
+        typer.secho(f"{label}  ", fg=color, bold=True, nl=False)
+        typer.secho(finding.title.replace("`", ""), bold=True)
+        typer.echo(wrap(finding.detail.replace("`", "")))
+        if finding.hint:
+            typer.secho(wrap(f"→ {finding.hint.replace('`', '')}"), dim=True)
+
+
+def warn_if_password_on_command_line(ctx: typer.Context, dsn: str | None) -> None:
+    """Passwords given as arguments are visible to other users through ps and shell history."""
+    source = ctx.get_parameter_source("dsn")
+    # Compared by name: typer may bundle its own copy of click's ParameterSource enum.
+    if dsn and source is not None and source.name == "COMMANDLINE":
+        from psycopg.conninfo import conninfo_to_dict
+
+        try:
+            has_password = bool(conninfo_to_dict(dsn).get("password"))
+        except Exception:  # an invalid string is reported later, where it's parsed
+            return
+        if has_password:
+            typer.secho(
+                "Warning: the --dsn password is visible to other users of this machine and "
+                "saved in shell history. Prefer the VECSHIFT_DSN environment variable.",
+                err=True,
+                fg=typer.colors.YELLOW,
+            )

vecshift/connectors/__init__.py

@@ -0,0 +1 @@
+"""Connectors for vector stores."""

vecshift/connectors/pgvector/__init__.py

@@ -0,0 +1,29 @@
+"""PostgreSQL with the pgvector extension, including Supabase."""
+
+from vecshift.connectors.pgvector.connection import (
+    ConnectError,
+    ConnectionSettings,
+    SupabaseMode,
+    connect,
+    connect_writer,
+    prepare,
+)
+from vecshift.connectors.pgvector.inspect import (
+    TargetSelectionError,
+    VectorColumn,
+    find_vector_columns,
+    inspect,
+)
+
+__all__ = [
+    "ConnectError",
+    "ConnectionSettings",
+    "SupabaseMode",
+    "TargetSelectionError",
+    "VectorColumn",
+    "connect",
+    "connect_writer",
+    "find_vector_columns",
+    "inspect",
+    "prepare",
+]

vecshift/connectors/pgvector/connection.py

@@ -0,0 +1,155 @@
+"""Connecting to PostgreSQL safely, with Supabase's connection modes in mind.
+
+Supabase offers a direct connection and Supavisor poolers. The ones that matter here:
+
+- ``db.<ref>.supabase.co:5432``: direct. IPv6 only unless the project has the IPv4 add-on.
+- ``*.pooler.supabase.com:5432``: session pooler. Works over IPv4.
+- ``*.pooler.supabase.com:6543`` and ``db.<ref>.supabase.co:6543``: transaction poolers,
+  which don't support prepared statements.
+
+VecShift always disables prepared statements and runs inside one read-only transaction,
+so every mode works.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import StrEnum
+
+import psycopg
+from psycopg import sql
+from psycopg.conninfo import conninfo_to_dict, make_conninfo
+
+APPLICATION_NAME = "vecshift"
+_TRANSACTION_POOLER_PORT = "6543"
+
+
+class ConnectError(Exception):
+    """A connection failed. The message is safe to show: it never includes a password."""
+
+    def __init__(self, message: str, hint: str | None = None) -> None:
+        super().__init__(message)
+        self.hint = hint
+
+
+class SupabaseMode(StrEnum):
+    DIRECT = "direct"
+    SESSION_POOLER = "session pooler"
+    TRANSACTION_POOLER = "transaction pooler"
+
+
+@dataclass(frozen=True, slots=True)
+class ConnectionSettings:
+    """A validated connection string plus what we learned about where it points."""
+
+    conninfo: str
+    display: str
+    """The connection string with the password removed, safe to print."""
+    supabase: SupabaseMode | None
+
+    @property
+    def description(self) -> str:
+        return f"Supabase {self.supabase}" if self.supabase else "PostgreSQL"
+
+
+def _supabase_mode(host: str, port: str) -> SupabaseMode | None:
+    host = host.lower().rstrip(".")
+    if host.endswith(".pooler.supabase.com"):
+        if port == _TRANSACTION_POOLER_PORT:
+            return SupabaseMode.TRANSACTION_POOLER
+        return SupabaseMode.SESSION_POOLER
+    if host.endswith(".supabase.co"):
+        if port == _TRANSACTION_POOLER_PORT:
+            return SupabaseMode.TRANSACTION_POOLER
+        return SupabaseMode.DIRECT
+    return None
+
+
+def prepare(dsn: str) -> ConnectionSettings:
+    """Parse and harden a connection string (a ``postgresql://`` URL or ``key=value`` pairs).
+
+    For Supabase hosts, TLS is required unless the string already sets ``sslmode``.
+    """
+    try:
+        params = conninfo_to_dict(dsn)
+    except psycopg.ProgrammingError as exc:
+        raise ConnectError(
+            "That doesn't look like a valid PostgreSQL connection string.",
+            hint="Expected something like postgresql://user:password@host:5432/dbname",
+        ) from exc
+
+    host = str(params.get("host") or "")
+    port = str(params.get("port") or "5432")
+    mode = _supabase_mode(host, port)
+    if mode is not None:
+        params.setdefault("sslmode", "require")
+    params.setdefault("application_name", APPLICATION_NAME)
+    params.setdefault("connect_timeout", "10")
+
+    shown = {k: v for k, v in params.items() if k == "user" or k == "dbname"}
+    user = f"{shown['user']}@" if shown.get("user") else ""
+    display = f"postgresql://{user}{host or 'localhost'}:{port}/{shown.get('dbname', '')}"
+
+    return ConnectionSettings(conninfo=make_conninfo("", **params), display=display, supabase=mode)
+
+
+def _hint_for(settings: ConnectionSettings, error: str) -> str | None:
+    lowered = error.lower()
+    if settings.supabase is SupabaseMode.DIRECT and any(
+        s in lowered
+        for s in ("network is unreachable", "cannot assign requested address", "translate host")
+    ):
+        return (
+            "Supabase's direct connection is IPv6 only unless your project has the IPv4 "
+            "add-on. Use the session pooler string from the dashboard's Connect button "
+            "instead (host ending in .pooler.supabase.com, port 5432)."
+        )
+    if "password authentication failed" in lowered and settings.supabase in (
+        SupabaseMode.SESSION_POOLER,
+        SupabaseMode.TRANSACTION_POOLER,
+    ):
+        return (
+            "Supabase poolers expect the user as postgres.<project-ref>, not just postgres. "
+            "Copy the exact string from the dashboard's Connect button."
+        )
+    if "timeout" in lowered:
+        return "Check the host and port, and that your network allows outbound connections."
+    return None
+
+
+def connect(settings: ConnectionSettings, statement_timeout_s: int = 60) -> psycopg.Connection:
+    """Open a read-only connection with a statement timeout.
+
+    Prepared statements are disabled so transaction poolers work. The caller should
+    roll back and close the connection when done.
+    """
+    try:
+        conn = psycopg.connect(settings.conninfo, prepare_threshold=None)
+    except psycopg.OperationalError as exc:
+        message = str(exc).strip().splitlines()[0] if str(exc).strip() else "Connection failed."
+        raise ConnectError(message, hint=_hint_for(settings, str(exc))) from exc
+
+    conn.read_only = True
+    # SET LOCAL lasts for the current transaction only, which also suits transaction poolers.
+    conn.execute(
+        sql.SQL("SET LOCAL statement_timeout = {}").format(sql.Literal(statement_timeout_s * 1000))
+    )
+    return conn
+
+
+def connect_writer(settings: ConnectionSettings) -> psycopg.Connection:
+    """An autocommit connection for apply, which manages its own short transactions.
+
+    Apply holds a session-level advisory lock and builds indexes concurrently, so it needs a
+    real session: a transaction pooler won't do.
+    """
+    if settings.supabase is SupabaseMode.TRANSACTION_POOLER:
+        raise ConnectError(
+            "apply can't run through a transaction pooler.",
+            hint="Use the session pooler (port 5432) or the direct connection string.",
+        )
+    try:
+        return psycopg.connect(settings.conninfo, prepare_threshold=None, autocommit=True)
+    except psycopg.OperationalError as exc:
+        message = str(exc).strip().splitlines()[0] if str(exc).strip() else "Connection failed."
+        raise ConnectError(message, hint=_hint_for(settings, str(exc))) from exc

vecshift/connectors/pgvector/documents.py

@@ -0,0 +1,96 @@
+"""Read a sample of documents (ID and text) from a pgvector table, for benchmarks."""
+
+from __future__ import annotations
+
+from collections.abc import Collection
+
+import psycopg
+from psycopg import sql
+
+from vecshift.connectors.pgvector.inspect import (
+    TEXT_COLUMNS,
+    TEXT_TYPES,
+    TargetSelectionError,
+    _columns,
+    _pick,
+    find_vector_columns,
+    select_column,
+)
+
+
+def _id_column(conn: psycopg.Connection, relid: int) -> str | None:
+    row = conn.execute(
+        """
+        SELECT a.attname FROM pg_index i
+        JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = i.indkey[0]
+        WHERE i.indrelid = %s AND i.indisprimary AND i.indnatts = 1
+        """,
+        (relid,),
+    ).fetchone()
+    return str(row[0]) if row else None
+
+
+def sample_documents(
+    conn: psycopg.Connection,
+    table: str | None,
+    column: str | None,
+    text_column: str | None,
+    size: int,
+    include_ids: Collection[str] = (),
+) -> tuple[list[tuple[str, str]], str]:
+    """Up to ``size`` (id, text) rows spread across the table, plus any ``include_ids``.
+
+    Returns the rows and the table's qualified name.
+    """
+    col = select_column(find_vector_columns(conn), table, column)
+    columns = _columns(conn, col.relid)
+    text = text_column or _pick(columns, TEXT_COLUMNS, TEXT_TYPES)
+    if text is None or text not in columns:
+        raise TargetSelectionError(
+            f"No text column found in {col.schema}.{col.table}. Pass --text-column."
+        )
+    id_col = _id_column(conn, col.relid)
+    id_expr = sql.Identifier(id_col) if id_col else sql.SQL("ctid")
+    if include_ids and not id_col:
+        raise TargetSelectionError(
+            "Labeled queries need a single-column primary key to match document IDs."
+        )
+    table_id = sql.Identifier(col.schema, col.table)
+    text_id = sql.Identifier(text)
+    nonempty = sql.SQL("{t} IS NOT NULL AND btrim({t}::text) <> ''").format(t=text_id)
+
+    est_row = conn.execute(
+        "SELECT reltuples::bigint FROM pg_class WHERE oid = %s", (col.relid,)
+    ).fetchone()
+    estimated = int(est_row[0]) if est_row and est_row[0] > 0 else 0
+    tablesample: sql.Composable = sql.SQL("")
+    if estimated > size * 2:
+        percent = min(100.0, 100.0 * size * 1.5 / estimated)
+        tablesample = sql.SQL("TABLESAMPLE SYSTEM ({}) REPEATABLE (7)").format(sql.Literal(percent))
+
+    rows = conn.execute(
+        sql.SQL(
+            "SELECT {id}::text, {t}::text FROM {table} {ts} WHERE {nonempty} "
+            "ORDER BY md5(ctid::text) LIMIT {n}"
+        ).format(
+            id=id_expr,
+            t=text_id,
+            table=table_id,
+            ts=tablesample,
+            nonempty=nonempty,
+            n=sql.Literal(size),
+        )
+    ).fetchall()
+    docs: dict[str, str] = {r[0]: r[1] for r in rows}
+
+    wanted = [i for i in include_ids if i not in docs]
+    if wanted:
+        extra = conn.execute(
+            sql.SQL("SELECT {id}::text, {t}::text FROM {table} WHERE {id}::text = ANY(%s)").format(
+                id=id_expr, t=text_id, table=table_id
+            ),
+            (wanted,),
+        ).fetchall()
+        for doc_id, body in extra:
+            docs[doc_id] = body or ""
+    return list(docs.items()), f"{col.schema}.{col.table}"