From 604846b05e391ea3350e11508ce793a01a5840db Mon Sep 17 00:00:00 2001
From: Mark Paluch <mpaluch@pivotal.io>
Date: Mon, 30 Sep 2019 13:42:03 +0200
Subject: [PATCH] =?UTF-8?q?Consistently=20use=20doWithSession(=E2=80=A6)?=
 =?UTF-8?q?=20from=20within=20VaultTemplate.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Closes gh-481.
---
 .../vault/core/ReactiveVaultTemplate.java     | 35 +++++++------
 .../vault/core/VaultTemplate.java             | 52 ++++++++++---------
 2 files changed, 48 insertions(+), 39 deletions(-)

diff --git a/spring-vault-core/src/main/java/org/springframework/vault/core/ReactiveVaultTemplate.java b/spring-vault-core/src/main/java/org/springframework/vault/core/ReactiveVaultTemplate.java
index d7c1b464..d6d53fab 100644
--- a/spring-vault-core/src/main/java/org/springframework/vault/core/ReactiveVaultTemplate.java
+++ b/spring-vault-core/src/main/java/org/springframework/vault/core/ReactiveVaultTemplate.java
@@ -162,10 +162,13 @@ public class ReactiveVaultTemplate implements ReactiveVaultOperations {
 	@Override
 	public <T> Mono<VaultResponseSupport<T>> read(String path, Class<T> responseType) {
 
-		ParameterizedTypeReference<VaultResponseSupport<T>> ref = VaultResponses
-				.getTypeReference(responseType);
+		return doWithSession(webClient -> {
 
-		return sessionClient.get().uri(path).exchange().flatMap(mapResponse(ref, path));
+			ParameterizedTypeReference<VaultResponseSupport<T>> ref = VaultResponses
+					.getTypeReference(responseType);
+
+			return webClient.get().uri(path).exchange().flatMap(mapResponse(ref, path));
+		});
 	}
 
 	@Override
@@ -181,7 +184,6 @@ public class ReactiveVaultTemplate implements ReactiveVaultOperations {
 		return read
 				.filter(response -> response.getData() != null
 						&& response.getData().containsKey("keys"))
-				//
 				.flatMapIterable(
 						response -> (List<String>) response.getRequiredData().get("keys"));
 	}
@@ -191,15 +193,18 @@ public class ReactiveVaultTemplate implements ReactiveVaultOperations {
 
 		Assert.hasText(path, "Path must not be empty");
 
-		RequestBodySpec uri = sessionClient.post().uri(path);
-		Mono<ClientResponse> exchange;
-		if (body != null) {
-			exchange = uri.syncBody(body).exchange();
-		}
-		else {
-			exchange = uri.exchange();
-		}
-		return exchange.flatMap(mapResponse(VaultResponse.class, path));
+		return doWithSession(webClient -> {
+
+			RequestBodySpec uri = webClient.post().uri(path);
+			Mono<ClientResponse> exchange;
+			if (body != null) {
+				exchange = uri.syncBody(body).exchange();
+			}
+			else {
+				exchange = uri.exchange();
+			}
+			return exchange.flatMap(mapResponse(VaultResponse.class, path));
+		});
 	}
 
 	@Override
@@ -207,8 +212,8 @@ public class ReactiveVaultTemplate implements ReactiveVaultOperations {
 
 		Assert.hasText(path, "Path must not be empty");
 
-		return sessionClient.delete().uri(path).exchange()
-				.flatMap(mapResponse(String.class, path)).then();
+		return doWithSession(webClient -> webClient.delete().uri(path).exchange()
+				.flatMap(mapResponse(String.class, path)).then());
 	}
 
 	@Override
diff --git a/spring-vault-core/src/main/java/org/springframework/vault/core/VaultTemplate.java b/spring-vault-core/src/main/java/org/springframework/vault/core/VaultTemplate.java
index 9672cdfb..7fea8a5b 100644
--- a/spring-vault-core/src/main/java/org/springframework/vault/core/VaultTemplate.java
+++ b/spring-vault-core/src/main/java/org/springframework/vault/core/VaultTemplate.java
@@ -276,20 +276,23 @@ public class VaultTemplate implements InitializingBean, VaultOperations, Disposa
 		ParameterizedTypeReference<VaultResponseSupport<T>> ref = VaultResponses
 				.getTypeReference(responseType);
 
-		try {
-			ResponseEntity<VaultResponseSupport<T>> exchange = sessionTemplate.exchange(
+		return doWithSession(restOperations -> {
+
+			try {
+				ResponseEntity<VaultResponseSupport<T>> exchange = restOperations.exchange(
 					path, HttpMethod.GET, null, ref);
 
-			return exchange.getBody();
-		}
-		catch (HttpStatusCodeException e) {
-
-			if (e.getStatusCode() == HttpStatus.NOT_FOUND) {
-				return null;
+				return exchange.getBody();
 			}
+			catch (HttpStatusCodeException e) {
 
-			throw VaultResponses.buildException(e, path);
-		}
+				if (e.getStatusCode() == HttpStatus.NOT_FOUND) {
+					return null;
+				}
+
+				throw VaultResponses.buildException(e, path);
+			}
+		});
 	}
 
 	@Override
@@ -315,12 +318,8 @@ public class VaultTemplate implements InitializingBean, VaultOperations, Disposa
 
 		Assert.hasText(path, "Path must not be empty");
 
-		try {
-			return sessionTemplate.postForObject(path, body, VaultResponse.class);
-		}
-		catch (HttpStatusCodeException e) {
-			throw VaultResponses.buildException(e, path);
-		}
+		return doWithSession(restOperations -> restOperations.postForObject(path, body,
+				VaultResponse.class));
 	}
 
 	@Override
@@ -328,17 +327,22 @@ public class VaultTemplate implements InitializingBean, VaultOperations, Disposa
 
 		Assert.hasText(path, "Path must not be empty");
 
-		try {
-			sessionTemplate.delete(path);
-		}
-		catch (HttpStatusCodeException e) {
+		doWithSession(restOperations -> {
 
-			if (e.getStatusCode() == HttpStatus.NOT_FOUND) {
-				return;
+			try {
+				restOperations.delete(path);
+			}
+			catch (HttpStatusCodeException e) {
+
+				if (e.getStatusCode() == HttpStatus.NOT_FOUND) {
+					return null;
+				}
+
+				throw VaultResponses.buildException(e, path);
 			}
 
-			throw VaultResponses.buildException(e, path);
-		}
+			return null;
+		});
 	}
 
 	@Override