refactor: add prefix support

Author: Samk13

invenio_config/env.py

@@ -9,8 +9,6 @@
 
 """Invenio environment configuration."""
 
-from __future__ import absolute_import, print_function
-
 import ast
 import os
 
@@ -48,12 +46,21 @@ def init_app(self, app):
             app.config[varname] = value
 
 
-def _get_env_var(prefix, keys):
+def _get_env_var(prefix: str, keys: list) -> dict:
     """Retrieve environment variables with a given prefix."""
     return {k: os.environ.get(f"{prefix}_{k.upper()}") for k in keys}
 
 
-def build_db_uri():
+def _check_prefix(prefix: str) -> str:
+    """Check if prefix ends with an underscore and remove it."""
+    if not prefix:
+        return "INVENIO"
+    if prefix.endswith("_"):
+        prefix = prefix[:-1]
+    return prefix
+
+
+def build_db_uri(prefix="INVENIO"):
     """
     Build database URI from environment variables or use default.
 
@@ -66,14 +73,15 @@ def build_db_uri():
     Note: For option 3, to assert that the INVENIO_DB_* settings take effect,
     you need to set SQLALCHEMY_DATABASE_URI="" in your environment.
     """
+    prefix = _check_prefix(prefix)
     default_uri = "postgresql+psycopg2://invenio-app-rdm:invenio-app-rdm@localhost/invenio-app-rdm"
 
-    for key in ["INVENIO_SQLALCHEMY_DATABASE_URI", "SQLALCHEMY_DATABASE_URI"]:
+    for key in [f"{prefix}_SQLALCHEMY_DATABASE_URI", "SQLALCHEMY_DATABASE_URI"]:
         if uri := os.environ.get(key):
             return uri
 
     db_params = _get_env_var(
-        "INVENIO_DB", ["user", "password", "host", "port", "name", "protocol"]
+        f"{prefix}_DB", ["user", "password", "host", "port", "name", "protocol"]
     )
     if all(db_params.values()):
         uri = f"{db_params['protocol']}://{db_params['user']}:{db_params['password']}@{db_params['host']}:{db_params['port']}/{db_params['name']}"
@@ -82,7 +90,7 @@ def build_db_uri():
     return default_uri
 
 
-def build_broker_url():
+def build_broker_url(prefix="INVENIO"):
     """
     Build broker URL from environment variables or use default.
 
@@ -93,23 +101,25 @@ def build_broker_url():
     4. Default URL
     Note: see: https://docs.celeryq.dev/en/stable/userguide/configuration.html#new-lowercase-settings
     """
+    prefix = _check_prefix(prefix)
     default_url = "amqp://guest:guest@localhost:5672/"
 
-    for key in ["INVENIO_BROKER_URL", "BROKER_URL"]:
+    for key in [f"{prefix}_BROKER_URL", "BROKER_URL"]:
         if broker_url := os.environ.get(key):
             return broker_url
 
     broker_params = _get_env_var(
-        "INVENIO_AMQP_BROKER", ["user", "password", "host", "port", "protocol"]
+        f"{prefix}_AMQP_BROKER", ["user", "password", "host", "port", "protocol"]
     )
     if all(broker_params.values()):
-        vhost = f"{os.environ.get('INVENIO_AMQP_BROKER_VHOST').removeprefix('/')}"
+        broker_env_var = f"{prefix}_AMQP_BROKER_VHOST"
+        vhost = f"{os.environ.get(broker_env_var).removeprefix('/')}"
         broker_url = f"{broker_params['protocol']}://{broker_params['user']}:{broker_params['password']}@{broker_params['host']}:{broker_params['port']}/{vhost}"
         return broker_url
     return default_url
 
 
-def build_redis_url(db=None):
+def build_redis_url(db=None, prefix="INVENIO"):
     """
     Build Redis URL from environment variables or use default.
 
@@ -119,16 +129,17 @@ def build_redis_url(db=None):
     3. INVENIO_KV_CACHE_* specific environment variables
     4. Default URL
     """
+    prefix = _check_prefix(prefix)
     db = db if db is not None else 0
     default_url = f"redis://localhost:6379/{db}"
 
-    for key in ["INVENIO_CACHE_REDIS_URL", "CACHE_REDIS_URL"]:
+    for key in [f"{prefix}_CACHE_REDIS_URL", "CACHE_REDIS_URL"]:
         if cache_url := os.environ.get(key):
             if cache_url.startswith(("redis://", "rediss://", "unix://")):
                 return cache_url
 
     redis_params = _get_env_var(
-        "INVENIO_KV_CACHE", ["host", "port", "password", "protocol"]
+        f"{prefix}_KV_CACHE", ["host", "port", "password", "protocol"]
     )
 
     if redis_params["host"] and redis_params["port"]: