Added write_to_env_file

env_manager/env_manager.py

@@ -1,20 +1,13 @@
 import os
 from arango import ArangoClient
 from dotenv import load_dotenv
+
 def get_env():
     """
-    Initializes a connection to an ArangoDB database and retrieves all documents from the 'enviroment' collection.
-
-    This function attempts to load environment variables from a .env file to get the ArangoDB password. If the .env file
-    is not found, it prompts the user to input the password manually. It then connects to the ArangoDB instance using
-    the provided credentials and returns the first document from the 'enviroment' collection.
+    Load environment variables from ArangoDB 'enviroment' collection.
 
     Returns:
-        dict: The first document from the 'enviroment' collection in the ArangoDB database.
-
-    Raises:
-        FileNotFoundError: If the .env file is not found.
-        KeyError: If the 'ARANGO_PWD_ENV_MANAGER' environment variable is not set.
+        dict: Key-value pairs from the 'enviroment' collection.
     """
     try:
         load_dotenv(".env")
@@ -41,6 +34,9 @@ def get_env():
 
 
 def set_env():
+    """
+    Set environment variables in the current process from ArangoDB 'enviroment' collection.
+    """
         # Set the environment variables
     for key, value in get_env().items():
         if isinstance(value, dict):
@@ -48,15 +44,44 @@ def set_env():
         os.environ[key] = str(value)
 
 def print_env():
+    """
+    Print all environment variables from ArangoDB 'enviroment' collection to the terminal.
+    """
         for key, value in get_env().items():
             if isinstance(value, dict):
                 continue
             print(f"{key}: {value}")
 
 def print_info():
+    """
+    Print the 'INFO' section from the ArangoDB 'enviroment' collection in a formatted way.
+    """
     env_info = get_env()['INFO']
     for i, (k, v) in enumerate(env_info.items()):
         if i % 2 == 0:
             print(f"\033[90m{k}: {v}\033[0m")  # Grey color for even indices
         else:
             print(f"{k}: {v}")  # Default color for odd indices
+
+def write_to_env_file():
+    """
+    Append environment variables from ArangoDB 'enviroment' collection to the local .env file.
+    Does not overwrite existing keys or write the 'INFO' key.
+    """
+    env_doc = get_env()
+    # Only append new keys, do not overwrite .env
+    existing_keys = set()
+    try:
+        with open(".env", "r") as f:
+            for line in f:
+                if "=" in line:
+                    k = line.split("=", 1)[0].strip()
+                    existing_keys.add(k)
+    except FileNotFoundError:
+        pass  # If .env doesn't exist, we'll just create it
+
+    with open(".env", "a") as f:
+        for key, value in env_doc.items():
+            if isinstance(value, dict) or key == "INFO" or key in existing_keys:
+                continue
+            f.write(f"{key}={value}\n")