import os import re from typing import Any, Dict, Optional from arango import ArangoClient from dotenv import load_dotenv load_dotenv() #TODO Should I make it async? https://python-arango-async.readthedocs.io/en/latest/migration.html 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, user=None, password=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= user or os.environ.get("ARANGO_USERNAME"), password= password or 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 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 import re if len(sys.argv) > 1 and sys.argv[1] == "admin": admin() else: arango = Arango(db_name="riksdagen", user="riksdagen") print("Connected to ArangoDB") print(arango.db.name) print(arango.db.collections()) collection = arango.db.collection("talks") update_docs = [] for doc in collection.all(): anforandetext = doc["anforandetext"] # Regex to find "- " as in slut- satsen anforandetext = re.sub(r"\w+-\s\w+", '', anforandetext)