diff --git a/src/main/java/com/databricks/jdbc/api/impl/BatchParameterSet.java b/src/main/java/com/databricks/jdbc/api/impl/BatchParameterSet.java new file mode 100644 index 000000000..6edc670db --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/BatchParameterSet.java @@ -0,0 +1,88 @@ +package com.databricks.jdbc.api.impl; + +import java.sql.Date; +import java.sql.Time; +import java.sql.Timestamp; +import java.util.Collections; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.stream.Collectors; + +/** + * Immutable, position-ordered snapshot of one prepared-statement parameter set. + * + *

This model preserves JDBC's one-based parameter indexes. Transport adapters are responsible + * for converting them to protocol-specific wire ordinals. It does not validate parameter + * completeness, index continuity, or consistency with other parameter sets; those validations + * remain the backend's responsibility. + */ +public final class BatchParameterSet { + + private final List parameters; + private final Map parameterBindings; + + private BatchParameterSet(List parameters) { + this.parameters = List.copyOf(parameters); + Map bindings = new LinkedHashMap<>(); + this.parameters.forEach(parameter -> bindings.put(parameter.cardinal(), parameter)); + this.parameterBindings = Collections.unmodifiableMap(bindings); + } + + public static BatchParameterSet from(Map parameterBindings) { + Objects.requireNonNull(parameterBindings, "parameterBindings"); + List orderedParameters = + parameterBindings.entrySet().stream() + .sorted(Comparator.comparingInt(Map.Entry::getKey)) + .map(BatchParameterSet::snapshotParameter) + .collect(Collectors.toList()); + return new BatchParameterSet(orderedParameters); + } + + public List getParameters() { + return parameters; + } + + public Map getParameterBindings() { + return parameterBindings; + } + + public int size() { + return parameters.size(); + } + + public boolean isEmpty() { + return parameters.isEmpty(); + } + + private static ImmutableSqlParameter snapshotParameter( + Map.Entry entry) { + ImmutableSqlParameter parameter = entry.getValue(); + return ImmutableSqlParameter.builder() + .cardinal(entry.getKey()) + .type(parameter.type()) + .value(snapshotValue(parameter.value())) + .build(); + } + + private static Object snapshotValue(Object value) { + if (value instanceof Timestamp) { + Timestamp timestamp = (Timestamp) value; + Timestamp copy = new Timestamp(timestamp.getTime()); + copy.setNanos(timestamp.getNanos()); + return copy; + } + if (value instanceof Date) { + return new Date(((Date) value).getTime()); + } + if (value instanceof Time) { + return new Time(((Time) value).getTime()); + } + if (value instanceof byte[]) { + return ((byte[]) value).clone(); + } + return value; + } +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java index dfa4b70f9..62233e70c 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -1480,6 +1480,11 @@ public boolean isBatchedInsertsEnabled() { return getParameter(DatabricksJdbcUrlParams.ENABLE_BATCHED_INSERTS).equals("1"); } + @Override + public boolean isNativeBatchingEnabled() { + return getParameter(DatabricksJdbcUrlParams.ENABLE_NATIVE_BATCHING).equals("1"); + } + @Override public List getNonRowcountQueryPrefixes() { String prefixesStr = getParameter(DatabricksJdbcUrlParams.NON_ROWCOUNT_QUERY_PREFIXES); diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatement.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatement.java index 32e856aac..27aea28d6 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatement.java @@ -33,7 +33,7 @@ public class DatabricksPreparedStatement extends DatabricksStatement implements JdbcLoggerFactory.getLogger(DatabricksPreparedStatement.class); private final String sql; private DatabricksParameterMetaData databricksParameterMetaData; - private List databricksBatchParameterMetaData; + private List batchParameterSets; private final boolean interpolateParameters; private final int CHUNK_SIZE = 8192; @@ -43,7 +43,7 @@ public DatabricksPreparedStatement(DatabricksConnection connection, String sql) this.sql = sql; this.interpolateParameters = connection.getConnectionContext().supportManyParameters(); this.databricksParameterMetaData = new DatabricksParameterMetaData(sql); - this.databricksBatchParameterMetaData = new ArrayList<>(); + this.batchParameterSets = new ArrayList<>(); // Cache whether this statement should return a ResultSet (based on SQL and config) this.shouldReturnResultSet = shouldReturnResultSetWithConfig(sql); } @@ -58,7 +58,7 @@ public DatabricksPreparedStatement(DatabricksConnection connection, String sql) this.sql = sql; this.interpolateParameters = interpolateParameters; this.databricksParameterMetaData = databricksParameterMetaData; - this.databricksBatchParameterMetaData = new ArrayList<>(); + this.batchParameterSets = new ArrayList<>(); // Cache whether this statement should return a ResultSet (based on SQL and config) this.shouldReturnResultSet = shouldReturnResultSetWithConfig(sql); } @@ -94,7 +94,7 @@ public int executeUpdate() throws SQLException { } @Override - public int[] executeBatch() throws DatabricksBatchUpdateException { + public int[] executeBatch() throws SQLException { LOGGER.debug("public int executeBatch()"); long[] largeUpdateCount = executeLargeBatch(); int[] updateCount = new int[largeUpdateCount.length]; @@ -107,10 +107,10 @@ public int[] executeBatch() throws DatabricksBatchUpdateException { } @Override - public long[] executeLargeBatch() throws DatabricksBatchUpdateException { + public long[] executeLargeBatch() throws SQLException { LOGGER.debug("public long executeLargeBatch()"); - if (databricksBatchParameterMetaData.isEmpty()) { + if (batchParameterSets.isEmpty()) { return new long[0]; } @@ -121,18 +121,41 @@ public long[] executeLargeBatch() throws DatabricksBatchUpdateException { connection, interpolateParameters, (sqlToExecute, params, statementType, closeStatement) -> - executeInternal(sqlToExecute, params, statementType, closeStatement)); - - long[] updateCounts = batchExecutor.executeBatch(databricksBatchParameterMetaData); + executeInternal(sqlToExecute, params, statementType, closeStatement), + new PreparedStatementBatchExecutor.NativeBatchExecutor() { + @Override + public boolean isSupported() { + return supportsNativeParameterBatching(); + } + + @Override + public long[] execute(String sql, List parameterSets) + throws SQLException { + return executeNativeBatchInternal(sql, parameterSets); + } + }); + + long[] updateCounts; + try { + updateCounts = batchExecutor.executeBatch(batchParameterSets); + } catch (NativeBatchResultException e) { + // The backend already completed the batch. Clear it before propagating the count-read error + // so a caller retry cannot insert the same rows again. + clearBatchAfterExecution(); + throw e; + } // Clear the batch after successful execution per JDBC spec + clearBatchAfterExecution(); + return updateCounts; + } + + private void clearBatchAfterExecution() { try { clearBatch(); } catch (SQLException e) { - LOGGER.error("Failed to clear batch after successful execution", e); + LOGGER.error("Failed to clear batch after execution", e); } - - return updateCounts; } @Override @@ -371,7 +394,8 @@ public boolean execute() throws SQLException { @Override public void addBatch() { LOGGER.debug("public void addBatch()"); - this.databricksBatchParameterMetaData.add(databricksParameterMetaData); + this.batchParameterSets.add( + BatchParameterSet.from(databricksParameterMetaData.getParameterBindings())); this.databricksParameterMetaData = new DatabricksParameterMetaData(sql); } @@ -380,7 +404,7 @@ public void clearBatch() throws DatabricksSQLException { LOGGER.debug("public void clearBatch()"); checkIfClosed(); this.databricksParameterMetaData = new DatabricksParameterMetaData(sql); - this.databricksBatchParameterMetaData = new ArrayList<>(); + this.batchParameterSets = new ArrayList<>(); } @Override @@ -755,7 +779,7 @@ private void checkLength(long targetLength, long sourceLength) throws SQLExcepti } private void checkIfBatchOperation() throws DatabricksSQLException { - if (!this.databricksBatchParameterMetaData.isEmpty()) { + if (!this.batchParameterSets.isEmpty()) { String errorMessage = "Batch must either be executed with executeBatch() or cleared with clearBatch()"; LOGGER.error(errorMessage); diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java index cde481ccf..e0fbb4ccf 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -62,6 +62,7 @@ enum ResultSetType { private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(DatabricksResultSet.class); protected static final String AFFECTED_ROWS_COUNT = "num_affected_rows"; + private static final String REPEAT_COUNT = "repeat"; private final ExecutionStatus executionStatus; private final StatementId statementId; private final IExecutionResult executionResult; @@ -2310,6 +2311,44 @@ public long getUpdateCount() throws SQLException { return updateCount; } + long[] getBatchUpdateCounts(int expectedCount) throws SQLException { + checkIfClosed(); + if (resultSetMetaData.getColumnNameIndex(AFFECTED_ROWS_COUNT) < 1) { + throw new DatabricksSQLException( + "Native batch result is missing column " + AFFECTED_ROWS_COUNT, + DatabricksDriverErrorCode.RESULT_SET_ERROR); + } + + long[] counts = new long[expectedCount]; + int index = 0; + boolean hasRepeatCount = resultSetMetaData.getColumnNameIndex(REPEAT_COUNT) > 0; + countingUpdateRows = true; + try { + while (next()) { + long repeatCount = hasRepeatCount ? getLong(REPEAT_COUNT) : 1; + if (repeatCount < 1 || repeatCount > expectedCount - index) { + throw new DatabricksSQLException( + "Native batch returned an invalid repeat count: " + repeatCount, + DatabricksDriverErrorCode.RESULT_SET_ERROR); + } + long affectedRows = getLong(AFFECTED_ROWS_COUNT); + for (long repeated = 0; repeated < repeatCount; repeated++) { + counts[index++] = affectedRows; + } + } + } finally { + countingUpdateRows = false; + } + + if (index != expectedCount) { + throw new DatabricksSQLException( + String.format( + "Native batch returned %d update counts for %d parameter sets", index, expectedCount), + DatabricksDriverErrorCode.RESULT_SET_ERROR); + } + return counts; + } + @Override public boolean hasUpdateCount() throws SQLException { checkIfClosed(); diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java index d1dd0d30e..f9493c50b 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java @@ -866,6 +866,15 @@ DatabricksResultSet executeInternal( LOGGER.debug(stackTraceMessage); CompletableFuture futureResultSet = getFutureResult(sql, params, statementType); + return waitForExecutionResult(sql, stackTraceMessage, futureResultSet, closeStatement); + } + + private DatabricksResultSet waitForExecutionResult( + String sql, + String stackTraceMessage, + CompletableFuture futureResultSet, + boolean closeStatement) + throws SQLException { try { resultSet = timeoutInSeconds == 0 @@ -938,6 +947,38 @@ DatabricksResultSet executeInternal( return result; } + boolean supportsNativeParameterBatching() { + try { + IDatabricksClient client = connection.getSession().getDatabricksClient(); + return client.supportsNativeParameterBatching(connection.getSession().getComputeResource()); + } catch (DatabricksSQLException e) { + LOGGER.warn("Unable to determine native batch capability, using legacy execution", e); + return false; + } + } + + long[] executeNativeBatchInternal(String sql, List parameterSets) + throws SQLException { + resetForNewExecution(); + DatabricksThreadContextHolder.setStatementType(StatementType.UPDATE); + String stackTraceMessage = + format( + "DatabricksResultSet executeNativeBatchInternal(String sql = %s, parameterSetCount = %s)", + sql, parameterSets.size()); + LOGGER.debug(stackTraceMessage); + DatabricksResultSet result = + waitForExecutionResult( + sql, + stackTraceMessage, + getFutureBatchResult(sql, parameterSets, StatementType.UPDATE), + true); + try { + return result.getBatchUpdateCounts(parameterSets.size()); + } catch (SQLException e) { + throw new NativeBatchResultException(e); + } + } + CompletableFuture getFutureResult( String sql, Map params, StatementType statementType) { return CompletableFuture.supplyAsync( @@ -954,6 +995,21 @@ CompletableFuture getFutureResult( executor); } + private CompletableFuture getFutureBatchResult( + String sql, List parameterSets, StatementType statementType) { + return CompletableFuture.supplyAsync( + () -> { + try { + String sqlString = escapeProcessing ? StringUtil.convertJdbcEscapeSequences(sql) : sql; + sqlString = StringUtil.removeRedundantEscapeClause(sqlString); + return getBatchResultFromClient(sqlString, parameterSets, statementType); + } catch (SQLException e) { + throw new RuntimeException(e); + } + }, + executor); + } + DatabricksResultSet getResultFromClient( String sql, Map params, StatementType statementType) throws SQLException { @@ -968,6 +1024,19 @@ DatabricksResultSet getResultFromClient( null /* metadataOperationType */); } + private DatabricksResultSet getBatchResultFromClient( + String sql, List parameterSets, StatementType statementType) + throws SQLException { + IDatabricksClient client = connection.getSession().getDatabricksClient(); + return client.executeStatementBatch( + sql, + connection.getSession().getComputeResource(), + parameterSets, + statementType, + connection.getSession(), + this); + } + void checkIfClosed() throws DatabricksSQLException { if (isClosed) { throw new DatabricksSQLException( diff --git a/src/main/java/com/databricks/jdbc/api/impl/LegacyPreparedStatementBatchExecutor.java b/src/main/java/com/databricks/jdbc/api/impl/LegacyPreparedStatementBatchExecutor.java new file mode 100644 index 000000000..05663b1c9 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/LegacyPreparedStatementBatchExecutor.java @@ -0,0 +1,207 @@ +package com.databricks.jdbc.api.impl; + +import com.databricks.jdbc.common.DatabricksJdbcConstants; +import com.databricks.jdbc.common.StatementType; +import com.databricks.jdbc.common.util.InsertStatementParser; +import com.databricks.jdbc.exception.DatabricksBatchUpdateException; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.log.JdbcLogger; +import com.databricks.jdbc.log.JdbcLoggerFactory; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.sql.Statement; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Executes prepared-statement batches using the legacy client-side strategies. + * + *

This class intentionally preserves the existing behavior: eligible INSERT statements may be + * rewritten into chunked multi-row INSERTs, while all other statements execute one parameter set at + * a time. + */ +class LegacyPreparedStatementBatchExecutor { + + private static final JdbcLogger LOGGER = + JdbcLoggerFactory.getLogger(LegacyPreparedStatementBatchExecutor.class); + + private final String sql; + private final DatabricksConnection connection; + private final boolean interpolateParameters; + private final PreparedStatementBatchExecutor.StatementExecutor statementExecutor; + + LegacyPreparedStatementBatchExecutor( + String sql, + DatabricksConnection connection, + boolean interpolateParameters, + PreparedStatementBatchExecutor.StatementExecutor statementExecutor) { + this.sql = sql; + this.connection = connection; + this.interpolateParameters = interpolateParameters; + this.statementExecutor = statementExecutor; + } + + long[] executeBatch(List batchParameterSets) + throws DatabricksBatchUpdateException { + if (batchParameterSets.isEmpty()) { + return new long[0]; + } + + // Try to optimize INSERT statements with multi-row batching + if (canUseBatchedInsert()) { + return executeBatchedInsert(batchParameterSets); + } else { + // Fall back to individual execution for non-INSERT or incompatible statements + return executeIndividualStatements(batchParameterSets); + } + } + + long[] executeIndividually(List batchParameterSets) + throws DatabricksBatchUpdateException { + return executeIndividualStatements(batchParameterSets); + } + + private boolean canUseBatchedInsert() { + // Check if batched inserts are enabled via connection property + if (!connection.getConnectionContext().isBatchedInsertsEnabled()) { + return false; + } + + // Use strict exception-based parsing for better error handling + try { + InsertStatementParser.parseInsertStrict(sql); + return true; + } catch (Exception e) { + // Not a valid INSERT statement suitable for batching + LOGGER.warn( + "EnableBatchedInserts is enabled but the INSERT statement could not be parsed for" + + " batching, falling back to individual execution: {}", + e.getMessage()); + return false; + } + } + + private long[] executeBatchedInsert(List batchParameterSets) + throws DatabricksBatchUpdateException { + LOGGER.debug("Executing batched INSERT with {} rows", batchParameterSets.size()); + + try { + InsertStatementParser.InsertInfo insertInfo = InsertStatementParser.parseInsertStrict(sql); + + // Calculate how many rows we can fit in one chunk + int parametersPerRow = insertInfo.getColumnCount(); + int maxRowsPerChunk; + + if (interpolateParameters) { + // When parameter interpolation is enabled (supportManyParameters=1), there is no + // parameter limit since values are interpolated directly into the SQL string. + // Try to execute all rows in a single batch, only limited by configured BatchInsertSize + // which users can set based on their data to avoid exceeding the 16MB statement limit. + int configuredBatchSize = connection.getConnectionContext().getBatchInsertSize(); + if (configuredBatchSize < 1) { + throw new DatabricksSQLException( + "BatchInsertSize must be at least 1, got: " + configuredBatchSize, + DatabricksDriverErrorCode.INVALID_STATE); + } + maxRowsPerChunk = Math.min(configuredBatchSize, batchParameterSets.size()); + } else { + // When using parameterized queries, respect the 256 parameter limit from Databricks + // backend + int maxRowsByParameterLimit = + DatabricksJdbcConstants.MAX_QUERY_PARAMETERS / parametersPerRow; + + // Ensure we have at least 1 row per chunk + if (maxRowsByParameterLimit < 1) { + maxRowsPerChunk = 1; + } else { + maxRowsPerChunk = maxRowsByParameterLimit; + } + } + + long[] allUpdateCounts = new long[batchParameterSets.size()]; + + // Process batches in chunks + for (int startIndex = 0; + startIndex < batchParameterSets.size(); + startIndex += maxRowsPerChunk) { + int endIndex = Math.min(startIndex + maxRowsPerChunk, batchParameterSets.size()); + int chunkSize = endIndex - startIndex; + + // Build multi-row INSERT for this chunk + String multiRowSql = InsertStatementParser.generateMultiRowInsert(insertInfo, chunkSize); + Map chunkParams = new HashMap<>(); + int paramIndex = 1; + + for (int i = startIndex; i < endIndex; i++) { + BatchParameterSet batchParams = batchParameterSets.get(i); + Map rowParams = batchParams.getParameterBindings(); + for (int j = 1; j <= rowParams.size(); j++) { + if (rowParams.containsKey(j)) { + chunkParams.put(paramIndex++, rowParams.get(j)); + } + } + } + + // Execute this chunk + String sqlToExecute = + interpolateParameters + ? com.databricks.jdbc.common.util.SQLInterpolator.interpolateSQL( + multiRowSql, chunkParams) + : multiRowSql; + Map paramsToSend = + interpolateParameters ? new HashMap<>() : chunkParams; + statementExecutor.execute(sqlToExecute, paramsToSend, StatementType.UPDATE, false); + + // Set update counts for this chunk (each row typically affects 1 row) + for (int i = startIndex; i < endIndex; i++) { + allUpdateCounts[i] = 1; + } + } + + return allUpdateCounts; + + } catch (DatabricksBatchUpdateException e) { + // Re-throw batch update exceptions (these already have proper update counts) + throw e; + } catch (Exception e) { + // Unexpected exception - mark all as failed + LOGGER.error("Unexpected error executing batched INSERT: {}", e.getMessage(), e); + long[] failedCounts = new long[batchParameterSets.size()]; + for (int i = 0; i < failedCounts.length; i++) { + failedCounts[i] = Statement.EXECUTE_FAILED; + } + throw new DatabricksBatchUpdateException( + e.getMessage(), DatabricksDriverErrorCode.BATCH_EXECUTE_EXCEPTION, failedCounts); + } + } + + private long[] executeIndividualStatements(List batchParameterSets) + throws DatabricksBatchUpdateException { + LOGGER.debug("Executing batch individually with {} statements", batchParameterSets.size()); + long[] largeUpdateCount = new long[batchParameterSets.size()]; + + for (int sqlQueryIndex = 0; sqlQueryIndex < batchParameterSets.size(); sqlQueryIndex++) { + BatchParameterSet batchParameterSet = batchParameterSets.get(sqlQueryIndex); + try { + DatabricksResultSet resultSet = + statementExecutor.execute( + sql, batchParameterSet.getParameterBindings(), StatementType.UPDATE, false); + largeUpdateCount[sqlQueryIndex] = resultSet.getUpdateCount(); + } catch (Exception e) { + LOGGER.error( + "Error executing batch update for index {}: {}", sqlQueryIndex, e.getMessage(), e); + // Set the current failed statement's count + largeUpdateCount[sqlQueryIndex] = Statement.EXECUTE_FAILED; + // Set all remaining statements as failed + for (int i = sqlQueryIndex + 1; i < largeUpdateCount.length; i++) { + largeUpdateCount[i] = Statement.EXECUTE_FAILED; + } + // WARNING: Due to lack of transaction support, any successfully executed statements + // before this failure have already been committed and cannot be rolled back + throw new DatabricksBatchUpdateException( + e.getMessage(), DatabricksDriverErrorCode.BATCH_EXECUTE_EXCEPTION, largeUpdateCount); + } + } + return largeUpdateCount; + } +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/NativeBatchResultException.java b/src/main/java/com/databricks/jdbc/api/impl/NativeBatchResultException.java new file mode 100644 index 000000000..e121b0bb3 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/NativeBatchResultException.java @@ -0,0 +1,22 @@ +package com.databricks.jdbc.api.impl; + +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.sql.SQLException; + +/** + * Indicates that a native batch succeeded but its JDBC update counts could not be read. + * + *

This is intentionally not a {@code BatchUpdateException}: backend execution did not fail. + */ +class NativeBatchResultException extends DatabricksSQLException { + + NativeBatchResultException(SQLException cause) { + super( + "Native batch execution succeeded, but JDBC update counts could not be read. " + + "Inserted rows may already be committed. Cause: " + + cause.getMessage(), + cause, + DatabricksDriverErrorCode.RESULT_SET_ERROR); + } +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutor.java b/src/main/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutor.java index 7387cf67f..3247eecc2 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutor.java +++ b/src/main/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutor.java @@ -1,28 +1,33 @@ package com.databricks.jdbc.api.impl; -import com.databricks.jdbc.common.DatabricksJdbcConstants; import com.databricks.jdbc.common.StatementType; import com.databricks.jdbc.common.util.InsertStatementParser; import com.databricks.jdbc.exception.DatabricksBatchUpdateException; -import com.databricks.jdbc.exception.DatabricksSQLException; -import com.databricks.jdbc.log.JdbcLogger; -import com.databricks.jdbc.log.JdbcLoggerFactory; -import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; import java.sql.SQLException; import java.sql.Statement; -import java.util.HashMap; +import java.util.Arrays; import java.util.List; import java.util.Map; class PreparedStatementBatchExecutor { - private static final JdbcLogger LOGGER = - JdbcLoggerFactory.getLogger(PreparedStatementBatchExecutor.class); + private static final NativeBatchExecutor UNSUPPORTED_NATIVE_EXECUTOR = + new NativeBatchExecutor() { + @Override + public boolean isSupported() { + return false; + } + + @Override + public long[] execute(String sql, List parameterSets) { + throw new IllegalStateException("Native batch execution is not supported"); + } + }; private final String sql; private final DatabricksConnection connection; - private final boolean interpolateParameters; - private final StatementExecutor statementExecutor; + private final LegacyPreparedStatementBatchExecutor legacyExecutor; + private final NativeBatchExecutor nativeExecutor; @FunctionalInterface interface StatementExecutor { @@ -34,178 +39,63 @@ DatabricksResultSet execute( throws SQLException; } + interface NativeBatchExecutor { + boolean isSupported(); + + long[] execute(String sql, List parameterSets) throws SQLException; + } + PreparedStatementBatchExecutor( String sql, DatabricksConnection connection, boolean interpolateParameters, StatementExecutor statementExecutor) { + this(sql, connection, interpolateParameters, statementExecutor, UNSUPPORTED_NATIVE_EXECUTOR); + } + + PreparedStatementBatchExecutor( + String sql, + DatabricksConnection connection, + boolean interpolateParameters, + StatementExecutor statementExecutor, + NativeBatchExecutor nativeExecutor) { this.sql = sql; this.connection = connection; - this.interpolateParameters = interpolateParameters; - this.statementExecutor = statementExecutor; + this.legacyExecutor = + new LegacyPreparedStatementBatchExecutor( + sql, connection, interpolateParameters, statementExecutor); + this.nativeExecutor = nativeExecutor; } - long[] executeBatch(List batchParameterMetaData) - throws DatabricksBatchUpdateException { - if (batchParameterMetaData.isEmpty()) { + long[] executeBatch(List batchParameterSets) throws SQLException { + if (batchParameterSets.isEmpty()) { return new long[0]; } - - // Try to optimize INSERT statements with multi-row batching - if (canUseBatchedInsert()) { - return executeBatchedInsert(batchParameterMetaData); - } else { - // Fall back to individual execution for non-INSERT or incompatible statements - return executeIndividualStatements(batchParameterMetaData); - } - } - - private boolean canUseBatchedInsert() { - // Check if batched inserts are enabled via connection property - if (!connection.getConnectionContext().isBatchedInsertsEnabled()) { - return false; + if (!InsertStatementParser.isParametrizedInsert(sql)) { + return legacyExecutor.executeIndividually(batchParameterSets); } - - // Use strict exception-based parsing for better error handling - try { - InsertStatementParser.parseInsertStrict(sql); - return true; - } catch (Exception e) { - // Not a valid INSERT statement suitable for batching - LOGGER.warn( - "EnableBatchedInserts is enabled but the INSERT statement could not be parsed for" - + " batching, falling back to individual execution: {}", - e.getMessage()); - return false; + if (!connection.getConnectionContext().isNativeBatchingEnabled() + || !nativeExecutor.isSupported()) { + return legacyExecutor.executeBatch(batchParameterSets); } - } - - private long[] executeBatchedInsert(List batchParameterMetaData) - throws DatabricksBatchUpdateException { - LOGGER.debug("Executing batched INSERT with {} rows", batchParameterMetaData.size()); - try { - InsertStatementParser.InsertInfo insertInfo = InsertStatementParser.parseInsertStrict(sql); - - // Calculate how many rows we can fit in one chunk - int parametersPerRow = insertInfo.getColumnCount(); - int maxRowsPerChunk; - - if (interpolateParameters) { - // When parameter interpolation is enabled (supportManyParameters=1), there is no - // parameter limit since values are interpolated directly into the SQL string. - // Try to execute all rows in a single batch, only limited by configured BatchInsertSize - // which users can set based on their data to avoid exceeding the 16MB statement limit. - int configuredBatchSize = connection.getConnectionContext().getBatchInsertSize(); - if (configuredBatchSize < 1) { - throw new DatabricksSQLException( - "BatchInsertSize must be at least 1, got: " + configuredBatchSize, - DatabricksDriverErrorCode.INVALID_STATE); - } - maxRowsPerChunk = Math.min(configuredBatchSize, batchParameterMetaData.size()); - } else { - // When using parameterized queries, respect the 256 parameter limit from Databricks - // backend - int maxRowsByParameterLimit = - DatabricksJdbcConstants.MAX_QUERY_PARAMETERS / parametersPerRow; - - // Ensure we have at least 1 row per chunk - if (maxRowsByParameterLimit < 1) { - maxRowsPerChunk = 1; - } else { - maxRowsPerChunk = maxRowsByParameterLimit; - } - } - - long[] allUpdateCounts = new long[batchParameterMetaData.size()]; - - // Process batches in chunks - for (int startIndex = 0; - startIndex < batchParameterMetaData.size(); - startIndex += maxRowsPerChunk) { - int endIndex = Math.min(startIndex + maxRowsPerChunk, batchParameterMetaData.size()); - int chunkSize = endIndex - startIndex; - - // Build multi-row INSERT for this chunk - String multiRowSql = InsertStatementParser.generateMultiRowInsert(insertInfo, chunkSize); - Map chunkParams = new HashMap<>(); - int paramIndex = 1; - - for (int i = startIndex; i < endIndex; i++) { - DatabricksParameterMetaData batchParams = batchParameterMetaData.get(i); - Map rowParams = batchParams.getParameterBindings(); - for (int j = 1; j <= rowParams.size(); j++) { - if (rowParams.containsKey(j)) { - chunkParams.put(paramIndex++, rowParams.get(j)); - } - } - } - - // Execute this chunk - String sqlToExecute = - interpolateParameters - ? com.databricks.jdbc.common.util.SQLInterpolator.interpolateSQL( - multiRowSql, chunkParams) - : multiRowSql; - Map paramsToSend = - interpolateParameters ? new HashMap<>() : chunkParams; - statementExecutor.execute(sqlToExecute, paramsToSend, StatementType.UPDATE, false); - - // Set update counts for this chunk (each row typically affects 1 row) - for (int i = startIndex; i < endIndex; i++) { - allUpdateCounts[i] = 1; - } - } - - return allUpdateCounts; - - } catch (DatabricksBatchUpdateException e) { - // Re-throw batch update exceptions (these already have proper update counts) + return nativeExecutor.execute(sql, batchParameterSets); + } catch (NativeBatchResultException e) { throw e; - } catch (Exception e) { - // Unexpected exception - mark all as failed - LOGGER.error("Unexpected error executing batched INSERT: {}", e.getMessage(), e); - long[] failedCounts = new long[batchParameterMetaData.size()]; - for (int i = 0; i < failedCounts.length; i++) { - failedCounts[i] = Statement.EXECUTE_FAILED; + } catch (SQLException e) { + if (isUnsupportedNativeBatching(e)) { + return legacyExecutor.executeBatch(batchParameterSets); } + long[] failedCounts = new long[batchParameterSets.size()]; + Arrays.fill(failedCounts, Statement.EXECUTE_FAILED); throw new DatabricksBatchUpdateException( - e.getMessage(), DatabricksDriverErrorCode.BATCH_EXECUTE_EXCEPTION, failedCounts); + e.getMessage(), e.getSQLState(), e.getErrorCode(), failedCounts, e); } } - private long[] executeIndividualStatements( - List batchParameterMetaData) - throws DatabricksBatchUpdateException { - LOGGER.debug("Executing batch individually with {} statements", batchParameterMetaData.size()); - long[] largeUpdateCount = new long[batchParameterMetaData.size()]; - - for (int sqlQueryIndex = 0; sqlQueryIndex < batchParameterMetaData.size(); sqlQueryIndex++) { - DatabricksParameterMetaData databricksParameterMetaData = - batchParameterMetaData.get(sqlQueryIndex); - try { - DatabricksResultSet resultSet = - statementExecutor.execute( - sql, - databricksParameterMetaData.getParameterBindings(), - StatementType.UPDATE, - false); - largeUpdateCount[sqlQueryIndex] = resultSet.getUpdateCount(); - } catch (Exception e) { - LOGGER.error( - "Error executing batch update for index {}: {}", sqlQueryIndex, e.getMessage(), e); - // Set the current failed statement's count - largeUpdateCount[sqlQueryIndex] = Statement.EXECUTE_FAILED; - // Set all remaining statements as failed - for (int i = sqlQueryIndex + 1; i < largeUpdateCount.length; i++) { - largeUpdateCount[i] = Statement.EXECUTE_FAILED; - } - // WARNING: Due to lack of transaction support, any successfully executed statements - // before this failure have already been committed and cannot be rolled back - throw new DatabricksBatchUpdateException( - e.getMessage(), DatabricksDriverErrorCode.BATCH_EXECUTE_EXCEPTION, largeUpdateCount); - } - } - return largeUpdateCount; + private boolean isUnsupportedNativeBatching(SQLException exception) { + return "42P02".equals(exception.getSQLState()) + && exception.getMessage() != null + && exception.getMessage().contains("[UNBOUND_SQL_PARAMETER]"); } } diff --git a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java index fb0d745a7..c3851296c 100644 --- a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java @@ -426,6 +426,9 @@ default int getHeartbeatIntervalSeconds() { /** Returns whether batched INSERT optimization is enabled */ boolean isBatchedInsertsEnabled(); + /** Returns whether native parameter batch execution is enabled */ + boolean isNativeBatchingEnabled(); + /** Returns whether transaction-related method calls should be ignored */ boolean getIgnoreTransactions(); diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java index 7fea2fe2c..bf6531432 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java @@ -194,6 +194,7 @@ public enum DatabricksJdbcUrlParams { "Timeout in seconds for metadata polling operations (e.g. GetTables, GetColumns). 0 means no timeout", "300"), ENABLE_BATCHED_INSERTS("EnableBatchedInserts", "Enable batched INSERT optimization", "0"), + ENABLE_NATIVE_BATCHING("EnableNativeBatching", "Enable native parameter batch execution", "0"), ENABLE_SQL_VALIDATION_FOR_IS_VALID( "EnableSQLValidationForIsValid", "Enable SQL query execution for connection validation in isValid() method", diff --git a/src/main/java/com/databricks/jdbc/common/util/ProtocolFeatureUtil.java b/src/main/java/com/databricks/jdbc/common/util/ProtocolFeatureUtil.java index a451d7f0e..1143771fc 100644 --- a/src/main/java/com/databricks/jdbc/common/util/ProtocolFeatureUtil.java +++ b/src/main/java/com/databricks/jdbc/common/util/ProtocolFeatureUtil.java @@ -140,6 +140,16 @@ public static boolean supportsAsyncMetadataOperations(TProtocolVersion protocolV return protocolVersion.compareTo(TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V9) >= 0; } + /** + * Checks if the given protocol version supports native parameter batches. + * + * @param protocolVersion The protocol version to check + * @return true if native parameter batches are supported, false otherwise + */ + public static boolean supportsNativeParameterBatching(TProtocolVersion protocolVersion) { + return protocolVersion.compareTo(TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V10) >= 0; + } + /** * Checks if the given protocol version indicates a non-Databricks compute. * diff --git a/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java b/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java index 71e790074..03b2083e3 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java @@ -15,6 +15,8 @@ import com.databricks.jdbc.telemetry.latency.DatabricksMetricsTimed; import com.databricks.sdk.core.DatabricksConfig; import java.sql.SQLException; +import java.sql.SQLFeatureNotSupportedException; +import java.util.List; import java.util.Map; /** Interface for Databricks client which abstracts the integration with Databricks server. */ @@ -71,6 +73,37 @@ DatabricksResultSet executeStatement( MetadataOperationType metadataOperationType) throws SQLException; + /** + * Returns whether this client can execute a native parameter batch for the given compute. + * + * @param computeResource underlying SQL warehouse or all-purpose cluster + */ + default boolean supportsNativeParameterBatching(IDatabricksComputeResource computeResource) { + return false; + } + + /** + * Executes one statement with multiple ordered parameter sets in a single backend request. + * + * @param sql SQL statement that needs to be executed + * @param computeResource underlying SQL warehouse or all-purpose cluster + * @param parameterSets ordered parameter sets for the statement + * @param statementType type of statement + * @param session underlying session + * @param parentStatement statement instance + */ + @DatabricksMetricsTimed + default DatabricksResultSet executeStatementBatch( + String sql, + IDatabricksComputeResource computeResource, + List parameterSets, + StatementType statementType, + IDatabricksSession session, + IDatabricksStatementInternal parentStatement) + throws SQLException { + throw new SQLFeatureNotSupportedException("Native parameter batching is not supported"); + } + /** * Executes a statement in Databricks server asynchronously * diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java index 1eb7c82aa..9de6a3b34 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java @@ -13,9 +13,11 @@ import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.api.internal.IDatabricksSession; import com.databricks.jdbc.api.internal.IDatabricksStatementInternal; +import com.databricks.jdbc.common.AllPurposeCluster; import com.databricks.jdbc.common.IDatabricksComputeResource; import com.databricks.jdbc.common.MetadataOperationType; import com.databricks.jdbc.common.StatementType; +import com.databricks.jdbc.common.Warehouse; import com.databricks.jdbc.common.util.DatabricksThreadContextHolder; import com.databricks.jdbc.common.util.DriverUtil; import com.databricks.jdbc.common.util.ProtocolFeatureUtil; @@ -172,6 +174,33 @@ public DatabricksResultSet executeStatement( return thriftAccessor.execute(request, parentStatement, session, statementType); } + @Override + public boolean supportsNativeParameterBatching(IDatabricksComputeResource computeResource) { + if (computeResource instanceof AllPurposeCluster) { + return ProtocolFeatureUtil.supportsNativeParameterBatching(serverProtocolVersion); + } + return computeResource instanceof Warehouse; + } + + @Override + public DatabricksResultSet executeStatementBatch( + String sql, + IDatabricksComputeResource computeResource, + List parameterSets, + StatementType statementType, + IDatabricksSession session, + IDatabricksStatementInternal parentStatement) + throws SQLException { + LOGGER.debug( + "Executing native parameter batch with {} parameter sets on {}", + parameterSets.size(), + computeResource); + DatabricksThreadContextHolder.setStatementType(statementType); + TExecuteStatementReq request = + getBatchRequest(sql, parameterSets, session, parentStatement, statementType); + return thriftAccessor.execute(request, parentStatement, session, statementType); + } + @Override public DatabricksResultSet executeStatementAsync( String sql, @@ -194,17 +223,47 @@ public DatabricksResultSet executeStatementAsync( @VisibleForTesting TSparkParameter mapToSparkParameterListItem(ImmutableSqlParameter parameter) { + return mapToSparkParameterListItem(parameter, parameter.cardinal()); + } + + private TSparkParameter mapToSparkParameterListItem( + ImmutableSqlParameter parameter, int ordinal) { Object value = parameter.value(); String typeString = parameter.type().name(); if (typeString.equals(DECIMAL) && value instanceof BigDecimal) { typeString = getDecimalTypeString((BigDecimal) value); } return new TSparkParameter() - .setOrdinal(parameter.cardinal()) + .setOrdinal(ordinal) .setType(typeString) .setValue(value != null ? TSparkParameterValue.stringValue(value.toString()) : null); } + private TExecuteStatementReq getBatchRequest( + String sql, + List parameterSets, + IDatabricksSession session, + IDatabricksStatementInternal parentStatement, + StatementType statementType) + throws SQLException { + TExecuteStatementReq request = + getRequest(sql, Collections.emptyMap(), session, parentStatement, false, statementType); + request.unsetParameters(); + request.unsetResultRowLimit(); + List> batchParameters = + parameterSets.stream() + .map( + parameterSet -> + parameterSet.getParameters().stream() + .map( + parameter -> + mapToSparkParameterListItem(parameter, parameter.cardinal() - 1)) + .collect(Collectors.toList())) + .collect(Collectors.toList()); + request.setBatchParameters(batchParameters); + return request; + } + private TExecuteStatementReq getRequest( String sql, Map parameters, diff --git a/src/test/java/com/databricks/jdbc/api/impl/BatchParameterSetTest.java b/src/test/java/com/databricks/jdbc/api/impl/BatchParameterSetTest.java new file mode 100644 index 000000000..dccb2e79c --- /dev/null +++ b/src/test/java/com/databricks/jdbc/api/impl/BatchParameterSetTest.java @@ -0,0 +1,113 @@ +package com.databricks.jdbc.api.impl; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.databricks.jdbc.model.core.ColumnInfoTypeName; +import java.sql.Timestamp; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; + +class BatchParameterSetTest { + + @Test + void ordersParametersAndPreservesJdbcIndexes() { + Map bindings = new HashMap<>(); + bindings.put(3, parameter(99, "third", ColumnInfoTypeName.STRING)); + bindings.put(1, parameter(99, "first", ColumnInfoTypeName.STRING)); + bindings.put(2, parameter(99, "second", ColumnInfoTypeName.STRING)); + + BatchParameterSet parameterSet = BatchParameterSet.from(bindings); + + assertEquals(List.of("first", "second", "third"), values(parameterSet)); + assertEquals(List.of(1, 2, 3), indexes(parameterSet)); + assertEquals(List.of(1, 2, 3), List.copyOf(parameterSet.getParameterBindings().keySet())); + } + + @Test + void preservesSparseIndexesWithoutValidation() { + Map bindings = new HashMap<>(); + bindings.put(3, parameter(3, "third", ColumnInfoTypeName.STRING)); + bindings.put(1, parameter(1, "first", ColumnInfoTypeName.STRING)); + + BatchParameterSet parameterSet = BatchParameterSet.from(bindings); + + assertEquals(List.of("first", "third"), values(parameterSet)); + assertEquals(List.of(1, 3), indexes(parameterSet)); + } + + @Test + void allowsEmptyParameterSet() { + BatchParameterSet parameterSet = BatchParameterSet.from(Map.of()); + + assertTrue(parameterSet.isEmpty()); + assertEquals(0, parameterSet.size()); + } + + @Test + void snapshotsBindingsAndMutableValues() { + Timestamp timestamp = Timestamp.valueOf("2026-08-10 12:34:56.123456789"); + byte[] bytes = new byte[] {1, 2, 3}; + Map bindings = new HashMap<>(); + bindings.put(1, parameter(1, timestamp, ColumnInfoTypeName.TIMESTAMP)); + bindings.put(2, parameter(2, bytes, ColumnInfoTypeName.BINARY)); + + BatchParameterSet parameterSet = BatchParameterSet.from(bindings); + bindings.clear(); + timestamp.setTime(0); + bytes[0] = 9; + + assertFalse(parameterSet.isEmpty()); + assertEquals( + Timestamp.valueOf("2026-08-10 12:34:56.123456789"), + parameterSet.getParameters().get(0).value()); + assertArrayEquals(new byte[] {1, 2, 3}, (byte[]) parameterSet.getParameters().get(1).value()); + assertThrows( + UnsupportedOperationException.class, + () -> parameterSet.getParameters().add(parameter(3, "extra", ColumnInfoTypeName.STRING))); + assertThrows( + UnsupportedOperationException.class, + () -> + parameterSet + .getParameterBindings() + .put(3, parameter(3, "extra", ColumnInfoTypeName.STRING))); + } + + @Test + void preservesNullValueAndType() { + BatchParameterSet parameterSet = + BatchParameterSet.from(Map.of(1, parameter(1, null, ColumnInfoTypeName.DECIMAL))); + + ImmutableSqlParameter parameter = parameterSet.getParameters().get(0); + assertNull(parameter.value()); + assertEquals(ColumnInfoTypeName.DECIMAL, parameter.type()); + assertEquals(1, parameter.cardinal()); + } + + private ImmutableSqlParameter parameter( + int cardinal, Object value, ColumnInfoTypeName columnInfoTypeName) { + return ImmutableSqlParameter.builder() + .cardinal(cardinal) + .value(value) + .type(columnInfoTypeName) + .build(); + } + + private List values(BatchParameterSet parameterSet) { + return parameterSet.getParameters().stream() + .map(ImmutableSqlParameter::value) + .collect(java.util.stream.Collectors.toList()); + } + + private List indexes(BatchParameterSet parameterSet) { + return parameterSet.getParameters().stream() + .map(ImmutableSqlParameter::cardinal) + .collect(java.util.stream.Collectors.toList()); + } +} diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksCallableStatementTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksCallableStatementTest.java index e132b2e1b..81fbe20b8 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksCallableStatementTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksCallableStatementTest.java @@ -3,6 +3,7 @@ import static com.databricks.jdbc.TestConstants.*; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyMap; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.when; @@ -304,7 +305,7 @@ void testBatchExecution() throws Exception { when(client.executeStatement( eq(CALL_SQL_AS_EXECUTED), eq(new Warehouse(WAREHOUSE_ID)), - any(HashMap.class), + anyMap(), eq(StatementType.UPDATE), any(IDatabricksSession.class), eq(stmt), diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java index 2d68683de..ef5af0f61 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java @@ -2081,4 +2081,23 @@ public void testRefreshTokenFlowDoesNotReadCredsFromUserPassword() throws Databr assertNull(ctx.getClientSecret()); assertNull(ctx.getNullableClientId()); } + + @Test + public void testNativeBatchingDisabledByDefault() throws DatabricksSQLException { + IDatabricksConnectionContext context = + DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, new Properties()); + + assertFalse(context.isNativeBatchingEnabled()); + } + + @ParameterizedTest + @CsvSource({"0, false", "1, true", "true, false"}) + public void testNativeBatchingConnectionProperty(String value, boolean expected) + throws DatabricksSQLException { + String url = TestConstants.VALID_URL_1 + ";EnableNativeBatching=" + value; + + IDatabricksConnectionContext context = DatabricksConnectionContext.parse(url, new Properties()); + + assertEquals(expected, context.isNativeBatchingEnabled()); + } } diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatementTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatementTest.java index 7552e2a93..7714736bd 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatementTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatementTest.java @@ -4,6 +4,7 @@ import static java.sql.JDBCType.DECIMAL; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; @@ -30,6 +31,8 @@ import java.sql.*; import java.util.Calendar; import java.util.HashMap; +import java.util.List; +import java.util.Map; import java.util.Properties; import java.util.TimeZone; import java.util.stream.Stream; @@ -439,6 +442,118 @@ public void testExecuteLargeBatchStatementThrowsError() throws Exception { } } + @Test + public void testAddBatchSnapshotsMutableParameterValues() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksPreparedStatement statement = + new DatabricksPreparedStatement(connection, "INSERT INTO events (created_at) VALUES (?)"); + Timestamp timestamp = Timestamp.valueOf("2026-08-10 12:34:56.123456789"); + Timestamp expectedTimestamp = Timestamp.valueOf(timestamp.toString()); + + statement.setTimestamp(1, timestamp); + statement.addBatch(); + timestamp.setTime(0); + + @SuppressWarnings("unchecked") + ArgumentCaptor> parametersCaptor = + ArgumentCaptor.forClass(Map.class); + when(client.executeStatement( + anyString(), + eq(new Warehouse(WAREHOUSE_ID)), + parametersCaptor.capture(), + eq(StatementType.UPDATE), + any(IDatabricksSession.class), + eq(statement), + any())) + .thenReturn(resultSet); + when(resultSet.getUpdateCount()).thenReturn(1L); + + assertArrayEquals(new int[] {1}, statement.executeBatch()); + Object snapshottedValue = parametersCaptor.getValue().get(1).value(); + assertEquals(expectedTimestamp, snapshottedValue); + assertNotSame(timestamp, snapshottedValue); + } + + @Test + public void testExecuteBatchUsesSupportedNativeClient() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL + "EnableNativeBatching=1;", new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, thriftClient); + DatabricksPreparedStatement statement = + new DatabricksPreparedStatement(connection, "INSERT INTO target (id, name) VALUES (?, ?)"); + statement.setInt(1, 1); + statement.setString(2, "first"); + statement.addBatch(); + statement.setInt(1, 2); + statement.setString(2, "second"); + statement.addBatch(); + when(thriftClient.supportsNativeParameterBatching(any())).thenReturn(true); + when(thriftClient.executeStatementBatch( + anyString(), + any(), + any(), + eq(StatementType.UPDATE), + any(IDatabricksSession.class), + eq(statement))) + .thenReturn(resultSet); + when(resultSet.getBatchUpdateCounts(2)).thenReturn(new long[] {1, 1}); + + assertArrayEquals(new int[] {1, 1}, statement.executeBatch()); + + @SuppressWarnings("unchecked") + ArgumentCaptor> parameterSetsCaptor = + ArgumentCaptor.forClass(List.class); + verify(thriftClient) + .executeStatementBatch( + eq("INSERT INTO target (id, name) VALUES (?, ?)"), + any(), + parameterSetsCaptor.capture(), + eq(StatementType.UPDATE), + any(IDatabricksSession.class), + eq(statement)); + assertEquals(2, parameterSetsCaptor.getValue().size()); + } + + @Test + public void testExecuteBatchThrowsResultErrorWhenNativeCountsCannotBeRead() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL + "EnableNativeBatching=1;", new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, thriftClient); + DatabricksPreparedStatement statement = + new DatabricksPreparedStatement(connection, "INSERT INTO target (id) VALUES (?)"); + statement.setInt(1, 1); + statement.addBatch(); + when(thriftClient.supportsNativeParameterBatching(any())).thenReturn(true); + when(thriftClient.executeStatementBatch( + anyString(), + any(), + any(), + eq(StatementType.UPDATE), + any(IDatabricksSession.class), + eq(statement))) + .thenReturn(resultSet); + SQLException countError = new SQLException("Missing update-count column", "RESULT_SET_ERROR"); + when(resultSet.getBatchUpdateCounts(1)).thenThrow(countError); + + NativeBatchResultException exception = + assertThrows(NativeBatchResultException.class, statement::executeBatch); + + assertEquals("RESULT_SET_ERROR", exception.getSQLState()); + assertSame(countError, exception.getCause()); + assertTrue(exception.getMessage().contains("Inserted rows may already be committed")); + assertArrayEquals(new int[0], statement.executeBatch()); + verify(thriftClient, times(1)) + .executeStatementBatch( + anyString(), + any(), + any(), + eq(StatementType.UPDATE), + any(IDatabricksSession.class), + eq(statement)); + } + public static ImmutableSqlParameter getSqlParam( int parameterIndex, Object x, String databricksType) { return ImmutableSqlParameter.builder() diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java index e6718275f..a12735d7b 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java @@ -1231,6 +1231,86 @@ void testGetUpdateCountForUpdateStatementMultipleRows() throws SQLException { assertEquals(5L, resultSet.getUpdateCount()); } + @Test + void testGetBatchUpdateCountsPreservesOrder() throws SQLException { + when(mockedResultSetMetadata.getColumnType(1)).thenReturn(Types.BIGINT); + when(mockedResultSetMetadata.getColumnNameIndex(AFFECTED_ROWS_COUNT)).thenReturn(1); + when(mockedExecutionResult.next()).thenReturn(true, true, true, false); + when(mockedExecutionResult.getObject(0)).thenReturn(3L, 1L, 2L); + DatabricksResultSet resultSet = + new DatabricksResultSet( + new StatementStatus().setState(StatementState.SUCCEEDED), + STATEMENT_ID, + StatementType.UPDATE, + null, + mockedExecutionResult, + mockedResultSetMetadata, + false); + + assertArrayEquals(new long[] {3, 1, 2}, resultSet.getBatchUpdateCounts(3)); + } + + @Test + void testGetBatchUpdateCountsExpandsRepeatColumn() throws SQLException { + when(mockedResultSetMetadata.getColumnType(1)).thenReturn(Types.BIGINT); + when(mockedResultSetMetadata.getColumnType(2)).thenReturn(Types.BIGINT); + when(mockedResultSetMetadata.getColumnNameIndex(AFFECTED_ROWS_COUNT)).thenReturn(1); + when(mockedResultSetMetadata.getColumnNameIndex("repeat")).thenReturn(2); + when(mockedExecutionResult.next()).thenReturn(true, false); + when(mockedExecutionResult.getObject(0)).thenReturn(1L); + when(mockedExecutionResult.getObject(1)).thenReturn(3L); + DatabricksResultSet resultSet = + new DatabricksResultSet( + new StatementStatus().setState(StatementState.SUCCEEDED), + STATEMENT_ID, + StatementType.UPDATE, + null, + mockedExecutionResult, + mockedResultSetMetadata, + false); + + assertArrayEquals(new long[] {1, 1, 1}, resultSet.getBatchUpdateCounts(3)); + } + + @Test + void testGetBatchUpdateCountsRejectsWrongCardinality() throws SQLException { + when(mockedResultSetMetadata.getColumnType(1)).thenReturn(Types.BIGINT); + when(mockedResultSetMetadata.getColumnNameIndex(AFFECTED_ROWS_COUNT)).thenReturn(1); + when(mockedExecutionResult.next()).thenReturn(true, false); + when(mockedExecutionResult.getObject(0)).thenReturn(1L); + DatabricksResultSet resultSet = + new DatabricksResultSet( + new StatementStatus().setState(StatementState.SUCCEEDED), + STATEMENT_ID, + StatementType.UPDATE, + null, + mockedExecutionResult, + mockedResultSetMetadata, + false); + + DatabricksSQLException exception = + assertThrows(DatabricksSQLException.class, () -> resultSet.getBatchUpdateCounts(2)); + assertTrue(exception.getMessage().contains("1 update counts for 2 parameter sets")); + } + + @Test + void testGetBatchUpdateCountsRejectsMissingAffectedRowsColumn() throws SQLException { + when(mockedResultSetMetadata.getColumnNameIndex(AFFECTED_ROWS_COUNT)).thenReturn(-1); + DatabricksResultSet resultSet = + new DatabricksResultSet( + new StatementStatus().setState(StatementState.SUCCEEDED), + STATEMENT_ID, + StatementType.UPDATE, + null, + mockedExecutionResult, + mockedResultSetMetadata, + false); + + DatabricksSQLException exception = + assertThrows(DatabricksSQLException.class, () -> resultSet.getBatchUpdateCounts(1)); + assertTrue(exception.getMessage().contains(AFFECTED_ROWS_COUNT)); + } + @Test void testGetUpdateCountForClosedResultSet() throws SQLException { DatabricksResultSet resultSet = getResultSet(StatementState.SUCCEEDED, null); @@ -1634,4 +1714,31 @@ void testGetUpdateCountBypassesMaxRows() throws Exception { // getUpdateCount() must iterate all 5 rows despite maxRows=2 assertEquals(5L, resultSet.getUpdateCount()); } + + @Test + void testGetBatchUpdateCountsBypassesMaxRows() throws Exception { + InlineJsonResult mockExec = mock(InlineJsonResult.class); + when(mockExec.next()).thenReturn(true, false); + when(mockExec.getObject(0)).thenReturn(1L); + when(mockExec.getObject(1)).thenReturn(5L); + + DatabricksResultSetMetaData mockMeta = mock(DatabricksResultSetMetaData.class); + when(mockMeta.getColumnType(1)).thenReturn(Types.BIGINT); + when(mockMeta.getColumnType(2)).thenReturn(Types.BIGINT); + when(mockMeta.getColumnNameIndex(AFFECTED_ROWS_COUNT)).thenReturn(1); + when(mockMeta.getColumnNameIndex("repeat")).thenReturn(2); + IDatabricksStatementInternal stmt = mock(IDatabricksStatementInternal.class); + when(stmt.getLargeMaxRows()).thenReturn(2L); + DatabricksResultSet resultSet = + new DatabricksResultSet( + new StatementStatus().setState(StatementState.SUCCEEDED), + STATEMENT_ID, + StatementType.UPDATE, + stmt, + mockExec, + mockMeta, + false); + + assertArrayEquals(new long[] {1, 1, 1, 1, 1}, resultSet.getBatchUpdateCounts(5)); + } } diff --git a/src/test/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutorTest.java b/src/test/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutorTest.java new file mode 100644 index 000000000..f2a1ba937 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutorTest.java @@ -0,0 +1,346 @@ +package com.databricks.jdbc.api.impl; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; +import com.databricks.jdbc.common.StatementType; +import com.databricks.jdbc.exception.DatabricksBatchUpdateException; +import com.databricks.jdbc.model.core.ColumnInfoTypeName; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +class PreparedStatementBatchExecutorTest { + + private static final String INSERT_SQL = "INSERT INTO target (id, name) VALUES (?, ?)"; + private static final String UPDATE_SQL = "UPDATE target SET name = ? WHERE id = ?"; + + @Mock private DatabricksConnection connection; + @Mock private IDatabricksConnectionContext connectionContext; + @Mock private PreparedStatementBatchExecutor.StatementExecutor statementExecutor; + @Mock private PreparedStatementBatchExecutor.NativeBatchExecutor nativeBatchExecutor; + @Mock private DatabricksResultSet firstResultSet; + @Mock private DatabricksResultSet secondResultSet; + + @Test + void emptyBatchDoesNotExecuteStatements() throws Exception { + PreparedStatementBatchExecutor executor = newExecutor(INSERT_SQL, false); + + assertArrayEquals(new long[0], executor.executeBatch(List.of())); + verifyNoInteractions(connection, statementExecutor); + } + + @Test + void disabledBatchedInsertsExecuteEachParameterSetIndividually() throws Exception { + setBatchedInsertsEnabled(false); + List batch = createBatch(2); + when(statementExecutor.execute(eq(INSERT_SQL), anyMap(), eq(StatementType.UPDATE), eq(false))) + .thenReturn(firstResultSet, secondResultSet); + when(firstResultSet.getUpdateCount()).thenReturn(3L); + when(secondResultSet.getUpdateCount()).thenReturn(5L); + + long[] counts = newExecutor(INSERT_SQL, false, nativeBatchExecutor).executeBatch(batch); + + assertArrayEquals(new long[] {3, 5}, counts); + verify(statementExecutor) + .execute(INSERT_SQL, batch.get(0).getParameterBindings(), StatementType.UPDATE, false); + verify(statementExecutor) + .execute(INSERT_SQL, batch.get(1).getParameterBindings(), StatementType.UPDATE, false); + verify(nativeBatchExecutor, never()).isSupported(); + } + + @Test + void ineligibleSqlFallsBackToIndividualExecution() throws Exception { + List batch = createBatch(1); + when(statementExecutor.execute(eq(UPDATE_SQL), anyMap(), eq(StatementType.UPDATE), eq(false))) + .thenReturn(firstResultSet); + when(firstResultSet.getUpdateCount()).thenReturn(7L); + + long[] counts = newExecutor(UPDATE_SQL, false, nativeBatchExecutor).executeBatch(batch); + + assertArrayEquals(new long[] {7}, counts); + verify(statementExecutor) + .execute(UPDATE_SQL, batch.get(0).getParameterBindings(), StatementType.UPDATE, false); + verify(nativeBatchExecutor, never()).isSupported(); + } + + @Test + void nativeBatchingHandsOrderedParameterSetsToNativeExecutor() throws Exception { + when(connection.getConnectionContext()).thenReturn(connectionContext); + when(connectionContext.isNativeBatchingEnabled()).thenReturn(true); + when(nativeBatchExecutor.isSupported()).thenReturn(true); + List batch = createBatch(2); + when(nativeBatchExecutor.execute(INSERT_SQL, batch)).thenReturn(new long[] {2, 3}); + + long[] counts = newExecutor(INSERT_SQL, false, nativeBatchExecutor).executeBatch(batch); + + assertArrayEquals(new long[] {2, 3}, counts); + assertEquals(List.of(1, 2), indexes(batch.get(0))); + assertEquals(List.of(1, 2), indexes(batch.get(1))); + verify(nativeBatchExecutor).execute(INSERT_SQL, batch); + verifyNoInteractions(statementExecutor); + } + + @Test + void unsupportedNativeExecutorFallsBackToLegacyExecution() throws Exception { + setBatchedInsertsEnabled(false); + when(connectionContext.isNativeBatchingEnabled()).thenReturn(true); + when(nativeBatchExecutor.isSupported()).thenReturn(false); + List batch = createBatch(1); + when(statementExecutor.execute(eq(INSERT_SQL), anyMap(), eq(StatementType.UPDATE), eq(false))) + .thenReturn(firstResultSet); + when(firstResultSet.getUpdateCount()).thenReturn(6L); + + long[] counts = newExecutor(INSERT_SQL, false, nativeBatchExecutor).executeBatch(batch); + + assertArrayEquals(new long[] {6}, counts); + verify(nativeBatchExecutor, never()).execute(anyString(), eq(batch)); + } + + @Test + void unboundParameterCompatibilityErrorFallsBackToLegacyExecution() throws Exception { + setBatchedInsertsEnabled(false); + when(connectionContext.isNativeBatchingEnabled()).thenReturn(true); + when(nativeBatchExecutor.isSupported()).thenReturn(true); + List batch = createBatch(1); + when(nativeBatchExecutor.execute(INSERT_SQL, batch)) + .thenThrow( + new SQLException("[UNBOUND_SQL_PARAMETER] Native batching is unsupported", "42P02", 0)); + when(statementExecutor.execute(eq(INSERT_SQL), anyMap(), eq(StatementType.UPDATE), eq(false))) + .thenReturn(firstResultSet); + when(firstResultSet.getUpdateCount()).thenReturn(4L); + + long[] counts = newExecutor(INSERT_SQL, false, nativeBatchExecutor).executeBatch(batch); + + assertArrayEquals(new long[] {4}, counts); + verify(statementExecutor) + .execute(INSERT_SQL, batch.get(0).getParameterBindings(), StatementType.UPDATE, false); + } + + @Test + void nonCompatibilityNativeErrorDoesNotFallback() throws Exception { + when(connection.getConnectionContext()).thenReturn(connectionContext); + when(connectionContext.isNativeBatchingEnabled()).thenReturn(true); + when(nativeBatchExecutor.isSupported()).thenReturn(true); + List batch = createBatch(2); + SQLException cause = + new SQLException("[PARAMETER_BATCH_ERROR] Too many parameters", "22023", 7); + when(nativeBatchExecutor.execute(INSERT_SQL, batch)).thenThrow(cause); + + DatabricksBatchUpdateException exception = + assertThrows( + DatabricksBatchUpdateException.class, + () -> newExecutor(INSERT_SQL, false, nativeBatchExecutor).executeBatch(batch)); + + assertEquals("22023", exception.getSQLState()); + assertEquals(7, exception.getErrorCode()); + assertSame(cause, exception.getCause()); + assertArrayEquals( + new long[] {Statement.EXECUTE_FAILED, Statement.EXECUTE_FAILED}, + exception.getLargeUpdateCounts()); + verifyNoInteractions(statementExecutor); + } + + @Test + void resultExtractionFailureThrowsDedicatedException() throws Exception { + when(connection.getConnectionContext()).thenReturn(connectionContext); + when(connectionContext.isNativeBatchingEnabled()).thenReturn(true); + when(nativeBatchExecutor.isSupported()).thenReturn(true); + List batch = createBatch(2); + SQLException cause = new SQLException("Missing update-count column", "RESULT_SET_ERROR", 11); + when(nativeBatchExecutor.execute(INSERT_SQL, batch)) + .thenThrow(new NativeBatchResultException(cause)); + + NativeBatchResultException exception = + assertThrows( + NativeBatchResultException.class, + () -> newExecutor(INSERT_SQL, false, nativeBatchExecutor).executeBatch(batch)); + + assertEquals("RESULT_SET_ERROR", exception.getSQLState()); + assertSame(cause, exception.getCause()); + assertTrue(exception.getMessage().contains("Inserted rows may already be committed")); + verifyNoInteractions(statementExecutor); + } + + @Test + void unboundSqlStateWithoutCompatibilityMarkerDoesNotFallback() throws Exception { + when(connection.getConnectionContext()).thenReturn(connectionContext); + when(connectionContext.isNativeBatchingEnabled()).thenReturn(true); + when(nativeBatchExecutor.isSupported()).thenReturn(true); + List batch = createBatch(1); + SQLException cause = new SQLException("A different unbound parameter error", "42P02", 3); + when(nativeBatchExecutor.execute(INSERT_SQL, batch)).thenThrow(cause); + + DatabricksBatchUpdateException exception = + assertThrows( + DatabricksBatchUpdateException.class, + () -> newExecutor(INSERT_SQL, false, nativeBatchExecutor).executeBatch(batch)); + + assertSame(cause, exception.getCause()); + verifyNoInteractions(statementExecutor); + } + + @Test + void eligibleInsertIsRewrittenWithFlattenedParameters() throws Exception { + setBatchedInsertsEnabled(true); + List batch = createBatch(2); + ArgumentCaptor sqlCaptor = ArgumentCaptor.forClass(String.class); + @SuppressWarnings("unchecked") + ArgumentCaptor> parametersCaptor = + ArgumentCaptor.forClass(Map.class); + when(statementExecutor.execute( + sqlCaptor.capture(), parametersCaptor.capture(), eq(StatementType.UPDATE), eq(false))) + .thenReturn(firstResultSet); + + long[] counts = newExecutor(INSERT_SQL, false).executeBatch(batch); + + assertArrayEquals(new long[] {1, 1}, counts); + assertEquals("INSERT INTO target (`id`, `name`) VALUES (?, ?), (?, ?)", sqlCaptor.getValue()); + assertEquals(4, parametersCaptor.getValue().size()); + assertEquals(1, parametersCaptor.getValue().get(1).cardinal()); + assertEquals(2, parametersCaptor.getValue().get(2).cardinal()); + assertEquals(1, parametersCaptor.getValue().get(3).cardinal()); + assertEquals(2, parametersCaptor.getValue().get(4).cardinal()); + } + + @Test + void parameterizedRewriteUsesTheExisting256ParameterChunkLimit() throws Exception { + setBatchedInsertsEnabled(true); + when(statementExecutor.execute(anyString(), anyMap(), eq(StatementType.UPDATE), eq(false))) + .thenReturn(firstResultSet); + + long[] counts = newExecutor(INSERT_SQL, false).executeBatch(createBatch(129)); + + assertEquals(129, counts.length); + verify(statementExecutor) + .execute(eq(multiRowInsert(128)), anyMap(), eq(StatementType.UPDATE), eq(false)); + verify(statementExecutor) + .execute(eq(multiRowInsert(1)), anyMap(), eq(StatementType.UPDATE), eq(false)); + } + + @Test + void interpolatedRewriteUsesConfiguredBatchInsertSize() throws Exception { + setBatchedInsertsEnabled(true); + when(connectionContext.getBatchInsertSize()).thenReturn(2); + ArgumentCaptor sqlCaptor = ArgumentCaptor.forClass(String.class); + @SuppressWarnings("unchecked") + ArgumentCaptor> parametersCaptor = + ArgumentCaptor.forClass(Map.class); + when(statementExecutor.execute( + sqlCaptor.capture(), parametersCaptor.capture(), eq(StatementType.UPDATE), eq(false))) + .thenReturn(firstResultSet); + + long[] counts = newExecutor(INSERT_SQL, true).executeBatch(createBatch(3)); + + assertArrayEquals(new long[] {1, 1, 1}, counts); + assertEquals(2, sqlCaptor.getAllValues().size()); + assertEquals( + "INSERT INTO target (`id`, `name`) VALUES (1, 'name-1'), (2, 'name-2')", + sqlCaptor.getAllValues().get(0)); + assertEquals( + "INSERT INTO target (`id`, `name`) VALUES (3, 'name-3')", sqlCaptor.getAllValues().get(1)); + assertTrue(parametersCaptor.getAllValues().stream().allMatch(Map::isEmpty)); + } + + @Test + void rewrittenBatchFailureMarksEveryParameterSetFailed() throws Exception { + setBatchedInsertsEnabled(true); + when(statementExecutor.execute(anyString(), anyMap(), eq(StatementType.UPDATE), eq(false))) + .thenThrow(new SQLException("rewrite failed")); + + DatabricksBatchUpdateException exception = + assertThrows( + DatabricksBatchUpdateException.class, + () -> newExecutor(INSERT_SQL, false).executeBatch(createBatch(3))); + + assertArrayEquals( + new long[] {Statement.EXECUTE_FAILED, Statement.EXECUTE_FAILED, Statement.EXECUTE_FAILED}, + exception.getLargeUpdateCounts()); + } + + @Test + void individualFailurePreservesEarlierCountAndMarksRemainingSetsFailed() throws Exception { + setBatchedInsertsEnabled(false); + when(statementExecutor.execute(eq(INSERT_SQL), anyMap(), eq(StatementType.UPDATE), eq(false))) + .thenReturn(firstResultSet) + .thenThrow(new SQLException("individual failed")); + when(firstResultSet.getUpdateCount()).thenReturn(4L); + + DatabricksBatchUpdateException exception = + assertThrows( + DatabricksBatchUpdateException.class, + () -> newExecutor(INSERT_SQL, false).executeBatch(createBatch(3))); + + assertArrayEquals( + new long[] {4, Statement.EXECUTE_FAILED, Statement.EXECUTE_FAILED}, + exception.getLargeUpdateCounts()); + } + + private PreparedStatementBatchExecutor newExecutor(String sql, boolean interpolateParameters) { + return new PreparedStatementBatchExecutor( + sql, connection, interpolateParameters, statementExecutor); + } + + private PreparedStatementBatchExecutor newExecutor( + String sql, + boolean interpolateParameters, + PreparedStatementBatchExecutor.NativeBatchExecutor nativeExecutor) { + return new PreparedStatementBatchExecutor( + sql, connection, interpolateParameters, statementExecutor, nativeExecutor); + } + + private void setBatchedInsertsEnabled(boolean enabled) { + when(connection.getConnectionContext()).thenReturn(connectionContext); + when(connectionContext.isBatchedInsertsEnabled()).thenReturn(enabled); + } + + private List createBatch(int rowCount) { + List batch = new ArrayList<>(); + for (int row = 1; row <= rowCount; row++) { + DatabricksParameterMetaData parameterMetaData = new DatabricksParameterMetaData(INSERT_SQL); + parameterMetaData.put(1, parameter(1, row, ColumnInfoTypeName.INT)); + parameterMetaData.put(2, parameter(2, "name-" + row, ColumnInfoTypeName.STRING)); + batch.add(BatchParameterSet.from(parameterMetaData.getParameterBindings())); + } + return batch; + } + + private ImmutableSqlParameter parameter( + int cardinal, Object value, ColumnInfoTypeName columnInfoTypeName) { + return ImmutableSqlParameter.builder() + .cardinal(cardinal) + .value(value) + .type(columnInfoTypeName) + .build(); + } + + private String multiRowInsert(int rows) { + return "INSERT INTO target (`id`, `name`) VALUES " + + String.join(", ", java.util.Collections.nCopies(rows, "(?, ?)")); + } + + private List indexes(BatchParameterSet parameterSet) { + return parameterSet.getParameters().stream() + .map(ImmutableSqlParameter::cardinal) + .collect(java.util.stream.Collectors.toList()); + } +} diff --git a/src/test/java/com/databricks/jdbc/common/util/ProtocolFeatureUtilTest.java b/src/test/java/com/databricks/jdbc/common/util/ProtocolFeatureUtilTest.java index 016b5e850..ed537ec4b 100644 --- a/src/test/java/com/databricks/jdbc/common/util/ProtocolFeatureUtilTest.java +++ b/src/test/java/com/databricks/jdbc/common/util/ProtocolFeatureUtilTest.java @@ -34,6 +34,8 @@ public class ProtocolFeatureUtilTest { private static final TProtocolVersion MIN_VERSION_PARAMETERIZED = SPARK_CLI_SERVICE_PROTOCOL_V8; private static final TProtocolVersion MIN_VERSION_ASYNC_OPERATIONS = SPARK_CLI_SERVICE_PROTOCOL_V9; + private static final TProtocolVersion MIN_VERSION_NATIVE_PARAMETER_BATCHING = + SPARK_CLI_SERVICE_PROTOCOL_V10; private static Stream protocolVersionProvider() { return Stream.of( @@ -154,6 +156,14 @@ public void testSupportsAsyncMetadataOperations(TProtocolVersion version) { assertEquals(expected, actual); } + @ParameterizedTest + @MethodSource("protocolVersionProvider") + public void testSupportsNativeParameterBatching(TProtocolVersion version) { + boolean expected = version.compareTo(MIN_VERSION_NATIVE_PARAMETER_BATCHING) >= 0; + boolean actual = ProtocolFeatureUtil.supportsNativeParameterBatching(version); + assertEquals(expected, actual); + } + @ParameterizedTest @MethodSource("protocolVersionProvider") public void testIsNonDatabricksCompute(TProtocolVersion version) { diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClientTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClientTest.java index c1b5aa857..c4a737e44 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClientTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClientTest.java @@ -158,6 +158,85 @@ void testCloseSession() throws SQLException { assertDoesNotThrow(() -> client.deleteSession(SESSION_INFO)); } + @Test + void testNativeBatchCapabilityUsesProtocolOnlyForAllPurposeClusters() { + DatabricksThriftServiceClient client = + new DatabricksThriftServiceClient(thriftAccessor, connectionContext); + + client.setServerProtocolVersion(TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V9); + assertFalse(client.supportsNativeParameterBatching(CLUSTER_COMPUTE)); + assertTrue(client.supportsNativeParameterBatching(WAREHOUSE_COMPUTE)); + + client.setServerProtocolVersion(TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V10); + assertTrue(client.supportsNativeParameterBatching(CLUSTER_COMPUTE)); + assertTrue(client.supportsNativeParameterBatching(WAREHOUSE_COMPUTE)); + } + + @Test + void testExecuteStatementBatchBuildsNativeThriftRequest() throws SQLException { + when(connectionContext.shouldEnableArrow()).thenReturn(true); + lenient().when(connectionContext.isCloudFetchEnabled()).thenReturn(true); + when(session.getSessionInfo()).thenReturn(SESSION_INFO); + when(parentStatement.getStatement()).thenReturn(statement); + when(parentStatement.getMaxRows()).thenReturn(10); + when(statement.getQueryTimeout()).thenReturn(15); + DatabricksThriftServiceClient client = + new DatabricksThriftServiceClient(thriftAccessor, connectionContext); + client.setServerProtocolVersion(TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V9); + List parameterSets = + List.of( + BatchParameterSet.from( + Map.of( + 1, + ImmutableSqlParameter.builder().cardinal(1).type(INT).value(1).build(), + 2, + ImmutableSqlParameter.builder() + .cardinal(2) + .type(STRING) + .value("first") + .build())), + BatchParameterSet.from( + Map.of( + 1, + ImmutableSqlParameter.builder().cardinal(1).type(INT).value(2).build(), + 2, + ImmutableSqlParameter.builder() + .cardinal(2) + .type(STRING) + .value("second") + .build()))); + when(thriftAccessor.execute( + any(TExecuteStatementReq.class), + eq(parentStatement), + eq(session), + eq(StatementType.UPDATE))) + .thenReturn(resultSet); + + DatabricksResultSet actual = + client.executeStatementBatch( + "INSERT INTO target VALUES (?, ?)", + WAREHOUSE_COMPUTE, + parameterSets, + StatementType.UPDATE, + session, + parentStatement); + + assertSame(resultSet, actual); + ArgumentCaptor requestCaptor = + ArgumentCaptor.forClass(TExecuteStatementReq.class); + verify(thriftAccessor) + .execute( + requestCaptor.capture(), eq(parentStatement), eq(session), eq(StatementType.UPDATE)); + TExecuteStatementReq request = requestCaptor.getValue(); + assertFalse(request.isSetParameters()); + assertFalse(request.isSetResultRowLimit()); + assertEquals(2, request.getBatchParametersSize()); + assertEquals(0, request.getBatchParameters().get(0).get(0).getOrdinal()); + assertEquals(1, request.getBatchParameters().get(0).get(1).getOrdinal()); + assertEquals("first", request.getBatchParameters().get(0).get(1).getValue().getStringValue()); + assertEquals("second", request.getBatchParameters().get(1).get(1).getValue().getStringValue()); + } + private static Stream protocolVersionProvider() { return Stream.of( Arguments.of(TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V1),