Search in sources :

Example 1 with DrillStatement

use of org.apache.drill.jdbc.DrillStatement in project drill by apache.

the class DrillCursor method nextRowInternally.

/**
 * ...
 * <p>
 *   Is to be called (once) from {@link #loadInitialSchema} for
 *   {@link DrillResultSetImpl#execute()}, and then (repeatedly) from
 *   {@link #next()} for {@link org.apache.calcite.avatica.AvaticaResultSet#next()}.
 * </p>
 *
 * @return  whether cursor is positioned at a row (false when after end of
 *   results)
 */
private boolean nextRowInternally() throws SQLException {
    if (currentRecordNumber + 1 < currentBatchHolder.getRecordCount()) {
        // Have next row in current batch--just advance index and report "at a row."
        currentRecordNumber++;
        return true;
    } else {
        try {
            QueryDataBatch qrb = resultsListener.getNext();
            // the (initial) schema but no rows)).
            if (afterFirstBatch) {
                while (qrb != null && (qrb.getHeader().getRowCount() == 0 && qrb.getData() == null)) {
                    // Empty message--dispose of and try to get another.
                    logger.warn("Spurious batch read: {}", qrb);
                    qrb.release();
                    qrb = resultsListener.getNext();
                }
            }
            afterFirstBatch = true;
            if (qrb == null) {
                // End of batches--clean up, set state to done, report after last row.
                // (We load it so we clear it.)
                currentBatchHolder.clear();
                afterLastRow = true;
                return false;
            } else {
                // Got next (or first) batch--reset record offset to beginning;
                // assimilate schema if changed; set up return value for first call
                // to next().
                currentRecordNumber = 0;
                if (qrb.getHeader().hasAffectedRowsCount()) {
                    int updateCount = qrb.getHeader().getAffectedRowsCount();
                    int currentUpdateCount = statement.getUpdateCount() == -1 ? 0 : statement.getUpdateCount();
                    ((DrillStatement) statement).setUpdateCount(updateCount + currentUpdateCount);
                    ((DrillStatement) statement).setResultSet(null);
                }
                final boolean schemaChanged;
                try {
                    schemaChanged = currentBatchHolder.load(qrb.getHeader().getDef(), qrb.getData());
                } finally {
                    qrb.release();
                }
                schema = currentBatchHolder.getSchema();
                if (schemaChanged) {
                    updateColumns();
                }
                if (returnTrueForNextCallToNext && currentBatchHolder.getRecordCount() == 0) {
                    returnTrueForNextCallToNext = false;
                }
                return true;
            }
        } catch (UserException e) {
            // error type is accessible, of course. :-()
            throw new SQLException(e.getMessage(), e);
        } catch (InterruptedException e) {
            // but JDBC client certainly could.
            throw new SQLException("Interrupted.", e);
        } catch (RuntimeException e) {
            throw new SQLException("Unexpected RuntimeException: " + e.toString(), e);
        }
    }
}
Also used : QueryDataBatch(org.apache.drill.exec.rpc.user.QueryDataBatch) SQLException(java.sql.SQLException) DrillStatement(org.apache.drill.jdbc.DrillStatement) UserException(org.apache.drill.common.exceptions.UserException)

Aggregations

SQLException (java.sql.SQLException)1 UserException (org.apache.drill.common.exceptions.UserException)1 QueryDataBatch (org.apache.drill.exec.rpc.user.QueryDataBatch)1 DrillStatement (org.apache.drill.jdbc.DrillStatement)1