1 files changed, 207 insertions, 0 deletions
diff --git a/vespaclient-java/src/main/java/com/yahoo/vespaget/DocumentRetriever.java b/vespaclient-java/src/main/java/com/yahoo/vespaget/DocumentRetriever.java
new file mode 100644
index 00000000000..6e52e89c580
--- /dev/null
+++ b/vespaclient-java/src/main/java/com/yahoo/vespaget/DocumentRetriever.java
@@ -0,0 +1,207 @@
+// Copyright 2017 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespaget;
+
+import com.yahoo.document.Document;
+import com.yahoo.document.DocumentId;
+import com.yahoo.document.json.JsonWriter;
+import com.yahoo.documentapi.SyncParameters;
+import com.yahoo.documentapi.messagebus.MessageBusDocumentAccess;
+import com.yahoo.documentapi.messagebus.MessageBusParams;
+import com.yahoo.documentapi.messagebus.MessageBusSyncSession;
+import com.yahoo.documentapi.messagebus.loadtypes.LoadType;
+import com.yahoo.documentapi.messagebus.loadtypes.LoadTypeSet;
+import com.yahoo.documentapi.messagebus.protocol.GetDocumentMessage;
+import com.yahoo.documentapi.messagebus.protocol.GetDocumentReply;
+import com.yahoo.messagebus.Message;
+import com.yahoo.messagebus.Reply;
+import com.yahoo.messagebus.Trace;
+import com.yahoo.text.Utf8;
+import com.yahoo.vespaclient.ClusterDef;
+import com.yahoo.vespaclient.ClusterList;
+
+import java.util.Iterator;
+import java.util.Map;
+
+/**
+ * The document retriever is responsible for retrieving documents using the Document API and printing the result to standard out.
+ *
+ * @author bjorncs
+ */
+public class DocumentRetriever {
+
+    private final ClusterList clusterList;
+    private final DocumentAccessFactory documentAccessFactory;
+    private final ClientParameters params;
+    private final LoadTypeSet loadTypeSet;
+
+    private MessageBusSyncSession session;
+    private MessageBusDocumentAccess documentAccess;
+
+    public DocumentRetriever(ClusterList clusterList,
+                             DocumentAccessFactory documentAccessFactory,
+                             LoadTypeSet loadTypeSet,
+                             ClientParameters params) {
+        this.clusterList = clusterList;
+        this.documentAccessFactory = documentAccessFactory;
+        this.loadTypeSet = loadTypeSet;
+        this.params = params;
+    }
+
+    public void shutdown() {
+        try {
+            if (session != null) {
+                session.destroy();
+            }
+        } catch (IllegalStateException e) {
+            // Ignore exception on shutdown
+        }
+        try {
+            if (documentAccess != null) {
+                documentAccess.shutdown();
+            }
+        } catch (IllegalStateException e) {
+            // Ignore exception on shutdown
+        }
+    }
+
+    public void retrieveDocuments() throws DocumentRetrieverException {
+        boolean first = true;
+        String route = params.cluster.isEmpty() ? params.route : resolveClusterRoute(params.cluster);
+        LoadType loadType = params.loadTypeName.isEmpty() ? null : resolveLoadType(params.loadTypeName);
+
+        MessageBusParams messageBusParams = createMessageBusParams(params.configId, params.timeout, route);
+        documentAccess = documentAccessFactory.createDocumentAccess(messageBusParams);
+        session = documentAccess.createSyncSession(new SyncParameters());
+        int trace = params.traceLevel;
+        if (trace > 0) {
+            session.setTraceLevel(trace);
+        }
+
+        Iterator<String> iter = params.documentIds;
+        if (params.jsonOutput && !params.printIdsOnly) {
+            System.out.println('[');
+        }
+        while (iter.hasNext()) {
+            if (params.jsonOutput && !params.printIdsOnly) {
+                if (!first) {
+                    System.out.println(',');
+                } else {
+                    first = false;
+                }
+            }
+            String docid = iter.next();
+            Message msg = createDocumentRequest(docid, loadType);
+            Reply reply = session.syncSend(msg);
+            printReply(reply);
+        }
+        if (params.jsonOutput && !params.printIdsOnly) {
+            System.out.println(']');
+        }
+    }
+
+    private String resolveClusterRoute(String clusterName) throws DocumentRetrieverException {
+        if (clusterList.getStorageClusters().isEmpty()) {
+            throw new DocumentRetrieverException("The Vespa cluster does not have any content clusters declared.");
+        }
+
+        ClusterDef clusterDef = null;
+        for (ClusterDef c : clusterList.getStorageClusters()) {
+            if (c.getName().equals(clusterName)) {
+                clusterDef = c;
+            }
+        }
+        if (clusterDef == null) {
+            String names = createClusterNamesString();
+            throw new DocumentRetrieverException(String.format(
+                    "The Vespa cluster contains the content clusters %s, not %s. Please select a valid vespa cluster.",
+                    names, clusterName));
+        }
+        return String.format("[Storage:cluster=%s;clusterconfigid=%s]", clusterDef.getName(), clusterDef.getConfigId());
+    }
+
+    private LoadType resolveLoadType(String loadTypeName) throws DocumentRetrieverException {
+        Map<String, LoadType> loadTypesNameMap = loadTypeSet.getNameMap();
+        if (!loadTypesNameMap.containsKey(loadTypeName)) {
+            throw new DocumentRetrieverException(String.format("Loadtype with name '%s' does not exist.\n", loadTypeName));
+        } else {
+            return loadTypesNameMap.get(loadTypeName);
+        }
+    }
+
+    private MessageBusParams createMessageBusParams(String configId, double timeout, String route) {
+        MessageBusParams messageBusParams = new MessageBusParams(loadTypeSet);
+        messageBusParams.setRoute(route);
+        messageBusParams.setProtocolConfigId(configId);
+        messageBusParams.setRoutingConfigId(configId);
+        messageBusParams.setDocumentManagerConfigId(configId);
+
+        if (timeout > 0) {
+            messageBusParams.getSourceSessionParams().setTimeout(timeout);
+        }
+        return messageBusParams;
+    }
+
+    private Message createDocumentRequest(String docid, LoadType loadType) {
+        GetDocumentMessage msg = new GetDocumentMessage(new DocumentId(docid), params.fieldSet);
+        msg.setPriority(params.priority);
+        msg.setRetryEnabled(!params.noRetry);
+
+        if (loadType != null) {
+            msg.setLoadType(loadType);
+        }
+        return msg;
+    }
+
+    private void printReply(Reply reply) {
+        Trace trace = reply.getTrace();
+        if (!trace.getRoot().isEmpty()) {
+            System.out.println(trace);
+        }
+
+        if (reply.hasErrors()) {
+            System.err.print("Request failed: ");
+            for (int i = 0; i < reply.getNumErrors(); i++) {
+                System.err.printf("\n  %s", reply.getError(i));
+            }
+            System.err.println();
+            return;
+        }
+
+        if (!(reply instanceof GetDocumentReply)) {
+            System.err.printf("Unexpected reply %s: '%s'\n", reply.getType(), reply.toString());
+            return;
+        }
+
+        GetDocumentReply documentReply = (GetDocumentReply) reply;
+        Document document = documentReply.getDocument();
+
+        if (document == null) {
+            System.out.println("Document not found.");
+            return;
+        }
+
+        if (params.showDocSize) {
+            System.out.printf("Document size: %d bytes.\n", document.getSerializedSize());
+        }
+        if (params.printIdsOnly) {
+            System.out.println(document.getId());
+        } else {
+            if (params.jsonOutput) {
+                System.out.print(Utf8.toString(JsonWriter.toByteArray(document)));
+            } else {
+                System.out.print(document.toXML("  "));
+            }
+        }
+    }
+
+    private String createClusterNamesString() {
+        StringBuilder names = new StringBuilder();
+        for (ClusterDef c : clusterList.getStorageClusters()) {
+            if (names.length() > 0) {
+                names.append(", ");
+            }
+            names.append(c.getName());
+        }
+        return names.toString();
+    }
+}