import os
from math import ceil
from typing import Any, Dict, List, Optional
[docs]class ArangoGraph:
"""ArangoDB图操作的包装器。
*安全提示*: 确保数据库连接使用的凭据范围狭窄，仅包括必要的权限。
如果未能这样做，可能会导致数据损坏或丢失，因为调用代码可能尝试执行会导致删除、变异数据（如果适当提示）或读取敏感数据（如果数据库中存在此类数据）的命令。
防止出现这种负面结果的最佳方法是（视情况）限制授予此工具使用的凭据的权限。
有关更多信息，请参见 https://python.langchain.com/docs/security。"""
[docs] def __init__(self, db: Any) -> None:
"""创建一个新的ArangoDB图形包装器实例。"""
self.set_db(db)
self.set_schema()
@property
def db(self) -> Any:
return self.__db
@property
def schema(self) -> Dict[str, Any]:
return self.__schema
[docs] def set_db(self, db: Any) -> None:
from arango.database import Database
if not isinstance(db, Database):
msg = "**db** parameter must inherit from arango.database.Database"
raise TypeError(msg)
self.__db: Database = db
self.set_schema()
[docs] def set_schema(self, schema: Optional[Dict[str, Any]] = None) -> None:
"""设置ArangoDB数据库的模式。
如果 **schema** 为None，则自动生成模式。
"""
self.__schema = self.generate_schema() if schema is None else schema
[docs] def generate_schema(
self, sample_ratio: float = 0
) -> Dict[str, List[Dict[str, Any]]]:
"""生成ArangoDB数据库的模式并返回
用户可以指定一个 **sample_ratio** （0到1）来确定
用于呈现每个Collection模式的文档/边的比例（相对于Collection大小）。
"""
if not 0 <= sample_ratio <= 1:
raise ValueError("**sample_ratio** value must be in between 0 to 1")
# Stores the Edge Relationships between each ArangoDB Document Collection
graph_schema: List[Dict[str, Any]] = [
{"graph_name": g["name"], "edge_definitions": g["edge_definitions"]}
for g in self.db.graphs()
]
# Stores the schema of every ArangoDB Document/Edge collection
collection_schema: List[Dict[str, Any]] = []
for collection in self.db.collections():
if collection["system"]:
continue
# Extract collection name, type, and size
col_name: str = collection["name"]
col_type: str = collection["type"]
col_size: int = self.db.collection(col_name).count()
# Skip collection if empty
if col_size == 0:
continue
# Set number of ArangoDB documents/edges to retrieve
limit_amount = ceil(sample_ratio * col_size) or 1
aql = f"""
FOR doc in {col_name}
LIMIT {limit_amount}
RETURN doc
"""
doc: Dict[str, Any]
properties: List[Dict[str, str]] = []
for doc in self.__db.aql.execute(aql):
for key, value in doc.items():
properties.append({"name": key, "type": type(value).__name__})
collection_schema.append(
{
"collection_name": col_name,
"collection_type": col_type,
f"{col_type}_properties": properties,
f"example_{col_type}": doc,
}
)
return {"Graph Schema": graph_schema, "Collection Schema": collection_schema}
[docs] def query(
self, query: str, top_k: Optional[int] = None, **kwargs: Any
) -> List[Dict[str, Any]]:
"""查询ArangoDB数据库。"""
import itertools
cursor = self.__db.aql.execute(query, **kwargs)
return [doc for doc in itertools.islice(cursor, top_k)]
[docs] @classmethod
def from_db_credentials(
cls,
url: Optional[str] = None,
dbname: Optional[str] = None,
username: Optional[str] = None,
password: Optional[str] = None,
) -> Any:
"""方便的构造函数，用于从凭据构建Arango DB。
参数：
url: Arango DB的URL。可以作为命名参数传递，也可以设置为环境变量``ARANGODB_URL``。默认为"http://localhost:8529"。
dbname: Arango DB的名称。可以作为命名参数传递，也可以设置为环境变量``ARANGODB_DBNAME``。默认为"_system"。
username: 可以作为命名参数传递，也可以设置为环境变量``ARANGODB_USERNAME``。默认为"root"。
password: 可以作为命名参数传递，也可以设置为环境变量``ARANGODB_PASSWORD``。默认为""。
返回：
一个arango.database.StandardDatabase。
"""
db = get_arangodb_client(
url=url, dbname=dbname, username=username, password=password
)
return cls(db)
[docs]def get_arangodb_client(
url: Optional[str] = None,
dbname: Optional[str] = None,
username: Optional[str] = None,
password: Optional[str] = None,
) -> Any:
"""从凭据中获取Arango DB客户端。
参数:
url: Arango DB的URL。可以作为命名参数传递，也可以设置为环境变量``ARANGODB_URL``。默认为"http://localhost:8529"。
dbname: Arango DB的名称。可以作为命名参数传递，也可以设置为环境变量``ARANGODB_DBNAME``。默认为"_system"。
username: 可以作为命名参数传递，也可以设置为环境变量``ARANGODB_USERNAME``。默认为"root"。
password: 可以作为命名参数传递，也可以设置为环境变量``ARANGODB_PASSWORD``。默认为""。
返回:
一个arango.database.StandardDatabase。
"""
try:
from arango import ArangoClient
except ImportError as e:
raise ImportError(
"Unable to import arango, please install with `pip install python-arango`."
) from e
_url: str = url or os.environ.get("ARANGODB_URL", "http://localhost:8529") # type: ignore[assignment] # noqa: E501
_dbname: str = dbname or os.environ.get("ARANGODB_DBNAME", "_system") # type: ignore[assignment] # noqa: E501
_username: str = username or os.environ.get("ARANGODB_USERNAME", "root") # type: ignore[assignment] # noqa: E501
_password: str = password or os.environ.get("ARANGODB_PASSWORD", "") # type: ignore[assignment] # noqa: E501
return ArangoClient(_url).db(_dbname, _username, _password, verify=True)