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
# A VID names its tag (`util.tag_vid`), so removing one removes one vertex.
supports_instance_delete: 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_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), unique=True
)
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(source_class, src_doc, match_keys_source)
dst_vid = make_vid(target_class, dst_doc, 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"
)
def delete_vertices(
self,
class_name: str,
key_docs: list[dict[str, Any]],
match_keys: tuple[str, ...],
*,
chunk_size: int = DEFAULT_DELETE_CHUNK_SIZE,
) -> None:
"""Remove vertices of one tag by VID, ``WITH EDGE`` removing their edges."""
vids = [
f'"{escape_nebula_string(make_vid(class_name, doc, match_keys))}"'
for doc in deletable_docs(key_docs, match_keys)
]
for start in range(0, len(vids), chunk_size):
chunk = ", ".join(vids[start : start + chunk_size])
self._execute(f"DELETE VERTEX {chunk} WITH EDGE")
def delete_edges(
self,
source_class: str,
target_class: str,
relation_name: str | None,
endpoints: list[tuple[dict[str, Any], dict[str, Any]]],
match_keys_source: tuple[str, ...],
match_keys_target: tuple[str, ...],
*,
collection_name: str | None = None,
chunk_size: int = DEFAULT_DELETE_CHUNK_SIZE,
) -> None:
"""Remove edges of one type between VID pairs.
Every edge GraFlo writes has rank 0, the rank ``DELETE EDGE`` removes
when none is given.
"""
if relation_name is None:
raise ValueError("NebulaGraph delete_edges requires the edge type")
pairs = [
f'"{escape_nebula_string(make_vid(source_class, s, match_keys_source))}"'
f'->"{escape_nebula_string(make_vid(target_class, t, match_keys_target))}"'
for s, t in deletable_endpoints(
endpoints, match_keys_source, match_keys_target
)
]
for start in range(0, len(pairs), chunk_size):
chunk = ", ".join(pairs[start : start + chunk_size])
self._execute(f"DELETE EDGE `{relation_name}` {chunk}")
# ------------------------------------------------------------------
# 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: 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:
"""Every identity value joined as the write path joins them, without the tag.
:meth:`fetch_edges` adds the tag it is given. 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 vertex_key(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``.
*from_id* and *to_id* are addresses (:meth:`vertex_address`), and
``_src`` / ``_dst`` in the rows are too: the tag is added and removed
here. 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
# Joined with the far-end predicate by AND in the query builder.
fc = render_conjunct(ff, doc_name="e", kind=self._expression_flavor())
q = fetch_edges_ngql(
from_type,
tag_vid(from_type, from_id),
edge_type=edge_type,
to_tag=to_type,
to_vid=tag_vid(to_type, to_id) if to_type and to_id else None,
to_key=to_id if to_id and not to_type else None,
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"] = key_of_vid(str(row.get("src", "")))
entry["_dst"] = key_of_vid(str(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(class_name, doc, match_keys)
try:
rs = self._execute(
f'FETCH PROP ON `{class_name}` "{escape_nebula_string(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