use of org.apache.gobblin.source.extractor.extract.Command in project incubator-gobblin by apache.
the class RestApiExtractor method extractMetadata.
@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();
// 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);
}
}
use of org.apache.gobblin.source.extractor.extract.Command in project incubator-gobblin by apache.
the class JdbcExtractor method getRecordSet.
@Override
public Iterator<JsonElement> getRecordSet(String schema, String entity, WorkUnit workUnit, List<Predicate> predicateList) throws DataRecordException, IOException {
Iterator<JsonElement> rs = null;
List<Command> cmds;
try {
if (isFirstPull()) {
this.log.info("Get data recordset using JDBC");
cmds = this.getDataMetadata(schema, entity, workUnit, predicateList);
this.dataResponse = this.executePreparedSql(cmds);
this.setFirstPull(false);
}
rs = this.getData(this.dataResponse);
return rs;
} catch (Exception e) {
throw new DataRecordException("Failed to get record set using JDBC; error - " + e.getMessage(), e);
}
}
use of org.apache.gobblin.source.extractor.extract.Command in project incubator-gobblin by apache.
the class JdbcExtractor method executeSql.
/**
* Execute query using JDBC simple Statement Set fetch size
*
* @param cmds commands - query, fetch size
* @return JDBC ResultSet
* @throws Exception
*/
private CommandOutput<?, ?> executeSql(List<Command> cmds) {
String query = null;
int fetchSize = 0;
for (Command cmd : cmds) {
if (cmd instanceof JdbcCommand) {
JdbcCommandType type = (JdbcCommandType) cmd.getCommandType();
switch(type) {
case QUERY:
query = cmd.getParams().get(0);
break;
case FETCHSIZE:
fetchSize = Integer.parseInt(cmd.getParams().get(0));
break;
default:
this.log.error("Command " + type.toString() + " not recognized");
break;
}
}
}
this.log.info("Executing query:" + query);
ResultSet resultSet = null;
try {
this.jdbcSource = createJdbcSource();
this.dataConnection = this.jdbcSource.getConnection();
Statement statement = this.dataConnection.createStatement();
if (fetchSize != 0 && this.getExpectedRecordCount() > 2000) {
statement.setFetchSize(fetchSize);
}
final boolean status = statement.execute(query);
if (status == false) {
this.log.error("Failed to execute sql:" + query);
}
resultSet = statement.getResultSet();
} catch (Exception e) {
this.log.error("Failed to execute sql:" + query + " ;error-" + e.getMessage(), e);
}
CommandOutput<JdbcCommand, ResultSet> output = new JdbcCommandOutput();
output.put((JdbcCommand) cmds.get(0), resultSet);
return output;
}
use of org.apache.gobblin.source.extractor.extract.Command in project incubator-gobblin by apache.
the class JdbcExtractor method executePreparedSql.
/**
* Execute query using JDBC PreparedStatement to pass query parameters Set
* fetch size
*
* @param cmds commands - query, fetch size, query parameters
* @return JDBC ResultSet
* @throws Exception
*/
private CommandOutput<?, ?> executePreparedSql(List<Command> cmds) {
String query = null;
List<String> queryParameters = null;
int fetchSize = 0;
for (Command cmd : cmds) {
if (cmd instanceof JdbcCommand) {
JdbcCommandType type = (JdbcCommandType) cmd.getCommandType();
switch(type) {
case QUERY:
query = cmd.getParams().get(0);
break;
case QUERYPARAMS:
queryParameters = cmd.getParams();
break;
case FETCHSIZE:
fetchSize = Integer.parseInt(cmd.getParams().get(0));
break;
default:
this.log.error("Command " + type.toString() + " not recognized");
break;
}
}
}
this.log.info("Executing query:" + query);
ResultSet resultSet = null;
try {
this.jdbcSource = createJdbcSource();
this.dataConnection = this.jdbcSource.getConnection();
PreparedStatement statement = this.dataConnection.prepareStatement(query, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY);
int parameterPosition = 1;
if (queryParameters != null && queryParameters.size() > 0) {
for (String parameter : queryParameters) {
statement.setString(parameterPosition, parameter);
parameterPosition++;
}
}
if (fetchSize != 0) {
statement.setFetchSize(fetchSize);
}
final boolean status = statement.execute();
if (status == false) {
this.log.error("Failed to execute sql:" + query);
}
resultSet = statement.getResultSet();
} catch (Exception e) {
this.log.error("Failed to execute sql:" + query + " ;error-" + e.getMessage(), e);
}
CommandOutput<JdbcCommand, ResultSet> output = new JdbcCommandOutput();
output.put((JdbcCommand) cmds.get(0), resultSet);
return output;
}
use of org.apache.gobblin.source.extractor.extract.Command in project incubator-gobblin by apache.
the class JdbcExtractor method extractMetadata.
@Override
public void extractMetadata(String schema, String entity, WorkUnit workUnit) throws SchemaException, IOException {
this.log.info("Extract metadata using JDBC");
String inputQuery = workUnitState.getProp(ConfigurationKeys.SOURCE_QUERYBASED_QUERY);
if (hasJoinOperation(inputQuery)) {
throw new RuntimeException("Query across multiple tables not supported");
}
String watermarkColumn = workUnitState.getProp(ConfigurationKeys.EXTRACT_DELTA_FIELDS_KEY);
this.enableDelimitedIdentifier = workUnitState.getPropAsBoolean(ConfigurationKeys.ENABLE_DELIMITED_IDENTIFIER, ConfigurationKeys.DEFAULT_ENABLE_DELIMITED_IDENTIFIER);
JsonObject defaultWatermark = this.getDefaultWatermark();
String derivedWatermarkColumnName = defaultWatermark.get("columnName").getAsString();
this.setSampleRecordCount(this.exractSampleRecordCountFromQuery(inputQuery));
inputQuery = this.removeSampleClauseFromQuery(inputQuery);
JsonArray targetSchema = new JsonArray();
List<String> headerColumns = new ArrayList<>();
try {
List<Command> cmds = this.getSchemaMetadata(schema, entity);
CommandOutput<?, ?> response = this.executePreparedSql(cmds);
JsonArray array = this.getSchema(response);
this.buildMetadataColumnMap(array);
this.parseInputQuery(inputQuery);
List<String> sourceColumns = this.getMetadataColumnList();
for (ColumnAttributes colMap : this.columnAliasMap) {
String alias = colMap.getAliasName();
String columnName = colMap.getColumnName();
String sourceColumnName = colMap.getSourceColumnName();
if (this.isMetadataColumn(columnName, sourceColumns)) {
String targetColumnName = this.getTargetColumnName(columnName, alias);
Schema obj = this.getUpdatedSchemaObject(columnName, alias, targetColumnName);
String jsonStr = gson.toJson(obj);
JsonObject jsonObject = gson.fromJson(jsonStr, JsonObject.class).getAsJsonObject();
targetSchema.add(jsonObject);
headerColumns.add(targetColumnName);
sourceColumnName = getLeftDelimitedIdentifier() + sourceColumnName + getRightDelimitedIdentifier();
this.columnList.add(sourceColumnName);
}
}
if (this.hasMultipleWatermarkColumns(watermarkColumn)) {
derivedWatermarkColumnName = getLeftDelimitedIdentifier() + derivedWatermarkColumnName + getRightDelimitedIdentifier();
this.columnList.add(derivedWatermarkColumnName);
headerColumns.add(derivedWatermarkColumnName);
targetSchema.add(defaultWatermark);
this.workUnitState.setProp(ConfigurationKeys.EXTRACT_DELTA_FIELDS_KEY, derivedWatermarkColumnName);
}
String outputColProjection = Joiner.on(",").useForNull("null").join(this.columnList);
outputColProjection = outputColProjection.replace(derivedWatermarkColumnName, Utils.getCoalesceColumnNames(watermarkColumn) + " AS " + derivedWatermarkColumnName);
this.setOutputColumnProjection(outputColProjection);
String extractQuery = this.getExtractQuery(schema, entity, inputQuery);
this.setHeaderRecord(headerColumns);
this.setOutputSchema(targetSchema);
this.setExtractSql(extractQuery);
// this.workUnit.getProp(ConfigurationKeys.EXTRACT_TABLE_NAME_KEY,
// this.escapeCharsInColumnName(this.workUnit.getProp(ConfigurationKeys.SOURCE_ENTITY),
// ConfigurationKeys.ESCAPE_CHARS_IN_COLUMN_NAME, "_"));
this.log.info("Schema:" + targetSchema);
this.log.info("Extract query: " + this.getExtractSql());
} catch (RuntimeException | IOException | SchemaException e) {
throw new SchemaException("Failed to get metadata using JDBC; error - " + e.getMessage(), e);
}
}
Aggregations