Repository: peopledoc/vault-cli

Issue:
Add cache
Make it so that when using as CLI, a same secret won't be read twice.

When using as lib, a context manager will allow you to have the same behaviour.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/peopledoc/vault-cli/pull/103?src=pr&el=h1) Report
> Merging [#103](https://codecov.io/gh/peopledoc/vault-cli/pull/103?src=pr&el=desc) into [master](https://codecov.io/gh/peopledoc/vault-cli/commit/3434698d553bea16f6e8a1f4dacd8398297543bd?src=pr&el=desc) will **not change** coverage.
> The diff coverage is `100%`.

[![Impacted file tree graph](https://codecov.io/gh/peopledoc/vault-cli/pull/103/graphs/tree.svg?width=650&token=fjUxXKxpOh&height=150&src=pr)](https://codecov.io/gh/peopledoc/vault-cli/pull/103?src=pr&el=tree)

```diff
@@          Coverage Diff          @@
##           master   #103   +/-   ##
=====================================
  Coverage     100%   100%           
=====================================
  Files          12     12           
  Lines         689    719   +30     
  Branches       86     88    +2     
=====================================
+ Hits          689    719   +30
```

| Flag | Coverage Δ | |
|---|---|---|
| #integration | `68.84% <87.09%> (+0.77%)` | :arrow_up: |
| #unit | `100% <100%> (ø)` | :arrow_up: |

| [Impacted Files](https://codecov.io/gh/peopledoc/vault-cli/pull/103?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [vault\_cli/client.py](https://codecov.io/gh/peopledoc/vault-cli/pull/103/diff?src=pr&el=tree#diff-dmF1bHRfY2xpL2NsaWVudC5weQ==) | `100% <100%> (ø)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/peopledoc/vault-cli/pull/103?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/peopledoc/vault-cli/pull/103?src=pr&el=footer). Last update [3434698...d90b6a8](https://codecov.io/gh/peopledoc/vault-cli/pull/103?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/.gitignore b/.gitignore
index 4b93a6c..7686c12 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,6 +7,5 @@ dist
 htmlcov
 .tox
 *.pyc
-/.vault-token
-/.vault.yml
-/.mypy_cache
\ No newline at end of file
+/vault.yml
+/.mypy_cache
diff --git a/vault_cli/client.py b/vault_cli/client.py
index 9ee3e72..622ea6c 100644
--- a/vault_cli/client.py
+++ b/vault_cli/client.py
@@ -1,8 +1,9 @@
 import contextlib
+import functools
 import json
 import logging
 import pathlib
-from typing import Dict, Iterable, Optional, Tuple, Type
+from typing import Dict, Iterable, Optional, Set, Tuple, Type
 
 import hvac
 import jinja2
@@ -73,6 +74,21 @@ def get_client_class() -> Type["VaultClientBase"]:
     return VaultClient
 
 
+def caching(method):
+    @functools.wraps(method)
+    def wrapper(self, *args, **kwargs):
+        override_cache = self.cache is None
+        if override_cache:
+            self.cache = {}
+        try:
+            return method(self, *args, **kwargs)
+        finally:
+            if override_cache:
+                self.cache = None
+
+    return wrapper
+
+
 class VaultClientBase:
 
     saved_settings: Optional[types.SettingsDict] = None
@@ -105,6 +121,8 @@ class VaultClientBase:
         self.password = password
         self.safe_write = safe_write
         self.render = render
+        self.cache: Optional[Dict[str, types.JSONDict]] = None
+        self._currently_fetching: Set[str] = set()
 
     @property
     def base_path(self):
@@ -113,7 +131,7 @@ class VaultClientBase:
     @base_path.setter
     def base_path(self, path: str):
         # ensure the base_path ends with a single '/'
-        self._base_path = (path.rstrip("/") + "/") if path else ""
+        self._base_path = (f"/{path.strip('/')}/") if path else ""
 
     def auth(self):
         verify_ca_bundle = self.verify
@@ -204,6 +222,7 @@ class VaultClientBase:
             for sub_path in self._browse_recursive_secrets(key_url, render=render):
                 yield sub_path
 
+    @caching
     def get_all_secrets(self, *paths: str, render: bool = True) -> types.JSONDict:
         """
         Takes several paths, return the nested dict of all secrets below
@@ -219,6 +238,7 @@ class VaultClientBase:
 
         return result
 
+    @caching
     def get_secrets(self, path: str, render: bool = True) -> types.JSONDict:
         """
         Takes a single path an return a path dict with all the secrets
@@ -237,17 +257,33 @@ class VaultClientBase:
     def list_secrets(self, path: str) -> Iterable[str]:
         return self._list_secrets(path=self._build_full_path(path))
 
+    @caching
     def get_secret(self, path: str, render: bool = True) -> types.JSONValue:
-        data = self._get_secret(path=self._build_full_path(path))
-        if len(data) == 1 and "value" in data:
-            # secrets set using vault-cli are in a key named "value".
-            # But some secrets (rabbitmq engine, secrets set from other clients) don't
-            # follow this rule.
-            secret = data["value"]
-        else:
-            secret = data
-        if render and self.render:
-            secret = self._render_template_value(secret)
+        full_path = self._build_full_path(path)
+        if full_path in self._currently_fetching:
+            return f'<recursive value "{path}">'
+
+        self._currently_fetching.add(full_path)
+        try:
+            assert self.cache is not None
+            try:
+                data = self.cache[full_path]
+            except KeyError:
+                data = self.cache[full_path] = self._get_secret(path=full_path)
+
+            if len(data) == 1 and "value" in data:
+                # secrets set using vault-cli are in a key named "value".
+                # But some secrets (rabbitmq engine, secrets set from other clients) don't
+                # follow this rule.
+                secret = data["value"]
+            else:
+                secret = data
+            if render and self.render:
+                secret = self._render_template_value(secret)
+
+        finally:
+            self._currently_fetching.remove(full_path)
+
         return secret
 
     def delete_secret(self, path: str) -> None:
@@ -269,6 +305,7 @@ class VaultClientBase:
             return iterator
         return list(iterator)
 
+    @caching
     def move_secrets_iter(
         self, source: str, dest: str, force: Optional[bool] = None
     ) -> Iterable[Tuple[str, str]]:
@@ -305,8 +342,9 @@ class VaultClientBase:
         if not secret.startswith(self.template_prefix):
             return secret
 
-        return self.render_template(secret[len(self.template_prefix) :], render=False)
+        return self.render_template(secret[len(self.template_prefix) :])
 
+    @caching
     def render_template(self, template: str, render: bool = True) -> str:
         def vault(path):
             try:
@@ -316,6 +354,7 @@ class VaultClientBase:
 
         return jinja2.Template(template).render(vault=vault)
 
+    @caching
     def set_secret(
         self, path: str, value: types.JSONValue, force: Optional[bool] = None
     ) -> None:
@@ -363,6 +402,14 @@ class VaultClientBase:
 
         self._set_secret(path=self._build_full_path(path), secret={"value": value})
 
+    @contextlib.contextmanager
+    def caching(self):
+        old_cache, self.cache = self.cache, {}
+        try:
+            yield
+        finally:
+            self.cache = old_cache
+
     def _init_client(
         self,
         url: str,


Test patch:
diff --git a/tests/unit/test_client_base.py b/tests/unit/test_client_base.py
index 48cc376..e345cb9 100644
--- a/tests/unit/test_client_base.py
+++ b/tests/unit/test_client_base.py
@@ -1,6 +1,8 @@
+import itertools
+
 import pytest
 
-from vault_cli import client, exceptions
+from vault_cli import client, exceptions, testing
 
 
 def test_get_client(mocker):
@@ -352,15 +354,26 @@ def test_vault_client_base_render_template_path_not_found(vault):
         ({"a": {"value": "!template!b"}, "b": {"value": "c"}}, "b"),
         # Secret is a template
         ({"a": {"value": "!template!{{ vault('b') }}"}, "b": {"value": "c"}}, "c"),
-        # No recursion
+        # Finite recursion
         (
             {
                 "a": {"value": "!template!{{ vault('b') }}"},
                 "b": {"value": "!template!{{ vault('c') }}"},
                 "c": {"value": "d"},
             },
-            "!template!{{ vault('c') }}",
+            "d",
         ),
+        # Infinite Recursion
+        (
+            {
+                "a": {"value": "!template!{{ vault('b') }}"},
+                "b": {"value": "!template!{{ vault('c') }}"},
+                "c": {"value": "!template!{{ vault('a') }}"},
+            },
+            '<recursive value "a">',
+        ),
+        # Direct Recursion
+        ({"a": {"value": "!template!{{ vault('a') }}"}}, '<recursive value "a">'),
     ],
 )
 def test_vault_client_base_get_secret(vault, vault_contents, expected):
@@ -369,6 +382,15 @@ def test_vault_client_base_get_secret(vault, vault_contents, expected):
     assert vault.get_secret("a") == expected
 
 
+def test_vault_client_base_get_secret_template_root(vault):
+    vault.base_path = "base"
+    vault.db = {"/base/a": {"value": '!template!{{vault("a")}} yay'}}
+
+    # In case of erroneous caching, e.g. a different cache entry
+    # for /base/a and base/a, we would find '<recursive value "a"> yay yay'
+    assert vault.get_secret("/base/a") == '<recursive value "a"> yay'
+
+
 def test_vault_client_base_get_secret_no_value(vault):
     # the secret has no "value" key. This can not happen when the variable is
     # set with vault-cli but can if set from another client or if we are using
@@ -412,16 +434,16 @@ def test_vault_client_base_get_secrets_error(vault):
 @pytest.mark.parametrize(
     "method, params, expected",
     [
-        ("get_secret", ["foo"], {"path": "base/foo"}),
+        ("get_secret", ["foo"], {"path": "/base/foo"}),
         ("get_secret", ["/foo"], {"path": "/foo"}),
-        ("delete_secret", ["foo"], {"path": "base/foo"}),
+        ("delete_secret", ["foo"], {"path": "/base/foo"}),
         ("delete_secret", ["/foo"], {"path": "/foo"}),
-        ("list_secrets", ["foo"], {"path": "base/foo"}),
+        ("list_secrets", ["foo"], {"path": "/base/foo"}),
         ("list_secrets", ["/foo"], {"path": "/foo"}),
         (
             "set_secret",
             ["foo", "value"],
-            {"path": "base/foo", "secret": {"value": "value"}},
+            {"path": "/base/foo", "secret": {"value": "value"}},
         ),
         (
             "set_secret",
@@ -438,15 +460,84 @@ def test_vault_client_base_absolute_path(vault, mocker, method, params, expected
     mocked.assert_called_with(**expected)
 
 
-@pytest.mark.parametrize("path, expected", [("foo", "base/foo"), ("/foo", "/foo")])
+@pytest.mark.parametrize("path, expected", [("foo", "/base/foo"), ("/foo", "/foo")])
 def test_vault_client_base_build_full_path(vault, path, expected):
     vault.base_path = "base/"
     assert vault._build_full_path(path) == expected
 
 
 @pytest.mark.parametrize(
-    "path, expected", [("foo", "foo/"), ("foo/", "foo/"), ("foo//", "foo/")]
+    "path, expected",
+    [
+        ("foo", "/foo/"),
+        ("foo/", "/foo/"),
+        ("foo//", "/foo/"),
+        ("/foo", "/foo/"),
+        ("/foo/", "/foo/"),
+        ("/foo//", "/foo/"),
+    ],
 )
 def test_vault_client_base_base_path(vault, path, expected):
     vault.base_path = path
     assert vault.base_path == expected
+
+
+def test_vault_client_base_get_secret_implicit_cache_ends(vault):
+    vault.db = {"a": {"value": "b"}}
+    assert vault.get_secret("a") == "b"
+    vault.db = {"a": {"value": "c"}}
+    # Value updated. Cache was just for the duration of the call
+    assert vault.get_secret("a") == "c"
+
+
+class RaceConditionTestVaultClient(testing.TestVaultClient):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.counter = itertools.count()
+
+    def _get_secret(self, path):
+        if path == "a":
+            val = next(self.counter)
+            return {"b": f"b{val}", "c": f"c{val}"}
+        return super()._get_secret(path)
+
+
+def test_vault_client_base_get_secret_implicit_cache_no_race_condition():
+    # In this test we check that if a value is read several times by
+    # a template, implicit caching makes sure we have the same value
+    # every time.
+
+    # Values returned by this client keep changing
+
+    vault = RaceConditionTestVaultClient()
+
+    assert vault.get_secret("a") == {"b": "b0", "c": "c0"}
+    assert vault.get_secret("a") == {"b": "b1", "c": "c1"}
+
+    vault.db = {"d": {"value": """!template!{{ vault("a").b }}-{{ vault("a").c }}"""}}
+
+    # b2-c3 would be the value if caching didn't work.
+    assert vault.get_secret("d") == "b2-c2"
+
+
+def test_vault_client_base_get_secrets_implicit_cache_no_race_condition():
+    # In this test, the same value is read twice by get-all and template
+    # We check that 2 values are consistent
+
+    vault = RaceConditionTestVaultClient()
+
+    vault.db = {
+        "a": {},
+        "d": {"value": """!template!{{ vault("a").b }}-{{ vault("a").c }}"""},
+    }
+
+    assert vault.get_secrets("") == {"a": {"b": "b0", "c": "c0"}, "d": "b0-c0"}
+
+
+def test_vault_client_base_get_secret_explicit_cache(vault):
+    vault.db = {"a": {"value": "b"}}
+    with vault.caching():
+        assert vault.get_secret("a") == "b"
+        vault.db = {"a": {"value": "c"}}
+        # Value not updated
+        assert vault.get_secret("a") == "b"