Add function to update existing RabbitMQ configuration with latest te…

…mplates (#45)

Includes `update-rabbitmq-config` CLI entrypoint
Includes Unit Tests

Co-authored-by: Daniel McKnight <daniel@neon.ai>

neon_diana_utils/cli.py

@@ -41,6 +41,15 @@ def neon_diana_cli(version: bool = False):
         click.echo(f"Diana version {__version__}")
 
 
+# Generic Utilities
+@neon_diana_cli.command(help="Update RabbitMQ Configuration")
+@click.argument("rabbitmq_json_path", default=None, required=False)
+def update_rabbitmq_config(rabbitmq_json_path):
+    from neon_diana_utils.configuration import update_rmq_config
+    updated_file = update_rmq_config(rabbitmq_json_path)
+    click.echo(f"Updated configuration at: {updated_file}")
+
+
 # Core
 @neon_diana_cli.command(help="Configure Neon Core")
 @click.option("--username", "-u", help="RabbitMQ username for Neon AI")

neon_diana_utils/configuration.py

@@ -195,6 +195,41 @@ def make_keys_config(write_config: bool,
     return config
 
 
+def update_rmq_config(config_file: str = None) -> str:
+    """
+    Update an existing RabbitMQ configuration with new definitions from DIANA.
+    This can be used to handle added service users without changing existing
+    ones and to reset/update permissions and vhosts.
+    @param config_file: Path to file to be updated
+    @returns: Path to updated file
+    """
+    if not config_file:
+        config_file = join(xdg_config_home(), "diana", "diana-backend",
+                           "rabbitmq.json")
+    else:
+        config_file = expanduser(config_file)
+
+    if not isfile(config_file):
+        raise FileNotFoundError(config_file)
+
+    with open(config_file) as f:
+        real_config = json.load(f)
+    new_config = generate_rmq_config("", "")
+    existing_users = (user['name'] for user in real_config['users'])
+    for user in new_config['users']:
+        if user['name'] in existing_users:
+            continue
+        LOG.info(f"Adding user: {user['name']}")
+        real_config['users'].append(user)
+    real_config['vhosts'] = new_config['vhosts']
+    real_config['permissions'] = new_config['permissions']
+
+    shutil.move(config_file, f"{config_file}.old")
+    with open(config_file, 'w+') as f:
+        json.dump(real_config, f, indent=2)
+    return config_file
+
+
 def generate_rmq_config(admin_username: str, admin_password: str,
                         output_file: str = None) -> dict:
     """
@@ -215,10 +250,12 @@ def generate_rmq_config(admin_username: str, admin_password: str,
             continue
         user['password'] = secrets.token_urlsafe(32)
 
-    base_config['users'].append({'name': admin_username,
-                                 'password': admin_password,
-                                 'tags': ['administrator']})
-
+    if admin_username and admin_password:
+        base_config['users'].append({'name': admin_username,
+                                     'password': admin_password,
+                                     'tags': ['administrator']})
+    else:
+        LOG.debug("Not adding unconfigured admin user")
     if output_file and validate_output_path(output_file):
         with open(output_file, 'w+') as f:
             json.dump(base_config, f, indent=2)
@@ -261,23 +298,24 @@ def update_env_file(env_file: str):
 
 
 def _get_neon_mq_user_config(mq_user: Optional[str], mq_pass: Optional[str],
-                             backend_config: str) -> dict:
+                             rmq_config: str) -> dict:
     """
     Get MQ config for neon core.
     @param mq_user: RabbitMQ Neon username
     @param mq_pass: RabbitMQ Neon password
-    @param backend_config: Path to Diana Backend configuration file to import
+    @param rmq_config: Path to RabbitMQ configuration file to import
     @returns dict user config to connect Neon Core to an MQ instance
     """
     # Check for passed or previously configured MQ user
-    if not all((mq_user, mq_pass)) and isfile(backend_config):
-        if click.confirm(f"Import Neon MQ user from {backend_config}?"):
-            with open(backend_config) as f:
-                config = yaml.safe_load(f)
-            user_config = config.get('MQ', {}).get('users',
-                                                   {}).get('chat_api_proxy', {})
-            mq_user = user_config.get('user')
-            mq_pass = user_config.get('password')
+    if not all((mq_user, mq_pass)) and isfile(rmq_config):
+        if click.confirm(f"Import Neon MQ user from {rmq_config}?"):
+            with open(rmq_config) as f:
+                config = json.load(f)
+            for user in config['users']:
+                if "core" in user['tags']:
+                    mq_user = user['name']
+                    mq_pass = user['password']
+                    break
 
     # Interactively configure MQ authentication
     user_config = {"user": mq_user, "password": mq_pass}
@@ -399,7 +437,7 @@ def configure_neon_core(mq_user: str = None,
 
     # Validate output paths
     output_path = expanduser(output_path or join(xdg_config_home(), "diana"))
-    backend_config = join(output_path, "diana-backend", "diana.yaml")
+    rmq_config = join(output_path, "diana-backend", "rabbitmq.json")
     # Output to `core` subdirectory
     if not validate_output_path(join(output_path, "neon-core")):
         click.echo(f"Path exists: {output_path}")
@@ -426,7 +464,7 @@ def configure_neon_core(mq_user: str = None,
 
     try:
         # Get MQ User Configuration
-        user_config = _get_neon_mq_user_config(mq_user, mq_pass, backend_config)
+        user_config = _get_neon_mq_user_config(mq_user, mq_pass, rmq_config)
         if not all((user_config['user'], user_config['password'])):
             # TODO: Prompt to configure MQ server/port?
             mq_config = dict()

tests/test_diana_utils.py

@@ -76,6 +76,24 @@ def test_make_keys_config(self, confirm):
 
         os.remove(test_output_file)
 
+    def test_update_rmq_config(self):
+        from neon_diana_utils.configuration import update_rmq_config
+        test_file = join(dirname(__file__), "test_rabbitmq.json")
+        update_rmq_config(test_file)
+        self.assertTrue(isfile(test_file))
+        self.assertTrue(isfile(f"{test_file}.old"))
+        with open(test_file) as f:
+            new_config = json.load(f)
+        with open(f"{test_file}.old") as f:
+            old_config = json.load(f)
+        self.assertIsInstance(new_config, dict)
+        self.assertIsInstance(old_config, dict)
+        for user in old_config['users']:
+            self.assertIn(user, new_config['users'])
+
+        os.remove(test_file)
+        shutil.move(f"{test_file}.old", test_file)
+
     def test_generate_rmq_config(self):
         from neon_diana_utils.configuration import generate_rmq_config
         test_output_file = join(dirname(__file__), "test_rmq.json")

tests/test_rabbitmq.json

@@ -0,0 +1,210 @@
+{
+  "users": [
+    {
+      "name": "neon_api_utils",
+      "password": "Klatchat2021",
+      "tags": [
+        "backend",
+        "user"
+      ]
+    },
+    {
+      "name": "neon_metrics",
+      "password": "old_metrics",
+      "tags": [
+        "backend",
+        "service"
+      ]
+    },
+    {
+      "name": "neon_coupons",
+      "password": "old_coupons",
+      "tags": [
+        "backend",
+        "service"
+      ]
+    },
+    {
+      "name": "neon_email",
+      "password": "old_email",
+      "tags": [
+        "backend",
+        "service"
+      ]
+    },
+    {
+      "name": "neon_script_parser",
+      "password": "old_script_parser",
+      "tags": [
+        "backend",
+        "service"
+      ]
+    },
+    {
+      "name": "neon_api",
+      "password": "old_neon_api",
+      "tags": [
+        "backend",
+        "service"
+      ]
+    },
+    {
+      "name": "neon_libretranslate",
+      "password": "old_libretranslate",
+      "tags": [
+        "backend",
+        "service"
+      ]
+    }
+  ],
+  "vhosts": [
+    {
+      "name": "/neon_emails"
+    },
+    {
+      "name": "/neon_api"
+    },
+    {
+      "name": "/neon_script_parser"
+    },
+    {
+      "name": "/neon_metrics"
+    },
+    {
+      "name": "/neon_coupons"
+    },
+    {
+      "name": "/neon_testing"
+    },
+    {
+      "name": "/translation"
+    },
+    {
+      "name": "/llm"
+    },
+    {
+      "name": "/neon_chat_api"
+    }
+  ],
+  "permissions": [
+    {
+      "user": "neon_core",
+      "vhost": "/neon_chat_api",
+      "configure": ".*",
+      "write": ".*",
+      "read": ".*"
+    },
+    {
+      "user": "neon_api_utils",
+      "vhost": "/neon_chat_api",
+      "configure": ".*",
+      "write": ".*",
+      "read": ".*(?!_request).*"
+    },
+    {
+      "user": "neon_llm_chatgpt",
+      "vhost": "/llm",
+      "configure": ".*",
+      "write": ".*",
+      "read": "chat_gpt_input"
+    },
+    {
+      "user": "neon_llm_fastchat",
+      "vhost": "/llm",
+      "configure": ".*",
+      "write": ".*",
+      "read": "fastchat_input"
+    },
+    {
+      "user": "neon_api_utils",
+      "vhost": "/llm",
+      "configure": ".*",
+      "write": ".*",
+      "read": ".*(?!_input).*"
+    },
+    {
+      "user": "neon_libretranslate",
+      "vhost": "/translation",
+      "configure": ".*",
+      "write": ".*",
+      "read": ".*"
+    },
+    {
+      "user": "neon_api_utils",
+      "vhost": "/neon_coupons",
+      "configure": ".*",
+      "write": ".*",
+      "read": ".*"
+    },
+    {
+      "user": "neon_email",
+      "vhost": "/neon_emails",
+      "configure": ".*",
+      "write": ".*",
+      "read": ".*"
+    },
+    {
+      "user": "neon_api_utils",
+      "vhost": "/neon_emails",
+      "configure": ".*",
+      "write": ".*",
+      "read": "^(?!neon_emails_input).*"
+    },
+    {
+      "user": "neon_api",
+      "vhost": "/neon_api",
+      "configure": ".*",
+      "write": ".*",
+      "read": ".*"
+    },
+    {
+      "user": "neon_api_utils",
+      "vhost": "/neon_metrics",
+      "configure": ".*",
+      "write": ".*",
+      "read": ""
+    },
+    {
+      "user": "neon_api_utils",
+      "vhost": "/neon_api",
+      "configure": "./*",
+      "write": "./*",
+      "read": "./*"
+    },
+    {
+      "user": "neon_api_utils",
+      "vhost": "/neon_script_parser",
+      "configure": ".*",
+      "write": ".*",
+      "read": ".*"
+    },
+    {
+      "user": "neon_script_parser",
+      "vhost": "/neon_script_parser",
+      "configure": ".*",
+      "write": ".*",
+      "read": ".*"
+    },
+    {
+      "user": "neon_metrics",
+      "vhost": "/neon_metrics",
+      "configure": ".*",
+      "write": ".*",
+      "read": ".*"
+    },
+    {
+      "user": "neon_api_utils",
+      "vhost": "/neon_testing",
+      "configure": "./*",
+      "write": "./*",
+      "read": "./*"
+    },
+    {
+      "user": "neon_coupons",
+      "vhost": "/neon_coupons",
+      "configure": ".*",
+      "write": ".*",
+      "read": ".*"
+    }
+  ]
+}