Skip to content

Commit 4c66db2

Browse files
Implement executeBatch (#532)
- Created `BatchExecutor` to manage batch SQL command execution - Serially execute SQL commands until success or a failure - Returns the update counts until last successful command - Throws batch exception as per JDBC specification - Introduced `maxBatchSize` to restrict the number of commands in a batch. Throws `SQLException` when attempting to add commands beyond the specified limit - Logs telemetry data at the debug level to facilitate performance monitoring and debugging - Lint changes
1 parent 4cb784a commit 4c66db2

15 files changed

+638
-158
lines changed

src/main/java/com/databricks/jdbc/api/IDatabricksConnection.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@ public interface IDatabricksConnection {
2727
/** Opens the connection and initiates the underlying session */
2828
void open() throws DatabricksSQLException;
2929

30+
/** Returns the connection context associated with the connection. */
31+
IDatabricksConnectionContext getConnectionContext();
32+
3033
/** Returns the statement handle for given statement-Id */
3134
Statement getStatement(String statementId) throws SQLException;
3235
}

src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -213,4 +213,7 @@ public static AuthMech parseAuthMech(String authMech) {
213213

214214
/** Returns the SSL trust store type of the trust store file. */
215215
String getSSLTrustStoreType();
216+
217+
/** Returns the maximum number of commands that can be executed in a single batch. */
218+
int getMaxBatchSize();
216219
}

src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,10 +30,11 @@
3030
/** Implementation for Databricks specific connection. */
3131
public class DatabricksConnection implements IDatabricksConnection, Connection {
3232
private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(DatabricksConnection.class);
33-
private IDatabricksSession session;
33+
private final IDatabricksSession session;
3434
private final Set<IDatabricksStatementInternal> statementSet = ConcurrentHashMap.newKeySet();
3535
private SQLWarning warnings = null;
3636
private volatile IDatabricksUCVolumeClient ucVolumeClient = null;
37+
private final IDatabricksConnectionContext connectionContext;
3738

3839
/**
3940
* Creates an instance of Databricks connection for given connection context.
@@ -42,13 +43,15 @@ public class DatabricksConnection implements IDatabricksConnection, Connection {
4243
*/
4344
public DatabricksConnection(IDatabricksConnectionContext connectionContext)
4445
throws DatabricksSQLException {
46+
this.connectionContext = connectionContext;
4547
this.session = new DatabricksSession(connectionContext);
4648
}
4749

4850
@VisibleForTesting
4951
public DatabricksConnection(
5052
IDatabricksConnectionContext connectionContext, IDatabricksClient testDatabricksClient)
5153
throws DatabricksSQLException {
54+
this.connectionContext = connectionContext;
5255
this.session = new DatabricksSession(connectionContext, testDatabricksClient);
5356
UserAgentManager.setUserAgent(connectionContext);
5457
}
@@ -526,6 +529,11 @@ public IDatabricksUCVolumeClient getUCVolumeClient() {
526529
return ucVolumeClient;
527530
}
528531

532+
@Override
533+
public IDatabricksConnectionContext getConnectionContext() {
534+
return connectionContext;
535+
}
536+
529537
/**
530538
* This function creates the exception message for the failed setClientInfo command
531539
*

src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -581,6 +581,11 @@ public String getSSLTrustStoreType() {
581581
return getParameter(DatabricksJdbcUrlParams.SSL_TRUST_STORE_TYPE);
582582
}
583583

584+
@Override
585+
public int getMaxBatchSize() {
586+
return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.MAX_BATCH_SIZE));
587+
}
588+
584589
private static boolean nullOrEmptyString(String s) {
585590
return s == null || s.isEmpty();
586591
}

src/main/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatement.java

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import com.databricks.jdbc.common.util.DatabricksTypeUtil;
1010
import com.databricks.jdbc.exception.DatabricksSQLException;
1111
import com.databricks.jdbc.exception.DatabricksSQLFeatureNotImplementedException;
12+
import com.databricks.jdbc.exception.DatabricksSQLFeatureNotSupportedException;
1213
import com.databricks.jdbc.log.JdbcLogger;
1314
import com.databricks.jdbc.log.JdbcLoggerFactory;
1415
import java.io.ByteArrayOutputStream;
@@ -38,9 +39,8 @@ public DatabricksPreparedStatement(DatabricksConnection connection, String sql)
3839
super(connection);
3940
this.sql = sql;
4041
this.interpolateParameters =
41-
connection.getSession().getConnectionContext().supportManyParameters()
42-
|| connection.getSession().getConnectionContext().getComputeResource()
43-
instanceof AllPurposeCluster;
42+
connection.getConnectionContext().supportManyParameters()
43+
|| connection.getConnectionContext().getComputeResource() instanceof AllPurposeCluster;
4444
this.databricksParameterMetaData = new DatabricksParameterMetaData();
4545
this.databricksBatchParameterMetaData = new ArrayList<>();
4646
}
@@ -552,6 +552,15 @@ public boolean execute(String sql, String[] columnNames) throws SQLException {
552552
throw new DatabricksSQLException("Method not supported in PreparedStatement");
553553
}
554554

555+
/** {@inheritDoc} */
556+
@Override
557+
public void addBatch(String sql) throws SQLException {
558+
LOGGER.debug(String.format("public void addBatch(String sql = {%s})", sql));
559+
checkIfClosed();
560+
throw new DatabricksSQLFeatureNotSupportedException(
561+
"Method not supported: addBatch(String sql)");
562+
}
563+
555564
private void checkLength(long targetLength, long sourceLength) throws SQLException {
556565
if (targetLength != sourceLength) {
557566
String errorMessage =

0 commit comments

Comments
 (0)