From a2479ef10b36055dbd66568dc2701b5f8de24ec4 Mon Sep 17 00:00:00 2001
From: Yoann Fouquet <y.fouquet@criteo.com>
Date: Thu, 23 Nov 2017 23:26:06 +0100
Subject: [PATCH] Allow specify datacenter for prepared queries

For each query, the possibility to specify a datacenter is added.
For the last query (execute), as datacenter is in QueryOptions, the behavior remains the same.
The behavior of existing methods remains unchanged.
---
 .../orbitz/consul/PreparedQueryClient.java    | 62 ++++++++++++++++---
 1 file changed, 55 insertions(+), 7 deletions(-)

diff --git a/src/main/java/com/orbitz/consul/PreparedQueryClient.java b/src/main/java/com/orbitz/consul/PreparedQueryClient.java
index 1e45693a..a5e39dc3 100644
--- a/src/main/java/com/orbitz/consul/PreparedQueryClient.java
+++ b/src/main/java/com/orbitz/consul/PreparedQueryClient.java
@@ -1,6 +1,7 @@
 package com.orbitz.consul;
 
 import com.google.common.base.Optional;
+import com.google.common.collect.ImmutableMap;
 import com.orbitz.consul.async.Callback;
 import com.orbitz.consul.model.query.PreparedQuery;
 import com.orbitz.consul.model.query.QueryId;
@@ -40,7 +41,27 @@ public class PreparedQueryClient {
      * @return The ID of the created query.
      */
     public String createPreparedQuery(PreparedQuery preparedQuery) {
-        return extract(api.createPreparedQuery(preparedQuery)).getId();
+        return createPreparedQuery(preparedQuery, null);
+    }
+
+    /**
+     * Creates a prepared query.
+     *
+     * POST /v1/query
+     *
+     * @param preparedQuery The prepared query to create.
+     * @param dc The data center.
+     * @return The ID of the created query.
+     */
+    public String createPreparedQuery(PreparedQuery preparedQuery, final String dc) {
+        return extract(api.createPreparedQuery(preparedQuery, dcQuery(dc))).getId();
+    }
+
+    private Map<String, String> dcQuery(String dc) {
+        if (dc != null) {
+            return ImmutableMap.of("dc", dc);
+        }
+        return Collections.emptyMap();
     }
     
     /**
@@ -51,7 +72,19 @@ public String createPreparedQuery(PreparedQuery preparedQuery) {
      * @return The list of prepared queries.
      */
     public List<StoredQuery> getPreparedQueries() {
-	return extract(api.getPreparedQueries());
+        return getPreparedQueries(null);
+    }
+
+    /**
+     * Retrieves the list of prepared queries.
+     *
+     * GET /v1/query
+     *
+     * @param dc The data center.
+     * @return The list of prepared queries.
+     */
+    public List<StoredQuery> getPreparedQueries(final String dc) {
+        return extract(api.getPreparedQueries(dcQuery(dc)));
     }
 
     /**
@@ -63,7 +96,20 @@ public List<StoredQuery> getPreparedQueries() {
      * @return The store prepared query.
      */
     public Optional<StoredQuery> getPreparedQuery(String id) {
-        List<StoredQuery> result = extract(api.getPreparedQuery(id));
+        return getPreparedQuery(id, null);
+    }
+
+    /**
+     * Retrieves a prepared query by its ID.
+     *
+     * GET /v1/query/{id}
+     *
+     * @param id The query ID.
+     * @param dc The data center.
+     * @return The store prepared query.
+     */
+    public Optional<StoredQuery> getPreparedQuery(String id, final String dc) {
+        List<StoredQuery> result = extract(api.getPreparedQuery(id, dcQuery(dc)));
 
         return result.isEmpty() ? Optional.<StoredQuery>absent() : Optional.of(result.get(0));
     }
@@ -94,14 +140,16 @@ public void execute(String nameOrId, QueryOptions options, final Callback<QueryR
      */
     interface Api {
 	
-	@GET("query")
-	Call<List<StoredQuery>> getPreparedQueries();
+        @GET("query")
+        Call<List<StoredQuery>> getPreparedQueries(@QueryMap Map<String, String> queryMap);
 
         @POST("query")
-        Call<QueryId> createPreparedQuery(@Body PreparedQuery preparedQuery);
+        Call<QueryId> createPreparedQuery(@Body PreparedQuery preparedQuery,
+                                          @QueryMap Map<String, String> queryMap);
 
         @GET("query/{id}")
-        Call<List<StoredQuery>> getPreparedQuery(@Path("id") String id);
+        Call<List<StoredQuery>> getPreparedQuery(@Path("id") String id,
+                                                 @QueryMap Map<String, String> queryMap);
 
         @GET("query/{nameOrId}/execute")
         Call<QueryResults> execute(@Path("nameOrId") String nameOrId,