added db_name

main
lasseedfast 5 months ago
parent bae27c758b
commit fa175b9c10
  1. 1
      .gitignore
  2. 2
      __init__.py
  3. BIN
      __pycache__/arango.cpython-310.pyc
  4. 99
      _arango.py

1
.gitignore vendored

@ -0,0 +1 @@
.env

@ -1 +1 @@
from .arango import arango
from ._arango import arango

@ -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"
Loading…
Cancel
Save