Source code for langchain_community.graphs.nebula_graph

import logging
from string import Template
from typing import Any, Dict, Optional

logger = logging.getLogger(__name__)

rel_query = Template(
    """
MATCH ()-[e:`$edge_type`]->()
  WITH e limit 1
MATCH (m)-[:`$edge_type`]->(n) WHERE id(m) == src(e) AND id(n) == dst(e)
RETURN "(:" + tags(m)[0] + ")-[:$edge_type]->(:" + tags(n)[0] + ")" AS rels
"""
)

RETRY_TIMES = 3


[docs] class NebulaGraph: """NebulaGraph wrapper for graph operations. NebulaGraph inherits methods from Neo4jGraph to bring ease to the user space. *Security note*: Make sure that the database connection uses credentials that are narrowly-scoped to only include necessary permissions. Failure to do so may result in data corruption or loss, since the calling code may attempt commands that would result in deletion, mutation of data if appropriately prompted or reading sensitive data if such data is present in the database. The best way to guard against such negative outcomes is to (as appropriate) limit the permissions granted to the credentials used with this tool. See https://python.langchain.com/docs/security for more information. """
[docs] def __init__( self, space: str, username: str = "root", password: str = "nebula", address: str = "127.0.0.1", port: int = 9669, session_pool_size: int = 30, ) -> None: """Create a new NebulaGraph wrapper instance.""" try: import nebula3 # noqa: F401 import pandas # noqa: F401 except ImportError: raise ImportError( "Please install NebulaGraph Python client and pandas first: " "`pip install nebula3-python pandas`" ) self.username = username self.password = password self.address = address self.port = port self.space = space self.session_pool_size = session_pool_size self.session_pool = self._get_session_pool() self.schema = "" # Set schema try: self.refresh_schema() except Exception as e: raise ValueError(f"Could not refresh schema. Error: {e}")
def _get_session_pool(self) -> Any: assert all( [self.username, self.password, self.address, self.port, self.space] ), ( "Please provide all of the following parameters: " "username, password, address, port, space" ) from nebula3.Config import SessionPoolConfig from nebula3.Exception import AuthFailedException, InValidHostname from nebula3.gclient.net.SessionPool import SessionPool config = SessionPoolConfig() config.max_size = self.session_pool_size try: session_pool = SessionPool( self.username, self.password, self.space, [(self.address, self.port)], ) except InValidHostname: raise ValueError( "Could not connect to NebulaGraph database. " "Please ensure that the address and port are correct" ) try: session_pool.init(config) except AuthFailedException: raise ValueError( "Could not connect to NebulaGraph database. " "Please ensure that the username and password are correct" ) except RuntimeError as e: raise ValueError(f"Error initializing session pool. Error: {e}") return session_pool def __del__(self) -> None: try: self.session_pool.close() except Exception as e: logger.warning(f"Could not close session pool. Error: {e}") @property def get_schema(self) -> str: """Returns the schema of the NebulaGraph database""" return self.schema
[docs] def execute(self, query: str, params: Optional[dict] = None, retry: int = 0) -> Any: """Query NebulaGraph database.""" from nebula3.Exception import IOErrorException, NoValidSessionException from nebula3.fbthrift.transport.TTransport import TTransportException params = params or {} try: result = self.session_pool.execute_parameter(query, params) if not result.is_succeeded(): logger.warning( f"Error executing query to NebulaGraph. " f"Error: {result.error_msg()}\n" f"Query: {query} \n" ) return result except NoValidSessionException: logger.warning( f"No valid session found in session pool. " f"Please consider increasing the session pool size. " f"Current size: {self.session_pool_size}" ) raise ValueError( f"No valid session found in session pool. " f"Please consider increasing the session pool size. " f"Current size: {self.session_pool_size}" ) except RuntimeError as e: if retry < RETRY_TIMES: retry += 1 logger.warning( f"Error executing query to NebulaGraph. " f"Retrying ({retry}/{RETRY_TIMES})...\n" f"query: {query} \n" f"Error: {e}" ) return self.execute(query, params, retry) else: raise ValueError(f"Error executing query to NebulaGraph. Error: {e}") except (TTransportException, IOErrorException): # connection issue, try to recreate session pool if retry < RETRY_TIMES: retry += 1 logger.warning( f"Connection issue with NebulaGraph. " f"Retrying ({retry}/{RETRY_TIMES})...\n to recreate session pool" ) self.session_pool = self._get_session_pool() return self.execute(query, params, retry)
[docs] def refresh_schema(self) -> None: """ Refreshes the NebulaGraph schema information. """ tags_schema, edge_types_schema, relationships = [], [], [] for tag in self.execute("SHOW TAGS").column_values("Name"): tag_name = tag.cast() tag_schema = {"tag": tag_name, "properties": []} r = self.execute(f"DESCRIBE TAG `{tag_name}`") props, types = r.column_values("Field"), r.column_values("Type") for i in range(r.row_size()): tag_schema["properties"].append((props[i].cast(), types[i].cast())) tags_schema.append(tag_schema) for edge_type in self.execute("SHOW EDGES").column_values("Name"): edge_type_name = edge_type.cast() edge_schema = {"edge": edge_type_name, "properties": []} r = self.execute(f"DESCRIBE EDGE `{edge_type_name}`") props, types = r.column_values("Field"), r.column_values("Type") for i in range(r.row_size()): edge_schema["properties"].append((props[i].cast(), types[i].cast())) edge_types_schema.append(edge_schema) # build relationships types r = self.execute( rel_query.substitute(edge_type=edge_type_name) ).column_values("rels") if len(r) > 0: relationships.append(r[0].cast()) self.schema = ( f"Node properties: {tags_schema}\n" f"Edge properties: {edge_types_schema}\n" f"Relationships: {relationships}\n" )
[docs] def query(self, query: str, retry: int = 0) -> Dict[str, Any]: result = self.execute(query, retry=retry) columns = result.keys() d: Dict[str, list] = {} for col_num in range(result.col_size()): col_name = columns[col_num] col_list = result.column_values(col_name) d[col_name] = [x.cast() for x in col_list] return d