chore(lint): fix ruff violations on Neo4j files
Pre-existing ruff violations surfaced by CI on PR #188: - TC001: move runtime `Neo4jDriver` imports into `TYPE_CHECKING` blocks (queries.py, chunk_writer.py, schema.py, tree_reader.py, tree_writer.py) - SIM117: combine nested `async with` in chunk_writer.write_chunks and tree_writer.write_document - SIM105: replace try/except/pass with `contextlib.suppress` in tree_writer._element_props - F401: remove unused `Neo4jDriver` import in main._init_neo4j - RUF100: remove unused `# noqa: E402` in tests/neo4j/conftest.py - I001: sort imports in tests/neo4j/test_chunk_writer.py and test_tree_writer.py Zero behaviour change. `ruff check .` now passes cleanly.
This commit is contained in:
parent
aa60fbb768
commit
5bc98ee483
9 changed files with 209 additions and 193 deletions
|
|
@ -16,9 +16,10 @@ from __future__ import annotations
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
from infra.neo4j.driver import Neo4jDriver
|
if TYPE_CHECKING:
|
||||||
|
from infra.neo4j.driver import Neo4jDriver
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
@ -63,23 +64,23 @@ async def write_chunks(
|
||||||
if ref:
|
if ref:
|
||||||
derived_rows.append({"chunk_id": cid, "doc_id": doc_id, "self_ref": ref})
|
derived_rows.append({"chunk_id": cid, "doc_id": doc_id, "self_ref": ref})
|
||||||
|
|
||||||
async with neo.driver.session(database=neo.database) as session:
|
async with (
|
||||||
async with await session.begin_transaction() as tx:
|
neo.driver.session(database=neo.database) as session,
|
||||||
# Replace existing chunks.
|
await session.begin_transaction() as tx,
|
||||||
await tx.run(
|
):
|
||||||
"""
|
# Replace existing chunks.
|
||||||
|
await tx.run(
|
||||||
|
"""
|
||||||
MATCH (d:Document {id: $doc_id})-[:HAS_CHUNK]->(c:Chunk)
|
MATCH (d:Document {id: $doc_id})-[:HAS_CHUNK]->(c:Chunk)
|
||||||
DETACH DELETE c
|
DETACH DELETE c
|
||||||
""",
|
""",
|
||||||
doc_id=doc_id,
|
doc_id=doc_id,
|
||||||
)
|
)
|
||||||
await tx.run(
|
await tx.run("MATCH (c:Chunk {doc_id: $doc_id}) DETACH DELETE c", doc_id=doc_id)
|
||||||
"MATCH (c:Chunk {doc_id: $doc_id}) DETACH DELETE c", doc_id=doc_id
|
|
||||||
)
|
|
||||||
|
|
||||||
if chunk_rows:
|
if chunk_rows:
|
||||||
await tx.run(
|
await tx.run(
|
||||||
"""
|
"""
|
||||||
MATCH (d:Document {id: $doc_id})
|
MATCH (d:Document {id: $doc_id})
|
||||||
UNWIND $rows AS r
|
UNWIND $rows AS r
|
||||||
CREATE (c:Chunk {
|
CREATE (c:Chunk {
|
||||||
|
|
@ -92,33 +93,33 @@ async def write_chunks(
|
||||||
})
|
})
|
||||||
MERGE (d)-[:HAS_CHUNK]->(c)
|
MERGE (d)-[:HAS_CHUNK]->(c)
|
||||||
""",
|
""",
|
||||||
doc_id=doc_id,
|
doc_id=doc_id,
|
||||||
rows=chunk_rows,
|
rows=chunk_rows,
|
||||||
)
|
)
|
||||||
|
|
||||||
if derived_rows:
|
if derived_rows:
|
||||||
await tx.run(
|
await tx.run(
|
||||||
"""
|
"""
|
||||||
UNWIND $rows AS r
|
UNWIND $rows AS r
|
||||||
MATCH (c:Chunk {id: r.chunk_id})
|
MATCH (c:Chunk {id: r.chunk_id})
|
||||||
MATCH (e:Element {doc_id: r.doc_id, self_ref: r.self_ref})
|
MATCH (e:Element {doc_id: r.doc_id, self_ref: r.self_ref})
|
||||||
MERGE (c)-[:DERIVED_FROM]->(e)
|
MERGE (c)-[:DERIVED_FROM]->(e)
|
||||||
""",
|
""",
|
||||||
rows=derived_rows,
|
rows=derived_rows,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Flag the Document with the new stage.
|
# Flag the Document with the new stage.
|
||||||
await tx.run(
|
await tx.run(
|
||||||
"""
|
"""
|
||||||
MATCH (d:Document {id: $doc_id})
|
MATCH (d:Document {id: $doc_id})
|
||||||
SET d.stages_applied = [s IN coalesce(d.stages_applied, []) WHERE s <> 'chunks']
|
SET d.stages_applied = [s IN coalesce(d.stages_applied, []) WHERE s <> 'chunks']
|
||||||
+ ['chunks'],
|
+ ['chunks'],
|
||||||
d.last_chunk_write = datetime()
|
d.last_chunk_write = datetime()
|
||||||
""",
|
""",
|
||||||
doc_id=doc_id,
|
doc_id=doc_id,
|
||||||
)
|
)
|
||||||
|
|
||||||
await tx.commit()
|
await tx.commit()
|
||||||
|
|
||||||
logger.info(
|
logger.info(
|
||||||
"Neo4j: wrote %d chunks (%d DERIVED_FROM) for doc %s",
|
"Neo4j: wrote %d chunks (%d DERIVED_FROM) for doc %s",
|
||||||
|
|
|
||||||
|
|
@ -3,9 +3,10 @@
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
from infra.neo4j.driver import Neo4jDriver
|
if TYPE_CHECKING:
|
||||||
|
from infra.neo4j.driver import Neo4jDriver
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
|
|
||||||
|
|
@ -7,21 +7,21 @@ this multiple times is safe — it's the contract integration tests rely on.
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from infra.neo4j.driver import Neo4jDriver
|
if TYPE_CHECKING:
|
||||||
|
from infra.neo4j.driver import Neo4jDriver
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
CONSTRAINTS: tuple[str, ...] = (
|
CONSTRAINTS: tuple[str, ...] = (
|
||||||
"CREATE CONSTRAINT document_id IF NOT EXISTS "
|
"CREATE CONSTRAINT document_id IF NOT EXISTS FOR (d:Document) REQUIRE d.id IS UNIQUE",
|
||||||
"FOR (d:Document) REQUIRE d.id IS UNIQUE",
|
|
||||||
"CREATE CONSTRAINT element_composite IF NOT EXISTS "
|
"CREATE CONSTRAINT element_composite IF NOT EXISTS "
|
||||||
"FOR (e:Element) REQUIRE (e.doc_id, e.self_ref) IS UNIQUE",
|
"FOR (e:Element) REQUIRE (e.doc_id, e.self_ref) IS UNIQUE",
|
||||||
"CREATE CONSTRAINT page_composite IF NOT EXISTS "
|
"CREATE CONSTRAINT page_composite IF NOT EXISTS "
|
||||||
"FOR (p:Page) REQUIRE (p.doc_id, p.page_no) IS UNIQUE",
|
"FOR (p:Page) REQUIRE (p.doc_id, p.page_no) IS UNIQUE",
|
||||||
"CREATE CONSTRAINT chunk_id IF NOT EXISTS "
|
"CREATE CONSTRAINT chunk_id IF NOT EXISTS FOR (c:Chunk) REQUIRE c.id IS UNIQUE",
|
||||||
"FOR (c:Chunk) REQUIRE c.id IS UNIQUE",
|
|
||||||
)
|
)
|
||||||
|
|
||||||
INDEXES: tuple[str, ...] = (
|
INDEXES: tuple[str, ...] = (
|
||||||
|
|
@ -30,8 +30,7 @@ INDEXES: tuple[str, ...] = (
|
||||||
)
|
)
|
||||||
|
|
||||||
FULLTEXT_INDEXES: tuple[str, ...] = (
|
FULLTEXT_INDEXES: tuple[str, ...] = (
|
||||||
"CREATE FULLTEXT INDEX element_text IF NOT EXISTS "
|
"CREATE FULLTEXT INDEX element_text IF NOT EXISTS FOR (e:Element) ON EACH [e.text]",
|
||||||
"FOR (e:Element) ON EACH [e.text]",
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -10,8 +10,10 @@ nodes directly.
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from infra.neo4j.driver import Neo4jDriver
|
if TYPE_CHECKING:
|
||||||
|
from infra.neo4j.driver import Neo4jDriver
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
@ -57,10 +59,6 @@ async def delete_document(neo: Neo4jDriver, doc_id: str) -> int:
|
||||||
)
|
)
|
||||||
record = await result.single()
|
record = await result.single()
|
||||||
# Also clean up orphan elements and pages tagged with this doc_id.
|
# Also clean up orphan elements and pages tagged with this doc_id.
|
||||||
await session.run(
|
await session.run("MATCH (e:Element {doc_id: $doc_id}) DETACH DELETE e", doc_id=doc_id)
|
||||||
"MATCH (e:Element {doc_id: $doc_id}) DETACH DELETE e", doc_id=doc_id
|
await session.run("MATCH (p:Page {doc_id: $doc_id}) DETACH DELETE p", doc_id=doc_id)
|
||||||
)
|
|
||||||
await session.run(
|
|
||||||
"MATCH (p:Page {doc_id: $doc_id}) DETACH DELETE p", doc_id=doc_id
|
|
||||||
)
|
|
||||||
return int(record["removed"]) if record else 0
|
return int(record["removed"]) if record else 0
|
||||||
|
|
|
||||||
|
|
@ -9,13 +9,15 @@ graph nodes is deferred to v0.6 (see docs/design/neo4j-integration.md §2).
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import contextlib
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from datetime import datetime, timezone
|
from datetime import UTC, datetime
|
||||||
from typing import Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
from infra.neo4j.driver import Neo4jDriver
|
if TYPE_CHECKING:
|
||||||
|
from infra.neo4j.driver import Neo4jDriver
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
@ -96,10 +98,8 @@ def _element_props(item: dict[str, Any], doc_id: str) -> dict[str, Any]:
|
||||||
props["caption"] = item.get("caption")
|
props["caption"] = item.get("caption")
|
||||||
if item.get("data") and isinstance(item["data"], dict):
|
if item.get("data") and isinstance(item["data"], dict):
|
||||||
# Tables carry cell layout under data; stringify to keep the schema flat.
|
# Tables carry cell layout under data; stringify to keep the schema flat.
|
||||||
try:
|
with contextlib.suppress(TypeError, ValueError):
|
||||||
props["cells_json"] = json.dumps(item["data"])
|
props["cells_json"] = json.dumps(item["data"])
|
||||||
except (TypeError, ValueError):
|
|
||||||
pass
|
|
||||||
return props
|
return props
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -145,7 +145,7 @@ async def write_document(
|
||||||
Fails fast (exception propagates) if Neo4j is unavailable — per design §8.5.
|
Fails fast (exception propagates) if Neo4j is unavailable — per design §8.5.
|
||||||
"""
|
"""
|
||||||
doc_data = json.loads(document_json)
|
doc_data = json.loads(document_json)
|
||||||
ingested_at = datetime.now(tz=timezone.utc).isoformat()
|
ingested_at = datetime.now(tz=UTC).isoformat()
|
||||||
|
|
||||||
elements: list[dict[str, Any]] = []
|
elements: list[dict[str, Any]] = []
|
||||||
for _, item in _iter_items(doc_data):
|
for _, item in _iter_items(doc_data):
|
||||||
|
|
@ -179,26 +179,24 @@ async def write_document(
|
||||||
|
|
||||||
reading_order = _dfs_order(doc_data)
|
reading_order = _dfs_order(doc_data)
|
||||||
|
|
||||||
async with neo.driver.session(database=neo.database) as session:
|
async with (
|
||||||
async with await session.begin_transaction() as tx:
|
neo.driver.session(database=neo.database) as session,
|
||||||
# 1. Wipe existing graph for this doc_id (replace strategy).
|
await session.begin_transaction() as tx,
|
||||||
await tx.run(
|
):
|
||||||
"MATCH (d:Document {id: $doc_id}) "
|
# 1. Wipe existing graph for this doc_id (replace strategy).
|
||||||
"OPTIONAL MATCH (d)-[:HAS_ROOT|HAS_CHUNK*0..]->(n) "
|
await tx.run(
|
||||||
"DETACH DELETE d, n",
|
"MATCH (d:Document {id: $doc_id}) "
|
||||||
doc_id=doc_id,
|
"OPTIONAL MATCH (d)-[:HAS_ROOT|HAS_CHUNK*0..]->(n) "
|
||||||
)
|
"DETACH DELETE d, n",
|
||||||
# Also wipe orphan elements/chunks that may still reference this doc.
|
doc_id=doc_id,
|
||||||
await tx.run(
|
)
|
||||||
"MATCH (e:Element {doc_id: $doc_id}) DETACH DELETE e", doc_id=doc_id
|
# Also wipe orphan elements/chunks that may still reference this doc.
|
||||||
)
|
await tx.run("MATCH (e:Element {doc_id: $doc_id}) DETACH DELETE e", doc_id=doc_id)
|
||||||
await tx.run(
|
await tx.run("MATCH (p:Page {doc_id: $doc_id}) DETACH DELETE p", doc_id=doc_id)
|
||||||
"MATCH (p:Page {doc_id: $doc_id}) DETACH DELETE p", doc_id=doc_id
|
|
||||||
)
|
|
||||||
|
|
||||||
# 2. Document node (carries the verbatim JSON for TreeReader).
|
# 2. Document node (carries the verbatim JSON for TreeReader).
|
||||||
await tx.run(
|
await tx.run(
|
||||||
"""
|
"""
|
||||||
CREATE (d:Document {
|
CREATE (d:Document {
|
||||||
id: $doc_id,
|
id: $doc_id,
|
||||||
title: $title,
|
title: $title,
|
||||||
|
|
@ -211,32 +209,32 @@ async def write_document(
|
||||||
document_json: $document_json
|
document_json: $document_json
|
||||||
})
|
})
|
||||||
""",
|
""",
|
||||||
doc_id=doc_id,
|
doc_id=doc_id,
|
||||||
title=filename,
|
title=filename,
|
||||||
source_uri=source_uri or "",
|
source_uri=source_uri or "",
|
||||||
ingested_at=ingested_at,
|
ingested_at=ingested_at,
|
||||||
docling_version=docling_version or "",
|
docling_version=docling_version or "",
|
||||||
tenant_id=tenant_id,
|
tenant_id=tenant_id,
|
||||||
document_json=document_json,
|
document_json=document_json,
|
||||||
|
)
|
||||||
|
|
||||||
|
# 3. Page nodes.
|
||||||
|
if pages:
|
||||||
|
await tx.run(
|
||||||
|
"UNWIND $pages AS p "
|
||||||
|
"CREATE (:Page {doc_id: p.doc_id, page_no: p.page_no, "
|
||||||
|
"width: p.width, height: p.height})",
|
||||||
|
pages=pages,
|
||||||
)
|
)
|
||||||
|
|
||||||
# 3. Page nodes.
|
# 4. Element nodes — use dynamic :Element:<specific> labels via APOC-free trick.
|
||||||
if pages:
|
# We split by specific label so the CREATE statement is static (no APOC).
|
||||||
await tx.run(
|
by_specific: dict[str, list[dict[str, Any]]] = {}
|
||||||
"UNWIND $pages AS p "
|
for e in elements:
|
||||||
"CREATE (:Page {doc_id: p.doc_id, page_no: p.page_no, "
|
by_specific.setdefault(e["specific_label"], []).append(e)
|
||||||
"width: p.width, height: p.height})",
|
for specific, batch in by_specific.items():
|
||||||
pages=pages,
|
await tx.run(
|
||||||
)
|
f"""
|
||||||
|
|
||||||
# 4. Element nodes — use dynamic :Element:<specific> labels via APOC-free trick.
|
|
||||||
# We split by specific label so the CREATE statement is static (no APOC).
|
|
||||||
by_specific: dict[str, list[dict[str, Any]]] = {}
|
|
||||||
for e in elements:
|
|
||||||
by_specific.setdefault(e["specific_label"], []).append(e)
|
|
||||||
for specific, batch in by_specific.items():
|
|
||||||
await tx.run(
|
|
||||||
f"""
|
|
||||||
UNWIND $batch AS e
|
UNWIND $batch AS e
|
||||||
CREATE (n:Element:{specific} {{
|
CREATE (n:Element:{specific} {{
|
||||||
doc_id: e.doc_id,
|
doc_id: e.doc_id,
|
||||||
|
|
@ -250,83 +248,83 @@ async def write_document(
|
||||||
cells_json: e.cells_json
|
cells_json: e.cells_json
|
||||||
}})
|
}})
|
||||||
""",
|
""",
|
||||||
batch=batch,
|
batch=batch,
|
||||||
)
|
)
|
||||||
|
|
||||||
# 5. PARENT_OF relations (tree structure). Order tracked inline.
|
# 5. PARENT_OF relations (tree structure). Order tracked inline.
|
||||||
parent_rows = [
|
parent_rows = [
|
||||||
{
|
{
|
||||||
"doc_id": doc_id,
|
"doc_id": doc_id,
|
||||||
"parent_ref": e["parent_ref"],
|
"parent_ref": e["parent_ref"],
|
||||||
"child_ref": e["self_ref"],
|
"child_ref": e["self_ref"],
|
||||||
"order": idx,
|
"order": idx,
|
||||||
}
|
}
|
||||||
for idx, e in enumerate(elements)
|
for idx, e in enumerate(elements)
|
||||||
if e["parent_ref"] and e["parent_ref"] != "#/body"
|
if e["parent_ref"] and e["parent_ref"] != "#/body"
|
||||||
]
|
]
|
||||||
if parent_rows:
|
if parent_rows:
|
||||||
await tx.run(
|
await tx.run(
|
||||||
"""
|
"""
|
||||||
UNWIND $rows AS r
|
UNWIND $rows AS r
|
||||||
MATCH (p:Element {doc_id: r.doc_id, self_ref: r.parent_ref})
|
MATCH (p:Element {doc_id: r.doc_id, self_ref: r.parent_ref})
|
||||||
MATCH (c:Element {doc_id: r.doc_id, self_ref: r.child_ref})
|
MATCH (c:Element {doc_id: r.doc_id, self_ref: r.child_ref})
|
||||||
MERGE (p)-[rel:PARENT_OF]->(c)
|
MERGE (p)-[rel:PARENT_OF]->(c)
|
||||||
SET rel.order = r.order
|
SET rel.order = r.order
|
||||||
""",
|
""",
|
||||||
rows=parent_rows,
|
rows=parent_rows,
|
||||||
)
|
)
|
||||||
|
|
||||||
# 6. HAS_ROOT for top-level children of the document body.
|
# 6. HAS_ROOT for top-level children of the document body.
|
||||||
root_rows = [
|
root_rows = [
|
||||||
{"doc_id": doc_id, "child_ref": e["self_ref"]}
|
{"doc_id": doc_id, "child_ref": e["self_ref"]}
|
||||||
for e in elements
|
for e in elements
|
||||||
if e["parent_ref"] == "#/body"
|
if e["parent_ref"] == "#/body"
|
||||||
]
|
]
|
||||||
if root_rows:
|
if root_rows:
|
||||||
await tx.run(
|
await tx.run(
|
||||||
"""
|
"""
|
||||||
UNWIND $rows AS r
|
UNWIND $rows AS r
|
||||||
MATCH (d:Document {id: r.doc_id})
|
MATCH (d:Document {id: r.doc_id})
|
||||||
MATCH (c:Element {doc_id: r.doc_id, self_ref: r.child_ref})
|
MATCH (c:Element {doc_id: r.doc_id, self_ref: r.child_ref})
|
||||||
MERGE (d)-[:HAS_ROOT]->(c)
|
MERGE (d)-[:HAS_ROOT]->(c)
|
||||||
""",
|
""",
|
||||||
rows=root_rows,
|
rows=root_rows,
|
||||||
)
|
)
|
||||||
|
|
||||||
# 7. ON_PAGE from first provenance.
|
# 7. ON_PAGE from first provenance.
|
||||||
on_page_rows = [
|
on_page_rows = [
|
||||||
{"doc_id": doc_id, "self_ref": e["self_ref"], "page_no": e["prov_page"]}
|
{"doc_id": doc_id, "self_ref": e["self_ref"], "page_no": e["prov_page"]}
|
||||||
for e in elements
|
for e in elements
|
||||||
if e["prov_page"] is not None
|
if e["prov_page"] is not None
|
||||||
]
|
]
|
||||||
if on_page_rows:
|
if on_page_rows:
|
||||||
await tx.run(
|
await tx.run(
|
||||||
"""
|
"""
|
||||||
UNWIND $rows AS r
|
UNWIND $rows AS r
|
||||||
MATCH (e:Element {doc_id: r.doc_id, self_ref: r.self_ref})
|
MATCH (e:Element {doc_id: r.doc_id, self_ref: r.self_ref})
|
||||||
MATCH (p:Page {doc_id: r.doc_id, page_no: r.page_no})
|
MATCH (p:Page {doc_id: r.doc_id, page_no: r.page_no})
|
||||||
MERGE (e)-[:ON_PAGE]->(p)
|
MERGE (e)-[:ON_PAGE]->(p)
|
||||||
""",
|
""",
|
||||||
rows=on_page_rows,
|
rows=on_page_rows,
|
||||||
)
|
)
|
||||||
|
|
||||||
# 8. NEXT chain in DFS pre-order.
|
# 8. NEXT chain in DFS pre-order.
|
||||||
if len(reading_order) > 1:
|
if len(reading_order) > 1:
|
||||||
pairs = [
|
pairs = [
|
||||||
{"doc_id": doc_id, "a": reading_order[i], "b": reading_order[i + 1]}
|
{"doc_id": doc_id, "a": reading_order[i], "b": reading_order[i + 1]}
|
||||||
for i in range(len(reading_order) - 1)
|
for i in range(len(reading_order) - 1)
|
||||||
]
|
]
|
||||||
await tx.run(
|
await tx.run(
|
||||||
"""
|
"""
|
||||||
UNWIND $pairs AS p
|
UNWIND $pairs AS p
|
||||||
MATCH (a:Element {doc_id: p.doc_id, self_ref: p.a})
|
MATCH (a:Element {doc_id: p.doc_id, self_ref: p.a})
|
||||||
MATCH (b:Element {doc_id: p.doc_id, self_ref: p.b})
|
MATCH (b:Element {doc_id: p.doc_id, self_ref: p.b})
|
||||||
MERGE (a)-[:NEXT]->(b)
|
MERGE (a)-[:NEXT]->(b)
|
||||||
""",
|
""",
|
||||||
pairs=pairs,
|
pairs=pairs,
|
||||||
)
|
)
|
||||||
|
|
||||||
await tx.commit()
|
await tx.commit()
|
||||||
|
|
||||||
logger.info(
|
logger.info(
|
||||||
"Neo4j: wrote doc %s (%d elements, %d pages)",
|
"Neo4j: wrote doc %s (%d elements, %d pages)",
|
||||||
|
|
|
||||||
|
|
@ -101,7 +101,7 @@ async def _init_neo4j():
|
||||||
logger.info("Neo4j disabled (NEO4J_URI not set)")
|
logger.info("Neo4j disabled (NEO4J_URI not set)")
|
||||||
return None
|
return None
|
||||||
|
|
||||||
from infra.neo4j import Neo4jDriver, bootstrap_schema, get_driver
|
from infra.neo4j import bootstrap_schema, get_driver
|
||||||
|
|
||||||
try:
|
try:
|
||||||
neo = await get_driver(
|
neo = await get_driver(
|
||||||
|
|
|
||||||
|
|
@ -15,7 +15,7 @@ import pytest
|
||||||
# (e.g. local dev without the dependency installed).
|
# (e.g. local dev without the dependency installed).
|
||||||
pytest.importorskip("neo4j")
|
pytest.importorskip("neo4j")
|
||||||
|
|
||||||
from infra.neo4j import close_driver, get_driver # noqa: E402
|
from infra.neo4j import close_driver, get_driver
|
||||||
|
|
||||||
|
|
||||||
def _cfg() -> tuple[str, str, str]:
|
def _cfg() -> tuple[str, str, str]:
|
||||||
|
|
|
||||||
|
|
@ -12,7 +12,6 @@ from infra.neo4j import fetch_graph, write_chunks, write_document
|
||||||
from infra.neo4j.schema import bootstrap_schema
|
from infra.neo4j.schema import bootstrap_schema
|
||||||
from tests.neo4j.test_tree_writer import FIXTURE
|
from tests.neo4j.test_tree_writer import FIXTURE
|
||||||
|
|
||||||
|
|
||||||
CHUNKS = [
|
CHUNKS = [
|
||||||
{
|
{
|
||||||
"text": "Introduction. First paragraph on page 1.",
|
"text": "Introduction. First paragraph on page 1.",
|
||||||
|
|
@ -66,8 +65,7 @@ async def test_write_chunks_and_derived_from(neo4j_driver):
|
||||||
for idx, expected in [(0, 2), (1, 1)]:
|
for idx, expected in [(0, 2), (1, 1)]:
|
||||||
cnt = await (
|
cnt = await (
|
||||||
await s.run(
|
await s.run(
|
||||||
"MATCH (c:Chunk {id: $cid})-[:DERIVED_FROM]->(e:Element) "
|
"MATCH (c:Chunk {id: $cid})-[:DERIVED_FROM]->(e:Element) RETURN count(e) AS n",
|
||||||
"RETURN count(e) AS n",
|
|
||||||
cid=f"doc-fixture::chunk::{idx}",
|
cid=f"doc-fixture::chunk::{idx}",
|
||||||
)
|
)
|
||||||
).single()
|
).single()
|
||||||
|
|
|
||||||
|
|
@ -13,7 +13,6 @@ import json
|
||||||
from infra.neo4j import read_document_json, write_document
|
from infra.neo4j import read_document_json, write_document
|
||||||
from infra.neo4j.schema import bootstrap_schema
|
from infra.neo4j.schema import bootstrap_schema
|
||||||
|
|
||||||
|
|
||||||
FIXTURE = {
|
FIXTURE = {
|
||||||
"name": "fixture.pdf",
|
"name": "fixture.pdf",
|
||||||
"pages": {
|
"pages": {
|
||||||
|
|
@ -89,41 +88,59 @@ async def test_write_creates_expected_structure(neo4j_driver):
|
||||||
assert result.pages_written == 2
|
assert result.pages_written == 2
|
||||||
|
|
||||||
async with neo4j_driver.driver.session(database=neo4j_driver.database) as s:
|
async with neo4j_driver.driver.session(database=neo4j_driver.database) as s:
|
||||||
assert await _count(
|
assert (
|
||||||
s,
|
await _count(
|
||||||
"MATCH (d:Document {id: $id}) RETURN count(d) AS n",
|
s,
|
||||||
id="doc-fixture",
|
"MATCH (d:Document {id: $id}) RETURN count(d) AS n",
|
||||||
) == 1
|
id="doc-fixture",
|
||||||
assert await _count(
|
)
|
||||||
s,
|
== 1
|
||||||
"MATCH (:Document {id: $id})-[:HAS_ROOT]->(e:Element) RETURN count(e) AS n",
|
)
|
||||||
id="doc-fixture",
|
assert (
|
||||||
) == 4
|
await _count(
|
||||||
assert await _count(
|
s,
|
||||||
s,
|
"MATCH (:Document {id: $id})-[:HAS_ROOT]->(e:Element) RETURN count(e) AS n",
|
||||||
"MATCH (e:Element:SectionHeader {doc_id: $id, self_ref: '#/texts/0'}) "
|
id="doc-fixture",
|
||||||
"RETURN count(e) AS n",
|
)
|
||||||
id="doc-fixture",
|
== 4
|
||||||
) == 1
|
)
|
||||||
assert await _count(
|
assert (
|
||||||
s,
|
await _count(
|
||||||
"MATCH (e:Element:Table {doc_id: $id}) RETURN count(e) AS n",
|
s,
|
||||||
id="doc-fixture",
|
"MATCH (e:Element:SectionHeader {doc_id: $id, self_ref: '#/texts/0'}) "
|
||||||
) == 1
|
"RETURN count(e) AS n",
|
||||||
|
id="doc-fixture",
|
||||||
|
)
|
||||||
|
== 1
|
||||||
|
)
|
||||||
|
assert (
|
||||||
|
await _count(
|
||||||
|
s,
|
||||||
|
"MATCH (e:Element:Table {doc_id: $id}) RETURN count(e) AS n",
|
||||||
|
id="doc-fixture",
|
||||||
|
)
|
||||||
|
== 1
|
||||||
|
)
|
||||||
# Reading-order chain: 3 NEXT edges for 4 elements.
|
# Reading-order chain: 3 NEXT edges for 4 elements.
|
||||||
assert await _count(
|
assert (
|
||||||
s,
|
await _count(
|
||||||
"MATCH (a:Element {doc_id: $id})-[:NEXT]->(b:Element {doc_id: $id}) "
|
s,
|
||||||
"RETURN count(*) AS n",
|
"MATCH (a:Element {doc_id: $id})-[:NEXT]->(b:Element {doc_id: $id}) "
|
||||||
id="doc-fixture",
|
"RETURN count(*) AS n",
|
||||||
) == 3
|
id="doc-fixture",
|
||||||
|
)
|
||||||
|
== 3
|
||||||
|
)
|
||||||
# ON_PAGE: one per element with prov.
|
# ON_PAGE: one per element with prov.
|
||||||
assert await _count(
|
assert (
|
||||||
s,
|
await _count(
|
||||||
"MATCH (:Element {doc_id: $id})-[:ON_PAGE]->(:Page {doc_id: $id}) "
|
s,
|
||||||
"RETURN count(*) AS n",
|
"MATCH (:Element {doc_id: $id})-[:ON_PAGE]->(:Page {doc_id: $id}) "
|
||||||
id="doc-fixture",
|
"RETURN count(*) AS n",
|
||||||
) == 4
|
id="doc-fixture",
|
||||||
|
)
|
||||||
|
== 4
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
async def test_rewrite_is_idempotent_replace(neo4j_driver):
|
async def test_rewrite_is_idempotent_replace(neo4j_driver):
|
||||||
|
|
@ -145,14 +162,18 @@ async def test_rewrite_is_idempotent_replace(neo4j_driver):
|
||||||
)
|
)
|
||||||
|
|
||||||
async with neo4j_driver.driver.session(database=neo4j_driver.database) as s:
|
async with neo4j_driver.driver.session(database=neo4j_driver.database) as s:
|
||||||
assert await _count(
|
assert (
|
||||||
s, "MATCH (d:Document {id: $id}) RETURN count(d) AS n", id="doc-fixture"
|
await _count(s, "MATCH (d:Document {id: $id}) RETURN count(d) AS n", id="doc-fixture")
|
||||||
) == 1
|
== 1
|
||||||
assert await _count(
|
)
|
||||||
s,
|
assert (
|
||||||
"MATCH (e:Element {doc_id: $id}) RETURN count(e) AS n",
|
await _count(
|
||||||
id="doc-fixture",
|
s,
|
||||||
) == 4
|
"MATCH (e:Element {doc_id: $id}) RETURN count(e) AS n",
|
||||||
|
id="doc-fixture",
|
||||||
|
)
|
||||||
|
== 4
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
async def test_reader_returns_verbatim_json(neo4j_driver):
|
async def test_reader_returns_verbatim_json(neo4j_driver):
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue