WIP for adding bucket space to stat-related messages in Java

10 files changed, 126 insertions, 7 deletions

diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/DocumentProtocol.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/DocumentProtocol.java
index 7cebdbb9abf..d8ee2a39ecd 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/DocumentProtocol.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/DocumentProtocol.java
@@ -357,8 +357,9 @@ public class DocumentProtocol implements Protocol {
         putRoutableFactory(MESSAGE_REMOVEDOCUMENT, new RoutableFactories52.RemoveDocumentMessageFactory(), from52);
 
         // 6.x serialization
-        // TODO stat, listbuckets
         putRoutableFactory(MESSAGE_CREATEVISITOR, new RoutableFactories60.CreateVisitorMessageFactory(), from6);
+        putRoutableFactory(MESSAGE_STATBUCKET, new RoutableFactories60.StatBucketMessageFactory(), from6);
+        putRoutableFactory(MESSAGE_GETBUCKETLIST, new RoutableFactories60.GetBucketListMessageFactory(), from6);
     }
 
     /**
diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/GetBucketListMessage.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/GetBucketListMessage.java
index 74629018e55..5e5d79f86b7 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/GetBucketListMessage.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/GetBucketListMessage.java
@@ -6,13 +6,19 @@ import com.yahoo.document.BucketId;
 public class GetBucketListMessage extends DocumentMessage {
 
     private BucketId bucketId;
+    private String bucketSpace = "default";
 
     GetBucketListMessage() {
         // must be deserialized into
     }
 
     public GetBucketListMessage(BucketId bucketId) {
+        this(bucketId, "default");
+    }
+
+    public GetBucketListMessage(BucketId bucketId, String bucketSpace) {
         this.bucketId = bucketId;
+        this.bucketSpace = bucketSpace;
     }
 
     public BucketId getBucketId() {
@@ -23,6 +29,14 @@ public class GetBucketListMessage extends DocumentMessage {
         bucketId = id;
     }
 
+    public String getBucketSpace() {
+        return bucketSpace;
+    }
+
+    public void setBucketSpace(String bucketSpace) {
+        this.bucketSpace = bucketSpace;
+    }
+
     @Override
     public DocumentReply createReply() {
         return new StatBucketReply();
@@ -30,7 +44,7 @@ public class GetBucketListMessage extends DocumentMessage {
 
     @Override
     public int getApproxSize() {
-        return super.getApproxSize() + 8;
+        return super.getApproxSize() + 8 + bucketSpace.length();
     }
 
     @Override
diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories50.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories50.java
index 7417e63480e..0ea82ba0d83 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories50.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories50.java
@@ -392,18 +392,28 @@ public abstract class RoutableFactories50 {
 
     public static class GetBucketListMessageFactory extends DocumentMessageFactory {
 
+        protected String decodeBucketSpace(Deserializer deserializer) {
+            return "default"; // TODO fixed space repo
+        }
+
         @Override
         protected DocumentMessage doDecode(DocumentDeserializer buf) {
             GetBucketListMessage msg = new GetBucketListMessage();
             msg.setBucketId(new BucketId(buf.getLong(null)));
+            msg.setBucketSpace(decodeBucketSpace(buf));
             return msg;
         }
 
+        protected boolean encodeBucketSpace(String bucketSpace, DocumentSerializer buf) {
+            // TODO fixed space repo
+            return "default".equals(bucketSpace);
+        }
+
         @Override
         protected boolean doEncode(DocumentMessage obj, DocumentSerializer buf) {
             GetBucketListMessage msg = (GetBucketListMessage)obj;
             buf.putLong(null, msg.getBucketId().getRawId());
-            return true;
+            return encodeBucketSpace(msg.getBucketSpace(), buf);
         }
     }
 
@@ -831,20 +841,30 @@ public abstract class RoutableFactories50 {
 
     public static class StatBucketMessageFactory extends DocumentMessageFactory {
 
+        protected String decodeBucketSpace(Deserializer deserializer) {
+            return "default"; // TODO fixed space repo
+        }
+
         @Override
         protected DocumentMessage doDecode(DocumentDeserializer buf) {
             StatBucketMessage msg = new StatBucketMessage();
             msg.setBucketId(new BucketId(buf.getLong(null)));
             msg.setDocumentSelection(decodeString(buf));
+            msg.setBucketSpace(decodeBucketSpace(buf));
             return msg;
         }
 
+        protected boolean encodeBucketSpace(String bucketSpace, DocumentSerializer buf) {
+            // TODO fixed space repo
+            return "default".equals(bucketSpace);
+        }
+
         @Override
         protected boolean doEncode(DocumentMessage obj, DocumentSerializer buf) {
             StatBucketMessage msg = (StatBucketMessage)obj;
             buf.putLong(null, msg.getBucketId().getRawId());
             encodeString(msg.getDocumentSelection(), buf);
-            return true;
+            return encodeBucketSpace(msg.getBucketSpace(), buf);
         }
     }
 
diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories60.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories60.java
index ef63fe8b295..723760df117 100644
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories60.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories60.java
@@ -19,4 +19,30 @@ public class RoutableFactories60 extends RoutableFactories52 {
         }
     }
 
+    public static class StatBucketMessageFactory extends RoutableFactories52.StatBucketMessageFactory {
+        @Override
+        protected String decodeBucketSpace(Deserializer deserializer) {
+            return decodeString(deserializer);
+        }
+
+        @Override
+        protected boolean encodeBucketSpace(String bucketSpace, DocumentSerializer buf) {
+            encodeString(bucketSpace, buf);
+            return true;
+        }
+    }
+
+    public static class GetBucketListMessageFactory extends RoutableFactories52.GetBucketListMessageFactory {
+        @Override
+        protected String decodeBucketSpace(Deserializer deserializer) {
+            return decodeString(deserializer);
+        }
+
+        @Override
+        protected boolean encodeBucketSpace(String bucketSpace, DocumentSerializer buf) {
+            encodeString(bucketSpace, buf);
+            return true;
+        }
+    }
+
 }
diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/StatBucketMessage.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/StatBucketMessage.java
index 9b2176d85f6..8fb4ceb6d41 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/StatBucketMessage.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/StatBucketMessage.java
@@ -6,6 +6,7 @@ import com.yahoo.document.BucketId;
 public class StatBucketMessage extends DocumentMessage {
 
     private BucketId bucketId;
+    private String bucketSpace = "default";
     private String documentSelection;
 
     StatBucketMessage() {
@@ -13,7 +14,12 @@ public class StatBucketMessage extends DocumentMessage {
     }
 
     public StatBucketMessage(BucketId bucket, String documentSelection) {
-        this.bucketId = bucket;
+        this(bucket, "default", documentSelection);
+    }
+
+    public StatBucketMessage(BucketId bucketId, String bucketSpace, String documentSelection) {
+        this.bucketId = bucketId;
+        this.bucketSpace = bucketSpace;
         this.documentSelection = documentSelection;
     }
 
@@ -33,6 +39,14 @@ public class StatBucketMessage extends DocumentMessage {
         this.documentSelection = documentSelection;
     }
 
+    public String getBucketSpace() {
+        return bucketSpace;
+    }
+
+    public void setBucketSpace(String bucketSpace) {
+        this.bucketSpace = bucketSpace;
+    }
+
     @Override
     public DocumentReply createReply() {
         return new StatBucketReply();
@@ -40,7 +54,7 @@ public class StatBucketMessage extends DocumentMessage {
 
     @Override
     public int getApproxSize() {
-        return super.getApproxSize() + 8 + documentSelection.length();
+        return super.getApproxSize() + 8 + bucketSpace.length() + documentSelection.length();
     }
 
     @Override
diff --git a/vespaclient-container-plugin/pom.xml b/vespaclient-container-plugin/pom.xml
index 3472f3fcdc5..1de6930307e 100644
--- a/vespaclient-container-plugin/pom.xml
+++ b/vespaclient-container-plugin/pom.xml
@@ -37,6 +37,11 @@
       <version>${project.version}</version>
       <scope>provided</scope>
     </dependency>
+    <dependency>
+      <groupId>com.yahoo.vespa</groupId>
+      <artifactId>configdefinitions</artifactId>
+      <version>${project.version}</version>
+    </dependency>
 
     <dependency>
       <groupId>org.apache.httpcomponents</groupId>
diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/BucketSpaceEnumerator.java b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/BucketSpaceEnumerator.java
index 4f3768e3740..24692859266 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/BucketSpaceEnumerator.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/BucketSpaceEnumerator.java
@@ -1,3 +1,4 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.document.restapi;
 
 import com.yahoo.config.subscription.ConfigGetter;
diff --git a/vespaclient-java/src/main/java/com/yahoo/vespastat/ClientParameters.java b/vespaclient-java/src/main/java/com/yahoo/vespastat/ClientParameters.java
index 47a0418e2d8..a05e4ed1e48 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespastat/ClientParameters.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespastat/ClientParameters.java
@@ -17,6 +17,7 @@ public class ClientParameters {
     public final SelectionType selectionType;
     // The selection id
     public final String id;
+    public final String bucketSpace;
 
     public ClientParameters(
             boolean help,
@@ -24,11 +25,22 @@ public class ClientParameters {
             String route,
             SelectionType selectionType,
             String id) {
+        this(help, dumpData, route, selectionType, id, "default");
+    }
+
+    public ClientParameters(
+            boolean help,
+            boolean dumpData,
+            String route,
+            SelectionType selectionType,
+            String id,
+            String bucketSpace) {
         this.help = help;
         this.dumpData = dumpData;
         this.route = route;
         this.selectionType = selectionType;
         this.id = id;
+        this.bucketSpace = bucketSpace;
     }
 
     public enum SelectionType {USER, GROUP, BUCKET, GID, DOCUMENT}
@@ -39,6 +51,7 @@ public class ClientParameters {
         private String route;
         private SelectionType selectionType;
         private String id;
+        private String bucketSpace = "default"; // TODO repo etc etc
 
         public Builder setHelp(boolean help) {
             this.help = help;
@@ -65,8 +78,13 @@ public class ClientParameters {
             return this;
         }
 
+        public Builder setBucketSpace(String bucketSpace) {
+            this.bucketSpace = bucketSpace;
+            return this;
+        }
+
         public ClientParameters build() {
-            return new ClientParameters(help, dumpData, route, selectionType, id);
+            return new ClientParameters(help, dumpData, route, selectionType, id, bucketSpace);
         }
     }
 
diff --git a/vespaclient-java/src/main/java/com/yahoo/vespastat/CommandLineOptions.java b/vespaclient-java/src/main/java/com/yahoo/vespastat/CommandLineOptions.java
index 59ab17af8d6..b6783efb25b 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespastat/CommandLineOptions.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespastat/CommandLineOptions.java
@@ -25,6 +25,7 @@ public class CommandLineOptions {
     private static final String BUCKET_OPTION = "bucket";
     private static final String GID_OPTION = "gid";
     private static final String DOCUMENT_OPTION = "document";
+    private static final String BUCKET_SPACE_OPTION = "bucketspace";
 
     private final Options options = createOptions();
 
@@ -51,6 +52,13 @@ public class CommandLineOptions {
                 .longOpt(ROUTE_OPTION)
                 .build());
 
+        options.addOption(Option.builder("s")
+                .hasArg(true)
+                .desc("Stat buckets within the given bucket space. If not provided, 'default' is used.")
+                .argName("space")
+                .longOpt(BUCKET_SPACE_OPTION)
+                .build());
+
         // A group of mutually exclusive options for user, group, bucket, gid and document.
         OptionGroup optionGroup = new OptionGroup();
         optionGroup.setRequired(false);
@@ -109,6 +117,7 @@ public class CommandLineOptions {
             builder.setHelp(cl.hasOption(HELP_OPTION));
             builder.setDumpData(cl.hasOption(DUMP_OPTION));
             builder.setRoute(cl.getOptionValue(ROUTE_OPTION, "default"));
+            builder.setBucketSpace(cl.getOptionValue(BUCKET_SPACE_OPTION, "default"));
 
             if (cl.hasOption(USER_OPTION)) {
                 builder.setSelectionType(ClientParameters.SelectionType.USER);
diff --git a/vespaclient-java/src/test/java/com/yahoo/vespastat/CommandLineOptionsTest.java b/vespaclient-java/src/test/java/com/yahoo/vespastat/CommandLineOptionsTest.java
index 31cef2a6a86..a6bd7280c36 100644
--- a/vespaclient-java/src/test/java/com/yahoo/vespastat/CommandLineOptionsTest.java
+++ b/vespaclient-java/src/test/java/com/yahoo/vespastat/CommandLineOptionsTest.java
@@ -70,6 +70,17 @@ public class CommandLineOptionsTest {
         }
     }
 
+    // TODO fixed bucket space repo
+    @Test
+    public void bucket_space_is_default_unless_specified() {
+        assertEquals("default", getParsedOptions("--user", "id").bucketSpace);
+    }
+
+    @Test
+    public void can_specify_explicit_bucket_space() {
+        assertEquals("global", getParsedOptions("--user", "id", "--bucketspace", "global").bucketSpace);
+    }
+
     @Test
     public void testDefaultRoute() {
         assertEquals("default", getParsedOptions("--user", "dummyuser").route);