class NebulaConnection(Connection):
"""NebulaGraph implementation of the ``Connection`` interface.
Automatically selects the correct Python driver and query language based on
``config.version``:
* **v3.x** -- ``nebula3-python``, nGQL
* **v5.x** -- ``nebula5-python``, ISO GQL / Cypher
"""
flavor: ClassVar[DBType] = DBType.NEBULA
supports_schema_introspection: ClassVar[bool] = True
# Properties and their types come from `DESCRIBE`, but nGQL records no
# endpoint tags on an edge type, so those must be observed on real edges.
# Half-sampled is still sampled: an edge type with no stored edges is lost.
schema_introspection_is_sampled: ClassVar[bool] = True
def __init__(self, config: NebulaConfig):
super().__init__()
self.config = config
self._adapter: NebulaClientAdapter = create_adapter(config)
self._space_name: str | None = None
self._tag_fields: dict[str, list[str]] = {}
if config.schema_name:
# A missing space is legitimate here — define_schema creates it —
# so this is not fatal. But swallowing it silently also hides a bad
# credential or an unreachable graphd behind the first confusing
# error much later, so say what happened.
try:
self._use_space(config.schema_name)
except Exception as e:
logger.debug(
"Could not select space '%s' at connect time: %s. "
"It will be created by define_schema if missing.",
config.schema_name,
e,
)
# ------------------------------------------------------------------
# Expression flavour override (instance-level, depends on version)
# ------------------------------------------------------------------
@classmethod
def expression_flavor(cls) -> ExpressionFlavor:
return ExpressionFlavor.NGQL
def _expression_flavor(self) -> ExpressionFlavor:
"""Instance-level flavour dispatch."""
if self.config.is_v3:
return ExpressionFlavor.NGQL
return ExpressionFlavor.CYPHER
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
def _execute(self, statement: str) -> NebulaResultSet:
return self._adapter.execute(statement)
def _use_space(self, space_name: str) -> None:
self._adapter.use_space(space_name)
self._space_name = space_name
self._load_tag_fields()
def _load_tag_fields(self) -> None:
"""Discover existing tags and their fields from the current space."""
try:
rs = self._adapter.execute("SHOW TAGS")
tag_names = [r.get("Name", r.get("name", "")) for r in rs.rows_as_dicts()]
except Exception as e:
# Expected on a space that does not exist yet; anything else leaves
# _tag_fields empty, which downstream reads as "tag has no fields".
logger.debug("Could not list tags in space '%s': %s", self._space_name, e)
return
for tag in tag_names:
if not tag:
continue
try:
desc = self._adapter.execute(f"DESCRIBE TAG `{tag}`")
self._tag_fields[tag] = [
r.get("Field", r.get("field", "")) for r in desc.rows_as_dicts()
]
except Exception as e:
# SHOW TAGS listed it, so failing to describe it is not a
# missing-space race — the tag stays absent from _tag_fields and
# every write against it will build an empty property list.
logger.warning(
"Could not describe tag '%s' in space '%s': %s. "
"Writes to this tag may omit properties.",
tag,
self._space_name,
e,
)
def _wait_for_dml_ready(self, tag_name: str) -> None:
"""Wait until DML operations are possible on a tag.
NebulaGraph's storaged schema cache may lag behind graphd's metadata
cache by several heartbeat cycles (~10 s with default settings).
``DESCRIBE TAG`` succeeds immediately, but DML like ``FETCH PROP`` or
``UPSERT VERTEX`` fails until the storage cache is warm.
"""
import time
check = f'FETCH PROP ON `{tag_name}` "__dml_check__" YIELD properties(vertex)'
for attempt in range(_SCHEMA_WAIT_RETRIES):
try:
self._adapter.execute(check)
logger.debug(
"DML ready for tag '%s' after %d attempt(s)", tag_name, attempt + 1
)
return
except Exception:
if attempt == _SCHEMA_WAIT_RETRIES - 1:
logger.warning(
"DML readiness check for tag '%s' did not succeed "
"after %d attempts",
tag_name,
_SCHEMA_WAIT_RETRIES,
)
time.sleep(_SCHEMA_WAIT_INTERVAL)
def _wait_for_edge_dml_ready(self, edge_type: str) -> None:
"""Wait until DML operations are possible on an edge type."""
import time
check = (
f'FETCH PROP ON `{edge_type}` "__src__"->"__dst__" YIELD properties(edge)'
)
for attempt in range(_SCHEMA_WAIT_RETRIES):
try:
self._adapter.execute(check)
logger.debug(
"DML ready for edge '%s' after %d attempt(s)",
edge_type,
attempt + 1,
)
return
except Exception:
if attempt == _SCHEMA_WAIT_RETRIES - 1:
logger.warning(
"DML readiness check for edge '%s' did not succeed "
"after %d attempts",
edge_type,
_SCHEMA_WAIT_RETRIES,
)
time.sleep(_SCHEMA_WAIT_INTERVAL)
def _render_filter(self, filters: Any, doc_name: str) -> str:
if self.config.is_v3:
return render_filters_ngql(filters, doc_name)
return render_filters_cypher(filters, doc_name)
def _tag_field_names(self, tag_name: str) -> list[str]:
return self._tag_fields.get(tag_name, [])
# ------------------------------------------------------------------
# Connection ABC – lifecycle
# ------------------------------------------------------------------
def execute(self, query: str | Any, **kwargs: Any) -> Any:
rs = self._execute(str(query))
return rs
def close(self) -> None:
self._adapter.close()
# ------------------------------------------------------------------
# Database (space) management
# ------------------------------------------------------------------
def _ensure_storage_hosts(self) -> None:
"""Register storaged hosts if not already present (v3.x requirement)."""
if not self.config.storaged_addresses:
return
try:
rs = self._adapter.execute("SHOW HOSTS")
existing = {
f"{r.get('Host', '')}:{r.get('Port', '')}" for r in rs.rows_as_dicts()
}
except Exception:
existing = set()
for addr in self.config.storaged_addresses:
if addr not in existing:
try:
host, port = addr.rsplit(":", 1)
self._adapter.execute(f'ADD HOSTS "{host}":{port}')
logger.info("Registered storage host %s", addr)
except Exception:
logger.debug("ADD HOSTS %s (may already exist)", addr)
import time
for _ in range(30):
try:
rs = self._adapter.execute("SHOW HOSTS")
statuses = [r.get("Status", "") for r in rs.rows_as_dicts()]
if statuses and all(s == "ONLINE" for s in statuses):
return
except Exception:
pass
time.sleep(1)
logger.warning("Storage hosts may not all be ONLINE yet")
def create_database(self, name: str) -> None:
self._ensure_storage_hosts()
stmt = create_space_ngql(
name,
vid_type=self.config.vid_type,
partition_num=self.config.partition_num,
replica_factor=self.config.replica_factor,
)
self._execute(stmt)
wait_for_space_ready(
self._adapter,
name,
max_retries=_SCHEMA_WAIT_RETRIES,
interval=_SCHEMA_WAIT_INTERVAL,
)
self._use_space(name)
logger.info("Created NebulaGraph space '%s'", name)
def delete_database(self, name: str) -> None:
self._execute(drop_space_ngql(name))
logger.info("Dropped NebulaGraph space '%s'", name)
# ------------------------------------------------------------------
# Schema definition
# ------------------------------------------------------------------
def define_schema(self, schema: Schema) -> None:
assert_schema_field_types_supported(self.flavor, schema)
self.define_vertex_classes(schema)
edges = list(schema.core_schema.edge_config.values())
self.define_edge_classes(edges)
def define_vertex_classes(self, schema: Schema) -> None:
"""Create one tag per vertex type, named as the rest of this class reads it.
The tag name comes from the db-aware projection, not the logical vertex
name: reads, writes and ``clear_data`` all resolve through
``vertex_dbname``, so declaring a tag under the logical name instead
leaves a schema nothing writes to, and every UPSERT fails with
``No schema found`` on any manifest that sets ``vertex_storage_names``.
"""
logical = schema.core_schema.vertex_config
db_aware = schema.resolve_db_aware(DBType.NEBULA).vertex_config
for vname in logical.vertex_set:
tag = db_aware.vertex_dbname(vname)
fields = logical.properties(vname)
stmt = create_tag_ngql(tag, fields)
self._execute(stmt)
self._tag_fields[tag] = [f.name for f in fields]
logger.debug("Created tag '%s' for vertex '%s'", tag, vname)
if logical.vertex_set:
sample = next(iter(logical.vertex_set))
self._wait_for_dml_ready(db_aware.vertex_dbname(sample))
def define_edge_classes(self, edges: list[Edge]) -> None:
created: set[str] = set()
for edge in edges:
rel = edge.relation or f"{edge.source}_{edge.target}"
if rel in created:
continue
edge_fields = []
if edge.properties:
edge_fields = list(edge.properties)
stmt = create_edge_type_ngql(rel, edge_fields)
self._execute(stmt)
created.add(rel)
logger.debug("Created edge type '%s'", rel)
if created:
sample_et = next(iter(created))
self._wait_for_edge_dml_ready(sample_et)
# ------------------------------------------------------------------
# Index management
# ------------------------------------------------------------------
def define_vertex_indexes(
self, vertex_config: VertexConfig, schema: Schema | None = None
) -> None:
if schema is None:
logger.warning(
"Schema is None: identity indexes cannot be ensured without schema"
)
# Index the tag under the name it was actually created with; the
# db_profile lookups below stay keyed on the logical name.
db_aware = (
schema.resolve_db_aware(DBType.NEBULA).vertex_config
if schema is not None
else None
)
for vname in vertex_config.vertex_set:
tag = db_aware.vertex_dbname(vname) if db_aware is not None else vname
fields = vertex_config.properties(vname)
string_fields = {f.name for f in fields if is_nebula_string_field(f)}
index_list = (
schema.db_profile.vertex_secondary_indexes(vname)
if schema is not None
else []
)
# Nebula requires TAG indexes for LOOKUP and many property-filtered MATCH
# plans. Keep identity index creation implicit so schemas without
# explicit database_features remain queryable/clearable.
identity_idx = Index(fields=vertex_config.identity_fields(vname))
all_indexes = [identity_idx, *index_list]
seen: set[tuple[str, ...]] = set()
for idx in all_indexes:
key = tuple(str(f) for f in idx.fields)
if not key or key in seen:
continue
seen.add(key)
self._add_tag_index(tag, idx, string_fields=string_fields)
def define_edge_indexes(
self, edges: list[Edge], schema: Schema | None = None
) -> None:
for edge in edges:
rel = edge.relation or f"{edge.source}_{edge.target}"
index_list = (
schema.db_profile.edge_secondary_indexes(edge.edge_id)
if schema is not None
else []
)
for idx in index_list:
self._add_edge_index(rel, idx)
def _add_tag_index(
self,
tag_name: str,
index: Index,
string_fields: set[str] | None = None,
) -> None:
idx_fields = [str(f) for f in index.fields]
idx_name = f"idx_{tag_name}_{'_'.join(idx_fields)}"
stmt = create_tag_index_ngql(
idx_name, tag_name, idx_fields, string_fields=string_fields
)
try:
self._execute(stmt)
self._rebuild_index(idx_name, kind="TAG")
logger.debug("Created tag index '%s'", idx_name)
except Exception as e:
# Not debug: without this index every filtered read on the tag fails
# with IndexNotFound, and the cause is three layers away.
logger.warning(
"Tag index '%s' on '%s' was not created: %s. Filtered reads on "
"these fields will fail with IndexNotFound.",
idx_name,
tag_name,
e,
)
def _add_edge_index(self, edge_type: str, index: Index) -> None:
idx_fields = [str(f) for f in index.fields]
idx_name = f"idx_{edge_type}_{'_'.join(idx_fields)}"
stmt = create_edge_index_ngql(idx_name, edge_type, idx_fields)
try:
self._execute(stmt)
self._rebuild_index(idx_name, kind="EDGE")
logger.debug("Created edge index '%s'", idx_name)
except Exception as e:
logger.warning(
"Edge index '%s' on '%s' was not created: %s. Filtered reads on "
"these fields will fail with IndexNotFound.",
idx_name,
edge_type,
e,
)
def _rebuild_index(self, idx_name: str, kind: str = "TAG") -> None:
"""Rebuild an index, waiting for propagation first, then for completion."""
import time
rebuild_stmt = f"REBUILD {kind} INDEX `{idx_name}`"
for attempt in range(_SCHEMA_WAIT_RETRIES):
try:
self._adapter.execute(rebuild_stmt)
break
except Exception:
if attempt == _SCHEMA_WAIT_RETRIES - 1:
logger.warning("Could not start rebuild for '%s'", idx_name)
return
time.sleep(_SCHEMA_WAIT_INTERVAL)
for _ in range(_SCHEMA_WAIT_RETRIES):
try:
rs = self._adapter.execute(f"SHOW {kind} INDEX STATUS")
for row in rs.rows_as_dicts():
name = row.get("Name", row.get("name", ""))
status = row.get("Index Status", row.get("index_status", ""))
if name == idx_name and status.upper() == "FINISHED":
return
except Exception:
pass
time.sleep(_SCHEMA_WAIT_INTERVAL)
logger.warning("Index rebuild for '%s' may not be complete", idx_name)
# ------------------------------------------------------------------
# init_db
# ------------------------------------------------------------------
def _resolve_space_name(self, schema: Schema) -> str:
space_name = self.config.schema_name
if not space_name:
space_name = schema.effective_namespace(DBType.NEBULA)
self.config.schema_name = space_name
return space_name
def _space_exists(self, space_name: str) -> bool:
try:
rs = self._adapter.execute("SHOW SPACES")
names = {row.get("Name", row.get("name", "")) for row in rs.rows_as_dicts()}
return space_name in names
except Exception:
try:
self._adapter.use_space(space_name)
return True
except Exception:
return False
def ensure_target_namespace(self, schema: Schema, *, create: bool) -> None:
"""Ensure the NebulaGraph space exists."""
space_name = self._resolve_space_name(schema)
if self._space_exists(space_name):
wait_for_space_ready(
self._adapter,
space_name,
max_retries=_SCHEMA_WAIT_RETRIES,
interval=_SCHEMA_WAIT_INTERVAL,
)
self._use_space(space_name)
return
if not create:
raise NamespaceNotFoundError(
f"NebulaGraph space '{space_name}' does not exist. "
"Create it manually or call with create_namespace=True."
)
self.create_database(space_name)
def _schema_has_artifacts(self) -> bool:
try:
rs = self._execute("SHOW TAGS")
return bool(rs.rows_as_dicts())
except Exception:
return False
def apply_target_schema(
self,
schema: Schema,
*,
recreate: bool,
create_namespace: bool = True,
) -> None:
"""Define tags, edge types, and indexes in the current space."""
self.report_edge_direction_support(schema)
space_name = self._resolve_space_name(schema)
if recreate:
if create_namespace:
try:
self.delete_database(space_name)
except Exception:
pass
self.create_database(space_name)
else:
try:
rs = self._execute("SHOW TAGS")
for row in rs.rows_as_dicts():
tag = row.get("Name", row.get("name", ""))
if tag:
self._execute(f"DROP TAG IF EXISTS `{tag}`")
except Exception as e:
logger.warning("Failed to drop tags during recreate: %s", e)
elif self._schema_has_artifacts():
raise SchemaExistsError(
f"Schema already exists in space '{space_name}'. "
"Set recreate_schema=True to replace."
)
self.define_schema(schema)
self.define_indexes(schema)
def init_db(
self,
schema: Schema,
recreate_schema: bool = False,
*,
create_namespace: bool = True,
) -> None:
"""Convenience wrapper: ensure space then apply schema."""
self.ensure_target_namespace(schema, create=create_namespace)
self.apply_target_schema(
schema, recreate=recreate_schema, create_namespace=create_namespace
)
# ------------------------------------------------------------------
# Data clearing
# ------------------------------------------------------------------
def clear_data(self, schema: Schema) -> None:
vc = schema.resolve_db_aware(DBType.NEBULA).vertex_config
for vertex in vc.vertex_set:
vname = vc.vertex_dbname(vertex)
try:
self._execute(
f"LOOKUP ON `{vname}` YIELD id(vertex) AS vid "
f"| DELETE VERTEX $-.vid"
)
except Exception as e:
logger.debug("clear_data for tag '%s': %s", vname, e)
def delete_graph_structure(
self,
vertex_types: tuple[str, ...] | list[str] = (),
graph_names: tuple[str, ...] | list[str] = (),
delete_all: bool = False,
) -> None:
if delete_all:
space_name = self._space_name or self.config.schema_name
if space_name:
self.delete_database(space_name)
return
for vt in vertex_types:
try:
self._execute(f"DROP TAG IF EXISTS `{vt}`")
except Exception as e:
logger.warning("Failed to drop tag '%s': %s", vt, e)
for gn in graph_names:
try:
self._execute(drop_space_ngql(gn))
except Exception as e:
logger.warning("Failed to drop space '%s': %s", gn, e)
# ------------------------------------------------------------------
# Document operations
# ------------------------------------------------------------------
def upsert_docs_batch(
self,
docs: list[dict[str, Any]],
class_name: str,
match_keys: list[str] | tuple[str, ...],
**kwargs: Any,
) -> None:
dry = kwargs.pop("dry", False)
if not docs:
return
match_keys_list = list(match_keys)
tag_fields = self._tag_field_names(class_name)
if not tag_fields:
tag_fields = list({k for doc in docs for k in doc})
statements = batch_upsert_vertices_ngql(
class_name, docs, match_keys_list, tag_fields
)
if dry or not statements:
return
# Execute in batches to avoid hitting statement-size limits
batch_size = 50
for i in range(0, len(statements), batch_size):
chunk = statements[i : i + batch_size]
combined = "; ".join(chunk)
self._execute(combined)
def insert_edges_batch(
self,
docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
source_class: str,
target_class: str,
relation_name: str,
match_keys_source: tuple[str, ...],
match_keys_target: tuple[str, ...],
filter_uniques: bool = True,
head: int | None = None,
**kwargs: Any,
) -> None:
opts = consume_insert_edges_kwargs(kwargs)
dry = opts.dry
if not docs_edges:
return
if head is not None:
docs_edges = docs_edges[:head]
# Build (src_vid, dst_vid, props) tuples
edge_tuples: list[tuple[str, str, dict[str, Any]]] = []
for edge_doc in docs_edges:
if not isinstance(edge_doc, (list, tuple)) or len(edge_doc) < 2:
continue
src_doc = edge_doc[0] if isinstance(edge_doc[0], dict) else {}
dst_doc = edge_doc[1] if isinstance(edge_doc[1], dict) else {}
props = (
edge_doc[2]
if len(edge_doc) > 2 and isinstance(edge_doc[2], dict)
else {}
)
src_vid = make_vid(src_doc, list(match_keys_source))
dst_vid = make_vid(dst_doc, list(match_keys_target))
edge_tuples.append((src_vid, dst_vid, props))
if dry or not edge_tuples:
return
# Determine edge property fields from schema or from data
all_prop_keys: set[str] = set()
for _, _, p in edge_tuples:
all_prop_keys.update(p.keys())
edge_fields = sorted(all_prop_keys) if all_prop_keys else None
batch_size = 200
for i in range(0, len(edge_tuples), batch_size):
chunk = edge_tuples[i : i + batch_size]
stmt = insert_edges_ngql(relation_name, chunk, edge_fields)
if stmt:
self._execute(stmt)
def insert_return_batch(
self, docs: list[dict[str, Any]], class_name: str
) -> list[dict[str, Any]] | str:
raise NotImplementedError(
"insert_return_batch is not implemented for NebulaGraph"
)
# ------------------------------------------------------------------
# Fetch operations
# ------------------------------------------------------------------
def fetch_docs(
self,
class_name: str,
filters: list[Any] | dict[str, Any] | None = None,
limit: int | None = None,
return_keys: list[str] | None = None,
unset_keys: list[str] | None = None,
**kwargs: Any,
) -> list[dict[str, Any]]:
if self.config.is_v3:
doc_name = f"v.`{class_name}`"
fc = self._render_filter(filters, doc_name)
q = fetch_docs_ngql(class_name, fc, limit, return_keys)
else:
fc = self._render_filter(filters, "v")
q = fetch_docs_gql(class_name, fc, limit, return_keys)
rs = self._execute(q)
rows = rs.rows_as_dicts()
if return_keys:
return rows
result: list[dict[str, Any]] = []
for row in rows:
v = row.get("v", row)
if isinstance(v, dict) and "tags" in v:
props: dict[str, Any] = {}
for tag_props in v["tags"].values():
props.update(tag_props)
result.append(props)
elif isinstance(v, dict):
result.append(v)
else:
result.append(row)
return result
def _describe_properties(
self, kind: str, name: str
) -> tuple[list[str], dict[str, FieldType]]:
"""Read a tag's or edge type's declared columns via ``DESCRIBE``."""
try:
rows = self._execute(f"DESCRIBE {kind} `{name}`").rows_as_dicts()
except Exception:
logger.debug("DESCRIBE %s `%s` failed", kind, name, exc_info=True)
return [], {}
properties: list[str] = []
types: dict[str, FieldType] = {}
for row in rows:
field = row.get("Field") or row.get("field")
if not field:
continue
properties.append(field)
declared = field_type_from_nebula(row.get("Type") or row.get("type"))
if declared is not None:
types[field] = declared
return properties, types
def _sample_edge_endpoints(
self, edge_type: str, sample_limit: int
) -> list[tuple[str, str]]:
"""Sample the ``(source_tag, target_tag)`` pairs an edge type connects.
An nGQL edge type carries no endpoint tags in its DDL -- unlike a
TigerGraph ``CREATE ... EDGE``, which names them -- so the pairs have to
be observed on real edges. An edge type present in the catalogue but with
no stored edges yields nothing and is dropped, which is the honest
outcome: there is no way to tell what it would connect.
"""
query = (
f"MATCH (a)-[e:`{edge_type}`]->(b) "
f"RETURN DISTINCT tags(a) AS s, tags(b) AS t LIMIT {int(sample_limit)}"
)
try:
rows = self._execute(query).rows_as_dicts()
except Exception:
logger.debug("endpoint sampling failed for `%s`", edge_type, exc_info=True)
return []
pairs: list[tuple[str, str]] = []
for row in rows:
source, target = row.get("s"), row.get("t")
if not isinstance(source, list) or not isinstance(target, list):
continue
if not source or not target:
continue
pair = (str(source[0]), str(target[0]))
if pair not in pairs:
pairs.append(pair)
return pairs
def introspect_graph_schema(
self,
schema_name: str | None = None,
*,
sample_limit: int = 100,
) -> Schema:
"""Infer a graflo Schema from NebulaGraph's tag and edge-type catalogue.
Properties and their declared types come from ``DESCRIBE``, so they are
complete rather than sampled. Edge *endpoints* are the exception -- nGQL
does not record them -- so those are sampled, and ``sample_limit`` bounds
only that half.
"""
from graflo.db.graph_introspection import (
GraphEdgeIntrospection,
GraphIntrospectionResult,
GraphSchemaInferencer,
GraphVertexIntrospection,
infer_identity_fields,
)
resolved_name = (
schema_name or self._space_name or self.config.schema_name or "nebula"
)
try:
tag_rows = self._execute("SHOW TAGS").rows_as_dicts()
except Exception as error:
raise RuntimeError(
f"Cannot introspect NebulaGraph space {resolved_name!r}: SHOW TAGS failed"
) from error
vertices: list[GraphVertexIntrospection] = []
for row in tag_rows:
tag = row.get("Name") or row.get("name")
if not tag:
continue
properties, types = self._describe_properties("TAG", tag)
vertices.append(
GraphVertexIntrospection(
name=tag,
properties=properties,
identity=infer_identity_fields(properties),
property_types=types,
)
)
try:
edge_rows = self._execute("SHOW EDGES").rows_as_dicts()
except Exception:
logger.debug("SHOW EDGES failed", exc_info=True)
edge_rows = []
edges: list[GraphEdgeIntrospection] = []
for row in edge_rows:
edge_type = row.get("Name") or row.get("name")
if not edge_type:
continue
properties, types = self._describe_properties("EDGE", edge_type)
for source, target in self._sample_edge_endpoints(edge_type, sample_limit):
edges.append(
GraphEdgeIntrospection(
source=source,
target=target,
relation=edge_type,
properties=properties,
property_types=types,
collection_name=edge_type,
)
)
introspection = GraphIntrospectionResult(
name=resolved_name,
vertices=vertices,
edges=edges,
sample_limit=sample_limit,
)
return GraphSchemaInferencer(db_flavor=DBType.NEBULA).infer_schema(
introspection, schema_name=resolved_name
)
def vertex_address(
self, doc: dict[str, Any], identity_fields: Sequence[str]
) -> str | None:
"""Compose the VID exactly as the write path does.
Nebula addresses a vertex by VID, and :func:`make_vid` joins *all*
identity-field values with ``::``. Falling back to the base
implementation would address a composite-identity vertex by its first
field alone — a VID that exists nowhere, so every edge query anchored
on it returns empty instead of raising.
"""
keys = list(identity_fields)
if not keys or any(doc.get(k) is None for k in keys):
return None
return make_vid(doc, keys)
def fetch_edges(
self,
from_type: str,
from_id: str,
edge_type: str | None = None,
to_type: str | None = None,
to_id: str | None = None,
filters: list[Any] | dict[str, Any] | None = None,
limit: int | None = None,
return_keys: list[str] | None = None,
unset_keys: list[str] | None = None,
direction: EdgeDirection = EdgeDirection.OUT,
**kwargs: Any,
) -> list[dict[str, Any]]:
"""Fetch edges incident to one vertex via ``GO``.
Nebula keeps an out-key and an in-key per edge, so ``IN`` / ``ANY`` are
as cheap as ``OUT`` once ``REVERSELY`` / ``BIDIRECT`` is asked for.
"""
fc = ""
if filters is not None:
if not isinstance(filters, FilterExpression):
ff = FilterExpression.from_dict(filters)
else:
ff = filters
fc = str(ff(doc_name="e", kind=self._expression_flavor()))
q = fetch_edges_ngql(
from_type,
from_id,
edge_type=edge_type,
to_tag=to_type,
to_vid=to_id,
filter_clause=fc,
limit=limit,
direction=direction,
)
rs = self._execute(q)
rows = rs.rows_as_dicts()
result: list[dict[str, Any]] = []
for row in rows:
entry = row.get("props", row)
if isinstance(entry, dict):
entry["_src"] = row.get("src", "")
entry["_dst"] = row.get("dst", "")
entry["_type"] = row.get("edge_type", "")
if return_keys and isinstance(entry, dict):
entry = {k: entry.get(k) for k in return_keys}
result.append(entry)
return result
# ------------------------------------------------------------------
# Presence / absence checks
# ------------------------------------------------------------------
def fetch_present_documents(
self,
batch: list[dict[str, Any]],
class_name: str,
match_keys: list[str] | tuple[str, ...],
keep_keys: list[str] | tuple[str, ...] | None = None,
flatten: bool = False,
filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
if not batch:
return []
results: list[dict[str, Any]] = []
for doc in batch:
vid = make_vid(doc, list(match_keys))
try:
rs = self._execute(
f'FETCH PROP ON `{class_name}` "{vid}" '
f"YIELD properties(vertex) AS props"
)
rows = rs.rows_as_dicts()
for row in rows:
props = row.get("props", row)
if isinstance(props, dict):
if keep_keys:
props = {k: props.get(k) for k in keep_keys}
results.append(props)
except Exception as e:
logger.debug("fetch_present_documents error for vid '%s': %s", vid, e)
return results
def keep_absent_documents(
self,
batch: list[dict[str, Any]],
class_name: str,
match_keys: list[str] | tuple[str, ...],
keep_keys: list[str] | tuple[str, ...] | None = None,
filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
if not batch:
return []
present = self.fetch_present_documents(
batch, class_name, match_keys, list(match_keys), filters=filters
)
present_keys: set[tuple[Any, ...]] = set()
for doc in present:
key_tuple = tuple(doc.get(k) for k in match_keys)
present_keys.add(key_tuple)
absent: list[dict[str, Any]] = []
for doc in batch:
key_tuple = tuple(doc.get(k) for k in match_keys)
if key_tuple not in present_keys:
if keep_keys:
absent.append({k: doc.get(k) for k in keep_keys})
else:
absent.append(doc)
return absent
# ------------------------------------------------------------------
# Aggregation
# ------------------------------------------------------------------
def aggregate(
self,
class_name: str,
aggregation_function: AggregationType,
discriminant: str | None = None,
aggregated_field: str | None = None,
filters: FilterExpression | list[Any] | dict[str, Any] | None = None,
) -> int | float | list[dict[str, Any]] | dict[str, int | float] | None:
agg_name = (
aggregation_function.value
if isinstance(aggregation_function, AggregationType)
else str(aggregation_function)
)
if agg_name == "AVERAGE":
agg_name = "AVG"
if self.config.is_v3:
doc_name = f"v.`{class_name}`"
fc = self._render_filter(filters, doc_name)
q = aggregate_ngql(class_name, agg_name, discriminant, aggregated_field, fc)
else:
fc = self._render_filter(filters, "v")
q = aggregate_gql(class_name, agg_name, discriminant, aggregated_field, fc)
rs = self._execute(q)
rows = rs.rows_as_dicts()
if agg_name == "COUNT" and discriminant:
return {row["key"]: row["count"] for row in rows}
if agg_name == "COUNT":
return rows[0]["count"] if rows else 0
if agg_name == "SORTED_UNIQUE":
return [row["val"] for row in rows]
if rows:
return rows[0].get("val")
return None