_arango/arango.py

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")