import os import re from typing import Any, Dict, Optional from arango import ArangoClient from dotenv import load_dotenv load_dotenv() if "ARANGO_HOST" not in os.environ: import env_manager env_manager.set_env() class Arango: """ Provides a simple interface for connecting to and interacting with an ArangoDB database. Environment variables required: - ARANGO_HOST: The host URL of the ArangoDB server. - ARANGO_DB: The name of the database to connect to. - ARANGO_USERNAME: Username for authentication. - ARANGO_PWD: Password for authentication. Example usage: arango = Arango() results = arango.execute_aql("FOR doc IN my_collection RETURN doc") """ def __init__(self): """ Initializes the ArangoDB client and connects to the specified database using environment variables. """ self.client = ArangoClient(hosts=os.environ.get("ARANGO_HOST")) self.db = self.client.db( os.environ.get("ARANGO_DB"), username=os.environ.get("ARANGO_USERNAME"), password=os.environ.get("ARANGO_PWD"), ) def fix_key(self, _key: str) -> str: """ Sanitizes a given key for use in ArangoDB by replacing disallowed characters with underscores. Allowed characters: alphanumeric, underscore, hyphen, dot, at symbol, parentheses, plus, equals, semicolon, dollar sign, asterisk, single quote, percent, or colon. Args: _key (str): The key to be sanitized. Returns: str: The sanitized key. """ return re.sub(r"[^A-Za-z0-9_\-\.@()+=;\$!*\'%:]", "_", _key) def clear_collections(self): """ Truncates (empties) all non-system collections in the connected database. System collections (names starting with '_') are skipped. """ for db in self.db.collections(): if not db["name"].startswith("_"): col = self.db.collection(db["name"]) col.truncate() print(f"Truncated collection: {db['name']}") def execute_aql( self, query: str, bind_vars: Optional[dict] = None, batch_size: Optional[int] = None, ) -> list[dict]: """ Executes an AQL (Arango Query Language) query and returns the results as a list of dictionaries. Args: query (str): The AQL query string. bind_vars (Optional[dict]): Optional dictionary of bind variables for the query. batch_size (Optional[int]): Optional batch size for fetching results. Returns: list[dict]: The query results. """ cursor = self.db.aql.execute( query, bind_vars=bind_vars or {}, batch_size=batch_size ) return list(cursor) def create_collection(self, name: str): """ Creates a new collection in the database if it does not already exist. Args: name (str): The name of the collection to create. """ if not self.db.has_collection(name): self.db.create_collection(name) print(f"Created collection: {name}") else: print(f"Collection '{name}' already exists.") load_dotenv() class AdminArango: """ Provides administrative functionalities for ArangoDB, such as managing databases and users. Connects to the '_system' database for admin operations. Example usage: admin = AdminArango() admin.create_database("mydb") admin.create_user("username", password="secret") admin.set_user_permission("username", "rw", "mydb") """ def __init__(self): """ Initializes the AdminArango client and connects to the '_system' database using environment variables. """ self.client = ArangoClient(hosts=os.environ.get("ARANGO_HOST")) self.sys_db = self.client.db( "_system", username=os.environ.get("ARANGO_USERNAME"), password=os.environ.get("ARANGO_PWD"), ) # --- Database Management --- def create_database( self, db_name: str, users: Optional[list[Dict[str, Any]]] = None ): """ Creates a new database in ArangoDB. Args: db_name (str): Name of the database to create. users (Optional[list[dict]]): List of user dicts to create with the database. """ if not self.sys_db.has_database(db_name): self.sys_db.create_database(db_name, users=users) print(f"Created database: {db_name}") else: print(f"Database '{db_name}' already exists.") def delete_database(self, db_name: str): """ Deletes a database from ArangoDB. Args: db_name (str): Name of the database to delete. """ if self.sys_db.has_database(db_name): self.sys_db.delete_database(db_name) print(f"Deleted database: {db_name}") else: print(f"Database '{db_name}' does not exist.") def list_databases(self): """ Lists all databases in the ArangoDB server. Returns: list[str]: List of database names. """ return self.sys_db.databases() # --- User Management --- def create_user( self, username: str, password: Optional[str] = None, active: Optional[bool] = True, extra: Optional[dict] = None, ): """ Creates a new user in ArangoDB. Args: username (str): Username for the new user. password (Optional[str]): Password for the user. active (Optional[bool]): Whether the user is active. extra (Optional[dict]): Extra user attributes. """ if not self.sys_db.has_user(username): self.sys_db.create_user( username, password=password, active=active, extra=extra ) print(f"Created user: {username}") else: print(f"User '{username}' already exists.") def delete_user(self, username: str): """ Deletes a user from ArangoDB. Args: username (str): Username of the user to delete. """ if self.sys_db.has_user(username): self.sys_db.delete_user(username) print(f"Deleted user: {username}") else: print(f"User '{username}' does not exist.") def set_user_permission( self, username: str, permission: str, database: str, collection: Optional[str] = None, ): """ Sets user permissions for a database or collection. Args: username (str): Username to set permissions for. permission (str): "rw" (read/write), "ro" (read-only), "none" (no access). database (str): Database name. collection (Optional[str]): Collection name (optional). """ self.sys_db.update_permission(username, permission, database, collection) print( f"Set permission '{permission}' for user '{username}' on database '{database}'" + (f", collection '{collection}'" if collection else "") ) arango = Arango() if __name__ == "__main__": arango = Arango() print("Connected to ArangoDB")