_arango/arango.py

from arango import ArangoClient
from dotenv import load_dotenv
import os
import re
from typing import Optional

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

    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)


arango = Arango()

if __name__ == "__main__":
    arango = Arango()
    print("Connected to ArangoDB")