From a5467c6e1997e3d6bd4ee81748411cd4b870ce0e Mon Sep 17 00:00:00 2001 From: Jens L Date: Tue, 21 May 2024 20:15:49 +0200 Subject: [PATCH] root: add primary-replica db router (#9479) * root: add primary-replica db router Signed-off-by: Jens Langhammer * copy all settings for database replicas Signed-off-by: Marc 'risson' Schmitt * refresh read replicas config, switch to using a dict instead of a list for easier refresh Signed-off-by: Marc 'risson' Schmitt * add test for get_keys Signed-off-by: Marc 'risson' Schmitt * fix getting override Signed-off-by: Marc 'risson' Schmitt * lint Signed-off-by: Marc 'risson' Schmitt * nosec Signed-off-by: Marc 'risson' Schmitt * small fixes Signed-off-by: Marc 'risson' Schmitt * fix replica settings Signed-off-by: Marc 'risson' Schmitt * generate config: add a dummy read replica Signed-off-by: Marc 'risson' Schmitt * add doc Signed-off-by: Marc 'risson' Schmitt * add healthchecks for replicas Signed-off-by: Marc 'risson' Schmitt * fix Signed-off-by: Marc 'risson' Schmitt * add note about hot reloading Signed-off-by: Marc 'risson' Schmitt --------- Signed-off-by: Jens Langhammer Signed-off-by: Marc 'risson' Schmitt Co-authored-by: Marc 'risson' Schmitt --- authentik/lib/config.py | 6 +++++ authentik/lib/default.yml | 4 +++ authentik/lib/tests/test_config.py | 6 +++++ authentik/root/db/base.py | 13 ++++++--- authentik/root/monitoring.py | 4 +-- authentik/root/settings.py | 20 ++++++++++++-- authentik/tenants/db.py | 29 +++++++++++++++++++++ scripts/generate_config.py | 3 +++ website/docs/installation/configuration.mdx | 16 ++++++++++++ 9 files changed, 94 insertions(+), 7 deletions(-) create mode 100644 authentik/tenants/db.py diff --git a/authentik/lib/config.py b/authentik/lib/config.py index b35d0d6eef..f4a808c549 100644 --- a/authentik/lib/config.py +++ b/authentik/lib/config.py @@ -304,6 +304,12 @@ class ConfigLoader: """Wrapper for get that converts value into boolean""" return str(self.get(path, default)).lower() == "true" + def get_keys(self, path: str, sep=".") -> list[str]: + """List attribute keys by using yaml path""" + root = self.raw + attr: Attr = get_path_from_dict(root, path, sep=sep, default=Attr({})) + return attr.keys() + def get_dict_from_b64_json(self, path: str, default=None) -> dict: """Wrapper for get that converts value from Base64 encoded string into dictionary""" config_value = self.get(path) diff --git a/authentik/lib/default.yml b/authentik/lib/default.yml index c2b2a72163..183ad8341c 100644 --- a/authentik/lib/default.yml +++ b/authentik/lib/default.yml @@ -10,6 +10,10 @@ postgresql: use_pgpool: false test: name: test_authentik + read_replicas: {} + # For example + # 0: + # host: replica1.example.com listen: listen_http: 0.0.0.0:9000 diff --git a/authentik/lib/tests/test_config.py b/authentik/lib/tests/test_config.py index d436787375..81c1150839 100644 --- a/authentik/lib/tests/test_config.py +++ b/authentik/lib/tests/test_config.py @@ -169,3 +169,9 @@ class TestConfig(TestCase): self.assertEqual(config.get("cache.timeout_flows"), "32m") self.assertEqual(config.get("cache.timeout_policies"), "3920ns") self.assertEqual(config.get("cache.timeout_reputation"), "298382us") + + def test_get_keys(self): + """Test get_keys""" + config = ConfigLoader() + config.set("foo.bar", "baz") + self.assertEqual(list(config.get_keys("foo")), ["bar"]) diff --git a/authentik/root/db/base.py b/authentik/root/db/base.py index c1297998e6..2385e1b972 100644 --- a/authentik/root/db/base.py +++ b/authentik/root/db/base.py @@ -10,8 +10,15 @@ class DatabaseWrapper(BaseDatabaseWrapper): def get_connection_params(self): """Refresh DB credentials before getting connection params""" - CONFIG.refresh("postgresql.password") conn_params = super().get_connection_params() - conn_params["user"] = CONFIG.get("postgresql.user") - conn_params["password"] = CONFIG.get("postgresql.password") + + prefix = "postgresql" + if self.alias.startswith("replica_"): + prefix = f"postgresql.read_replicas.{self.alias.removeprefix('replica_')}" + + for setting in ("host", "port", "user", "password"): + conn_params[setting] = CONFIG.refresh(f"{prefix}.{setting}") + if conn_params[setting] is None and self.alias.startswith("replica_"): + conn_params[setting] = CONFIG.refresh(f"postgresql.{setting}") + return conn_params diff --git a/authentik/root/monitoring.py b/authentik/root/monitoring.py index 5234a81d6d..b6ce04306e 100644 --- a/authentik/root/monitoring.py +++ b/authentik/root/monitoring.py @@ -47,8 +47,8 @@ class ReadyView(View): def dispatch(self, request: HttpRequest) -> HttpResponse: try: - db_conn = connections["default"] - _ = db_conn.cursor() + for db_conn in connections.all(): + _ = db_conn.cursor() except OperationalError: # pragma: no cover return HttpResponse(status=503) try: diff --git a/authentik/root/settings.py b/authentik/root/settings.py index 12d4350104..2b416cf95f 100644 --- a/authentik/root/settings.py +++ b/authentik/root/settings.py @@ -293,7 +293,7 @@ DATABASES = { "NAME": CONFIG.get("postgresql.name"), "USER": CONFIG.get("postgresql.user"), "PASSWORD": CONFIG.get("postgresql.password"), - "PORT": CONFIG.get_int("postgresql.port"), + "PORT": CONFIG.get("postgresql.port"), "SSLMODE": CONFIG.get("postgresql.sslmode"), "SSLROOTCERT": CONFIG.get("postgresql.sslrootcert"), "SSLCERT": CONFIG.get("postgresql.sslcert"), @@ -313,7 +313,23 @@ if CONFIG.get_bool("postgresql.use_pgbouncer", False): # https://docs.djangoproject.com/en/4.0/ref/databases/#persistent-connections DATABASES["default"]["CONN_MAX_AGE"] = None # persistent -DATABASE_ROUTERS = ("django_tenants.routers.TenantSyncRouter",) +for replica in CONFIG.get_keys("postgresql.read_replicas"): + _database = DATABASES["default"].copy() + for setting in DATABASES["default"].keys(): + default = object() + if setting in ("TEST",): + continue + override = CONFIG.get( + f"postgresql.read_replicas.{replica}.{setting.lower()}", default=default + ) + if override is not default: + _database[setting] = override + DATABASES[f"replica_{replica}"] = _database + +DATABASE_ROUTERS = ( + "authentik.tenants.db.FailoverRouter", + "django_tenants.routers.TenantSyncRouter", +) # Email # These values should never actually be used, emails are only sent from email stages, which diff --git a/authentik/tenants/db.py b/authentik/tenants/db.py new file mode 100644 index 0000000000..e6735a01de --- /dev/null +++ b/authentik/tenants/db.py @@ -0,0 +1,29 @@ +from random import choice + +from django.conf import settings + + +class FailoverRouter: + """Support an primary/read-replica PostgreSQL setup (reading from replicas + and write to primary only)""" + + def __init__(self) -> None: + super().__init__() + self.database_aliases = set(settings.DATABASES.keys()) + self.read_replica_aliases = list(self.database_aliases - {"default"}) + self.replica_enabled = len(self.read_replica_aliases) > 0 + + def db_for_read(self, model, **hints): + if not self.replica_enabled: + return "default" + return choice(self.read_replica_aliases) # nosec + + def db_for_write(self, model, **hints): + return "default" + + def allow_relation(self, obj1, obj2, **hints): + """Relations between objects are allowed if both objects are + in the primary/replica pool.""" + if obj1._state.db in self.database_aliases and obj2._state.db in self.database_aliases: + return True + return None diff --git a/scripts/generate_config.py b/scripts/generate_config.py index 1f81aa224d..2fb164aa90 100644 --- a/scripts/generate_config.py +++ b/scripts/generate_config.py @@ -12,6 +12,9 @@ with open("local.env.yml", "w", encoding="utf-8") as _config: "secret_key": generate_id(), "postgresql": { "user": "postgres", + "read_replicas": { + "0": {}, + }, }, "outposts": { "container_image_base": "ghcr.io/goauthentik/dev-%(type)s:gh-%(build_hash)s", diff --git a/website/docs/installation/configuration.mdx b/website/docs/installation/configuration.mdx index ddcabd51e4..2fb7c9c667 100644 --- a/website/docs/installation/configuration.mdx +++ b/website/docs/installation/configuration.mdx @@ -77,6 +77,22 @@ To check if your config has been applied correctly, you can run the following co - `AUTHENTIK_POSTGRESQL__SSLCERT`: Path to x509 client certificate to authenticate to server - `AUTHENTIK_POSTGRESQL__SSLKEY`: Path to private key of `SSLCERT` certificate +Additionally, databases used only for read operations can be configured. Increase the number in the following configuration variables for each read replica. + +- `AUTHENTIK_POSTGRESQL__READ_REPLICAS__0__HOST`: same as above +- `AUTHENTIK_POSTGRESQL__READ_REPLICAS__0__NAME`: same as above +- `AUTHENTIK_POSTGRESQL__READ_REPLICAS__0__USER`: same as above +- `AUTHENTIK_POSTGRESQL__READ_REPLICAS__0__PORT`: same as above +- `AUTHENTIK_POSTGRESQL__READ_REPLICAS__0__PASSWORD`: same as above +- `AUTHENTIK_POSTGRESQL__READ_REPLICAS__0__SSLMODE`: same as above +- `AUTHENTIK_POSTGRESQL__READ_REPLICAS__0__SSLROOTCERT`: same as above +- `AUTHENTIK_POSTGRESQL__READ_REPLICAS__0__SSLCERT`: same as above +- `AUTHENTIK_POSTGRESQL__READ_REPLICAS__0__SSLKEY`: same as above + +Note that `USE_PGBOUNCER` and `USE_PGPOOL` are inherited from the main database configuration and are not overridable per read replica. By default, if read replicas are configured, the main database is not used for reads. If you'd like it to be included for reads, add it as a read replica. + +All PostgreSQL settings, apart from `USE_PGBOUNCER` and `USE_PGPOOL`, support hot-reloading. Adding and removing read replicas doesn't support hot-reloading. + ## Redis Settings - `AUTHENTIK_REDIS__HOST`: Redis server host when not using configuration URL