gobblin.source.extractor.extract.restapi.RestApiConnector.java Source code

Java tutorial

Introduction

Here is the source code for gobblin.source.extractor.extract.restapi.RestApiConnector.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package gobblin.source.extractor.extract.restapi;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;

import org.apache.commons.io.IOUtils;
import org.apache.http.HttpEntity;
import org.apache.http.HttpResponse;
import org.apache.http.StatusLine;
import org.apache.http.client.HttpClient;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpRequestBase;
import org.apache.http.util.EntityUtils;

import com.google.common.base.Charsets;
import com.google.gson.Gson;
import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;

import gobblin.configuration.ConfigurationKeys;
import gobblin.configuration.State;
import gobblin.http.HttpClientConfiguratorLoader;
import gobblin.source.extractor.exception.RestApiConnectionException;
import gobblin.source.extractor.exception.RestApiProcessingException;
import gobblin.source.extractor.extract.Command;
import gobblin.source.extractor.extract.CommandOutput;
import gobblin.source.extractor.extract.restapi.RestApiCommand.RestApiCommandType;

import lombok.Setter;
import lombok.extern.slf4j.Slf4j;

/**
 * A class for connecting to Rest APIs, construct queries and getting responses.
 */
@Slf4j
public abstract class RestApiConnector {

    public static final String REST_API_CONNECTOR_CLASS = "rest.api.connector.class";

    protected static final Gson GSON = new Gson();

    protected HttpClient httpClient = null;
    protected boolean autoEstablishAuthToken = false;

    @Setter
    protected long authTokenTimeout;
    protected String accessToken = null;
    protected long createdAt;
    protected String instanceUrl;
    protected String updatedQuery;

    protected final State state;

    public RestApiConnector(State state) {
        this.state = state;
        this.authTokenTimeout = state.getPropAsInt(ConfigurationKeys.SOURCE_CONN_TIMEOUT,
                ConfigurationKeys.DEFAULT_CONN_TIMEOUT);
    }

    /**
     * get http connection
     * @return true if the connection is success else false
     */
    public boolean connect() throws RestApiConnectionException {
        if (this.autoEstablishAuthToken) {
            if (this.authTokenTimeout <= 0) {
                return false;
            } else if ((System.currentTimeMillis() - this.createdAt) > this.authTokenTimeout) {
                return false;
            }
        }

        HttpEntity httpEntity = null;
        try {
            httpEntity = getAuthentication();

            if (httpEntity != null) {
                JsonElement json = GSON.fromJson(EntityUtils.toString(httpEntity), JsonObject.class);
                if (json == null) {
                    log.error("Http entity: " + httpEntity);
                    log.error("entity class: " + httpEntity.getClass().getName());
                    log.error("entity string size: " + EntityUtils.toString(httpEntity).length());
                    log.error("content length: " + httpEntity.getContentLength());
                    log.error("content: " + IOUtils.toString(httpEntity.getContent(), Charsets.UTF_8));
                    throw new RestApiConnectionException(
                            "JSON is NULL ! Failed on authentication with the following HTTP response received:\n"
                                    + EntityUtils.toString(httpEntity));
                }

                JsonObject jsonRet = json.getAsJsonObject();
                log.info("jsonRet: " + jsonRet.toString());
                parseAuthenticationResponse(jsonRet);
            }
        } catch (IOException e) {
            throw new RestApiConnectionException("Failed to get rest api connection; error - " + e.getMessage(), e);
        } finally {
            if (httpEntity != null) {
                try {
                    EntityUtils.consume(httpEntity);
                } catch (IOException e) {
                    throw new RestApiConnectionException("Failed to consume httpEntity; error - " + e.getMessage(),
                            e);
                }
            }
        }

        return true;
    }

    protected HttpClient getHttpClient() {
        if (this.httpClient == null) {
            HttpClientConfiguratorLoader configuratorLoader = new HttpClientConfiguratorLoader(this.state);
            this.httpClient = configuratorLoader.getConfigurator()
                    .setStatePropertiesPrefix(ConfigurationKeys.SOURCE_CONN_PREFIX).configure(this.state)
                    .createClient();
        }
        return this.httpClient;
    }

    private static boolean hasId(JsonObject json) {
        if (json.has("id") || json.has("Id") || json.has("ID") || json.has("iD")) {
            return true;
        }
        return false;
    }

    /**
     * get http response in json format using url
     * @return json string with the response
     */
    public CommandOutput<?, ?> getResponse(List<Command> cmds) throws RestApiProcessingException {
        String url = cmds.get(0).getParams().get(0);

        log.info("URL: " + url);
        String jsonStr = null;
        HttpRequestBase httpRequest = new HttpGet(url);
        addHeaders(httpRequest);
        HttpEntity httpEntity = null;
        HttpResponse httpResponse = null;
        try {
            httpResponse = this.httpClient.execute(httpRequest);
            StatusLine status = httpResponse.getStatusLine();
            httpEntity = httpResponse.getEntity();

            if (httpEntity != null) {
                jsonStr = EntityUtils.toString(httpEntity);
            }

            if (status.getStatusCode() >= 400) {
                log.info("Unable to get response using: " + url);
                JsonElement jsonRet = GSON.fromJson(jsonStr, JsonArray.class);
                throw new RestApiProcessingException(
                        getFirstErrorMessage("Failed to retrieve response from", jsonRet));
            }
        } catch (Exception e) {
            throw new RestApiProcessingException("Failed to process rest api request; error - " + e.getMessage(),
                    e);
        } finally {
            try {
                if (httpEntity != null) {
                    EntityUtils.consume(httpEntity);
                }
                // httpResponse.close();
            } catch (Exception e) {
                throw new RestApiProcessingException("Failed to consume httpEntity; error - " + e.getMessage(), e);
            }
        }
        CommandOutput<RestApiCommand, String> output = new RestApiCommandOutput();
        output.put((RestApiCommand) cmds.get(0), jsonStr);
        return output;
    }

    protected void addHeaders(HttpRequestBase httpRequest) {
        if (this.accessToken != null) {
            httpRequest.addHeader("Authorization", "OAuth " + this.accessToken);
        }
        httpRequest.addHeader("Content-Type", "application/json");
        //httpRequest.addHeader("Accept-Encoding", "zip");
        //httpRequest.addHeader("Content-Encoding", "gzip");
        //httpRequest.addHeader("Connection", "Keep-Alive");
        //httpRequest.addHeader("Keep-Alive", "timeout=60000");
    }

    /**
     * get error message while executing http url
     * @return error message
     */
    private static String getFirstErrorMessage(String defaultMessage, JsonElement json) {
        if (json == null) {
            return defaultMessage;
        }

        JsonObject jsonObject = null;

        if (!json.isJsonArray()) {
            jsonObject = json.getAsJsonObject();
        } else {
            JsonArray jsonArray = json.getAsJsonArray();
            if (jsonArray.size() != 0) {
                jsonObject = jsonArray.get(0).getAsJsonObject();
            }
        }

        if (jsonObject != null) {
            if (jsonObject.has("error_description")) {
                defaultMessage = defaultMessage + jsonObject.get("error_description").getAsString();
            } else if (jsonObject.has("message")) {
                defaultMessage = defaultMessage + jsonObject.get("message").getAsString();
            }
        }

        return defaultMessage;
    }

    /**
     * Build a list of {@link Command}s given a String Rest query.
     */
    public static List<Command> constructGetCommand(String restQuery) {
        return Arrays.asList(new RestApiCommand().build(Arrays.asList(restQuery), RestApiCommandType.GET));
    }

    public boolean isConnectionClosed() {
        return this.httpClient == null;
    }

    /**
     * To be overridden by subclasses that require authentication.
     */
    public abstract HttpEntity getAuthentication() throws RestApiConnectionException;

    protected void parseAuthenticationResponse(JsonObject jsonRet) throws RestApiConnectionException {
        if (!hasId(jsonRet)) {
            throw new RestApiConnectionException(
                    "Failed on authentication with the following HTTP response received:" + jsonRet.toString());
        }

        this.instanceUrl = jsonRet.get("instance_url").getAsString();
        this.accessToken = jsonRet.get("access_token").getAsString();
        this.createdAt = System.currentTimeMillis();
    }
}