From cb1e07625f0b196408a1cd12d120af269ed9d0a4 Mon Sep 17 00:00:00 2001
From: Russ Allbery <rra@lsst.org>
Date: Wed, 22 Jan 2025 08:50:51 -0800
Subject: [PATCH] Allow SecretStr for EncryptedPydanticRedisStorage

Allow the encryption key for `EncryptedPydanticRedisStorage` to be
passed in as a `SecretStr` instead of `str`. This allows easier
integration with a secret that comes from Pydantic-parsed settings
that use `SecretStr` for secrets.
---
 changelog.d/20250122_084940_rra_DM_48495.md |  3 +++
 safir/src/safir/redis/_storage.py           |  6 ++++--
 safir/tests/redis_test.py                   | 23 ++++++++++++++++++++-
 3 files changed, 29 insertions(+), 3 deletions(-)
 create mode 100644 changelog.d/20250122_084940_rra_DM_48495.md

diff --git a/changelog.d/20250122_084940_rra_DM_48495.md b/changelog.d/20250122_084940_rra_DM_48495.md
new file mode 100644
index 00000000..58d3ccc1
--- /dev/null
+++ b/changelog.d/20250122_084940_rra_DM_48495.md
@@ -0,0 +1,3 @@
+### New features
+
+- Allow the encryption key to be passed to `safir.redis.EncryptedPydanticRedisStorage` as a `pydantic.SecretStr` instead of `str`. This allows easier integration with secrets that come from Pydantic-parsed settings.
diff --git a/safir/src/safir/redis/_storage.py b/safir/src/safir/redis/_storage.py
index c2637f3d..23e2df97 100644
--- a/safir/src/safir/redis/_storage.py
+++ b/safir/src/safir/redis/_storage.py
@@ -13,7 +13,7 @@
         "Install it with `pip install safir[redis]`."
     ) from e
 from cryptography.fernet import Fernet
-from pydantic import BaseModel
+from pydantic import BaseModel, SecretStr
 
 from safir.slack.blockkit import (
     SlackCodeBlock,
@@ -241,10 +241,12 @@ def __init__(
         *,
         datatype: type[S],
         redis: redis.Redis,
-        encryption_key: str,
+        encryption_key: str | SecretStr,
         key_prefix: str = "",
     ) -> None:
         super().__init__(datatype=datatype, redis=redis, key_prefix=key_prefix)
+        if isinstance(encryption_key, SecretStr):
+            encryption_key = encryption_key.get_secret_value()
         self._fernet = Fernet(encryption_key.encode())
 
     def _serialize(self, obj: S) -> bytes:
diff --git a/safir/tests/redis_test.py b/safir/tests/redis_test.py
index b973a32f..14819d21 100644
--- a/safir/tests/redis_test.py
+++ b/safir/tests/redis_test.py
@@ -5,7 +5,7 @@
 import pytest
 import redis.asyncio as redis
 from cryptography.fernet import Fernet
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, SecretStr
 
 from safir.redis import (
     DeserializeError,
@@ -77,6 +77,27 @@ async def test_encrypted_pydantic_redis_storage(
     await basic_testing(storage)
 
 
+@pytest.mark.asyncio
+async def test_encrypted_pydantic_redis_storage_secretstr(
+    redis_client: redis.Redis,
+) -> None:
+    """Test that the secret can be passed as str or SecretStr."""
+    encryption_key = Fernet.generate_key().decode()
+    storage = EncryptedPydanticRedisStorage(
+        datatype=DemoModel, redis=redis_client, encryption_key=encryption_key
+    )
+
+    await storage.store("mark42", DemoModel(name="Mark", value=42))
+
+    storage = EncryptedPydanticRedisStorage(
+        datatype=DemoModel,
+        redis=redis_client,
+        encryption_key=SecretStr(encryption_key),
+    )
+
+    assert await storage.get("mark42") == DemoModel(name="Mark", value=42)
+
+
 class PetModel(BaseModel):
     id: int
     name: str