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.source.extractor.extract.restapi; import com.google.common.collect.ImmutableList; import gobblin.source.extractor.exception.RestApiConnectionException; import gobblin.source.extractor.exception.RestApiProcessingException; import gobblin.source.extractor.utils.Utils; import java.io.IOException; import java.util.Iterator; import java.util.List; import com.google.common.base.Joiner; import com.google.common.base.Strings; import com.google.gson.Gson; import com.google.gson.JsonArray; import com.google.gson.JsonElement; import com.google.gson.JsonObject; import com.google.common.base.Splitter; import gobblin.configuration.ConfigurationKeys; import gobblin.configuration.WorkUnitState; import gobblin.source.extractor.watermark.Predicate; import gobblin.source.extractor.DataRecordException; import gobblin.source.extractor.exception.HighWatermarkException; import gobblin.source.extractor.exception.RecordCountException; import gobblin.source.extractor.exception.SchemaException; import gobblin.source.extractor.extract.QueryBasedExtractor; import gobblin.source.extractor.extract.Command; import gobblin.source.extractor.extract.CommandOutput; import gobblin.source.extractor.extract.SourceSpecificLayer; import gobblin.source.extractor.schema.Schema; import gobblin.source.workunit.WorkUnit; import lombok.extern.slf4j.Slf4j; /** * An implementation of rest api extractor for the sources that are using rest api * * @param <D> type of data record * @param <S> type of schema */ @Slf4j public abstract class RestApiExtractor extends QueryBasedExtractor<JsonArray, JsonElement> implements SourceSpecificLayer<JsonArray, JsonElement>, RestApiSpecificLayer { private static final Gson GSON = new Gson(); protected String instanceUrl; protected String updatedQuery; protected final RestApiConnector connector; public RestApiExtractor(WorkUnitState state) { super(state); this.connector = getConnector(state); } protected abstract RestApiConnector getConnector(WorkUnitState state); protected String buildDataQuery(String inputQuery, String entity) { String dataQuery = null; if (inputQuery == null && this.columnList.size() != 0) { // if input query is null, build the query from metadata dataQuery = "SELECT " + Joiner.on(",").join(this.columnList) + " FROM " + entity; } else { // if input query is not null, build the query with intersection of columns from input query and columns from Metadata if (inputQuery != null) { String queryLowerCase = inputQuery.toLowerCase(); int columnsStartIndex = queryLowerCase.indexOf("select ") + 7; int columnsEndIndex = queryLowerCase.indexOf(" from "); if (columnsStartIndex > 0 && columnsEndIndex > 0) { String givenColumnList = inputQuery.substring(columnsStartIndex, columnsEndIndex); dataQuery = inputQuery.replace(givenColumnList, Joiner.on(",").join(this.columnList)); } else { dataQuery = inputQuery; } } } log.info("Updated data query: " + dataQuery); return dataQuery; } @Override public void extractMetadata(String schema, String entity, WorkUnit workUnit) throws SchemaException { log.info("Extract Metadata using Rest Api"); JsonArray columnArray = new JsonArray(); String inputQuery = workUnitState.getProp(ConfigurationKeys.SOURCE_QUERYBASED_QUERY); List<String> columnListInQuery = null; JsonArray array = null; if (!Strings.isNullOrEmpty(inputQuery)) { columnListInQuery = Utils.getColumnListFromQuery(inputQuery); } String excludedColumns = workUnitState.getProp(ConfigurationKeys.SOURCE_QUERYBASED_EXCLUDED_COLUMNS); List<String> columnListExcluded = ImmutableList.<String>of(); if (Strings.isNullOrEmpty(inputQuery) && !Strings.isNullOrEmpty(excludedColumns)) { Splitter splitter = Splitter.on(",").omitEmptyStrings().trimResults(); columnListExcluded = splitter.splitToList(excludedColumns.toLowerCase()); } try { boolean success = this.connector.connect(); if (!success) { throw new SchemaException("Failed to connect."); } log.debug("Connected successfully."); List<Command> cmds = this.getSchemaMetadata(schema, entity); CommandOutput<?, ?> response = this.connector.getResponse(cmds); array = this.getSchema(response); for (JsonElement columnElement : array) { Schema obj = GSON.fromJson(columnElement, Schema.class); String columnName = obj.getColumnName(); obj.setWaterMark(this.isWatermarkColumn(workUnitState.getProp("extract.delta.fields"), columnName)); if (this.isWatermarkColumn(workUnitState.getProp("extract.delta.fields"), columnName)) { obj.setNullable(false); } else if (this.getPrimarykeyIndex(workUnitState.getProp("extract.primary.key.fields"), columnName) == 0) { // set all columns as nullable except primary key and watermark columns obj.setNullable(true); } obj.setPrimaryKey( this.getPrimarykeyIndex(workUnitState.getProp("extract.primary.key.fields"), columnName)); String jsonStr = GSON.toJson(obj); JsonObject jsonObject = GSON.fromJson(jsonStr, JsonObject.class).getAsJsonObject(); // If input query is null or provided '*' in the query select all columns. // Else, consider only the columns mentioned in the column list if (inputQuery == null || columnListInQuery == null || (columnListInQuery.size() == 1 && columnListInQuery.get(0).equals("*")) || (columnListInQuery.size() >= 1 && this.isMetadataColumn(columnName, columnListInQuery))) { if (!columnListExcluded.contains(columnName.trim().toLowerCase())) { this.columnList.add(columnName); columnArray.add(jsonObject); } } } this.updatedQuery = buildDataQuery(inputQuery, entity); log.info("Schema:" + columnArray); this.setOutputSchema(columnArray); } catch (RuntimeException | RestApiConnectionException | RestApiProcessingException | IOException | SchemaException e) { throw new SchemaException("Failed to get schema using rest api; error - " + e.getMessage(), e); } } @Override public long getMaxWatermark(String schema, String entity, String watermarkColumn, List<Predicate> predicateList, String watermarkSourceFormat) throws HighWatermarkException { log.info("Get high watermark using Rest Api"); long CalculatedHighWatermark = -1; try { boolean success = this.connector.connect(); if (!success) { throw new HighWatermarkException("Failed to connect."); } log.debug("Connected successfully."); List<Command> cmds = this.getHighWatermarkMetadata(schema, entity, watermarkColumn, predicateList); CommandOutput<?, ?> response = this.connector.getResponse(cmds); CalculatedHighWatermark = this.getHighWatermark(response, watermarkColumn, watermarkSourceFormat); log.info("High watermark:" + CalculatedHighWatermark); return CalculatedHighWatermark; } catch (Exception e) { throw new HighWatermarkException( "Failed to get high watermark using rest api; error - " + e.getMessage(), e); } } @Override public long getSourceCount(String schema, String entity, WorkUnit workUnit, List<Predicate> predicateList) throws RecordCountException { log.info("Get source record count using Rest Api"); long count = 0; try { boolean success = this.connector.connect(); if (!success) { throw new RecordCountException("Failed to connect."); } log.debug("Connected successfully."); List<Command> cmds = this.getCountMetadata(schema, entity, workUnit, predicateList); CommandOutput<?, ?> response = this.connector.getResponse(cmds); count = getCount(response); log.info("Source record count:" + count); return count; } catch (Exception e) { throw new RecordCountException("Failed to get record count using rest api; error - " + e.getMessage(), e); } } @Override public Iterator<JsonElement> getRecordSet(String schema, String entity, WorkUnit workUnit, List<Predicate> predicateList) throws DataRecordException { log.debug("Get data records using Rest Api"); Iterator<JsonElement> rs = null; List<Command> cmds; try { boolean success = true; if (this.connector.isConnectionClosed()) { success = this.connector.connect(); } if (!success) { throw new DataRecordException("Failed to connect."); } log.debug("Connected successfully."); if (this.getPullStatus() == false) { return null; } if (this.getNextUrl() == null) { cmds = this.getDataMetadata(schema, entity, workUnit, predicateList); } else { cmds = RestApiConnector.constructGetCommand(this.getNextUrl()); } CommandOutput<?, ?> response = this.connector.getResponse(cmds); rs = this.getData(response); return rs; } catch (Exception e) { throw new DataRecordException("Failed to get records using rest api; error - " + e.getMessage(), e); } } @Override public void setTimeOut(int timeOut) { this.connector.setAuthTokenTimeout(timeOut); } }