class DBWriter:
"""Push :class:`GraphContainer` data to the target graph database.
The orchestrator (e.g. :class:`Caster`) must initialize ``schema`` and
``ingestion_model`` for the target database (``db_profile.db_flavor``,
:meth:`Schema.finish_init`, :meth:`IngestionModel.finish_init`) before
calling :meth:`write`; this class does not repeat that work on every batch.
Concurrency contract: one instance may serve concurrent :meth:`write`
calls (sibling batches in flight) from a single event loop. All per-batch
mutation happens on the caller's ``gc``; instance state is limited to the
cached db-aware schema (pre-warm it via :meth:`_db_aware_for` before
fanning out) and one shared semaphore, so ``max_concurrent`` bounds DB
operations across every in-flight batch, not per call.
Naming contract: documents in ``gc`` carry logical property names, and
stay that way. Storage, relation and property names the target stores
differently come from the profile (completed by
:func:`~graflo.architecture.evolution.sanitize.with_physical_names`), and
every document and field list is translated to them at the backend call,
then translated back for anything the backend returns.
Attributes:
schema: Schema configuration providing vertex/edge metadata.
dry: When ``True`` no database mutations are performed.
max_concurrent: Upper bound on concurrent DB operations (semaphore size).
A file-backend target is written by one operation at a time.
stats: What attaching records and resolving edge endpoints found,
accumulated over every :meth:`write` of this instance.
"""
def __init__(
self,
schema: Schema,
ingestion_model: IngestionModel,
*,
dry: bool = False,
max_concurrent: int = 1,
):
self.schema = schema
self.ingestion_model = ingestion_model
self.dry = dry
self.max_concurrent = max_concurrent
self._context: WriteContext | None = None
self._semaphore: asyncio.Semaphore | None = None
self._semaphore_loop: asyncio.AbstractEventLoop | None = None
self._semaphore_bound: int | None = None
self._collection_locks: dict[tuple[str | None, ...], asyncio.Lock] = {}
self._collection_locks_loop: asyncio.AbstractEventLoop | None = None
self._reported_undeclared_edges: set[tuple] = set()
# Idle connections per config object, as (manager, connection) pairs.
self._pools: dict[int, tuple[DBConfig, queue.SimpleQueue[tuple[Any, Any]]]] = {}
self._pools_lock = threading.Lock()
self.stats = WriteStats()
def close(self) -> None:
"""Close the connections this writer keeps for its run.
A connection still in use is closed when its work returns. A failure to
close one is logged, not raised, so it neither stops the others closing
nor replaces the error a failed run is raising.
"""
with self._pools_lock:
pools, self._pools = self._pools, {}
for _, idle in pools.values():
while True:
try:
manager, _ = idle.get_nowait()
except queue.Empty:
break
self._close_quietly(manager)
@staticmethod
def _close_quietly(manager: Any) -> None:
try:
manager.__exit__(None, None, None)
except Exception:
logger.warning("Could not close a database connection", exc_info=True)
@contextmanager
def _connection(self, conn_conf: DBConfig) -> Generator[Any, None, None]:
"""A connection to *conn_conf* for one unit of work, kept for the run.
Connections are opened lazily and reused, so a run holds at most as
many as it ran operations at once (``max_concurrent`` for writes). Each
is used by one thread at a time; one whose work raised is closed rather
than reused. :meth:`close` closes the rest.
"""
if conn_conf.connection_type in _UNPOOLED_FLAVORS:
with ConnectionManager(connection_config=conn_conf) as db:
yield db
return
with self._pools_lock:
entry = self._pools.get(id(conn_conf))
if entry is None:
entry = (conn_conf, queue.SimpleQueue())
self._pools[id(conn_conf)] = entry
idle = entry[1]
try:
manager, db = idle.get_nowait()
except queue.Empty:
manager = ConnectionManager(connection_config=conn_conf)
db = manager.__enter__()
try:
yield db
if isinstance(db, Connection):
db.reset_after_use()
except BaseException:
self._close_quietly(manager)
raise
with self._pools_lock:
# A writer closed while this work ran has dropped the pool.
kept = self._pools.get(id(conn_conf)) is entry
if kept:
idle.put((manager, db))
if not kept:
self._close_quietly(manager)
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
async def write(
self,
gc: GraphContainer,
conn_conf: DBConfig,
resource_name: str | None,
*,
stamp: Mapping[str, Any] | None = None,
guard: FilterExpression | None = None,
) -> None:
"""Push *gc* to the database (vertices, attached records, extra weights,
then edges).
Every reserved field of every element written is set from *stamp*,
whatever value a document carries.
*guard* is a trusted read guard every write-time lookup must pass —
secondary-identity edge endpoints, attached-record matches and
extra-weight vertices — so a write never links to, attaches onto or
copies from a vertex the guard hides.
A container with no documents opens no connection and writes nothing;
a class with no documents is skipped the same way.
.. note::
*gc* is mutated in-place: blank-vertex keys are updated and blank
edges are extended after the vertex round-trip.
Raises:
ValueError: If an element of a type with a reserved field and
*stamp* does not set it.
NotImplementedError: If *guard* is given and the backend does not
apply read guards; nothing is written.
"""
# Checks that hold whatever the container holds run first, so a batch
# that cast to nothing is refused like any other.
if guard is not None:
require_read_guard(conn_conf)
if gc.is_empty():
return
check_stamp(gc, self.schema, stamp)
refuse_reserved_attributes(gc, self.schema)
self._check_union(gc, conn_conf)
# A container that no resource produced (``migrate_graph`` writes an
# exported graph through an empty ingestion model) has no resource to
# consult, and needs none: extra weights, endpoint selections and the
# blank-sibling join below are resource-level ingestion features.
resource = (
self.ingestion_model.fetch_resource(resource_name)
if resource_name is not None or self.ingestion_model.resources
else None
)
apply_stamp(gc, self.schema, stamp, edges=False)
await self._push_vertices(gc, conn_conf)
await self._attach_vertices(gc, conn_conf, guard=guard)
if resource is not None:
# Joins each blank vertex to its siblings from the same cast. A
# resource-less container carries its edges explicitly, and the join
# would pair unrelated documents. O(blank vertices x edges x
# documents) of pure Python, so off the event loop.
await asyncio.to_thread(self._resolve_blank_edges, gc, conn_conf)
if resource is not None:
await self._enrich_extra_weights(gc, conn_conf, resource, guard=guard)
# Edges last: blank resolution and extra weights add rows and
# attributes, and the stamp must win over both.
refuse_reserved_attributes(gc, self.schema)
apply_stamp(gc, self.schema, stamp, vertices=False)
await self._push_edges(gc, conn_conf, resource, guard=guard)
def delete_vertices(
self, conn_conf: DBConfig, vertex: str, key_docs: list[dict[str, Any]]
) -> None:
"""Remove vertices of *vertex*, with every edge incident to them.
*key_docs* carry the vertex's identity fields, under logical names; a
document missing one removes nothing. Nothing is removed in a dry run.
Raises:
ValueError: If the target does not support deleting vertices.
"""
ConnectionManager.require(conn_conf, ConnectionCapability.INSTANCE_DELETE)
if self.dry:
return
vc = self._db_aware_for(conn_conf).vertex_config
keys = self._keys_for(conn_conf)
match_keys = tuple(keys.vertex_fields(vertex, vc.identity_fields(vertex)))
with ConnectionManager(connection_config=conn_conf) as db:
db.delete_vertices(
vc.vertex_dbname(vertex), keys.vertex_docs(vertex, key_docs), match_keys
)
def delete_edges(
self,
conn_conf: DBConfig,
edge_id: EdgeId,
endpoints: list[tuple[dict[str, Any], dict[str, Any]]],
) -> None:
"""Remove the edges *edge_id* between the given ``(source, target)`` pairs.
Endpoints carry their vertices' identity fields, under logical names.
Nothing is removed in a dry run.
Raises:
ValueError: If the target does not support deleting edges, or the
schema does not declare *edge_id*.
"""
ConnectionManager.require(conn_conf, ConnectionCapability.INSTANCE_DELETE)
core_ec = self.schema.core_schema.edge_config
if edge_id not in core_ec:
raise ValueError(f"Edge {edge_id} is not declared in the schema")
if self.dry:
return
schema_db = self._db_aware_for(conn_conf)
vc = schema_db.vertex_config
keys = self._keys_for(conn_conf)
edge = core_ec.edge_for(edge_id)
runtime = schema_db.edge_config.runtime(edge)
_, relation_name = project_edge_docs(
self.context(conn_conf), [], edge.relation, runtime
)
source_keys = vc.identity_fields(edge.source)
target_keys = vc.identity_fields(edge.target)
stored = [
(keys.vertex_doc(edge.source, source), keys.vertex_doc(edge.target, target))
for source, target in endpoints
]
with ConnectionManager(connection_config=conn_conf) as db:
db.delete_edges(
vc.vertex_dbname(edge.source),
vc.vertex_dbname(edge.target),
relation_name,
stored,
tuple(keys.vertex_fields(edge.source, source_keys)),
tuple(keys.vertex_fields(edge.target, target_keys)),
collection_name=runtime.storage_name(),
)
def delete_scoped(
self,
conn_conf: DBConfig,
field: str,
value: Any,
*,
vertices: Sequence[str] = (),
edges: Sequence[EdgeId] = (),
) -> None:
"""Remove the elements of the given types whose *field* equals *value*.
Names are logical. Edges go first; a removed vertex takes every edge
incident to it, scoped or not. Nothing is removed in a dry run.
Args:
conn_conf: The database to remove them from.
field: The scope field, declared on every type named.
value: The scope; never ``None``.
vertices: Logical vertex types.
edges: Logical ``(source, target, relation)`` of declared edges.
Raises:
ValueError: If the target does not support scoped deletes, *value*
is ``None``, an edge is not declared, or a type does not
declare *field*.
"""
ConnectionManager.require(conn_conf, ConnectionCapability.SCOPED_DELETE)
if value is None:
raise ValueError("a scoped delete needs a scope value, not None")
core = self.schema.core_schema
for vertex in vertices:
if field not in core.vertex_config.property_names(vertex):
raise ValueError(f"Vertex {vertex!r} declares no field {field!r}")
for edge_id in edges:
if edge_id not in core.edge_config:
raise ValueError(f"Edge {edge_id} is not declared in the schema")
if field not in core.edge_config.edge_for(edge_id).property_names:
raise ValueError(f"Edge {edge_id} declares no field {field!r}")
if self.dry:
return
schema_db = self._db_aware_for(conn_conf)
vc = schema_db.vertex_config
keys = self._keys_for(conn_conf)
with ConnectionManager(connection_config=conn_conf) as db:
for edge_id in edges:
edge = core.edge_config.edge_for(edge_id)
runtime = schema_db.edge_config.runtime(edge)
_, relation_name = project_edge_docs(
self.context(conn_conf), [], edge.relation, runtime
)
db.delete_scoped_edges(
vc.vertex_dbname(edge.source),
vc.vertex_dbname(edge.target),
relation_name,
keys.edge_fields(edge_id, [field])[0],
value,
collection_name=runtime.storage_name(),
)
for vertex in vertices:
db.delete_scoped_vertices(
vc.vertex_dbname(vertex),
keys.vertex_fields(vertex, [field])[0],
value,
)
# ------------------------------------------------------------------
# Vertices
# ------------------------------------------------------------------
async def _push_vertices(self, gc: GraphContainer, conn_conf: DBConfig) -> None:
"""Upsert all vertex collections in *gc*.
Pre-write hooks depend on :attr:`~graflo.architecture.schema.vertex.Vertex.identity_mode`:
``hash`` vertices get deterministic SHA256 ids;
``assigned`` vertices get idempotent uuid4 fill (usually already minted at assemble);
``blank`` vertices get random UUIDs;
``natural`` vertices upsert directly on ``identity`` fields (one or many).
UUID-typed natural identity fields are validated when present.
"""
ctx = self.context(conn_conf)
vc = ctx.db_aware.vertex_config
def _prepare(vcol: str, data: list[dict]) -> StoredVertices:
writable = prepare_vertices(ctx, vcol, data)
identity_fields = vc.identity_fields(vcol)
if identity_fields:
# One row per vertex per batch, later wins: a backend that
# replaces a row or refuses a second touch would otherwise lose
# or fail on the earlier document.
writable = fuse_doc_basis(
writable, tuple(identity_fields), union_fields(ctx, vcol)
)
return project_vertices(ctx, vcol, writable)
async def _push_one(vcol: str, data: list[dict]) -> None:
# Prepared before a write slot or a connection is taken: nothing in
# it reads the database.
batch = await asyncio.to_thread(_prepare, vcol, data)
async with AsyncExitStack() as stack:
await self._acquire_write_slot(
stack, conn_conf, ("vertex", vc.vertex_dbname(vcol))
)
def _sync() -> None:
with self._connection(conn_conf) as db:
self._upsert(db, batch, conn_conf)
await asyncio.to_thread(_sync)
await asyncio.gather(
*[_push_one(vcol, data) for vcol, data in gc.vertices.items() if data]
)
def _check_union(self, gc: GraphContainer, conn_conf: DBConfig) -> None:
"""Refuse a write with ``merge: union`` fields the target cannot union.
Checked once, before anything is written, on every path.
Raises:
ValueError: A union field on a vertex or an edge when the target
lacks union support.
"""
core = self.schema.core_schema
vc = core.vertex_config
on_vertices = any(
vcol in vc.vertex_set
and docs
and any(f.merge == "union" for f in vc.properties(vcol))
for docs_by_type in (gc.vertices, gc.attached)
for vcol, docs in docs_by_type.items()
)
on_edges = any(
rows
and (edge := core.edge_config.declared(edge_id)) is not None
and edge_union_fields(edge)
for edge_id, rows in gc.edges.items()
)
if on_vertices or on_edges:
ConnectionManager.require(conn_conf, ConnectionCapability.UNION_MERGE)
def _upsert(self, db: Any, batch: StoredVertices, conn_conf: DBConfig) -> None:
"""Upsert one vertex type's stored documents on their identity.
Raises:
ValueError: If the batch has union fields the target cannot write
as a union, rather than writing a replace in its place.
"""
union_kw: dict[str, Any] = {}
if batch.union_fields:
ConnectionManager.require(conn_conf, ConnectionCapability.UNION_MERGE)
union_kw["union_fields"] = batch.union_fields
db.upsert_docs_batch(
batch.docs,
batch.class_name,
batch.match_keys,
update_keys="doc",
filter_uniques=True,
dry=self.dry,
**union_kw,
)
# ------------------------------------------------------------------
# Attached records
# ------------------------------------------------------------------
async def _attach_vertices(
self,
gc: GraphContainer,
conn_conf: DBConfig,
*,
guard: FilterExpression | None = None,
) -> None:
"""Write attached records onto the vertices their secondary identity finds.
Each record is looked up by the identity ``gc.attached_by`` names for its
class, takes the primary identity of the vertex it matches, and is
upserted on it (records landing on one vertex fused into one row) — in a call of its own, after the class's owners, so a
record never nulls a property the owners wrote in the same batch. A
record that matches nothing is counted and not written: attaching never
creates a vertex. Several matches follow ``endpoints_on_ambiguous``,
writing the record onto each vertex it keeps. A vertex *guard* hides is
never matched. A dry run only counts.
Raises:
ValueError: If an attached class has no selector in ``gc.attached_by``.
AmbiguousEndpointError: on several matches under the ``error`` policy;
no class's records are written then.
"""
vc = self._db_aware_for(conn_conf).vertex_config
policy = self.ingestion_model.endpoints_on_ambiguous
async def _resolve_one(
vcol: str, docs: list[dict]
) -> tuple[str, str, list[dict[str, Any]], AttachStats]:
selector = gc.attached_by.get(vcol)
if selector is None:
raise ValueError(
f"Attached '{vcol}' records name no secondary identity to be "
"found by."
)
if self.dry:
return vcol, selector, [], AttachStats(documents=len(docs))
async with self._db_semaphore(conn_conf):
def _sync() -> tuple[list[dict[str, Any]], AttachStats]:
with self._connection(conn_conf) as db:
return self._resolve_attached(
db,
vcol,
docs,
selector,
conn_conf=conn_conf,
policy=policy,
guard=guard,
)
filled, stats = await asyncio.to_thread(_sync)
return vcol, selector, filled, stats
async def _upsert_one(
vcol: str, selector: str, filled: list[dict[str, Any]], stats: AttachStats
) -> None:
if filled:
async with AsyncExitStack() as stack:
await self._acquire_write_slot(
stack, conn_conf, ("vertex", vc.vertex_dbname(vcol))
)
def _sync() -> None:
with self._connection(conn_conf) as db:
self._upsert_attached(db, vcol, filled, conn_conf=conn_conf)
await asyncio.to_thread(_sync)
stats.written = len(filled)
self._record_attach(vcol, selector, stats)
# Every class resolves before any writes, so an ``error`` raised for one
# class leaves the others unwritten too.
resolved = await _gather_settled(
[_resolve_one(vcol, docs) for vcol, docs in gc.attached.items() if docs]
)
await asyncio.gather(*[_upsert_one(*entry) for entry in resolved])
def _resolve_attached(
self,
db: Any,
vcol: str,
docs: list[dict],
selector: str,
*,
conn_conf: DBConfig,
policy: EndpointAmbiguityPolicy,
guard: FilterExpression | None = None,
) -> tuple[list[dict[str, Any]], AttachStats]:
"""Resolve one class's attached records to the rows to upsert.
Returns the rows, fused one per vertex and keyed like their owners, and
the counts; ``written`` is left for the upsert to set.
Raises:
AmbiguousEndpointError: on several matches under the ``error`` policy.
"""
vc = self._db_aware_for(conn_conf).vertex_config
keys = self._keys_for(conn_conf)
stats = AttachStats(documents=len(docs))
match_fields = vc.match_fields(vcol, selector)
identity_fields = vc.identity_fields(vcol)
# Resolution reads the database, so it runs on stored names.
match_keys = tuple(keys.vertex_fields(vcol, match_fields))
identity_keys = tuple(keys.vertex_fields(vcol, identity_fields))
stored_docs = keys.vertex_docs(vcol, docs)
resolvable = [key_tuple(doc, match_keys) is not None for doc in stored_docs]
found: dict[int, list[dict[str, Any]]] = {}
if any(resolvable):
found = db.resolve_vertices(
vc.vertex_dbname(vcol),
stored_docs,
match_keys,
identity_keys,
**_guard_kwargs(guard),
)
filled: list[dict[str, Any]] = []
for position, doc in enumerate(docs):
if not resolvable[position]:
# A partial composite key must never be partially matched.
stats.unresolvable += 1
continue
# A match without its primary identity cannot be upserted onto: keyed
# on null, the write would create a vertex.
matches = [
match
for match in found.get(position) or []
if key_tuple(match, identity_keys) is not None
]
if not matches:
stats.unmatched += 1
continue
if len(matches) > 1:
stats.ambiguous += 1
selected = select_matches(
matches,
policy=policy,
identity_fields=identity_keys,
describe=lambda count, doc=stored_docs[position]: (
f"Attached '{vcol}' record matched {count} vertices on "
f"{selector} {list(match_keys)}={key_tuple(doc, match_keys)}."
),
)
if selected:
stats.attached += 1
filled.extend(
{**doc, **keys.logical_vertex_doc(vcol, identity)}
for identity in selected
)
if not filled:
return filled, stats
# Records landing on one vertex become one row, later records winning:
# a batch upsert may not touch the same row twice.
ctx = self.context(conn_conf)
filled = fuse_doc_basis(filled, tuple(identity_fields), union_fields(ctx, vcol))
# Keyed as their owners were when written.
mirror_default_identity(ctx, vcol, filled)
return filled, stats
def _upsert_attached(
self, db: Any, vcol: str, filled: list[dict[str, Any]], *, conn_conf: DBConfig
) -> None:
"""Upsert one class's resolved attached rows on their primary identity."""
ctx = self.context(conn_conf)
derive_vertex_fields(ctx, vcol, filled)
self._upsert(db, project_vertices(ctx, vcol, filled), conn_conf)
def _record_attach(self, vcol: str, selector: str, stats: AttachStats) -> None:
self.stats.add_attached(vcol, stats)
if stats.has_findings():
logger.warning(
"Attached '%s' records found by %s: %s",
vcol,
selector,
stats.summary(),
)
else:
logger.debug(
"Attached '%s' records found by %s: %s", vcol, selector, stats.summary()
)
# ------------------------------------------------------------------
# Blank-edge resolution
# ------------------------------------------------------------------
def _resolve_blank_edges(self, gc: GraphContainer, conn_conf: DBConfig) -> None:
"""Extend edge lists for blank vertices after their keys are resolved."""
vc = self._db_aware_for(conn_conf).vertex_config
# Only edges between types this batch carries can be joined.
present = self.schema.core_schema.edge_config.ids_between(set(gc.vertices))
for vcol in vc.blank_vertices:
for edge_id in present:
vfrom, vto, _relation = edge_id
if vcol == vfrom or vcol == vto:
if edge_id not in gc.edges:
gc.edges[edge_id] = []
source_docs = gc.vertices[vfrom]
target_docs = gc.vertices[vto]
source_id_fields = vc.identity_fields(vfrom)
target_id_fields = vc.identity_fields(vto)
shared_fields = [
f for f in source_id_fields if f in target_id_fields
]
if shared_fields:
target_by_key: dict[tuple, list[dict]] = {}
for target_doc in target_docs:
key = tuple(target_doc.get(f) for f in shared_fields)
if any(item is None for item in key):
continue
target_by_key.setdefault(key, []).append(target_doc)
for source_doc in source_docs:
key = tuple(source_doc.get(f) for f in shared_fields)
if any(item is None for item in key):
continue
for target_doc in target_by_key.get(key, []):
gc.edges[edge_id].append((source_doc, target_doc, {}))
else:
gc.edges[edge_id].extend(
(x, y, {}) for x, y in zip(source_docs, target_docs)
)
# ------------------------------------------------------------------
# Extra weights
# ------------------------------------------------------------------
async def _enrich_extra_weights(
self,
gc: GraphContainer,
conn_conf: DBConfig,
resource,
*,
guard: FilterExpression | None = None,
) -> None:
"""Copy stored vertex fields onto the edges of the documents that name the vertex.
For each ``extra_weights`` rule, the weight vertex a document produced
is read back from the database, and its fields become attributes of
that document's edges. A document that produced several takes the first
one the database holds; one that produced none, or none the database
holds, keeps its edges as cast. A vertex *guard* hides is not read.
"""
rules = [
(entry.edge, weight)
for entry in resource.config.extra_weights
for weight in entry.vertex_weights
]
if self.dry or not rules:
return
vc = self._db_aware_for(conn_conf).vertex_config
def _sync() -> None:
with self._connection(conn_conf) as db:
for edge, weight in rules:
if not vc.has_vertex(weight.name):
logger.error(f"{weight.name} not a valid vertex")
continue
self._enrich_edges_from_stored_vertex(
db, gc, conn_conf, edge=edge, weight=weight, guard=guard
)
await asyncio.to_thread(_sync)
def _enrich_edges_from_stored_vertex(
self,
db: Any,
gc: GraphContainer,
conn_conf: DBConfig,
*,
edge: Edge,
weight: Weight,
guard: FilterExpression | None = None,
) -> None:
"""Apply one ``extra_weights`` rule to every document of *gc*."""
name = weight.name
source_fields = _weight_source_fields(weight)
if name is None or not source_fields:
return
vc = self._db_aware_for(conn_conf).vertex_config
keys = self._keys_for(conn_conf)
# One lookup for the batch; `owners` says which document asked for each key.
key_docs: list[dict[str, Any]] = []
owners: list[int] = []
pending: dict[int, list[tuple]] = {}
for position, item in enumerate(gc.linear):
triples = _document_edges(item, edge)
if not triples:
continue
pending[position] = triples
for doc in item.get(name, ()):
key_docs.append(doc)
owners.append(position)
if not key_docs:
return
resolved = db.resolve_vertices(
vc.vertex_dbname(name),
keys.vertex_docs(name, key_docs),
tuple(keys.vertex_fields(name, vc.identity_fields(name))),
tuple(keys.vertex_fields(name, source_fields)),
**_guard_kwargs(guard),
)
for index in sorted(resolved):
matches = resolved[index]
if not matches:
continue
# Popped, so a document's first stored vertex is the one it takes.
triples = pending.pop(owners[index], None)
if triples is None:
continue
attributes = _weight_attributes(
weight, keys.logical_vertex_doc(name, matches[0])
)
for triple in triples:
# The attribute dict is the one `gc.edges` holds, so the edge
# write sees it.
triple[2].update(attributes)
# ------------------------------------------------------------------
# Edges
# ------------------------------------------------------------------
async def _push_edges(
self,
gc: GraphContainer,
conn_conf: DBConfig,
resource: Any | None = None,
*,
guard: FilterExpression | None = None,
) -> None:
"""Insert all edges in *gc*.
Each key in ``gc.edges`` is a concrete ``(source, target, relation)``
triple produced by the extraction pipeline. We look up the matching
schema :class:`Edge` for each key (trying an exact match first, then a
``relation=None`` schema entry for dynamic-relation edges) and fire one
async task per key — one DB write per concrete relation, no inner loop.
Endpoints declared by a secondary identity are resolved to their primary
identity first, so the write itself stays a plain primary-key operation
on every backend; a vertex *guard* hides is never matched.
Raises:
AmbiguousEndpointError: on several matches under the ``error`` policy;
no edge of *gc* is inserted then.
"""
ctx = self.context(conn_conf)
vc = ctx.db_aware.vertex_config
core_ec = self.schema.core_schema.edge_config
def _schema_edge_for(edge_id: tuple) -> Edge | None:
"""Return the schema Edge for a gc edge key, or None if not declared."""
return core_ec.declared(edge_id)
endpoint_match_for = self._endpoint_match_lookup(resource)
async def _resolve_one(
edge_id: tuple, docs: list
) -> tuple[tuple, Edge, list] | None:
edge = _schema_edge_for(edge_id)
if edge is None:
self._report_undeclared_edge(edge_id, len(docs))
return None
endpoint_match = endpoint_match_for(edge_id)
if endpoint_match is None:
return edge_id, edge, docs
async with self._db_semaphore(conn_conf):
def _sync() -> list:
with self._connection(conn_conf) as db:
return self._resolve_endpoints(
db=db,
docs=docs,
edge=edge,
edge_id=edge_id,
match=endpoint_match,
ctx=ctx,
guard=guard,
)
edge_docs = await asyncio.to_thread(_sync)
return (edge_id, edge, edge_docs) if edge_docs else None
async def _insert_one(edge_id: tuple, edge: Edge, edge_docs: list) -> None:
source, target, relation = edge_id
# Projected before a write slot or a connection is taken.
stored = (
None
if self.dry
else await asyncio.to_thread(
project_edges, ctx, edge, relation, edge_docs
)
)
async with AsyncExitStack() as stack:
# Cypher relationship MERGE has the same concurrent
# check-then-create race as node MERGE. Two edge ids never MERGE
# one pattern (their labels or relation type differ), so the
# lock is per concrete id -- the container key, since a
# ``relation_field`` template's schema relation is None.
lock_key: tuple[str | None, ...] = (
("edge",)
if conn_conf.connection_type in _SERIAL_EDGE_FLAVORS
else (
"edge",
vc.vertex_dbname(source),
vc.vertex_dbname(target),
relation,
)
)
await self._acquire_write_slot(stack, conn_conf, lock_key)
def _sync() -> None:
with self._connection(conn_conf) as db:
if stored is not None:
db.insert_edges_batch(
docs_edges=stored.docs,
source_class=stored.source_class,
target_class=stored.target_class,
relation_name=stored.relation_name,
match_keys_source=stored.match_keys_source,
match_keys_target=stored.match_keys_target,
**self._edge_kwargs(stored, conn_conf),
)
await asyncio.to_thread(_sync)
# Every edge id resolves before any inserts, so an ``error`` raised for
# one leaves the others unwritten too.
resolved = await _gather_settled(
[_resolve_one(edge_id, docs) for edge_id, docs in gc.edges.items() if docs]
)
await asyncio.gather(
*[_insert_one(*entry) for entry in resolved if entry is not None]
)
def _edge_kwargs(self, stored: StoredEdges, conn_conf: DBConfig) -> dict[str, Any]:
"""The flavor's options for inserting *stored*.
Raises:
ValueError: If the edge has union fields the target cannot write as
a union, rather than writing a replace in its place.
"""
merge_props = stored.merge_properties
flavor = conn_conf.connection_type
edge_kw: dict[str, Any] = {
"filter_uniques": False,
"dry": self.dry,
"collection_name": stored.collection_name,
}
if flavor in CYPHER_FLAVORS or flavor in (DBType.POSTGRES, DBType.NEBULA):
# On PostgreSQL the conflict target must match the edge table's
# unique key, built from the same names; on NebulaGraph they rank
# the edge.
if merge_props is not None:
edge_kw["relationship_merge_properties"] = merge_props
elif flavor == DBType.ARANGO and merge_props is not None:
edge_kw["uniq_weight_fields"] = list(merge_props)
if stored.union_fields:
ConnectionManager.require(conn_conf, ConnectionCapability.UNION_MERGE)
edge_kw["union_fields"] = stored.union_fields
return edge_kw
def _report_undeclared_edge(self, edge_id: tuple, count: int) -> None:
"""Say, once per edge id, that its edges are not written."""
if edge_id in self._reported_undeclared_edges:
return
self._reported_undeclared_edges.add(edge_id)
logger.warning(
"Edge %s is not written: the schema does not declare it (%s in this "
"batch; reported once). Declare the edge, or name the declared "
"relation on the step that produces it.",
edge_id,
count,
)
def _endpoint_match_lookup(self, resource: Any | None) -> Any:
"""Return a lookup for a resource's endpoint identity selections.
Only edges that select a secondary identity have an entry, so edges
matched on the primary identity never touch the resolution path.
"""
registry = getattr(resource, "edge_derivation", None) if resource else None
if registry is None:
return lambda edge_id: None
return registry.endpoint_match_for
def _resolve_endpoints(
self,
*,
db: Any,
docs: list,
edge: Edge,
edge_id: tuple,
match: Any,
ctx: WriteContext,
guard: FilterExpression | None = None,
) -> list:
"""Map secondary-identity endpoints to primary identities before writing."""
vertex_config = ctx.db_aware.vertex_config
source_fields = vertex_config.match_fields(edge.source, match.source)
target_fields = vertex_config.match_fields(edge.target, match.target)
source_identity = vertex_config.identity_fields(edge.source)
target_identity = vertex_config.identity_fields(edge.target)
policy = match.on_ambiguous or self.ingestion_model.endpoints_on_ambiguous
keys = ctx.keys
# Resolution reads the database, so it runs on stored names; the
# resolved endpoints come back to logical names for the rest of the write.
stored_docs = docs
if keys.active:
stored_docs = [
(
keys.vertex_doc(edge.source, doc[0]),
keys.vertex_doc(edge.target, doc[1]),
*doc[2:],
)
for doc in docs
]
def _stored(vertex: str, fields: list[str]) -> list[str]:
return keys.vertex_fields(vertex, fields)
resolved, stats = resolve_edge_endpoints(
db,
stored_docs,
source_class=vertex_config.vertex_dbname(edge.source),
target_class=vertex_config.vertex_dbname(edge.target),
source_match_fields=_stored(edge.source, source_fields),
target_match_fields=_stored(edge.target, target_fields),
source_identity_fields=_stored(edge.source, source_identity),
target_identity_fields=_stored(edge.target, target_identity),
resolve_source=list(source_fields) != list(source_identity),
resolve_target=list(target_fields) != list(target_identity),
policy=policy,
guard=guard,
)
self.stats.add_endpoints(edge_id, stats)
if keys.active:
resolved = [
(
keys.logical_vertex_doc(edge.source, doc[0]),
keys.logical_vertex_doc(edge.target, doc[1]),
*doc[2:],
)
for doc in resolved
]
if stats.has_findings():
logger.warning(
"Edge %s endpoint resolution (policy=%s): %s",
edge_id,
policy,
stats.summary(),
)
else:
logger.debug("Edge %s endpoint resolution: %s", edge_id, stats.summary())
return resolved
def _db_semaphore(self, conn_conf: DBConfig) -> asyncio.Semaphore:
"""Shared semaphore so ``max_concurrent`` bounds the whole run.
Created lazily per event loop: a writer reused across separate
``asyncio.run`` calls must not carry a semaphore bound to a closed loop.
"""
loop = asyncio.get_running_loop()
bound = self.max_concurrent
if (
self._semaphore is None
or self._semaphore_loop is not loop
or self._semaphore_bound != bound
):
self._semaphore = asyncio.Semaphore(bound)
self._semaphore_loop = loop
self._semaphore_bound = bound
return self._semaphore
async def _acquire_write_slot(
self,
stack: AsyncExitStack,
conn_conf: DBConfig,
collection: tuple[str | None, ...],
) -> None:
"""Enter the semaphore and, when the backend's upsert can race with
itself (Cypher MERGE has no cross-transaction atomicity), a
per-collection lock so the same collection is written by one batch at a
time while distinct collections still proceed in parallel.
The lock is taken before the semaphore, so a write queued behind its
collection does not hold a ``max_concurrent`` slot while it waits.
"""
if conn_conf.connection_type not in _CONCURRENT_UPSERT_SAFE_FLAVORS:
loop = asyncio.get_running_loop()
if self._collection_locks_loop is not loop:
self._collection_locks = {}
self._collection_locks_loop = loop
lock = self._collection_locks.setdefault(collection, asyncio.Lock())
await stack.enter_async_context(lock)
await stack.enter_async_context(self._db_semaphore(conn_conf))
def context(self, conn_conf: DBConfig) -> WriteContext:
"""The schema resolved for *conn_conf*'s flavor, cached for the run.
Physical names are filled in, so the storage and relation names written
to are the ones DDL declared. Resolve it before fanning out concurrent
:meth:`write` calls, so they never race the cache fill.
"""
flavor = conn_conf.connection_type
if self._context is None or self._context.flavor != flavor:
self._context = WriteContext.for_flavor(self.schema, flavor)
return self._context
def _db_aware_for(self, conn_conf: DBConfig) -> SchemaDBAware:
return self.context(conn_conf).db_aware
def _keys_for(self, conn_conf: DBConfig) -> PhysicalKeys:
return self.context(conn_conf).keys