Skip to content

Commit

Permalink
config: granular env-based solution for connection strings
Browse files Browse the repository at this point in the history
* build db uri
* build redis url
* build mq url

partially closes: inveniosoftware/helm-invenio#112
  • Loading branch information
Samk13 committed Nov 21, 2024
1 parent f96a443 commit 85be1aa
Show file tree
Hide file tree
Showing 3 changed files with 220 additions and 8 deletions.
18 changes: 10 additions & 8 deletions invenio_app_rdm/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,8 @@
)
from werkzeug.local import LocalProxy

from invenio_app_rdm.utils.utils import build_broker_url, build_db_uri, build_redis_url

from .theme.views import notification_settings
from .users.schemas import NotificationsUserSchema, UserPreferencesNotificationsSchema

Expand Down Expand Up @@ -214,7 +216,8 @@ def _(x):
# =============
# https://flask-limiter.readthedocs.io/en/stable/#configuration

RATELIMIT_STORAGE_URI = "redis://localhost:6379/3"
RATELIMIT_STORAGE_URI = build_redis_url(db=3)

"""Storage for ratelimiter."""

# Increase defaults
Expand Down Expand Up @@ -380,7 +383,7 @@ def files_rest_permission_factory(obj, action):
# See https://invenio-accounts.readthedocs.io/en/latest/configuration.html
# See https://flask-security.readthedocs.io/en/3.0.0/configuration.html

ACCOUNTS_SESSION_REDIS_URL = "redis://localhost:6379/1"
ACCOUNTS_SESSION_REDIS_URL = build_redis_url(db=1)
"""Redis session storage URL."""

ACCOUNTS_USERINFO_HEADERS = True
Expand Down Expand Up @@ -413,7 +416,7 @@ def files_rest_permission_factory(obj, action):
# See docs.celeryproject.org/en/latest/userguide/configuration.html
# See https://flask-celeryext.readthedocs.io/en/latest/

BROKER_URL = "amqp://guest:guest@localhost:5672/"
BROKER_URL = build_broker_url()
"""URL of message broker for Celery 3 (default is RabbitMQ)."""

CELERY_BEAT_SCHEDULE = {
Expand Down Expand Up @@ -487,16 +490,15 @@ def files_rest_permission_factory(obj, action):
CELERY_BROKER_URL = BROKER_URL
"""Same as BROKER_URL to support Celery 4."""

CELERY_RESULT_BACKEND = "redis://localhost:6379/2"
CELERY_RESULT_BACKEND = build_redis_url(db=2)
"""URL of backend for result storage (default is Redis)."""

# Flask-SQLAlchemy
# ================
# See https://flask-sqlalchemy.palletsprojects.com/en/2.x/config/

SQLALCHEMY_DATABASE_URI = (
"postgresql+psycopg2://invenio-app-rdm:invenio-app-rdm@localhost/invenio-app-rdm"
)
SQLALCHEMY_DATABASE_URI = build_db_uri()

"""Database URI including user and password.
Default value is provided to make module testing easier.
Expand Down Expand Up @@ -688,7 +690,7 @@ def files_rest_permission_factory(obj, action):
# =============
# See https://flask-caching.readthedocs.io/en/latest/index.html#configuring-flask-caching # noqa

CACHE_REDIS_URL = "redis://localhost:6379/0"
CACHE_REDIS_URL = build_redis_url()
"""URL to connect to Redis server."""

CACHE_TYPE = "flask_caching.backends.redis"
Expand Down
86 changes: 86 additions & 0 deletions invenio_app_rdm/utils/utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
# -*- coding: utf-8 -*-
#
# Copyright (C) 2024 CERN.
# Copyright (C) 2024 KTH Royal Institute of Technology.
#
# Invenio App RDM is free software; you can redistribute it and/or modify it
# under the terms of the MIT License; see LICENSE file for more details.

"""Utilities for building connection strings."""

import os

from click import secho


def build_db_uri():
"""Build the database URI."""
DEFAULT_URI = "postgresql+psycopg2://invenio-app-rdm:invenio-app-rdm@localhost/invenio-app-rdm"
params = {
k: os.environ.get(f"DB_{k.upper()}")
for k in ["user", "password", "host", "port", "name"]
}

if all(params.values()):
uri = f"postgresql+psycopg2://{params['user']}:{params['password']}@{params['host']}:{params['port']}/{params['name']}"
secho(
f"Constructed database URI: '{params['user']}:***@{params['host']}:{params['port']}/{params['name']}'",
fg="blue",
)
return uri

uri = os.environ.get("SQLALCHEMY_DATABASE_URI")
if uri:
secho(f"Using SQLALCHEMY_DATABASE_URI: '{uri}'", fg="blue")
return uri

secho(f"Falling back to the default URI: '{DEFAULT_URI}'", fg="blue")
return DEFAULT_URI


def build_broker_url():
"""Build the broker URL."""
DEFAULT_BROKER_URL = "amqp://guest:guest@localhost:5672/"
params = {
k: os.environ.get(f"BROKER_{k.upper()}")
for k in ["user", "password", "host", "port"]
}

if all(params.values()):
uri = f"amqp://{params['user']}:{params['password']}@{params['host']}:{params['port']}/"
secho(
f"Constructed AMQP URL: '{params['user']}:***@{params['host']}:{params['port']}/'",
fg="blue",
)
return uri

uri = os.environ.get("BROKER_URL")
if uri:
secho(f"AMQP URI: '{uri}'", fg="blue")
return uri

secho(f"Falling back to the default URI: '{DEFAULT_BROKER_URL}'", fg="blue")
return DEFAULT_BROKER_URL


def build_redis_url(db=None):
"""Build the Redis broker URL."""
redis_host = os.environ.get("REDIS_HOST")
redis_port = os.environ.get("REDIS_PORT")
redis_password = os.environ.get("REDIS_PASSWORD")
db = db if db is not None else 0
DEFAULT_BROKER_URL = f"redis://localhost:6379/{db}"

if redis_host and redis_port:
password = f":{redis_password}@" if redis_password else ""
uri = f"redis://{password}{redis_host}:{redis_port}/{db}"
secho(f"Constructed Redis URL: '{uri}'", fg="blue")
return uri

uri = os.environ.get("BROKER_URL")
if uri and uri.startswith(("redis://", "rediss://", "unix://")):
secho(f"Using Redis BROKER_URL: '{uri}'", fg="blue")
return uri

secho(f"Falling back to the default Redis URL: '{DEFAULT_BROKER_URL}'", fg="blue")
return DEFAULT_BROKER_URL
124 changes: 124 additions & 0 deletions tests/test_utils.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# -*- coding: utf-8 -*-
#
# Copyright (C) 2021 TU Wien.
# Copyright (C) 2024 KTH Royal Institute of Technology.
#
# Invenio App RDM is free software; you can redistribute it and/or modify it
# under the terms of the MIT License; see LICENSE file for more details.
Expand All @@ -9,7 +10,10 @@

from datetime import datetime

import pytest

from invenio_app_rdm.records_ui.utils import set_default_value
from invenio_app_rdm.utils.utils import build_broker_url, build_db_uri, build_redis_url


def test_set_default_value__value():
Expand Down Expand Up @@ -65,3 +69,123 @@ def test_set_default_value__explicit_and_automatic_prefix():
dict1["metadata"]["publication_date"] == dict2["metadata"]["publication_date"]
)
assert dict1["metadata"]["publication_date"] == value


@pytest.mark.parametrize(
"env_vars, expected_uri",
[
(
{
"DB_USER": "testuser",
"DB_PASSWORD": "testpassword",
"DB_HOST": "testhost",
"DB_PORT": "5432",
"DB_NAME": "testdb",
},
"postgresql+psycopg2://testuser:testpassword@testhost:5432/testdb",
),
(
{"SQLALCHEMY_DATABASE_URI": "sqlite:///test.db"},
"sqlite:///test.db",
),
(
{},
"postgresql+psycopg2://invenio-app-rdm:invenio-app-rdm@localhost/invenio-app-rdm",
),
],
)
def test_build_db_uri(monkeypatch, env_vars, expected_uri):
"""Test building database URI."""
for key in [
"DB_USER",
"DB_PASSWORD",
"DB_HOST",
"DB_PORT",
"DB_NAME",
"SQLALCHEMY_DATABASE_URI",
]:
monkeypatch.delenv(key, raising=False)
for key, value in env_vars.items():
monkeypatch.setenv(key, value)

assert build_db_uri() == expected_uri


@pytest.mark.parametrize(
"env_vars, expected_url",
[
(
{
"BROKER_USER": "testuser",
"BROKER_PASSWORD": "testpassword",
"BROKER_HOST": "testhost",
"BROKER_PORT": "5672",
},
"amqp://testuser:testpassword@testhost:5672/",
),
(
{"BROKER_URL": "amqp://guest:guest@localhost:5672/"},
"amqp://guest:guest@localhost:5672/",
),
(
{},
"amqp://guest:guest@localhost:5672/",
),
],
)
def test_build_broker_url(monkeypatch, env_vars, expected_url):
"""Test building broker URL."""
for key in [
"BROKER_USER",
"BROKER_PASSWORD",
"BROKER_HOST",
"BROKER_PORT",
"BROKER_URL",
]:
monkeypatch.delenv(key, raising=False)
for key, value in env_vars.items():
monkeypatch.setenv(key, value)

assert build_broker_url() == expected_url


@pytest.mark.parametrize(
"env_vars, db, expected_url",
[
(
{
"REDIS_HOST": "testhost",
"REDIS_PORT": "6379",
"REDIS_PASSWORD": "testpassword",
},
2,
"redis://:testpassword@testhost:6379/2",
),
(
{
"REDIS_HOST": "testhost",
"REDIS_PORT": "6379",
},
1,
"redis://testhost:6379/1",
),
(
{"BROKER_URL": "redis://localhost:6379/0"},
None,
"redis://localhost:6379/0",
),
(
{},
4,
"redis://localhost:6379/4",
),
],
)
def test_build_redis_url(monkeypatch, env_vars, db, expected_url):
"""Test building Redis URL."""
for key in ["REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "BROKER_URL"]:
monkeypatch.delenv(key, raising=False)
for key, value in env_vars.items():
monkeypatch.setenv(key, value)

assert build_redis_url(db=db) == expected_url

0 comments on commit 85be1aa

Please sign in to comment.