From 2df8074499505ecc20a45f09aef57a0d8588849e Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Mon, 10 Aug 2026 19:11:40 +0000 Subject: [PATCH 1/7] Add native batch capture and routing boundary Snapshot prepared parameter sets at addBatch time and route eligible batches through an injectable native executor while preserving the legacy production path. Signed-off-by: Sreekanth Vadigi --- .../jdbc/api/impl/BatchParameterSet.java | 19 ++++-- .../api/impl/DatabricksPreparedStatement.java | 17 ++--- .../LegacyPreparedStatementBatchExecutor.java | 41 +++++------ .../impl/PreparedStatementBatchExecutor.java | 50 +++++++++++++- .../jdbc/api/impl/BatchParameterSetTest.java | 17 +++-- .../impl/DatabricksPreparedStatementTest.java | 36 ++++++++++ .../PreparedStatementBatchExecutorTest.java | 68 ++++++++++++++++--- 7 files changed, 198 insertions(+), 50 deletions(-) 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/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()); + } } From 597237e946403890862ba875b54121b489cfec12 Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Mon, 10 Aug 2026 19:35:13 +0000 Subject: [PATCH 2/7] Fix callable batch test parameter matcher Accept the immutable parameter map produced by the shared batch model instead of requiring a concrete HashMap. Signed-off-by: Sreekanth Vadigi --- .../jdbc/api/impl/DatabricksCallableStatementTest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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), From 74304f7ea5ed1ce1d184a32c928aa2741a398d50 Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Mon, 10 Aug 2026 21:25:31 +0000 Subject: [PATCH 3/7] Implement native Thrift parameter batching Route eligible prepared batches through compute-aware Thrift batchParameters, preserve legacy compatibility fallback, and map repeat-compressed backend results to ordered JDBC update counts. Signed-off-by: Sreekanth Vadigi --- .../api/impl/DatabricksPreparedStatement.java | 14 ++- .../jdbc/api/impl/DatabricksResultSet.java | 39 +++++++ .../jdbc/api/impl/DatabricksStatement.java | 65 +++++++++++ .../LegacyPreparedStatementBatchExecutor.java | 5 + .../impl/PreparedStatementBatchExecutor.java | 35 ++++-- .../jdbc/common/util/ProtocolFeatureUtil.java | 10 ++ .../jdbc/dbclient/IDatabricksClient.java | 33 ++++++ .../thrift/DatabricksThriftServiceClient.java | 61 +++++++++- .../impl/DatabricksPreparedStatementTest.java | 41 +++++++ .../api/impl/DatabricksResultSetTest.java | 107 ++++++++++++++++++ .../PreparedStatementBatchExecutorTest.java | 65 ++++++++++- .../common/util/ProtocolFeatureUtilTest.java | 10 ++ .../DatabricksThriftServiceClientTest.java | 79 +++++++++++++ 13 files changed, 551 insertions(+), 13 deletions(-) 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 9111a8f77..f0225125d 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatement.java @@ -121,7 +121,19 @@ public long[] executeLargeBatch() throws DatabricksBatchUpdateException { connection, interpolateParameters, (sqlToExecute, params, statementType, closeStatement) -> - executeInternal(sqlToExecute, params, statementType, closeStatement)); + 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 = batchExecutor.executeBatch(batchParameterSets); 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..a069d8f94 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,34 @@ 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); + return result.getBatchUpdateCounts(parameterSets.size()); + } + CompletableFuture getFutureResult( String sql, Map params, StatementType statementType) { return CompletableFuture.supplyAsync( @@ -954,6 +991,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 +1020,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 index bf9a178dd..05663b1c9 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/LegacyPreparedStatementBatchExecutor.java +++ b/src/main/java/com/databricks/jdbc/api/impl/LegacyPreparedStatementBatchExecutor.java @@ -56,6 +56,11 @@ long[] executeBatch(List 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()) { 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 062f5c32c..bd257a71f 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutor.java +++ b/src/main/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutor.java @@ -4,6 +4,8 @@ import com.databricks.jdbc.common.util.InsertStatementParser; import com.databricks.jdbc.exception.DatabricksBatchUpdateException; import java.sql.SQLException; +import java.sql.Statement; +import java.util.Arrays; import java.util.List; import java.util.Map; @@ -40,8 +42,7 @@ DatabricksResultSet execute( interface NativeBatchExecutor { boolean isSupported(); - long[] execute(String sql, List parameterSets) - throws DatabricksBatchUpdateException; + long[] execute(String sql, List parameterSets) throws SQLException; } PreparedStatementBatchExecutor( @@ -68,16 +69,32 @@ long[] execute(String sql, List parameterSets) long[] executeBatch(List batchParameterSets) throws DatabricksBatchUpdateException { - if (canUseNativeBatching(batchParameterSets)) { + if (batchParameterSets.isEmpty()) { + return new long[0]; + } + if (!InsertStatementParser.isParametrizedInsert(sql)) { + return legacyExecutor.executeIndividually(batchParameterSets); + } + if (!connection.getConnectionContext().isNativeBatchingEnabled() + || !nativeExecutor.isSupported()) { + return legacyExecutor.executeBatch(batchParameterSets); + } + try { return nativeExecutor.execute(sql, batchParameterSets); + } 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(), e.getSQLState(), e.getErrorCode(), failedCounts, e); } - return legacyExecutor.executeBatch(batchParameterSets); } - private boolean canUseNativeBatching(List batchParameterSets) { - return !batchParameterSets.isEmpty() - && connection.getConnectionContext().isNativeBatchingEnabled() - && InsertStatementParser.isParametrizedInsert(sql) - && nativeExecutor.isSupported(); + 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/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/DatabricksPreparedStatementTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatementTest.java index 3b908ba78..1e6df2a88 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatementTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatementTest.java @@ -31,6 +31,7 @@ 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; @@ -475,6 +476,46 @@ public void testAddBatchSnapshotsMutableParameterValues() throws Exception { 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()); + } + 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 index fdd79858f..b1b552da3 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutorTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutorTest.java @@ -2,6 +2,7 @@ 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; @@ -69,8 +70,6 @@ void disabledBatchedInsertsExecuteEachParameterSetIndividually() throws Exceptio @Test void ineligibleSqlFallsBackToIndividualExecution() throws Exception { - setBatchedInsertsEnabled(true); - when(connectionContext.isNativeBatchingEnabled()).thenReturn(true); List batch = createBatch(1); when(statementExecutor.execute(eq(UPDATE_SQL), anyMap(), eq(StatementType.UPDATE), eq(false))) .thenReturn(firstResultSet); @@ -117,6 +116,68 @@ void unsupportedNativeExecutorFallsBackToLegacyExecution() throws Exception { 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 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); 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), From b88cd288f12efdd3126506dc37a443dfea7c1842 Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Mon, 10 Aug 2026 22:33:59 +0000 Subject: [PATCH 4/7] Distinguish native batch result extraction failures Surface post-success count-read errors as RESULT_SET_ERROR and clear committed batches before propagating so retries cannot duplicate inserted rows. Signed-off-by: Sreekanth Vadigi --- .../api/impl/DatabricksPreparedStatement.java | 23 ++++++++--- .../jdbc/api/impl/DatabricksStatement.java | 6 ++- .../api/impl/NativeBatchResultException.java | 22 +++++++++++ .../impl/PreparedStatementBatchExecutor.java | 5 ++- .../impl/DatabricksPreparedStatementTest.java | 38 +++++++++++++++++++ .../PreparedStatementBatchExecutorTest.java | 21 ++++++++++ 6 files changed, 106 insertions(+), 9 deletions(-) create mode 100644 src/main/java/com/databricks/jdbc/api/impl/NativeBatchResultException.java 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 f0225125d..27aea28d6 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatement.java @@ -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,7 +107,7 @@ public int[] executeBatch() throws DatabricksBatchUpdateException { } @Override - public long[] executeLargeBatch() throws DatabricksBatchUpdateException { + public long[] executeLargeBatch() throws SQLException { LOGGER.debug("public long executeLargeBatch()"); if (batchParameterSets.isEmpty()) { @@ -135,16 +135,27 @@ public long[] execute(String sql, List parameterSets) } }); - long[] updateCounts = batchExecutor.executeBatch(batchParameterSets); + 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 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 a069d8f94..f9493c50b 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java @@ -972,7 +972,11 @@ long[] executeNativeBatchInternal(String sql, List parameterS stackTraceMessage, getFutureBatchResult(sql, parameterSets, StatementType.UPDATE), true); - return result.getBatchUpdateCounts(parameterSets.size()); + try { + return result.getBatchUpdateCounts(parameterSets.size()); + } catch (SQLException e) { + throw new NativeBatchResultException(e); + } } CompletableFuture getFutureResult( 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 bd257a71f..3247eecc2 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutor.java +++ b/src/main/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutor.java @@ -67,8 +67,7 @@ interface NativeBatchExecutor { this.nativeExecutor = nativeExecutor; } - long[] executeBatch(List batchParameterSets) - throws DatabricksBatchUpdateException { + long[] executeBatch(List batchParameterSets) throws SQLException { if (batchParameterSets.isEmpty()) { return new long[0]; } @@ -81,6 +80,8 @@ long[] executeBatch(List batchParameterSets) } try { return nativeExecutor.execute(sql, batchParameterSets); + } catch (NativeBatchResultException e) { + throw e; } catch (SQLException e) { if (isUnsupportedNativeBatching(e)) { return legacyExecutor.executeBatch(batchParameterSets); 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 1e6df2a88..7714736bd 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatementTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatementTest.java @@ -516,6 +516,44 @@ public void testExecuteBatchUsesSupportedNativeClient() throws Exception { 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/PreparedStatementBatchExecutorTest.java b/src/test/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutorTest.java index b1b552da3..f2a1ba937 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutorTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/PreparedStatementBatchExecutorTest.java @@ -160,6 +160,27 @@ void nonCompatibilityNativeErrorDoesNotFallback() throws Exception { 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); From 07f31353e7ff8c94a644c74a4ace3d970dd63bac Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Tue, 11 Aug 2026 11:56:02 +0000 Subject: [PATCH 5/7] Implement native SEA parameter batching Submit ordered parameter_sets through the existing SEA execution lifecycle so opt-in prepared batches share the same routing, fallback, and JDBC result handling as Thrift. Signed-off-by: Sreekanth Vadigi --- .../impl/sqlexec/DatabricksSdkClient.java | 111 ++++++++++++++++-- .../sqlexec/ExecuteStatementRequest.java | 26 ++++ .../client/sqlexec/StatementParameterSet.java | 47 ++++++++ .../impl/sqlexec/DatabricksSdkClientTest.java | 100 ++++++++++++++++ 4 files changed, 275 insertions(+), 9 deletions(-) create mode 100644 src/main/java/com/databricks/jdbc/model/client/sqlexec/StatementParameterSet.java diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java index a62a4c631..1ab04e85e 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java @@ -191,6 +191,27 @@ public DatabricksResultSet executeStatement( IDatabricksStatementInternal parentStatement, MetadataOperationType metadataOperationType) throws SQLException { + return executeStatement( + sql, + computeResource, + parameters, + statementType, + session, + parentStatement, + metadataOperationType, + null); + } + + private DatabricksResultSet executeStatement( + String sql, + IDatabricksComputeResource computeResource, + Map parameters, + StatementType statementType, + IDatabricksSession session, + IDatabricksStatementInternal parentStatement, + MetadataOperationType metadataOperationType, + ExecuteStatementRequest suppliedRequest) + throws SQLException { LOGGER.debug( "public DatabricksResultSet executeStatement(String sql = {}, compute resource = {}, Map parameters = {}, StatementType statementType = {}, IDatabricksSession session = {}, parentStatement = {}, metadataOperationType = {})", sql, @@ -205,14 +226,16 @@ public DatabricksResultSet executeStatement( long executionStartTime = Instant.now().toEpochMilli(); DatabricksThreadContextHolder.setStatementType(statementType); ExecuteStatementRequest request = - getRequest( - statementType, - sql, - ((Warehouse) computeResource).getWarehouseId(), - session, - parameters, - parentStatement, - false); + suppliedRequest != null + ? suppliedRequest + : getRequest( + statementType, + sql, + ((Warehouse) computeResource).getWarehouseId(), + session, + parameters, + parentStatement, + false); ExecuteStatementResponse response; try { Request req = new Request(Request.POST, STATEMENT_PATH, apiClient.serialize(request)); @@ -355,6 +378,39 @@ public DatabricksResultSet executeStatement( return resultSet; } + @Override + public boolean supportsNativeParameterBatching(IDatabricksComputeResource computeResource) { + return computeResource instanceof Warehouse; + } + + @Override + public DatabricksResultSet executeStatementBatch( + String sql, + IDatabricksComputeResource computeResource, + List parameterSets, + StatementType statementType, + IDatabricksSession session, + IDatabricksStatementInternal parentStatement) + throws SQLException { + ExecuteStatementRequest request = + getBatchRequest( + statementType, + sql, + ((Warehouse) computeResource).getWarehouseId(), + session, + parameterSets, + parentStatement); + return executeStatement( + sql, + computeResource, + Collections.emptyMap(), + statementType, + session, + parentStatement, + null, + request); + } + @Override public DatabricksResultSet executeStatementAsync( String sql, @@ -756,15 +812,52 @@ private ExecuteStatementRequest getRequest( return request; } + private ExecuteStatementRequest getBatchRequest( + StatementType statementType, + String sql, + String warehouseId, + IDatabricksSession session, + List parameterSets, + IDatabricksStatementInternal parentStatement) + throws SQLException { + ExecuteStatementRequest request = + getRequest( + statementType, + sql, + warehouseId, + session, + Collections.emptyMap(), + parentStatement, + false); + List seaParameterSets = + parameterSets.stream() + .map( + parameterSet -> + new StatementParameterSet() + .setParameters( + parameterSet.getParameters().stream() + .map( + parameter -> + mapToParameterListItem(parameter, parameter.cardinal() - 1)) + .collect(Collectors.toList()))) + .collect(Collectors.toList()); + return request.setParameterSets(seaParameterSets).unsetRowLimit(); + } + @VisibleForTesting StatementParameterListItem mapToParameterListItem(ImmutableSqlParameter parameter) { + return mapToParameterListItem(parameter, parameter.cardinal()); + } + + private StatementParameterListItem mapToParameterListItem( + 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 PositionalStatementParameterListItem() - .setOrdinal(parameter.cardinal()) + .setOrdinal(ordinal) .setType(typeString) .setValue(value != null ? value.toString() : null); } diff --git a/src/main/java/com/databricks/jdbc/model/client/sqlexec/ExecuteStatementRequest.java b/src/main/java/com/databricks/jdbc/model/client/sqlexec/ExecuteStatementRequest.java index 24968d2cf..12e173cd8 100644 --- a/src/main/java/com/databricks/jdbc/model/client/sqlexec/ExecuteStatementRequest.java +++ b/src/main/java/com/databricks/jdbc/model/client/sqlexec/ExecuteStatementRequest.java @@ -43,6 +43,9 @@ public class ExecuteStatementRequest { @JsonProperty("parameters") private Collection parameters; + @JsonProperty("parameter_sets") + private Collection parameterSets; + @JsonProperty("result_compression") private CompressionCodec resultCompression; @@ -82,6 +85,10 @@ public Collection getParameters() { return parameters; } + public Collection getParameterSets() { + return parameterSets; + } + public CompressionCodec getResultCompression() { return resultCompression; } @@ -107,6 +114,11 @@ public ExecuteStatementRequest setRowLimit(long rowLimit) { return this; } + public ExecuteStatementRequest unsetRowLimit() { + this.rowLimit = null; + return this; + } + public ExecuteStatementRequest setSessionId(String sessionId) { this.sessionId = sessionId; return this; @@ -135,6 +147,17 @@ public ExecuteStatementRequest setWaitTimeout(String waitTimeout) { public ExecuteStatementRequest setParameters(Collection parameters) { this.parameters = parameters; + if (parameters != null) { + this.parameterSets = null; + } + return this; + } + + public ExecuteStatementRequest setParameterSets(Collection parameterSets) { + this.parameterSets = parameterSets; + if (parameterSets != null) { + this.parameters = null; + } return this; } @@ -145,6 +168,7 @@ public String toString() { .add("format", format) .add("onWaitTimeout", onWaitTimeout) .add("parameters", parameters) + .add("parameterSetCount", parameterSets == null ? 0 : parameterSets.size()) .add("statement", statement) .add("sessionId", sessionId) .add("waitTimeout", waitTimeout) @@ -160,6 +184,7 @@ public int hashCode() { format, onWaitTimeout, parameters, + parameterSets, rowLimit, statement, waitTimeout, @@ -177,6 +202,7 @@ public boolean equals(Object o) { && Objects.equals(format, that.format) && Objects.equals(onWaitTimeout, that.onWaitTimeout) && Objects.equals(parameters, that.parameters) + && Objects.equals(parameterSets, that.parameterSets) && Objects.equals(rowLimit, that.rowLimit) && Objects.equals(statement, that.statement) && Objects.equals(waitTimeout, that.waitTimeout) diff --git a/src/main/java/com/databricks/jdbc/model/client/sqlexec/StatementParameterSet.java b/src/main/java/com/databricks/jdbc/model/client/sqlexec/StatementParameterSet.java new file mode 100644 index 000000000..7cc586978 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/model/client/sqlexec/StatementParameterSet.java @@ -0,0 +1,47 @@ +package com.databricks.jdbc.model.client.sqlexec; + +import com.databricks.sdk.service.sql.StatementParameterListItem; +import com.fasterxml.jackson.annotation.JsonProperty; +import java.util.Collection; +import java.util.List; +import java.util.Objects; + +/** Driver-owned SEA model for one ordered statement parameter set. */ +public class StatementParameterSet { + + @JsonProperty("parameters") + private Collection parameters; + + public Collection getParameters() { + return parameters; + } + + public StatementParameterSet setParameters(Collection parameters) { + this.parameters = parameters == null ? null : List.copyOf(parameters); + return this; + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (!(other instanceof StatementParameterSet)) { + return false; + } + StatementParameterSet that = (StatementParameterSet) other; + return Objects.equals(parameters, that.parameters); + } + + @Override + public int hashCode() { + return Objects.hash(parameters); + } + + @Override + public String toString() { + return "StatementParameterSet{parameterCount=" + + (parameters == null ? 0 : parameters.size()) + + '}'; + } +} diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java index 4a6e26629..ce1370d82 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java @@ -15,11 +15,13 @@ import com.databricks.jdbc.api.impl.*; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; +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.DatabricksTypeUtil; +import com.databricks.jdbc.common.util.JsonUtil; import com.databricks.jdbc.dbclient.impl.common.ConfiguratorUtilsTest; import com.databricks.jdbc.dbclient.impl.common.StatementId; import com.databricks.jdbc.exception.DatabricksSQLException; @@ -38,6 +40,7 @@ import com.databricks.sdk.core.DatabricksError; import com.databricks.sdk.core.http.Request; import com.databricks.sdk.service.sql.*; +import com.fasterxml.jackson.databind.JsonNode; import java.io.File; import java.io.IOException; import java.math.BigDecimal; @@ -227,6 +230,73 @@ public void testExecuteStatement() throws Exception { eq(ExecuteStatementResponse.class)); } + @Test + public void testNativeBatchCapabilityIsWarehouseOnly() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksSdkClient databricksSdkClient = + new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient); + + assertTrue(databricksSdkClient.supportsNativeParameterBatching(warehouse)); + assertFalse( + databricksSdkClient.supportsNativeParameterBatching( + new AllPurposeCluster("org", "cluster"))); + } + + @Test + public void testExecuteStatementBatchBuildsSeaParameterSets() throws Exception { + setupClientMocks(true, false); + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksSdkClient databricksSdkClient = + new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient); + DatabricksConnection connection = + new DatabricksConnection(connectionContext, databricksSdkClient); + connection.open(); + DatabricksStatement statement = new DatabricksStatement(connection); + statement.setMaxRows(100); + List parameterSets = + List.of( + BatchParameterSet.from( + Map.of( + 1, + getSqlParam(1, 1, DatabricksTypeUtil.INT), + 2, + getSqlParam(2, "first", DatabricksTypeUtil.STRING))), + BatchParameterSet.from( + Map.of( + 1, + getSqlParam(1, 2, DatabricksTypeUtil.INT), + 2, + getSqlParam(2, "second", DatabricksTypeUtil.STRING)))); + + databricksSdkClient.executeStatementBatch( + "INSERT INTO target VALUES (?, ?)", + warehouse, + parameterSets, + StatementType.UPDATE, + connection.getSession(), + statement); + + ArgumentCaptor captor = + ArgumentCaptor.forClass(ExecuteStatementRequest.class); + verify(apiClient, atLeastOnce()).serialize(captor.capture()); + ExecuteStatementRequest request = captor.getValue(); + assertNull(request.getParameters()); + assertNull(request.getRowLimit()); + assertEquals(2, request.getParameterSets().size()); + List capturedSets = new ArrayList<>(request.getParameterSets()); + List firstSet = + new ArrayList<>(capturedSets.get(0).getParameters()); + assertEquals(0, ((PositionalStatementParameterListItem) firstSet.get(0)).getOrdinal()); + assertEquals(1, ((PositionalStatementParameterListItem) firstSet.get(1)).getOrdinal()); + assertEquals("first", firstSet.get(1).getValue()); + JsonNode requestJson = JsonUtil.getMapper().valueToTree(request); + assertTrue(requestJson.has("parameter_sets")); + assertTrue(requestJson.get("parameters").isNull()); + assertTrue(requestJson.get("row_limit").isNull()); + } + @Test public void testExecuteStatementAsync() throws Exception { setupClientMocks(false, true); @@ -340,6 +410,36 @@ public void testHandleFailedExecution_FailedState_ThrowsWithoutHY008() throws Ex assertTrue(exception.getMessage().contains("execution failed")); } + @Test + public void testHandleFailedExecutionPreservesUnboundParameterFallbackSignal() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksSdkClient databricksSdkClient = + new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient); + StatementStatus failedStatus = + new StatementStatus() + .setState(StatementState.FAILED) + .setSqlState("42P02") + .setError( + new ServiceError() + .setMessage("[UNBOUND_SQL_PARAMETER] Found an unbound parameter") + .setErrorCode(ServiceErrorCode.BAD_REQUEST)); + ExecuteStatementResponse response = + new ExecuteStatementResponse() + .setStatementId(STATEMENT_ID.toSQLExecStatementId()) + .setStatus(failedStatus); + + DatabricksSQLException exception = + assertThrows( + DatabricksSQLException.class, + () -> + databricksSdkClient.handleFailedExecution( + response, STATEMENT_ID.toSQLExecStatementId(), STATEMENT)); + + assertEquals("42P02", exception.getSQLState()); + assertTrue(exception.getMessage().contains("[UNBOUND_SQL_PARAMETER]")); + } + @Test public void testHandleFailedExecution_unityCatalogError_remapsToCommunicationLinkFailure() throws Exception { From 63149c87d86fa4c6ede892d627b726d4def1183e Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Tue, 11 Aug 2026 19:33:47 +0000 Subject: [PATCH 6/7] Add native batch integration coverage Record real SEA and Thrift responses to guard native parameter batching end to end. Signed-off-by: Sreekanth Vadigi --- .../impl/sqlexec/DatabricksSdkClient.java | 4 + ...reparedStatementBatchIntegrationTests.java | 110 ++++++++++++++++++ ...-1222e668-6d52-4a06-ba2d-803be63ceffe.json | 35 ++++++ ...-51bc1dc1-0639-4083-ae04-6e471f2e8318.json | 36 ++++++ ...-f6212d0a-dace-46ca-bcfd-e3ade1993c79.json | 38 ++++++ ...-20ca48a8-b2a7-416d-a303-03c60356c816.json | 32 +++++ ...-3c461cc1-34d7-49ee-81fa-f752f52f27d9.json | 40 +++++++ ...-7bd94372-a5ea-4b40-b971-754ac62f31bf.json | 38 ++++++ ...-8ee062f9-b6b4-46b2-abd9-3c5d89d24124.json | 41 +++++++ ...-b1bd5e4e-ca97-4b66-8dd9-6085a7af56f3.json | 38 ++++++ ...-573ae363-718a-4356-a9e9-7c386002c0fa.json | 33 ++++++ ...-0ddda2c0-d42a-440a-9d98-8da2246d0625.json | 38 ++++++ ...-6dbb293f-340b-44f9-8a53-e8ed6ebec150.json | 38 ++++++ ...-845f43ca-6a6e-4c14-97c9-98c4428f211f.json | 38 ++++++ ...-b33f2be4-6601-4ab7-9d09-f2df83f83fd8.json | 37 ++++++ ...-b57c5383-59c9-467b-b5ef-7d233f1e2e30.json | 37 ++++++ ...-f02e77e1-4324-41e0-b558-006f609d39ff.json | 38 ++++++ 17 files changed, 671 insertions(+) create mode 100644 src/test/java/com/databricks/jdbc/integration/fakeservice/tests/NativePreparedStatementBatchIntegrationTests.java create mode 100644 src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/.well-known_databricks-config-1222e668-6d52-4a06-ba2d-803be63ceffe.json create mode 100644 src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/.well-known_databricks-config-51bc1dc1-0639-4083-ae04-6e471f2e8318.json create mode 100644 src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_sessions-f6212d0a-dace-46ca-bcfd-e3ade1993c79.json create mode 100644 src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_sessions_01f195b7-5ac0-1831-b357-698025591969-20ca48a8-b2a7-416d-a303-03c60356c816.json create mode 100644 src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-3c461cc1-34d7-49ee-81fa-f752f52f27d9.json create mode 100644 src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-7bd94372-a5ea-4b40-b971-754ac62f31bf.json create mode 100644 src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-8ee062f9-b6b4-46b2-abd9-3c5d89d24124.json create mode 100644 src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-b1bd5e4e-ca97-4b66-8dd9-6085a7af56f3.json create mode 100644 src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/.well-known_databricks-config-573ae363-718a-4356-a9e9-7c386002c0fa.json create mode 100644 src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-0ddda2c0-d42a-440a-9d98-8da2246d0625.json create mode 100644 src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-6dbb293f-340b-44f9-8a53-e8ed6ebec150.json create mode 100644 src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-845f43ca-6a6e-4c14-97c9-98c4428f211f.json create mode 100644 src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-b33f2be4-6601-4ab7-9d09-f2df83f83fd8.json create mode 100644 src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-b57c5383-59c9-467b-b5ef-7d233f1e2e30.json create mode 100644 src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-f02e77e1-4324-41e0-b558-006f609d39ff.json diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java index 1ab04e85e..64b4d1020 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java @@ -392,6 +392,10 @@ public DatabricksResultSet executeStatementBatch( IDatabricksSession session, IDatabricksStatementInternal parentStatement) throws SQLException { + LOGGER.debug( + "Executing native parameter batch with {} parameter sets on {}", + parameterSets.size(), + computeResource); ExecuteStatementRequest request = getBatchRequest( statementType, diff --git a/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/NativePreparedStatementBatchIntegrationTests.java b/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/NativePreparedStatementBatchIntegrationTests.java new file mode 100644 index 000000000..c09dff9db --- /dev/null +++ b/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/NativePreparedStatementBatchIntegrationTests.java @@ -0,0 +1,110 @@ +package com.databricks.jdbc.integration.fakeservice.tests; + +import static com.databricks.jdbc.common.DatabricksJdbcConstants.FAKE_SERVICE_URI_PROP_SUFFIX; +import static com.databricks.jdbc.dbclient.impl.sqlexec.PathConstants.STATEMENT_PATH; +import static com.databricks.jdbc.integration.IntegrationTestUtil.deleteTable; +import static com.databricks.jdbc.integration.IntegrationTestUtil.getFullyQualifiedTableName; +import static com.databricks.jdbc.integration.IntegrationTestUtil.getValidJDBCConnection; +import static com.databricks.jdbc.integration.IntegrationTestUtil.setupDatabaseTable; +import static com.github.tomakehurst.wiremock.client.WireMock.equalToJson; +import static com.github.tomakehurst.wiremock.client.WireMock.postRequestedFor; +import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; + +import com.databricks.jdbc.api.impl.DatabricksConnection; +import com.databricks.jdbc.common.DatabricksClientType; +import com.databricks.jdbc.common.DatabricksJdbcUrlParams; +import com.databricks.jdbc.integration.fakeservice.AbstractFakeServiceIntegrationTests; +import com.databricks.jdbc.integration.fakeservice.FakeServiceExtension; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.SQLException; +import java.util.Properties; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Replays real backend responses for native PreparedStatement batching. */ +public class NativePreparedStatementBatchIntegrationTests + extends AbstractFakeServiceIntegrationTests { + + private Connection connection; + + @BeforeEach + void setUp() throws SQLException { + if (!isSqlExecSdkClient()) { + String targetUri = System.getProperty("thrift_server.targetURI"); + String routeProperty = targetUri + FAKE_SERVICE_URI_PROP_SUFFIX; + System.setProperty( + routeProperty, System.getProperty(routeProperty).replace("localhost", "127.0.0.1")); + } + Properties properties = new Properties(); + properties.setProperty(DatabricksJdbcUrlParams.ENABLE_NATIVE_BATCHING.getParamName(), "1"); + properties.setProperty(DatabricksJdbcUrlParams.ENABLE_BATCHED_INSERTS.getParamName(), "0"); + connection = getValidJDBCConnection(properties); + } + + @AfterEach + void cleanUp() throws SQLException { + if (connection == null) { + return; + } + if (((DatabricksConnection) connection).getConnectionContext().getClientType() + != DatabricksClientType.THRIFT + || getFakeServiceMode() != FakeServiceExtension.FakeServiceMode.REPLAY) { + connection.close(); + } + } + + @Test + void testNativeBatchSuccess() throws SQLException { + String tableName = "native_batch_success_table"; + setupDatabaseTable(connection, tableName); + String sql = + "INSERT INTO " + + getFullyQualifiedTableName(tableName) + + " (id, col1, col2) VALUES (?, ?, ?)"; + + try (PreparedStatement statement = connection.prepareStatement(sql)) { + addRow(statement, 1, "first-a", "first-b"); + addRow(statement, 2, "second-a", "second-b"); + assertArrayEquals(new int[] {1, 1}, statement.executeBatch()); + } finally { + deleteTable(connection, tableName); + } + + if (isSqlExecSdkClient()) { + getDatabricksApiExtension() + .verify( + 1, + postRequestedFor(urlEqualTo(STATEMENT_PATH)) + .withRequestBody( + equalToJson( + "{" + + "\"statement\":\"" + + sql + + "\"," + + "\"parameter_sets\":[" + + "{\"parameters\":[" + + "{\"ordinal\":0,\"type\":\"INT\"}," + + "{\"ordinal\":1,\"type\":\"STRING\"}," + + "{\"ordinal\":2,\"type\":\"STRING\"}]}," + + "{\"parameters\":[" + + "{\"ordinal\":0,\"type\":\"INT\"}," + + "{\"ordinal\":1,\"type\":\"STRING\"}," + + "{\"ordinal\":2,\"type\":\"STRING\"}]}" + + "]" + + "}", + true, + true))); + } + } + + private void addRow(PreparedStatement statement, int id, String col1, String col2) + throws SQLException { + statement.setInt(1, id); + statement.setString(2, col1); + statement.setString(3, col2); + statement.addBatch(); + } +} diff --git a/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/.well-known_databricks-config-1222e668-6d52-4a06-ba2d-803be63ceffe.json b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/.well-known_databricks-config-1222e668-6d52-4a06-ba2d-803be63ceffe.json new file mode 100644 index 000000000..84a767687 --- /dev/null +++ b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/.well-known_databricks-config-1222e668-6d52-4a06-ba2d-803be63ceffe.json @@ -0,0 +1,35 @@ +{ + "id" : "1222e668-6d52-4a06-ba2d-803be63ceffe", + "name" : ".well-known_databricks-config", + "request" : { + "url" : "/.well-known/databricks-config", + "method" : "GET" + }, + "response" : { + "status" : 200, + "body" : "{\"oidc_endpoint\":\"https://e2-dogfood.staging.cloud.databricks.com/oidc\",\"account_id\":\"7a99b43c-b46c-432b-b0a7-814217701909\",\"workspace_id\":\"6051921418418893\",\"cloud_provider\":\"aws\",\"host_type\":\"workspace\",\"token_federation_default_oidc_audiences\":[\"https://e2-dogfood.staging.cloud.databricks.com/oidc/v1/token\"]}", + "headers" : { + "x-request-id" : "fbcdacd1-034f-4dd0-93e2-e43bed4b3fba", + "date" : "Tue, 11 Aug 2026 19:03:35 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "vary" : "Accept-Encoding", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-content-type-options" : "nosniff", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/json", + "server-timing" : "request_id;dur=0;desc=\"fbcdacd1-034f-4dd0-93e2-e43bed4b3fba\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "1222e668-6d52-4a06-ba2d-803be63ceffe", + "scenarioName" : "scenario-2-.well-known-databricks-config", + "requiredScenarioState" : "scenario-2-.well-known-databricks-config-2", + "insertionIndex" : 16 +} \ No newline at end of file diff --git a/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/.well-known_databricks-config-51bc1dc1-0639-4083-ae04-6e471f2e8318.json b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/.well-known_databricks-config-51bc1dc1-0639-4083-ae04-6e471f2e8318.json new file mode 100644 index 000000000..f37e7c081 --- /dev/null +++ b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/.well-known_databricks-config-51bc1dc1-0639-4083-ae04-6e471f2e8318.json @@ -0,0 +1,36 @@ +{ + "id" : "51bc1dc1-0639-4083-ae04-6e471f2e8318", + "name" : ".well-known_databricks-config", + "request" : { + "url" : "/.well-known/databricks-config", + "method" : "GET" + }, + "response" : { + "status" : 200, + "body" : "{\"oidc_endpoint\":\"https://e2-dogfood.staging.cloud.databricks.com/oidc\",\"account_id\":\"7a99b43c-b46c-432b-b0a7-814217701909\",\"workspace_id\":\"6051921418418893\",\"cloud_provider\":\"aws\",\"host_type\":\"workspace\",\"token_federation_default_oidc_audiences\":[\"https://e2-dogfood.staging.cloud.databricks.com/oidc/v1/token\"]}", + "headers" : { + "x-request-id" : "7fac2c32-4bca-4ae3-a67a-85778f489de9", + "date" : "Tue, 11 Aug 2026 19:03:35 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "vary" : "Accept-Encoding", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-content-type-options" : "nosniff", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/json", + "server-timing" : "request_id;dur=0;desc=\"7fac2c32-4bca-4ae3-a67a-85778f489de9\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "51bc1dc1-0639-4083-ae04-6e471f2e8318", + "scenarioName" : "scenario-2-.well-known-databricks-config", + "requiredScenarioState" : "Started", + "newScenarioState" : "scenario-2-.well-known-databricks-config-2", + "insertionIndex" : 17 +} \ No newline at end of file diff --git a/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_sessions-f6212d0a-dace-46ca-bcfd-e3ade1993c79.json b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_sessions-f6212d0a-dace-46ca-bcfd-e3ade1993c79.json new file mode 100644 index 000000000..5f0833d6e --- /dev/null +++ b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_sessions-f6212d0a-dace-46ca-bcfd-e3ade1993c79.json @@ -0,0 +1,38 @@ +{ + "id" : "f6212d0a-dace-46ca-bcfd-e3ade1993c79", + "name" : "api_2.0_sql_sessions", + "request" : { + "url" : "/api/2.0/sql/sessions/", + "method" : "POST", + "bodyPatterns" : [ { + "equalToJson" : "{\"warehouse_id\":\"dd43ee29fedd958d\",\"schema\":\"default\",\"catalog\":\"SPARK\"}", + "ignoreArrayOrder" : true, + "ignoreExtraElements" : true + } ] + }, + "response" : { + "status" : 200, + "body" : "{\"session_id\":\"01f195b7-5ac0-1831-b357-698025591969\"}", + "headers" : { + "x-request-id" : "2bed7242-4859-4259-96a5-b55d8e4f158d", + "date" : "Tue, 11 Aug 2026 19:03:35 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "vary" : "Accept-Encoding", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-content-type-options" : "nosniff", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/json", + "server-timing" : "request_id;dur=0;desc=\"2bed7242-4859-4259-96a5-b55d8e4f158d\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "f6212d0a-dace-46ca-bcfd-e3ade1993c79", + "insertionIndex" : 15 +} \ No newline at end of file diff --git a/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_sessions_01f195b7-5ac0-1831-b357-698025591969-20ca48a8-b2a7-416d-a303-03c60356c816.json b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_sessions_01f195b7-5ac0-1831-b357-698025591969-20ca48a8-b2a7-416d-a303-03c60356c816.json new file mode 100644 index 000000000..e027f7eed --- /dev/null +++ b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_sessions_01f195b7-5ac0-1831-b357-698025591969-20ca48a8-b2a7-416d-a303-03c60356c816.json @@ -0,0 +1,32 @@ +{ + "id" : "20ca48a8-b2a7-416d-a303-03c60356c816", + "name" : "api_2.0_sql_sessions_01f195b7-5ac0-1831-b357-698025591969", + "request" : { + "url" : "/api/2.0/sql/sessions/01f195b7-5ac0-1831-b357-698025591969?warehouse_id=dd43ee29fedd958d", + "method" : "DELETE" + }, + "response" : { + "status" : 200, + "body" : "{}", + "headers" : { + "x-request-id" : "612043ea-79c1-464d-b388-4b13b9f65913", + "date" : "Tue, 11 Aug 2026 19:03:40 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-content-type-options" : "nosniff", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/json", + "server-timing" : "request_id;dur=0;desc=\"612043ea-79c1-464d-b388-4b13b9f65913\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "20ca48a8-b2a7-416d-a303-03c60356c816", + "insertionIndex" : 10 +} \ No newline at end of file diff --git a/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-3c461cc1-34d7-49ee-81fa-f752f52f27d9.json b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-3c461cc1-34d7-49ee-81fa-f752f52f27d9.json new file mode 100644 index 000000000..7236aa471 --- /dev/null +++ b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-3c461cc1-34d7-49ee-81fa-f752f52f27d9.json @@ -0,0 +1,40 @@ +{ + "id" : "3c461cc1-34d7-49ee-81fa-f752f52f27d9", + "name" : "api_2.0_sql_statements", + "request" : { + "url" : "/api/2.0/sql/statements/", + "method" : "POST", + "bodyPatterns" : [ { + "equalToJson" : "{\"statement\":\"DROP TABLE IF EXISTS main.jdbc_test_schema.native_batch_success_table\",\"warehouse_id\":\"dd43ee29fedd958d\",\"session_id\":\"01f195b7-5ac0-1831-b357-698025591969\",\"disposition\":\"INLINE_OR_EXTERNAL_LINKS\",\"format\":\"ARROW_STREAM\",\"on_wait_timeout\":\"CONTINUE\",\"parameters\":[],\"result_compression\":\"LZ4_FRAME\"}", + "ignoreArrayOrder" : true, + "ignoreExtraElements" : true + } ] + }, + "response" : { + "status" : 200, + "body" : "{\"statement_id\":\"01f195b7-5d07-1ef5-b645-4698e6cf2595\",\"status\":{\"state\":\"CLOSED\"},\"manifest\":{\"format\":\"ARROW_STREAM\",\"schema\":{\"column_count\":0},\"total_chunk_count\":0,\"total_row_count\":0,\"total_byte_count\":0,\"truncated\":false,\"result_compression\":\"LZ4_FRAME\"},\"result\":{}}", + "headers" : { + "x-request-id" : "647764e0-88f5-458e-8b19-f5375f120775", + "date" : "Tue, 11 Aug 2026 19:03:40 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "vary" : "Accept-Encoding", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-content-type-options" : "nosniff", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/json", + "server-timing" : "request_id;dur=0;desc=\"647764e0-88f5-458e-8b19-f5375f120775\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "3c461cc1-34d7-49ee-81fa-f752f52f27d9", + "scenarioName" : "scenario-1-api-2.0-sql-statements", + "requiredScenarioState" : "scenario-1-api-2.0-sql-statements-2", + "insertionIndex" : 11 +} \ No newline at end of file diff --git a/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-7bd94372-a5ea-4b40-b971-754ac62f31bf.json b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-7bd94372-a5ea-4b40-b971-754ac62f31bf.json new file mode 100644 index 000000000..71acb7a3e --- /dev/null +++ b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-7bd94372-a5ea-4b40-b971-754ac62f31bf.json @@ -0,0 +1,38 @@ +{ + "id" : "7bd94372-a5ea-4b40-b971-754ac62f31bf", + "name" : "api_2.0_sql_statements", + "request" : { + "url" : "/api/2.0/sql/statements/", + "method" : "POST", + "bodyPatterns" : [ { + "equalToJson" : "{\"statement\":\"CREATE TABLE IF NOT EXISTS main.jdbc_test_schema.native_batch_success_table (id INT PRIMARY KEY, col1 VARCHAR(255), col2 VARCHAR(255))\",\"warehouse_id\":\"dd43ee29fedd958d\",\"session_id\":\"01f195b7-5ac0-1831-b357-698025591969\",\"disposition\":\"INLINE_OR_EXTERNAL_LINKS\",\"format\":\"ARROW_STREAM\",\"on_wait_timeout\":\"CONTINUE\",\"parameters\":[],\"result_compression\":\"LZ4_FRAME\"}", + "ignoreArrayOrder" : true, + "ignoreExtraElements" : true + } ] + }, + "response" : { + "status" : 200, + "body" : "{\"statement_id\":\"01f195b7-5aef-1597-ae00-b00b714a87f1\",\"status\":{\"state\":\"CLOSED\"},\"manifest\":{\"format\":\"ARROW_STREAM\",\"schema\":{\"column_count\":0},\"total_chunk_count\":0,\"total_row_count\":0,\"total_byte_count\":0,\"truncated\":false,\"result_compression\":\"LZ4_FRAME\"},\"result\":{}}", + "headers" : { + "x-request-id" : "6446af60-e8d7-4c42-bf94-461e48a335cc", + "date" : "Tue, 11 Aug 2026 19:03:37 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "vary" : "Accept-Encoding", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-content-type-options" : "nosniff", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/json", + "server-timing" : "request_id;dur=0;desc=\"6446af60-e8d7-4c42-bf94-461e48a335cc\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "7bd94372-a5ea-4b40-b971-754ac62f31bf", + "insertionIndex" : 13 +} \ No newline at end of file diff --git a/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-8ee062f9-b6b4-46b2-abd9-3c5d89d24124.json b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-8ee062f9-b6b4-46b2-abd9-3c5d89d24124.json new file mode 100644 index 000000000..bbcb6d381 --- /dev/null +++ b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-8ee062f9-b6b4-46b2-abd9-3c5d89d24124.json @@ -0,0 +1,41 @@ +{ + "id" : "8ee062f9-b6b4-46b2-abd9-3c5d89d24124", + "name" : "api_2.0_sql_statements", + "request" : { + "url" : "/api/2.0/sql/statements/", + "method" : "POST", + "bodyPatterns" : [ { + "equalToJson" : "{\"statement\":\"DROP TABLE IF EXISTS main.jdbc_test_schema.native_batch_success_table\",\"warehouse_id\":\"dd43ee29fedd958d\",\"session_id\":\"01f195b7-5ac0-1831-b357-698025591969\",\"disposition\":\"INLINE_OR_EXTERNAL_LINKS\",\"format\":\"ARROW_STREAM\",\"on_wait_timeout\":\"CONTINUE\",\"parameters\":[],\"result_compression\":\"LZ4_FRAME\"}", + "ignoreArrayOrder" : true, + "ignoreExtraElements" : true + } ] + }, + "response" : { + "status" : 200, + "body" : "{\"statement_id\":\"01f195b7-5acb-1773-aa8f-b32cb9fa583e\",\"status\":{\"state\":\"CLOSED\"},\"manifest\":{\"format\":\"ARROW_STREAM\",\"schema\":{\"column_count\":0},\"total_chunk_count\":0,\"total_row_count\":0,\"total_byte_count\":0,\"truncated\":false,\"result_compression\":\"LZ4_FRAME\"},\"result\":{}}", + "headers" : { + "x-request-id" : "ade2f6af-341b-4e86-870c-9a00a0c1fa60", + "date" : "Tue, 11 Aug 2026 19:03:36 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "vary" : "Accept-Encoding", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-content-type-options" : "nosniff", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/json", + "server-timing" : "request_id;dur=0;desc=\"ade2f6af-341b-4e86-870c-9a00a0c1fa60\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "8ee062f9-b6b4-46b2-abd9-3c5d89d24124", + "scenarioName" : "scenario-1-api-2.0-sql-statements", + "requiredScenarioState" : "Started", + "newScenarioState" : "scenario-1-api-2.0-sql-statements-2", + "insertionIndex" : 14 +} \ No newline at end of file diff --git a/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-b1bd5e4e-ca97-4b66-8dd9-6085a7af56f3.json b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-b1bd5e4e-ca97-4b66-8dd9-6085a7af56f3.json new file mode 100644 index 000000000..faa01e12c --- /dev/null +++ b/src/test/resources/sqlexecapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/api_2.0_sql_statements-b1bd5e4e-ca97-4b66-8dd9-6085a7af56f3.json @@ -0,0 +1,38 @@ +{ + "id" : "b1bd5e4e-ca97-4b66-8dd9-6085a7af56f3", + "name" : "api_2.0_sql_statements", + "request" : { + "url" : "/api/2.0/sql/statements/", + "method" : "POST", + "bodyPatterns" : [ { + "equalToJson" : "{\"statement\":\"INSERT INTO main.jdbc_test_schema.native_batch_success_table (id, col1, col2) VALUES (?, ?, ?)\",\"warehouse_id\":\"dd43ee29fedd958d\",\"session_id\":\"01f195b7-5ac0-1831-b357-698025591969\",\"disposition\":\"INLINE_OR_EXTERNAL_LINKS\",\"format\":\"ARROW_STREAM\",\"on_wait_timeout\":\"CONTINUE\",\"parameter_sets\":[{\"parameters\":[{\"type\":\"INT\",\"value\":\"1\",\"ordinal\":0},{\"type\":\"STRING\",\"value\":\"first-a\",\"ordinal\":1},{\"type\":\"STRING\",\"value\":\"first-b\",\"ordinal\":2}]},{\"parameters\":[{\"type\":\"INT\",\"value\":\"2\",\"ordinal\":0},{\"type\":\"STRING\",\"value\":\"second-a\",\"ordinal\":1},{\"type\":\"STRING\",\"value\":\"second-b\",\"ordinal\":2}]}],\"result_compression\":\"LZ4_FRAME\"}", + "ignoreArrayOrder" : true, + "ignoreExtraElements" : true + } ] + }, + "response" : { + "status" : 200, + "body" : "{\"statement_id\":\"01f195b7-5bf3-1ed4-bcd6-66efbeddd07a\",\"status\":{\"state\":\"CLOSED\"},\"manifest\":{\"format\":\"ARROW_STREAM\",\"schema\":{\"column_count\":3,\"columns\":[{\"name\":\"num_affected_rows\",\"type_text\":\"BIGINT\",\"type_name\":\"LONG\",\"position\":0},{\"name\":\"num_inserted_rows\",\"type_text\":\"BIGINT\",\"type_name\":\"LONG\",\"position\":1},{\"name\":\"repeat\",\"type_text\":\"BIGINT\",\"type_name\":\"LONG\",\"position\":2}]},\"total_chunk_count\":1,\"chunks\":[{\"chunk_index\":0,\"row_offset\":0,\"row_count\":1,\"byte_count\":288}],\"total_row_count\":1,\"total_byte_count\":288,\"truncated\":false,\"result_compression\":\"LZ4_FRAME\"},\"result\":{\"chunk_index\":0,\"row_offset\":0,\"row_count\":1,\"byte_count\":172,\"attachment\":\"BCJNGHRwjs8AAACg/////yABAAAQAAEAsAoADgAGAA0ACAAKDwAiAAQYAKEBCgAMAAAACAAEGAARCAwAASEAoQADAAAAmAAAAEAuAKQAAIb///8UAAAABABQAAACARgHAAMCAIB0////AAAAATAA3wYAAAByZXBlYXQAAL44AAwXrDgA9RARAAAAbnVtX2luc2VydGVkX3Jvd3MAEgAYABQAEwASrAAXEogAERyAADECASAIAAICAHQIAAwACAAHzAABlAAEXADwAWFmZmVjdGVkX3Jvd3MAAAAI9gXPAAAAAIJR3acEIk0YdHCOlQAAAKL/////6AAAABQAAQDyAgwAFgAOABUAEAAEAAwAAAAwFwBgAAAABAAQCADDAwoAGAAMAAgABAAKOABQeAAAAAEcAAMCABMGCAAAAgAEGAAQCA0ABwgAAUwABzgAERgOAAYgABAgDQAHIAARKA4ABiAAAAIAFwMgAAQCAARwAAQCAA8QAAUEGAAPCAAFgAIAAAAAAAAAJybOwwAAAADB25yBBCJNGHRwjggAAID/////AAAAAIaGkggAAAAAhoaSCA==\"}}", + "headers" : { + "x-request-id" : "bfa8e2ed-6eaa-4012-a82e-3b77806d36b4", + "date" : "Tue, 11 Aug 2026 19:03:39 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "vary" : "Accept-Encoding", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-content-type-options" : "nosniff", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/json", + "server-timing" : "request_id;dur=0;desc=\"bfa8e2ed-6eaa-4012-a82e-3b77806d36b4\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "b1bd5e4e-ca97-4b66-8dd9-6085a7af56f3", + "insertionIndex" : 12 +} \ No newline at end of file diff --git a/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/.well-known_databricks-config-573ae363-718a-4356-a9e9-7c386002c0fa.json b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/.well-known_databricks-config-573ae363-718a-4356-a9e9-7c386002c0fa.json new file mode 100644 index 000000000..4b209f189 --- /dev/null +++ b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/.well-known_databricks-config-573ae363-718a-4356-a9e9-7c386002c0fa.json @@ -0,0 +1,33 @@ +{ + "id" : "573ae363-718a-4356-a9e9-7c386002c0fa", + "name" : ".well-known_databricks-config", + "request" : { + "url" : "/.well-known/databricks-config", + "method" : "GET" + }, + "response" : { + "status" : 200, + "body" : "{\"oidc_endpoint\":\"https://e2-dogfood.staging.cloud.databricks.com/oidc\",\"account_id\":\"7a99b43c-b46c-432b-b0a7-814217701909\",\"workspace_id\":\"6051921418418893\",\"cloud_provider\":\"aws\",\"host_type\":\"workspace\",\"token_federation_default_oidc_audiences\":[\"https://e2-dogfood.staging.cloud.databricks.com/oidc/v1/token\"]}", + "headers" : { + "x-request-id" : "5ca803a7-9228-4809-b556-4d0d0d6536fb", + "date" : "Tue, 11 Aug 2026 19:05:24 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "vary" : "Accept-Encoding", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-content-type-options" : "nosniff", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/json", + "server-timing" : "request_id;dur=0;desc=\"5ca803a7-9228-4809-b556-4d0d0d6536fb\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "573ae363-718a-4356-a9e9-7c386002c0fa", + "insertionIndex" : 15 +} \ No newline at end of file diff --git a/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-0ddda2c0-d42a-440a-9d98-8da2246d0625.json b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-0ddda2c0-d42a-440a-9d98-8da2246d0625.json new file mode 100644 index 000000000..4cf7b352f --- /dev/null +++ b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-0ddda2c0-d42a-440a-9d98-8da2246d0625.json @@ -0,0 +1,38 @@ +{ + "id" : "0ddda2c0-d42a-440a-9d98-8da2246d0625", + "name" : "sql_protocolv1_o_6051921418418893_0819-204509-hill72", + "request" : { + "url" : "/sql/protocolv1/o/6051921418418893/0819-204509-hill72", + "method" : "POST", + "bodyPatterns" : [ { + "binaryEqualTo" : "gAEAAQAAABBFeGVjdXRlU3RhdGVtZW50AAAABQwAAQwAAQwAAQsAAQAAABCn8FflpIBPTqXI403tRARHCwACAAAAELmzugmzUkwDt1BFEq66qUEAAAsAAgAAAEVEUk9QIFRBQkxFIElGIEVYSVNUUyBtYWluLmpkYmNfdGVzdF9zY2hlbWEubmF0aXZlX2JhdGNoX3N1Y2Nlc3NfdGFibGUKAAUAAAAAAAAAAAwFAQoAAQAAAAAAHoSACgACAAAAABghowAAAgUCAQIFAwECBQQBDAUGAgABAQIAAgECAAMBAgAEAQIABQEADwUIDAAAAAAAAA==" + } ] + }, + "response" : { + "status" : 200, + "base64Body" : "gAEAAgAAABBFeGVjdXRlU3RhdGVtZW50AAAABQwAAAwAAQgAAQAAAAAADAACDAABCwABAAAAEHSeg6AYPkjgt2GcgS9dCEILAAIAAAAQQl8z+8v7TWaweJDePVJ2FAYNAQAAAAgAAgAAAAACAAMBAAwFAQwAAQwAAQgAAQAAAAAACAACAAAAAggNAgAAAAAADAACDAABCAABAAAAAAAMAAIPAAEMAAAAAQsAAQAAAAZSZXN1bHQMAAIPAAEMAAAAAQwAAQgAAQAAAAcAAAAIAAMAAAABCwAEAAAAAAAACAUBAAAAAAIFAgELBQMAAABI/////0AAAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAAAAAAAAAAACAUEAAAAAAoFBQAAAAAAAAAACgUGAAAAAAAAAAACBQcACA0BAAAABQoNBQAAAAAAAAAAAg0RAAINEwEADAADDAABCAABAAAAAAACAAIADAADCgABAAAAAAAAAAAPAAIMAAAAAA8FAQwAAAAAAAAMAAQMAAEIAAEAAAAAAAAAAAA=", + "headers" : { + "x-request-id" : "0b5c70b9-f82f-4ec4-a872-23749c0aec52", + "date" : "Tue, 11 Aug 2026 19:05:36 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-frame-options" : "SAMEORIGIN", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-databricks-query-rejection-reason" : "{\"query_stats\":{}}", + "x-content-type-options" : "nosniff", + "x-xss-protection" : "1; mode=block", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/x-thrift", + "server-timing" : "request_id;dur=0;desc=\"0b5c70b9-f82f-4ec4-a872-23749c0aec52\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "0ddda2c0-d42a-440a-9d98-8da2246d0625", + "insertionIndex" : 10 +} \ No newline at end of file diff --git a/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-6dbb293f-340b-44f9-8a53-e8ed6ebec150.json b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-6dbb293f-340b-44f9-8a53-e8ed6ebec150.json new file mode 100644 index 000000000..2efa63868 --- /dev/null +++ b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-6dbb293f-340b-44f9-8a53-e8ed6ebec150.json @@ -0,0 +1,38 @@ +{ + "id" : "6dbb293f-340b-44f9-8a53-e8ed6ebec150", + "name" : "sql_protocolv1_o_6051921418418893_0819-204509-hill72", + "request" : { + "url" : "/sql/protocolv1/o/6051921418418893/0819-204509-hill72", + "method" : "POST", + "bodyPatterns" : [ { + "binaryEqualTo" : "gAEAAQAAABBFeGVjdXRlU3RhdGVtZW50AAAAAwwAAQwAAQwAAQsAAQAAABCn8FflpIBPTqXI403tRARHCwACAAAAELmzugmzUkwDt1BFEq66qUEAAAsAAgAAAIZDUkVBVEUgVEFCTEUgSUYgTk9UIEVYSVNUUyBtYWluLmpkYmNfdGVzdF9zY2hlbWEubmF0aXZlX2JhdGNoX3N1Y2Nlc3NfdGFibGUgKGlkIElOVCBQUklNQVJZIEtFWSwgY29sMSBWQVJDSEFSKDI1NSksIGNvbDIgVkFSQ0hBUigyNTUpKQoABQAAAAAAAAAADAUBCgABAAAAAAAehIAKAAIAAAAAGCGjAAACBQIBAgUDAQIFBAEMBQYCAAEBAgACAQIAAwECAAQBAgAFAQAPBQgMAAAAAAAA" + } ] + }, + "response" : { + "status" : 200, + "base64Body" : "gAEAAgAAABBFeGVjdXRlU3RhdGVtZW50AAAAAwwAAAwAAQgAAQAAAAAADAACDAABCwABAAAAEPGPS9IHmUoQi8QyCrXXhOkLAAIAAAAQUCprFWvcQomFTntH6fbUbAYNAQAAAAgAAgAAAAACAAMBAAwFAQwAAQwAAQgAAQAAAAAACAACAAAAAggNAgAAAAAADAACDAABCAABAAAAAAAMAAIPAAEMAAAAAQsAAQAAAAZSZXN1bHQMAAIPAAEMAAAAAQwAAQgAAQAAAAcAAAAIAAMAAAABCwAEAAAAAAAACAUBAAAAAAIFAgELBQMAAABI/////0AAAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAAAAAAAAAAACAUEAAAAAAoFBQAAAAAAAAAACgUGAAAAAAAAAAACBQcACA0BAAAABQoNBQAAAAAAAAAAAg0RAAINEwEADAADDAABCAABAAAAAAACAAIADAADCgABAAAAAAAAAAAPAAIMAAAAAA8FAQwAAAAAAAAMAAQMAAEIAAEAAAAAAAAAAAA=", + "headers" : { + "x-request-id" : "7892ec11-9abe-4236-9bf2-09a8c8ee58e9", + "date" : "Tue, 11 Aug 2026 19:05:30 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-frame-options" : "SAMEORIGIN", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-databricks-query-rejection-reason" : "{\"query_stats\":{}}", + "x-content-type-options" : "nosniff", + "x-xss-protection" : "1; mode=block", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/x-thrift", + "server-timing" : "request_id;dur=0;desc=\"7892ec11-9abe-4236-9bf2-09a8c8ee58e9\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "6dbb293f-340b-44f9-8a53-e8ed6ebec150", + "insertionIndex" : 12 +} \ No newline at end of file diff --git a/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-845f43ca-6a6e-4c14-97c9-98c4428f211f.json b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-845f43ca-6a6e-4c14-97c9-98c4428f211f.json new file mode 100644 index 000000000..eed1da082 --- /dev/null +++ b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-845f43ca-6a6e-4c14-97c9-98c4428f211f.json @@ -0,0 +1,38 @@ +{ + "id" : "845f43ca-6a6e-4c14-97c9-98c4428f211f", + "name" : "sql_protocolv1_o_6051921418418893_0819-204509-hill72", + "request" : { + "url" : "/sql/protocolv1/o/6051921418418893/0819-204509-hill72", + "method" : "POST", + "bodyPatterns" : [ { + "binaryEqualTo" : "gAEAAQAAABBFeGVjdXRlU3RhdGVtZW50AAAABAwAAQwAAQwAAQsAAQAAABCn8FflpIBPTqXI403tRARHCwACAAAAELmzugmzUkwDt1BFEq66qUEAAAsAAgAAAF5JTlNFUlQgSU5UTyBtYWluLmpkYmNfdGVzdF9zY2hlbWEubmF0aXZlX2JhdGNoX3N1Y2Nlc3NfdGFibGUgKGlkLCBjb2wxLCBjb2wyKSBWQUxVRVMgKD8sID8sID8pCgAFAAAAAAAAAAAMBQEKAAEAAAAAAB6EgAoAAgAAAAAYIaMAAAIFAgECBQMBAgUEAQwFBgIAAQECAAIBAgADAQIABAECAAUBAA8FEQ8AAAACDAAAAAMIAAEAAAAACwADAAAAA0lOVAwABAsAAQAAAAExAAAIAAEAAAABCwADAAAABlNUUklORwwABAsAAQAAAAdmaXJzdC1hAAAIAAEAAAACCwADAAAABlNUUklORwwABAsAAQAAAAdmaXJzdC1iAAAMAAAAAwgAAQAAAAALAAMAAAADSU5UDAAECwABAAAAATIAAAgAAQAAAAELAAMAAAAGU1RSSU5HDAAECwABAAAACHNlY29uZC1hAAAIAAEAAAACCwADAAAABlNUUklORwwABAsAAQAAAAhzZWNvbmQtYgAAAAA=" + } ] + }, + "response" : { + "status" : 200, + "base64Body" : "gAEAAgAAABBFeGVjdXRlU3RhdGVtZW50AAAABAwAAAwAAQgAAQAAAAAADAACDAABCwABAAAAEC9iwT392kaRp/T/i/IjyFsLAAIAAAAQ1FI2ry6BTESNjwxq0AoEEQYNAQAAAAgAAgAAAAACAAMBAAwFAQwAAQwAAQgAAQAAAAAACAACAAAAAgoACwAAAAAAAAACCA0CAAAAAAAMAAIMAAEIAAEAAAAAAAwAAg8AAQwAAAADCwABAAAAEW51bV9hZmZlY3RlZF9yb3dzDAACDwABDAAAAAEMAAEIAAEAAAAEAAAACAADAAAAAQsABAAAAAAACwABAAAAEW51bV9pbnNlcnRlZF9yb3dzDAACDwABDAAAAAEMAAEIAAEAAAAEAAAACAADAAAAAgsABAAAAAAACwABAAAABnJlcGVhdAwAAg8AAQwAAAABDAABCAABAAAABAAAAAgAAwAAAAMLAAQAAAAAAAAIBQEAAAAAAgUCAQsFAwAAAoD/////eAIAABAAAAAAAAoADgAGAA0ACAAKAAAAAAAEABAAAAAAAQoADAAAAAgABAAKAAAACAAAAAgAAAAAAAAAAwAAAHgBAACwAAAABAAAAKb+//8UAAAAhAAAAIQAAAAAAAIBiAAAAAIAAAA8AAAABAAAAGD+//8IAAAAEAAAAAYAAAAibG9uZyIAABcAAABTcGFyazpEYXRhVHlwZTpKc29uVHlwZQCU/v//CAAAABAAAAAGAAAAQklHSU5UAAAWAAAAU3Bhcms6RGF0YVR5cGU6U3FsTmFtZQAAAAAAAIz+//8AAAABQAAAAAYAAAByZXBlYXQAAE7///8UAAAAhAAAAIQAAAAAAAIBiAAAAAIAAAA8AAAABAAAAAj///8IAAAAEAAAAAYAAAAibG9uZyIAABcAAABTcGFyazpEYXRhVHlwZTpKc29uVHlwZQA8////CAAAABAAAAAGAAAAQklHSU5UAAAWAAAAU3Bhcms6RGF0YVR5cGU6U3FsTmFtZQAAAAAAADT///8AAAABQAAAABEAAABudW1faW5zZXJ0ZWRfcm93cwASABgAFAATABIADAAAAAgABAASAAAAFAAAAIwAAACUAAAAAAACAZgAAAACAAAARAAAAAQAAADM////CAAAABAAAAAGAAAAImxvbmciAAAXAAAAU3Bhcms6RGF0YVR5cGU6SnNvblR5cGUACAAMAAgABAAIAAAACAAAABAAAAAGAAAAQklHSU5UAAAWAAAAU3Bhcms6RGF0YVR5cGU6U3FsTmFtZQAAAAAAAAgADAAIAAcACAAAAAAAAAFAAAAAEQAAAG51bV9hZmZlY3RlZF9yb3dzAAAACAUEAAAAAAoFBQAAAAAAAAEgCgUGAAAAAAAAAKwCBQcACA0BAAAABQoNBQAAAAAAAAAAAg0RAAINEwEADAADDAABCAABAAAAAAACAAIADAADCgABAAAAAAAAAAAPAAIMAAAAAA8FAQwAAAABCwABAAAArAQiTRh0cI6VAAAAov/////oAAAAFAABAPICDAAWAA4AFQAQAAQADAAAADAXAGAAAAAEABAIAMMDCgAYAAwACAAEAAo4AFB4AAAAARwAAwIAEwYIAAACAAQYABAIDQAHCAABTAAHOAARGA4ABiAAECANAAcgABEoDgAGIAAAAgAXAyAABAIABHAABAIADxAABQQYAA8IAAWAAgAAAAAAAAAnJs7DAAAAAMHbnIEKAAIAAAAAAAAAAQAAAAwABAwAAQgAAQAAAAAAAAAAAA==", + "headers" : { + "x-request-id" : "0f446c10-8760-4cb2-9e27-973e1fb31eab", + "date" : "Tue, 11 Aug 2026 19:05:35 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-frame-options" : "SAMEORIGIN", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-databricks-query-rejection-reason" : "{\"query_stats\":{}}", + "x-content-type-options" : "nosniff", + "x-xss-protection" : "1; mode=block", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/x-thrift", + "server-timing" : "request_id;dur=0;desc=\"0f446c10-8760-4cb2-9e27-973e1fb31eab\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "845f43ca-6a6e-4c14-97c9-98c4428f211f", + "insertionIndex" : 11 +} \ No newline at end of file diff --git a/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-b33f2be4-6601-4ab7-9d09-f2df83f83fd8.json b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-b33f2be4-6601-4ab7-9d09-f2df83f83fd8.json new file mode 100644 index 000000000..79149b958 --- /dev/null +++ b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-b33f2be4-6601-4ab7-9d09-f2df83f83fd8.json @@ -0,0 +1,37 @@ +{ + "id" : "b33f2be4-6601-4ab7-9d09-f2df83f83fd8", + "name" : "sql_protocolv1_o_6051921418418893_0819-204509-hill72", + "request" : { + "url" : "/sql/protocolv1/o/6051921418418893/0819-204509-hill72", + "method" : "POST", + "bodyPatterns" : [ { + "binaryEqualTo" : "gAEAAQAAAAxDbG9zZVNlc3Npb24AAAAGDAABDAABDAABCwABAAAAEKfwV+WkgE9OpcjjTe1EBEcLAAIAAAAQubO6CbNSTAO3UEUSrrqpQQAAAAA=" + } ] + }, + "response" : { + "status" : 200, + "base64Body" : "gAEAAgAAAAxDbG9zZVNlc3Npb24AAAAGDAAADAABCAABAAAAAAAAAA==", + "headers" : { + "x-request-id" : "54217fcb-de65-4401-a6df-7725d1f850f2", + "date" : "Tue, 11 Aug 2026 19:05:36 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-frame-options" : "SAMEORIGIN", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-content-type-options" : "nosniff", + "x-xss-protection" : "1; mode=block", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/x-thrift", + "server-timing" : "request_id;dur=0;desc=\"54217fcb-de65-4401-a6df-7725d1f850f2\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "b33f2be4-6601-4ab7-9d09-f2df83f83fd8", + "insertionIndex" : 9 +} \ No newline at end of file diff --git a/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-b57c5383-59c9-467b-b5ef-7d233f1e2e30.json b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-b57c5383-59c9-467b-b5ef-7d233f1e2e30.json new file mode 100644 index 000000000..d411eb967 --- /dev/null +++ b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-b57c5383-59c9-467b-b5ef-7d233f1e2e30.json @@ -0,0 +1,37 @@ +{ + "id" : "b57c5383-59c9-467b-b5ef-7d233f1e2e30", + "name" : "sql_protocolv1_o_6051921418418893_0819-204509-hill72", + "request" : { + "url" : "/sql/protocolv1/o/6051921418418893/0819-204509-hill72", + "method" : "POST", + "bodyPatterns" : [ { + "binaryEqualTo" : "gAEAAQAAAAtPcGVuU2Vzc2lvbgAAAAEMAAEIAAH////5DQAECwsAAAAACgUCAAAAAAAApQoMBQQLAAEAAAAFU1BBUksLAAIAAAAHZGVmYXVsdAACBQUBAAA=" + } ] + }, + "response" : { + "status" : 200, + "base64Body" : "gAEAAgAAAAtPcGVuU2Vzc2lvbgAAAAEMAAAMAAEIAAEAAAAAAAgAAgAApQoMAAMMAAELAAEAAAAQp/BX5aSAT06lyONN7UQERwsAAgAAABC5s7oJs1JMA7dQRRKuuqlBBg0BAAAACA0BAAClCgANAAQLCwAAAAAMBQQLAAEAAAAFc3BhcmsLAAIAAAAHZGVmYXVsdAACBQUBAAA=", + "headers" : { + "x-request-id" : "2ba47cf3-0d04-4a9d-baef-a60799395848", + "date" : "Tue, 11 Aug 2026 19:05:24 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-frame-options" : "SAMEORIGIN", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-content-type-options" : "nosniff", + "x-xss-protection" : "1; mode=block", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/x-thrift", + "server-timing" : "request_id;dur=0;desc=\"2ba47cf3-0d04-4a9d-baef-a60799395848\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "b57c5383-59c9-467b-b5ef-7d233f1e2e30", + "insertionIndex" : 14 +} \ No newline at end of file diff --git a/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-f02e77e1-4324-41e0-b558-006f609d39ff.json b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-f02e77e1-4324-41e0-b558-006f609d39ff.json new file mode 100644 index 000000000..7908ae2e9 --- /dev/null +++ b/src/test/resources/thriftserverapi/nativepreparedstatementbatchintegrationtests/testnativebatchsuccess/mappings/sql_protocolv1_o_6051921418418893_0819-204509-hill72-f02e77e1-4324-41e0-b558-006f609d39ff.json @@ -0,0 +1,38 @@ +{ + "id" : "f02e77e1-4324-41e0-b558-006f609d39ff", + "name" : "sql_protocolv1_o_6051921418418893_0819-204509-hill72", + "request" : { + "url" : "/sql/protocolv1/o/6051921418418893/0819-204509-hill72", + "method" : "POST", + "bodyPatterns" : [ { + "binaryEqualTo" : "gAEAAQAAABBFeGVjdXRlU3RhdGVtZW50AAAAAgwAAQwAAQwAAQsAAQAAABCn8FflpIBPTqXI403tRARHCwACAAAAELmzugmzUkwDt1BFEq66qUEAAAsAAgAAAEVEUk9QIFRBQkxFIElGIEVYSVNUUyBtYWluLmpkYmNfdGVzdF9zY2hlbWEubmF0aXZlX2JhdGNoX3N1Y2Nlc3NfdGFibGUKAAUAAAAAAAAAAAwFAQoAAQAAAAAAHoSACgACAAAAABghowAAAgUCAQIFAwECBQQBDAUGAgABAQIAAgECAAMBAgAEAQIABQEADwUIDAAAAAAAAA==" + } ] + }, + "response" : { + "status" : 200, + "base64Body" : "gAEAAgAAABBFeGVjdXRlU3RhdGVtZW50AAAAAgwAAAwAAQgAAQAAAAAADAACDAABCwABAAAAEI4LEETD+0HChPVYgjxZzKkLAAIAAAAQJ+Nqe/3oQve+tBbaWY40HgYNAQAAAAgAAgAAAAACAAMBAAwFAQwAAQwAAQgAAQAAAAAACAACAAAAAggNAgAAAAAADAACDAABCAABAAAAAAAMAAIPAAEMAAAAAQsAAQAAAAZSZXN1bHQMAAIPAAEMAAAAAQwAAQgAAQAAAAcAAAAIAAMAAAABCwAEAAAAAAAACAUBAAAAAAIFAgELBQMAAABI/////0AAAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAAAAAAAAAAACAUEAAAAAAoFBQAAAAAAAAAACgUGAAAAAAAAAAACBQcACA0BAAAABQoNBQAAAAAAAAAAAg0RAAINEwEADAADDAABCAABAAAAAAACAAIADAADCgABAAAAAAAAAAAPAAIMAAAAAA8FAQwAAAAAAAAMAAQMAAEIAAEAAAAAAAAAAAA=", + "headers" : { + "x-request-id" : "0ede753a-7357-4bc6-906d-0315f81a8f0a", + "date" : "Tue, 11 Aug 2026 19:05:28 GMT", + "server" : "databricks", + "x-databricks-popp-response-code-details" : "via_upstream", + "x-databricks-shard-debug" : "oregon-staging", + "x-databricks-popp-fast-path-routing-reason" : "not_eligible", + "x-frame-options" : "SAMEORIGIN", + "x-databricks-popp-shadow-routing-reason" : "spog-domain-checker-false", + "x-databricks-upstream-cluster" : "oregon-staging-h2", + "x-databricks-org-id" : "6051921418418893", + "strict-transport-security" : "max-age=31536000; includeSubDomains; preload", + "x-databricks-query-rejection-reason" : "{\"query_stats\":{}}", + "x-content-type-options" : "nosniff", + "x-xss-protection" : "1; mode=block", + "x-databricks-popp-routing-reason" : "deployment-name", + "content-type" : "application/x-thrift", + "server-timing" : "request_id;dur=0;desc=\"0ede753a-7357-4bc6-906d-0315f81a8f0a\", client_protocol;dur=0;desc=\"HTTP/1.1\"", + "alt-svc" : "h3=\":443\"; ma=86400, h3-29=\":443\"; ma=86400", + "x-databricks-apiproxy-response-code-details" : "via_upstream" + } + }, + "uuid" : "f02e77e1-4324-41e0-b558-006f609d39ff", + "insertionIndex" : 13 +} \ No newline at end of file From 2d0f4da302b183f15de9804af1c12346b5531bf4 Mon Sep 17 00:00:00 2001 From: Sreekanth Vadigi Date: Tue, 11 Aug 2026 19:52:08 +0000 Subject: [PATCH 7/7] Strengthen native batch order assertion Validate exact parameter values and array order so replay coverage detects reordered native batches. Signed-off-by: Sreekanth Vadigi --- ...tivePreparedStatementBatchIntegrationTests.java | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/NativePreparedStatementBatchIntegrationTests.java b/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/NativePreparedStatementBatchIntegrationTests.java index c09dff9db..906dc7b6d 100644 --- a/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/NativePreparedStatementBatchIntegrationTests.java +++ b/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/NativePreparedStatementBatchIntegrationTests.java @@ -86,16 +86,16 @@ void testNativeBatchSuccess() throws SQLException { + "\"," + "\"parameter_sets\":[" + "{\"parameters\":[" - + "{\"ordinal\":0,\"type\":\"INT\"}," - + "{\"ordinal\":1,\"type\":\"STRING\"}," - + "{\"ordinal\":2,\"type\":\"STRING\"}]}," + + "{\"ordinal\":0,\"type\":\"INT\",\"value\":\"1\"}," + + "{\"ordinal\":1,\"type\":\"STRING\",\"value\":\"first-a\"}," + + "{\"ordinal\":2,\"type\":\"STRING\",\"value\":\"first-b\"}]}," + "{\"parameters\":[" - + "{\"ordinal\":0,\"type\":\"INT\"}," - + "{\"ordinal\":1,\"type\":\"STRING\"}," - + "{\"ordinal\":2,\"type\":\"STRING\"}]}" + + "{\"ordinal\":0,\"type\":\"INT\",\"value\":\"2\"}," + + "{\"ordinal\":1,\"type\":\"STRING\",\"value\":\"second-a\"}," + + "{\"ordinal\":2,\"type\":\"STRING\",\"value\":\"second-b\"}]}" + "]" + "}", - true, + false, true))); } }