Make vespa-stat use bucketSpace.

5 files changed, 20 insertions, 18 deletions

diff --git a/vespaclient-java/src/main/java/com/yahoo/vespastat/BucketStatsPrinter.java b/vespaclient-java/src/main/java/com/yahoo/vespastat/BucketStatsPrinter.java
index c412ed48bb5..2b271af5289 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespastat/BucketStatsPrinter.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespastat/BucketStatsPrinter.java
@@ -23,19 +23,19 @@ public class BucketStatsPrinter {
         this.out = out;
     }
 
-    public void retrieveAndPrintBucketStats(ClientParameters.SelectionType type, String id, boolean dumpData) throws BucketStatsException {
+    public void retrieveAndPrintBucketStats(ClientParameters.SelectionType type, String id, boolean dumpData, String bucketSpace) throws BucketStatsException {
         BucketId bucketId = retriever.getBucketIdForType(type, id);
         if (type == ClientParameters.SelectionType.GROUP || type == ClientParameters.SelectionType.USER) {
             out.printf("Generated 32-bit bucket id: %s\n", bucketId);
         }
 
-        List<GetBucketListReply.BucketInfo> bucketList = retriever.retrieveBucketList(bucketId);
+        List<GetBucketListReply.BucketInfo> bucketList = retriever.retrieveBucketList(bucketId, bucketSpace);
         printBucketList(bucketList);
 
         if (dumpData) {
             for (GetBucketListReply.BucketInfo bucketInfo : bucketList) {
                 BucketId bucket = bucketInfo.getBucketId();
-                String bucketStats = retriever.retrieveBucketStats(type, id, bucket);
+                String bucketStats = retriever.retrieveBucketStats(type, id, bucket, bucketSpace);
                 printBucketStats(bucket, bucketStats);
             }
         }
diff --git a/vespaclient-java/src/main/java/com/yahoo/vespastat/BucketStatsRetriever.java b/vespaclient-java/src/main/java/com/yahoo/vespastat/BucketStatsRetriever.java
index f0d79e577c9..1c855455a37 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespastat/BucketStatsRetriever.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespastat/BucketStatsRetriever.java
@@ -87,15 +87,15 @@ public class BucketStatsRetriever {
         }
     }
 
-    public String retrieveBucketStats(ClientParameters.SelectionType type, String id, BucketId bucketId) throws BucketStatsException {
+    public String retrieveBucketStats(ClientParameters.SelectionType type, String id, BucketId bucketId, String bucketSpace) throws BucketStatsException {
         String documentSelection = createDocumentSelection(type, id);
-        StatBucketMessage msg = new StatBucketMessage(bucketId, documentSelection);
+        StatBucketMessage msg = new StatBucketMessage(bucketId, bucketSpace, documentSelection);
         StatBucketReply statBucketReply = sendMessage(msg, StatBucketReply.class);
         return statBucketReply.getResults();
     }
 
-    public List<GetBucketListReply.BucketInfo> retrieveBucketList(BucketId bucketId) throws BucketStatsException {
-        GetBucketListMessage msg = new GetBucketListMessage(bucketId);
+    public List<GetBucketListReply.BucketInfo> retrieveBucketList(BucketId bucketId, String bucketSpace) throws BucketStatsException {
+        GetBucketListMessage msg = new GetBucketListMessage(bucketId, bucketSpace);
         GetBucketListReply bucketListReply = sendMessage(msg, GetBucketListReply.class);
         return bucketListReply.getBuckets();
     }
diff --git a/vespaclient-java/src/main/java/com/yahoo/vespastat/Main.java b/vespaclient-java/src/main/java/com/yahoo/vespastat/Main.java
index 4e53fdc6a97..57f3ab20186 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespastat/Main.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespastat/Main.java
@@ -24,7 +24,7 @@ public class Main {
                     params.route,
                     createShutdownHookRegistrar());
             BucketStatsPrinter printer = new BucketStatsPrinter(retriever, System.out);
-            printer.retrieveAndPrintBucketStats(params.selectionType, params.id, params.dumpData);
+            printer.retrieveAndPrintBucketStats(params.selectionType, params.id, params.dumpData, params.bucketSpace);
         } catch (IllegalArgumentException e) {
             System.err.printf("Failed to parse command line arguments: %s.\n", e.getMessage());
         } catch (BucketStatsException e) {
diff --git a/vespaclient-java/src/test/java/com/yahoo/vespastat/BucketStatsPrinterTest.java b/vespaclient-java/src/test/java/com/yahoo/vespastat/BucketStatsPrinterTest.java
index 874a2a382ea..aaeada5be25 100644
--- a/vespaclient-java/src/test/java/com/yahoo/vespastat/BucketStatsPrinterTest.java
+++ b/vespaclient-java/src/test/java/com/yahoo/vespastat/BucketStatsPrinterTest.java
@@ -22,13 +22,14 @@ public class BucketStatsPrinterTest {
 
     private BucketStatsRetriever retriever;
     private final ByteArrayOutputStream out = new ByteArrayOutputStream();
+    private final String bucketSpace = "default";
 
     @Before
     public void mockBucketStatsRetriever() throws BucketStatsException {
         retriever = mock(BucketStatsRetriever.class);
         when(retriever.getBucketIdForType(any(), any())).thenReturn(new BucketId(0x42));
-        when(retriever.retrieveBucketList(any())).thenReturn(Collections.emptyList());
-        when(retriever.retrieveBucketStats(any(), any(), any())).thenReturn("");
+        when(retriever.retrieveBucketList(any(), any())).thenReturn(Collections.emptyList());
+        when(retriever.retrieveBucketStats(any(), any(), any(), any())).thenReturn("");
     }
 
     @After
@@ -44,7 +45,7 @@ public class BucketStatsPrinterTest {
 
     private String retreiveAndPrintBucketStats(ClientParameters.SelectionType type, String id, boolean dumpData) throws BucketStatsException {
         BucketStatsPrinter printer = new BucketStatsPrinter(retriever, new PrintStream(out));
-        printer.retrieveAndPrintBucketStats(type, id, dumpData);
+        printer.retrieveAndPrintBucketStats(type, id, dumpData, bucketSpace);
         return getOutputString();
     }
 
@@ -68,7 +69,7 @@ public class BucketStatsPrinterTest {
         List<GetBucketListReply.BucketInfo> bucketList = new ArrayList<>();
         String dummyInfoString = "dummyinformation";
         bucketList.add(new GetBucketListReply.BucketInfo(new BucketId(0), dummyInfoString));
-        when(retriever.retrieveBucketList(any())).thenReturn(bucketList);
+        when(retriever.retrieveBucketList(any(), any())).thenReturn(bucketList);
 
         String output = retreiveAndPrintBucketStats(ClientParameters.SelectionType.USER, "1234", false);
         assertTrue(output.contains(dummyInfoString));
@@ -78,8 +79,8 @@ public class BucketStatsPrinterTest {
     public void testShouldPrintBucketStats() throws BucketStatsException {
         String dummyBucketStats = "dummystats";
         GetBucketListReply.BucketInfo bucketInfo = new GetBucketListReply.BucketInfo(new BucketId(0), "dummy");
-        when(retriever.retrieveBucketList(any())).thenReturn(Collections.singletonList(bucketInfo));
-        when(retriever.retrieveBucketStats(any(), any(), any())).thenReturn(dummyBucketStats);
+        when(retriever.retrieveBucketList(any(), any())).thenReturn(Collections.singletonList(bucketInfo));
+        when(retriever.retrieveBucketStats(any(), any(), any(), any())).thenReturn(dummyBucketStats);
 
         String output = retreiveAndPrintBucketStats(ClientParameters.SelectionType.USER, "1234", true);
         assertTrue(output.contains(dummyBucketStats));
diff --git a/vespaclient-java/src/test/java/com/yahoo/vespastat/BucketStatsRetrieverTest.java b/vespaclient-java/src/test/java/com/yahoo/vespastat/BucketStatsRetrieverTest.java
index 557cefd8504..a34e3a73205 100644
--- a/vespaclient-java/src/test/java/com/yahoo/vespastat/BucketStatsRetrieverTest.java
+++ b/vespaclient-java/src/test/java/com/yahoo/vespastat/BucketStatsRetrieverTest.java
@@ -31,6 +31,7 @@ public class BucketStatsRetrieverTest {
     private DocumentAccessFactory mockedFactory;
     private MessageBusDocumentAccess mockedDocumentAccess;
     private MessageBusSyncSession mockedSession;
+    private final String bucketSpace = "default";
 
 
     @Before
@@ -67,7 +68,7 @@ public class BucketStatsRetrieverTest {
         reply.getBuckets().add(new GetBucketListReply.BucketInfo(bucketId, bucketInfo));
         when(mockedSession.syncSend(any())).thenReturn(reply);
 
-        List<GetBucketListReply.BucketInfo> bucketList = createRetriever().retrieveBucketList(bucketId);
+        List<GetBucketListReply.BucketInfo> bucketList = createRetriever().retrieveBucketList(bucketId, bucketSpace);
 
         verify(mockedSession, times(1)).syncSend(any());
         assertEquals(1, bucketList.size());
@@ -83,7 +84,7 @@ public class BucketStatsRetrieverTest {
         StatBucketReply reply = new StatBucketReply();
         reply.setResults(bucketInfo);
         when(mockedSession.syncSend(any())).thenReturn(reply);
-        String result = createRetriever().retrieveBucketStats(ClientParameters.SelectionType.DOCUMENT, docId, bucketId);
+        String result = createRetriever().retrieveBucketStats(ClientParameters.SelectionType.DOCUMENT, docId, bucketId, bucketSpace);
 
         verify(mockedSession, times(1)).syncSend(any());
         assertEquals(bucketInfo, result);
@@ -112,7 +113,7 @@ public class BucketStatsRetrieverTest {
         reply.addError(new Error(0, "errormsg"));
         when(mockedSession.syncSend(any())).thenReturn(reply);
 
-        createRetriever().retrieveBucketList(new BucketId(1));
+        createRetriever().retrieveBucketList(new BucketId(1), bucketSpace);
     }
 
     @Test
@@ -124,7 +125,7 @@ public class BucketStatsRetrieverTest {
         when(mockedSession.syncSend(any())).thenReturn(reply);
 
         BucketStatsRetriever retriever = new BucketStatsRetriever(mockedFactory, route, t -> {});
-        retriever.retrieveBucketList(new BucketId(0));
+        retriever.retrieveBucketList(new BucketId(0), bucketSpace);
 
         verify(mockedSession).syncSend(argThat(new ArgumentMatcher<Message>() {
             @Override