diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2eea525 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +.env \ No newline at end of file diff --git a/__init__.py b/__init__.py index cda57b5..3f332e1 100644 --- a/__init__.py +++ b/__init__.py @@ -1 +1 @@ -from .arango import arango +from ._arango import arango diff --git a/__pycache__/arango.cpython-310.pyc b/__pycache__/arango.cpython-310.pyc new file mode 100644 index 0000000..7fe8652 Binary files /dev/null and b/__pycache__/arango.cpython-310.pyc differ diff --git a/arango.py b/_arango.py similarity index 71% rename from arango.py rename to _arango.py index 99802ec..a2080a0 100644 --- a/arango.py +++ b/_arango.py @@ -7,12 +7,10 @@ from dotenv import load_dotenv load_dotenv() -if "ARANGO_HOST" not in os.environ: +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. @@ -28,13 +26,17 @@ class Arango: results = arango.execute_aql("FOR doc IN my_collection RETURN doc") """ - def __init__(self): + 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_HOST")) + + 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( - os.environ.get("ARANGO_DB"), + db_name, username=os.environ.get("ARANGO_USERNAME"), password=os.environ.get("ARANGO_PWD"), ) @@ -102,8 +104,6 @@ class Arango: print(f"Collection '{name}' already exists.") -load_dotenv() - class AdminArango: """ @@ -121,11 +121,14 @@ class AdminArango: """ Initializes the AdminArango client and connects to the '_system' database using environment variables. """ - self.client = ArangoClient(hosts=os.environ.get("ARANGO_HOST")) + 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_USERNAME"), - password=os.environ.get("ARANGO_PWD"), + username=os.environ.get("ARANGO_ADMIN_USERNAME"), + password=os.environ.get("ARANGO_ADMIN_PWD"), ) # --- Database Management --- @@ -140,11 +143,55 @@ class AdminArango: 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.") + # 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): """ @@ -237,6 +284,7 @@ def admin(): 1) Create a database 2) Create a user and set permissions """ + import getpass GREY = "\033[90m" RESET = "\033[0m" @@ -252,15 +300,24 @@ def admin(): users = [] if add_user == "y": username = input("Enter username: ").strip() - password = input("Enter password: ").strip() - users.append({"username": username, "password": password}) + # 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() - password = input("Enter password: ").strip() - active = input("Should the user be active? (y/n): ").strip().lower() == "y" - admin.create_user(username, password=password, active=active) + # 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"