Examples with RecordCountException org.apache.gobblin.source.extractor.exception.RecordCountException used on opensource projects

Example 1 with RecordCountException

use of org.apache.gobblin.source.extractor.exception.RecordCountException in project incubator-gobblin by apache.

the class JdbcExtractor method getCount.

@Override
public long getCount(CommandOutput<?, ?> response) throws RecordCountException {
    this.log.debug("Extract source record count from resultset");
    ResultSet resultset = null;
    long count = 0;
    Iterator<ResultSet> itr = (Iterator<ResultSet>) response.getResults().values().iterator();
    if (itr.hasNext()) {
        resultset = itr.next();
        try {
            if (resultset.next()) {
                count = resultset.getLong(1);
            }
        } catch (Exception e) {
            throw new RecordCountException("Failed to get source record count from database; error - " + e.getMessage(), e);
        }
    } else {
        throw new RuntimeException("Failed to get source record count from database - Resultset has no records");
    }
    return count;
}

Example 2 with RecordCountException

use of org.apache.gobblin.source.extractor.exception.RecordCountException in project incubator-gobblin by apache.

the class SalesforceExtractor method getCount.

@Override
public long getCount(CommandOutput<?, ?> response) throws RecordCountException {
    log.info("Get source record count from salesforce");
    String output;
    Iterator<String> itr = (Iterator<String>) response.getResults().values().iterator();
    if (itr.hasNext()) {
        output = itr.next();
    } else {
        throw new RecordCountException("Failed to get count from salesforce; REST response has no output");
    }
    JsonElement element = GSON.fromJson(output, JsonObject.class);
    long count;
    try {
        JsonObject jsonObject = element.getAsJsonObject();
        count = jsonObject.get("totalSize").getAsLong();
    } catch (Exception e) {
        throw new RecordCountException("Failed to get record count from salesforce; error - " + e.getMessage(), e);
    }
    return count;
}

Example 3 with RecordCountException

use of org.apache.gobblin.source.extractor.exception.RecordCountException in project incubator-gobblin by apache.

the class SalesforceExtractor method getCountMetadata.

@Override
public List<Command> getCountMetadata(String schema, String entity, WorkUnit workUnit, List<Predicate> predicateList) throws RecordCountException {
    log.debug("Build url to retrieve source record count");
    String existingPredicate = "";
    if (this.updatedQuery != null) {
        String queryLowerCase = this.updatedQuery.toLowerCase();
        int startIndex = queryLowerCase.indexOf(" where ");
        if (startIndex > 0) {
            existingPredicate = this.updatedQuery.substring(startIndex);
        }
    }
    String query = "SELECT COUNT() FROM " + entity + existingPredicate;
    String limitString = getLimitFromInputQuery(query);
    query = query.replace(limitString, "");
    try {
        if (isNullPredicate(predicateList)) {
            log.info("QUERY with null predicate: " + query);
            return constructGetCommand(this.sfConnector.getFullUri(getSoqlUrl(query)));
        }
        Iterator<Predicate> i = predicateList.listIterator();
        while (i.hasNext()) {
            Predicate predicate = i.next();
            query = SqlQueryUtils.addPredicate(query, predicate.getCondition());
        }
        query = query + getLimitFromInputQuery(this.updatedQuery);
        log.info("QUERY: " + query);
        return constructGetCommand(this.sfConnector.getFullUri(getSoqlUrl(query)));
    } catch (Exception e) {
        throw new RecordCountException("Failed to get salesforce url for record count; error - " + e.getMessage(), e);
    }
}

Example 4 with RecordCountException

use of org.apache.gobblin.source.extractor.exception.RecordCountException in project incubator-gobblin by apache.

the class RestApiExtractor method getSourceCount.

@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);
    }
}

Example 5 with RecordCountException

use of org.apache.gobblin.source.extractor.exception.RecordCountException in project incubator-gobblin by apache.

the class QueryBasedExtractor method build.

/**
 * build schema, record count and high water mark
 */
public Extractor<S, D> build() throws ExtractPrepareException {
    String watermarkColumn = this.workUnitState.getProp(ConfigurationKeys.EXTRACT_DELTA_FIELDS_KEY);
    long lwm = partition.getLowWatermark();
    long hwm = partition.getHighWatermark();
    log.info("Low water mark: " + lwm + "; and High water mark: " + hwm);
    WatermarkType watermarkType;
    if (StringUtils.isBlank(this.workUnitState.getProp(ConfigurationKeys.SOURCE_QUERYBASED_WATERMARK_TYPE))) {
        watermarkType = null;
    } else {
        watermarkType = WatermarkType.valueOf(this.workUnitState.getProp(ConfigurationKeys.SOURCE_QUERYBASED_WATERMARK_TYPE).toUpperCase());
    }
    log.info("Source Entity is " + this.entity);
    try {
        this.setTimeOut(this.workUnitState.getPropAsInt(ConfigurationKeys.SOURCE_CONN_TIMEOUT, ConfigurationKeys.DEFAULT_CONN_TIMEOUT));
        this.extractMetadata(this.schema, this.entity, this.workUnit);
        if (StringUtils.isNotBlank(watermarkColumn)) {
            if (partition.isLastPartition()) {
                // Get a more accurate high watermark from the source
                long adjustedHighWatermark = this.getLatestWatermark(watermarkColumn, watermarkType, lwm, hwm);
                log.info("High water mark from source: " + adjustedHighWatermark);
                // Else, consider the low watermark as high water mark(with no delta).i.e, don't move the pointer
                if (adjustedHighWatermark == ConfigurationKeys.DEFAULT_WATERMARK_VALUE) {
                    adjustedHighWatermark = getLowWatermarkWithNoDelta(lwm);
                }
                this.highWatermark = adjustedHighWatermark;
            } else {
                this.highWatermark = hwm;
            }
            log.info("High water mark for the current run: " + highWatermark);
            this.setRangePredicates(watermarkColumn, watermarkType, lwm, highWatermark);
        }
        // if it is set to true, skip count calculation and set source count to -1
        if (!Boolean.valueOf(this.workUnitState.getProp(ConfigurationKeys.SOURCE_QUERYBASED_SKIP_COUNT_CALC))) {
            this.sourceRecordCount = this.getSourceCount(this.schema, this.entity, this.workUnit, this.predicateList);
        } else {
            log.info("Skip count calculation");
            this.sourceRecordCount = -1;
        }
        if (this.sourceRecordCount == 0) {
            log.info("Record count is 0; Setting fetch status to false to skip readRecord()");
            this.setFetchStatus(false);
        }
    } catch (SchemaException e) {
        throw new ExtractPrepareException("Failed to get schema for this object; error - " + e.getMessage(), e);
    } catch (HighWatermarkException e) {
        throw new ExtractPrepareException("Failed to get high watermark; error - " + e.getMessage(), e);
    } catch (RecordCountException e) {
        throw new ExtractPrepareException("Failed to get record count; error - " + e.getMessage(), e);
    } catch (Exception e) {
        throw new ExtractPrepareException("Failed to prepare the extract build; error - " + e.getMessage(), e);
    }
    return this;
}