Examples with CancelQueryTask com.mysql.cj.CancelQueryTask used on opensource projects

Example 6 with CancelQueryTask

use of com.mysql.cj.CancelQueryTask in project aws-mysql-jdbc by awslabs.

the class ClientPreparedStatement method executePreparedBatchAsMultiStatement.

/**
 * Rewrites the already prepared statement into a multi-statement
 * query of 'statementsPerBatch' values and executes the entire batch
 * using this new statement.
 *
 * @param batchTimeout
 *            timeout for the batch execution
 * @return update counts in the same fashion as executeBatch()
 *
 * @throws SQLException
 *             if a database access error occurs or this method is called on a closed PreparedStatement
 */
protected long[] executePreparedBatchAsMultiStatement(int batchTimeout) throws SQLException {
    synchronized (checkClosed().getConnectionMutex()) {
        // This is kind of an abuse, but it gets the job done
        if (this.batchedValuesClause == null) {
            this.batchedValuesClause = ((PreparedQuery<?>) this.query).getOriginalSql() + ";";
        }
        JdbcConnection locallyScopedConn = this.connection;
        boolean multiQueriesEnabled = locallyScopedConn.getPropertySet().getBooleanProperty(PropertyKey.allowMultiQueries).getValue();
        CancelQueryTask timeoutTask = null;
        try {
            clearWarnings();
            int numBatchedArgs = this.query.getBatchedArgs().size();
            if (this.retrieveGeneratedKeys) {
                this.batchedGeneratedKeys = new ArrayList<>(numBatchedArgs);
            }
            int numValuesPerBatch = ((PreparedQuery<?>) this.query).computeBatchSize(numBatchedArgs);
            if (numBatchedArgs < numValuesPerBatch) {
                numValuesPerBatch = numBatchedArgs;
            }
            java.sql.PreparedStatement batchedStatement = null;
            int batchedParamIndex = 1;
            int numberToExecuteAsMultiValue = 0;
            int batchCounter = 0;
            int updateCountCounter = 0;
            long[] updateCounts = new long[numBatchedArgs * getParseInfo().getNumberOfQueries()];
            SQLException sqlEx = null;
            try {
                if (!multiQueriesEnabled) {
                    ((NativeSession) locallyScopedConn.getSession()).enableMultiQueries();
                }
                batchedStatement = this.retrieveGeneratedKeys ? ((Wrapper) locallyScopedConn.prepareStatement(generateMultiStatementForBatch(numValuesPerBatch), RETURN_GENERATED_KEYS)).unwrap(java.sql.PreparedStatement.class) : ((Wrapper) locallyScopedConn.prepareStatement(generateMultiStatementForBatch(numValuesPerBatch))).unwrap(java.sql.PreparedStatement.class);
                timeoutTask = startQueryTimer((StatementImpl) batchedStatement, batchTimeout);
                numberToExecuteAsMultiValue = numBatchedArgs < numValuesPerBatch ? numBatchedArgs : numBatchedArgs / numValuesPerBatch;
                int numberArgsToExecute = numberToExecuteAsMultiValue * numValuesPerBatch;
                for (int i = 0; i < numberArgsToExecute; i++) {
                    if (i != 0 && i % numValuesPerBatch == 0) {
                        try {
                            batchedStatement.execute();
                        } catch (SQLException ex) {
                            sqlEx = handleExceptionForBatch(batchCounter, numValuesPerBatch, updateCounts, ex);
                        }
                        updateCountCounter = processMultiCountsAndKeys((StatementImpl) batchedStatement, updateCountCounter, updateCounts);
                        batchedStatement.clearParameters();
                        batchedParamIndex = 1;
                    }
                    batchedParamIndex = setOneBatchedParameterSet(batchedStatement, batchedParamIndex, this.query.getBatchedArgs().get(batchCounter++));
                }
                try {
                    batchedStatement.execute();
                } catch (SQLException ex) {
                    sqlEx = handleExceptionForBatch(batchCounter - 1, numValuesPerBatch, updateCounts, ex);
                }
                updateCountCounter = processMultiCountsAndKeys((StatementImpl) batchedStatement, updateCountCounter, updateCounts);
                batchedStatement.clearParameters();
                numValuesPerBatch = numBatchedArgs - batchCounter;
                if (timeoutTask != null) {
                    // we need to check the cancel state now because we loose if after the following batchedStatement.close()
                    ((JdbcPreparedStatement) batchedStatement).checkCancelTimeout();
                }
            } finally {
                if (batchedStatement != null) {
                    batchedStatement.close();
                    batchedStatement = null;
                }
            }
            try {
                if (numValuesPerBatch > 0) {
                    batchedStatement = this.retrieveGeneratedKeys ? locallyScopedConn.prepareStatement(generateMultiStatementForBatch(numValuesPerBatch), RETURN_GENERATED_KEYS) : locallyScopedConn.prepareStatement(generateMultiStatementForBatch(numValuesPerBatch));
                    if (timeoutTask != null) {
                        timeoutTask.setQueryToCancel((Query) batchedStatement);
                    }
                    batchedParamIndex = 1;
                    while (batchCounter < numBatchedArgs) {
                        batchedParamIndex = setOneBatchedParameterSet(batchedStatement, batchedParamIndex, this.query.getBatchedArgs().get(batchCounter++));
                    }
                    try {
                        batchedStatement.execute();
                    } catch (SQLException ex) {
                        sqlEx = handleExceptionForBatch(batchCounter - 1, numValuesPerBatch, updateCounts, ex);
                    }
                    updateCountCounter = processMultiCountsAndKeys((StatementImpl) batchedStatement, updateCountCounter, updateCounts);
                    batchedStatement.clearParameters();
                }
                if (timeoutTask != null) {
                    stopQueryTimer(timeoutTask, true, true);
                    timeoutTask = null;
                }
                if (sqlEx != null) {
                    throw SQLError.createBatchUpdateException(sqlEx, updateCounts, this.exceptionInterceptor);
                }
                return updateCounts;
            } finally {
                if (batchedStatement != null) {
                    batchedStatement.close();
                }
            }
        } finally {
            stopQueryTimer(timeoutTask, false, false);
            resetCancelledState();
            if (!multiQueriesEnabled) {
                ((NativeSession) locallyScopedConn.getSession()).disableMultiQueries();
            }
            clearBatch();
        }
    }
}

Example 7 with CancelQueryTask

use of com.mysql.cj.CancelQueryTask in project aws-mysql-jdbc by awslabs.

the class StatementImpl method executeBatchInternal.

protected long[] executeBatchInternal() throws SQLException {
    JdbcConnection locallyScopedConn = checkClosed();
    synchronized (locallyScopedConn.getConnectionMutex()) {
        if (locallyScopedConn.isReadOnly()) {
            throw SQLError.createSQLException(Messages.getString("Statement.34") + Messages.getString("Statement.35"), MysqlErrorNumbers.SQL_STATE_ILLEGAL_ARGUMENT, getExceptionInterceptor());
        }
        implicitlyCloseAllOpenResults();
        List<Object> batchedArgs = this.query.getBatchedArgs();
        if (batchedArgs == null || batchedArgs.size() == 0) {
            return new long[0];
        }
        // we timeout the entire batch, not individual statements
        int individualStatementTimeout = getTimeoutInMillis();
        setTimeoutInMillis(0);
        CancelQueryTask timeoutTask = null;
        try {
            resetCancelledState();
            statementBegins();
            try {
                // The JDBC spec doesn't forbid this, but doesn't provide for it either...we do..
                this.retrieveGeneratedKeys = true;
                long[] updateCounts = null;
                if (batchedArgs != null) {
                    int nbrCommands = batchedArgs.size();
                    this.batchedGeneratedKeys = new ArrayList<>(batchedArgs.size());
                    boolean multiQueriesEnabled = locallyScopedConn.getPropertySet().getBooleanProperty(PropertyKey.allowMultiQueries).getValue();
                    if (multiQueriesEnabled || (locallyScopedConn.getPropertySet().getBooleanProperty(PropertyKey.rewriteBatchedStatements).getValue() && nbrCommands > 4)) {
                        return executeBatchUsingMultiQueries(multiQueriesEnabled, nbrCommands, individualStatementTimeout);
                    }
                    timeoutTask = startQueryTimer(this, individualStatementTimeout);
                    updateCounts = new long[nbrCommands];
                    for (int i = 0; i < nbrCommands; i++) {
                        updateCounts[i] = -3;
                    }
                    SQLException sqlEx = null;
                    int commandIndex = 0;
                    for (commandIndex = 0; commandIndex < nbrCommands; commandIndex++) {
                        try {
                            String sql = (String) batchedArgs.get(commandIndex);
                            updateCounts[commandIndex] = executeUpdateInternal(sql, true, true);
                            if (timeoutTask != null) {
                                // we need to check the cancel state on each iteration to generate timeout exception if needed
                                checkCancelTimeout();
                            }
                            // limit one generated key per OnDuplicateKey statement
                            getBatchedGeneratedKeys(this.results.getFirstCharOfQuery() == 'I' && containsOnDuplicateKeyInString(sql) ? 1 : 0);
                        } catch (SQLException ex) {
                            updateCounts[commandIndex] = EXECUTE_FAILED;
                            if (this.continueBatchOnError && !(ex instanceof MySQLTimeoutException) && !(ex instanceof MySQLStatementCancelledException) && !hasDeadlockOrTimeoutRolledBackTx(ex)) {
                                sqlEx = ex;
                            } else {
                                long[] newUpdateCounts = new long[commandIndex];
                                if (hasDeadlockOrTimeoutRolledBackTx(ex)) {
                                    for (int i = 0; i < newUpdateCounts.length; i++) {
                                        newUpdateCounts[i] = java.sql.Statement.EXECUTE_FAILED;
                                    }
                                } else {
                                    System.arraycopy(updateCounts, 0, newUpdateCounts, 0, commandIndex);
                                }
                                sqlEx = ex;
                                break;
                            // throw SQLError.createBatchUpdateException(ex, newUpdateCounts, getExceptionInterceptor());
                            }
                        }
                    }
                    if (sqlEx != null) {
                        throw SQLError.createBatchUpdateException(sqlEx, updateCounts, getExceptionInterceptor());
                    }
                }
                if (timeoutTask != null) {
                    stopQueryTimer(timeoutTask, true, true);
                    timeoutTask = null;
                }
                return (updateCounts != null) ? updateCounts : new long[0];
            } finally {
                this.query.getStatementExecuting().set(false);
            }
        } finally {
            stopQueryTimer(timeoutTask, false, false);
            resetCancelledState();
            setTimeoutInMillis(individualStatementTimeout);
            clearBatch();
        }
    }
}

Example 8 with CancelQueryTask

use of com.mysql.cj.CancelQueryTask in project aws-mysql-jdbc by awslabs.

the class StatementImpl method executeBatchUsingMultiQueries.

/**
 * Rewrites batch into a single query to send to the server. This method
 * will constrain each batch to be shorter than max_allowed_packet on the
 * server.
 *
 * @param multiQueriesEnabled
 *            is multi-queries syntax allowed?
 * @param nbrCommands
 *            number of queries in a batch
 * @param individualStatementTimeout
 *            timeout for a single query in a batch
 *
 * @return update counts in the same manner as executeBatch()
 * @throws SQLException
 *             if a database access error occurs or this method is called on a closed PreparedStatement
 */
private long[] executeBatchUsingMultiQueries(boolean multiQueriesEnabled, int nbrCommands, int individualStatementTimeout) throws SQLException {
    JdbcConnection locallyScopedConn = checkClosed();
    synchronized (locallyScopedConn.getConnectionMutex()) {
        if (!multiQueriesEnabled) {
            this.session.enableMultiQueries();
        }
        java.sql.Statement batchStmt = null;
        CancelQueryTask timeoutTask = null;
        try {
            long[] updateCounts = new long[nbrCommands];
            for (int i = 0; i < nbrCommands; i++) {
                updateCounts[i] = JdbcStatement.EXECUTE_FAILED;
            }
            int commandIndex = 0;
            StringBuilder queryBuf = new StringBuilder();
            batchStmt = locallyScopedConn.createStatement();
            JdbcStatement jdbcBatchedStmt = (JdbcStatement) batchStmt;
            getQueryAttributesBindings().runThroughAll(a -> jdbcBatchedStmt.setAttribute(a.getName(), a.getValue()));
            timeoutTask = startQueryTimer((StatementImpl) batchStmt, individualStatementTimeout);
            int counter = 0;
            String connectionEncoding = locallyScopedConn.getPropertySet().getStringProperty(PropertyKey.characterEncoding).getValue();
            int numberOfBytesPerChar = StringUtils.startsWithIgnoreCase(connectionEncoding, "utf") ? 3 : (this.session.getServerSession().getCharsetSettings().isMultibyteCharset(connectionEncoding) ? 2 : 1);
            int escapeAdjust = 1;
            batchStmt.setEscapeProcessing(this.doEscapeProcessing);
            if (this.doEscapeProcessing) {
                // We assume packet _could_ grow by this amount, as we're not sure how big statement will end up after escape processing
                escapeAdjust = 2;
            }
            SQLException sqlEx = null;
            int argumentSetsInBatchSoFar = 0;
            for (commandIndex = 0; commandIndex < nbrCommands; commandIndex++) {
                String nextQuery = (String) this.query.getBatchedArgs().get(commandIndex);
                if (((((queryBuf.length() + nextQuery.length()) * numberOfBytesPerChar) + 1 + /* for semicolon */
                NativeConstants.HEADER_LENGTH) * escapeAdjust) + 32 > this.maxAllowedPacket.getValue()) {
                    try {
                        batchStmt.execute(queryBuf.toString(), java.sql.Statement.RETURN_GENERATED_KEYS);
                    } catch (SQLException ex) {
                        sqlEx = handleExceptionForBatch(commandIndex, argumentSetsInBatchSoFar, updateCounts, ex);
                    }
                    counter = processMultiCountsAndKeys((StatementImpl) batchStmt, counter, updateCounts);
                    queryBuf = new StringBuilder();
                    argumentSetsInBatchSoFar = 0;
                }
                queryBuf.append(nextQuery);
                queryBuf.append(";");
                argumentSetsInBatchSoFar++;
            }
            if (queryBuf.length() > 0) {
                try {
                    batchStmt.execute(queryBuf.toString(), java.sql.Statement.RETURN_GENERATED_KEYS);
                } catch (SQLException ex) {
                    sqlEx = handleExceptionForBatch(commandIndex - 1, argumentSetsInBatchSoFar, updateCounts, ex);
                }
                counter = processMultiCountsAndKeys((StatementImpl) batchStmt, counter, updateCounts);
            }
            if (timeoutTask != null) {
                stopQueryTimer(timeoutTask, true, true);
                timeoutTask = null;
            }
            if (sqlEx != null) {
                throw SQLError.createBatchUpdateException(sqlEx, updateCounts, getExceptionInterceptor());
            }
            return (updateCounts != null) ? updateCounts : new long[0];
        } finally {
            stopQueryTimer(timeoutTask, false, false);
            resetCancelledState();
            try {
                if (batchStmt != null) {
                    batchStmt.close();
                }
            } finally {
                if (!multiQueriesEnabled) {
                    this.session.disableMultiQueries();
                }
            }
        }
    }
}

Example 9 with CancelQueryTask

use of com.mysql.cj.CancelQueryTask in project aws-mysql-jdbc by awslabs.

the class StatementImpl method executeUpdateInternal.

protected long executeUpdateInternal(String sql, boolean isBatch, boolean returnGeneratedKeys) throws SQLException {
    synchronized (checkClosed().getConnectionMutex()) {
        JdbcConnection locallyScopedConn = this.connection;
        checkNullOrEmptyQuery(sql);
        resetCancelledState();
        char firstStatementChar = ParseInfo.firstCharOfStatementUc(sql, this.session.getServerSession().isNoBackslashEscapesSet());
        if (!isNonResultSetProducingQuery(sql)) {
            throw SQLError.createSQLException(Messages.getString("Statement.46"), "01S03", getExceptionInterceptor());
        }
        this.retrieveGeneratedKeys = returnGeneratedKeys;
        this.lastQueryIsOnDupKeyUpdate = returnGeneratedKeys && firstStatementChar == 'I' && containsOnDuplicateKeyInString(sql);
        ResultSetInternalMethods rs = null;
        if (this.doEscapeProcessing) {
            Object escapedSqlResult = EscapeProcessor.escapeSQL(sql, this.session.getServerSession().getSessionTimeZone(), this.session.getServerSession().getCapabilities().serverSupportsFracSecs(), this.session.getServerSession().isServerTruncatesFracSecs(), getExceptionInterceptor());
            sql = escapedSqlResult instanceof String ? (String) escapedSqlResult : ((EscapeProcessorResult) escapedSqlResult).escapedSql;
        }
        if (locallyScopedConn.isReadOnly(false)) {
            throw SQLError.createSQLException(Messages.getString("Statement.42") + Messages.getString("Statement.43"), MysqlErrorNumbers.SQL_STATE_ILLEGAL_ARGUMENT, getExceptionInterceptor());
        }
        implicitlyCloseAllOpenResults();
        // The checking and changing of databases must happen in sequence, so synchronize on the same mutex that _conn is using
        CancelQueryTask timeoutTask = null;
        String oldDb = null;
        try {
            timeoutTask = startQueryTimer(this, getTimeoutInMillis());
            if (!locallyScopedConn.getDatabase().equals(getCurrentDatabase())) {
                oldDb = locallyScopedConn.getDatabase();
                locallyScopedConn.setDatabase(getCurrentDatabase());
            }
            // 
            // Only apply max_rows to selects
            // 
            locallyScopedConn.setSessionMaxRows(-1);
            statementBegins();
            // null database: force read of field info on DML
            rs = ((NativeSession) locallyScopedConn.getSession()).execSQL(this, sql, -1, null, false, getResultSetFactory(), null, isBatch);
            if (timeoutTask != null) {
                stopQueryTimer(timeoutTask, true, true);
                timeoutTask = null;
            }
        } catch (CJTimeoutException | OperationCancelledException e) {
            throw SQLExceptionsMapping.translateException(e, this.exceptionInterceptor);
        } finally {
            stopQueryTimer(timeoutTask, false, false);
            if (oldDb != null) {
                locallyScopedConn.setDatabase(oldDb);
            }
            if (!isBatch) {
                this.query.getStatementExecuting().set(false);
            }
        }
        this.results = rs;
        rs.setFirstCharOfQuery(firstStatementChar);
        this.updateCount = rs.getUpdateCount();
        this.lastInsertId = rs.getUpdateID();
        return this.updateCount;
    }
}

Example 10 with CancelQueryTask

use of com.mysql.cj.CancelQueryTask in project aws-mysql-jdbc by awslabs.

the class StatementImpl method executeInternal.

private boolean executeInternal(String sql, boolean returnGeneratedKeys) throws SQLException {
    JdbcConnection locallyScopedConn = checkClosed();
    synchronized (locallyScopedConn.getConnectionMutex()) {
        checkClosed();
        checkNullOrEmptyQuery(sql);
        resetCancelledState();
        implicitlyCloseAllOpenResults();
        if (sql.charAt(0) == '/') {
            if (sql.startsWith(PING_MARKER)) {
                doPingInstead();
                return true;
            }
        }
        this.retrieveGeneratedKeys = returnGeneratedKeys;
        this.lastQueryIsOnDupKeyUpdate = returnGeneratedKeys && ParseInfo.firstCharOfStatementUc(sql, this.session.getServerSession().isNoBackslashEscapesSet()) == 'I' && containsOnDuplicateKeyInString(sql);
        if (!ParseInfo.isReadOnlySafeQuery(sql, this.session.getServerSession().isNoBackslashEscapesSet()) && locallyScopedConn.isReadOnly()) {
            throw SQLError.createSQLException(Messages.getString("Statement.27") + Messages.getString("Statement.28"), MysqlErrorNumbers.SQL_STATE_ILLEGAL_ARGUMENT, getExceptionInterceptor());
        }
        try {
            setupStreamingTimeout(locallyScopedConn);
            if (this.doEscapeProcessing) {
                Object escapedSqlResult = EscapeProcessor.escapeSQL(sql, this.session.getServerSession().getSessionTimeZone(), this.session.getServerSession().getCapabilities().serverSupportsFracSecs(), this.session.getServerSession().isServerTruncatesFracSecs(), getExceptionInterceptor());
                sql = escapedSqlResult instanceof String ? (String) escapedSqlResult : ((EscapeProcessorResult) escapedSqlResult).escapedSql;
            }
            CachedResultSetMetaData cachedMetaData = null;
            ResultSetInternalMethods rs = null;
            this.batchedGeneratedKeys = null;
            if (useServerFetch()) {
                rs = createResultSetUsingServerFetch(sql);
            } else {
                CancelQueryTask timeoutTask = null;
                String oldDb = null;
                try {
                    timeoutTask = startQueryTimer(this, getTimeoutInMillis());
                    if (!locallyScopedConn.getDatabase().equals(getCurrentDatabase())) {
                        oldDb = locallyScopedConn.getDatabase();
                        locallyScopedConn.setDatabase(getCurrentDatabase());
                    }
                    // Check if we have cached metadata for this query...
                    if (locallyScopedConn.getPropertySet().getBooleanProperty(PropertyKey.cacheResultSetMetadata).getValue()) {
                        cachedMetaData = locallyScopedConn.getCachedMetaData(sql);
                    }
                    // Only apply max_rows to selects
                    locallyScopedConn.setSessionMaxRows(isResultSetProducingQuery(sql) ? this.maxRows : -1);
                    statementBegins();
                    rs = ((NativeSession) locallyScopedConn.getSession()).execSQL(this, sql, this.maxRows, null, createStreamingResultSet(), getResultSetFactory(), cachedMetaData, false);
                    if (timeoutTask != null) {
                        stopQueryTimer(timeoutTask, true, true);
                        timeoutTask = null;
                    }
                } catch (CJTimeoutException | OperationCancelledException e) {
                    throw SQLExceptionsMapping.translateException(e, this.exceptionInterceptor);
                } finally {
                    stopQueryTimer(timeoutTask, false, false);
                    if (oldDb != null) {
                        locallyScopedConn.setDatabase(oldDb);
                    }
                }
            }
            if (rs != null) {
                this.lastInsertId = rs.getUpdateID();
                this.results = rs;
                rs.setFirstCharOfQuery(ParseInfo.firstCharOfStatementUc(sql, this.session.getServerSession().isNoBackslashEscapesSet()));
                if (rs.hasRows()) {
                    if (cachedMetaData != null) {
                        locallyScopedConn.initializeResultsMetadataFromCache(sql, cachedMetaData, this.results);
                    } else if (this.session.getPropertySet().getBooleanProperty(PropertyKey.cacheResultSetMetadata).getValue()) {
                        locallyScopedConn.initializeResultsMetadataFromCache(sql, null, /* will be created */
                        this.results);
                    }
                }
            }
            return ((rs != null) && rs.hasRows());
        } finally {
            this.query.getStatementExecuting().set(false);
        }
    }
}