_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_HOSTS" 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, db_name: str = None):
        """
        Initializes the ArangoDB client and connects to the specified database using environment variables.
        """

        self.client = ArangoClient(hosts=os.environ.get("ARANGO_HOSTS"))

        if db_name is None:
            db_name = os.environ.get("ARANGO_DB")
        self.db = self.client.db(
            db_name,
            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.")



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.
        """
        arango_hosts = os.environ.get("ARANGO_HOSTS")
        if not arango_hosts:
            raise RuntimeError("Missing ARANGO_HOSTS environment variable.")
        self.client = ArangoClient(hosts=arango_hosts)
        self.sys_db = self.client.db(
            "_system",
            username=os.environ.get("ARANGO_ADMIN_USERNAME"),
            password=os.environ.get("ARANGO_ADMIN_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.
        """
        # Check if database exists
        try:
            if self.sys_db.has_database(db_name):
                print(f"Database '{db_name}' already exists.")
                return
        except Exception as e:
            print(f"Error checking database existence: {e}")
            return

        # Prepare users list for database creation
        user_objs = []
        if users:
            for user in users:
                # Check if user exists
                try:
                    if self.sys_db.has_user(user['username']):
                        print(f"User '{user['username']}' already exists. Will not create, but will grant permissions.")
                    else:
                        self.sys_db.create_user(
                            username=user['username'],
                            password=user.get('password'),
                            active=True
                        )
                        print(f"User '{user['username']}' created.")
                except Exception as e:
                    print(f"Error checking/creating user '{user['username']}': {e}")
                # Add user object for database creation (even if already exists)
                user_objs.append({
                    'username': user['username'],
                    'password': user.get('password')
                })

        # Create database
        try:
            self.sys_db.create_database(db_name, users=user_objs if user_objs else None)
            print(f"Database '{db_name}' created.")
        except Exception as e:
            print(f"Error creating database: {e}")
            return

        # Grant user permissions on the new database
        if users:
            try:
                db = self.client.db(db_name, username='root', password=os.environ.get("ARANGO_PWD"))
                for user in users:
                    self.sys_db.update_permission(user['username'], 'rw', db_name)
                    print(f"Granted 'rw' permission to user '{user['username']}' on database '{db_name}'.")
            except Exception as e:
                print(f"Error setting permissions: {e}")

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


def admin():
    """
    Interactive terminal guide for ArangoDB admin tasks:
    1) Create a database
    2) Create a user and set permissions
    """
    import getpass
    GREY = "\033[90m"
    RESET = "\033[0m"

    admin = AdminArango()
    print("\nArangoDB Admin Guide")
    print("1) Create a database")
    print("2) Create user and set permissions")
    choice = input("Choose an option (1/2): ").strip()

    if choice == "1":
        db_name = input("Enter new database name: ").strip()
        add_user = input("Add a user to this database? (y/n): ").strip().lower()
        users = []
        if add_user == "y":
            username = input("Enter username: ").strip()
            # Check if user exists before asking for password
            if admin.sys_db.has_user(username):
                print(f"User '{username}' already exists. Will not prompt for password.")
                users.append({"username": username})
            else:
                password = getpass.getpass("Enter password: ")
                users.append({"username": username, "password": password})
        admin.create_database(db_name, users=users if users else None)

    elif choice == "2":
        username = input("Enter new username: ").strip()
        # Check if user exists before asking for password
        if admin.sys_db.has_user(username):
            print(f"User '{username}' already exists. Skipping creation and password prompt.")
        else:
            password = getpass.getpass("Enter password: ")
            active = input("Should the user be active? (y/n): ").strip().lower() == "y"
            admin.create_user(username, password=password, active=active)
        db_name = input("Set permissions for which database?: ").strip()
        print(
            f"{GREY}Permission options:\n"
            "  rw   - read and write access\n"
            "  ro   - read-only access\n"
            "  none - no access\n"
            "See: https://docs.python-arango.com/en/main/specs.html#user-permissions"
            f"{RESET}"
        )
        permission = input("Permission (rw/ro/none): ").strip()
        print(
            f"{GREY}Leave collection blank to set permissions for the whole database.{RESET}"
        )
        collection = input("Collection (leave blank for database-level): ").strip()
        collection = collection if collection else None
        admin.set_user_permission(username, permission, db_name, collection)
    else:
        print("Invalid choice.")

# You can call admin_terminal_guide() from __main__ or elsewhere as needed.
arango = Arango()

if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "admin":
        admin()
    else:
        arango = Arango()
        print("Connected to ArangoDB")