Java tutorial
/* * 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.writer.http; import java.io.IOException; import java.net.URI; import java.net.URISyntaxException; import org.apache.commons.lang3.builder.ToStringBuilder; import org.apache.commons.lang3.builder.ToStringStyle; import org.apache.commons.lang3.StringUtils; import org.apache.http.HttpHeaders; import org.apache.http.client.methods.CloseableHttpResponse; import org.apache.http.client.methods.HttpUriRequest; import org.apache.http.client.methods.RequestBuilder; import org.apache.http.entity.ContentType; import org.apache.http.entity.StringEntity; import org.apache.http.util.EntityUtils; import org.apache.oltu.oauth2.client.OAuthClient; import org.apache.oltu.oauth2.client.URLConnectionClient; import org.apache.oltu.oauth2.client.request.OAuthClientRequest; import org.apache.oltu.oauth2.client.response.OAuthJSONAccessTokenResponse; import org.apache.oltu.oauth2.common.OAuth; import org.apache.oltu.oauth2.common.exception.OAuthProblemException; import org.apache.oltu.oauth2.common.exception.OAuthSystemException; import org.apache.oltu.oauth2.common.message.types.GrantType; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Optional; import com.google.common.base.Preconditions; import com.google.gson.JsonArray; import com.google.gson.JsonElement; import com.google.gson.JsonObject; import com.google.gson.JsonParser; import gobblin.converter.http.RestEntry; import gobblin.writer.exception.NonTransientException; /** * Writes to Salesforce via RESTful API, supporting INSERT_ONLY_NOT_EXIST, and UPSERT. * */ public class SalesforceRestWriter extends RestJsonWriter { public static enum Operation { INSERT_ONLY_NOT_EXIST, UPSERT } static final String DUPLICATE_VALUE_ERR_CODE = "DUPLICATE_VALUE"; private String accessToken; private final URI oauthEndPoint; private final String clientId; private final String clientSecret; private final String userId; private final String password; private final String securityToken; private final Operation operation; private final int batchSize; private final Optional<String> batchResourcePath; private Optional<JsonArray> batchRecords = Optional.absent(); private long numRecordsWritten = 0L; public SalesforceRestWriter(SalesForceRestWriterBuilder builder) { super(builder); this.oauthEndPoint = builder.getSvcEndpoint().get(); //Set oauth end point this.clientId = builder.getClientId(); this.clientSecret = builder.getClientSecret(); this.userId = builder.getUserId(); this.password = builder.getPassword(); this.securityToken = builder.getSecurityToken(); this.operation = builder.getOperation(); this.batchSize = builder.getBatchSize(); this.batchResourcePath = builder.getBatchResourcePath(); Preconditions.checkArgument(batchSize == 1 || batchResourcePath.isPresent(), "Batch resource path is missing"); if (batchSize > 1) { getLog().info("Batch api will be used with batch size " + batchSize); } try { onConnect(oauthEndPoint); } catch (IOException e) { throw new RuntimeException(e); } } @VisibleForTesting SalesforceRestWriter(SalesForceRestWriterBuilder builder, String accessToken) { super(builder); this.oauthEndPoint = builder.getSvcEndpoint().get(); //Set oauth end point this.clientId = builder.getClientId(); this.clientSecret = builder.getClientSecret(); this.userId = builder.getUserId(); this.password = builder.getPassword(); this.securityToken = builder.getSecurityToken(); this.operation = builder.getOperation(); this.batchSize = builder.getBatchSize(); this.batchResourcePath = builder.getBatchResourcePath(); Preconditions.checkArgument(batchSize == 1 || batchResourcePath.isPresent(), "Batch resource path is missing"); this.accessToken = accessToken; } /** * Retrieve access token, if needed, retrieve instance url, and set server host URL * {@inheritDoc} * @see gobblin.writer.http.HttpWriter#onConnect(org.apache.http.HttpHost) */ @Override public void onConnect(URI serverHost) throws IOException { if (!StringUtils.isEmpty(accessToken)) { return; //No need to be called if accessToken is active. } try { getLog().info("Getting Oauth2 access token."); OAuthClientRequest request = OAuthClientRequest.tokenLocation(serverHost.toString()) .setGrantType(GrantType.PASSWORD).setClientId(clientId).setClientSecret(clientSecret) .setUsername(userId).setPassword(password + securityToken).buildQueryMessage(); OAuthClient client = new OAuthClient(new URLConnectionClient()); OAuthJSONAccessTokenResponse response = client.accessToken(request, OAuth.HttpMethod.POST); accessToken = response.getAccessToken(); setCurServerHost(new URI(response.getParam("instance_url"))); } catch (OAuthProblemException e) { throw new NonTransientException("Error while authenticating with Oauth2", e); } catch (OAuthSystemException e) { throw new RuntimeException("Failed getting access token", e); } catch (URISyntaxException e) { throw new RuntimeException("Failed due to invalid instance url", e); } } /** * For single request, creates HttpUriRequest and decides post/patch operation based on input parameter. * * For batch request, add the record into JsonArray as a subrequest and only creates HttpUriRequest with POST method if it filled the batch size. * {@inheritDoc} * @see gobblin.writer.http.RestJsonWriter#onNewRecord(gobblin.converter.rest.RestEntry) */ @Override public Optional<HttpUriRequest> onNewRecord(RestEntry<JsonObject> record) { Preconditions.checkArgument(!StringUtils.isEmpty(accessToken), "Access token has not been acquired."); Preconditions.checkNotNull(record, "Record should not be null"); RequestBuilder builder = null; JsonObject payload = null; if (batchSize > 1) { if (!batchRecords.isPresent()) { batchRecords = Optional.of(new JsonArray()); } batchRecords.get().add(newSubrequest(record)); if (batchRecords.get().size() < batchSize) { //No need to send. Return absent. return Optional.absent(); } payload = newPayloadForBatch(); builder = RequestBuilder.post().setUri(combineUrl(getCurServerHost(), batchResourcePath)); } else { switch (operation) { case INSERT_ONLY_NOT_EXIST: builder = RequestBuilder.post(); break; case UPSERT: builder = RequestBuilder.patch(); break; default: throw new IllegalArgumentException(operation + " is not supported."); } builder.setUri(combineUrl(getCurServerHost(), record.getResourcePath())); payload = record.getRestEntryVal(); } return Optional.of(newRequest(builder, payload)); } /** * Create batch subrequest. For more detail @link https://developer.salesforce.com/docs/atlas.en-us.api_rest.meta/api_rest/requests_composite_batch.htm * * @param record * @return */ private JsonObject newSubrequest(RestEntry<JsonObject> record) { Preconditions.checkArgument(record.getResourcePath().isPresent(), "Resource path is not defined"); JsonObject subReq = new JsonObject(); subReq.addProperty("url", record.getResourcePath().get()); subReq.add("richInput", record.getRestEntryVal()); switch (operation) { case INSERT_ONLY_NOT_EXIST: subReq.addProperty("method", "POST"); break; case UPSERT: subReq.addProperty("method", "PATCH"); break; default: throw new IllegalArgumentException(operation + " is not supported."); } return subReq; } /** * @return JsonObject contains batch records */ private JsonObject newPayloadForBatch() { JsonObject payload = new JsonObject(); payload.add("batchRequests", batchRecords.get()); return payload; } private HttpUriRequest newRequest(RequestBuilder builder, JsonElement payload) { try { builder.addHeader(HttpHeaders.CONTENT_TYPE, ContentType.APPLICATION_JSON.getMimeType()) .addHeader(HttpHeaders.AUTHORIZATION, "OAuth " + accessToken) .setEntity(new StringEntity(payload.toString(), ContentType.APPLICATION_JSON)); } catch (Exception e) { throw new RuntimeException(e); } if (getLog().isDebugEnabled()) { getLog().debug("Request builder: " + ToStringBuilder.reflectionToString(builder, ToStringStyle.SHORT_PREFIX_STYLE)); } return builder.build(); } @Override public void flush() { try { if (isRetry()) { //flushing failed and it should be retried. super.writeImpl(null); return; } if (batchRecords.isPresent() && batchRecords.get().size() > 0) { getLog().info( "Flusing remaining subrequest of batch. # of subrequests: " + batchRecords.get().size()); curRequest = Optional.of( newRequest(RequestBuilder.post().setUri(combineUrl(getCurServerHost(), batchResourcePath)), newPayloadForBatch())); super.writeImpl(null); } } catch (IOException e) { throw new RuntimeException(e); } } /** * Make it fail (throw exception) if status code is greater or equal to 400 except, * the status code is 400 and error code is duplicate value, regard it as success(do not throw exception). * * If status code is 401 or 403, re-acquire access token before make it fail -- retry will take care of rest. * * {@inheritDoc} * @see gobblin.writer.http.HttpWriter#processResponse(org.apache.http.HttpResponse) */ @Override public void processResponse(CloseableHttpResponse response) throws IOException, UnexpectedResponseException { if (getLog().isDebugEnabled()) { getLog().debug("Received response " + ToStringBuilder.reflectionToString(response, ToStringStyle.SHORT_PREFIX_STYLE)); } int statusCode = response.getStatusLine().getStatusCode(); if (statusCode == 401 || statusCode == 403) { getLog().info("Reacquiring access token."); accessToken = null; onConnect(oauthEndPoint); throw new RuntimeException( "Access denied. Access token has been reacquired and retry may solve the problem. " + ToStringBuilder.reflectionToString(response, ToStringStyle.SHORT_PREFIX_STYLE)); } if (batchSize > 1) { processBatchRequestResponse(response); numRecordsWritten += batchRecords.get().size(); batchRecords = Optional.absent(); } else { processSingleRequestResponse(response); numRecordsWritten++; } } private void processSingleRequestResponse(CloseableHttpResponse response) throws IOException, UnexpectedResponseException { int statusCode = response.getStatusLine().getStatusCode(); if (statusCode < 400) { return; } String entityStr = EntityUtils.toString(response.getEntity()); if (statusCode == 400 && Operation.INSERT_ONLY_NOT_EXIST.equals(operation) && entityStr != null) { //Ignore if it's duplicate entry error code JsonArray jsonArray = new JsonParser().parse(entityStr).getAsJsonArray(); JsonObject jsonObject = jsonArray.get(0).getAsJsonObject(); if (isDuplicate(jsonObject, statusCode)) { return; } } throw new RuntimeException("Failed due to " + entityStr + " (Detail: " + ToStringBuilder.reflectionToString(response, ToStringStyle.SHORT_PREFIX_STYLE) + " )"); } /** * Check results from batch response, if any of the results is failure throw exception. * @param response * @throws IOException * @throws UnexpectedResponseException */ private void processBatchRequestResponse(CloseableHttpResponse response) throws IOException, UnexpectedResponseException { String entityStr = EntityUtils.toString(response.getEntity()); int statusCode = response.getStatusLine().getStatusCode(); if (statusCode >= 400) { throw new RuntimeException("Failed due to " + entityStr + " (Detail: " + ToStringBuilder.reflectionToString(response, ToStringStyle.SHORT_PREFIX_STYLE) + " )"); } JsonObject jsonBody = new JsonParser().parse(entityStr).getAsJsonObject(); if (!jsonBody.get("hasErrors").getAsBoolean()) { return; } JsonArray results = jsonBody.get("results").getAsJsonArray(); for (JsonElement jsonElem : results) { JsonObject json = jsonElem.getAsJsonObject(); int subStatusCode = json.get("statusCode").getAsInt(); if (subStatusCode < 400) { continue; } else if (subStatusCode == 400 && Operation.INSERT_ONLY_NOT_EXIST.equals(operation)) { JsonElement resultJsonElem = json.get("result"); Preconditions.checkNotNull(resultJsonElem, "Error response should contain result property"); JsonObject resultJsonObject = resultJsonElem.getAsJsonArray().get(0).getAsJsonObject(); if (isDuplicate(resultJsonObject, subStatusCode)) { continue; } } throw new RuntimeException("Failed due to " + jsonBody + " (Detail: " + ToStringBuilder.reflectionToString(response, ToStringStyle.SHORT_PREFIX_STYLE) + " )"); } } private boolean isDuplicate(JsonObject responseJsonObject, int statusCode) { return statusCode == 400 && Operation.INSERT_ONLY_NOT_EXIST.equals(operation) && DUPLICATE_VALUE_ERR_CODE.equals(responseJsonObject.get("errorCode").getAsString()); } /** * {@inheritDoc} */ @Override public long recordsWritten() { return this.numRecordsWritten; } }