Postgresql optional parameters (#344)

* support for optional psql params Signed-off-by: Rafał Safin <[email protected]> * add additional conn params for pg Signed-off-by: rafsaf <[email protected]> --------- Signed-off-by: rafsaf <[email protected]>
rafsaf · Dec 15, 2024 · 10f1608 · 10f1608
1 parent 6ae82bf
commit 10f1608
Show file tree

Hide file tree

Showing 8 changed files with 114 additions and 68 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -5,12 +5,12 @@ repos:
       - id: check-yaml
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.8.2
+    rev: v0.8.3
     hooks:
       - id: ruff-format
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.8.2
+    rev: v0.8.3
     hooks:
       - id: ruff
         args: [--fix]

diff --git a/docs/backup_targets/postgresql.md b/docs/backup_targets/postgresql.md
@@ -27,6 +27,17 @@ POSTGRESQL_SOME_STRING="host=... password=... cron_rule=..."
 | max_backups        | int                  | Soft limit how many backups can live at once for backup target. Defaults to `7`. This must makes sense with cron expression you use. For example if you want to have `7` day retention, and make backups at 5:00, `max_backups=7` is fine, but if you make `4` backups per day, you would need `max_backups=28`. Limit is soft and can be exceeded if no backup is older than value specified in min_retention_days. Min `1` and max `998`. Defaults to enviornment variable BACKUP_MAX_NUMBER, see [Configuration](./../configuration.md). | BACKUP_MAX_NUMBER         |
 | min_retention_days | int                  | Hard minimum backups lifetime in days. Ogion won't ever delete files before, regardles of other options. Min `0` and max `36600`. Defaults to enviornment variable BACKUP_MIN_RETENTION_DAYS, see [Configuration](./../configuration.md).                                                                                                                                                                                                                                                                                                   | BACKUP_MIN_RETENTION_DAYS |
 
+## Additional connection params
+
+Extra variables that starts with `conn_` will be passed AS IS to psql command underthehood as url-encoded connection params:
+
+For example you can use it for SSL setup:
+
+- `conn_sslmode=verify-ca`
+- `conn_sslrootcert=path-to-mounted-server-ca-file`
+- `conn_sslcert=path-to-mounted-client-ca-file`
+- `conn_sslkey=path-to-mounted-client-key-file`
+
 ## Examples
 
 ```bash
@@ -38,6 +49,9 @@ POSTGRESQL_SECOND_DB='host=10.0.0.1 port=5432 user=foo password=change_me! db=ba
 
 # 3. PostgreSQL in local network with backup on every 6 hours at '15 with max number of backups of 20
 POSTGRESQL_THIRD_DB='host=192.168.1.5 port=5432 user=root password=change_me_please! db=project cron_rule=15 */3 * * * max_backups=20'
+
+# 4. PostgreSQL connected using sslmode require
+POSTGRESQL_4_DB_SSL='host=localhost port=5432 password=secret cron_rule=* * * * * conn_sslmode=require'
 ```
 
 <br>

diff --git a/ogion/backup_targets/postgresql.py b/ogion/backup_targets/postgresql.py
@@ -62,10 +62,23 @@ def _get_escaped_conn_uri(self) -> str:
         pgpass_file = self._init_pgpass_file()
         encoded_user = urllib.parse.quote_plus(self.target_model.user)
         encoded_db = urllib.parse.quote_plus(self.target_model.db)
+
+        params = {"passfile": pgpass_file}
+        if self.target_model.model_extra is not None:
+            for param, value in self.target_model.model_extra.items():
+                if not param.startswith("conn_"):
+                    continue
+
+                params[param.removeprefix("conn_")] = value
+
+        log.debug("psql connection params: %s", params)
+
         uri = (
             f"postgresql://{encoded_user}@{self.target_model.host}:{self.target_model.port}/{encoded_db}?"
-            f"passfile={pgpass_file}"
-        )
+        ) + urllib.parse.urlencode(params)
+
+        log.debug("psql connection url: %s", uri)
+
         escaped_uri = shlex.quote(uri)
         return escaped_uri
 
@@ -100,8 +113,8 @@ def _postgres_connection(self) -> str:
             break
         if version is None:  # pragma: no cover
             msg = (
-                "postgres_connection error processing sql result, "
-                "version unknown: {result}"
+                f"postgres_connection error processing sql result, "
+                f"version unknown: {result}"
             )
             log.error(msg)
             raise ValueError(msg)

diff --git a/ogion/core.py b/ogion/core.py
@@ -21,6 +21,7 @@
 
 SAFE_LETTER_PATTERN = re.compile(r"[^A-Za-z0-9_]*")
 DATETIME_BACKUP_FILE_PATTERN = re.compile(r"_[0-9]{8}_[0-9]{4}_")
+MODEL_SPLIT_EQUATION_PATTERN = re.compile(r"( \w*\=|^\w*\=)")
 
 _BM = TypeVar("_BM", bound=BaseModel)
 
@@ -116,26 +117,24 @@ def _validate_model(
     env_name: str,
     env_value: str,
     target: type[_BM],
-    value_whitespace_split: bool = False,
 ) -> _BM:
     target_name: str = target.__name__.lower()
     log.info("validating %s variable: `%s`", target_name, env_name)
     log.debug("%s=%s", target_name, env_value)
     try:
         env_value_parts = env_value.strip()
+        fields_matches = [
+            match.group()
+            for match in MODEL_SPLIT_EQUATION_PATTERN.finditer(env_value_parts)
+        ]
         target_kwargs: dict[str, Any] = {"env_name": env_name}
-        for field_name in target.model_fields.keys():
-            if env_value_parts.startswith(f"{field_name}="):
-                f = f"{field_name}="
-            else:
-                f = f" {field_name}="
-            if f in env_value_parts:
-                _, val = env_value_parts.split(f, maxsplit=1)
-                for other_field in target.model_fields.keys():
-                    val = val.split(f" {other_field}=")[0]
-                if value_whitespace_split:
-                    val = val.split()[0]
-                target_kwargs[field_name] = val
+
+        while fields_matches:
+            field_match = fields_matches.pop()
+            rest, value = env_value_parts.split(field_match, maxsplit=1)
+            env_value_parts = rest.rstrip()
+            target_kwargs[field_match.removesuffix("=").strip()] = value
+
         log.debug("calculated arguments: %s", target_kwargs)
         validated_target = target.model_validate(target_kwargs)
     except Exception:
@@ -173,7 +172,6 @@ def create_provider_model() -> upload_provider_models.ProviderModel:
         "backup_provider",
         config.options.BACKUP_PROVIDER,
         upload_provider_models.ProviderModel,
-        value_whitespace_split=True,
     )
     target_model_cls = provider_map[base_provider.name]
     return _validate_model(

diff --git a/ogion/models/backup_target_models.py b/ogion/models/backup_target_models.py
@@ -45,6 +45,10 @@ class PostgreSQLTargetModel(TargetModel):
     db: str = "postgres"
     password: SecretStr
 
+    model_config = ConfigDict(
+        extra="allow",
+    )
+
 
 class MariaDBTargetModel(TargetModel):
     name: config.BackupTargetEnum = config.BackupTargetEnum.MARIADB
@@ -54,6 +58,10 @@ class MariaDBTargetModel(TargetModel):
     db: str = "mariadb"
     password: SecretStr
 
+    model_config = ConfigDict(
+        extra="allow",
+    )
+
 
 class SingleFileTargetModel(TargetModel):
     name: config.BackupTargetEnum = config.BackupTargetEnum.FILE

diff --git a/poetry.lock b/poetry.lock
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -48,18 +48,22 @@ def _to_target_model(
     model: type[TM],
 ) -> TM:
     DB_VERSION_BY_ENV_VAR[compose_db.name] = compose_db.version
-    return model(
-        env_name=compose_db.name,
-        cron_rule="* * * * *",
-        host=compose_db.name if DOCKER_TESTS else "localhost",
-        port=(
-            int(compose_db.ports[0].split(":")[1])
-            if DOCKER_TESTS
-            else int(compose_db.ports[0].split(":")[0])
-        ),
-        password=SecretStr(DB_PWD),
-        db=DB_NAME,
-        user=DB_USERNAME,
+    return model.model_validate(
+        {
+            "env_name": compose_db.name,
+            "cron_rule": "* * * * *",
+            "host": compose_db.name if DOCKER_TESTS else "localhost",
+            "port": (
+                int(compose_db.ports[0].split(":")[1])
+                if DOCKER_TESTS
+                else int(compose_db.ports[0].split(":")[0])
+            ),
+            "password": SecretStr(DB_PWD),
+            "db": DB_NAME,
+            "user": DB_USERNAME,
+            "conn_sslmode": "prefer",
+            "dummy_extra": "test",
+        }
     )
 
 

diff --git a/tests/test_core.py b/tests/test_core.py
@@ -195,6 +195,16 @@ def test_run_create_age_archive_can_be_decrypted(
         ],
         True,
     ),
+    (
+        [
+            (
+                "POSTGRESQL_FIRST_DB",
+                "host=localhost port5432 password=secret "
+                "cron_rule=* * * * * ssl_mode=require",
+            ),
+        ],
+        True,
+    ),
 ]
-Original file line number
+Diff line change
@@ Expand Up / @@ -195,6 +195,16 @@ def test_run_create_age_archive_can_be_decrypted( @@
             ],
             True,
         ),
+        (
+            [
+                (
+                    "POSTGRESQL_FIRST_DB",
+                    "host=localhost port5432 password=secret "
+                    "cron_rule=* * * * * ssl_mode=require",
+                ),
+            ],
+            True,
+        ),
     ]
@@ Expand Down @@