com.blackbirdai.client.AsyncUploadClient Maven / Gradle / Ivy

Go to download
package com.blackbirdai.client;

import com.blackbirdai.client.model.DeleteRequest;
import com.blackbirdai.client.model.UploadRequest;
import com.blackbirdai.client.model.UploadResponse;
import com.twitter.finagle.Http;
import com.twitter.finagle.Service;
import com.twitter.util.Future;
import com.twitter.util.FutureTransformer;
import org.jboss.netty.handler.codec.http.*;

/**
 * Asynchronous client for uploading data to blackbird search.
 */
public class AsyncUploadClient {

    private final String host;
    private final String user;
    private final String token;
    private final String instanceId;
    private final Service httpClient;

    private static final String APPLICATION_JSON = "application/json";

    /**
     * Constructor for AsyncClient.
     * @param host Address of Blackbird search endpoint.
     * @param user User name for your Blackbird account.
     * @param token Authentication token for upload. Available on admin page.
     * @param instanceId Id of search instance.
     */
    public AsyncUploadClient(String host, String user, String token, String instanceId) {
        this.host = host;
        this.user = user;
        this.token = token;
        this.instanceId = instanceId;
        this.httpClient = Http.client().withTls(host).newService(host + ":443");
    }

    /**
     * Add documents.
     */
    public Future add(UploadRequest request) { return uploadInternal(request, "add"); }

    /**
     * Update documents.
     */
    public Future update(UploadRequest request) {
        return uploadInternal(request, "update");
    }

    /**
     * Upsert documents.
     */
    public Future upsert(UploadRequest request) {
        return uploadInternal(request, "upsert");
    }

    /**
     * Delete documents.
     */
    public Future delete(DeleteRequest request) {
        return uploadInternal(request, "delete");
    }

    private Future uploadInternal(UploadRequest request, String uploadType) {
        try {
            String uri = request.toQueryStr(uploadType, instanceId, user, token);
            HttpRequest httpRequest = new DefaultHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.POST, uri);
            httpRequest.headers().set(HttpHeaders.Names.HOST, host);
            httpRequest.headers().set(HttpHeaders.Names.CONNECTION, HttpHeaders.Values.KEEP_ALIVE);
            httpRequest.headers().set(HttpHeaders.Names.CONTENT_TYPE, APPLICATION_JSON);
            httpRequest.setContent(request.toRequestBody());
            return httpClient.apply(httpRequest).transformedBy(new FutureTransformer() {
                @Override
                public UploadResponse map(HttpResponse httpResponse) {
                    return new UploadResponse(httpResponse);
                }

                @Override
                public Future rescue(Throwable t) {
                    return Future.value(new UploadResponse(false, "Request failed: " + t.getMessage()));
                }
            });
        } catch (InvalidRequestException e) {
            return Future.value(new UploadResponse(false, "Invalid request: " + e.getMessage()));
        }
    }

}