Merge remote-tracking branch 'databricks/main' into fetch-requests-columnar-view

Author: jayantsing-db

NEXT_CHANGELOG.md

@@ -8,11 +8,15 @@
 - **SQL Scripting support**: Added support for [SQL Scripting](https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-scripting)
 - Added a client property `enableVolumeOperations` to enable  GET/PUT/REMOVE volume operations on a stream. For backward compatibility, allowedVolumeIngestionPaths can also be used for REMOVE operation.
 - Support for fetching schemas across all catalogs (when catalog is specified as null or a wildcard) in `DatabaseMetaData#getSchemas` API in SQL Execution mode.
+- **Configurable SQL validation in isValid()**: Added `EnableSQLValidationForIsValid` connection property to control whether `isValid()` method executes an actual SQL query for server-side validation. Default value is 0.
+- Implement multi-row INSERT batching optimization for prepared statements to improve performance when executing large batches of INSERT operations.
+- Implement lazy/incremental fetching for columnar results when using Databricks JDBC in Thrift mode without Arrow support. The change modifies the behavior from buffering entire result sets in memory to maintaining only a limited number of rows at a time, reducing peak heap memory usage and preventing OutOfMemory errors.
 
 ### Updated
 - Databricks SDK dependency upgraded to latest version 0.60.0
 
 ### Fixed
--
+- Integrated Azure U2M flow into driver for improved stability.
+- Fixed `ResultSet.getString` for Boolean columns in Metadata result set.
 ---
 *Note: When making changes, please add your change under the appropriate section with a brief description.* 

README.md

@@ -91,6 +91,7 @@ Optional parameters:
 - `OAuth2RedirectUrlPort` - Ports for redirect URL (default: 8020)
 - `EnableOIDCDiscovery` - Enable OIDC discovery (default: 1)
 - `OAuthDiscoveryURL` - OIDC discovery endpoint (default: /oidc/.well-known/oauth-authorization-server)
+- `EnableSQLValidationForIsValid` - Enable SQL query based validation in `isValid()` connection checks (default: 0)
 
 ### Logging
 

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

@@ -414,7 +414,21 @@ public SQLXML createSQLXML() throws SQLException {
   @Override
   public boolean isValid(int timeout) throws SQLException {
     ValidationUtil.checkIfNonNegative(timeout, "timeout");
-    return !isClosed();
+    if (isClosed()) {
+      return false;
+    }
+    if (connectionContext.getEnableSQLValidationForIsValid()) {
+      try (Statement stmt = createStatement()) {
+        stmt.setQueryTimeout(timeout);
+        // This is a lightweight query to check if the connection is valid
+        stmt.execute("SELECT VERSION()");
+        return true;
+      } catch (Exception e) {
+        LOGGER.debug("Validation failed for isValid(): {}", e.getMessage());
+        return false;
+      }
+    }
+    return true;
   }
 
   /**

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

@@ -215,6 +215,11 @@ public String getHostUrl() throws DatabricksParsingException {
     }
   }
 
+  @Override
+  public String getHost() {
+    return this.host;
+  }
+
   @Override
   public IDatabricksComputeResource getComputeResource() {
     return computeResource;
@@ -225,6 +230,12 @@ public String getHttpPath() {
     return getParameter(DatabricksJdbcUrlParams.HTTP_PATH);
   }
 
+  public boolean getEnableSQLValidationForIsValid() {
+    LOGGER.debug("String getEnableSQLValidationForIsValid()");
+    return getParameter(DatabricksJdbcUrlParams.ENABLE_SQL_VALIDATION_FOR_IS_VALID, "0")
+        .equals("1");
+  }
+
   @Override
   public String getHostForOAuth() {
     return this.host;
@@ -1019,4 +1030,9 @@ public int getTelemetryFlushIntervalInMilliseconds() {
     return Math.max(
         1000, Integer.parseInt(getParameter(DatabricksJdbcUrlParams.TELEMETRY_FLUSH_INTERVAL)));
   }
+
+  @Override
+  public boolean isBatchedInsertsEnabled() {
+    return getParameter(DatabricksJdbcUrlParams.ENABLE_BATCHED_INSERTS).equals("1");
+  }
 }

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

@@ -6,7 +6,7 @@
 import com.databricks.jdbc.common.util.WrapperUtil;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
-import com.databricks.sdk.service.sql.ColumnInfoTypeName;
+import com.databricks.jdbc.model.core.ColumnInfoTypeName;
 import java.sql.ParameterMetaData;
 import java.sql.SQLException;
 import java.util.HashMap;

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

@@ -7,8 +7,10 @@
 import static com.databricks.jdbc.common.util.SQLInterpolator.surroundPlaceholdersWithQuotes;
 import static com.databricks.jdbc.common.util.ValidationUtil.throwErrorIfNull;
 
+import com.databricks.jdbc.common.DatabricksJdbcConstants;
 import com.databricks.jdbc.common.StatementType;
 import com.databricks.jdbc.common.util.DatabricksTypeUtil;
+import com.databricks.jdbc.common.util.InsertStatementParser;
 import com.databricks.jdbc.exception.*;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
@@ -88,6 +90,138 @@ public int[] executeBatch() throws DatabricksBatchUpdateException {
   @Override
   public long[] executeLargeBatch() throws DatabricksBatchUpdateException {
     LOGGER.debug("public long executeLargeBatch()");
+
+    if (databricksBatchParameterMetaData.isEmpty()) {
+      return new long[0];
+    }
+
+    // Try to optimize INSERT statements with multi-row batching
+    if (canUseBatchedInsert()) {
+      return executeBatchedInsert();
+    } else {
+      // Fall back to individual execution for non-INSERT or incompatible statements
+      return executeIndividualStatements();
+    }
+  }
+
+  /**
+   * Checks if the current batch can be optimized using multi-row INSERT. All statements must be
+   * compatible INSERT operations.
+   *
+   * <p>A batch is eligible for multi-row INSERT optimization when:
+   *
+   * <ul>
+   *   <li>The EnableBatchedInserts connection property is enabled (default: true)
+   *   <li>The SQL statement is an INSERT operation
+   *   <li>The INSERT can be parsed successfully (has table name and column list)
+   *   <li>The batch contains parameter sets for multiple rows
+   * </ul>
+   *
+   * <p>Compatible INSERT operations target the same table with the same columns in the same order.
+   * When compatible, multiple individual INSERTs like:
+   *
+   * <pre>
+   *   INSERT INTO users (id, name) VALUES (?, ?)  -- with parameters [1, "Alice"]
+   *   INSERT INTO users (id, name) VALUES (?, ?)  -- with parameters [2, "Bob"]
+   * </pre>
+   *
+   * Are combined into a single multi-row INSERT:
+   *
+   * <pre>
+   *   INSERT INTO users (id, name) VALUES (?, ?), (?, ?)  -- with parameters [1, "Alice", 2, "Bob"]
+   * </pre>
+   */
+  private boolean canUseBatchedInsert() {
+    // Check if batched inserts are enabled via connection property
+    if (!connection.getConnectionContext().isBatchedInsertsEnabled()) {
+      return false;
+    }
+
+    // Use strict exception-based parsing for better error handling
+    try {
+      InsertStatementParser.parseInsertStrict(sql);
+      return !databricksBatchParameterMetaData.isEmpty();
+    } catch (Exception e) {
+      // Not a valid INSERT statement suitable for batching
+      return false;
+    }
+  }
+
+  /** Executes the batch as a single multi-row INSERT statement. */
+  private long[] executeBatchedInsert() throws DatabricksBatchUpdateException {
+    LOGGER.debug("Executing batched INSERT with {} rows", databricksBatchParameterMetaData.size());
+
+    try {
+      InsertStatementParser.InsertInfo insertInfo = InsertStatementParser.parseInsertStrict(sql);
+
+      // Calculate how many rows we can fit in one chunk based on parameter limit
+      int parametersPerRow = insertInfo.getColumnCount();
+      int maxRowsPerChunk = DatabricksJdbcConstants.MAX_QUERY_PARAMETERS / parametersPerRow;
+
+      // Ensure we have at least 1 row per chunk
+      if (maxRowsPerChunk < 1) {
+        maxRowsPerChunk = 1;
+      }
+
+      long[] allUpdateCounts = new long[databricksBatchParameterMetaData.size()];
+      int processedRows = 0;
+
+      // Process batches in chunks
+      for (int startIndex = 0;
+          startIndex < databricksBatchParameterMetaData.size();
+          startIndex += maxRowsPerChunk) {
+        int endIndex =
+            Math.min(startIndex + maxRowsPerChunk, databricksBatchParameterMetaData.size());
+        int chunkSize = endIndex - startIndex;
+
+        LOGGER.debug("Processing chunk {}-{} ({} rows)", startIndex + 1, endIndex, chunkSize);
+
+        // Generate multi-row SQL for this chunk
+        String multiRowSql = InsertStatementParser.generateMultiRowInsert(insertInfo, chunkSize);
+
+        // Combine parameters for this chunk
+        Map<Integer, ImmutableSqlParameter> chunkParams = new HashMap<>();
+        int paramIndex = 1;
+
+        for (int i = startIndex; i < endIndex; i++) {
+          DatabricksParameterMetaData batchParams = databricksBatchParameterMetaData.get(i);
+          Map<Integer, ImmutableSqlParameter> rowParams = batchParams.getParameterBindings();
+          for (int j = 1; j <= rowParams.size(); j++) {
+            if (rowParams.containsKey(j)) {
+              chunkParams.put(paramIndex++, rowParams.get(j));
+            }
+          }
+        }
+
+        // Execute this chunk
+        executeInternal(multiRowSql, chunkParams, StatementType.UPDATE, false);
+
+        // Set update counts for this chunk (each row typically affects 1 row)
+        for (int i = startIndex; i < endIndex; i++) {
+          allUpdateCounts[i] = 1;
+        }
+
+        processedRows += chunkSize;
+      }
+
+      LOGGER.debug("Successfully processed {} rows in chunks", processedRows);
+      return allUpdateCounts;
+
+    } catch (Exception e) {
+      LOGGER.error("Error executing batched INSERT: {}", e.getMessage(), e);
+      long[] failedCounts = new long[databricksBatchParameterMetaData.size()];
+      for (int i = 0; i < failedCounts.length; i++) {
+        failedCounts[i] = Statement.EXECUTE_FAILED;
+      }
+      throw new DatabricksBatchUpdateException(
+          e.getMessage(), DatabricksDriverErrorCode.BATCH_EXECUTE_EXCEPTION, failedCounts);
+    }
+  }
+
+  /** Executes batch statements individually (fallback method). */
+  private long[] executeIndividualStatements() throws DatabricksBatchUpdateException {
+    LOGGER.debug(
+        "Executing batch individually with {} statements", databricksBatchParameterMetaData.size());
     long[] largeUpdateCount = new long[databricksBatchParameterMetaData.size()];
 
     for (int sqlQueryIndex = 0;

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

@@ -624,7 +624,7 @@ public boolean isFirst() throws SQLException {
   public boolean isLast() throws SQLException {
     checkIfClosed();
     if (executionResult instanceof LazyThriftResult) {
-      return !executionResult.hasNext();
+      return executionResult.getCurrentRow() >= 0 && !executionResult.hasNext();
     }
     return executionResult.getCurrentRow() == resultSetMetaData.getTotalRows() - 1;
   }

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

@@ -20,10 +20,10 @@
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.client.thrift.generated.*;
+import com.databricks.jdbc.model.core.ColumnInfo;
+import com.databricks.jdbc.model.core.ColumnInfoTypeName;
 import com.databricks.jdbc.model.core.ColumnMetadata;
 import com.databricks.jdbc.model.core.ResultManifest;
-import com.databricks.sdk.service.sql.ColumnInfo;
-import com.databricks.sdk.service.sql.ColumnInfoTypeName;
 import com.google.common.collect.ImmutableList;
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;

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

@@ -692,6 +692,14 @@ static boolean isSelectQuery(String query) {
     return SELECT_PATTERN.matcher(trimmedQuery).find();
   }
 
+  static boolean isInsertQuery(String query) {
+    if (query == null || query.trim().isEmpty()) {
+      return false;
+    }
+    String trimmedQuery = trimCommentsAndWhitespaces(query);
+    return INSERT_PATTERN.matcher(trimmedQuery).find();
+  }
+
   DatabricksResultSet executeInternal(
       String sql,
       Map<Integer, ImmutableSqlParameter> params,

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

@@ -218,6 +218,11 @@ private void loadCurrentBatch() throws DatabricksSQLException {
         totalRowsFetched);
   }
 
+  /**
+   * Fetches the next batch of data from the server and creates columnar views.
+   *
+   * @throws DatabricksSQLException if the fetch operation fails
+   */
   private void fetchNextBatch() throws DatabricksSQLException {
     try {
       LOGGER.debug("Fetching next batch, current total rows fetched: {}", totalRowsFetched);