Close the source adapter one-shot ingestion builds
create_document_from_source resolves a fetcher per call and never closed it, so every one-shot URL or WebDAV ingest leaked an httpx connection pool. Close it, but only when we built it: resolve_adhoc_fetcher returns a caller-supplied source when one matches the URI, and the ingester keeps those open across jobs. Directory ingestion also yielded symlinked files resolving outside the directory it was given. rglob does not recurse into symlinked directories, so a symlinked file was the only way out of the tree; skip those, as FSSource.discover already does. The chunk repository docstring named client._ensure_chunks_embedded, which does not exist.
This commit is contained in:
Yiorgis Gozadinos
parent
895b5f0532
commit
73944f91ea
7 changed files with 350 additions and 63 deletions
|
|
@ -25,6 +25,8 @@
|
|||
|
||||
### Fixed
|
||||
|
||||
- `create_document_from_source` closes the source adapter it builds for the call; adapters passed in through `sources` are left to their owner.
|
||||
- Directory ingestion skips symlinked files resolving outside the given directory, matching `FSSource.discover`.
|
||||
- A FULL rebuild no longer deletes a source-backed document before re-ingesting it: it refreshes the document in place, so the document id is preserved and a failed fetch or conversion falls back to rebuilding from stored content instead of losing the document.
|
||||
- `DocumentRepository.delete_all` recreated `document_items` from `DocumentItemRecord` instead of `get_document_items_arrow_schema()`, returning `picture_data` as `binary` rather than `large_binary`.
|
||||
- `server.json` runtime arguments are `mcp --stdio`, was `serve --mcp`.
|
||||
|
|
|
|||
|
|
@ -726,6 +726,14 @@ async def create_document_from_source(
|
|||
documents: list[Document] = []
|
||||
filter = FileFilter()
|
||||
for child in local_path.rglob("*"):
|
||||
# rglob does not recurse into symlinked directories, but it does
|
||||
# yield symlinked files. Skip the ones resolving outside the
|
||||
# directory the caller named, as FSSource.discover does.
|
||||
if child.is_symlink():
|
||||
resolved = child.resolve(strict=False)
|
||||
if not resolved.is_relative_to(local_path.resolve()):
|
||||
continue
|
||||
child = resolved
|
||||
if child.is_file() and filter.include_file(str(child)):
|
||||
doc = await create_document_from_source(
|
||||
client,
|
||||
|
|
@ -768,80 +776,91 @@ async def create_document_from_source(
|
|||
source_str, sources=sources, storage_options=storage_options
|
||||
)
|
||||
|
||||
# The stored URI is what we look up + persist by. For an explicit uri
|
||||
# override, use it as-is. For a file:// input the source string is
|
||||
# already canonical (URL-encoded); round-tripping via Path.as_uri()
|
||||
# would double-encode any escapes like %5B. For bare paths, canonicalize.
|
||||
if uri is not None:
|
||||
stored_uri = uri
|
||||
elif parsed_url.scheme == "file":
|
||||
stored_uri = source_str
|
||||
elif parsed_url.scheme == "":
|
||||
stored_uri = Path(source_str).absolute().as_uri()
|
||||
else:
|
||||
stored_uri = source_str
|
||||
# A fetcher built for this call holds its own httpx pool (HTTP, WebDAV) and
|
||||
# has to be closed here. One handed in through `sources` belongs to the
|
||||
# caller: the ingester keeps its sources open across jobs.
|
||||
owns_fetcher = all(fetcher is not configured for configured in sources or ())
|
||||
|
||||
existing_doc = await client.get_document_by_uri(stored_uri)
|
||||
try:
|
||||
# The stored URI is what we look up + persist by. For an explicit uri
|
||||
# override, use it as-is. For a file:// input the source string is
|
||||
# already canonical (URL-encoded); round-tripping via Path.as_uri()
|
||||
# would double-encode any escapes like %5B. For bare paths, canonicalize.
|
||||
if uri is not None:
|
||||
stored_uri = uri
|
||||
elif parsed_url.scheme == "file":
|
||||
stored_uri = source_str
|
||||
elif parsed_url.scheme == "":
|
||||
stored_uri = Path(source_str).absolute().as_uri()
|
||||
else:
|
||||
stored_uri = source_str
|
||||
|
||||
# Cheap revision-based short-circuit: only worth a HEAD when we have a
|
||||
# stored revision to compare against. All sources persist their native
|
||||
# revision (mtime_ns for FS, ETag for S3, ETag/Last-Modified for HTTP)
|
||||
# under the canonical "source_revision" metadata key.
|
||||
stored_revision = (
|
||||
(existing_doc.metadata or {}).get("source_revision") if existing_doc else None
|
||||
)
|
||||
if existing_doc and stored_revision and not force:
|
||||
current_revision = await fetcher.head(source_str)
|
||||
if current_revision == stored_revision:
|
||||
existing_doc = await client.get_document_by_uri(stored_uri)
|
||||
|
||||
# Cheap revision-based short-circuit: only worth a HEAD when we have a
|
||||
# stored revision to compare against. All sources persist their native
|
||||
# revision (mtime_ns for FS, ETag for S3, ETag/Last-Modified for HTTP)
|
||||
# under the canonical "source_revision" metadata key.
|
||||
stored_revision = (
|
||||
(existing_doc.metadata or {}).get("source_revision")
|
||||
if existing_doc
|
||||
else None
|
||||
)
|
||||
if existing_doc and stored_revision and not force:
|
||||
current_revision = await fetcher.head(source_str)
|
||||
if current_revision == stored_revision:
|
||||
return await _refresh_doc_metadata(
|
||||
client,
|
||||
existing_doc,
|
||||
title=title,
|
||||
user_metadata=metadata,
|
||||
source_metadata=None,
|
||||
)
|
||||
|
||||
with logfire.span("document.fetch", uri=source_str) as fetch_span:
|
||||
result = await fetcher.fetch(source_str)
|
||||
fetch_span.set_attribute("bytes", len(result.body))
|
||||
fetch_span.set_attribute("content_hash", result.content_hash)
|
||||
|
||||
provider_metadata = await _provider_metadata(
|
||||
metadata_provider, source_id or fetcher.source_id, source_str, result
|
||||
)
|
||||
user_metadata = {**metadata, **provider_metadata}
|
||||
|
||||
# MD5 short-circuit: the bytes are unchanged even if the revision wasn't.
|
||||
# Refresh the source-derived metadata (revision may have rolled) but skip
|
||||
# convert/embed/store entirely.
|
||||
if (
|
||||
existing_doc
|
||||
and not force
|
||||
and existing_doc.metadata.get("md5") == result.content_hash
|
||||
):
|
||||
source_meta: dict = {
|
||||
"content_type": result.content_type,
|
||||
"md5": result.content_hash,
|
||||
**result.extra_metadata,
|
||||
}
|
||||
if result.revision is not None:
|
||||
source_meta["source_revision"] = result.revision
|
||||
return await _refresh_doc_metadata(
|
||||
client,
|
||||
existing_doc,
|
||||
title=title,
|
||||
user_metadata=metadata,
|
||||
source_metadata=None,
|
||||
user_metadata=user_metadata,
|
||||
source_metadata=source_meta,
|
||||
)
|
||||
|
||||
with logfire.span("document.fetch", uri=source_str) as fetch_span:
|
||||
result = await fetcher.fetch(source_str)
|
||||
fetch_span.set_attribute("bytes", len(result.body))
|
||||
fetch_span.set_attribute("content_hash", result.content_hash)
|
||||
|
||||
provider_metadata = await _provider_metadata(
|
||||
metadata_provider, source_id or fetcher.source_id, source_str, result
|
||||
)
|
||||
user_metadata = {**metadata, **provider_metadata}
|
||||
|
||||
# MD5 short-circuit: the bytes are unchanged even if the revision wasn't.
|
||||
# Refresh the source-derived metadata (revision may have rolled) but skip
|
||||
# convert/embed/store entirely.
|
||||
if (
|
||||
existing_doc
|
||||
and not force
|
||||
and existing_doc.metadata.get("md5") == result.content_hash
|
||||
):
|
||||
source_meta: dict = {
|
||||
"content_type": result.content_type,
|
||||
"md5": result.content_hash,
|
||||
**result.extra_metadata,
|
||||
}
|
||||
if result.revision is not None:
|
||||
source_meta["source_revision"] = result.revision
|
||||
return await _refresh_doc_metadata(
|
||||
return await _ingest_fetch_result(
|
||||
client,
|
||||
existing_doc,
|
||||
result,
|
||||
title=title,
|
||||
user_metadata=user_metadata,
|
||||
source_metadata=source_meta,
|
||||
stored_uri=stored_uri,
|
||||
existing_doc=existing_doc,
|
||||
)
|
||||
|
||||
return await _ingest_fetch_result(
|
||||
client,
|
||||
result,
|
||||
title=title,
|
||||
user_metadata=user_metadata,
|
||||
stored_uri=stored_uri,
|
||||
existing_doc=existing_doc,
|
||||
)
|
||||
finally:
|
||||
if owns_fetcher:
|
||||
await fetcher.aclose()
|
||||
|
||||
|
||||
async def update_document(
|
||||
|
|
|
|||
|
|
@ -60,7 +60,8 @@ class ChunkRepository:
|
|||
"""Create one or more chunks in the database.
|
||||
|
||||
Chunks must have embeddings set before calling this method.
|
||||
Use client._ensure_chunks_embedded() to embed chunks if needed.
|
||||
Use haiku.rag.client.processing.ensure_chunks_embedded() to embed
|
||||
chunks if needed.
|
||||
"""
|
||||
self.store._assert_writable()
|
||||
# Handle single chunk
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -19,6 +19,7 @@ from haiku.rag.client.documents import (
|
|||
)
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.embeddings import EmbedderWrapper
|
||||
from haiku.rag.ingester.sources.base import FetchResult
|
||||
from haiku.rag.store.compression import decompress_json
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
from haiku.rag.store.models.document import Document
|
||||
|
|
@ -444,6 +445,32 @@ async def test_client_create_document_from_directory(temp_db_path):
|
|||
assert not any("unsupported.xyz" in uri for uri in uris)
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_directory_ingest_skips_symlinks_escaping_the_tree(temp_db_path):
|
||||
"""A symlinked file resolving outside the named directory is not ingested;
|
||||
one resolving inside it is. Matches FSSource.discover."""
|
||||
import os
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
root = Path(temp_dir) / "tree"
|
||||
outside = Path(temp_dir) / "outside"
|
||||
root.mkdir()
|
||||
outside.mkdir()
|
||||
|
||||
(root / "real.txt").write_text("Inside the tree.")
|
||||
(outside / "secret.txt").write_text("Outside the tree.")
|
||||
os.symlink(outside / "secret.txt", root / "escape.txt")
|
||||
os.symlink(root / "real.txt", root / "inside_link.txt")
|
||||
|
||||
result = await client.create_document_from_source(root)
|
||||
|
||||
assert isinstance(result, list)
|
||||
uris = sorted(uri for doc in result if (uri := doc.uri))
|
||||
assert not any("secret" in uri or "escape" in uri for uri in uris)
|
||||
assert any("real.txt" in uri for uri in uris)
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_client_create_document_from_url(temp_db_path):
|
||||
"""Test creating a document from a URL."""
|
||||
|
|
@ -2408,6 +2435,74 @@ def test_check_source_accessible_file_uri(tmp_path):
|
|||
assert check_source_accessible((tmp_path / "gone.txt").as_uri()) is False
|
||||
|
||||
|
||||
class _CountingSource:
|
||||
"""A real Source over one in-memory document that counts its closes."""
|
||||
|
||||
def __init__(self, uri: str, body: bytes) -> None:
|
||||
self.source_id = "counting"
|
||||
self.supported_extensions = None
|
||||
self.max_file_size = None
|
||||
self._uri = uri
|
||||
self._body = body
|
||||
self.closes = 0
|
||||
|
||||
def supports(self, uri: str) -> bool:
|
||||
return uri == self._uri
|
||||
|
||||
async def head(self, uri: str) -> str | None:
|
||||
return "v1"
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.closes += 1
|
||||
|
||||
async def fetch(self, uri: str) -> "FetchResult":
|
||||
import hashlib
|
||||
|
||||
return FetchResult(
|
||||
uri=uri,
|
||||
body=self._body,
|
||||
content_type="text/markdown",
|
||||
content_hash=hashlib.md5(self._body).hexdigest(),
|
||||
revision="v1",
|
||||
)
|
||||
|
||||
def discover(self, since=None, *, known_uris=None):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_adhoc_source_is_closed_after_ingest(temp_db_path, monkeypatch):
|
||||
"""An ad-hoc fetcher is built for this one call, so this call has to close
|
||||
it — HTTP and WebDAV adapters hold an httpx connection pool."""
|
||||
from haiku.rag.ingester import sources as sources_module
|
||||
|
||||
uri = "https://example.com/counting.md"
|
||||
fetcher = _CountingSource(uri, b"# Counting\n\nAd-hoc fetched body.")
|
||||
monkeypatch.setattr(
|
||||
sources_module, "resolve_adhoc_fetcher", lambda *a, **kw: fetcher
|
||||
)
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document_from_source(uri)
|
||||
assert not isinstance(doc, list)
|
||||
|
||||
assert fetcher.closes == 1
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_configured_source_is_not_closed_after_ingest(temp_db_path):
|
||||
"""A source handed in by the caller (the ingester's long-lived pool) is not
|
||||
ours to close: closing it would tear down the pool mid-run."""
|
||||
uri = "https://example.com/configured.md"
|
||||
fetcher = _CountingSource(uri, b"# Configured\n\nCaller-owned body.")
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document_from_source(uri, sources=[fetcher])
|
||||
assert not isinstance(doc, list)
|
||||
|
||||
assert fetcher.closes == 0
|
||||
|
||||
|
||||
def _bbox_doc(*, with_page_image: bool, pages: tuple[int, ...] = (1,)):
|
||||
"""DoclingDocument with one paragraph per page, each carrying a bbox.
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue