com.blackbirdai.client.model.UploadRequest Maven / Gradle / Ivy

Go to download
package com.blackbirdai.client.model;

import com.blackbirdai.client.InvalidRequestException;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.JsonNodeFactory;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.common.base.Charsets;
import com.google.common.collect.Lists;
import org.jboss.netty.buffer.ChannelBuffer;
import org.jboss.netty.buffer.ChannelBuffers;

import java.util.List;

/**
 * Container for the parameters of a upload request.
 */
public class UploadRequest {

    protected final String subject;
    protected List docs = null;

    /**
     * Constructor for UploadRequest.
     */
    public UploadRequest(String subject) {
        this.subject = subject;
    }

    /**
     * Default empty constructor for UploadRequest.
     */
    public UploadRequest() {
        this.subject = "products";
    }

    /**
     * Update documents.
     *
     * @return Updated request object.
     */
    public UploadRequest withDocs(List docs) {
        this.docs = Lists.newArrayListWithCapacity(docs.size());
        for (Document d : docs) {
            this.docs.add(d.underlying);
        }
        return this;
    }

    /**
     * Append documents.
     *
     * @return Updated request object.
     */
    public UploadRequest addDocs(List docs) {
        if (this.docs == null) {
            this.docs = Lists.newArrayListWithCapacity(docs.size());
        }
        for (Document d : docs) {
            this.docs.add(d.underlying);
        }
        return this;
    }

    /**
     * Append a document.
     *
     * @return Updated request object.
     */
    public UploadRequest addDoc(Document doc) {
        if (this.docs == null) {
            this.docs = Lists.newArrayList();
        }
        this.docs.add(doc.underlying);
        return this;
    }

    /**
     * Update documents in JsonNode format.
     *
     * @return Updated request object.
     */
    public UploadRequest withDocJsons(List docs) {
        this.docs = Lists.newArrayListWithCapacity(docs.size());
        this.docs.addAll(docs);
        return this;
    }

    /**
     * Append documents in JsonNode format.
     *
     * @return Updated request object.
     */
    public UploadRequest addDocJsons(List docs) {
        if (this.docs == null) {
            this.docs = Lists.newArrayListWithCapacity(docs.size());
        }
        this.docs.addAll(docs);
        return this;
    }

    /**
     * Append a document in JsonNode format.
     *
     * @return Updated request object.
     */
    public UploadRequest addDocJson(JsonNode doc) {
        if (this.docs == null) {
            this.docs = Lists.newArrayList();
        }
        this.docs.add(doc);
        return this;
    }

    /**
     * Update documents in json format.
     *
     * @return Updated request object.
     */
    public UploadRequest withDocStrings(List docStrings) throws InvalidRequestException {
        this.docs = Lists.newArrayListWithCapacity(docStrings.size());
        for (String s : docStrings) {
            try {
                JsonNode json = Document.Reader.readValue(s);
                this.docs.add(json);
            } catch (Exception e) {
                throw new InvalidRequestException("Document is in invalid JSON format", e);
            }
        }
        return this;
    }

    /**
     * Append documents in json format.
     *
     * @return Updated request object.
     */
    public UploadRequest addDocStrings(List docStrings) throws InvalidRequestException {
        if (this.docs == null) {
            this.docs = Lists.newArrayListWithCapacity(docStrings.size());
        }
        for (String s : docStrings) {
            try {
                JsonNode json = Document.Reader.readValue(s);
                this.docs.add(json);
            } catch (Exception e) {
                throw new InvalidRequestException("Document is in invalid JSON format", e);
            }
        }
        return this;
    }

    /**
     * Append a document in json format.
     *
     * @return Updated request object.
     */
    public UploadRequest addDocString(String docString) throws InvalidRequestException {
        if (this.docs == null) {
            this.docs = Lists.newArrayList();
        }
        try {
            JsonNode json = Document.Reader.readValue(docString);
            this.docs.add(json);
        } catch (Exception e) {
            throw new InvalidRequestException("Document is in invalid JSON format", e);
        }
        return this;
    }

    /**
     * For internal use. Returns query-string for http request.
     * @exclude
     */
    public String toQueryStr(String uploadType, String instanceId, String user, String token) throws InvalidRequestException {
        String[] split = instanceId.split("\\.");
        if (split.length != 3) {
            throw new InvalidRequestException("instanceId is not in valid format");
        } else {
            StringBuilder builder = new StringBuilder("/");
            builder.append(uploadType);
            builder.append("?company_id=");
            builder.append(split[0]);
            builder.append("&env=");
            builder.append(split[1]);
            builder.append("&instance_name=");
            builder.append(split[2]);
            builder.append("&user=");
            builder.append(user);
            builder.append("&token=");
            builder.append(token);
            return builder.toString();
        }
    }

    public ChannelBuffer toRequestBody() throws InvalidRequestException {
        ArrayNode dataJson = new ArrayNode(JsonNodeFactory.instance).addAll(docs);
        ObjectNode subjectJson = new ObjectNode(JsonNodeFactory.instance);
        subjectJson.put("name", subject);
        subjectJson.set("data", dataJson);
        ObjectNode bodyJson = new ObjectNode(JsonNodeFactory.instance);
        bodyJson.set("subjects", (new ArrayNode(JsonNodeFactory.instance)).add(subjectJson));
        try {
            return ChannelBuffers.copiedBuffer(Document.Writer.writeValueAsString(bodyJson), Charsets.UTF_8);
        } catch (JsonProcessingException e) {
            throw new InvalidRequestException("Could not encode JSON body for upload request", e);
        }
    }
}