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..1ed4eb9fa --- /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 batchParameterMetaData) + throws DatabricksBatchUpdateException { + if (batchParameterMetaData.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; + } + + // 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 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) + 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; + } + throw new DatabricksBatchUpdateException( + e.getMessage(), DatabricksDriverErrorCode.BATCH_EXECUTE_EXCEPTION, failedCounts); + } + } + + 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; + } +} 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..cd5a725a9 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,14 @@ 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.List; import java.util.Map; class PreparedStatementBatchExecutor { - private static final JdbcLogger LOGGER = - JdbcLoggerFactory.getLogger(PreparedStatementBatchExecutor.class); - - private final String sql; - private final DatabricksConnection connection; - private final boolean interpolateParameters; - private final StatementExecutor statementExecutor; + private final LegacyPreparedStatementBatchExecutor legacyExecutor; @FunctionalInterface interface StatementExecutor { @@ -39,173 +25,13 @@ DatabricksResultSet execute( DatabricksConnection connection, boolean interpolateParameters, StatementExecutor statementExecutor) { - this.sql = sql; - this.connection = connection; - this.interpolateParameters = interpolateParameters; - this.statementExecutor = statementExecutor; + this.legacyExecutor = + new LegacyPreparedStatementBatchExecutor( + sql, connection, interpolateParameters, statementExecutor); } long[] executeBatch(List batchParameterMetaData) throws DatabricksBatchUpdateException { - if (batchParameterMetaData.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; - } - - // 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 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) - 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; - } - throw new DatabricksBatchUpdateException( - e.getMessage(), DatabricksDriverErrorCode.BATCH_EXECUTE_EXCEPTION, failedCounts); - } - } - - 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; + return legacyExecutor.executeBatch(batchParameterMetaData); } } 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..9d911322f --- /dev/null +++ b/src/test/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutorTest.java @@ -0,0 +1,212 @@ +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.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.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 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).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); + } + + @Test + void ineligibleSqlFallsBackToIndividualExecution() throws Exception { + setBatchedInsertsEnabled(true); + 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).executeBatch(batch); + + assertArrayEquals(new long[] {7}, counts); + verify(statementExecutor) + .execute(UPDATE_SQL, batch.get(0).getParameterBindings(), StatementType.UPDATE, false); + } + + @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 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(parameterMetaData); + } + 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, "(?, ?)")); + } +}