Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 18 additions & 2 deletions apps/api/alembic/env.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,10 @@ def run_migrations_offline() -> None:
include_object=include_object,
literal_binds=True,
dialect_opts={"paramstyle": "named"},
# Keep each migration in its own transaction so migrations that
# require an autocommit block (for example CREATE INDEX
# CONCURRENTLY) can safely commit only their own predecessor.
transaction_per_migration=True,
# Pass through configured SSL connect args.
connect_args=ssl_connect_args,
)
Expand All @@ -119,13 +123,25 @@ def run_migrations_online() -> None:
configured_connection = config.attributes.get("connection")

def run_with_connection(connection: Connection) -> None:
caller_owned_transaction = connection.in_transaction()
if settings.API_STANDALONE_MODE_ENABLED:
ensure_better_auth_user_table(connection)
# The standalone bootstrap query starts SQLAlchemy's implicit
# transaction before Alembic begins tracking migration
# transactions. End only that transaction; never commit a
# transaction supplied by the caller.
if not caller_owned_transaction:
connection.commit()

context.configure(
connection=connection,
target_metadata=target_metadata,
include_object=include_object,
# Required for migrations that use autocommit_block().
transaction_per_migration=True,
# Concurrent DDL cannot run inside a transaction owned by the
# caller. Migrations use regular DDL for that compatibility path.
knowhere_external_transaction=caller_owned_transaction,
)

with context.begin_transaction():
Expand All @@ -136,7 +152,7 @@ def run_with_connection(connection: Connection) -> None:
return

if isinstance(configured_connection, Engine):
with configured_connection.begin() as connection:
with configured_connection.connect() as connection:
run_with_connection(connection)
return

Expand All @@ -149,7 +165,7 @@ def run_with_connection(connection: Connection) -> None:
connect_args=ssl_connect_args,
)

with connectable.begin() as connection:
with connectable.connect() as connection:
run_with_connection(connection)


Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
"""Add the chunk snapshot pagination order index.

Revision ID: fbf0c1d2e3f4
Revises: f0d85d209e68, fbe1c2d3e4f5
Create Date: 2026-08-27 00:00:00.000000
"""

from __future__ import annotations

from typing import Sequence, Union

from alembic import op
from sqlalchemy import text


revision: str = "fbf0c1d2e3f4"
down_revision: Union[str, Sequence[str], None] = (
"f0d85d209e68",
"fbe1c2d3e4f5",
)
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

_INDEX_NAME = "idx_document_chunks_revision_snapshot_order"
_INDEX_COLUMNS = "(document_id, job_result_id, sort_order, chunk_id, id)"
_INDEX_METHOD_AND_COLUMNS = f"using btree {_INDEX_COLUMNS}"


def _read_index_definition() -> tuple[str, bool] | None:
row = op.get_bind().execute(
text(
"""
SELECT indexdef, index_metadata.indisvalid
FROM pg_indexes AS indexes
JOIN pg_class AS index_class
ON index_class.relname = indexes.indexname
JOIN pg_namespace AS index_namespace
ON index_namespace.oid = index_class.relnamespace
AND index_namespace.nspname = indexes.schemaname
JOIN pg_index AS index_metadata
ON index_metadata.indexrelid = index_class.oid
WHERE indexes.schemaname = current_schema()
AND indexes.tablename = 'document_chunks'
AND indexes.indexname = :index_name
"""
),
{"index_name": _INDEX_NAME},
).first()
if row is None:
return None
return str(row[0]), bool(row[1])


def _index_is_intended() -> bool:
definition = _read_index_definition()
if definition is None or not definition[1]:
return False
normalized_definition = " ".join(definition[0].lower().split())
return (
normalized_definition.startswith("create index ")
and _INDEX_METHOD_AND_COLUMNS in normalized_definition
and " where " not in normalized_definition
)


def _drop_index(*, concurrently: bool) -> None:
concurrent_clause = "CONCURRENTLY " if concurrently else ""
op.execute(f"DROP INDEX {concurrent_clause}IF EXISTS {_INDEX_NAME}")


def _create_index(*, concurrently: bool) -> None:
concurrent_clause = "CONCURRENTLY " if concurrently else ""
op.execute(
f"""
CREATE INDEX {concurrent_clause}IF NOT EXISTS {_INDEX_NAME}
ON document_chunks {_INDEX_COLUMNS}
"""
)


def upgrade() -> None:
external_transaction = bool(
op.get_context().opts.get("knowhere_external_transaction", False)
)
if external_transaction:
if _read_index_definition() is not None and not _index_is_intended():
_drop_index(concurrently=False)
if _read_index_definition() is None:
_create_index(concurrently=False)
return

# Index creation must not hold a write lock on document_chunks while the
# production corpus is being indexed. CONCURRENTLY cannot run inside the
# transaction Alembic normally opens, so switch to an autocommit block.
with op.get_context().autocommit_block():
if _read_index_definition() is not None and not _index_is_intended():
_drop_index(concurrently=True)
if _read_index_definition() is None:
_create_index(concurrently=True)


def downgrade() -> None:
external_transaction = bool(
op.get_context().opts.get("knowhere_external_transaction", False)
)
if external_transaction:
_drop_index(concurrently=False)
else:
with op.get_context().autocommit_block():
_drop_index(concurrently=True)
18 changes: 14 additions & 4 deletions apps/api/app/core/exception_handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@
"""

import uuid
from collections.abc import Awaitable, Callable
from collections.abc import Awaitable, Callable, Mapping
from typing import List, cast

from fastapi import FastAPI, HTTPException, Request
Expand Down Expand Up @@ -88,7 +88,10 @@ def _get_request_id(request: Request) -> str:


async def knowhere_exception_handler(
request: Request, exc: KnowhereException
request: Request,
exc: KnowhereException,
*,
response_headers: Mapping[str, str] | None = None,
) -> JSONResponse:
"""
This handler enforces the separation between:
Expand Down Expand Up @@ -119,7 +122,10 @@ async def knowhere_exception_handler(
exc.logging(request_id=request_id)

# Always include request ID header for client-side correlation
headers = {"X-Request-ID": request_id}
# Preserve framework-provided headers such as ``Allow: POST`` for 405
# responses and ``WWW-Authenticate`` for authentication challenges.
headers = dict(response_headers or {})
headers["X-Request-ID"] = request_id
retry_after = exc.details.get("retry_after")
if retry_after:
headers["Retry-After"] = str(retry_after)
Expand Down Expand Up @@ -208,7 +214,11 @@ async def http_exception_handler(request: Request, exc: HTTPException) -> JSONRe
)

# Delegate to central handler
return await knowhere_exception_handler(request, knowhere_exc)
return await knowhere_exception_handler(
request,
knowhere_exc,
response_headers=exc.headers,
)


async def validation_exception_handler(
Expand Down
55 changes: 55 additions & 0 deletions apps/api/tests/contract/test_exception_handlers_contract.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
from __future__ import annotations

import sys
from pathlib import Path
from typing import cast

from fastapi import FastAPI
from httpx import ASGITransport, AsyncClient

from tests.support.import_environment import (
configure_import_environment,
ensure_import_paths,
)

configure_import_environment()
ensure_import_paths()


def _prepare_api_app_imports() -> None:
api_root = str(Path(__file__).resolve().parents[2])
if api_root in sys.path:
sys.path.remove(api_root)
sys.path.insert(0, api_root)


def _create_post_only_app() -> FastAPI:
_prepare_api_app_imports()

from app.core.exception_handlers import setup_exception_handlers

app = FastAPI()

@app.post("/v2/retrieval/query")
async def query_retrieval() -> dict[str, bool]:
return {"ok": True}

setup_exception_handlers(app)
return app


async def test_get_to_post_only_route_returns_method_not_allowed() -> None:
app = _create_post_only_app()
transport = ASGITransport(app=app)

async with AsyncClient(transport=transport, base_url="http://test") as client:
response = await client.get("/v2/retrieval/query")

assert response.status_code == 405
assert response.headers["allow"] == "POST"

response_json = cast(dict[str, object], response.json())
error = cast(dict[str, object], response_json["error"])
assert response_json["success"] is False
assert error["code"] == "METHOD_NOT_ALLOWED"
assert error["message"] == "Method not allowed"
Loading
Loading