From 29691719faa0da3f5a790d6f24add039874ef79f Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Fri, 30 Oct 2026 17:25:06 -0700 Subject: [PATCH] feat: manage property graphs defined over tables Replaces #4322, keeping only its property graph defined over tables. A property graph maps node labels and relationship types to existing tables of its namespace, in the shape of SQL/PGQ's `CREATE PROPERTY GRAPH`. It stores only that definition: a query reads the tables as nodes and edges through SQL `GRAPH_TABLE`, as they are when it runs. - **Rust:** `Connection::create_property_graph` takes a `PropertyGraphDefinition` of `NodeTable`s and `EdgeTable`s. `describe_property_graph`, `drop_property_graph` and `list_property_graphs` manage it. - **Python:** `lancedb.graph` has `NodeTable`, `EdgeTable` and `Endpoint`. Sync and async connections gain `create_property_graph`, `open_property_graph`, `drop_property_graph` and `list_property_graphs`, and a `PropertyGraph` handle describes its graph. - **Support:** remote connections only. Local connections raise `NotImplementedError` in Python and return `NotSupported` in Rust. ```python import lancedb from lancedb.graph import EdgeTable, Endpoint, NodeTable db = lancedb.connect("db://demo", api_key="...") db.create_property_graph( "social", nodes=[NodeTable("person", key="person_id", label="Person")], edges=[ EdgeTable( "knows", label="KNOWS", source=Endpoint("src_id", references=("person", "person_id")), destination=Endpoint("dst_id", references=("person", "person_id")), ) ], ) db.execute_query(""" SELECT * FROM GRAPH_TABLE (social MATCH (p:Person)-[:KNOWS]->(f:Person) WHERE p.age > 30 COLUMNS (p.name AS person, f.name AS friend)) """).read_all() ``` --- docs/src/python/python.md | 33 +++ python/python/lancedb/__init__.py | 3 + python/python/lancedb/_lancedb.pyi | 19 ++ python/python/lancedb/db.py | 216 +++++++++++++++ python/python/lancedb/graph.py | 256 ++++++++++++++++++ python/python/lancedb/remote/db.py | 56 ++++ python/python/tests/test_remote_db.py | 224 ++++++++++++++++ python/src/connection.rs | 74 +++++ rust/lancedb/src/connection.rs | 138 +++++++++- rust/lancedb/src/database.rs | 51 ++++ rust/lancedb/src/graph.rs | 281 +++++++++++++++++++ rust/lancedb/src/lib.rs | 1 + rust/lancedb/src/remote/db.rs | 372 +++++++++++++++++++++++++- rust/lancedb/src/utils/mod.rs | 15 ++ 14 files changed, 1735 insertions(+), 4 deletions(-) create mode 100644 python/python/lancedb/graph.py create mode 100644 rust/lancedb/src/graph.rs diff --git a/docs/src/python/python.md b/docs/src/python/python.md index e737223bf..60746a0f0 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -493,6 +493,31 @@ automatic scheduling requests. ::: lancedb.view.ViewDescription +## Property Graphs + +A property graph is defined over tables of its namespace: which tables hold +nodes and which hold edges, each node table's key and label, and the node keys +each edge table's source and destination columns reference -- the shape of +SQL/PGQ's `CREATE PROPERTY GRAPH`. It stores no rows of its own; a query reads +its tables as nodes and edges through SQL `GRAPH_TABLE`. + +Create one with +[create_property_graph][lancedb.db.DBConnection.create_property_graph], and +manage it with +[open_property_graph][lancedb.db.DBConnection.open_property_graph], +[list_property_graphs][lancedb.db.DBConnection.list_property_graphs] and +[drop_property_graph][lancedb.db.DBConnection.drop_property_graph]. + +::: lancedb.graph.NodeTable + +::: lancedb.graph.EdgeTable + +::: lancedb.graph.Endpoint + +::: lancedb.graph.PropertyGraph + +::: lancedb.graph.PropertyGraphDescription + ## Expressions Type-safe expression builder for filters and projections. Use these instead @@ -840,6 +865,14 @@ Table hold your actual data as a collection of records / rows. ::: lancedb.materialized_view.AsyncMaterializedView +## Property Graphs (Asynchronous) + +`AsyncConnection` returns this handle from +[create_property_graph][lancedb.db.AsyncConnection.create_property_graph] and +[open_property_graph][lancedb.db.AsyncConnection.open_property_graph]. + +::: lancedb.graph.AsyncPropertyGraph + ## Indices (Asynchronous) Indices can be created on a table to speed up queries. This section diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 4c38f374c..623279bf0 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -48,6 +48,7 @@ from .materialized_view import ( MaterializedViewDefinition, ) from .view import ViewDescription as ViewDescription +from .graph import PropertyGraph, PropertyGraphDescription from .table import AsyncTable, CompactionOptions, Table from .types import BaseTokenizerType from ._lancedb import Session @@ -603,6 +604,8 @@ __all__ = [ "MaterializedView", "MaterializedViewDefinition", "ViewDescription", + "PropertyGraph", + "PropertyGraphDescription", "connect", "connect_async", "tokenize", diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 4c9ad38de..239e80cdf 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -225,6 +225,25 @@ class Connection(object): page_token: Optional[str] = None, page_limit: Optional[int] = None, ) -> NameListing: ... + async def create_property_graph( + self, + name: str, + definition_json: str, + namespace_path: Optional[List[str]] = None, + ) -> str: ... + async def describe_property_graph( + self, name: str, namespace_path: Optional[List[str]] = None + ) -> str: ... + async def drop_property_graph( + self, name: str, namespace_path: Optional[List[str]] = None + ) -> None: ... + def list_property_graphs( + self, + namespace_path: Optional[List[str]] = None, + *, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> NameListing: ... def list_jobs( self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None ) -> JobListing: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index a93444a9e..b9aae1a29 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -59,6 +59,13 @@ from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError from . import __version__ from ._lancedb import connect as lancedb_connect # type: ignore from .functions import FunctionVersion, UdfDefinition +from .graph import ( + AsyncPropertyGraph, + EdgeTable, + NodeTable, + PropertyGraph, + _definition_json, +) from .job import AsyncJob, Job, _typed_job from .storage_integration import SecretReference, StorageIntegrationInfo from .sql import AsyncQuery as AsyncSqlQuery @@ -1222,6 +1229,99 @@ class DBConnection(EnforceOverrides): "View operations are not supported for this connection type" ) + def create_property_graph( + self, + name: str, + nodes: Sequence[NodeTable], + edges: Sequence[EdgeTable] = (), + *, + namespace_path: Optional[List[str]] = None, + ) -> PropertyGraph: + """Create a property graph over tables in its namespace. + + Each [NodeTable][lancedb.graph.NodeTable] names the table holding one + label's nodes and its key column, and each + [EdgeTable][lancedb.graph.EdgeTable] the table holding one label's + edges and the node keys its source and destination columns reference + -- the shape of SQL/PGQ's ``CREATE PROPERTY GRAPH``. The graph stores + nothing else: a query reads its tables through SQL ``GRAPH_TABLE``, as + they are when it runs. + + The database checks the definition against the tables now, so one they + cannot satisfy is refused here rather than at the first query. A name + another property graph holds is an error. Local connections raise + ``NotImplementedError``. + + >>> import lancedb + >>> from lancedb.graph import EdgeTable, Endpoint, NodeTable + >>> db = lancedb.connect("db://my_database") # doctest: +SKIP + >>> graph = db.create_property_graph( + ... "social", + ... nodes=[NodeTable("person", key="person_id", label="Person")], + ... edges=[ + ... EdgeTable( + ... "knows", + ... label="KNOWS", + ... source=Endpoint("src_id", references=("person", "person_id")), + ... destination=Endpoint( + ... "dst_id", references=("person", "person_id") + ... ), + ... ) + ... ], + ... ) # doctest: +SKIP + >>> graph.describe().nodes[0].properties # doctest: +SKIP + ['person_id', 'name'] + """ + raise NotImplementedError( + "Property graph operations are not supported for this connection type" + ) + + def open_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> PropertyGraph: + """Open the property graph named ``name``. + + The graph is described to check that it exists, so a name no property + graph holds raises [HttpError][lancedb.remote.errors.HttpError] with + ``status_code`` 404. Local connections raise ``NotImplementedError``. + """ + raise NotImplementedError( + "Property graph operations are not supported for this connection type" + ) + + def drop_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> None: + """Drop a property graph and wait for its definition to be deleted. + + The tables it reads are untouched: a property graph holds no rows of + its own. Local connections raise ``NotImplementedError``. + """ + raise NotImplementedError( + "Property graph operations are not supported for this connection type" + ) + + def list_property_graphs( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[str]: + """The names of the property graphs in one namespace. + + Names only; a definition comes from + [PropertyGraph.describe][lancedb.graph.PropertyGraph.describe]. Local + connections raise ``NotImplementedError``. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + raise NotImplementedError( + "Property graph operations are not supported for this connection type" + ) + def open_job(self, job_id: str) -> Job: """Open a server-side job by id, returning a handle with its record already populated. @@ -2285,6 +2385,61 @@ class LanceDBConnection(DBConnection): ) ) + @override + def create_property_graph( + self, + name: str, + nodes: Sequence[NodeTable], + edges: Sequence[EdgeTable] = (), + *, + namespace_path: Optional[List[str]] = None, + ) -> PropertyGraph: + return PropertyGraph( + LOOP.run( + self._conn.create_property_graph( + name, nodes, edges, namespace_path=namespace_path + ) + ) + ) + + @override + def open_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> PropertyGraph: + return PropertyGraph( + LOOP.run( + self._conn.open_property_graph(name, namespace_path=namespace_path) + ) + ) + + @override + def drop_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> None: + LOOP.run(self._conn.drop_property_graph(name, namespace_path=namespace_path)) + + @override + def list_property_graphs( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[str]: + """List property graphs. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return Listing( + self._conn.list_property_graphs( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ) + ) + @override def list_jobs( self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None @@ -3597,6 +3752,67 @@ class AsyncConnection(object): ) ) + async def create_property_graph( + self, + name: str, + nodes: Sequence[NodeTable], + edges: Sequence[EdgeTable] = (), + *, + namespace_path: Optional[List[str]] = None, + ) -> AsyncPropertyGraph: + """Create a property graph over tables in its namespace. + + See + [DBConnection.create_property_graph][lancedb.DBConnection.create_property_graph]. + """ + definition = _definition_json(nodes, edges) + graph = AsyncPropertyGraph(self._inner, name, namespace_path) + await self._inner.create_property_graph(name, definition, graph.namespace_path) + return graph + + async def open_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> AsyncPropertyGraph: + """Open the property graph named ``name``, describing it to check that + it exists. + + See + [DBConnection.open_property_graph][lancedb.DBConnection.open_property_graph]. + """ + graph = AsyncPropertyGraph(self._inner, name, namespace_path) + await graph.describe() + return graph + + async def drop_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> None: + """Drop a property graph and wait for its definition to be deleted. + + The tables it reads are untouched. + """ + await self._inner.drop_property_graph(name, list(namespace_path or [])) + + def list_property_graphs( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> AsyncListing[str]: + """The names of the property graphs in one namespace. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return AsyncListing( + self._inner.list_property_graphs( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ) + ) + def list_jobs( self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None ) -> AsyncListing[JobInfo]: diff --git a/python/python/lancedb/graph.py b/python/python/lancedb/graph.py new file mode 100644 index 000000000..cbdcdb009 --- /dev/null +++ b/python/python/lancedb/graph.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Property graphs: graphs defined over tables a database already holds. + +A property graph stores no rows. Its definition says which tables of its +namespace hold nodes and which hold edges: each node table's key and label +(:class:`NodeTable`), and each edge table's label and the node keys its source +and destination columns reference (:class:`EdgeTable`) -- the shape of +SQL/PGQ's ``CREATE PROPERTY GRAPH``. A query reads the tables through SQL +``GRAPH_TABLE``, as they are when it runs. + +See ``DBConnection.create_property_graph``. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple + +from .background_loop import LOOP + +if TYPE_CHECKING: + from ._lancedb import Connection as LanceDbConnection + + +@dataclass(frozen=True) +class NodeTable: + """A table whose rows are the nodes of one label.""" + + table: str + """The table, in the graph's namespace.""" + key: str + """The column that identifies a node.""" + label: str + """The label its nodes carry.""" + properties: Optional[List[str]] = None + """The columns exposed as node properties; every column when ``None``. A + created graph reports the list it resolved to, the key among them.""" + + +@dataclass(frozen=True) +class Endpoint: + """One end of an edge: a column of the edge table and the node key it holds.""" + + column: str + """The edge table's column.""" + references: Tuple[str, str] + """The node table and its key column, as ``(table, column)``.""" + + +@dataclass(frozen=True) +class EdgeTable: + """A table whose rows are the edges of one label.""" + + table: str + """The table, in the graph's namespace.""" + label: str + """The label its edges carry.""" + source: Endpoint + """The column naming each edge's source node.""" + destination: Endpoint + """The column naming each edge's destination node.""" + properties: Optional[List[str]] = None + """The columns exposed as edge properties; every column when ``None``. A + created graph reports the list it resolved to, both endpoint columns among + them.""" + + +@dataclass(frozen=True) +class PropertyGraphDescription: + """What a database records about one property graph.""" + + name: str + """The graph's name within its namespace.""" + namespace_path: List[str] + """The namespace holding the graph; empty is the root namespace.""" + nodes: List[NodeTable] + """The node tables, each with the properties create resolved.""" + edges: List[EdgeTable] + """The edge tables, each with the properties create resolved.""" + + +def _definition_json(nodes: Sequence[NodeTable], edges: Sequence[EdgeTable]) -> str: + """The definition a create request carries, refusing elements of the wrong + type before any request is made.""" + nodes, edges = list(nodes), list(edges) + for element, expected in [ + *((node, NodeTable) for node in nodes), + *((edge, EdgeTable) for edge in edges), + ]: + if not isinstance(element, expected): + raise TypeError( + "A property graph's nodes are NodeTable and its edges EdgeTable, " + f"not {type(element).__name__}" + ) + return json.dumps( + { + "nodes": [_node_table_to_json(node) for node in nodes], + "edges": [_edge_table_to_json(edge) for edge in edges], + } + ) + + +def _description_from_json(text: str) -> PropertyGraphDescription: + described = json.loads(text) + return PropertyGraphDescription( + name=described["name"], + namespace_path=list(described.get("namespace", [])), + nodes=[_node_table_from_json(node) for node in described["nodes"]], + edges=[_edge_table_from_json(edge) for edge in described.get("edges", [])], + ) + + +def _properties_to_json( + encoded: Dict[str, Any], properties: Optional[List[str]] +) -> None: + # Absent means every column, which an empty list would not. + if properties is None: + return + if isinstance(properties, str): + raise TypeError("properties is a list of column names, not a string") + encoded["properties"] = list(properties) + + +def _node_table_to_json(node: NodeTable) -> Dict[str, Any]: + encoded: Dict[str, Any] = { + "table": node.table, + "key": node.key, + "label": node.label, + } + _properties_to_json(encoded, node.properties) + return encoded + + +def _edge_table_to_json(edge: EdgeTable) -> Dict[str, Any]: + encoded: Dict[str, Any] = { + "table": edge.table, + "label": edge.label, + "source": _endpoint_to_json(edge.source), + "destination": _endpoint_to_json(edge.destination), + } + _properties_to_json(encoded, edge.properties) + return encoded + + +def _endpoint_to_json(endpoint: Endpoint) -> Dict[str, Any]: + if not isinstance(endpoint, Endpoint): + raise TypeError( + "An edge's source and destination are Endpoint, " + f"not {type(endpoint).__name__}" + ) + table, column = endpoint.references + return { + "column": endpoint.column, + "references": {"table": table, "column": column}, + } + + +def _node_table_from_json(node: Dict[str, Any]) -> NodeTable: + return NodeTable( + table=node["table"], + key=node["key"], + label=node["label"], + properties=node.get("properties"), + ) + + +def _edge_table_from_json(edge: Dict[str, Any]) -> EdgeTable: + return EdgeTable( + table=edge["table"], + label=edge["label"], + source=_endpoint_from_json(edge["source"]), + destination=_endpoint_from_json(edge["destination"]), + properties=edge.get("properties"), + ) + + +def _endpoint_from_json(endpoint: Dict[str, Any]) -> Endpoint: + references = endpoint["references"] + return Endpoint( + column=endpoint["column"], + references=(references["table"], references["column"]), + ) + + +class AsyncPropertyGraph: + """A handle on one property graph: its name and namespace, and the + connection it is described through. + + Obtained from ``AsyncConnection.create_property_graph`` or + ``AsyncConnection.open_property_graph``. The graph is queried through SQL + ``GRAPH_TABLE``, not through this handle. + """ + + def __init__( + self, + inner: "LanceDbConnection", + name: str, + namespace_path: Optional[List[str]] = None, + ): + self._inner = inner + self._name = name + self._namespace_path = list(namespace_path or []) + + def __repr__(self) -> str: + return ( + f"AsyncPropertyGraph(name={self._name!r}, " + f"namespace_path={self._namespace_path!r})" + ) + + @property + def name(self) -> str: + """The graph's name within its namespace.""" + return self._name + + @property + def namespace_path(self) -> List[str]: + """The namespace holding the graph; empty is the root namespace.""" + return list(self._namespace_path) + + async def describe(self) -> PropertyGraphDescription: + """The graph's node and edge tables, each with the properties create + resolved.""" + return _description_from_json( + await self._inner.describe_property_graph(self._name, self._namespace_path) + ) + + +class PropertyGraph: + """Synchronous variant of + [AsyncPropertyGraph][lancedb.graph.AsyncPropertyGraph].""" + + def __init__(self, graph: AsyncPropertyGraph): + self._async = graph + + def __repr__(self) -> str: + return ( + f"PropertyGraph(name={self.name!r}, namespace_path={self.namespace_path!r})" + ) + + @property + def name(self) -> str: + """The graph's name within its namespace.""" + return self._async.name + + @property + def namespace_path(self) -> List[str]: + """The namespace holding the graph; empty is the root namespace.""" + return self._async.namespace_path + + def describe(self) -> PropertyGraphDescription: + """The graph's node and edge tables, each with the properties create + resolved.""" + return LOOP.run(self._async.describe()) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 48ce686b9..99bd9bca9 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -38,6 +38,7 @@ from ..common import DATA from lancedb.listing import Listing, NamespaceListing, TableListing from ..db import DBConnection, LOOP from ..functions import FunctionVersion, UdfDefinition +from ..graph import EdgeTable, NodeTable, PropertyGraph from ..job import AsyncJob, Job from ..sql import Query as SqlQuery from ..sql import QueryDescription, QueryParameters @@ -1187,6 +1188,61 @@ class RemoteDBConnection(DBConnection): ) ) + @override + def create_property_graph( + self, + name: str, + nodes: Sequence[NodeTable], + edges: Sequence[EdgeTable] = (), + *, + namespace_path: Optional[List[str]] = None, + ) -> PropertyGraph: + return PropertyGraph( + LOOP.run( + self._conn.create_property_graph( + name, nodes, edges, namespace_path=namespace_path + ) + ) + ) + + @override + def open_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> PropertyGraph: + return PropertyGraph( + LOOP.run( + self._conn.open_property_graph(name, namespace_path=namespace_path) + ) + ) + + @override + def drop_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> None: + LOOP.run(self._conn.drop_property_graph(name, namespace_path=namespace_path)) + + @override + def list_property_graphs( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[str]: + """List property graphs. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return Listing( + self._conn.list_property_graphs( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ) + ) + @override def list_jobs( self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 8a0e2ea92..acbe7542c 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors from concurrent.futures import ThreadPoolExecutor import contextlib +import copy +from dataclasses import replace from datetime import timedelta import http.server import io @@ -24,6 +26,14 @@ import lancedb import numpy as np from lancedb.conftest import MockNonNormTextEmbeddingFunction, MockTextEmbeddingFunction from lancedb.embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry +from lancedb.graph import ( + AsyncPropertyGraph, + EdgeTable, + Endpoint, + NodeTable, + PropertyGraph, + PropertyGraphDescription, +) from lancedb.query import ( AsyncQuery, ColumnOrdering, @@ -4583,3 +4593,217 @@ def test_remote_job_handle_reports_its_own_detail(): "limit": 500, "filter": "state = 'claim_complete'", } + + +def _property_graph_handler(calls, respond): + """A mock route handler recording ``(method, path, body)`` and answering + with the ``(status, body)`` that ``respond(path, body)`` returns, as JSON.""" + + def handler(request): + length = int(request.headers.get("Content-Length") or 0) + raw = request.rfile.read(length) if length else b"" + body = json.loads(raw) if raw else None + calls.append((request.command, request.path, body)) + status, payload = respond(request.path, body) + request.send_response(status) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(json.dumps(payload).encode()) + + return handler + + +SOCIAL_NODES = [ + NodeTable("person", key="person_id", label="Person", properties=["name"]) +] +SOCIAL_EDGES = [ + EdgeTable( + "knows", + label="KNOWS", + source=Endpoint("src_id", references=("person", "person_id")), + destination=Endpoint("dst_id", references=("person", "person_id")), + ) +] +# The create body for the tables above. The edge table leaves its properties +# out, which means every column. +SOCIAL_DEFINITION = { + "nodes": [ + { + "table": "person", + "key": "person_id", + "label": "Person", + "properties": ["name"], + } + ], + "edges": [ + { + "table": "knows", + "label": "KNOWS", + "source": { + "column": "src_id", + "references": {"table": "person", "column": "person_id"}, + }, + "destination": { + "column": "dst_id", + "references": {"table": "person", "column": "person_id"}, + }, + } + ], +} +# Create resolves each table's properties, the key and endpoint columns always +# among them. +RESOLVED_NODES = [replace(SOCIAL_NODES[0], properties=["person_id", "name"])] +RESOLVED_EDGES = [replace(SOCIAL_EDGES[0], properties=["src_id", "dst_id", "since"])] + + +def _social_description(namespace): + described = copy.deepcopy(SOCIAL_DEFINITION) + described["nodes"][0]["properties"] = RESOLVED_NODES[0].properties + described["edges"][0]["properties"] = RESOLVED_EDGES[0].properties + return {"name": "social", "namespace": namespace, **described} + + +def test_property_graph_crud_addresses_its_own_routes(): + # The definition is the create body and describe posts an empty object. + # An accepted drop hands back its cleanup job, which drop waits on. + calls = [] + + def respond(path, body): + if path.startswith("/v1/namespace/analytics/property_graph/list"): + if "page_token=p2" in path: + return 200, {"property_graphs": ["payments"]} + return 200, {"property_graphs": ["social"], "page_token": "p2"} + if path.endswith("/drop"): + return 202, {"job_id": "j1-pg-drop"} + if path == "/v1/jobs/describe": + return 200, {"job_id": body["job_id"], "job_state": "DONE"} + return 200, _social_description(["analytics"]) + + with mock_lancedb_connection(_property_graph_handler(calls, respond)) as db: + graph = db.create_property_graph( + "social", SOCIAL_NODES, SOCIAL_EDGES, namespace_path=["analytics"] + ) + assert isinstance(graph, PropertyGraph) + assert (graph.name, graph.namespace_path) == ("social", ["analytics"]) + assert graph.describe() == PropertyGraphDescription( + name="social", + namespace_path=["analytics"], + nodes=RESOLVED_NODES, + edges=RESOLVED_EDGES, + ) + + opened = db.open_property_graph("social", namespace_path=["analytics"]) + assert (opened.name, opened.namespace_path) == ("social", ["analytics"]) + + graphs = db.list_property_graphs(namespace_path=["analytics"], page_limit=1) + assert list(graphs) == ["social", "payments"] + + db.drop_property_graph("social", namespace_path=["analytics"]) + + assert calls == [ + ("POST", "/v1/property_graph/analytics$social/create", SOCIAL_DEFINITION), + ("POST", "/v1/property_graph/analytics$social/describe", {}), + ("POST", "/v1/property_graph/analytics$social/describe", {}), + ("GET", "/v1/namespace/analytics/property_graph/list?limit=1", None), + ( + "GET", + "/v1/namespace/analytics/property_graph/list?page_token=p2&limit=1", + None, + ), + ("POST", "/v1/property_graph/analytics$social/drop", None), + ("POST", "/v1/jobs/describe", {"job_id": "j1-pg-drop"}), + ] + + +@pytest.mark.asyncio +async def test_async_property_graph_handles(): + calls = [] + + def respond(path, body): + if path.endswith("/property_graph/list"): + return 200, {"property_graphs": ["social"]} + if path.endswith("/drop"): + return 200, {} + return 200, _social_description([]) + + async with mock_lancedb_connection_async( + _property_graph_handler(calls, respond) + ) as db: + graph = await db.create_property_graph("social", SOCIAL_NODES) + assert isinstance(graph, AsyncPropertyGraph) + described = await graph.describe() + assert (described.namespace_path, described.nodes) == ([], RESOLVED_NODES) + + opened = await db.open_property_graph("social") + assert isinstance(opened, AsyncPropertyGraph) + assert [name async for name in db.list_property_graphs()] == ["social"] + await db.drop_property_graph("social") + + assert calls == [ + ( + "POST", + "/v1/property_graph/social/create", + {"nodes": SOCIAL_DEFINITION["nodes"], "edges": []}, + ), + ("POST", "/v1/property_graph/social/describe", {}), + ("POST", "/v1/property_graph/social/describe", {}), + ("GET", "/v1/namespace/$/property_graph/list", None), + ("POST", "/v1/property_graph/social/drop", None), + ] + + +def test_a_missing_property_graph_is_not_found(): + # Opening describes the graph, so a name no graph holds fails there. + calls = [] + + def respond(path, body): + return 404, {"error": "property graph 'social' not found"} + + with mock_lancedb_connection(_property_graph_handler(calls, respond)) as db: + with pytest.raises(HttpError) as missing: + db.open_property_graph("social") + assert missing.value.status_code == 404 + with pytest.raises(HttpError) as missing: + db.drop_property_graph("social") + assert missing.value.status_code == 404 + + assert calls == [ + ("POST", "/v1/property_graph/social/describe", {}), + ("POST", "/v1/property_graph/social/drop", None), + ] + + +def test_a_property_graph_takes_node_and_edge_tables(): + # A definition of the wrong shape is refused before any request is made. + calls = [] + with mock_lancedb_connection( + _property_graph_handler(calls, lambda path, body: (200, {})) + ) as db: + with pytest.raises(TypeError, match="nodes are NodeTable"): + db.create_property_graph("social", SOCIAL_EDGES) + with pytest.raises(TypeError, match="not NodeTable"): + db.create_property_graph("social", SOCIAL_NODES, SOCIAL_NODES) + with pytest.raises(TypeError, match="are Endpoint, not tuple"): + db.create_property_graph( + "social", + SOCIAL_NODES, + [replace(SOCIAL_EDGES[0], source=("person", "person_id"))], + ) + with pytest.raises(TypeError, match="not a string"): + db.create_property_graph( + "social", [replace(SOCIAL_NODES[0], properties="name")] + ) + assert calls == [] + + +def test_local_connections_refuse_property_graphs(tmp_path): + db = lancedb.connect(tmp_path) + refused = "Property graph operations are not supported" + with pytest.raises(NotImplementedError, match=refused): + db.create_property_graph("social", SOCIAL_NODES, SOCIAL_EDGES) + with pytest.raises(NotImplementedError, match=refused): + db.open_property_graph("social") + with pytest.raises(NotImplementedError, match=refused): + list(db.list_property_graphs()) + with pytest.raises(NotImplementedError, match=refused): + db.drop_property_graph("social") diff --git a/python/src/connection.rs b/python/src/connection.rs index 004456737..d4516d405 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -24,6 +24,7 @@ use lancedb::{ connection::NamespaceClientPushdownOperation, database::namespace::LanceNamespaceDatabase, database::{CreateTableMode, Database, ReadConsistency}, + graph::PropertyGraphDefinition, }; use pyo3::{ Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python, @@ -1116,6 +1117,79 @@ impl Connection { )) } + /// Create a property graph from its definition's JSON, and return its + /// description's JSON. The Python layer names the fields; the JSON keeps + /// the binding free of classes kept in step with the Rust structs. + #[pyo3(signature = (name, definition_json, namespace_path=None))] + pub fn create_property_graph( + self_: PyRef<'_, Self>, + name: String, + definition_json: String, + namespace_path: Option>, + ) -> PyResult> { + let inner = self_.get_inner()?.clone(); + let namespace_path = namespace_path.unwrap_or_default(); + future_into_py(self_.py(), async move { + let definition = PropertyGraphDefinition::from_json(&definition_json).infer_error()?; + inner + .create_property_graph(name, &definition, &namespace_path) + .await + .infer_error()? + .to_json() + .infer_error() + }) + } + + /// The JSON description of the property graph named `name`. + #[pyo3(signature = (name, namespace_path=None))] + pub fn describe_property_graph( + self_: PyRef<'_, Self>, + name: String, + namespace_path: Option>, + ) -> PyResult> { + let inner = self_.get_inner()?.clone(); + let namespace_path = namespace_path.unwrap_or_default(); + future_into_py(self_.py(), async move { + inner + .describe_property_graph(name, &namespace_path) + .await + .infer_error()? + .to_json() + .infer_error() + }) + } + + #[pyo3(signature = (name, namespace_path=None))] + pub fn drop_property_graph( + self_: PyRef<'_, Self>, + name: String, + namespace_path: Option>, + ) -> PyResult> { + let inner = self_.get_inner()?.clone(); + let namespace_path = namespace_path.unwrap_or_default(); + future_into_py(self_.py(), async move { + inner + .drop_property_graph(name, &namespace_path) + .await + .infer_error() + }) + } + + #[pyo3(signature = (namespace_path=None, *, page_token=None, page_limit=None))] + pub fn list_property_graphs( + &self, + namespace_path: Option>, + page_token: Option, + page_limit: Option, + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_property_graphs( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, page_limit), + ), + )) + } + #[pyo3(signature = (*, page_token=None, page_limit=None))] pub fn list_jobs( &self, diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index d0f8c2f51..cea56a340 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -28,6 +28,7 @@ use crate::database::{ }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; +use crate::graph::{PropertyGraphDefinition, PropertyGraphDescription}; use crate::listing::{Listing, ListingOptions, ListingPage}; #[cfg(feature = "remote")] use crate::remote::{ @@ -38,7 +39,10 @@ use crate::remote::{ }, }; use crate::secrets::SecretInfo; -use crate::utils::{validate_namespace, validate_secret_reference, validate_view_reference}; +use crate::utils::{ + validate_namespace, validate_property_graph_reference, validate_secret_reference, + validate_view_reference, +}; use crate::view::ViewDescription; use lance::io::ObjectStoreParams; pub use lance_file::version::LanceFileVersion; @@ -1049,6 +1053,138 @@ impl Connection { }) } + /// Create a property graph over tables in its namespace. + /// + /// The definition names the tables that hold its nodes, with each one's + /// key column and label, and the tables that hold its edges, with the node + /// keys each one's source and destination columns reference -- the shape + /// of SQL/PGQ's `CREATE PROPERTY GRAPH`. The graph stores nothing else: a + /// query reads its tables through SQL `GRAPH_TABLE`, as they are when it + /// runs. + /// + /// The database checks the definition against the tables now, so one they + /// cannot satisfy is refused here rather than at the first query. The + /// returned description lists each table's properties as they resolved. + /// A name another property graph holds is an error. Local databases return + /// [`Error::NotSupported`]. + /// + /// # Example + /// + /// ```no_run + /// use futures::TryStreamExt; + /// use lancedb::graph::{EdgeTable, Endpoint, NodeTable, PropertyGraphDefinition}; + /// + /// # async fn graph_lifecycle( + /// # connection: &lancedb::Connection, + /// # ) -> Result<(), Box> { + /// let namespace = vec!["analytics".to_string()]; + /// + /// let definition = PropertyGraphDefinition { + /// nodes: vec![NodeTable::new("person", "person_id", "Person").properties(["name"])], + /// edges: vec![EdgeTable::new( + /// "knows", + /// "KNOWS", + /// Endpoint::new("src_id", "person", "person_id"), + /// Endpoint::new("dst_id", "person", "person_id"), + /// )], + /// }; + /// let graph = connection + /// .create_property_graph("social", &definition, &namespace) + /// .await?; + /// // The key is always a property, beside the columns asked for. + /// println!("{:?}", graph.definition.nodes[0].properties); + /// + /// let names: Vec<_> = connection + /// .list_property_graphs(&namespace, Default::default()) + /// .try_collect() + /// .await?; + /// assert!(names.iter().any(|name| name == "social")); + /// + /// // Dropping the graph leaves `person` and `knows` untouched. + /// connection.drop_property_graph("social", &namespace).await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn create_property_graph( + &self, + name: impl AsRef, + definition: &PropertyGraphDefinition, + namespace_path: &[String], + ) -> Result { + validate_property_graph_reference(name.as_ref(), namespace_path)?; + self.internal + .create_property_graph(name.as_ref(), definition, namespace_path) + .await + } + + /// What this database records about one property graph: its node and edge + /// tables, with each table's properties as create resolved them. + /// + /// Local databases return [`Error::NotSupported`]. + pub async fn describe_property_graph( + &self, + name: impl AsRef, + namespace_path: &[String], + ) -> Result { + validate_property_graph_reference(name.as_ref(), namespace_path)?; + self.internal + .describe_property_graph(name.as_ref(), namespace_path) + .await + } + + /// Drop a property graph and wait for its definition to be deleted. + /// + /// The tables it reads are untouched: a property graph holds no rows of + /// its own. Local databases return [`Error::NotSupported`]. + pub async fn drop_property_graph( + &self, + name: impl AsRef, + namespace_path: &[String], + ) -> Result<()> { + validate_property_graph_reference(name.as_ref(), namespace_path)?; + self.internal + .drop_property_graph(name.as_ref(), namespace_path) + .await + } + + /// The names of the property graphs in one namespace. + /// + /// Names only; a definition comes from [`Self::describe_property_graph`]. + /// Local databases return [`Error::NotSupported`]. + /// + /// Results are fetched lazily, one page at a time. See [`Listing`] + /// for cached-result and continuation-token semantics. Errors terminate iteration. + /// + /// ``` + /// use futures::TryStreamExt; + /// + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// let mut items = connection.list_property_graphs(&[], Default::default()); + /// while let Some(item) = items.try_next().await? { + /// println!("{item:?}"); + /// } + /// # Ok(()) + /// # } + /// ``` + pub fn list_property_graphs( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Listing { + let database = self.database().clone(); + let namespace_path = namespace_path.to_vec(); + Listing::new(options, move |options| { + let database = database.clone(); + let namespace_path = namespace_path.clone(); + async move { + validate_namespace(&namespace_path)?; + database + .list_property_graphs(&namespace_path, options) + .await + } + }) + } + /// Rename a table in the database. /// /// This is only supported in LanceDB Distributed. diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index f5fba221d..3dfdf444c 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -29,6 +29,7 @@ use lance_namespace::models::{ use crate::data::scannable::Scannable; use crate::error::Result; +use crate::graph::{PropertyGraphDefinition, PropertyGraphDescription}; use crate::job::Job; use crate::listing::{ListingOptions, ListingPage}; use crate::materialized_view::CreateMaterializedViewRequest; @@ -316,6 +317,12 @@ fn view_ops_not_supported() -> Result { }) } +fn property_graph_ops_not_supported() -> Result { + Err(crate::error::Error::NotSupported { + message: "Property graph operations are not supported by this database".to_string(), + }) +} + /// The `Database` trait defines the interface for database implementations. /// /// A database is responsible for managing tables and their metadata. @@ -586,6 +593,50 @@ pub trait Database: ) -> Result> { view_ops_not_supported() } + /// Create a property graph over tables in its namespace. The tables are + /// checked against the definition, and each table's properties resolved, + /// once, by the database, so a definition they cannot satisfy is refused + /// now rather than at the first query. + async fn create_property_graph( + &self, + _name: &str, + _definition: &PropertyGraphDefinition, + _namespace_path: &[String], + ) -> Result { + property_graph_ops_not_supported() + } + /// What the database records about one property graph: its node and edge + /// tables, with the properties create resolved. + async fn describe_property_graph( + &self, + _name: &str, + _namespace_path: &[String], + ) -> Result { + property_graph_ops_not_supported() + } + /// Drop a property graph and wait for its definition to be deleted. The + /// tables it reads are untouched. + async fn drop_property_graph(&self, _name: &str, _namespace_path: &[String]) -> Result<()> { + property_graph_ops_not_supported() + } + /// Drop a property graph and return the job deleting its definition, + /// without waiting. + #[doc(hidden)] + async fn drop_property_graph_async( + &self, + _name: &str, + _namespace_path: &[String], + ) -> Result { + property_graph_ops_not_supported() + } + /// Fetch one page of property graph names in one namespace. + async fn list_property_graphs( + &self, + _namespace_path: &[String], + _options: ListingOptions, + ) -> Result> { + property_graph_ops_not_supported() + } /// Open a job by id, returning a handle with its record already /// populated. Fails with [`crate::Error::JobNotFound`] when the server has /// no such job. diff --git a/rust/lancedb/src/graph.rs b/rust/lancedb/src/graph.rs new file mode 100644 index 000000000..510fd780d --- /dev/null +++ b/rust/lancedb/src/graph.rs @@ -0,0 +1,281 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Property graphs: a graph defined over tables a database already holds. +//! +//! A property graph stores no rows. Its definition says which tables of its +//! namespace hold nodes and which hold edges -- each node table's key and +//! label, each edge table's label and the node keys its source and destination +//! columns reference -- the shape of SQL/PGQ's `CREATE PROPERTY GRAPH`. A query +//! reads the tables through SQL `GRAPH_TABLE`, as they are when it runs. +//! +//! The verbs live on [`crate::connection::Connection`]. A graph is read +//! through SQL, by name, so there is no handle returning its rows. + +use serde::{Deserialize, Serialize}; + +use crate::error::{Error, Result}; + +/// A property graph's node and edge tables. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PropertyGraphDefinition { + /// The node tables, one per label. + pub nodes: Vec, + /// The edge tables, one per label. + #[serde(default)] + pub edges: Vec, +} + +impl PropertyGraphDefinition { + /// Parse a definition from its JSON form. + pub fn from_json(json: &str) -> Result { + serde_json::from_str(json).map_err(|source| Error::InvalidInput { + message: format!("invalid property graph definition: {source}"), + }) + } +} + +/// A table whose rows are the nodes of one label. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct NodeTable { + /// The table, in the graph's namespace. + pub table: String, + /// The column that identifies a node. + pub key: String, + /// The label its nodes carry. + pub label: String, + /// The columns exposed as node properties; every column when absent. A + /// created graph reports the list it resolved to, the key among them. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub properties: Option>, +} + +impl NodeTable { + /// The nodes labeled `label` are the rows of `table`, each identified by + /// its `key` column, with every column a property. + pub fn new(table: impl Into, key: impl Into, label: impl Into) -> Self { + Self { + table: table.into(), + key: key.into(), + label: label.into(), + properties: None, + } + } + + /// Expose only these columns as node properties. + pub fn properties(mut self, properties: impl IntoIterator>) -> Self { + self.properties = Some(properties.into_iter().map(Into::into).collect()); + self + } +} + +/// A table whose rows are the edges of one label. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct EdgeTable { + /// The table, in the graph's namespace. + pub table: String, + /// The label its edges carry. + pub label: String, + /// The column naming each edge's source node. + pub source: Endpoint, + /// The column naming each edge's destination node. + pub destination: Endpoint, + /// The columns exposed as edge properties; every column when absent. A + /// created graph reports the list it resolved to, both endpoint columns + /// among them. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub properties: Option>, +} + +impl EdgeTable { + /// The edges labeled `label` are the rows of `table`, running from the + /// node `source` names to the node `destination` names, with every column + /// a property. + pub fn new( + table: impl Into, + label: impl Into, + source: Endpoint, + destination: Endpoint, + ) -> Self { + Self { + table: table.into(), + label: label.into(), + source, + destination, + properties: None, + } + } + + /// Expose only these columns as edge properties. + pub fn properties(mut self, properties: impl IntoIterator>) -> Self { + self.properties = Some(properties.into_iter().map(Into::into).collect()); + self + } +} + +/// One end of an edge: a column of the edge table and the node key it holds. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Endpoint { + /// The edge table's column. + pub column: String, + /// The node table and key column its values match. + pub references: EndpointReference, +} + +impl Endpoint { + /// The edge table's `column`, holding keys of the node table `table`, + /// whose key column is `key`. + pub fn new( + column: impl Into, + table: impl Into, + key: impl Into, + ) -> Self { + Self { + column: column.into(), + references: EndpointReference { + table: table.into(), + column: key.into(), + }, + } + } +} + +/// The node table and key column an endpoint's values match. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct EndpointReference { + /// A node table of the same graph. + pub table: String, + /// That node table's key column. + pub column: String, +} + +/// What a database records about one property graph: its node and edge +/// tables, with each table's properties as create resolved them. +/// +/// Returned by [`crate::connection::Connection::describe_property_graph`], and +/// by `create_property_graph` so a caller has the resolved properties without +/// a second call. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PropertyGraphDescription { + /// The graph's name within its namespace. + pub name: String, + /// The namespace holding the graph; empty is the root namespace. + #[serde(rename = "namespace", default)] + pub namespace_path: Vec, + /// The graph's node and edge tables. + #[serde(flatten)] + pub definition: PropertyGraphDefinition, +} + +impl PropertyGraphDescription { + /// The description in its JSON form. + pub fn to_json(&self) -> Result { + serde_json::to_string(self).map_err(|source| Error::Runtime { + message: format!("could not encode a property graph description: {source}"), + }) + } +} + +#[cfg(test)] +mod tests { + use futures::TryStreamExt; + + use super::*; + + fn social() -> PropertyGraphDefinition { + PropertyGraphDefinition { + nodes: vec![NodeTable::new("person", "person_id", "Person").properties(["name"])], + edges: vec![EdgeTable::new( + "knows", + "KNOWS", + Endpoint::new("src_id", "person", "person_id"), + Endpoint::new("dst_id", "person", "person_id"), + )], + } + } + + #[test] + fn test_a_definition_is_the_sql_pgq_shape() { + let json = serde_json::to_value(social()).unwrap(); + // Properties left out are left out of the body too: absent means every + // column, which an empty list would not. + assert_eq!( + json, + serde_json::json!({ + "nodes": [{"table": "person", "key": "person_id", "label": "Person", + "properties": ["name"]}], + "edges": [{"table": "knows", "label": "KNOWS", + "source": {"column": "src_id", + "references": {"table": "person", "column": "person_id"}}, + "destination": {"column": "dst_id", + "references": {"table": "person", + "column": "person_id"}}}] + }) + ); + assert_eq!( + PropertyGraphDefinition::from_json(&json.to_string()).unwrap(), + social() + ); + assert_eq!( + PropertyGraphDefinition::from_json(r#"{"nodes": []}"#).unwrap(), + PropertyGraphDefinition { + nodes: vec![], + edges: vec![], + } + ); + assert!(matches!( + PropertyGraphDefinition::from_json(r#"{"edges": []}"#), + Err(Error::InvalidInput { .. }) + )); + } + + #[test] + fn test_a_description_carries_its_definition_at_the_top_level() { + let description = PropertyGraphDescription { + name: "social".to_string(), + namespace_path: vec!["analytics".to_string()], + definition: social(), + }; + let json: serde_json::Value = + serde_json::from_str(&description.to_json().unwrap()).unwrap(); + assert_eq!(json["name"], "social"); + assert_eq!(json["namespace"], serde_json::json!(["analytics"])); + assert_eq!(json["nodes"][0]["table"], "person"); + assert_eq!(json["edges"][0]["label"], "KNOWS"); + let decoded: PropertyGraphDescription = serde_json::from_value(json).unwrap(); + assert_eq!(decoded, description); + } + + #[tokio::test] + async fn test_local_databases_refuse_property_graphs() { + let dir = tempfile::tempdir().unwrap(); + let conn = crate::connect(dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let refused = |error: Error| { + assert!( + matches!(error, Error::NotSupported { ref message } + if message.contains("Property graph operations")), + "{error}" + ); + }; + refused( + conn.create_property_graph("social", &social(), &[]) + .await + .unwrap_err(), + ); + refused( + conn.describe_property_graph("social", &[]) + .await + .unwrap_err(), + ); + refused(conn.drop_property_graph("social", &[]).await.unwrap_err()); + refused( + conn.list_property_graphs(&[], Default::default()) + .try_collect::>() + .await + .unwrap_err(), + ); + } +} diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index b1abf58af..8d8574023 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -199,6 +199,7 @@ pub mod embeddings; pub mod error; pub mod expr; pub mod function; +pub mod graph; pub mod index; pub mod io; pub mod ipc; diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 9e37b304e..8098f70ba 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -26,6 +26,7 @@ use crate::database::{ }; use crate::error::Result; use crate::function::{FunctionRegistrationRequest, FunctionVersion}; +use crate::graph::{PropertyGraphDefinition, PropertyGraphDescription}; use crate::job::Job; use crate::listing::{ListingOptions, ListingPage}; use crate::materialized_view::CreateMaterializedViewRequest; @@ -34,8 +35,8 @@ use crate::remote::util::stream_as_body; use crate::secrets::SecretInfo; use crate::table::BaseTable; use crate::utils::{ - validate_database_name, validate_function_name, validate_namespace, validate_table_name, - validate_view_name, + validate_database_name, validate_function_name, validate_namespace, + validate_property_graph_name, validate_table_name, validate_view_name, }; use crate::view::ViewDescription; @@ -652,6 +653,14 @@ fn build_namespace_identifier(namespace: &[String]) -> Result { } } +/// Build a full path to a property graph. Note that we do not need to +/// percent-encode this when using it in a URL because the validation functions +/// ensure that only URL-safe characters are present. +fn build_property_graph_identifier(name: &str, namespace: &[String]) -> Result { + validate_property_graph_name(name)?; + build_compound_identifier(name, namespace) +} + /// Build a secure cache key using length prefixes. /// This format is completely unambiguous regardless of delimiter or content. /// Format: [u32_len][namespace1][u32_len][namespace2]...[u32_len][table_name] @@ -696,6 +705,14 @@ struct RemoteListJobsResponse { page_token: Option, } +#[derive(serde::Deserialize)] +struct RemoteListPropertyGraphsResponse { + #[serde(default)] + property_graphs: Vec, + #[serde(default)] + page_token: Option, +} + #[async_trait] impl Database for RemoteDatabase { fn uri(&self) -> &str { @@ -1032,6 +1049,103 @@ impl Database for RemoteDatabase { self.sql_client()?.list_views(namespace_path, options).await } + async fn create_property_graph( + &self, + name: &str, + definition: &PropertyGraphDefinition, + namespace_path: &[String], + ) -> Result { + let graph_id = build_property_graph_identifier(name, namespace_path)?; + let req = self + .client + .post(&format!("/v1/property_graph/{graph_id}/create")) + .json(definition); + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + response.json().await.err_to_http(request_id) + } + + async fn describe_property_graph( + &self, + name: &str, + namespace_path: &[String], + ) -> Result { + let graph_id = build_property_graph_identifier(name, namespace_path)?; + let req = self + .client + .post(&format!("/v1/property_graph/{graph_id}/describe")) + .json(&serde_json::json!({})); + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + response.json().await.err_to_http(request_id) + } + + async fn drop_property_graph(&self, name: &str, namespace_path: &[String]) -> Result<()> { + self.drop_property_graph_async(name, namespace_path) + .await? + .wait() + .await + } + + async fn drop_property_graph_async( + &self, + name: &str, + namespace_path: &[String], + ) -> Result { + let graph_id = build_property_graph_identifier(name, namespace_path)?; + let req = self + .client + .post(&format!("/v1/property_graph/{graph_id}/drop")); + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + let status = response.status(); + let body = response.text().await.err_to_http(request_id.clone())?; + match status { + // Nothing was bound to the name, so nothing is being deleted. + StatusCode::OK => Ok(Job::new_done()), + StatusCode::ACCEPTED => { + let job_id = extract_job_id(&body).ok_or_else(|| Error::Http { + source: "property graph drop response did not contain a valid job_id".into(), + request_id, + status_code: Some(status), + })?; + Ok(Job::new(Box::new(RemoteJob::new( + self.client.clone(), + job_id, + )))) + } + _ => Err(Error::Http { + source: "property graph drop must return 200 OK or 202 Accepted".into(), + request_id, + status_code: Some(status), + }), + } + } + + async fn list_property_graphs( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Result> { + let namespace_id = build_namespace_identifier(namespace_path)?; + let path = format!("/v1/namespace/{namespace_id}/property_graph/list"); + let mut req = self.client.get(&path); + if let Some(token) = options.page_token { + req = req.query(&[("page_token", token)]); + } + if let Some(limit) = options.page_limit { + req = req.query(&[("limit", limit)]); + } + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + let response: RemoteListPropertyGraphsResponse = + response.json().await.err_to_http(request_id)?; + Ok(ListingPage { + items: response.property_graphs, + page_token: response.page_token, + }) + } + async fn open_job(&self, job_id: &str) -> Result { let handle = super::job::RemoteJob::new(self.client.clone(), job_id.to_string()); match crate::job::JobHandle::describe(&handle).await { @@ -1743,6 +1857,7 @@ mod tests { use crate::connection::ConnectBuilder; use crate::database::Database; + use crate::graph::{EdgeTable, Endpoint, NodeTable, PropertyGraphDefinition}; use crate::materialized_view::CreateMaterializedViewRequest; use crate::{ Connection, Error, @@ -3666,7 +3781,7 @@ mod tests { assert!(items.try_next().await.unwrap().is_none()); assert!(items.try_next().await.unwrap().is_none()); } - for resource in ["jobs"] { + for resource in ["property_graph", "jobs"] { let count = Arc::new(AtomicUsize::new(0)); let conn = Connection::new_with_handler({ let count = count.clone(); @@ -3692,6 +3807,7 @@ mod tests { } let (key, item) = match resource { "secret" => ("secrets", serde_json::json!({"name": "a"})), + "property_graph" => ("property_graphs", serde_json::json!("a")), "jobs" => ( "jobs", serde_json::json!({"job_id": "a", "table": "t", "job_type": "create_index", "state": "in_progress", "created_at_millis": 1}), @@ -3713,6 +3829,9 @@ mod tests { match resource { "secret" => check(conn.list_secrets(&namespace, options), count).await, "view" => check(conn.list_views(&namespace, options), count).await, + "property_graph" => { + check(conn.list_property_graphs(&namespace, options), count).await + } _ => check(conn.list_jobs(options), count).await, } } @@ -4199,4 +4318,251 @@ mod tests { assert!(matches!(result, Err(Error::InvalidInput { .. }))); } } + + /// The node table names its properties and the edge table leaves them + /// out, which means every column. + fn social_graph() -> PropertyGraphDefinition { + PropertyGraphDefinition { + nodes: vec![NodeTable::new("person", "person_id", "Person").properties(["name"])], + edges: vec![EdgeTable::new( + "knows", + "KNOWS", + Endpoint::new("src_id", "person", "person_id"), + Endpoint::new("dst_id", "person", "person_id"), + )], + } + } + + /// `social_graph` as create resolves it: each table's properties listed, + /// the key and endpoint columns among them. + fn resolved_social_graph() -> PropertyGraphDefinition { + let mut graph = social_graph(); + graph.nodes[0].properties = Some(vec!["person_id".into(), "name".into()]); + graph.edges[0].properties = Some(vec!["src_id".into(), "dst_id".into(), "since".into()]); + graph + } + + fn property_graph_body(name: &str, namespace: &[&str]) -> String { + let mut body = serde_json::to_value(resolved_social_graph()).unwrap(); + let fields = body.as_object_mut().unwrap(); + fields.insert("name".into(), name.into()); + fields.insert("namespace".into(), namespace.into()); + body.to_string() + } + + #[tokio::test] + async fn test_create_property_graph_posts_the_definition_and_returns_the_resolved_one() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!( + request.url().path(), + "/v1/property_graph/analytics$social/create" + ); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + // The name and its namespace address the graph in the path, so the + // body is the definition alone. + assert_eq!( + body, + serde_json::json!({ + "nodes": [{"table": "person", "key": "person_id", "label": "Person", + "properties": ["name"]}], + "edges": [{"table": "knows", "label": "KNOWS", + "source": {"column": "src_id", + "references": {"table": "person", + "column": "person_id"}}, + "destination": {"column": "dst_id", + "references": {"table": "person", + "column": "person_id"}}}] + }) + ); + http::Response::builder() + .status(200) + .body(property_graph_body("social", &["analytics"])) + .unwrap() + }); + let graph = conn + .create_property_graph("social", &social_graph(), &["analytics".into()]) + .await + .unwrap(); + assert_eq!(graph.name, "social"); + assert_eq!(graph.namespace_path, vec!["analytics".to_string()]); + assert_eq!(graph.definition, resolved_social_graph()); + } + + #[tokio::test] + async fn test_describe_and_drop_address_the_property_graph_in_the_path() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/property_graph/social/describe"); + assert_eq!(request.body().unwrap().as_bytes().unwrap(), b"{}"); + http::Response::builder() + .status(200) + .body(property_graph_body("social", &[])) + .unwrap() + }); + let graph = conn.describe_property_graph("social", &[]).await.unwrap(); + assert!(graph.namespace_path.is_empty()); + assert_eq!(graph.definition, resolved_social_graph()); + + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!( + request.url().path(), + "/v1/property_graph/analytics$social/drop" + ); + assert!(request.body().is_none(), "{:?}", request.body()); + http::Response::builder().status(200).body("{}").unwrap() + }); + conn.drop_property_graph("social", &["analytics".into()]) + .await + .unwrap(); + } + + /// An accepted drop hands back the job deleting the definition, and the + /// waiting `drop_property_graph` waits on it. + #[tokio::test] + async fn test_drop_property_graph_waits_for_the_cleanup_job() { + let db = super::RemoteDatabase::new_mock(|_| { + http::Response::builder() + .status(202) + .body(r#"{"job_id":"j1-pg-drop"}"#) + .unwrap() + }); + let job = db.drop_property_graph_async("social", &[]).await.unwrap(); + assert_eq!(job.id(), Some("j1-pg-drop")); + + let requests = Arc::new(AtomicUsize::new(0)); + let conn = Connection::new_with_handler({ + let requests = requests.clone(); + move |request| match ( + requests.fetch_add(1, Ordering::SeqCst), + request.url().path(), + ) { + (0, "/v1/property_graph/social/drop") => http::Response::builder() + .status(202) + .body(r#"{"job_id":"j1-pg-drop"}"#) + .unwrap(), + (1, "/v1/jobs/describe") => { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(body, serde_json::json!({"job_id": "j1-pg-drop"})); + http::Response::builder() + .status(200) + .body(r#"{"job_id":"j1-pg-drop","job_state":"DONE"}"#) + .unwrap() + } + (n, path) => panic!("unexpected request {n}: {path}"), + } + }); + conn.drop_property_graph("social", &[]).await.unwrap(); + assert_eq!(requests.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn test_drop_property_graph_rejects_an_unexpected_response() { + let db = super::RemoteDatabase::new_mock(|_| { + http::Response::builder().status(200).body("{}").unwrap() + }); + let job = db.drop_property_graph_async("social", &[]).await.unwrap(); + assert_eq!(job.id(), None); + job.wait().await.unwrap(); + + for (status, body, expected) in [ + (202, "{}", "valid job_id"), + (202, r#"{"job_id":""}"#, "valid job_id"), + (204, "", "200 OK or 202 Accepted"), + ] { + let db = super::RemoteDatabase::new_mock(move |_| { + http::Response::builder().status(status).body(body).unwrap() + }); + let error = db + .drop_property_graph_async("social", &[]) + .await + .err() + .unwrap(); + assert!(error.to_string().contains(expected), "{error}"); + } + } + + /// A graph the database does not hold is the service's 404, kept as the + /// status of the error so a caller can tell it from a refused request. + #[tokio::test] + async fn test_a_missing_property_graph_reports_not_found() { + let conn = Connection::new_with_handler(|_| { + http::Response::builder() + .status(404) + .body("property graph 'social' not found") + .unwrap() + }); + let not_found = |error: Error| { + assert!( + matches!( + error, + Error::Http { + status_code: Some(http::StatusCode::NOT_FOUND), + .. + } + ), + "{error:?}" + ); + }; + not_found( + conn.describe_property_graph("social", &[]) + .await + .unwrap_err(), + ); + not_found(conn.drop_property_graph("social", &[]).await.unwrap_err()); + } + + #[tokio::test] + async fn test_list_property_graphs_walks_pages() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::GET); + assert_eq!(request.url().path(), "/v1/namespace/$/property_graph/list"); + let page = request + .url() + .query_pairs() + .find(|(key, _)| key == "page_token") + .map(|(_, value)| value.into_owned()); + let body = match page.as_deref() { + None => r#"{"property_graphs":["social"],"page_token":"p2"}"#, + Some("p2") => r#"{"property_graphs":["payments"]}"#, + Some(other) => panic!("unexpected page token: {other}"), + }; + http::Response::builder().status(200).body(body).unwrap() + }); + assert_eq!( + conn.list_property_graphs(&[], Default::default()) + .try_collect::>() + .await + .unwrap(), + vec!["social".to_string(), "payments".to_string()] + ); + } + + /// A name carrying the delimiter would split back apart as a different + /// graph, so it is refused before it reaches a route. + #[tokio::test] + async fn test_property_graph_names_that_would_resplit_are_refused() { + let conn = Connection::new_with_handler(|_| -> http::Response { + panic!("an invalid identifier must not reach the service") + }); + for (name, namespace) in [ + ("analytics$social", vec![]), + ("social", vec!["ana$lytics".to_string()]), + ("", vec![]), + ("..", vec![]), + ] { + let error = match conn.describe_property_graph(name, &namespace).await { + Ok(_) => panic!("accepted {name:?} in {namespace:?}"), + Err(error) => error.to_string(), + }; + assert!( + error.contains("property graph name") || error.contains("namespace component"), + "{error}" + ); + } + } } diff --git a/rust/lancedb/src/utils/mod.rs b/rust/lancedb/src/utils/mod.rs index ebf3f697b..061ca0f14 100644 --- a/rust/lancedb/src/utils/mod.rs +++ b/rust/lancedb/src/utils/mod.rs @@ -38,6 +38,9 @@ const MAX_SECRET_NAME_CHARS: usize = 128; /// The maximum length of a function name. const MAX_FUNCTION_NAME_CHARS: usize = 128; +/// The maximum length of a property graph name. +const MAX_PROPERTY_GRAPH_NAME_CHARS: usize = 128; + /// The maximum length of a database name. const MAX_DB_NAME_CHARS: usize = 64; @@ -214,6 +217,18 @@ pub fn validate_view_reference(name: &str, namespace_path: &[String]) -> Result< validate_view_name(name) } +/// Validate a property graph name. +pub fn validate_property_graph_name(name: &str) -> Result<()> { + validate_view_or_namespace_or_secret(name, "property graph name", MAX_PROPERTY_GRAPH_NAME_CHARS) +} + +/// Validate a property graph name and every segment of the namespace path +/// holding it. +pub fn validate_property_graph_reference(name: &str, namespace_path: &[String]) -> Result<()> { + validate_namespace(namespace_path)?; + validate_property_graph_name(name) +} + /// Validate a function name. These do not allow dashes or periods. pub fn validate_function_name(name: &str) -> Result<()> { if name.is_empty() {