CassandraStatement.java

/*
 *
 *   Licensed 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 com.ing.data.cassandra.jdbc;

import com.datastax.oss.driver.api.core.ConsistencyLevel;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.cql.AsyncResultSet;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.internal.core.cql.MultiPageResultSet;
import com.datastax.oss.driver.internal.core.cql.SinglePageResultSet;
import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
import com.google.common.collect.Lists;
import edu.umd.cs.findbugs.annotations.NonNull;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.CallableStatement;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.SQLFeatureNotSupportedException;
import java.sql.SQLNonTransientException;
import java.sql.SQLRecoverableException;
import java.sql.SQLSyntaxErrorException;
import java.sql.SQLTransientException;
import java.sql.SQLWarning;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletionStage;

/**
 * Cassandra statement: implementation class for {@link Statement}.
 * <p>
 *     It also implements {@link CassandraStatementExtras} interface providing extra methods not defined in JDBC API to
 *     manage some properties specific to the Cassandra statements (e.g. consistency level).
 * </p>
 */
public class CassandraStatement extends AbstractStatement
    implements CassandraStatementExtras, Comparable<Object>, Statement {

    /**
     * Maximal number of queries executable in a single batch.
     */
    public static final int MAX_ASYNC_QUERIES = 1000;
    /**
     * CQL statements separator: semi-colon ({@code ;}).
     */
    public static final String STATEMENTS_SEPARATOR_REGEX = ";";
    /**
     * The default fetch size.
     */
    protected static final int DEFAULT_FETCH_SIZE = 100;

    private static final Logger LOG = LoggerFactory.getLogger(CassandraStatement.class);

    /**
     * The Cassandra connection.
     */
    protected CassandraConnection connection;
    /**
     * The CQL statement.
     */
    protected String cql;
    /**
     * The list of CQL queries contained into a single batch.
     */
    protected ArrayList<String> batchQueries;
    /**
     * The direction for fetching rows from database. By default: {@link ResultSet#FETCH_FORWARD}.
     */
    protected int fetchDirection = ResultSet.FETCH_FORWARD;
    /**
     * The number of result set rows that is the default fetch size for {@link ResultSet} objects generated from this
     * statement. By default: {@value #DEFAULT_FETCH_SIZE}.
     */
    protected int fetchSize = DEFAULT_FETCH_SIZE;
    /**
     * The maximum number of bytes that can be returned for character and binary column values in a {@link ResultSet}
     * object produced by this statement. By default, there is no limit (0).
     */
    protected int maxFieldSize = 0;
    /**
     * The maximum number of rows that a {@link ResultSet} object produced by this statement. By default, there is no
     * limit (0).
     */
    protected int maxRows = 0;
    /**
     * The result set type for {@link ResultSet} objects generated by this statement.
     */
    protected int resultSetType;
    /**
     * The result set concurrency for {@link ResultSet} objects generated by this statement.
     */
    protected int resultSetConcurrency;
    /**
     * The result set holdability for {@link ResultSet} objects generated by this statement.
     */
    protected int resultSetHoldability;
    /**
     * The current result set for this statement.
     */
    protected ResultSet currentResultSet = null;
    /**
     * The update count.
     */
    protected int updateCount = -1;
    /**
     * Whether the escape processing is on or off. By default, it is on, even Cassandra implementation currently does
     * not take it into account.
     */
    protected boolean escapeProcessing = true;
    /**
     * The Datastax Java driver statement.
     */
    protected com.datastax.oss.driver.api.core.cql.Statement<?> statement;
    /**
     * The consistency level used for the statement.
     */
    protected ConsistencyLevel consistencyLevel;

    /**
     * Constructor. It instantiates a new Cassandra statement with default values and a {@code null} CQL statement for
     * a {@link CassandraConnection}.
     * <p>
     *     By default, the result set type is {@link ResultSet#TYPE_FORWARD_ONLY}, the result set concurrency is
     *     {@link ResultSet#CONCUR_READ_ONLY} and the result set holdability is
     *     {@link ResultSet#HOLD_CURSORS_OVER_COMMIT}.
     * </p>
     *
     * @param connection The Cassandra connection to the database.
     * @throws SQLException when something went wrong during the instantiation of the statement.
     */
    CassandraStatement(final CassandraConnection connection) throws SQLException {
        this(connection, null, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY,
            ResultSet.HOLD_CURSORS_OVER_COMMIT);
    }

    /**
     * Constructor. It instantiates a new Cassandra statement with default values for a {@link CassandraConnection}.
     * <p>
     *     By default, the result set type is {@link ResultSet#TYPE_FORWARD_ONLY}, the result set concurrency is
     *     {@link ResultSet#CONCUR_READ_ONLY} and the result set holdability is
     *     {@link ResultSet#HOLD_CURSORS_OVER_COMMIT}.
     * </p>
     *
     * @param connection The Cassandra connection to the database.
     * @param cql        The CQL statement.
     * @throws SQLException when something went wrong during the instantiation of the statement.
     */
    CassandraStatement(final CassandraConnection connection, final String cql) throws SQLException {
        this(connection, cql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY,
            ResultSet.HOLD_CURSORS_OVER_COMMIT);
    }

    /**
     * Constructor. It instantiates a new Cassandra statement with default holdability and specified result set type
     * and concurrency for a {@link CassandraConnection}.
     * <p>
     *     By default, the result set holdability is {@link ResultSet#HOLD_CURSORS_OVER_COMMIT}.
     * </p>
     *
     * @param connection            The Cassandra connection to the database.
     * @param cql                   The CQL statement.
     * @param resultSetType         The result set type.
     * @param resultSetConcurrency  The result set concurrency.
     * @throws SQLException when something went wrong during the instantiation of the statement.
     */
    CassandraStatement(final CassandraConnection connection, final String cql, final int resultSetType,
                       final int resultSetConcurrency) throws SQLException {
        this(connection, cql, resultSetType, resultSetConcurrency, ResultSet.HOLD_CURSORS_OVER_COMMIT);
    }

    /**
     * Constructor. It instantiates a new Cassandra statement with specified result set type, concurrency and
     * holdability for a {@link CassandraConnection}.
     *
     * @param connection            The Cassandra connection to the database.
     * @param cql                   The CQL statement.
     * @param resultSetType         The result set type.
     * @param resultSetConcurrency  The result set concurrency.
     * @param resultSetHoldability  The result set holdability.
     * @throws SQLException when something went wrong during the instantiation of the statement.
     */
    CassandraStatement(final CassandraConnection connection, final String cql, final int resultSetType,
                       final int resultSetConcurrency, final int resultSetHoldability) throws SQLException {
        this.connection = connection;
        this.cql = cql;
        this.batchQueries = Lists.newArrayList();
        this.consistencyLevel = connection.getDefaultConsistencyLevel();

        if (!(resultSetType == ResultSet.TYPE_FORWARD_ONLY
            || resultSetType == ResultSet.TYPE_SCROLL_INSENSITIVE
            || resultSetType == ResultSet.TYPE_SCROLL_SENSITIVE)) {
            throw new SQLSyntaxErrorException(Utils.BAD_TYPE_RS);
        }
        this.resultSetType = resultSetType;

        if (!(resultSetConcurrency == ResultSet.CONCUR_READ_ONLY
            || resultSetConcurrency == ResultSet.CONCUR_UPDATABLE)) {
            throw new SQLSyntaxErrorException(Utils.BAD_CONCURRENCY_RS);
        }
        this.resultSetConcurrency = resultSetConcurrency;

        if (!(resultSetHoldability == ResultSet.HOLD_CURSORS_OVER_COMMIT
            || resultSetHoldability == ResultSet.CLOSE_CURSORS_AT_COMMIT)) {
            throw new SQLSyntaxErrorException(Utils.BAD_HOLD_RS);
        }
        this.resultSetHoldability = resultSetHoldability;
    }

    @Override
    public void addBatch(final String query) throws SQLException {
        checkNotClosed();
        batchQueries.add(query);
    }

    /**
     * Checks that the statement is not closed.
     *
     * @throws SQLException when something went wrong during the checking of the statement status.
     */
    protected final void checkNotClosed() throws SQLException {
        if (isClosed()) {
            throw new SQLRecoverableException(Utils.WAS_CLOSED_STMT);
        }
    }

    @Override
    public void clearBatch() throws SQLException {
        checkNotClosed();
        this.batchQueries = new ArrayList<>();
    }

    @Override
    public void clearWarnings() throws SQLException {
        // This implementation does not support the collection of warnings so clearing is a no-op but it still throws
        // an exception when called on a closed statement.
        checkNotClosed();
    }

    @Override
    public void close() {
        this.connection = null;
        this.cql = null;
    }

    @Override
    public int compareTo(@NonNull final Object target) {
        if (this.equals(target)) {
            return 0;
        }
        if (this.hashCode() < target.hashCode()) {
            return -1;
        }
        return 1;
    }

    private void doExecute(final String cql) throws SQLException {
        final List<CompletionStage<AsyncResultSet>> futures = new ArrayList<>();

        try {
            final String[] cqlQueries = cql.split(STATEMENTS_SEPARATOR_REGEX);
            if (cqlQueries.length > 1 && !(cql.trim().toLowerCase().startsWith("begin")
                && cql.toLowerCase().contains("batch") && cql.toLowerCase().contains("apply"))) {
                // Several statements in the query to execute asynchronously...

                final ArrayList<com.datastax.oss.driver.api.core.cql.ResultSet> results = Lists.newArrayList();
                if (cqlQueries.length > MAX_ASYNC_QUERIES * 1.1) {
                    // Protect the cluster from receiving too many queries at once and force the dev to split the load
                    throw new SQLNonTransientException("Too many queries at once (" + cqlQueries.length
                        + "). You must split your queries into more batches !");
                }

                StringBuilder prevCqlQuery = new StringBuilder();
                for (final String cqlQuery : cqlQueries) {
                    if ((cqlQuery.contains("'") && ((StringUtils.countMatches(cqlQuery, "'") % 2 == 1
                        && prevCqlQuery.length() == 0)
                        || (StringUtils.countMatches(cqlQuery, "'") % 2 == 0 && prevCqlQuery.length() > 0)))
                        || (prevCqlQuery.toString().length() > 0 && !cqlQuery.contains("'"))) {
                        prevCqlQuery.append(cqlQuery).append(";");
                    } else {
                        prevCqlQuery.append(cqlQuery);
                        if (LOG.isTraceEnabled() || this.connection.isDebugMode()) {
                            LOG.debug("CQL: {}", prevCqlQuery);
                        }
                        final SimpleStatement stmt = SimpleStatement.newInstance(prevCqlQuery.toString())
                            .setConsistencyLevel(this.connection.getDefaultConsistencyLevel())
                            .setPageSize(this.fetchSize);
                        final CompletionStage<AsyncResultSet> resultSetFuture =
                            ((CqlSession) this.connection.getSession()).executeAsync(stmt);
                        futures.add(resultSetFuture);
                        prevCqlQuery = new StringBuilder();
                    }
                }

                for (final CompletionStage<AsyncResultSet> future : futures) {
                    final AsyncResultSet asyncResultSet = CompletableFutures.getUninterruptibly(future);
                    final com.datastax.oss.driver.api.core.cql.ResultSet rows;
                    if (asyncResultSet.hasMorePages()) {
                        rows = new MultiPageResultSet(asyncResultSet);
                    } else {
                        rows = new SinglePageResultSet(asyncResultSet);
                    }
                    results.add(rows);
                }

                this.currentResultSet = new CassandraResultSet(this, results);
            } else {
                // Only one statement to execute, so do it synchronously.
                if (LOG.isTraceEnabled() || this.connection.isDebugMode()) {
                    LOG.debug("CQL: " + cql);
                }
                final SimpleStatement stmt = SimpleStatement.newInstance(cql)
                    .setConsistencyLevel(this.connection.getDefaultConsistencyLevel())
                    .setPageSize(this.fetchSize);
                this.currentResultSet = new CassandraResultSet(this,
                    ((CqlSession) this.connection.getSession()).execute(stmt));
            }
        } catch (final Exception e) {
            for (final CompletionStage<AsyncResultSet> future : futures) {
                future.toCompletableFuture().cancel(true);
            }
            throw new SQLTransientException(e);
        }
    }

    @Override
    public boolean execute(final String query) throws SQLException {
        checkNotClosed();
        doExecute(query);
        // Return true if the first result is a non-null ResultSet object; false if the first result is an update count
        // or there is no result.
        return this.currentResultSet != null && ((CassandraResultSet) this.currentResultSet).hasMoreRows();
    }

    @Override
    public boolean execute(final String cql, final int autoGeneratedKeys) throws SQLException {
        checkNotClosed();
        if (!(autoGeneratedKeys == RETURN_GENERATED_KEYS || autoGeneratedKeys == NO_GENERATED_KEYS)) {
            throw new SQLSyntaxErrorException(Utils.BAD_AUTO_GEN);
        }
        if (autoGeneratedKeys == RETURN_GENERATED_KEYS) {
            throw new SQLFeatureNotSupportedException(Utils.NO_GEN_KEYS);
        }
        return execute(cql);
    }

    @Override
    public int[] executeBatch() throws SQLException {
        final int[] returnCounts = new int[this.batchQueries.size()];
        final List<CompletionStage<AsyncResultSet>> futures = new ArrayList<>();
        if (LOG.isTraceEnabled() || this.connection.isDebugMode()) {
            LOG.debug("CQL statements: " + this.batchQueries.size());
        }

        for (final String query : this.batchQueries) {
            if (LOG.isTraceEnabled() || this.connection.isDebugMode()) {
                LOG.debug("CQL: " + query);
            }
            final SimpleStatement stmt = SimpleStatement.newInstance(query)
                .setConsistencyLevel(this.connection.getDefaultConsistencyLevel());
            final CompletionStage<AsyncResultSet> resultSetFuture =
                ((CqlSession) this.connection.getSession()).executeAsync(stmt);
            futures.add(resultSetFuture);
        }

        int i = 0;
        for (final CompletionStage<AsyncResultSet> future : futures) {
            CompletableFutures.getUninterruptibly(future);
            returnCounts[i] = 1;
            i++;
        }

        return returnCounts;
    }

    @Override
    public ResultSet executeQuery(final String cql) throws SQLException {
        checkNotClosed();
        doExecute(cql);
        if (this.currentResultSet == null) {
            throw new SQLNonTransientException(Utils.NO_RESULT_SET);
        }
        return currentResultSet;
    }

    /**
     * Executes the given CQL statement, which may be an {@code INSERT}, {@code UPDATE}, or {@code DELETE} statement or
     * a CQL statement that returns nothing, such as a CQL DDL statement.
     * <p>
     *     <b>Note:</b> This method cannot be called on a {@link PreparedStatement} or {@link CallableStatement}.
     * </p>
     *
     * @param cql A CQL Data Manipulation Language (DML) statement, such as {@code INSERT}, {@code UPDATE}, or
     * {@code DELETE}; or a CQL statement that returns nothing, such as a DDL statement.
     * @return Always 0, for any statement. The rationale is that Datastax Java driver does not provide update count.
     * @throws SQLException when something went wrong during the execution of the statement.
     */
    @Override
    public int executeUpdate(final String cql) throws SQLException {
        checkNotClosed();
        doExecute(cql);
        return connection.getOptionSet().getSQLUpdateResponse();
    }

    @Override
    public int executeUpdate(final String cql, final int autoGeneratedKeys) throws SQLException {
        checkNotClosed();
        if (!(autoGeneratedKeys == RETURN_GENERATED_KEYS || autoGeneratedKeys == NO_GENERATED_KEYS)) {
            throw new SQLFeatureNotSupportedException(Utils.BAD_AUTO_GEN);
        }
        return executeUpdate(cql);
    }

    @Override
    public Connection getConnection() throws SQLException {
        checkNotClosed();
        return this.connection;
    }

    @Override
    public ConsistencyLevel getConsistencyLevel() {
        return this.consistencyLevel;
    }

    @Override
    public void setConsistencyLevel(final ConsistencyLevel consistencyLevel) {
        this.consistencyLevel = consistencyLevel;
        this.statement = this.statement.setConsistencyLevel(consistencyLevel);
    }

    @Override
    public int getFetchDirection() throws SQLException {
        checkNotClosed();
        return this.fetchDirection;
    }

    @Override
    public void setFetchDirection(final int direction) throws SQLException {
        checkNotClosed();
        if (direction == ResultSet.FETCH_FORWARD || direction == ResultSet.FETCH_REVERSE
            || direction == ResultSet.FETCH_UNKNOWN) {
            if (getResultSetType() == ResultSet.TYPE_FORWARD_ONLY && direction != ResultSet.FETCH_FORWARD) {
                throw new SQLSyntaxErrorException(String.format(Utils.BAD_FETCH_DIR, direction));
            }
            this.fetchDirection = direction;
        } else {
            throw new SQLSyntaxErrorException(String.format(Utils.BAD_FETCH_DIR, direction));
        }
    }

    @Override
    public int getFetchSize() throws SQLException {
        checkNotClosed();
        return this.fetchSize;
    }

    @Override
    public void setFetchSize(final int rows) throws SQLException {
        checkNotClosed();
        if (rows < 0) {
            throw new SQLSyntaxErrorException(String.format(Utils.BAD_FETCH_SIZE, rows));
        }
        this.fetchSize = rows;
    }

    @Override
    public int getMaxFieldSize() throws SQLException {
        checkNotClosed();
        return this.maxFieldSize;
    }

    /**
     * Sets the limit for the maximum number of bytes that can be returned for character and binary column values in a
     * {@link ResultSet} object produced by this {@link Statement} object.
     * <p>
     *     This setting is silently ignored. There is no such limit, so {@link #maxFieldSize} is always 0.
     * </p>
     *
     * @param max The new column size limit in bytes; zero means there is no limit.
     * @throws SQLException if a database access error occurs or this method is called on a closed {@link Statement}.
     */
    @Override
    public void setMaxFieldSize(final int max) throws SQLException {
        checkNotClosed();
    }

    @Override
    public int getMaxRows() throws SQLException {
        checkNotClosed();
        return this.maxRows;
    }

    /**
     * Sets the limit for the maximum number of rows that any {@link ResultSet} object generated by this
     * {@link Statement} object can contain to the given number. If the limit is exceeded, the excess rows are silently
     * dropped.
     * <p>
     *     This setting is silently ignored. There is no such limit, so {@link #maxRows} is always 0.
     * </p>
     *
     * @param max The new max rows limit; zero means there is no limit.
     * @throws SQLException if a database access error occurs or this method is called on a closed {@link Statement}.
     */
    @Override
    public void setMaxRows(final int max) throws SQLException {
        checkNotClosed();
    }

    @Override
    public boolean getMoreResults() throws SQLException {
        checkNotClosed();
        resetResults();
        // in the current Cassandra implementation there are never more results
        return false;
    }

    @Override
    public boolean getMoreResults(final int current) throws SQLException {
        checkNotClosed();

        switch (current) {
            case CLOSE_CURRENT_RESULT:
                resetResults();
                break;
            case CLOSE_ALL_RESULTS:
            case KEEP_CURRENT_RESULT:
                throw new SQLFeatureNotSupportedException(Utils.NO_MULTIPLE);
            default:
                throw new SQLSyntaxErrorException(String.format(Utils.BAD_KEEP_RS, current));
        }
        // In the current Cassandra implementation there are never more results.
        return false;
    }

    /**
     * Retrieves the number of seconds the driver will wait for a {@link Statement} object to execute. If the limit is
     * exceeded, a {@link SQLException} is thrown.
     * <p>
     *     The Cassandra implementation does not support timeouts on queries, so the query timeout is always 0.
     * </p>
     *
     * @return The current query timeout limit in seconds; zero means there is no limit.
     * @throws SQLException if a database access error occurs or this method is called on a closed {@link Statement}.
     */
    @Override
    public int getQueryTimeout() throws SQLException {
        return 0;
    }

    /**
     * Sets the number of seconds the driver will wait for a {@link Statement} object to execute to the given number of
     * seconds.
     * <p>
     *     This setting is silently ignored. The Cassandra implementation does not support timeouts on queries, so the
     *     query timeout is always 0.
     * </p>
     *
     * @param seconds The new query timeout limit in seconds; zero means there is no limit.
     * @throws SQLException if a database access error occurs, this method is called on a closed {@link Statement}.
     */
    @Override
    public void setQueryTimeout(final int seconds) throws SQLException {
        checkNotClosed();
    }

    @Override
    public ResultSet getResultSet() throws SQLException {
        checkNotClosed();
        return this.currentResultSet;
    }

    @Override
    public int getResultSetConcurrency() throws SQLException {
        checkNotClosed();
        return ResultSet.CONCUR_READ_ONLY;
    }

    @Override
    public int getResultSetHoldability() throws SQLException {
        checkNotClosed();
        // The Cassandra implementation does not support commits, so this is the closest match.
        return ResultSet.HOLD_CURSORS_OVER_COMMIT;
    }

    @Override
    public int getResultSetType() throws SQLException {
        checkNotClosed();
        return ResultSet.TYPE_FORWARD_ONLY;
    }

    @Override
    public int getUpdateCount() throws SQLException {
        checkNotClosed();
        return this.updateCount;
    }

    @Override
    public SQLWarning getWarnings() throws SQLException {
        // The rationale is there are no warnings to return in this implementation but it still throws an exception
        // when called on a closed result set.
        checkNotClosed();
        return null;
    }

    @Override
    public boolean isClosed() {
        return this.connection == null;
    }

    /**
     * Returns a value indicating whether the {@link Statement} is poolable or not.
     *
     * @return Always {@code false}.
     * @throws SQLException if this method is called on a closed {@link Statement}.
     */
    @Override
    public boolean isPoolable() throws SQLException {
        checkNotClosed();
        return false;
    }

    /**
     * Requests that a {@link Statement} be pooled or not pooled.
     * <p>
     *     This setting is silently ignored. The {@code CassandraStatement} are never poolable.
     * </p>
     *
     * @param poolable Requests that the statement be pooled if {@code true} and that the statement not be pooled if
     *                 {@code false}.
     * @throws SQLException if this method is called on a closed {@link Statement}.
     */
    @Override
    public void setPoolable(final boolean poolable) throws SQLException {
        checkNotClosed();
    }

    @Override
    public boolean isWrapperFor(final Class<?> iface) throws SQLException {
        return false;
    }

    /**
     * Resets the current result set for this statement.
     */
    protected final void resetResults() {
        this.currentResultSet = null;
        this.updateCount = -1;
    }

    @Override
    public void setEscapeProcessing(final boolean enable) throws SQLException {
        checkNotClosed();
        // The Cassandra implementation does not currently take this into account.
        this.escapeProcessing = enable;
    }

    @Override
    public <T> T unwrap(final Class<T> iface) throws SQLException {
        if (iface.isInstance(this)) {
            return iface.cast(this);
        }
        throw new SQLFeatureNotSupportedException(String.format(Utils.NO_INTERFACE, iface.getSimpleName()));
    }

}