diff --git a/src/main/java/com/databricks/jdbc/api/impl/BatchParameterSet.java b/src/main/java/com/databricks/jdbc/api/impl/BatchParameterSet.java index f2fe4dd15..6edc670db 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/BatchParameterSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/BatchParameterSet.java @@ -3,7 +3,9 @@ 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; @@ -12,16 +14,21 @@ /** * Immutable, position-ordered snapshot of one prepared-statement parameter set. * - *

This model normalizes JDBC's one-based parameter indexes to zero-based wire ordinals. It does - * not validate parameter completeness, index continuity, or consistency with other parameter sets; - * those validations remain the backend's responsibility. + *

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) { @@ -38,6 +45,10 @@ public List getParameters() { return parameters; } + public Map getParameterBindings() { + return parameterBindings; + } + public int size() { return parameters.size(); } @@ -50,7 +61,7 @@ private static ImmutableSqlParameter snapshotParameter( Map.Entry entry) { ImmutableSqlParameter parameter = entry.getValue(); return ImmutableSqlParameter.builder() - .cardinal(entry.getKey() - 1) + .cardinal(entry.getKey()) .type(parameter.type()) .value(snapshotValue(parameter.value())) .build(); 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..9111a8f77 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); } @@ -110,7 +110,7 @@ public int[] executeBatch() throws DatabricksBatchUpdateException { public long[] executeLargeBatch() throws DatabricksBatchUpdateException { LOGGER.debug("public long executeLargeBatch()"); - if (databricksBatchParameterMetaData.isEmpty()) { + if (batchParameterSets.isEmpty()) { return new long[0]; } @@ -123,7 +123,7 @@ public long[] executeLargeBatch() throws DatabricksBatchUpdateException { (sqlToExecute, params, statementType, closeStatement) -> executeInternal(sqlToExecute, params, statementType, closeStatement)); - long[] updateCounts = batchExecutor.executeBatch(databricksBatchParameterMetaData); + long[] updateCounts = batchExecutor.executeBatch(batchParameterSets); // Clear the batch after successful execution per JDBC spec try { @@ -371,7 +371,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 +381,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 +756,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/LegacyPreparedStatementBatchExecutor.java b/src/main/java/com/databricks/jdbc/api/impl/LegacyPreparedStatementBatchExecutor.java index 1ed4eb9fa..bf9a178dd 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/LegacyPreparedStatementBatchExecutor.java +++ b/src/main/java/com/databricks/jdbc/api/impl/LegacyPreparedStatementBatchExecutor.java @@ -41,18 +41,18 @@ class LegacyPreparedStatementBatchExecutor { this.statementExecutor = statementExecutor; } - long[] executeBatch(List batchParameterMetaData) + long[] executeBatch(List batchParameterSets) throws DatabricksBatchUpdateException { - if (batchParameterMetaData.isEmpty()) { + if (batchParameterSets.isEmpty()) { return new long[0]; } // Try to optimize INSERT statements with multi-row batching if (canUseBatchedInsert()) { - return executeBatchedInsert(batchParameterMetaData); + return executeBatchedInsert(batchParameterSets); } else { // Fall back to individual execution for non-INSERT or incompatible statements - return executeIndividualStatements(batchParameterMetaData); + return executeIndividualStatements(batchParameterSets); } } @@ -76,9 +76,9 @@ private boolean canUseBatchedInsert() { } } - private long[] executeBatchedInsert(List batchParameterMetaData) + private long[] executeBatchedInsert(List batchParameterSets) throws DatabricksBatchUpdateException { - LOGGER.debug("Executing batched INSERT with {} rows", batchParameterMetaData.size()); + LOGGER.debug("Executing batched INSERT with {} rows", batchParameterSets.size()); try { InsertStatementParser.InsertInfo insertInfo = InsertStatementParser.parseInsertStrict(sql); @@ -98,7 +98,7 @@ private long[] executeBatchedInsert(List batchParam "BatchInsertSize must be at least 1, got: " + configuredBatchSize, DatabricksDriverErrorCode.INVALID_STATE); } - maxRowsPerChunk = Math.min(configuredBatchSize, batchParameterMetaData.size()); + maxRowsPerChunk = Math.min(configuredBatchSize, batchParameterSets.size()); } else { // When using parameterized queries, respect the 256 parameter limit from Databricks // backend @@ -113,13 +113,13 @@ private long[] executeBatchedInsert(List batchParam } } - long[] allUpdateCounts = new long[batchParameterMetaData.size()]; + long[] allUpdateCounts = new long[batchParameterSets.size()]; // Process batches in chunks for (int startIndex = 0; - startIndex < batchParameterMetaData.size(); + startIndex < batchParameterSets.size(); startIndex += maxRowsPerChunk) { - int endIndex = Math.min(startIndex + maxRowsPerChunk, batchParameterMetaData.size()); + int endIndex = Math.min(startIndex + maxRowsPerChunk, batchParameterSets.size()); int chunkSize = endIndex - startIndex; // Build multi-row INSERT for this chunk @@ -128,7 +128,7 @@ private long[] executeBatchedInsert(List batchParam int paramIndex = 1; for (int i = startIndex; i < endIndex; i++) { - DatabricksParameterMetaData batchParams = batchParameterMetaData.get(i); + BatchParameterSet batchParams = batchParameterSets.get(i); Map rowParams = batchParams.getParameterBindings(); for (int j = 1; j <= rowParams.size(); j++) { if (rowParams.containsKey(j)) { @@ -161,7 +161,7 @@ private long[] executeBatchedInsert(List batchParam } 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()]; + long[] failedCounts = new long[batchParameterSets.size()]; for (int i = 0; i < failedCounts.length; i++) { failedCounts[i] = Statement.EXECUTE_FAILED; } @@ -170,22 +170,17 @@ private long[] executeBatchedInsert(List batchParam } } - private long[] executeIndividualStatements( - List batchParameterMetaData) + private long[] executeIndividualStatements(List batchParameterSets) throws DatabricksBatchUpdateException { - LOGGER.debug("Executing batch individually with {} statements", batchParameterMetaData.size()); - long[] largeUpdateCount = new long[batchParameterMetaData.size()]; + LOGGER.debug("Executing batch individually with {} statements", batchParameterSets.size()); + long[] largeUpdateCount = new long[batchParameterSets.size()]; - for (int sqlQueryIndex = 0; sqlQueryIndex < batchParameterMetaData.size(); sqlQueryIndex++) { - DatabricksParameterMetaData databricksParameterMetaData = - batchParameterMetaData.get(sqlQueryIndex); + for (int sqlQueryIndex = 0; sqlQueryIndex < batchParameterSets.size(); sqlQueryIndex++) { + BatchParameterSet batchParameterSet = batchParameterSets.get(sqlQueryIndex); try { DatabricksResultSet resultSet = statementExecutor.execute( - sql, - databricksParameterMetaData.getParameterBindings(), - StatementType.UPDATE, - false); + sql, batchParameterSet.getParameterBindings(), StatementType.UPDATE, false); largeUpdateCount[sqlQueryIndex] = resultSet.getUpdateCount(); } catch (Exception e) { LOGGER.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 cd5a725a9..062f5c32c 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutor.java +++ b/src/main/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutor.java @@ -1,6 +1,7 @@ package com.databricks.jdbc.api.impl; import com.databricks.jdbc.common.StatementType; +import com.databricks.jdbc.common.util.InsertStatementParser; import com.databricks.jdbc.exception.DatabricksBatchUpdateException; import java.sql.SQLException; import java.util.List; @@ -8,7 +9,23 @@ class PreparedStatementBatchExecutor { + 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 LegacyPreparedStatementBatchExecutor legacyExecutor; + private final NativeBatchExecutor nativeExecutor; @FunctionalInterface interface StatementExecutor { @@ -20,18 +37,47 @@ DatabricksResultSet execute( throws SQLException; } + interface NativeBatchExecutor { + boolean isSupported(); + + long[] execute(String sql, List parameterSets) + throws DatabricksBatchUpdateException; + } + 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.legacyExecutor = new LegacyPreparedStatementBatchExecutor( sql, connection, interpolateParameters, statementExecutor); + this.nativeExecutor = nativeExecutor; } - long[] executeBatch(List batchParameterMetaData) + long[] executeBatch(List batchParameterSets) throws DatabricksBatchUpdateException { - return legacyExecutor.executeBatch(batchParameterMetaData); + if (canUseNativeBatching(batchParameterSets)) { + return nativeExecutor.execute(sql, batchParameterSets); + } + return legacyExecutor.executeBatch(batchParameterSets); + } + + private boolean canUseNativeBatching(List batchParameterSets) { + return !batchParameterSets.isEmpty() + && connection.getConnectionContext().isNativeBatchingEnabled() + && InsertStatementParser.isParametrizedInsert(sql) + && nativeExecutor.isSupported(); } } diff --git a/src/test/java/com/databricks/jdbc/api/impl/BatchParameterSetTest.java b/src/test/java/com/databricks/jdbc/api/impl/BatchParameterSetTest.java index 5527324f4..dccb2e79c 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/BatchParameterSetTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/BatchParameterSetTest.java @@ -17,7 +17,7 @@ class BatchParameterSetTest { @Test - void ordersParametersByJdbcIndexAndUsesZeroBasedOrdinals() { + void ordersParametersAndPreservesJdbcIndexes() { Map bindings = new HashMap<>(); bindings.put(3, parameter(99, "third", ColumnInfoTypeName.STRING)); bindings.put(1, parameter(99, "first", ColumnInfoTypeName.STRING)); @@ -26,7 +26,8 @@ void ordersParametersByJdbcIndexAndUsesZeroBasedOrdinals() { BatchParameterSet parameterSet = BatchParameterSet.from(bindings); assertEquals(List.of("first", "second", "third"), values(parameterSet)); - assertEquals(List.of(0, 1, 2), ordinals(parameterSet)); + assertEquals(List.of(1, 2, 3), indexes(parameterSet)); + assertEquals(List.of(1, 2, 3), List.copyOf(parameterSet.getParameterBindings().keySet())); } @Test @@ -38,7 +39,7 @@ void preservesSparseIndexesWithoutValidation() { BatchParameterSet parameterSet = BatchParameterSet.from(bindings); assertEquals(List.of("first", "third"), values(parameterSet)); - assertEquals(List.of(0, 2), ordinals(parameterSet)); + assertEquals(List.of(1, 3), indexes(parameterSet)); } @Test @@ -70,6 +71,12 @@ void snapshotsBindingsAndMutableValues() { assertThrows( UnsupportedOperationException.class, () -> parameterSet.getParameters().add(parameter(3, "extra", ColumnInfoTypeName.STRING))); + assertThrows( + UnsupportedOperationException.class, + () -> + parameterSet + .getParameterBindings() + .put(3, parameter(3, "extra", ColumnInfoTypeName.STRING))); } @Test @@ -80,7 +87,7 @@ void preservesNullValueAndType() { ImmutableSqlParameter parameter = parameterSet.getParameters().get(0); assertNull(parameter.value()); assertEquals(ColumnInfoTypeName.DECIMAL, parameter.type()); - assertEquals(0, parameter.cardinal()); + assertEquals(1, parameter.cardinal()); } private ImmutableSqlParameter parameter( @@ -98,7 +105,7 @@ private List values(BatchParameterSet parameterSet) { .collect(java.util.stream.Collectors.toList()); } - private List ordinals(BatchParameterSet parameterSet) { + 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/DatabricksPreparedStatementTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatementTest.java index 7552e2a93..3b908ba78 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,7 @@ import java.sql.*; import java.util.Calendar; import java.util.HashMap; +import java.util.Map; import java.util.Properties; import java.util.TimeZone; import java.util.stream.Stream; @@ -439,6 +441,40 @@ 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); + } + public static ImmutableSqlParameter getSqlParam( int parameterIndex, Object x, String databricksType) { return ImmutableSqlParameter.builder() diff --git a/src/test/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutorTest.java b/src/test/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutorTest.java index 9d911322f..fdd79858f 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutorTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutorTest.java @@ -7,6 +7,7 @@ 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; @@ -35,6 +36,7 @@ class PreparedStatementBatchExecutorTest { @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; @@ -49,40 +51,76 @@ void emptyBatchDoesNotExecuteStatements() throws Exception { @Test void disabledBatchedInsertsExecuteEachParameterSetIndividually() throws Exception { setBatchedInsertsEnabled(false); - List batch = createBatch(2); + 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); + 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 { setBatchedInsertsEnabled(true); - List batch = createBatch(1); + when(connectionContext.isNativeBatchingEnabled()).thenReturn(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); + 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 eligibleInsertIsRewrittenWithFlattenedParameters() throws Exception { setBatchedInsertsEnabled(true); - List batch = createBatch(2); + List batch = createBatch(2); ArgumentCaptor sqlCaptor = ArgumentCaptor.forClass(String.class); @SuppressWarnings("unchecked") ArgumentCaptor> parametersCaptor = @@ -180,18 +218,26 @@ private PreparedStatementBatchExecutor newExecutor(String sql, boolean interpola 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<>(); + 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); + batch.add(BatchParameterSet.from(parameterMetaData.getParameterBindings())); } return batch; } @@ -209,4 +255,10 @@ 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()); + } }