feat: Labels for encrypted fields (#38075)

tests/unit_tests/db_engine_specs/test_base.py

@@ -360,6 +360,51 @@ def test_unmask_encrypted_extra() -> None:
     )
 
 
+@pytest.mark.parametrize(
+    "masked_encrypted_extra,expected_result",
+    [
+        (
+            {
+                "$.credentials_info.private_key": "Private Key",
+                "$.access_token": "Access Token",
+            },
+            {
+                "$.credentials_info.private_key",
+                "$.access_token",
+            },
+        ),
+        (
+            {
+                "$.credentials_info.private_key",
+                "$.access_token",
+            },
+            {
+                "$.credentials_info.private_key",
+                "$.access_token",
+            },
+        ),
+        (
+            None,
+            {"$.*"},
+        ),
+    ],
+)
+def test_encrypted_extra_sensitive_field_paths_from_dict(
+    masked_encrypted_extra: set[str] | dict[str, str] | None,
+    expected_result: set[str],
+) -> None:
+    """
+    Test that `encrypted_extra_sensitive_field_paths` extracts the keys
+    when `encrypted_extra_sensitive_fields` is a dict.
+    """
+
+    class DictFieldsSpec(BaseEngineSpec):
+        if masked_encrypted_extra:
+            encrypted_extra_sensitive_fields = masked_encrypted_extra
+
+    assert DictFieldsSpec.encrypted_extra_sensitive_field_paths() == expected_result
+
+
 def test_impersonate_user_backwards_compatible(mocker: MockerFixture) -> None:
     """
     Test that the `impersonate_user` method calls the original methods it replaced.