diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md index 6f102169fa..7854ec0e0c 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.md +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -26,7 +26,7 @@ If available, share redacted client side logs - OS: [e.g. Windows] - Java version [e.g. Java 21] - Java vendor [e.g. OpenJDK] - - Driver Version [e.g. 3.0.7] + - Driver Version [e.g. 3.1.1] - BI Tool (if used) [e.g. DBeaver] - BI Tool version (if applicable) [e.g. 24.3.5] diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a5b610619..6e1fd1cee7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,18 @@ # Version Changelog +## [v3.1.1] - 2026-01-07 + +### Added +- Added token caching for all authentication providers to reduce token endpoint calls. +- We will be rolling out the use of Databricks SQL Execution API by default for queries submitted on DBSQL. To continue using Databricks Thrift Server backend for execution, set `UseThriftClient` to `1`. + +### Updated +- Changed default value of `IgnoreTransactions` from `0` to `1` to disable multi-statement transactions by default. Preview participants can opt-in by setting `IgnoreTransactions=0`. Also updated `supportsTransactions()` to respect this flag. + +### Fixed +- [PECOBLR-1131] Fix incorrect refetching of expired CloudFetch links when using Thrift protocol. +- Fixed logging to respect params when the driver is shaded. +- Fixed `isWildcard` to return true only when the value is `*` + ## [v3.0.7] - 2025-12-18 ### Updated diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 7883b672ae..434c4a84f7 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -3,15 +3,27 @@ ## [Unreleased] ### Added -- Added token caching for all authentication providers to reduce token endpoint calls. +- Added streaming prefetch mode for Thrift inline results (columnar and Arrow) with background batch prefetching and configurable sliding window for improved throughput. +- Added `EnableInlineStreaming` connection parameter to enable/disable streaming mode (default: enabled). +- Added `ThriftMaxBatchesInMemory` connection parameter to control the sliding window size for streaming (default: 3). +- Added support for disabling CloudFetch via `EnableQueryResultDownload=0` to use inline Arrow results instead. +- Added `EnableMetricViewSupport` connection parameter to enable/disable Metric View table type (default: disabled). +- Added `NonRowcountQueryPrefixes` connection parameter to specify comma-separated query prefixes that should return result sets instead of row counts. ### Updated +- Enhanced error logging for token exchange failures. +- Geospatial column type names now include SRID information (e.g., `GEOMETRY(4326)` instead of `GEOMETRY`). +- Changed default value of `IgnoreTransactions` from `0` to `1` to disable multi-statement transactions by default. Preview participants can opt-in by setting `IgnoreTransactions=0`. Also updated `supportsTransactions()` to respect this flag. +- Implemented lazy loading for inline Arrow results, fetching arrow batches on demand instead of all at once. This improves memory usage and initial response time for large result sets when using the Thrift protocol with Arrow format. +- **Enhanced `enableMultipleCatalogSupport` behavior**: When this parameter is disabled (`enableMultipleCatalogSupport=0`), metadata operations (such as `getSchemas()`, `getTables()`, `getColumns()`, etc.) now return results only when the catalog parameter is either `null` or matches the current catalog. For any other catalog name, an empty result set is returned. This ensures metadata queries are restricted to the current catalog context. When enabled (`enableMultipleCatalogSupport=1`), metadata operations continue to work across all accessible catalogs. ### Fixed - -- [PECOBLR-1131] Fix incorrect refetching of expired CloudFetch links when using Thrift protocol. -- Fixed logging to respect params when the driver is shaded. -- Fixed `isWildcard` to return true only when the value is `*` +- Fixed timeout exception handling to throw `SQLTimeoutException` instead of `DatabricksHttpException` when queries timeout during result fetching phase. This completes the timeout exception fix to handle both query execution polling and result fetching phases. +- Fixed `getTypeInfo()` and `getClientInfoProperties()` to return fresh ResultSet instances on each call instead of shared static instances. This resolves issues where calling these methods multiple times would fail due to exhausted cursor state (Issue #1178). +- Fixed complex data type metadata support when retrieving 0 rows in Arrow format +- Normalized TIMESTAMP_NTZ to TIMESTAMP in Thrift path for consistency with SEA behavior +- Fixed complex types not being returned as objects in SEA Inline mode when `EnableComplexDatatypeSupport=true`. +- Fixed `StringIndexOutOfBoundsException` when parsing complex data types in Thrift CloudFetch mode. The issue occurred when metadata contained incomplete type information (e.g., "ARRAY" instead of "ARRAY"). Now retrieves complete type information from Arrow metadata. --- *Note: When making changes, please add your change under the appropriate section diff --git a/README.md b/README.md index 1bf6b699cf..28911d1daa 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ Add the following dependency to your `pom.xml`: com.databricks databricks-jdbc - 3.0.7 + 3.1.1 ``` diff --git a/development/.release-freeze.json b/development/.release-freeze.json index eb9f69be91..f12ea3ce8d 100644 --- a/development/.release-freeze.json +++ b/development/.release-freeze.json @@ -1,4 +1,4 @@ { - "freeze": false, - "reason": "Release freeze disabled" + "freeze": true, + "reason": "Release version 3.2.1" } diff --git a/pom.xml b/pom.xml index b5f6ae56bc..72e42dfe84 100644 --- a/pom.xml +++ b/pom.xml @@ -5,7 +5,7 @@ com.databricks databricks-jdbc - 3.0.7 + 3.1.1 jar Databricks JDBC Driver Databricks JDBC Driver. diff --git a/release-notes.txt b/release-notes.txt index bd8a876b74..1162ed0430 100644 --- a/release-notes.txt +++ b/release-notes.txt @@ -6,6 +6,21 @@ The release notes summarize enhancements, new features, known issues, workflow c Version History ============================================================== +3.1.1 ======================================================================== +Released 2026-01-07 + +Enhancements & New Features + + * Added token caching for all authentication providers to reduce token endpoint calls. + * We will be rolling out the use of Databricks SQL Execution API by default for queries submitted on DBSQL. To continue using Databricks Thrift Server backend for execution, set `UseThriftClient` to `1`. + * Changed default value of `IgnoreTransactions` from `0` to `1` to disable multi-statement transactions by default. Preview participants can opt-in by setting `IgnoreTransactions=0`. Also updated `supportsTransactions()` to respect this flag. + +Resolved Issues + + * [PECOBLR-1131] Fix incorrect refetching of expired CloudFetch links when using Thrift protocol. + * Fixed logging to respect params when the driver is shaded. + * Fixed `isWildcard` to return true only when the value is `*` + 3.0.7 ======================================================================== Released 2025-12-18 diff --git a/src/main/java/com/databricks/client/jdbc/DataSource.java b/src/main/java/com/databricks/client/jdbc/DataSource.java index 729703c1f8..7b9f8744ef 100644 --- a/src/main/java/com/databricks/client/jdbc/DataSource.java +++ b/src/main/java/com/databricks/client/jdbc/DataSource.java @@ -4,7 +4,6 @@ import com.databricks.jdbc.common.DatabricksJdbcConstants; import com.databricks.jdbc.common.DatabricksJdbcUrlParams; -import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.pooling.DatabricksPooledConnection; @@ -37,13 +36,13 @@ public DataSource(Driver driver) { } @Override - public Connection getConnection() throws DatabricksSQLException { + public Connection getConnection() throws SQLException { LOGGER.debug("public Connection getConnection()"); return getConnection(this.getUsername(), this.getPassword()); } @Override - public Connection getConnection(String username, String password) throws DatabricksSQLException { + public Connection getConnection(String username, String password) throws SQLException { LOGGER.debug("public Connection getConnection(String, String)"); if (username != null) { setUsername(username); @@ -55,14 +54,13 @@ public Connection getConnection(String username, String password) throws Databri } @Override - public PooledConnection getPooledConnection() throws DatabricksSQLException { + public PooledConnection getPooledConnection() throws SQLException { LOGGER.debug("public PooledConnection getPooledConnection()"); return new DatabricksPooledConnection(getConnection()); } @Override - public PooledConnection getPooledConnection(String user, String password) - throws DatabricksSQLException { + public PooledConnection getPooledConnection(String user, String password) throws SQLException { LOGGER.debug("public PooledConnection getPooledConnection(String, String)"); return new DatabricksPooledConnection(getConnection(user, password)); } diff --git a/src/main/java/com/databricks/client/jdbc/Driver.java b/src/main/java/com/databricks/client/jdbc/Driver.java index 83cb1861f7..31e0d7b2cb 100644 --- a/src/main/java/com/databricks/client/jdbc/Driver.java +++ b/src/main/java/com/databricks/client/jdbc/Driver.java @@ -47,7 +47,7 @@ public boolean acceptsURL(String url) { } @Override - public Connection connect(String url, Properties info) throws DatabricksSQLException { + public Connection connect(String url, Properties info) throws SQLException { if (!acceptsURL(url)) { // Return null connection if URL is not accepted - as per JDBC standard. return null; diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java index 06385ebcc6..248d7285bf 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java @@ -63,7 +63,7 @@ public DatabricksConnection( } @Override - public void open() throws DatabricksSQLException { + public void open() throws SQLException { this.session.open(); } @@ -252,13 +252,13 @@ private boolean fetchAutoCommitStateFromServer() throws SQLException { ResultSet rs = statement.executeQuery("SET AUTOCOMMIT"); if (rs.next()) { - // The result should contain the value = "true" or "false" + // The result may contain "true"/"false" or "1"/"0" depending on the server String value = rs.getString(1); // Column 1: value LOGGER.debug( "Fetched autoCommit state from server: value={}. Updating session cache.", value); - boolean autoCommitState = "true".equalsIgnoreCase(value); + boolean autoCommitState = "true".equalsIgnoreCase(value) || "1".equals(value); // Update the session cache with the server value session.setAutoCommit(autoCommitState); @@ -382,13 +382,20 @@ public void rollback() throws SQLException { return; } + // Rollback is not valid when auto-commit is enabled (no active transaction) + if (getAutoCommit()) { + throw new DatabricksTransactionException( + "Cannot use rollback while Connection is in auto-commit mode.", + new SQLException("Cannot use rollback while Connection is in auto-commit mode."), + DatabricksDriverErrorCode.TRANSACTION_ROLLBACK_ERROR); + } + // Execute ROLLBACK command Statement statement = null; try { statement = createStatement(); statement.execute("ROLLBACK"); // Note: Server auto-starts new transaction if autocommit=false - // Note: ROLLBACK is more forgiving - typically succeeds even on unexpected states } catch (SQLException e) { LOGGER.error(e, "Error {} while rolling back transaction", e.getMessage()); @@ -401,7 +408,7 @@ public void rollback() throws SQLException { } @Override - public void close() throws DatabricksSQLException { + public void close() throws SQLException { LOGGER.debug("public void close()"); for (IDatabricksStatementInternal statement : statementSet) { statement.close(false); diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java index eafbefde82..4ecedb432d 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -457,7 +457,7 @@ public DatabricksClientType getClientTypeFromContext() { // Check if circuit breaker is open due to recent 429 rate limit failures if (SeaCircuitBreakerManager.isCircuitOpen()) { long remainingMs = SeaCircuitBreakerManager.getTimeRemainingMs(); - LOGGER.info( + LOGGER.debug( "SEA circuit breaker is OPEN due to recent 429 rate limit failures. " + "Using THRIFT client. Circuit will close in {} ({}ms)", SeaCircuitBreakerManager.getTimeRemainingFormatted(), @@ -469,7 +469,7 @@ public DatabricksClientType getClientTypeFromContext() { return DatabricksClientType.THRIFT; } // Check if CloudFetch is disabled - Thrift is required for inline mode - if (!Objects.equals(getParameter(DatabricksJdbcUrlParams.ENABLE_CLOUD_FETCH), "1")) { + if (!isCloudFetchEnabled()) { return DatabricksClientType.THRIFT; } // Check feature flag to determine if SEA client should be enabled @@ -1162,9 +1162,19 @@ public boolean isBatchedInsertsEnabled() { return getParameter(DatabricksJdbcUrlParams.ENABLE_BATCHED_INSERTS).equals("1"); } + @Override + public List getNonRowcountQueryPrefixes() { + String prefixesStr = getParameter(DatabricksJdbcUrlParams.NON_ROWCOUNT_QUERY_PREFIXES); + return Arrays.stream(prefixesStr.split(",")) + .map(String::trim) + .map(String::toUpperCase) + .filter(s -> !s.isEmpty()) + .collect(Collectors.toList()); + } + @Override public boolean getIgnoreTransactions() { - return getParameter(DatabricksJdbcUrlParams.IGNORE_TRANSACTIONS, "0").equals("1"); + return getParameter(DatabricksJdbcUrlParams.IGNORE_TRANSACTIONS).equals("1"); } @Override @@ -1198,6 +1208,27 @@ public boolean isStreamingChunkProviderEnabled() { return getParameter(DatabricksJdbcUrlParams.ENABLE_STREAMING_CHUNK_PROVIDER).equals("1"); } + @Override + public boolean isInlineStreamingEnabled() { + return getParameter(DatabricksJdbcUrlParams.ENABLE_INLINE_STREAMING).equals("1"); + } + + @Override + public boolean isCloudFetchEnabled() { + return getParameter(DatabricksJdbcUrlParams.ENABLE_CLOUD_FETCH).equals("1"); + } + + @Override + public int getThriftMaxBatchesInMemory() { + try { + return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.THRIFT_MAX_BATCHES_IN_MEMORY)); + } catch (NumberFormatException e) { + LOGGER.warn("Invalid value for ThriftMaxBatchesInMemory, using default value"); + return Integer.parseInt( + DatabricksJdbcUrlParams.THRIFT_MAX_BATCHES_IN_MEMORY.getDefaultValue()); + } + } + @Override public int getLinkPrefetchWindow() { return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.LINK_PREFETCH_WINDOW)); diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaData.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaData.java index f79c581ec1..62170b80a2 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaData.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaData.java @@ -2,7 +2,6 @@ import static com.databricks.jdbc.common.MetadataResultConstants.*; import static com.databricks.jdbc.dbclient.impl.common.CommandConstants.METADATA_STATEMENT_ID; -import static com.databricks.jdbc.dbclient.impl.sqlexec.ResultConstants.CLIENT_INFO_PROPERTIES_RESULT; import com.databricks.jdbc.api.impl.converters.ConverterHelper; import com.databricks.jdbc.api.internal.IDatabricksConnectionInternal; @@ -834,7 +833,7 @@ public int getDefaultTransactionIsolation() throws SQLException { public boolean supportsTransactions() throws SQLException { LOGGER.debug("public boolean supportsTransactions()"); throwExceptionIfConnectionIsClosed(); - return true; + return !session.getConnectionContext().getIgnoreTransactions(); } @Override @@ -1529,7 +1528,7 @@ public boolean autoCommitFailureClosesAllResultSets() throws SQLException { public ResultSet getClientInfoProperties() throws SQLException { LOGGER.debug("public ResultSet getClientInfoProperties()"); throwExceptionIfConnectionIsClosed(); - return CLIENT_INFO_PROPERTIES_RESULT; + return metadataResultSetBuilder.getClientInfoPropertiesResult(); } @Override 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 4f250874b4..48ab4cc5d4 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatement.java @@ -339,7 +339,7 @@ public boolean execute() throws SQLException { checkIfClosed(); checkIfBatchOperation(); interpolateIfRequiredAndExecute(StatementType.SQL); - return shouldReturnResultSet(sql); + return shouldReturnResultSetWithConfig(sql); } @Override 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 8b41cb5df2..14f4912e1d 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -1,14 +1,18 @@ package com.databricks.jdbc.api.impl; import static com.databricks.jdbc.common.DatabricksJdbcConstants.EMPTY_STRING; +import static com.databricks.jdbc.common.util.DatabricksThriftUtil.getArrowMetadata; import static com.databricks.jdbc.common.util.DatabricksTypeUtil.*; import com.databricks.jdbc.api.IDatabricksResultSet; import com.databricks.jdbc.api.IExecutionStatus; import com.databricks.jdbc.api.impl.arrow.ArrowStreamResult; import com.databricks.jdbc.api.impl.arrow.ChunkProvider; +import com.databricks.jdbc.api.impl.arrow.LazyThriftInlineArrowResult; +import com.databricks.jdbc.api.impl.arrow.StreamingInlineArrowResult; import com.databricks.jdbc.api.impl.converters.ConverterHelper; import com.databricks.jdbc.api.impl.converters.ObjectConverter; +import com.databricks.jdbc.api.impl.thrift.StreamingColumnarResult; import com.databricks.jdbc.api.impl.volume.VolumeOperationResult; import com.databricks.jdbc.api.internal.IDatabricksResultSetInternal; import com.databricks.jdbc.api.internal.IDatabricksSession; @@ -24,12 +28,9 @@ import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; -import com.databricks.jdbc.model.core.ColumnMetadata; -import com.databricks.jdbc.model.core.ResultData; -import com.databricks.jdbc.model.core.ResultManifest; -import com.databricks.jdbc.model.core.StatementStatus; +import com.databricks.jdbc.model.core.*; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; -import com.databricks.jdbc.telemetry.latency.TelemetryCollector; +import com.databricks.jdbc.telemetry.TelemetryHelper; import com.databricks.sdk.support.ToStringer; import com.google.common.annotations.VisibleForTesting; import java.io.InputStream; @@ -83,7 +84,7 @@ public DatabricksResultSet( StatementType statementType, IDatabricksSession session, IDatabricksStatementInternal parentStatement) - throws DatabricksSQLException { + throws SQLException { this.executionStatus = new ExecutionStatus(statementStatus); this.statementId = statementId; if (resultData != null) { @@ -152,10 +153,7 @@ public DatabricksResultSet( this.executionResult = ExecutionResultFactory.getResultSet(resultsResp, session, parentStatement); long rowSize = executionResult.getRowCount(); - List arrowMetadata = null; - if (executionResult instanceof ArrowStreamResult) { - arrowMetadata = ((ArrowStreamResult) executionResult).getArrowMetadata(); - } + List arrowMetadata = getArrowMetadata(resultsResp.getResultSetMetadata()); this.resultSetMetaData = new DatabricksResultSetMetaData( statementId, @@ -268,9 +266,8 @@ public DatabricksResultSet( public boolean next() throws SQLException { checkIfClosed(); boolean hasNext = this.executionResult.next(); - TelemetryCollector.getInstance() - .recordResultSetIteration( - statementId.toSQLExecStatementId(), resultSetMetaData.getChunkCount(), hasNext); + TelemetryHelper.recordResultSetIteration( + parentStatement, statementId, resultSetMetaData.getChunkCount(), hasNext); return hasNext; } @@ -478,22 +475,6 @@ public ResultSetMetaData getMetaData() throws SQLException { return resultSetMetaData; } - /** - * Checks if the given type name represents a complex type (ARRAY, MAP, STRUCT, GEOMETRY, or - * GEOGRAPHY). - * - * @param typeName The type name to check - * @return true if the type name starts with ARRAY, MAP, STRUCT, GEOMETRY, or GEOGRAPHY, false - * otherwise - */ - private static boolean isComplexType(String typeName) { - return typeName.startsWith(ARRAY) - || typeName.startsWith(MAP) - || typeName.startsWith(STRUCT) - || typeName.startsWith(GEOMETRY) - || typeName.startsWith(GEOGRAPHY); - } - /** * Checks if the given type name represents a geospatial type (GEOMETRY or GEOGRAPHY). * @@ -529,27 +510,28 @@ public Object getObject(int columnIndex) throws SQLException { } private Object handleComplexDataTypes(Object obj, String columnName) - throws DatabricksParsingException { - if (complexDatatypeSupport) return obj; + throws DatabricksSQLException { if (resultSetType == ResultSetType.SEA_INLINE) { - return handleComplexDataTypesForSEAInline(obj, columnName); + obj = convertToComplexDataTypesForSEAInline(obj, columnName); } - return obj.toString(); + return complexDatatypeSupport ? obj : obj.toString(); } - private Object handleComplexDataTypesForSEAInline(Object obj, String columnName) - throws DatabricksParsingException { + private Object convertToComplexDataTypesForSEAInline(Object obj, String columnName) + throws DatabricksSQLException { ComplexDataTypeParser parser = new ComplexDataTypeParser(); if (columnName.startsWith(ARRAY)) { - return parser.parseJsonStringToDbArray(obj.toString(), columnName).toString(); + return parser.parseJsonStringToDbArray(obj.toString(), columnName); } else if (columnName.startsWith(MAP)) { - return parser.parseJsonStringToDbMap(obj.toString(), columnName).toString(); + return parser.parseJsonStringToDbMap(obj.toString(), columnName); } else if (columnName.startsWith(STRUCT)) { - return parser.parseJsonStringToDbStruct(obj.toString(), columnName).toString(); + return parser.parseJsonStringToDbStruct(obj.toString(), columnName); } else if (columnName.startsWith(GEOMETRY)) { - return obj; + return ConverterHelper.getConverterForColumnType(Types.OTHER, GEOMETRY) + .toDatabricksGeometry(obj); } else if (columnName.startsWith(GEOGRAPHY)) { - return obj; + return ConverterHelper.getConverterForColumnType(Types.OTHER, GEOGRAPHY) + .toDatabricksGeography(obj); } throw new DatabricksParsingException( "Unexpected metadata format. Type is not a COMPLEX: " + columnName, @@ -613,12 +595,11 @@ public boolean isBeforeFirst() throws SQLException { /** * {@inheritDoc} * - *

Limitation: For lazy-loaded result sets ({@link LazyThriftResult}), particularly - * those using {@link - * com.databricks.jdbc.model.client.thrift.generated.TSparkRowSetType#COLUMN_BASED_SET}, this - * method cannot reliably determine the cursor position. The total row count remains unknown until - * all rows are fetched, preventing accurate detection of whether the cursor is after the last - * row. This is specific to Databricks JDBC dialect. + *

Limitation: For lazy/streaming result sets ({@link LazyThriftResult}, {@link + * StreamingColumnarResult}, {@link LazyThriftInlineArrowResult}, {@link + * StreamingInlineArrowResult}), this method cannot reliably determine the cursor position. The + * total row count remains unknown until all rows are fetched, preventing accurate detection of + * whether the cursor is after the last row. This is specific to Databricks JDBC dialect. */ @Override public boolean isAfterLast() throws SQLException { @@ -638,8 +619,10 @@ public boolean isFirst() throws SQLException { *

This method uses different strategies based on the result set type: * *

    - *
  • For {@link LazyThriftResult} instances: Checks if there are no more rows available (using - * {@code hasNext()}), since the total row count is unknown until all rows are fetched. + *
  • For lazy/streaming result types ({@link LazyThriftResult}, {@link + * StreamingColumnarResult}, {@link LazyThriftInlineArrowResult}, {@link + * StreamingInlineArrowResult}): Checks if there are no more rows available (using {@code + * hasNext()}), since the total row count is unknown until all rows are fetched. *
  • For other result types: Compares the current row position against the known total row * count. *
@@ -650,7 +633,10 @@ public boolean isFirst() throws SQLException { @Override public boolean isLast() throws SQLException { checkIfClosed(); - if (executionResult instanceof LazyThriftResult) { + if (executionResult instanceof LazyThriftResult + || executionResult instanceof StreamingColumnarResult + || executionResult instanceof LazyThriftInlineArrowResult + || executionResult instanceof StreamingInlineArrowResult) { return executionResult.getCurrentRow() >= 0 && !executionResult.hasNext(); } return executionResult.getCurrentRow() == resultSetMetaData.getTotalRows() - 1; diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaData.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaData.java index 65eb95519c..933d059fdf 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaData.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaData.java @@ -107,6 +107,12 @@ public DatabricksResultSetMetaData( typeText = "STRING"; } + // store base type eg. DECIMAL instead of DECIMAL(7,2) except for geospatial datatypes + String finalTypeText = + isGeospatialType(columnTypeName) + ? typeText + : metadataResultSetBuilder.stripTypeName(typeText); + int columnType = DatabricksTypeUtil.getColumnType(columnTypeName); int[] precisionAndScale = getPrecisionAndScale(columnInfo, columnType); int precision = precisionAndScale[0]; @@ -116,9 +122,7 @@ public DatabricksResultSetMetaData( .columnName(columnInfo.getName()) .columnTypeClassName(DatabricksTypeUtil.getColumnTypeClassName(columnTypeName)) .columnType(columnType) - .columnTypeText( - metadataResultSetBuilder.stripTypeName( - typeText)) // store base type eg. DECIMAL instead of DECIMAL(7,2) + .columnTypeText(finalTypeText) .typePrecision(precision) .typeScale(scale) .displaySize(DatabricksTypeUtil.getDisplaySize(columnTypeName, precision, scale)) @@ -187,7 +191,11 @@ public DatabricksResultSetMetaData( columnIndex++) { TColumnDesc columnDesc = resultManifest.getSchema().getColumns().get(columnIndex); - ColumnInfo columnInfo = getColumnInfoFromTColumnDesc(columnDesc); + String columnArrowMetadata = + arrowMetadata != null && columnIndex < arrowMetadata.size() + ? arrowMetadata.get(columnIndex) + : null; + ColumnInfo columnInfo = getColumnInfoFromTColumnDesc(columnDesc, columnArrowMetadata); int[] precisionAndScale = getPrecisionAndScale(columnInfo); int precision = precisionAndScale[0]; int scale = precisionAndScale[1]; @@ -201,6 +209,12 @@ public DatabricksResultSetMetaData( && arrowMetadata.get(columnIndex) != null) ? arrowMetadata.get(columnIndex) : getTypeTextFromTypeDesc(columnDesc.getTypeDesc()); + + // Normalize TIMESTAMP_NTZ to TIMESTAMP for consistency with SEA path + if (columnTypeText != null && columnTypeText.equalsIgnoreCase(TIMESTAMP_NTZ)) { + columnTypeText = TIMESTAMP; + } + columnBuilder .columnName(columnInfo.getName()) .columnTypeClassName( @@ -222,20 +236,6 @@ public DatabricksResultSetMetaData( .columnTypeClassName("java.lang.String") .columnType(Types.OTHER) .columnTypeText(VARIANT); - } else if (isGeometryColumn(arrowMetadata, columnIndex) - && ctx.isGeoSpatialSupportEnabled()) { - // Only set GEOMETRY type if geospatial support is enabled - columnBuilder - .columnTypeClassName(GEOMETRY_CLASS_NAME) - .columnType(Types.OTHER) - .columnTypeText(GEOMETRY); - } else if (isGeographyColumn(arrowMetadata, columnIndex) - && ctx.isGeoSpatialSupportEnabled()) { - // Only set GEOGRAPHY type if geospatial support is enabled - columnBuilder - .columnTypeClassName(GEOGRAPHY_CLASS_NAME) - .columnType(Types.OTHER) - .columnTypeText(GEOGRAPHY); } else if ((isGeometryColumn(arrowMetadata, columnIndex) || isGeographyColumn(arrowMetadata, columnIndex)) && !ctx.isGeoSpatialSupportEnabled()) { diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java index b6038f6fb4..86a8327f62 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java @@ -28,6 +28,7 @@ import com.databricks.jdbc.telemetry.latency.DatabricksMetricsTimedProcessor; import com.databricks.sdk.support.ToStringer; import com.google.common.annotations.VisibleForTesting; +import java.sql.SQLException; import java.util.HashMap; import java.util.Map; import javax.annotation.Nullable; @@ -131,7 +132,7 @@ public boolean isOpen() { } @Override - public void open() throws DatabricksSQLException { + public void open() throws SQLException { LOGGER.debug("public void open()"); // Skip for tests, it would be already set @@ -205,7 +206,7 @@ public void open() throws DatabricksSQLException { } @Override - public void close() throws DatabricksSQLException { + public void close() throws SQLException { LOGGER.debug("public void close()"); synchronized (this) { if (isSessionOpen) { @@ -361,7 +362,7 @@ public void setEmptyMetadataClient() { public void forceClose() { try { this.close(); - } catch (DatabricksSQLException e) { + } catch (SQLException e) { LOGGER.error("Error closing session resources, but marking the session as closed."); } finally { this.isSessionOpen = false; 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 702201088d..96a833539b 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java @@ -22,6 +22,7 @@ import java.sql.*; import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.concurrent.*; import org.apache.http.entity.InputStreamEntity; @@ -75,7 +76,7 @@ public ResultSet executeQuery(String sql) throws SQLException { // TODO (PECO-1731): Can it fail fast without executing SQL query? checkIfClosed(); ResultSet rs = executeInternal(sql, new HashMap<>(), StatementType.QUERY); - if (!shouldReturnResultSet(sql)) { + if (!shouldReturnResultSetWithConfig(sql)) { String errorMessage = "A ResultSet was expected but not generated from query. However, query " + "execution was successful."; @@ -236,7 +237,7 @@ public void setCursorName(String name) throws SQLException { public boolean execute(String sql) throws SQLException { checkIfClosed(); resultSet = executeInternal(sql, new HashMap<>(), StatementType.SQL); - return shouldReturnResultSet(sql); + return shouldReturnResultSetWithConfig(sql); } @Override @@ -680,9 +681,15 @@ static String trimCommentsAndWhitespaces(String query) { } @VisibleForTesting - static boolean shouldReturnResultSet(String query) { + static boolean shouldReturnResultSet(String query, List nonRowcountQueryPrefixes) { String trimmedQuery = trimCommentsAndWhitespaces(query); + // Check configured non-rowcount prefixes first + String upperQuery = trimmedQuery.toUpperCase(); + if (nonRowcountQueryPrefixes.stream().anyMatch(upperQuery::startsWith)) { + return true; + } + // Check if the query matches any of the patterns that return a ResultSet return SELECT_PATTERN.matcher(trimmedQuery).find() || SHOW_PATTERN.matcher(trimmedQuery).find() @@ -707,6 +714,11 @@ static boolean shouldReturnResultSet(String query) { // Otherwise, it should not return a ResultSet } + protected boolean shouldReturnResultSetWithConfig(String query) { + return shouldReturnResultSet( + query, connection.getConnectionContext().getNonRowcountQueryPrefixes()); + } + static boolean isSelectQuery(String query) { String trimmedQuery = trimCommentsAndWhitespaces(query); return SELECT_PATTERN.matcher(trimmedQuery).find(); diff --git a/src/main/java/com/databricks/jdbc/api/impl/ExecutionResultFactory.java b/src/main/java/com/databricks/jdbc/api/impl/ExecutionResultFactory.java index c93a3f2879..b9aec3fc1b 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/ExecutionResultFactory.java +++ b/src/main/java/com/databricks/jdbc/api/impl/ExecutionResultFactory.java @@ -1,7 +1,11 @@ package com.databricks.jdbc.api.impl; import com.databricks.jdbc.api.impl.arrow.ArrowStreamResult; +import com.databricks.jdbc.api.impl.arrow.LazyThriftInlineArrowResult; +import com.databricks.jdbc.api.impl.arrow.StreamingInlineArrowResult; +import com.databricks.jdbc.api.impl.thrift.StreamingColumnarResult; import com.databricks.jdbc.api.impl.volume.VolumeOperationResult; +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.util.DatabricksThriftUtil; @@ -14,7 +18,7 @@ import com.databricks.jdbc.model.core.ResultData; import com.databricks.jdbc.model.core.ResultManifest; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; -import com.databricks.jdbc.telemetry.latency.TelemetryCollector; +import com.databricks.jdbc.telemetry.TelemetryHelper; import java.sql.SQLException; import java.util.List; @@ -28,7 +32,7 @@ static IExecutionResult getResultSet( StatementId statementId, IDatabricksSession session, IDatabricksStatementInternal statement) - throws DatabricksSQLException { + throws SQLException { IExecutionResult resultHandler = getResultHandler(data, manifest, statementId, session); if (manifest.getIsVolumeOperation() != null && manifest.getIsVolumeOperation()) { return new VolumeOperationResult( @@ -49,7 +53,8 @@ private static IExecutionResult getResultHandler( throw new DatabricksParsingException( "Empty response format", DatabricksDriverErrorCode.INVALID_STATE); } - TelemetryCollector.getInstance().setResultFormat(statementId, manifest.getFormat()); + TelemetryHelper.setResultFormat( + session.getConnectionContext(), statementId, manifest.getFormat()); LOGGER.info("Processing result of format {} from SQL Execution API", manifest.getFormat()); // We use JSON_ARRAY for metadata and update commands, and ARROW_STREAM for query results switch (manifest.getFormat()) { @@ -90,15 +95,15 @@ private static IExecutionResult getResultHandler( IDatabricksSession session) throws SQLException { TSparkRowSetType resultFormat = resultsResp.getResultSetMetadata().getResultFormat(); - TelemetryCollector.getInstance().setResultFormat(parentStatement, resultFormat); + TelemetryHelper.setResultFormat(session.getConnectionContext(), parentStatement, resultFormat); LOGGER.info("Processing result of format {} from Thrift server", resultFormat); switch (resultFormat) { case COLUMN_BASED_SET: - return new LazyThriftResult(resultsResp, parentStatement, session); + return createThriftColumnarResult(resultsResp, parentStatement, session); case ARROW_BASED_SET: - return new ArrowStreamResult(resultsResp, true, parentStatement, session); + return createInlineArrowResult(resultsResp, parentStatement, session); case URL_BASED_SET: - return new ArrowStreamResult(resultsResp, false, parentStatement, session); + return new ArrowStreamResult(resultsResp, parentStatement, session); case ROW_BASED_SET: throw new DatabricksSQLFeatureNotSupportedException( "Invalid state - row based set cannot be received"); @@ -108,6 +113,50 @@ private static IExecutionResult getResultHandler( } } + /** + * Creates the appropriate result handler for Thrift columnar results. Uses + * StreamingColumnarResult by default for improved throughput, otherwise falls back to + * LazyThriftResult if disabled. + */ + private static IExecutionResult createThriftColumnarResult( + TFetchResultsResp resultsResp, + IDatabricksStatementInternal parentStatement, + IDatabricksSession session) + throws DatabricksSQLException { + IDatabricksConnectionContext connectionContext = session.getConnectionContext(); + + // Streaming is enabled by default (ENABLE_INLINE_STREAMING defaults to "1") + if (connectionContext.isInlineStreamingEnabled()) { + LOGGER.info("Using StreamingColumnarResult for improved throughput (default)"); + return new StreamingColumnarResult(resultsResp, parentStatement, session); + } else { + LOGGER.info("Using LazyThriftResult (streaming explicitly disabled)"); + return new LazyThriftResult(resultsResp, parentStatement, session); + } + } + + /** + * Creates the appropriate result handler for inline Arrow results. Uses + * StreamingInlineArrowResult by default for improved throughput, otherwise falls back to + * LazyThriftInlineArrowResult. + */ + private static IExecutionResult createInlineArrowResult( + TFetchResultsResp resultsResp, + IDatabricksStatementInternal parentStatement, + IDatabricksSession session) + throws DatabricksSQLException { + IDatabricksConnectionContext connectionContext = session.getConnectionContext(); + + // Streaming is enabled by default (ENABLE_INLINE_STREAMING defaults to "1") + if (connectionContext.isInlineStreamingEnabled()) { + LOGGER.info("Using StreamingInlineArrowResult for improved throughput (default)"); + return new StreamingInlineArrowResult(resultsResp, parentStatement, session); + } else { + LOGGER.info("Using LazyThriftInlineArrowResult (streaming explicitly disabled)"); + return new LazyThriftInlineArrowResult(resultsResp, parentStatement, session); + } + } + static IExecutionResult getResultSet(Object[][] rows) { return new InlineJsonResult(rows); } diff --git a/src/main/java/com/databricks/jdbc/api/impl/IExecutionResult.java b/src/main/java/com/databricks/jdbc/api/impl/IExecutionResult.java index d17e43174a..efa8d07926 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/IExecutionResult.java +++ b/src/main/java/com/databricks/jdbc/api/impl/IExecutionResult.java @@ -1,6 +1,7 @@ package com.databricks.jdbc.api.impl; import com.databricks.jdbc.exception.DatabricksSQLException; +import java.sql.SQLException; /** Interface to provide methods over an underlying statement result */ public interface IExecutionResult { @@ -26,7 +27,7 @@ public interface IExecutionResult { * * @return true if cursor is moved at next row */ - boolean next() throws DatabricksSQLException; + boolean next() throws SQLException; /** Returns if there is next row in the result set */ boolean hasNext(); diff --git a/src/main/java/com/databricks/jdbc/api/impl/LazyThriftResult.java b/src/main/java/com/databricks/jdbc/api/impl/LazyThriftResult.java index 568b232aca..854da86566 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/LazyThriftResult.java +++ b/src/main/java/com/databricks/jdbc/api/impl/LazyThriftResult.java @@ -10,6 +10,7 @@ import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.sql.SQLException; public class LazyThriftResult implements IExecutionResult { private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(LazyThriftResult.class); @@ -99,10 +100,10 @@ public long getCurrentRow() { * Moves the cursor to the next row. Fetches additional data from server if needed. * * @return true if there is a next row, false if at the end - * @throws DatabricksSQLException if an error occurs while fetching data + * @throws SQLException if an error occurs while fetching data */ @Override - public boolean next() throws DatabricksSQLException { + public boolean next() throws SQLException { if (isClosed || hasReachedEnd) { return false; } @@ -221,9 +222,9 @@ private void loadCurrentBatch() throws DatabricksSQLException { /** * Fetches the next batch of data from the server and creates columnar views. * - * @throws DatabricksSQLException if the fetch operation fails + * @throws SQLException if the fetch operation fails */ - private void fetchNextBatch() throws DatabricksSQLException { + private void fetchNextBatch() throws SQLException { try { LOGGER.debug("Fetching next batch, current total rows fetched: {}", totalRowsFetched); currentResponse = session.getDatabricksClient().getMoreResults(statement); diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java index f06787673b..1a96642994 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java @@ -337,13 +337,23 @@ private ArrowData getRecordBatchList( long rowCount = 0L; try (ArrowStreamReader arrowStreamReader = new ArrowStreamReader(inputStream, rootAllocator)) { VectorSchemaRoot vectorSchemaRoot = arrowStreamReader.getVectorSchemaRoot(); - boolean fetchedMetadata = false; + + // Extract metadata from VectorSchemaRoot before loading any batches. + // The Arrow IPC format sends the schema first (before any record batches), + // so field metadata is available even when there are 0 rows. + // VectorSchemaRoot will contain field vectors with metadata, but rowCount will be 0. + if (vectorSchemaRoot != null && vectorSchemaRoot.getFieldVectors() != null) { + metadata = getMetadataInformationFromSchemaRoot(vectorSchemaRoot); + LOGGER.debug( + "Extracted metadata from VectorSchemaRoot before loading batches. " + + "Schema has {} fields. Statement: {}, Chunk: {}", + vectorSchemaRoot.getFieldVectors().size(), + statementId, + chunkIndex); + } + while (arrowStreamReader.loadNextBatch()) { rowCount += vectorSchemaRoot.getRowCount(); - if (!fetchedMetadata) { - metadata = getMetadataInformationFromSchemaRoot(vectorSchemaRoot); - fetchedMetadata = true; - } recordBatchList.add(getVectorsFromSchemaRoot(vectorSchemaRoot, rootAllocator)); vectorSchemaRoot.clear(); } diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractRemoteChunkProvider.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractRemoteChunkProvider.java index 5fc9dc1500..348a6f70ee 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractRemoteChunkProvider.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractRemoteChunkProvider.java @@ -17,8 +17,9 @@ import com.databricks.jdbc.model.core.ResultData; import com.databricks.jdbc.model.core.ResultManifest; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; -import com.databricks.jdbc.telemetry.latency.TelemetryCollector; +import com.databricks.jdbc.telemetry.TelemetryHelper; import com.databricks.sdk.service.sql.BaseChunkInfo; +import java.sql.SQLException; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ExecutionException; @@ -83,7 +84,7 @@ protected AbstractRemoteChunkProvider( chunkCount, chunkIndexToChunksMap, resultData.getExternalLinks() != null ? resultData.getExternalLinks().size() : 1); - TelemetryCollector.getInstance().recordTotalChunks(statementId, chunkCount); + TelemetryHelper.recordTotalChunks(session.getConnectionContext(), statementId, chunkCount); initializeData(); } @@ -94,7 +95,7 @@ protected AbstractRemoteChunkProvider( IDatabricksHttpClient httpClient, int maxParallelChunkDownloadsPerQuery, CompressionCodec compressionCodec) - throws DatabricksSQLException { + throws SQLException { this.chunkReadyTimeoutSeconds = session.getConnectionContext().getChunkReadyTimeoutSeconds(); this.maxParallelChunkDownloadsPerQuery = maxParallelChunkDownloadsPerQuery; this.session = session; @@ -264,14 +265,14 @@ private ConcurrentMap initializeChunksMap( TFetchResultsResp resultsResp, IDatabricksStatementInternal parentStatement, IDatabricksSession session) - throws DatabricksSQLException { + throws SQLException { ConcurrentMap chunkIndexMap = new ConcurrentHashMap<>(); populateChunkIndexMap(resultsResp.getResults(), chunkIndexMap); while (resultsResp.hasMoreRows) { resultsResp = session.getDatabricksClient().getMoreResults(parentStatement); populateChunkIndexMap(resultsResp.getResults(), chunkIndexMap); } - TelemetryCollector.getInstance().recordTotalChunks(statementId, chunkCount); + TelemetryHelper.recordTotalChunks(session.getConnectionContext(), statementId, chunkCount); return chunkIndexMap; } diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunk.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunk.java index 9aba1304dd..d7dfc9efe3 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunk.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunk.java @@ -4,6 +4,7 @@ import static com.databricks.jdbc.common.util.ValidationUtil.checkHTTPError; import static com.databricks.jdbc.telemetry.TelemetryHelper.getStatementIdString; +import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.common.CompressionCodec; import com.databricks.jdbc.common.DatabricksJdbcUrlParams; import com.databricks.jdbc.common.util.DecompressionUtil; @@ -15,7 +16,7 @@ import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.client.thrift.generated.TSparkArrowResultLink; import com.databricks.jdbc.model.core.ExternalLink; -import com.databricks.jdbc.telemetry.latency.TelemetryCollector; +import com.databricks.jdbc.telemetry.TelemetryHelper; import com.databricks.sdk.service.sql.BaseChunkInfo; import java.io.IOException; import java.io.InputStream; @@ -28,6 +29,7 @@ public class ArrowResultChunk extends AbstractArrowResultChunk { private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(ArrowResultChunk.class); + private final IDatabricksConnectionContext connectionContext; private ArrowResultChunk(Builder builder) throws DatabricksParsingException { super( @@ -39,6 +41,7 @@ private ArrowResultChunk(Builder builder) throws DatabricksParsingException { builder.chunkLink, builder.expiryTime, builder.chunkReadyTimeoutSeconds); + this.connectionContext = builder.connectionContext; if (builder.inputStream != null) { // Data is already available try { @@ -80,10 +83,9 @@ protected void downloadData( checkHTTPError(response); long downloadTimeMs = (System.nanoTime() - startTime) / 1_000_000; - // Add telemetry for the time to first byte for chunk download - TelemetryCollector.getInstance() - .recordChunkDownloadLatency( - getStatementIdString(statementId), chunkIndex, downloadTimeMs); + // Record chunk download latency telemetry + TelemetryHelper.recordChunkDownloadLatency( + connectionContext, getStatementIdString(statementId), chunkIndex, downloadTimeMs); // Read compressed stream fully (download latency excludes decompression) byte[] compressed = IOUtils.toByteArray(response.getEntity().getContent()); @@ -183,12 +185,18 @@ public static class Builder { private InputStream inputStream; private int chunkReadyTimeoutSeconds = Integer.parseInt(DatabricksJdbcUrlParams.CHUNK_READY_TIMEOUT_SECONDS.getDefaultValue()); + private IDatabricksConnectionContext connectionContext; public Builder withStatementId(StatementId statementId) { this.statementId = statementId; return this; } + public Builder withConnectionContext(IDatabricksConnectionContext connectionContext) { + this.connectionContext = connectionContext; + return this; + } + public Builder withChunkInfo(BaseChunkInfo baseChunkInfo) { this.chunkIndex = baseChunkInfo.getChunkIndex(); this.numRows = baseChunkInfo.getRowCount(); diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java index 1d7067a3e8..a10de8e1e3 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java @@ -1,7 +1,7 @@ package com.databricks.jdbc.api.impl.arrow; +import static com.databricks.jdbc.common.util.ArrowUtil.getColumnInfoList; import static com.databricks.jdbc.common.util.DatabricksThriftUtil.createExternalLink; -import static com.databricks.jdbc.common.util.DatabricksThriftUtil.getColumnInfoFromTColumnDesc; import com.databricks.jdbc.api.impl.ComplexDataTypeParser; import com.databricks.jdbc.api.impl.IExecutionResult; @@ -15,9 +15,7 @@ import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; -import com.databricks.jdbc.model.client.thrift.generated.TColumnDesc; import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; -import com.databricks.jdbc.model.client.thrift.generated.TGetResultSetMetadataResp; import com.databricks.jdbc.model.client.thrift.generated.TSparkArrowResultLink; import com.databricks.jdbc.model.core.ChunkLinkFetchResult; import com.databricks.jdbc.model.core.ColumnInfo; @@ -26,6 +24,7 @@ import com.databricks.jdbc.model.core.ResultData; import com.databricks.jdbc.model.core.ResultManifest; import com.google.common.annotations.VisibleForTesting; +import java.sql.SQLException; import java.util.ArrayList; import java.util.Collection; import java.util.List; @@ -128,6 +127,7 @@ private static ChunkProvider createRemoteChunkProvider( linkPrefetchWindow, chunkReadyTimeoutSeconds, cloudFetchSpeedThreshold, + connectionContext, initialLinks); } else { // Use the original RemoteChunkProvider @@ -143,13 +143,11 @@ private static ChunkProvider createRemoteChunkProvider( public ArrowStreamResult( TFetchResultsResp resultsResp, - boolean isInlineArrow, IDatabricksStatementInternal parentStatementId, IDatabricksSession session) - throws DatabricksSQLException { + throws SQLException { this( resultsResp, - isInlineArrow, parentStatementId, session, DatabricksHttpClientFactory.getInstance().getClient(session.getConnectionContext())); @@ -158,19 +156,14 @@ public ArrowStreamResult( @VisibleForTesting ArrowStreamResult( TFetchResultsResp resultsResp, - boolean isInlineArrow, IDatabricksStatementInternal parentStatement, IDatabricksSession session, IDatabricksHttpClient httpClient) - throws DatabricksSQLException { + throws SQLException { this.session = session; - setColumnInfo(resultsResp.getResultSetMetadata()); - if (isInlineArrow) { - this.chunkProvider = new InlineChunkProvider(resultsResp, parentStatement, session); - } else { - this.chunkProvider = - createThriftRemoteChunkProvider(resultsResp, parentStatement, session, httpClient); - } + this.columnInfos = getColumnInfoList(resultsResp.getResultSetMetadata()); + this.chunkProvider = + createThriftRemoteChunkProvider(resultsResp, parentStatement, session, httpClient); } /** @@ -187,7 +180,7 @@ private static ChunkProvider createThriftRemoteChunkProvider( IDatabricksStatementInternal parentStatement, IDatabricksSession session, IDatabricksHttpClient httpClient) - throws DatabricksSQLException { + throws SQLException { IDatabricksConnectionContext connectionContext = session.getConnectionContext(); CompressionCodec compressionCodec = @@ -215,6 +208,7 @@ private static ChunkProvider createThriftRemoteChunkProvider( linkPrefetchWindow, chunkReadyTimeoutSeconds, cloudFetchSpeedThreshold, + connectionContext, initialLinks); } else { // Use the original RemoteChunkProvider @@ -238,48 +232,15 @@ public List getArrowMetadata() throws DatabricksSQLException { /** {@inheritDoc} */ @Override public Object getObject(int columnIndex) throws DatabricksSQLException { - ColumnInfoTypeName requiredType = columnInfos.get(columnIndex).getTypeName(); + ColumnInfo columnInfo = columnInfos.get(columnIndex); + ColumnInfoTypeName requiredType = columnInfo.getTypeName(); String arrowMetadata = chunkIterator.getType(columnIndex); if (arrowMetadata == null) { - arrowMetadata = columnInfos.get(columnIndex).getTypeText(); - } - - // Handle complex type conversion when complex datatype support is disabled - boolean isComplexDatatypeSupportEnabled = - this.session.getConnectionContext().isComplexDatatypeSupportEnabled(); - boolean isGeoSpatialSupportEnabled = - this.session.getConnectionContext().isGeoSpatialSupportEnabled(); - - // Check if we need to convert geospatial types to string when geospatial support is disabled - // This check must come before the general complex type check - if (!isGeoSpatialSupportEnabled && isGeospatialType(requiredType)) { - LOGGER.debug("Geospatial support is disabled, converting {} to STRING", requiredType); - - Object result = - chunkIterator.getColumnObjectAtCurrentRow( - columnIndex, ColumnInfoTypeName.STRING, "STRING", columnInfos.get(columnIndex)); - if (result == null) { - return null; - } - // Return raw string for geospatial types when support is disabled - return result; - } - - if (!isComplexDatatypeSupportEnabled && isComplexType(requiredType)) { - LOGGER.debug("Complex datatype support is disabled, converting complex type to STRING"); - - Object result = - chunkIterator.getColumnObjectAtCurrentRow( - columnIndex, ColumnInfoTypeName.STRING, "STRING", columnInfos.get(columnIndex)); - if (result == null) { - return null; - } - ComplexDataTypeParser parser = new ComplexDataTypeParser(); - return parser.formatComplexTypeString(result.toString(), requiredType.name(), arrowMetadata); + arrowMetadata = columnInfo.getTypeText(); } - return chunkIterator.getColumnObjectAtCurrentRow( - columnIndex, requiredType, arrowMetadata, columnInfos.get(columnIndex)); + return getObjectWithComplexTypeHandling( + session, chunkIterator, columnIndex, requiredType, arrowMetadata, columnInfo); } /** @@ -374,14 +335,63 @@ public ChunkProvider getChunkProvider() { return chunkProvider; } - private void setColumnInfo(TGetResultSetMetadataResp resultManifest) { - columnInfos = new ArrayList<>(); - if (resultManifest.getSchema() == null) { - return; + /** + * Helper method to handle complex type and geospatial type conversion when support is disabled. + * + *

This method is also used by LazyThriftInlineArrowResult for consistent type handling. + * + * @param session The databricks session + * @param chunkIterator The chunk iterator + * @param columnIndex The column index + * @param requiredType The required column type + * @param arrowMetadata The arrow metadata + * @param columnInfo The column info + * @return The object value (converted if complex/geospatial type and support disabled) + * @throws DatabricksSQLException if an error occurs + */ + protected static Object getObjectWithComplexTypeHandling( + IDatabricksSession session, + ArrowResultChunkIterator chunkIterator, + int columnIndex, + ColumnInfoTypeName requiredType, + String arrowMetadata, + ColumnInfo columnInfo) + throws DatabricksSQLException { + boolean isComplexDatatypeSupportEnabled = + session.getConnectionContext().isComplexDatatypeSupportEnabled(); + boolean isGeoSpatialSupportEnabled = + session.getConnectionContext().isGeoSpatialSupportEnabled(); + + // Check if we need to convert geospatial types to string when geospatial support is disabled + // This check must come before the general complex type check + if (!isGeoSpatialSupportEnabled && isGeospatialType(requiredType)) { + LOGGER.debug("Geospatial support is disabled, converting {} to STRING", requiredType); + + Object result = + chunkIterator.getColumnObjectAtCurrentRow( + columnIndex, ColumnInfoTypeName.STRING, "STRING", columnInfo); + if (result == null) { + return null; + } + // Return raw string for geospatial types when support is disabled + return result; } - for (TColumnDesc tColumnDesc : resultManifest.getSchema().getColumns()) { - columnInfos.add(getColumnInfoFromTColumnDesc(tColumnDesc)); + + if (!isComplexDatatypeSupportEnabled && isComplexType(requiredType)) { + LOGGER.debug("Complex datatype support is disabled, converting complex type to STRING"); + Object result = + chunkIterator.getColumnObjectAtCurrentRow( + columnIndex, ColumnInfoTypeName.STRING, "STRING", columnInfo); + if (result == null) { + return null; + } + ComplexDataTypeParser parser = new ComplexDataTypeParser(); + + return parser.formatComplexTypeString(result.toString(), requiredType.name(), arrowMetadata); } + + return chunkIterator.getColumnObjectAtCurrentRow( + columnIndex, requiredType, arrowMetadata, columnInfo); } /** diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadService.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadService.java index 5d0f2f563b..96ac936c5c 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadService.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadService.java @@ -6,13 +6,13 @@ import com.databricks.jdbc.common.DatabricksClientType; import com.databricks.jdbc.common.util.DriverUtil; import com.databricks.jdbc.dbclient.impl.common.StatementId; -import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.exception.DatabricksValidationException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.core.ChunkLinkFetchResult; import com.databricks.jdbc.model.core.ExternalLink; import com.google.common.annotations.VisibleForTesting; +import java.sql.SQLException; import java.time.Instant; import java.util.Map; import java.util.concurrent.CompletableFuture; @@ -282,7 +282,7 @@ private void triggerNextBatchDownload() { triggerNextBatchDownload(); } } - } catch (DatabricksSQLException e) { + } catch (SQLException e) { // If the download fails, complete exceptionally all pending futures handleBatchDownloadError(batchStartIndex, e); } @@ -295,7 +295,7 @@ private void triggerNextBatchDownload() { *

Completes all pending futures exceptionally with the encountered error and resets the * download progress flag. */ - private void handleBatchDownloadError(long batchStartIndex, DatabricksSQLException e) { + private void handleBatchDownloadError(long batchStartIndex, SQLException e) { LOGGER.error( e, "Failed to download links for batch starting at {} : {}", diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkFetcher.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkFetcher.java index 6d6a88e566..d148bd02cd 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkFetcher.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkFetcher.java @@ -1,8 +1,8 @@ package com.databricks.jdbc.api.impl.arrow; -import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.model.core.ChunkLinkFetchResult; import com.databricks.jdbc.model.core.ExternalLink; +import java.sql.SQLException; /** * Abstraction for fetching chunk links from either SEA or Thrift backend. Implementations handle @@ -26,10 +26,9 @@ public interface ChunkLinkFetcher { * @param startRowOffset The row offset to start fetching from (used by Thrift with * FETCH_ABSOLUTE) * @return ChunkLinkFetchResult containing the fetched links and continuation information - * @throws DatabricksSQLException if the fetch operation fails + * @throws SQLException if the fetch operation fails */ - ChunkLinkFetchResult fetchLinks(long startChunkIndex, long startRowOffset) - throws DatabricksSQLException; + ChunkLinkFetchResult fetchLinks(long startChunkIndex, long startRowOffset) throws SQLException; /** * Refetches a specific chunk link that may have expired. @@ -42,9 +41,9 @@ ChunkLinkFetchResult fetchLinks(long startChunkIndex, long startRowOffset) * @param chunkIndex The specific chunk index to refetch (used by SEA) * @param rowOffset The row offset of the chunk to refetch (used by Thrift) * @return The refreshed ExternalLink with a new expiration time - * @throws DatabricksSQLException if the refetch operation fails + * @throws SQLException if the refetch operation fails */ - ExternalLink refetchLink(long chunkIndex, long rowOffset) throws DatabricksSQLException; + ExternalLink refetchLink(long chunkIndex, long rowOffset) throws SQLException; /** Closes any resources held by the fetcher. */ void close(); diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/InlineChunkProvider.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/InlineChunkProvider.java index e22d974a40..32f5e1b803 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/InlineChunkProvider.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/InlineChunkProvider.java @@ -1,31 +1,17 @@ package com.databricks.jdbc.api.impl.arrow; -import static com.databricks.jdbc.common.util.DatabricksTypeUtil.*; import static com.databricks.jdbc.common.util.DecompressionUtil.decompress; -import com.databricks.jdbc.api.internal.IDatabricksSession; -import com.databricks.jdbc.api.internal.IDatabricksStatementInternal; import com.databricks.jdbc.common.CompressionCodec; import com.databricks.jdbc.exception.DatabricksParsingException; import com.databricks.jdbc.exception.DatabricksSQLException; 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.ResultData; import com.databricks.jdbc.model.core.ResultManifest; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; import com.google.common.annotations.VisibleForTesting; import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; -import java.io.IOException; -import java.sql.SQLException; -import java.util.ArrayList; -import java.util.List; -import org.apache.arrow.vector.types.pojo.ArrowType; -import org.apache.arrow.vector.types.pojo.Field; -import org.apache.arrow.vector.types.pojo.FieldType; -import org.apache.arrow.vector.types.pojo.Schema; -import org.apache.arrow.vector.util.SchemaUtility; /** Class to manage inline Arrow chunks */ public class InlineChunkProvider implements ChunkProvider { @@ -37,23 +23,6 @@ public class InlineChunkProvider implements ChunkProvider { private final ArrowResultChunk arrowResultChunk; // There is only one packet of data in case of inline arrow - InlineChunkProvider( - TFetchResultsResp resultsResp, - IDatabricksStatementInternal parentStatement, - IDatabricksSession session) - throws DatabricksParsingException { - this.currentChunkIndex = -1; - this.totalRows = 0; - ByteArrayInputStream byteStream = initializeByteStream(resultsResp, session, parentStatement); - ArrowResultChunk.Builder builder = - ArrowResultChunk.builder().withInputStream(byteStream, totalRows); - - if (parentStatement != null) { - builder.withStatementId(parentStatement.getStatementId()); - } - arrowResultChunk = builder.build(); - } - /** * Constructor for inline arrow chunk provider from {@link ResultData} and {@link ResultManifest}. * @@ -123,97 +92,6 @@ public boolean isClosed() { return isClosed; } - private ByteArrayInputStream initializeByteStream( - TFetchResultsResp resultsResp, - IDatabricksSession session, - IDatabricksStatementInternal parentStatement) - throws DatabricksParsingException { - ByteArrayOutputStream baos = new ByteArrayOutputStream(); - CompressionCodec compressionType = - CompressionCodec.getCompressionMapping(resultsResp.getResultSetMetadata()); - try { - byte[] serializedSchema = getSerializedSchema(resultsResp.getResultSetMetadata()); - if (serializedSchema != null) { - baos.write(serializedSchema); - } - writeToByteOutputStream( - compressionType, parentStatement, resultsResp.getResults().getArrowBatches(), baos); - while (resultsResp.hasMoreRows) { - resultsResp = session.getDatabricksClient().getMoreResults(parentStatement); - writeToByteOutputStream( - compressionType, parentStatement, resultsResp.getResults().getArrowBatches(), baos); - } - return new ByteArrayInputStream(baos.toByteArray()); - } catch (DatabricksSQLException | IOException e) { - handleError(e); - } - return null; - } - - private void writeToByteOutputStream( - CompressionCodec compressionCodec, - IDatabricksStatementInternal parentStatement, - List arrowBatchList, - ByteArrayOutputStream baos) - throws DatabricksSQLException, IOException { - for (TSparkArrowBatch arrowBatch : arrowBatchList) { - byte[] decompressedBytes = - decompress( - arrowBatch.getBatch(), - compressionCodec, - String.format( - "Data fetch for inline arrow batch [%d] and statement [%s] with decompression algorithm : [%s]", - arrowBatch.getRowCount(), parentStatement, compressionCodec)); - totalRows += arrowBatch.getRowCount(); - baos.write(decompressedBytes); - } - } - - private byte[] getSerializedSchema(TGetResultSetMetadataResp metadata) - throws DatabricksSQLException { - if (metadata.getArrowSchema() != null) { - return metadata.getArrowSchema(); - } - Schema arrowSchema = hiveSchemaToArrowSchema(metadata.getSchema()); - try { - return SchemaUtility.serialize(arrowSchema); - } catch (IOException e) { - handleError(e); - } - // should never reach here; - return null; - } - - private Schema hiveSchemaToArrowSchema(TTableSchema hiveSchema) - throws DatabricksParsingException { - List fields = new ArrayList<>(); - if (hiveSchema == null) { - return new Schema(fields); - } - try { - hiveSchema - .getColumns() - .forEach( - columnDesc -> { - try { - fields.add(getArrowField(columnDesc)); - } catch (SQLException e) { - throw new RuntimeException(e); - } - }); - } catch (RuntimeException e) { - handleError(e); - } - return new Schema(fields); - } - - private Field getArrowField(TColumnDesc columnDesc) throws SQLException { - TPrimitiveTypeEntry primitiveTypeEntry = getTPrimitiveTypeOrDefault(columnDesc.getTypeDesc()); - ArrowType arrowType = mapThriftToArrowType(primitiveTypeEntry.getType()); - FieldType fieldType = new FieldType(true, arrowType, null); - return new Field(columnDesc.getColumnName(), fieldType, null); - } - @VisibleForTesting void handleError(Exception e) throws DatabricksParsingException { String errorMessage = diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResult.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResult.java new file mode 100644 index 0000000000..98d97a8992 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResult.java @@ -0,0 +1,352 @@ +package com.databricks.jdbc.api.impl.arrow; + +import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_RESULT_ROW_LIMIT; +import static com.databricks.jdbc.common.util.ArrowUtil.createArrowByteStream; +import static com.databricks.jdbc.common.util.ArrowUtil.getColumnInfoList; +import static com.databricks.jdbc.common.util.ArrowUtil.getSerializedSchema; +import static com.databricks.jdbc.common.util.ArrowUtil.getTotalRowsInResponse; + +import com.databricks.jdbc.api.impl.IExecutionResult; +import com.databricks.jdbc.api.internal.IDatabricksSession; +import com.databricks.jdbc.api.internal.IDatabricksStatementInternal; +import com.databricks.jdbc.exception.DatabricksParsingException; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.log.JdbcLogger; +import com.databricks.jdbc.log.JdbcLoggerFactory; +import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; +import com.databricks.jdbc.model.core.ColumnInfo; +import com.databricks.jdbc.model.core.ColumnInfoTypeName; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.io.ByteArrayInputStream; +import java.sql.SQLException; +import java.util.List; + +/** + * Lazy implementation for thrift-based inline Arrow results that fetches arrow batches on demand. + * Similar to LazyThriftResult but processes Arrow data instead of columnar thrift data. + */ +public class LazyThriftInlineArrowResult implements IExecutionResult { + private static final JdbcLogger LOGGER = + JdbcLoggerFactory.getLogger(LazyThriftInlineArrowResult.class); + + private TFetchResultsResp currentResponse; + private ArrowResultChunk currentChunk; + private ArrowResultChunkIterator currentChunkIterator; + private long globalRowIndex; + private final IDatabricksSession session; + private final IDatabricksStatementInternal statement; + private final int maxRows; + private boolean hasReachedEnd; + private boolean isClosed; + private long totalRowsFetched; + private List columnInfos; + private byte[] cachedSchema; // Cache schema from first response for subsequent batches + + /** + * Creates a new LazyThriftInlineArrowResult that lazily fetches arrow data on demand. + * + * @param initialResponse the initial response from the server + * @param statement the statement that generated this result + * @param session the session to use for fetching additional data + * @throws DatabricksSQLException if the initial response cannot be processed + */ + public LazyThriftInlineArrowResult( + TFetchResultsResp initialResponse, + IDatabricksStatementInternal statement, + IDatabricksSession session) + throws DatabricksSQLException { + this.currentResponse = initialResponse; + this.statement = statement; + this.session = session; + this.maxRows = statement != null ? statement.getMaxRows() : DEFAULT_RESULT_ROW_LIMIT; + this.globalRowIndex = -1; + this.hasReachedEnd = false; + this.isClosed = false; + this.totalRowsFetched = 0; + + // Initialize column info from metadata + this.columnInfos = getColumnInfoList(initialResponse.getResultSetMetadata()); + + // Cache the schema from the first response for use in subsequent batches + try { + this.cachedSchema = getSerializedSchema(initialResponse.getResultSetMetadata()); + } catch (DatabricksParsingException e) { + LOGGER.error("Failed to cache Arrow schema: {}", e.getMessage(), e); + throw new DatabricksSQLException( + "Failed to cache Arrow schema", e, DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR); + } + + // Load initial chunk + loadCurrentChunk(); + LOGGER.debug( + "LazyThriftInlineArrowResult initialized with {} rows in first chunk, hasMoreRows: {}", + currentChunk.numRows, + currentResponse.hasMoreRows); + } + + /** + * Gets the value at the specified column index for the current row. + * + * @param columnIndex the zero-based column index + * @return the value at the specified column + * @throws DatabricksSQLException if the result is closed, cursor is invalid, or column index is + * out of bounds + */ + @Override + public Object getObject(int columnIndex) throws DatabricksSQLException { + validateGetObjectState(columnIndex); + + ColumnInfo columnInfo = columnInfos.get(columnIndex); + ColumnInfoTypeName requiredType = columnInfo.getTypeName(); + String arrowMetadata = currentChunkIterator.getType(columnIndex); + if (arrowMetadata == null) { + arrowMetadata = columnInfo.getTypeText(); + } + + return ArrowStreamResult.getObjectWithComplexTypeHandling( + session, currentChunkIterator, columnIndex, requiredType, arrowMetadata, columnInfo); + } + + /** + * Validates the state before getting an object at the specified column index. + * + * @param columnIndex the zero-based column index to validate + * @throws DatabricksSQLException if the result is closed, cursor is invalid, or column index is + * out of bounds + */ + private void validateGetObjectState(int columnIndex) throws DatabricksSQLException { + if (isClosed) { + LOGGER.warn("Attempted to get object from closed result"); + throw new DatabricksSQLException( + "Result is already closed", DatabricksDriverErrorCode.STATEMENT_CLOSED); + } + if (globalRowIndex == -1) { + LOGGER.warn("Attempted to get object before calling next()"); + throw new DatabricksSQLException( + "Cursor is before first row", DatabricksDriverErrorCode.INVALID_STATE); + } + if (currentChunkIterator == null) { + LOGGER.warn("No current chunk available when getting object"); + throw new DatabricksSQLException( + "No current chunk available", DatabricksDriverErrorCode.INVALID_STATE); + } + if (columnIndex < 0 || columnIndex >= columnInfos.size()) { + LOGGER.warn("Column index {} out of bounds (size: {})", columnIndex, columnInfos.size()); + throw new DatabricksSQLException( + "Column index out of bounds " + columnIndex, DatabricksDriverErrorCode.INVALID_STATE); + } + } + + /** + * Gets the current row index (0-based). Returns -1 if before the first row. + * + * @return the current row index + */ + @Override + public long getCurrentRow() { + return globalRowIndex; + } + + /** + * Moves the cursor to the next row. Fetches additional data from server if needed. + * + * @return true if there is a next row, false if at the end + * @throws DatabricksSQLException if an error occurs while fetching data + */ + @Override + public boolean next() throws SQLException { + if (isClosed || hasReachedEnd) { + return false; + } + + if (!hasNext()) { + return false; + } + + // Check if we've reached the maxRows limit + boolean hasRowLimit = maxRows > 0; + if (hasRowLimit && globalRowIndex + 1 >= maxRows) { + hasReachedEnd = true; + return false; + } + + // Try to advance in current chunk + if (currentChunkIterator != null && currentChunkIterator.hasNextRow()) { + boolean advanced = currentChunkIterator.nextRow(); + if (advanced) { + globalRowIndex++; + return true; + } + } + + // Need to fetch next chunk + while (currentResponse.hasMoreRows) { + fetchNextChunk(); + + // If we got a chunk with data, advance to first row + if (currentChunkIterator != null && currentChunkIterator.hasNextRow()) { + boolean advanced = currentChunkIterator.nextRow(); + if (advanced) { + globalRowIndex++; + return true; + } + } + } + + // No more data available + hasReachedEnd = true; + return false; + } + + /** + * Checks if there are more rows available without advancing the cursor. + * + * @return true if there are more rows, false otherwise + */ + @Override + public boolean hasNext() { + if (isClosed || hasReachedEnd) { + return false; + } + + // Check maxRows limit + boolean hasRowLimit = maxRows > 0; + if (hasRowLimit && globalRowIndex + 1 >= maxRows) { + return false; + } + + // Check if there are more rows in current chunk + if (currentChunkIterator != null && currentChunkIterator.hasNextRow()) { + return true; + } + + // Check if there are more chunks to fetch + return currentResponse.hasMoreRows; + } + + /** Closes this result and releases associated resources. */ + @Override + public void close() { + this.isClosed = true; + if (currentChunk != null) { + currentChunk.releaseChunk(); + } + this.currentChunk = null; + this.currentChunkIterator = null; + this.currentResponse = null; + LOGGER.debug( + "LazyThriftInlineArrowResult closed after fetching {} total rows", totalRowsFetched); + } + + /** + * Gets the number of rows in the current chunk. + * + * @return the number of rows in the current chunk + */ + @Override + public long getRowCount() { + return currentChunk != null ? currentChunk.numRows : 0; + } + + /** + * Gets the chunk count. Always returns 0 for lazy thrift inline arrow results. + * + * @return 0 (lazy results don't use chunks in the same sense as buffered results) + */ + @Override + public long getChunkCount() { + return 0; + } + + /** + * Gets the Arrow metadata for the current chunk. + * + * @return list of arrow metadata strings, or null if no chunk is loaded + * @throws DatabricksSQLException if an error occurs + */ + public List getArrowMetadata() throws DatabricksSQLException { + if (currentChunk == null) { + return null; + } + return currentChunk.getArrowMetadata(); + } + + private void loadCurrentChunk() throws DatabricksSQLException { + try { + ByteArrayInputStream byteStream = + createArrowByteStream(cachedSchema, currentResponse, getClass()); + long rowCount = getTotalRowsInResponse(currentResponse); + + ArrowResultChunk.Builder builder = + ArrowResultChunk.builder().withInputStream(byteStream, rowCount); + + if (statement != null) { + builder.withStatementId(statement.getStatementId()); + } + + currentChunk = builder.build(); + currentChunkIterator = currentChunk.getChunkIterator(); + totalRowsFetched += rowCount; + + LOGGER.debug( + "Loaded arrow chunk with {} rows, total fetched: {}", rowCount, totalRowsFetched); + } catch (DatabricksParsingException e) { + LOGGER.error("Failed to load current chunk: {}", e.getMessage()); + // Clean up any partially loaded chunk to prevent memory leaks + if (currentChunk != null) { + currentChunk.releaseChunk(); + currentChunk = null; + } + currentChunkIterator = null; + hasReachedEnd = true; + throw new DatabricksSQLException( + "Failed to process arrow data", DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR); + } + } + + /** + * Fetches the next chunk of data from the server and creates arrow chunks. + * + * @throws SQLException if the fetch operation fails + */ + private void fetchNextChunk() throws SQLException { + try { + LOGGER.debug("Fetching next arrow chunk, current total rows fetched: {}", totalRowsFetched); + currentResponse = session.getDatabricksClient().getMoreResults(statement); + + // Release previous chunk to free memory + if (currentChunk != null) { + currentChunk.releaseChunk(); + } + + loadCurrentChunk(); + + LOGGER.debug( + "Fetched arrow chunk with {} rows, hasMoreRows: {}", + currentChunk.numRows, + currentResponse.hasMoreRows); + } catch (DatabricksSQLException e) { + LOGGER.error("Failed to fetch next arrow chunk: {}", e.getMessage()); + hasReachedEnd = true; + throw e; + } + } + + /** + * Gets the total number of rows fetched from the server so far. + * + * @return the total number of rows fetched from the server + */ + long getTotalRowsFetched() { + return totalRowsFetched; + } + + /** + * Checks if all data has been fetched from the server. + * + * @return true if all data has been fetched (either reached end or maxRows limit) + */ + boolean isCompletelyFetched() { + return hasReachedEnd || !currentResponse.hasMoreRows; + } +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/RemoteChunkProvider.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/RemoteChunkProvider.java index f8ddf0a2b5..3fccf62c02 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/RemoteChunkProvider.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/RemoteChunkProvider.java @@ -12,6 +12,7 @@ import com.databricks.jdbc.model.core.ResultData; import com.databricks.jdbc.model.core.ResultManifest; import com.databricks.sdk.service.sql.BaseChunkInfo; +import java.sql.SQLException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ThreadFactory; @@ -47,7 +48,7 @@ public class RemoteChunkProvider extends AbstractRemoteChunkProvider= MAX_RETRIES) { LOGGER.error( diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkProvider.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkProvider.java index 5df90961a7..54aa036615 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkProvider.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkProvider.java @@ -1,5 +1,6 @@ package com.databricks.jdbc.api.impl.arrow; +import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.common.CompressionCodec; import com.databricks.jdbc.dbclient.IDatabricksHttpClient; import com.databricks.jdbc.dbclient.impl.common.StatementId; @@ -10,6 +11,7 @@ import com.databricks.jdbc.model.core.ChunkLinkFetchResult; import com.databricks.jdbc.model.core.ExternalLink; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.sql.SQLException; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ExecutionException; @@ -60,6 +62,7 @@ public class StreamingChunkProvider implements ChunkProvider { private final CompressionCodec compressionCodec; private final StatementId statementId; private final double cloudFetchSpeedThreshold; + private final IDatabricksConnectionContext connectionContext; // Chunk storage private final ConcurrentMap chunks = new ConcurrentHashMap<>(); @@ -125,6 +128,7 @@ public StreamingChunkProvider( int linkPrefetchWindow, int chunkReadyTimeoutSeconds, double cloudFetchSpeedThreshold, + IDatabricksConnectionContext connectionContext, ChunkLinkFetchResult initialLinks) throws DatabricksParsingException { @@ -136,6 +140,7 @@ public StreamingChunkProvider( this.linkPrefetchWindow = linkPrefetchWindow; this.chunkReadyTimeoutSeconds = chunkReadyTimeoutSeconds; this.cloudFetchSpeedThreshold = cloudFetchSpeedThreshold; + this.connectionContext = connectionContext; LOGGER.info( "Creating StreamingChunkProvider for statement {}: maxChunksInMemory={}, linkPrefetchWindow={}", @@ -368,7 +373,7 @@ private void linkPrefetchLoop() { LOGGER.debug("Link prefetch thread exiting for statement {}", statementId); } - private void fetchNextLinkBatch() throws DatabricksSQLException { + private void fetchNextLinkBatch() throws SQLException { if (endOfStreamReached || closed) { return; } @@ -461,6 +466,7 @@ private void createChunkFromLink(ExternalLink link) throws DatabricksParsingExce .withStatementId(statementId) .withChunkMetadata(chunkIndex, rowCount, rowOffset) .withChunkReadyTimeoutSeconds(chunkReadyTimeoutSeconds) + .withConnectionContext(connectionContext) .build(); chunk.setChunkLink(link); diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResult.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResult.java new file mode 100644 index 0000000000..7c9025c3f0 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResult.java @@ -0,0 +1,339 @@ +package com.databricks.jdbc.api.impl.arrow; + +import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_RESULT_ROW_LIMIT; +import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_STREAMING_BATCH_TIMEOUT_SECONDS; +import static com.databricks.jdbc.common.util.ArrowUtil.getColumnInfoList; + +import com.databricks.jdbc.api.impl.IExecutionResult; +import com.databricks.jdbc.api.impl.streaming.StreamingBatch; +import com.databricks.jdbc.api.impl.streaming.ThriftStreamingProvider; +import com.databricks.jdbc.api.impl.thrift.ThriftBatchFetcher; +import com.databricks.jdbc.api.impl.thrift.ThriftBatchFetcherImpl; +import com.databricks.jdbc.api.internal.IDatabricksSession; +import com.databricks.jdbc.api.internal.IDatabricksStatementInternal; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.log.JdbcLogger; +import com.databricks.jdbc.log.JdbcLoggerFactory; +import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; +import com.databricks.jdbc.model.core.ColumnInfo; +import com.databricks.jdbc.model.core.ColumnInfoTypeName; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.util.List; + +/** + * High-throughput streaming implementation for inline Arrow results. + * + *

Uses {@link ThriftStreamingProvider} for proactive batch prefetching, achieving throughput + * comparable to eager loading while maintaining the memory benefits of lazy loading. + * + *

Key features: + * + *

    + *
  • Background prefetch thread fetches Arrow batches ahead of consumption + *
  • Sliding window limits memory usage to a configurable number of batches + *
  • Non-blocking iteration when prefetch keeps up with consumption + *
  • Automatic native memory cleanup via type-safe release actions + *
  • Type-safe: Uses generic {@code ThriftStreamingProvider} + *
+ * + *

This implementation replaces {@code LazyThriftInlineArrowResult} for improved throughput. + */ +public class StreamingInlineArrowResult implements IExecutionResult { + + private static final JdbcLogger LOGGER = + JdbcLoggerFactory.getLogger(StreamingInlineArrowResult.class); + + // Streaming infrastructure - type-safe generic provider + private final ThriftStreamingProvider provider; + private final IDatabricksSession session; + + // Current position + private StreamingBatch currentBatch; + private ArrowResultChunkIterator currentChunkIterator; + private long globalRowIndex; + + // Metadata + private List columnInfos; + private final int maxRows; + + // State + private boolean hasReachedEnd; + private volatile boolean isClosed; + + /** + * Creates a new StreamingInlineArrowResult. + * + *

Configuration values (maxBatchesInMemory, timeout) are read from the session's connection + * context. + * + * @param initialResponse The initial Thrift response containing the first Arrow batch + * @param statement The statement that generated this result + * @param session The session for fetching additional batches + * @throws DatabricksSQLException if initialization fails + */ + public StreamingInlineArrowResult( + TFetchResultsResp initialResponse, + IDatabricksStatementInternal statement, + IDatabricksSession session) + throws DatabricksSQLException { + + this.session = session; + this.maxRows = statement != null ? statement.getMaxRows() : DEFAULT_RESULT_ROW_LIMIT; + + this.globalRowIndex = -1; + this.hasReachedEnd = false; + this.isClosed = false; + + // Initialize column info from metadata + this.columnInfos = getColumnInfoList(initialResponse.getResultSetMetadata()); + + // Create batch fetcher and type-safe generic provider for Arrow + ThriftBatchFetcher fetcher = new ThriftBatchFetcherImpl(session, statement); + this.provider = + ThriftStreamingProvider.forInlineArrow( + fetcher, + initialResponse, + statement != null ? statement.getStatementId() : null, + session.getConnectionContext().getThriftMaxBatchesInMemory(), + DEFAULT_STREAMING_BATCH_TIMEOUT_SECONDS); + + // Move to first batch (check nextBatch() return value to handle empty initial batches) + if (provider.nextBatch()) { + currentBatch = provider.getCurrentBatch(); + // Type-safe: getData() returns ArrowResultChunk directly! + currentChunkIterator = currentBatch.getData().getChunkIterator(); + } + + LOGGER.debug( + "StreamingInlineArrowResult initialized - firstBatchRows={}, maxRows={}, maxBatchesInMemory={}", + currentBatch != null ? currentBatch.getRowCount() : 0, + maxRows, + session.getConnectionContext().getThriftMaxBatchesInMemory()); + } + + /** + * Gets the value at the specified column index for the current row. + * + * @param columnIndex the zero-based column index + * @return the value at the specified column + * @throws DatabricksSQLException if access fails + */ + @Override + public Object getObject(int columnIndex) throws DatabricksSQLException { + validateGetObjectState(columnIndex); + + ColumnInfo columnInfo = columnInfos.get(columnIndex); + ColumnInfoTypeName requiredType = columnInfo.getTypeName(); + String arrowMetadata = currentChunkIterator.getType(columnIndex); + if (arrowMetadata == null) { + arrowMetadata = columnInfo.getTypeText(); + } + + // Use shared complex type handling from ArrowStreamResult + return ArrowStreamResult.getObjectWithComplexTypeHandling( + session, currentChunkIterator, columnIndex, requiredType, arrowMetadata, columnInfo); + } + + /** Validates state before getting an object. */ + private void validateGetObjectState(int columnIndex) throws DatabricksSQLException { + if (isClosed) { + LOGGER.error("Attempted to access closed result"); + throw new DatabricksSQLException( + "Result is closed", DatabricksDriverErrorCode.STATEMENT_CLOSED); + } + if (globalRowIndex == -1) { + LOGGER.error("Attempted to access data before first row"); + throw new DatabricksSQLException( + "Cursor is before first row", DatabricksDriverErrorCode.INVALID_STATE); + } + if (currentChunkIterator == null) { + LOGGER.error("No current chunk available at row {}", globalRowIndex); + throw new DatabricksSQLException( + "No current chunk available", DatabricksDriverErrorCode.INVALID_STATE); + } + if (columnIndex < 0 || columnIndex >= columnInfos.size()) { + LOGGER.error("Column index {} out of bounds (0-{})", columnIndex, columnInfos.size() - 1); + throw new DatabricksSQLException( + "Column index out of bounds: " + columnIndex, DatabricksDriverErrorCode.INVALID_STATE); + } + } + + /** + * Gets the current row index (0-based). Returns -1 if before the first row. + * + * @return the current row index + */ + @Override + public long getCurrentRow() { + return globalRowIndex; + } + + /** + * Moves the cursor to the next row. Fetches additional batches from server if needed. + * + * @return true if there is a next row, false if at the end + * @throws DatabricksSQLException if an error occurs + */ + @Override + public boolean next() throws DatabricksSQLException { + if (isClosed || hasReachedEnd) { + return false; + } + + if (!hasNext()) { + return false; + } + + // Check maxRows limit + boolean hasRowLimit = maxRows > 0; + if (hasRowLimit && globalRowIndex + 1 >= maxRows) { + hasReachedEnd = true; + return false; + } + + globalRowIndex++; + + // Try to move to next row in current chunk + if (currentChunkIterator != null && currentChunkIterator.hasNextRow()) { + currentChunkIterator.nextRow(); + return true; + } + + // Need to move to next batch + if (provider.hasNextBatch()) { + provider.nextBatch(); + currentBatch = provider.getCurrentBatch(); + + if (currentBatch != null) { + // Type-safe: getData() returns ArrowResultChunk directly! + ArrowResultChunk chunk = currentBatch.getData(); + if (chunk == null) { + LOGGER.warn("Batch {} has null data", currentBatch.getBatchIndex()); + hasReachedEnd = true; + globalRowIndex--; + return false; + } + currentChunkIterator = chunk.getChunkIterator(); + currentChunkIterator.nextRow(); + + LOGGER.debug( + "Moved to batch {} - globalRow={}, batchesInMemory={}", + currentBatch.getBatchIndex(), + globalRowIndex, + provider.getBatchesInMemory()); + + return true; + } + } + + // No more data + hasReachedEnd = true; + globalRowIndex--; + return false; + } + + /** + * Checks if there are more rows available without advancing the cursor. + * + * @return true if there are more rows, false otherwise + */ + @Override + public boolean hasNext() { + if (isClosed || hasReachedEnd) { + return false; + } + + // Check maxRows limit + boolean hasRowLimit = maxRows > 0; + if (hasRowLimit && globalRowIndex + 1 >= maxRows) { + return false; + } + + // Check current chunk + if (currentChunkIterator != null && currentChunkIterator.hasNextRow()) { + return true; + } + + // Check if more batches available + return provider.hasNextBatch(); + } + + /** Closes this result and releases associated resources. */ + @Override + public void close() { + if (isClosed) { + return; + } + + long totalRows = provider.getTotalRowsFetched(); + isClosed = true; + currentBatch = null; + currentChunkIterator = null; + + // Provider will release all Arrow chunks using the type-safe Consumer + provider.close(); + + LOGGER.debug("Closed - totalRowsFetched={}, rowsConsumed={}", totalRows, globalRowIndex + 1); + } + + /** + * Gets the number of rows in the current batch. + * + * @return the number of rows in the current batch + */ + @Override + public long getRowCount() { + return currentBatch != null ? currentBatch.getRowCount() : 0; + } + + /** + * Gets the chunk count. Always returns 0 for streaming results. + * + * @return 0 + */ + @Override + public long getChunkCount() { + return 0; + } + + /** + * Gets the Arrow metadata for the current chunk. + * + * @return list of arrow metadata strings, or null if no chunk is loaded + * @throws DatabricksSQLException if an error occurs + */ + public List getArrowMetadata() throws DatabricksSQLException { + if (currentBatch == null) { + return null; + } + ArrowResultChunk chunk = currentBatch.getData(); + return chunk != null ? chunk.getArrowMetadata() : null; + } + + /** + * Gets the total number of rows fetched from the server so far. + * + * @return the total rows fetched + */ + public long getTotalRowsFetched() { + return provider.getTotalRowsFetched(); + } + + /** + * Checks if all data has been fetched from the server. + * + * @return true if end of stream reached + */ + public boolean isCompletelyFetched() { + return hasReachedEnd || provider.isEndOfStreamReached(); + } + + /** + * Gets the number of batches currently in memory. + * + * @return the batch count in memory + */ + public int getBatchesInMemory() { + return provider.getBatchesInMemory(); + } +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/ThriftChunkLinkFetcher.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/ThriftChunkLinkFetcher.java index d38653271f..819c1e7d9f 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/ThriftChunkLinkFetcher.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/ThriftChunkLinkFetcher.java @@ -8,6 +8,7 @@ import com.databricks.jdbc.model.core.ChunkLinkFetchResult; import com.databricks.jdbc.model.core.ExternalLink; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.sql.SQLException; /** * ChunkLinkFetcher implementation for the Thrift client. @@ -32,7 +33,7 @@ public ThriftChunkLinkFetcher(IDatabricksSession session, StatementId statementI @Override public ChunkLinkFetchResult fetchLinks(long startChunkIndex, long startRowOffset) - throws DatabricksSQLException { + throws SQLException { // Thrift uses startRowOffset with FETCH_ABSOLUTE; startChunkIndex is used for metadata LOGGER.debug( "Fetching links starting from chunk index {}, row offset {} for statement {}", @@ -46,7 +47,7 @@ public ChunkLinkFetchResult fetchLinks(long startChunkIndex, long startRowOffset } @Override - public ExternalLink refetchLink(long chunkIndex, long rowOffset) throws DatabricksSQLException { + public ExternalLink refetchLink(long chunkIndex, long rowOffset) throws SQLException { // Thrift uses rowOffset with FETCH_ABSOLUTE LOGGER.info( "Refetching expired link for chunk {}, row offset {} of statement {}", diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/incubator/RemoteChunkProviderV2.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/incubator/RemoteChunkProviderV2.java index e2321c2d2d..848b9be611 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/incubator/RemoteChunkProviderV2.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/incubator/RemoteChunkProviderV2.java @@ -14,6 +14,7 @@ import com.databricks.jdbc.model.core.ResultManifest; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; import com.databricks.sdk.service.sql.BaseChunkInfo; +import java.sql.SQLException; import java.util.concurrent.ExecutionException; /** @@ -51,7 +52,7 @@ public RemoteChunkProviderV2( IDatabricksHttpClient httpClient, int maxParallelChunkDownloadsPerQuery, CompressionCodec compressionCodec) - throws DatabricksSQLException { + throws SQLException { super( parentStatement, resultsResp, diff --git a/src/main/java/com/databricks/jdbc/api/impl/converters/ArrowToJavaObjectConverter.java b/src/main/java/com/databricks/jdbc/api/impl/converters/ArrowToJavaObjectConverter.java index 348e4f9225..5c985def2d 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/converters/ArrowToJavaObjectConverter.java +++ b/src/main/java/com/databricks/jdbc/api/impl/converters/ArrowToJavaObjectConverter.java @@ -14,6 +14,7 @@ import java.math.RoundingMode; import java.sql.Date; import java.sql.Timestamp; +import java.sql.Types; import java.time.*; import java.time.format.DateTimeFormatter; import java.time.format.DateTimeParseException; @@ -140,8 +141,11 @@ public static Object convert( IntervalConverter ic = new IntervalConverter(arrowMetadata); return ic.toLiteral(object); case GEOMETRY: + return ConverterHelper.getConverterForColumnType(Types.OTHER, GEOMETRY) + .toDatabricksGeometry(object); case GEOGRAPHY: - return convertToGeospatial(object, requiredType); + return ConverterHelper.getConverterForColumnType(Types.OTHER, GEOGRAPHY) + .toDatabricksGeography(object); case NULL: return null; default: @@ -169,21 +173,6 @@ private static Object convertToStruct(Object object, String arrowMetadata) return parser.parseJsonStringToDbStruct(object.toString(), arrowMetadata); } - private static AbstractDatabricksGeospatial convertToGeospatial( - Object object, ColumnInfoTypeName type) throws DatabricksSQLException { - String ewkt = convertToString(object); - - // Parse EWKT to extract SRID from data - // SRID is always present in EWKT unless it's 0, in which case it is handled in - // WKTConverter.extractSRIDFromEWKT() - int srid = WKTConverter.extractSRIDFromEWKT(ewkt); - String cleanWkt = WKTConverter.removeSRIDFromEWKT(ewkt); - - return type == ColumnInfoTypeName.GEOMETRY - ? new DatabricksGeometry(cleanWkt, srid) - : new DatabricksGeography(cleanWkt, srid); - } - private static Object convertToTimestamp(Object object, Optional timeZoneOpt) throws DatabricksSQLException { if (object instanceof Text) { diff --git a/src/main/java/com/databricks/jdbc/api/impl/streaming/ColumnarResponseProcessor.java b/src/main/java/com/databricks/jdbc/api/impl/streaming/ColumnarResponseProcessor.java new file mode 100644 index 0000000000..d850344fb1 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/streaming/ColumnarResponseProcessor.java @@ -0,0 +1,56 @@ +package com.databricks.jdbc.api.impl.streaming; + +import static com.databricks.jdbc.common.util.DatabricksThriftUtil.createColumnarView; + +import com.databricks.jdbc.api.impl.ColumnarRowView; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.log.JdbcLogger; +import com.databricks.jdbc.log.JdbcLoggerFactory; +import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; +import java.util.function.Consumer; + +/** + * Processes Thrift responses into columnar view data. + * + *

This processor converts {@link TFetchResultsResp} into {@link ColumnarRowView} for Thrift + * columnar result handling. + */ +public class ColumnarResponseProcessor implements ThriftResponseProcessor { + + private static final JdbcLogger LOGGER = + JdbcLoggerFactory.getLogger(ColumnarResponseProcessor.class); + + @Override + public StreamingBatch processInitialResponse(TFetchResultsResp response) + throws DatabricksSQLException { + LOGGER.debug("Processing initial columnar response"); + return processResponse(response, 0, 0); + } + + @Override + public StreamingBatch processResponse( + TFetchResultsResp response, long batchIndex, long rowOffset) throws DatabricksSQLException { + + StreamingBatch batch = + new StreamingBatch<>(batchIndex, rowOffset, getReleaseAction()); + + ColumnarRowView view = createColumnarView(response.getResults()); + batch.setData(view, view.getRowCount(), response.hasMoreRows); + + LOGGER.debug( + "Processed columnar batch {}: rows={}, hasMoreRows={}", + batchIndex, + view.getRowCount(), + response.hasMoreRows); + + return batch; + } + + @Override + public Consumer getReleaseAction() { + // ColumnarRowView doesn't need explicit cleanup - just let GC handle it + return view -> { + /* no-op */ + }; + } +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessor.java b/src/main/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessor.java new file mode 100644 index 0000000000..477359f2bc --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessor.java @@ -0,0 +1,107 @@ +package com.databricks.jdbc.api.impl.streaming; + +import static com.databricks.jdbc.common.util.ArrowUtil.createArrowByteStream; +import static com.databricks.jdbc.common.util.ArrowUtil.getSerializedSchema; +import static com.databricks.jdbc.common.util.ArrowUtil.getTotalRowsInResponse; + +import com.databricks.jdbc.api.impl.arrow.ArrowResultChunk; +import com.databricks.jdbc.dbclient.impl.common.StatementId; +import com.databricks.jdbc.exception.DatabricksParsingException; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.log.JdbcLogger; +import com.databricks.jdbc.log.JdbcLoggerFactory; +import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.io.ByteArrayInputStream; +import java.util.function.Consumer; + +/** + * Processes Thrift responses into Arrow chunks. + * + *

This processor converts {@link TFetchResultsResp} into {@link ArrowResultChunk} for inline + * Arrow result handling. It caches the Arrow schema from the first response for use in subsequent + * batches. + */ +public class InlineArrowResponseProcessor implements ThriftResponseProcessor { + + private static final JdbcLogger LOGGER = + JdbcLoggerFactory.getLogger(InlineArrowResponseProcessor.class); + + private final StatementId statementId; + private volatile byte[] + cachedSchema; // Cache schema from first response, volatile for visibility across threads + + /** + * Creates a new inline Arrow response processor. + * + * @param statementId The statement ID for logging and chunk creation + */ + public InlineArrowResponseProcessor(StatementId statementId) { + this.statementId = statementId; + } + + @Override + public StreamingBatch processInitialResponse(TFetchResultsResp response) + throws DatabricksSQLException { + LOGGER.debug("Processing initial inline Arrow response"); + // Cache the schema for subsequent batches + try { + this.cachedSchema = getSerializedSchema(response.getResultSetMetadata()); + } catch (DatabricksParsingException e) { + LOGGER.error("Failed to serialize Arrow schema: {}", e.getMessage(), e); + throw new DatabricksSQLException( + "Failed to serialize Arrow schema", + e, + DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR); + } + return processResponse(response, 0, 0); + } + + @Override + public StreamingBatch processResponse( + TFetchResultsResp response, long batchIndex, long rowOffset) throws DatabricksSQLException { + + StreamingBatch batch = + new StreamingBatch<>(batchIndex, rowOffset, getReleaseAction()); + + try { + ByteArrayInputStream byteStream = createArrowByteStream(cachedSchema, response, getClass()); + long rowCount = getTotalRowsInResponse(response); + + ArrowResultChunk.Builder builder = + ArrowResultChunk.builder().withInputStream(byteStream, rowCount); + + if (statementId != null) { + builder.withStatementId(statementId); + } + + ArrowResultChunk chunk = builder.build(); + batch.setData(chunk, rowCount, response.hasMoreRows); + + LOGGER.debug( + "Processed inline Arrow batch {}: rows={}, hasMoreRows={}", + batchIndex, + rowCount, + response.hasMoreRows); + + return batch; + + } catch (DatabricksParsingException e) { + LOGGER.error("Failed to process inline Arrow batch {}: {}", batchIndex, e.getMessage(), e); + batch.setError(e); + throw new DatabricksSQLException( + "Failed to process Arrow data", e, DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR); + } + } + + @Override + public Consumer getReleaseAction() { + // Arrow chunks require explicit native memory cleanup + return chunk -> { + if (chunk != null) { + chunk.releaseChunk(); + LOGGER.debug("Released Arrow chunk native memory"); + } + }; + } +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/streaming/StreamingBatch.java b/src/main/java/com/databricks/jdbc/api/impl/streaming/StreamingBatch.java new file mode 100644 index 0000000000..355ffe16a5 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/streaming/StreamingBatch.java @@ -0,0 +1,183 @@ +package com.databricks.jdbc.api.impl.streaming; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.function.Consumer; + +/** + * Type-safe batch container for streaming results. + * + *

This generic container holds batch data of any type (ColumnarRowView or ArrowResultChunk) and + * manages batch lifecycle including status tracking and resource cleanup. + * + * @param The type of data held (ColumnarRowView or ArrowResultChunk) + */ +public class StreamingBatch { + + /** Batch lifecycle status. */ + public enum Status { + PENDING, // Not yet fetched + FETCHING, // Fetch in progress + READY, // Data available + ERROR, // Fetch failed + RELEASED // Memory released + } + + private final long batchIndex; + private final long rowOffset; + private volatile Status status; + + // Type-safe data storage + private volatile T data; + private final Consumer releaseAction; + + // Metadata from response + private volatile long rowCount; + private volatile boolean hasMoreRows; + private volatile Throwable error; + private final CompletableFuture> readyFuture; + + /** + * Creates a new batch with the specified release action. + * + * @param batchIndex The batch index + * @param rowOffset The starting row offset + * @param releaseAction Action to call when releasing data (e.g., ArrowResultChunk::releaseChunk) + */ + public StreamingBatch(long batchIndex, long rowOffset, Consumer releaseAction) { + this.batchIndex = batchIndex; + this.rowOffset = rowOffset; + this.releaseAction = releaseAction; + this.status = Status.PENDING; + this.readyFuture = new CompletableFuture<>(); + } + + /** + * Sets the batch data and marks as ready. + * + * @param data The batch data + * @param rowCount Number of rows in this batch + * @param hasMoreRows Whether more batches are available after this one + */ + public void setData(T data, long rowCount, boolean hasMoreRows) { + this.data = data; + this.rowCount = rowCount; + this.hasMoreRows = hasMoreRows; + this.status = Status.READY; + this.readyFuture.complete(this); + } + + /** + * Gets the typed data. No casting required! + * + * @return The batch data + */ + public T getData() { + return data; + } + + /** Sets the batch as currently fetching. */ + public void setFetching() { + this.status = Status.FETCHING; + } + + /** + * Sets error state. + * + * @param error The error that occurred + */ + public void setError(Throwable error) { + this.error = error; + this.status = Status.ERROR; + this.readyFuture.completeExceptionally(error); + } + + /** + * Waits for the batch to be ready. + * + * @param timeoutSeconds Maximum time to wait in seconds + * @throws InterruptedException if waiting is interrupted + * @throws ExecutionException if the batch fetch failed + * @throws TimeoutException if the timeout is exceeded + */ + public void waitUntilReady(long timeoutSeconds) + throws InterruptedException, ExecutionException, TimeoutException { + readyFuture.get(timeoutSeconds, TimeUnit.SECONDS); + } + + /** + * Checks if the batch is ready. + * + * @return true if data is available + */ + public boolean isReady() { + return status == Status.READY; + } + + /** Releases the batch data using the type-specific release action. */ + public void release() { + if (data != null && releaseAction != null) { + releaseAction.accept(data); + } + this.data = null; + this.status = Status.RELEASED; + } + + // Getters + + /** + * Gets the batch index. + * + * @return The zero-based batch index + */ + public long getBatchIndex() { + return batchIndex; + } + + /** + * Gets the row offset. + * + * @return The starting row offset for this batch + */ + public long getRowOffset() { + return rowOffset; + } + + /** + * Gets the row count. + * + * @return The number of rows in this batch + */ + public long getRowCount() { + return rowCount; + } + + /** + * Checks if more rows are available after this batch. + * + * @return true if more batches are available + */ + public boolean hasMoreRows() { + return hasMoreRows; + } + + /** + * Gets the current status. + * + * @return The batch status + */ + public Status getStatus() { + return status; + } + + /** + * Gets the error if status is ERROR. + * + * @return The error, or null if no error occurred + */ + public Throwable getError() { + return error; + } +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/streaming/ThriftResponseProcessor.java b/src/main/java/com/databricks/jdbc/api/impl/streaming/ThriftResponseProcessor.java new file mode 100644 index 0000000000..1e88f3a518 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/streaming/ThriftResponseProcessor.java @@ -0,0 +1,54 @@ +package com.databricks.jdbc.api.impl.streaming; + +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; +import java.util.function.Consumer; + +/** + * Processes a Thrift fetch response into typed result data. + * + *

Implementations handle the conversion of {@link TFetchResultsResp} into specific data types + * (ColumnarRowView for columnar results, ArrowResultChunk for inline Arrow results). + * + * @param The type of data produced (ColumnarRowView or ArrowResultChunk) + */ +public interface ThriftResponseProcessor { + + /** + * Processes the initial response and returns batch data. + * + *

This method is called once when the streaming provider is created. It may perform additional + * initialization such as caching schema information. + * + * @param response The initial Thrift fetch response + * @return A StreamingBatch containing the processed data + * @throws DatabricksSQLException if processing fails + */ + StreamingBatch processInitialResponse(TFetchResultsResp response) + throws DatabricksSQLException; + + /** + * Processes a fetched response and returns batch data. + * + *

This method is called for each subsequent batch fetched from the server. + * + * @param response The Thrift fetch response + * @param batchIndex The zero-based batch index + * @param rowOffset The starting row offset for this batch + * @return A StreamingBatch containing the processed data + * @throws DatabricksSQLException if processing fails + */ + StreamingBatch processResponse(TFetchResultsResp response, long batchIndex, long rowOffset) + throws DatabricksSQLException; + + /** + * Creates a release action for the data type. + * + *

This action is called when the batch is released from memory. For types that require + * explicit cleanup (like ArrowResultChunk with native memory), this should perform that cleanup. + * For types that don't need cleanup (like ColumnarRowView), this can be a no-op. + * + * @return A Consumer that releases the data + */ + Consumer getReleaseAction(); +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/streaming/ThriftStreamingProvider.java b/src/main/java/com/databricks/jdbc/api/impl/streaming/ThriftStreamingProvider.java new file mode 100644 index 0000000000..fd0494e9d3 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/streaming/ThriftStreamingProvider.java @@ -0,0 +1,548 @@ +package com.databricks.jdbc.api.impl.streaming; + +import com.databricks.jdbc.api.impl.ColumnarRowView; +import com.databricks.jdbc.api.impl.arrow.ArrowResultChunk; +import com.databricks.jdbc.api.impl.thrift.ThriftBatchFetcher; +import com.databricks.jdbc.dbclient.impl.common.StatementId; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.log.JdbcLogger; +import com.databricks.jdbc.log.JdbcLoggerFactory; +import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.sql.SQLException; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.locks.Condition; +import java.util.concurrent.locks.ReentrantLock; + +/** + * Type-safe streaming provider for Thrift-based results. + * + *

This generic provider works with any data type through pluggable processors. It provides + * proactive prefetching with a sliding window to maximize throughput while bounding memory usage. + * + *

Key features: + * + *

    + *
  • Type-safe: {@code getData()} returns the correct type without casting + *
  • Prefetching: Background thread fetches ahead of consumer + *
  • Memory-bounded: Sliding window limits batches in memory + *
  • Pluggable: Processors handle type-specific data conversion + *
+ * + * @param The type of data in each batch (ColumnarRowView or ArrowResultChunk) + */ +public class ThriftStreamingProvider implements AutoCloseable { + + private static final JdbcLogger LOGGER = + JdbcLoggerFactory.getLogger(ThriftStreamingProvider.class); + private static final String PREFETCH_THREAD_NAME = "databricks-inline-prefetcher"; + + // Configuration + private final int maxBatchesInMemory; + private final int batchReadyTimeoutSeconds; + + // Dependencies - note the generic processor type + private final ThriftBatchFetcher batchFetcher; + private final ThriftResponseProcessor processor; + + // Type-safe batch storage + private final ConcurrentMap> batches = new ConcurrentHashMap<>(); + + // Position tracking + private final AtomicLong currentBatchIndex = new AtomicLong(-1); + private final AtomicLong highestFetchedBatchIndex = new AtomicLong(-1); + private final AtomicLong nextBatchToFetch = new AtomicLong(1); // 0 is initial batch + private final AtomicLong totalRowsFetched = new AtomicLong(0); + private final AtomicLong nextRowOffset = new AtomicLong(0); + + // State and synchronization + private volatile boolean endOfStreamReached = false; + private volatile boolean closed = false; + private volatile DatabricksSQLException prefetchError = null; + private final ReentrantLock prefetchLock = new ReentrantLock(); + private final Condition batchAvailable = prefetchLock.newCondition(); + private final Condition consumerAdvanced = prefetchLock.newCondition(); + private final AtomicInteger batchesInMemory = new AtomicInteger(0); + + // Prefetch thread + private final Thread prefetchThread; + + // ==================== Factory Methods ==================== + + /** + * Creates a streaming provider for Thrift Columnar results. + * + * @param fetcher The batch fetcher for retrieving data from server + * @param initialResponse The initial Thrift fetch response + * @param maxBatchesInMemory Maximum batches to keep in memory (sliding window) + * @param timeoutSeconds Timeout waiting for batch to be ready + * @return A type-safe provider that produces ColumnarRowView data + * @throws DatabricksSQLException if initialization fails + */ + public static ThriftStreamingProvider forColumnar( + ThriftBatchFetcher fetcher, + TFetchResultsResp initialResponse, + int maxBatchesInMemory, + int timeoutSeconds) + throws DatabricksSQLException { + return new ThriftStreamingProvider<>( + fetcher, + initialResponse, + new ColumnarResponseProcessor(), + maxBatchesInMemory, + timeoutSeconds); + } + + /** + * Creates a streaming provider for Inline Arrow results. + * + * @param fetcher The batch fetcher for retrieving data from server + * @param initialResponse The initial Thrift fetch response + * @param statementId The statement ID for chunk creation + * @param maxBatchesInMemory Maximum batches to keep in memory (sliding window) + * @param timeoutSeconds Timeout waiting for batch to be ready + * @return A type-safe provider that produces ArrowResultChunk data + * @throws DatabricksSQLException if initialization fails + */ + public static ThriftStreamingProvider forInlineArrow( + ThriftBatchFetcher fetcher, + TFetchResultsResp initialResponse, + StatementId statementId, + int maxBatchesInMemory, + int timeoutSeconds) + throws DatabricksSQLException { + return new ThriftStreamingProvider<>( + fetcher, + initialResponse, + new InlineArrowResponseProcessor(statementId), + maxBatchesInMemory, + timeoutSeconds); + } + + // ==================== Constructor ==================== + + private ThriftStreamingProvider( + ThriftBatchFetcher fetcher, + TFetchResultsResp initialResponse, + ThriftResponseProcessor processor, + int maxBatchesInMemory, + int timeoutSeconds) + throws DatabricksSQLException { + + // Validate required parameters + if (initialResponse == null) { + LOGGER.error("Cannot create ThriftStreamingProvider: initialResponse is null"); + throw new IllegalArgumentException("initialResponse cannot be null"); + } + if (fetcher == null) { + LOGGER.error("Cannot create ThriftStreamingProvider: fetcher is null"); + throw new IllegalArgumentException("fetcher cannot be null"); + } + if (processor == null) { + LOGGER.error("Cannot create ThriftStreamingProvider: processor is null"); + throw new IllegalArgumentException("processor cannot be null"); + } + + this.batchFetcher = fetcher; + this.processor = processor; + // We need at least 2 batches in memory to perform any kind of prefetching + int effectiveMaxBatchesInMemory = Math.max(2, maxBatchesInMemory); + if (maxBatchesInMemory < 2) { + LOGGER.warn( + "Configured maxBatchesInMemory={} is less than the minimum of 2; using 2 instead to enable prefetching.", + maxBatchesInMemory); + } + this.maxBatchesInMemory = effectiveMaxBatchesInMemory; + this.batchReadyTimeoutSeconds = timeoutSeconds; + + LOGGER.debug( + "Creating ThriftStreamingProvider: maxBatches={}, timeout={}s, processor={}", + this.maxBatchesInMemory, + timeoutSeconds, + processor.getClass().getSimpleName()); + + // Process initial batch using the generic processor + StreamingBatch initialBatch = processor.processInitialResponse(initialResponse); + batches.put(0L, initialBatch); + highestFetchedBatchIndex.set(0); + batchesInMemory.incrementAndGet(); + totalRowsFetched.addAndGet(initialBatch.getRowCount()); + nextRowOffset.set(initialBatch.getRowCount()); + + LOGGER.debug( + "Initial batch processed: rows={}, hasMoreRows={}", + initialBatch.getRowCount(), + initialBatch.hasMoreRows()); + + if (!initialBatch.hasMoreRows()) { + endOfStreamReached = true; + LOGGER.debug("Single batch result - all data in initial response"); + } + + // Start prefetch thread + this.prefetchThread = new Thread(this::prefetchLoop, PREFETCH_THREAD_NAME); + this.prefetchThread.setDaemon(true); + this.prefetchThread.start(); + + notifyConsumerAdvanced(); + } + + // ==================== Public API ==================== + + /** + * Checks if there are more batches available. + * + * @return true if more batches may be available + */ + public boolean hasNextBatch() { + if (closed) return false; + if (!endOfStreamReached) return true; + return currentBatchIndex.get() < highestFetchedBatchIndex.get(); + } + + /** + * Moves to the next batch. Releases the previous batch. + * + *

This method automatically skips empty batches (rowCount == 0), continuing to advance until a + * non-empty batch is found or no more batches are available. This ensures consumers never see + * empty batches and matches the behavior of lazy result implementations. + * + * @return true if moved to a non-empty batch, false if no more batches + * @throws DatabricksSQLException if an error occurred during prefetch + */ + public boolean nextBatch() throws DatabricksSQLException { + if (closed) return false; + + checkPrefetchError(); + + // Release previous batch (type-safe release via Consumer) + long prevIndex = currentBatchIndex.get(); + if (prevIndex >= 0) { + releaseBatch(prevIndex); + } + + // Keep advancing until we find a non-empty batch or run out of batches + while (hasNextBatch()) { + currentBatchIndex.incrementAndGet(); + notifyConsumerAdvanced(); + + StreamingBatch batch = getCurrentBatch(); + if (batch != null && batch.getRowCount() > 0) { + return true; // Found a non-empty batch + } + + // Empty batch - release it and try next + LOGGER.debug("Skipping empty batch {}", currentBatchIndex.get()); + releaseBatch(currentBatchIndex.get()); + } + + return false; // No more non-empty batches available + } + + /** + * Gets the current batch with type-safe data access. + * + *

No casting required - returns {@code StreamingBatch} with correctly typed data! + * + * @return The current batch, or null if before first batch + * @throws DatabricksSQLException if the batch cannot be retrieved + */ + public StreamingBatch getCurrentBatch() throws DatabricksSQLException { + long batchIdx = currentBatchIndex.get(); + if (batchIdx < 0) return null; + + checkPrefetchError(); + + StreamingBatch batch = batches.get(batchIdx); + + if (batch == null) { + LOGGER.debug("Batch {} not yet available, waiting for prefetch", batchIdx); + waitForBatchCreation(batchIdx); + batch = batches.get(batchIdx); + } + + if (batch == null) { + LOGGER.error("Batch {} not found after waiting", batchIdx); + throw new DatabricksSQLException( + "Batch " + batchIdx + " not found after waiting", + DatabricksDriverErrorCode.CHUNK_READY_ERROR); + } + + // Wait for batch to be ready + if (!batch.isReady()) { + try { + batch.waitUntilReady(batchReadyTimeoutSeconds); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.warn("Interrupted waiting for batch {}", batchIdx); + throw new DatabricksSQLException( + "Interrupted waiting for batch " + batchIdx, + e, + DatabricksDriverErrorCode.THREAD_INTERRUPTED_ERROR); + } catch (ExecutionException e) { + LOGGER.error("Failed to fetch batch {}: {}", batchIdx, e.getCause().getMessage(), e); + throw new DatabricksSQLException( + "Failed to fetch batch " + batchIdx, + e.getCause(), + DatabricksDriverErrorCode.CHUNK_READY_ERROR); + } catch (TimeoutException e) { + LOGGER.error( + "Timeout waiting for batch {} (timeout: {}s)", batchIdx, batchReadyTimeoutSeconds); + throw new DatabricksSQLException( + "Timeout waiting for batch " + + batchIdx + + " (timeout: " + + batchReadyTimeoutSeconds + + "s)", + DatabricksDriverErrorCode.CHUNK_READY_ERROR); + } + } + + return batch; + } + + /** + * Gets total rows fetched so far. + * + * @return The total row count + */ + public long getTotalRowsFetched() { + return totalRowsFetched.get(); + } + + /** + * Gets number of batches currently in memory. + * + * @return The batch count + */ + public int getBatchesInMemory() { + return batchesInMemory.get(); + } + + /** + * Checks if the end of stream has been reached. + * + * @return true if all batches have been fetched from the server + */ + public boolean isEndOfStreamReached() { + return endOfStreamReached; + } + + @Override + public void close() { + if (closed) return; + + LOGGER.debug("Closing ThriftStreamingProvider, total rows: {}", totalRowsFetched.get()); + closed = true; + + notifyConsumerAdvanced(); + notifyBatchAvailable(); + + // Interrupt and wait for prefetch thread to terminate before releasing resources + prefetchThread.interrupt(); + try { + prefetchThread.join(5000); // Wait up to 5s for clean shutdown + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.debug("Interrupted while waiting for prefetch thread to terminate"); + } + + // Release all batches using type-safe release action + // Continue releasing subsequent batches even if one fails + for (StreamingBatch batch : batches.values()) { + try { + batch.release(); + } catch (Exception e) { + LOGGER.warn("Error releasing batch during close: {}", e.getMessage(), e); + } + } + batches.clear(); + + try { + batchFetcher.close(); + } catch (Exception e) { + LOGGER.warn("Error closing batchFetcher: {}", e.getMessage(), e); + } + } + + // ==================== Prefetch Logic ==================== + + private void prefetchLoop() { + LOGGER.debug("Prefetch thread started"); + + while (!closed && !Thread.currentThread().isInterrupted()) { + try { + prefetchLock.lock(); + try { + while (!closed && !endOfStreamReached && batchesInMemory.get() >= maxBatchesInMemory) { + LOGGER.debug( + "Prefetch waiting: batches={}/{}", batchesInMemory.get(), maxBatchesInMemory); + consumerAdvanced.await(); + } + } finally { + prefetchLock.unlock(); + } + + if (closed || endOfStreamReached) break; + + fetchNextBatchInternal(); + + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.debug("Prefetch thread interrupted"); + break; + } catch (DatabricksSQLException e) { + LOGGER.error("Prefetch error: {}", e.getMessage()); + prefetchError = e; + notifyBatchAvailable(); + break; + } catch (Exception e) { + LOGGER.error("Unexpected prefetch error: {}", e.getMessage(), e); + prefetchError = + new DatabricksSQLException( + "Unexpected prefetch error: " + e.getMessage(), + e, + DatabricksDriverErrorCode.CHUNK_READY_ERROR); + notifyBatchAvailable(); + break; + } + } + + LOGGER.debug("Prefetch thread exiting"); + } + + private void fetchNextBatchInternal() throws SQLException { + long batchIndex = nextBatchToFetch.getAndIncrement(); + long rowOffset = nextRowOffset.get(); + + LOGGER.debug("Fetching batch {} at offset {}", batchIndex, rowOffset); + + // Fetch from server + TFetchResultsResp response = batchFetcher.fetchNextBatch(); + + // Process using type-safe generic processor + StreamingBatch batch = processor.processResponse(response, batchIndex, rowOffset); + + // Update state + batches.put(batchIndex, batch); + batchesInMemory.incrementAndGet(); + highestFetchedBatchIndex.updateAndGet(cur -> Math.max(cur, batchIndex)); + totalRowsFetched.addAndGet(batch.getRowCount()); + nextRowOffset.addAndGet(batch.getRowCount()); + + LOGGER.debug( + "Batch {} ready: rows={}, hasMore={}", + batchIndex, + batch.getRowCount(), + batch.hasMoreRows()); + + if (!batch.hasMoreRows()) { + endOfStreamReached = true; + LOGGER.debug("End of stream at batch {}", batchIndex); + } + + notifyBatchAvailable(); + } + + // ==================== Resource Management ==================== + + private void releaseBatch(long batchIndex) { + StreamingBatch batch = batches.remove(batchIndex); + if (batch != null) { + // Decrement counter BEFORE release to prevent prefetch stall if release() throws + batchesInMemory.decrementAndGet(); + try { + batch.release(); // Uses type-safe Consumer internally + } catch (Exception e) { + LOGGER.warn("Error releasing batch {}: {}", batchIndex, e.getMessage(), e); + } + LOGGER.debug("Released batch {}, batches in memory: {}", batchIndex, batchesInMemory.get()); + notifyConsumerAdvanced(); + } + } + + private void waitForBatchCreation(long batchIndex) throws DatabricksSQLException { + prefetchLock.lock(); + try { + long waitStartTime = System.currentTimeMillis(); + long timeoutMillis = batchReadyTimeoutSeconds * 1000L; + + while (!closed && !batches.containsKey(batchIndex)) { + checkPrefetchError(); + if (endOfStreamReached && batchIndex > highestFetchedBatchIndex.get()) { + LOGGER.error( + "Batch {} does not exist (highest fetched: {})", + batchIndex, + highestFetchedBatchIndex.get()); + throw new DatabricksSQLException( + "Batch " + + batchIndex + + " does not exist (highest: " + + highestFetchedBatchIndex.get() + + ")", + DatabricksDriverErrorCode.CHUNK_READY_ERROR); + } + + // Check for timeout + long elapsedMillis = System.currentTimeMillis() - waitStartTime; + if (elapsedMillis >= timeoutMillis) { + LOGGER.error( + "Timeout waiting for batch {} to be created (timeout: {}s)", + batchIndex, + batchReadyTimeoutSeconds); + throw new DatabricksSQLException( + "Timeout waiting for batch " + + batchIndex + + " to be created (timeout: " + + batchReadyTimeoutSeconds + + "s)", + DatabricksDriverErrorCode.CHUNK_READY_ERROR); + } + + try { + long remainingMillis = timeoutMillis - elapsedMillis; + batchAvailable.await(remainingMillis, TimeUnit.MILLISECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.warn("Interrupted waiting for batch {} creation", batchIndex); + throw new DatabricksSQLException( + "Interrupted waiting for batch", + e, + DatabricksDriverErrorCode.THREAD_INTERRUPTED_ERROR); + } + } + } finally { + prefetchLock.unlock(); + } + } + + private void checkPrefetchError() throws DatabricksSQLException { + if (prefetchError != null) { + LOGGER.error("Prefetch failed: {}", prefetchError.getMessage(), prefetchError); + throw new DatabricksSQLException( + "Prefetch failed: " + prefetchError.getMessage(), + prefetchError, + DatabricksDriverErrorCode.CHUNK_READY_ERROR); + } + } + + private void notifyConsumerAdvanced() { + prefetchLock.lock(); + try { + consumerAdvanced.signalAll(); + } finally { + prefetchLock.unlock(); + } + } + + private void notifyBatchAvailable() { + prefetchLock.lock(); + try { + batchAvailable.signalAll(); + } finally { + prefetchLock.unlock(); + } + } +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResult.java b/src/main/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResult.java new file mode 100644 index 0000000000..bf5a95d346 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResult.java @@ -0,0 +1,310 @@ +package com.databricks.jdbc.api.impl.thrift; + +import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_RESULT_ROW_LIMIT; +import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_STREAMING_BATCH_TIMEOUT_SECONDS; + +import com.databricks.jdbc.api.impl.ColumnarRowView; +import com.databricks.jdbc.api.impl.IExecutionResult; +import com.databricks.jdbc.api.impl.streaming.StreamingBatch; +import com.databricks.jdbc.api.impl.streaming.ThriftStreamingProvider; +import com.databricks.jdbc.api.internal.IDatabricksSession; +import com.databricks.jdbc.api.internal.IDatabricksStatementInternal; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.log.JdbcLogger; +import com.databricks.jdbc.log.JdbcLoggerFactory; +import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; + +/** + * High-throughput streaming implementation for Thrift columnar results. + * + *

Uses {@link ThriftStreamingProvider} for proactive batch prefetching, achieving throughput + * comparable to eager loading while maintaining the memory benefits of lazy loading. + * + *

Key features: + * + *

    + *
  • Background prefetch thread fetches batches ahead of consumption + *
  • Sliding window limits memory usage to a configurable number of batches + *
  • Non-blocking iteration when prefetch keeps up with consumption + *
  • Maintains correct row ordering through sequential fetch + *
  • Type-safe: Uses generic {@code ThriftStreamingProvider} + *
+ * + *

This implementation replaces {@code LazyThriftResult} for improved throughput while + * maintaining the same {@link IExecutionResult} interface. + */ +public class StreamingColumnarResult implements IExecutionResult { + + private static final JdbcLogger LOGGER = + JdbcLoggerFactory.getLogger(StreamingColumnarResult.class); + + // Streaming infrastructure - type-safe generic provider + private final ThriftStreamingProvider provider; + + // Current position within the current batch + private StreamingBatch currentBatch; + private int currentBatchRowIndex; + private long globalRowIndex; + + // Limits + private final int maxRows; + + // State + private boolean hasReachedEnd; + private volatile boolean isClosed; + + /** + * Creates a new StreamingColumnarResult. + * + *

Configuration values (maxBatchesInMemory, timeout) are read from the session's connection + * context. + * + * @param initialResponse The initial Thrift response containing the first batch + * @param statement The statement that generated this result + * @param session The session for fetching additional batches + * @throws DatabricksSQLException if initialization fails + */ + public StreamingColumnarResult( + TFetchResultsResp initialResponse, + IDatabricksStatementInternal statement, + IDatabricksSession session) + throws DatabricksSQLException { + + this.maxRows = statement != null ? statement.getMaxRows() : DEFAULT_RESULT_ROW_LIMIT; + + this.globalRowIndex = -1; + this.currentBatchRowIndex = -1; + this.hasReachedEnd = false; + this.isClosed = false; + + // Create batch fetcher and type-safe generic provider + ThriftBatchFetcher fetcher = new ThriftBatchFetcherImpl(session, statement); + this.provider = + ThriftStreamingProvider.forColumnar( + fetcher, + initialResponse, + session.getConnectionContext().getThriftMaxBatchesInMemory(), + DEFAULT_STREAMING_BATCH_TIMEOUT_SECONDS); + + // Move to first batch (check nextBatch() return value to handle empty initial batches) + if (provider.nextBatch()) { + currentBatch = provider.getCurrentBatch(); + } + + LOGGER.debug( + "StreamingColumnarResult initialized - firstBatchRows={}, maxRows={}, maxBatchesInMemory={}", + currentBatch != null ? currentBatch.getRowCount() : 0, + maxRows, + session.getConnectionContext().getThriftMaxBatchesInMemory()); + } + + /** + * Gets the value at the specified column index for the current row. + * + * @param columnIndex the zero-based column index + * @return the value at the specified column + * @throws DatabricksSQLException if access fails + */ + @Override + public Object getObject(int columnIndex) throws DatabricksSQLException { + if (isClosed) { + LOGGER.error("Attempted to access closed result"); + throw new DatabricksSQLException( + "Result is closed", DatabricksDriverErrorCode.STATEMENT_CLOSED); + } + if (globalRowIndex == -1) { + LOGGER.error("Attempted to access data before first row"); + throw new DatabricksSQLException( + "Cursor is before first row", DatabricksDriverErrorCode.INVALID_STATE); + } + if (currentBatch == null || currentBatchRowIndex < 0) { + LOGGER.error( + "Invalid cursor position: batch={}, rowIndex={}", currentBatch, currentBatchRowIndex); + throw new DatabricksSQLException( + "Invalid cursor position", DatabricksDriverErrorCode.INVALID_STATE); + } + + // Type-safe: getData() returns ColumnarRowView directly, no casting! + ColumnarRowView view = currentBatch.getData(); + if (view == null) { + LOGGER.error("Batch data not available at row {}", globalRowIndex); + throw new DatabricksSQLException( + "Batch data not available", DatabricksDriverErrorCode.INVALID_STATE); + } + if (columnIndex < 0 || columnIndex >= view.getColumnCount()) { + LOGGER.error("Column index {} out of bounds (0-{})", columnIndex, view.getColumnCount() - 1); + throw new DatabricksSQLException( + "Column index out of bounds: " + columnIndex, DatabricksDriverErrorCode.INVALID_STATE); + } + + return view.getValue(currentBatchRowIndex, columnIndex); + } + + /** + * Gets the current row index (0-based). Returns -1 if before the first row. + * + * @return the current row index + */ + @Override + public long getCurrentRow() { + return globalRowIndex; + } + + /** + * Moves the cursor to the next row. Fetches additional batches from server if needed. + * + * @return true if there is a next row, false if at the end + * @throws DatabricksSQLException if an error occurs + */ + @Override + public boolean next() throws DatabricksSQLException { + if (isClosed || hasReachedEnd) { + return false; + } + + if (!hasNext()) { + return false; + } + + // Check maxRows limit + boolean hasRowLimit = maxRows > 0; + if (hasRowLimit && globalRowIndex + 1 >= maxRows) { + hasReachedEnd = true; + return false; + } + + // Move to next row + currentBatchRowIndex++; + globalRowIndex++; + + // Check if we need to move to next batch + ColumnarRowView batchData = currentBatch != null ? currentBatch.getData() : null; + long batchRowCount = batchData != null ? batchData.getRowCount() : 0; + if (currentBatch != null && currentBatchRowIndex >= batchRowCount) { + + // Try to move to next batch + if (provider.hasNextBatch()) { + provider.nextBatch(); + currentBatch = provider.getCurrentBatch(); + currentBatchRowIndex = 0; + + if (currentBatch == null) { + LOGGER.warn("Got null batch after nextBatch()"); + hasReachedEnd = true; + globalRowIndex--; + currentBatchRowIndex--; + return false; + } + + // Log batch transition + LOGGER.debug( + "Moved to batch {} - globalRow={}, batchesInMemory={}", + currentBatch.getBatchIndex(), + globalRowIndex, + provider.getBatchesInMemory()); + } else { + // No more batches + hasReachedEnd = true; + globalRowIndex--; + currentBatchRowIndex--; + return false; + } + } + + return true; + } + + /** + * Checks if there are more rows available without advancing the cursor. + * + * @return true if there are more rows, false otherwise + */ + @Override + public boolean hasNext() { + if (isClosed || hasReachedEnd) { + return false; + } + + // Check maxRows limit + boolean hasRowLimit = maxRows > 0; + if (hasRowLimit && globalRowIndex + 1 >= maxRows) { + return false; + } + + // Check current batch - type-safe getData() returns ColumnarRowView + if (currentBatch != null) { + ColumnarRowView view = currentBatch.getData(); + if (view != null && currentBatchRowIndex + 1 < view.getRowCount()) { + return true; + } + } + + // Check if more batches available + return provider.hasNextBatch(); + } + + /** Closes this result and releases associated resources. */ + @Override + public void close() { + if (isClosed) { + return; + } + + long totalRows = provider.getTotalRowsFetched(); + isClosed = true; + currentBatch = null; + + provider.close(); + + LOGGER.debug("Closed - totalRowsFetched={}, rowsConsumed={}", totalRows, globalRowIndex + 1); + } + + /** + * Gets the number of rows in the current batch. + * + * @return the number of rows in the current batch + */ + @Override + public long getRowCount() { + return currentBatch != null ? currentBatch.getRowCount() : 0; + } + + /** + * Gets the chunk count. Always returns 0 for thrift columnar results (chunks are an Arrow + * concept). + * + * @return 0 + */ + @Override + public long getChunkCount() { + return 0; + } + + /** + * Gets the total number of rows fetched from the server so far. + * + * @return the total rows fetched + */ + public long getTotalRowsFetched() { + return provider.getTotalRowsFetched(); + } + + /** + * Checks if all data has been fetched from the server. + * + * @return true if end of stream reached + */ + public boolean isCompletelyFetched() { + return hasReachedEnd || provider.isEndOfStreamReached(); + } + + /** + * Gets the number of batches currently in memory. + * + * @return the batch count in memory + */ + public int getBatchesInMemory() { + return provider.getBatchesInMemory(); + } +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/thrift/ThriftBatchFetcher.java b/src/main/java/com/databricks/jdbc/api/impl/thrift/ThriftBatchFetcher.java new file mode 100644 index 0000000000..28ee2eea59 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/thrift/ThriftBatchFetcher.java @@ -0,0 +1,44 @@ +package com.databricks.jdbc.api.impl.thrift; + +import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; +import java.sql.SQLException; + +/** + * Interface for fetching Thrift columnar result batches from the server. + * + *

This abstraction enables testing with mock implementations and potential future enhancements + * like caching or instrumentation. + * + *

Implementations must be thread-safe as the fetcher may be called from the prefetch thread + * while the main thread is consuming data. + */ +public interface ThriftBatchFetcher { + + /** + * Fetches the next batch of results from the server. + * + *

This is a blocking network call that uses the Thrift FETCH_NEXT orientation. Each call + * returns the next sequential batch from the server's cursor. + * + *

Thread Safety: This method should only be called from a single thread at a time, as + * the server maintains a cursor that advances with each call. + * + * @return The fetch response containing the batch data and hasMoreRows flag + * @throws SQLException if the fetch fails due to network or server errors + */ + TFetchResultsResp fetchNextBatch() throws SQLException; + + /** + * Closes the fetcher and releases any associated resources. + * + *

After calling close(), any subsequent calls to fetchNextBatch() should throw an exception. + */ + void close(); + + /** + * Checks if the fetcher has been closed. + * + * @return true if close() has been called + */ + boolean isClosed(); +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/thrift/ThriftBatchFetcherImpl.java b/src/main/java/com/databricks/jdbc/api/impl/thrift/ThriftBatchFetcherImpl.java new file mode 100644 index 0000000000..11ae9c1bc5 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/thrift/ThriftBatchFetcherImpl.java @@ -0,0 +1,78 @@ +package com.databricks.jdbc.api.impl.thrift; + +import com.databricks.jdbc.api.internal.IDatabricksSession; +import com.databricks.jdbc.api.internal.IDatabricksStatementInternal; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.log.JdbcLogger; +import com.databricks.jdbc.log.JdbcLoggerFactory; +import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.sql.SQLException; + +/** + * Default implementation of ThriftBatchFetcher that uses the session's Databricks client to fetch + * batches. + * + *

This implementation delegates to {@code session.getDatabricksClient().getMoreResults()} which + * uses Thrift's FETCH_NEXT orientation to retrieve sequential batches. + */ +public class ThriftBatchFetcherImpl implements ThriftBatchFetcher { + + private static final JdbcLogger LOGGER = + JdbcLoggerFactory.getLogger(ThriftBatchFetcherImpl.class); + + private final IDatabricksSession session; + private final IDatabricksStatementInternal statement; + private volatile boolean closed; + + /** + * Creates a new ThriftBatchFetcherImpl. + * + * @param session The session to use for fetching + * @param statement The statement that generated the result + */ + public ThriftBatchFetcherImpl( + IDatabricksSession session, IDatabricksStatementInternal statement) { + this.session = session; + this.statement = statement; + this.closed = false; + LOGGER.debug( + "Created ThriftBatchFetcherImpl for statement {}", + statement != null ? statement.getStatementId() : "null"); + } + + @Override + public TFetchResultsResp fetchNextBatch() throws SQLException { + if (closed) { + LOGGER.error("Attempted to fetch batch from closed ThriftBatchFetcher"); + throw new DatabricksSQLException( + "ThriftBatchFetcher is closed", DatabricksDriverErrorCode.STATEMENT_CLOSED); + } + + LOGGER.debug( + "Fetching next batch for statement {}", + statement != null ? statement.getStatementId() : "null"); + long startTime = System.currentTimeMillis(); + + TFetchResultsResp response = session.getDatabricksClient().getMoreResults(statement); + + long duration = System.currentTimeMillis() - startTime; + LOGGER.debug( + "Fetched batch in {}ms, hasMoreRows: {}", + duration, + response != null && response.hasMoreRows); + + return response; + } + + @Override + public void close() { + LOGGER.debug("Closing ThriftBatchFetcherImpl"); + this.closed = true; + } + + @Override + public boolean isClosed() { + return closed; + } +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/volume/VolumeOperationResult.java b/src/main/java/com/databricks/jdbc/api/impl/volume/VolumeOperationResult.java index 186f4192ff..f7e8357bc9 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/volume/VolumeOperationResult.java +++ b/src/main/java/com/databricks/jdbc/api/impl/volume/VolumeOperationResult.java @@ -22,6 +22,7 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Strings; import java.io.IOException; +import java.sql.SQLException; import java.util.HashMap; import java.util.Map; import org.apache.http.HttpEntity; @@ -48,7 +49,7 @@ public VolumeOperationResult( IDatabricksSession session, IExecutionResult resultHandler, IDatabricksStatementInternal statement) - throws DatabricksSQLException { + throws SQLException { this.rowCount = totalRows; this.columnCount = totalColumns; this.session = session; @@ -68,7 +69,7 @@ public VolumeOperationResult( IExecutionResult resultHandler, IDatabricksHttpClient httpClient, IDatabricksStatementInternal statement) - throws DatabricksSQLException { + throws SQLException { this.rowCount = manifest.getTotalRowCount(); this.columnCount = manifest.getSchema().getColumnCount(); this.session = session; @@ -210,7 +211,7 @@ public boolean next() throws DatabricksSQLException { } } - private void completeVolumeOperation() throws DatabricksSQLException { + private void completeVolumeOperation() throws SQLException { while (resultHandler.hasNext()) { validateMetadata(); resultHandler.next(); diff --git a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java index f3a8c09114..95125968ac 100644 --- a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java @@ -179,6 +179,8 @@ public interface IDatabricksConnectionContext { int getIdleHttpConnectionExpiry(); + List getNonRowcountQueryPrefixes(); + boolean supportManyParameters(); String getConnectionURL(); @@ -426,6 +428,29 @@ public interface IDatabricksConnectionContext { /** Returns whether streaming chunk provider is enabled for result fetching. */ boolean isStreamingChunkProviderEnabled(); + /** + * Returns whether streaming mode is enabled for inline results (Thrift columnar and inline + * Arrow). + */ + boolean isInlineStreamingEnabled(); + + /** + * Returns whether CloudFetch (URL-based result download) is enabled. + * + *

When enabled (default), the server may return URL_BASED_SET results that are downloaded from + * cloud storage. When disabled, the server returns ARROW_BASED_SET with inline Arrow data. + * + * @return true if CloudFetch is enabled, false otherwise + */ + boolean isCloudFetchEnabled(); + + /** + * Returns the maximum number of batches to keep in memory for Thrift streaming. + * + * @return the max batches in memory (default: 3) + */ + int getThriftMaxBatchesInMemory(); + /** * Returns the number of chunk links to prefetch ahead of consumption. * diff --git a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionInternal.java b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionInternal.java index 07a9265e52..415e181dee 100644 --- a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionInternal.java +++ b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionInternal.java @@ -1,8 +1,8 @@ package com.databricks.jdbc.api.internal; import com.databricks.jdbc.api.IDatabricksStatement; -import com.databricks.jdbc.exception.DatabricksSQLException; import java.sql.Connection; +import java.sql.SQLException; /** Interface providing Databricks specific Connection APIs. */ public interface IDatabricksConnectionInternal extends Connection { @@ -21,7 +21,7 @@ public interface IDatabricksConnectionInternal extends Connection { Connection getConnection(); /** Opens the connection and initiates the underlying session */ - void open() throws DatabricksSQLException; + void open() throws SQLException; /** Returns the connection context associated with the connection. */ IDatabricksConnectionContext getConnectionContext(); diff --git a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksSession.java b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksSession.java index 547f361f4f..eaa87dc351 100644 --- a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksSession.java +++ b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksSession.java @@ -6,6 +6,7 @@ import com.databricks.jdbc.dbclient.IDatabricksClient; import com.databricks.jdbc.dbclient.IDatabricksMetadataClient; import com.databricks.jdbc.exception.DatabricksSQLException; +import java.sql.SQLException; import java.util.Map; import javax.annotation.Nullable; @@ -38,10 +39,10 @@ public interface IDatabricksSession { boolean isOpen(); /** Opens a new session. */ - void open() throws DatabricksSQLException; + void open() throws SQLException; /** Closes the session. */ - void close() throws DatabricksSQLException; + void close() throws SQLException; /** * Returns the client for connecting to Databricks server diff --git a/src/main/java/com/databricks/jdbc/auth/DatabricksTokenFederationProvider.java b/src/main/java/com/databricks/jdbc/auth/DatabricksTokenFederationProvider.java index ec898fd982..4b4381a86b 100644 --- a/src/main/java/com/databricks/jdbc/auth/DatabricksTokenFederationProvider.java +++ b/src/main/java/com/databricks/jdbc/auth/DatabricksTokenFederationProvider.java @@ -4,6 +4,7 @@ import com.databricks.jdbc.common.DatabricksJdbcConstants; import com.databricks.jdbc.common.util.DriverUtil; import com.databricks.jdbc.common.util.JsonUtil; +import com.databricks.jdbc.common.util.ValidationUtil; import com.databricks.jdbc.dbclient.IDatabricksHttpClient; import com.databricks.jdbc.dbclient.impl.http.DatabricksHttpClientFactory; import com.databricks.jdbc.exception.DatabricksDriverException; @@ -230,13 +231,29 @@ Token retrieveToken( StandardCharsets.UTF_8)); headers.forEach(postRequest::setHeader); HttpResponse response = hc.execute(postRequest); + + // Check for HTTP errors and build detailed error message + String httpError = ValidationUtil.checkHTTPErrorWithoutThrowingError(response); + if (!httpError.equals(DatabricksJdbcConstants.EMPTY_STRING)) { + String errorMessage = + String.format( + "Failed to retrieve the exchanged token from OIDC token endpoint. %s", httpError); + LOGGER.error(errorMessage); + throw new DatabricksDriverException(errorMessage, DatabricksDriverErrorCode.AUTH_ERROR); + } + OAuthResponse resp = JsonUtil.getMapper().readValue(response.getEntity().getContent(), OAuthResponse.class); return createToken(resp.getAccessToken(), resp.getTokenType()); + } catch (DatabricksDriverException e) { + // Already logged and has telemetry, just re-throw + throw e; } catch (Exception e) { - LOGGER.error(e, "Failed to retrieve the exchanged token"); - throw new DatabricksDriverException( - "Failed to retrieve the exchanged token", e, DatabricksDriverErrorCode.AUTH_ERROR); + String errorMessage = + String.format( + "Failed to parse token response from OIDC token endpoint. Error: %s", e.getMessage()); + LOGGER.error(e, errorMessage); + throw new DatabricksDriverException(errorMessage, e, DatabricksDriverErrorCode.AUTH_ERROR); } } diff --git a/src/main/java/com/databricks/jdbc/common/CommandName.java b/src/main/java/com/databricks/jdbc/common/CommandName.java index b950a2901f..3e4b3a85ea 100644 --- a/src/main/java/com/databricks/jdbc/common/CommandName.java +++ b/src/main/java/com/databricks/jdbc/common/CommandName.java @@ -29,5 +29,6 @@ public enum CommandName { GET_PSEUDO_COLUMNS, GET_IMPORTED_KEYS, GET_EXPORTED_KEYS, + GET_CLIENT_INFO_PROPERTIES, DEFAULT } diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java index de69f78a49..6eb8ef12ab 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java @@ -62,6 +62,7 @@ public final class DatabricksJdbcConstants { public static final String USER_HOME_PROPERTY = "user.home"; public static final String PORT = "port"; public static final int DEFAULT_PORT = 443; + public static final String REQUEST_ID_HEADER = "x-request-id"; public static final String THRIFT_ERROR_MESSAGE_HEADER = "X-Thriftserver-Error-Message"; public static final String ALLOWED_VOLUME_INGESTION_PATHS = "VolumeOperationAllowedLocalPaths"; public static final String ENABLE_VOLUME_OPERATIONS = "enableVolumeOperations"; @@ -108,6 +109,7 @@ public final class DatabricksJdbcConstants { public static final String QUERY_EXECUTION_TIMEOUT_SQLSTATE = "57KD0"; public static final int TEMPORARY_REDIRECT_STATUS_CODE = 307; public static final String REDACTED_TOKEN = "****"; + public static final String QUERY_TAGS = "query_tags"; public static final int MAX_DEFAULT_STRING_COLUMN_LENGTH = 32767; public static final int DEFUALT_STRING_COLUMN_LENGTH = 255; public static final int DEFAULT_MAX_CONCURRENT_PRESIGNED_REQUESTS = 50; diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java index ddbec923f8..1e0c5945d4 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java @@ -169,7 +169,7 @@ public enum DatabricksJdbcUrlParams { "EnableSQLValidationForIsValid", "Enable SQL query execution for connection validation in isValid() method", "0"), - IGNORE_TRANSACTIONS("IgnoreTransactions", "Ignore transaction-related method calls", "0"), + IGNORE_TRANSACTIONS("IgnoreTransactions", "Ignore transaction-related method calls", "1"), FETCH_AUTOCOMMIT_FROM_SERVER( "FetchAutoCommitFromServer", "Fetch auto-commit state from server using SQL query instead of using cached value", @@ -192,6 +192,14 @@ public enum DatabricksJdbcUrlParams { "EnableStreamingChunkProvider", "Enable streaming chunk provider for result fetching (experimental)", "0"), + ENABLE_INLINE_STREAMING( + "EnableInlineStreaming", + "Enable streaming mode with background prefetching for inline results (Thrift columnar and inline Arrow)", + "1"), + THRIFT_MAX_BATCHES_IN_MEMORY( + "ThriftMaxBatchesInMemory", + "Maximum number of batches to keep in memory for Thrift streaming (sliding window size)", + "3"), LINK_PREFETCH_WINDOW( "LinkPrefetchWindow", "Number of chunk links to prefetch ahead of consumption. " @@ -203,7 +211,11 @@ public enum DatabricksJdbcUrlParams { "Comma-separated list of HTTP status codes that should be retried irrespective of Retry-After header.", ""), API_RETRY_TIMEOUT( - "ApiRetryTimeout", "Timeout for retrying API retriable codes in seconds", "300"); + "ApiRetryTimeout", "Timeout for retrying API retriable codes in seconds", "300"), + NON_ROWCOUNT_QUERY_PREFIXES( + "NonRowcountQueryPrefixes", + "Comma-separated list of query prefixes (like INSERT,UPDATE,DELETE) that should return result sets instead of row counts", + ""); private final String paramName; private final String defaultValue; diff --git a/src/main/java/com/databricks/jdbc/common/EnvironmentVariables.java b/src/main/java/com/databricks/jdbc/common/EnvironmentVariables.java index 16c5f1e145..edece5003a 100644 --- a/src/main/java/com/databricks/jdbc/common/EnvironmentVariables.java +++ b/src/main/java/com/databricks/jdbc/common/EnvironmentVariables.java @@ -5,6 +5,7 @@ public final class EnvironmentVariables { public static final int DEFAULT_STATEMENT_TIMEOUT_SECONDS = 0; // Infinite timeout public static final int DEFAULT_RESULT_ROW_LIMIT = 0; // no limit + public static final int DEFAULT_STREAMING_BATCH_TIMEOUT_SECONDS = 300; // 5 minutes public static final int DEFAULT_ROW_LIMIT_PER_BLOCK = 2000000; // Setting a limit for resource and cost efficiency public static final int DEFAULT_BYTE_LIMIT = 404857600; diff --git a/src/main/java/com/databricks/jdbc/common/MetadataResultConstants.java b/src/main/java/com/databricks/jdbc/common/MetadataResultConstants.java index 9b1870b8fc..6adc4d9b3f 100644 --- a/src/main/java/com/databricks/jdbc/common/MetadataResultConstants.java +++ b/src/main/java/com/databricks/jdbc/common/MetadataResultConstants.java @@ -10,7 +10,9 @@ public class MetadataResultConstants { public static String NULL_STRING = "NULL"; public static final String PARSE_SYNTAX_ERROR_SQL_STATE = "42601"; - public static final String[] DEFAULT_TABLE_TYPES = {"TABLE", "VIEW", "SYSTEM TABLE"}; + public static final String[] DEFAULT_TABLE_TYPES = { + "TABLE", "VIEW", "SYSTEM TABLE", "METRIC_VIEW" + }; public static final ResultColumn CATALOG_COLUMN = new ResultColumn("TABLE_CAT", "catalogName", Types.VARCHAR); public static final ResultColumn CATALOG_FULL_COLUMN = @@ -95,6 +97,40 @@ public class MetadataResultConstants { new ResultColumn("IS_AUTOINCREMENT", "isAutoIncrement", Types.VARCHAR); public static final ResultColumn IS_GENERATED_COLUMN = new ResultColumn("IS_GENERATEDCOLUMN", "isGenerated", Types.VARCHAR); + + // TYPE_INFO columns + private static final ResultColumn LITERAL_PREFIX_COLUMN = + new ResultColumn("LITERAL_PREFIX", "literalPrefix", Types.VARCHAR); + private static final ResultColumn LITERAL_SUFFIX_COLUMN = + new ResultColumn("LITERAL_SUFFIX", "literalSuffix", Types.VARCHAR); + private static final ResultColumn CREATE_PARAMS_COLUMN = + new ResultColumn("CREATE_PARAMS", "createParams", Types.VARCHAR); + private static final ResultColumn CASE_SENSITIVE_COLUMN = + new ResultColumn("CASE_SENSITIVE", "caseSensitive", Types.BIT); + private static final ResultColumn SEARCHABLE_COLUMN = + new ResultColumn("SEARCHABLE", "searchable", Types.SMALLINT); + private static final ResultColumn UNSIGNED_ATTRIBUTE_COLUMN = + new ResultColumn("UNSIGNED_ATTRIBUTE", "unsignedAttribute", Types.BIT); + private static final ResultColumn FIXED_PREC_SCALE_COLUMN = + new ResultColumn("FIXED_PREC_SCALE", "fixedPrecScale", Types.BIT); + private static final ResultColumn AUTO_INCREMENT_COLUMN = + new ResultColumn("AUTO_INCREMENT", "autoIncrement", Types.BIT); + private static final ResultColumn LOCAL_TYPE_NAME_COLUMN = + new ResultColumn("LOCAL_TYPE_NAME", "localTypeName", Types.VARCHAR); + private static final ResultColumn MINIMUM_SCALE_COLUMN = + new ResultColumn("MINIMUM_SCALE", "minimumScale", Types.SMALLINT); + private static final ResultColumn MAXIMUM_SCALE_COLUMN = + new ResultColumn("MAXIMUM_SCALE", "maximumScale", Types.SMALLINT); + + // CLIENT_INFO_PROPERTIES columns + private static final ResultColumn NAME_COLUMN = new ResultColumn("NAME", "name", Types.VARCHAR); + private static final ResultColumn MAX_LEN_COLUMN = + new ResultColumn("MAX_LEN", "maxLen", Types.INTEGER); + private static final ResultColumn DEFAULT_VALUE_COLUMN = + new ResultColumn("DEFAULT_VALUE", "defaultValue", Types.VARCHAR); + private static final ResultColumn DESCRIPTION_COLUMN = + new ResultColumn("DESCRIPTION", "description", Types.VARCHAR); + private static final ResultColumn ATTR_NAME = new ResultColumn("ATTR_NAME", "attrName", Types.VARCHAR); private static final ResultColumn ATTR_TYPE_NAME = @@ -483,6 +519,30 @@ public class MetadataResultConstants { PK_NAME, DEFERRABILITY); + public static final List TYPE_INFO_COLUMNS = + List.of( + TYPE_NAME_COLUMN, + DATA_TYPE_COLUMN, + PRECISION_COLUMN, + LITERAL_PREFIX_COLUMN, + LITERAL_SUFFIX_COLUMN, + CREATE_PARAMS_COLUMN, + NULLABLE_COLUMN, + CASE_SENSITIVE_COLUMN, + SEARCHABLE_COLUMN, + UNSIGNED_ATTRIBUTE_COLUMN, + FIXED_PREC_SCALE_COLUMN, + AUTO_INCREMENT_COLUMN, + LOCAL_TYPE_NAME_COLUMN, + MINIMUM_SCALE_COLUMN, + MAXIMUM_SCALE_COLUMN, + SQL_DATA_TYPE_COLUMN, + SQL_DATETIME_SUB_COLUMN, + NUM_PREC_RADIX_COLUMN); + + public static final List CLIENT_INFO_PROPERTIES_COLUMNS = + List.of(NAME_COLUMN, MAX_LEN_COLUMN, DEFAULT_VALUE_COLUMN, DESCRIPTION_COLUMN); + public static final Map> NON_NULLABLE_COLUMNS_MAP = new HashMap<>() { { @@ -614,6 +674,7 @@ public class MetadataResultConstants { FKCOLUMN_NAME, KEY_SEQUENCE_COLUMN, DEFERRABILITY)); + put(CommandName.GET_CLIENT_INFO_PROPERTIES, List.of(NAME_COLUMN, MAX_LEN_COLUMN)); } }; } diff --git a/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java b/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java index 777581e5a9..56187a7476 100644 --- a/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java +++ b/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java @@ -4,6 +4,7 @@ import com.databricks.jdbc.common.DatabricksClientConfiguratorManager; import com.databricks.jdbc.common.util.DriverUtil; import com.databricks.jdbc.common.util.JsonUtil; +import com.databricks.jdbc.common.util.UserAgentManager; import com.databricks.jdbc.dbclient.IDatabricksHttpClient; import com.databricks.jdbc.dbclient.impl.http.DatabricksHttpClientFactory; import com.databricks.jdbc.exception.DatabricksHttpException; @@ -90,6 +91,12 @@ private void refreshAllFeatureFlags() { IDatabricksHttpClient httpClient = DatabricksHttpClientFactory.getInstance().getClient(connectionContext); HttpGet request = new HttpGet(featureFlagEndpoint); + + // Set custom User-Agent for connector service (includes custom user agent without client + // type) + String userAgent = UserAgentManager.buildUserAgentForConnectorService(connectionContext); + request.setHeader("User-Agent", userAgent); + DatabricksClientConfiguratorManager.getInstance() .getConfigurator(connectionContext) .getDatabricksConfig() diff --git a/src/main/java/com/databricks/jdbc/common/util/ArrowUtil.java b/src/main/java/com/databricks/jdbc/common/util/ArrowUtil.java new file mode 100644 index 0000000000..28f37ae7ea --- /dev/null +++ b/src/main/java/com/databricks/jdbc/common/util/ArrowUtil.java @@ -0,0 +1,255 @@ +package com.databricks.jdbc.common.util; + +import static com.databricks.jdbc.common.util.DatabricksThriftUtil.getColumnInfoFromTColumnDesc; +import static com.databricks.jdbc.common.util.DatabricksTypeUtil.getTPrimitiveTypeOrDefault; +import static com.databricks.jdbc.common.util.DatabricksTypeUtil.mapThriftToArrowType; +import static com.databricks.jdbc.common.util.DecompressionUtil.decompress; + +import com.databricks.jdbc.common.CompressionCodec; +import com.databricks.jdbc.exception.DatabricksParsingException; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.log.JdbcLogger; +import com.databricks.jdbc.log.JdbcLoggerFactory; +import com.databricks.jdbc.model.client.thrift.generated.TColumnDesc; +import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; +import com.databricks.jdbc.model.client.thrift.generated.TGetResultSetMetadataResp; +import com.databricks.jdbc.model.client.thrift.generated.TPrimitiveTypeEntry; +import com.databricks.jdbc.model.client.thrift.generated.TSparkArrowBatch; +import com.databricks.jdbc.model.client.thrift.generated.TTableSchema; +import com.databricks.jdbc.model.core.ColumnInfo; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.List; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.SchemaUtility; + +/** + * Utility class for Arrow operations. + * + *

Provides methods for: + * + *

    + *
  • Converting Thrift/Hive schemas to Arrow schemas and serialization + *
  • Creating Arrow IPC byte streams from Thrift responses + *
  • Processing Arrow batches with decompression + *
+ * + *

This consolidates Arrow handling logic used by both streaming and lazy inline Arrow result + * handlers. + */ +public final class ArrowUtil { + + private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(ArrowUtil.class); + + private ArrowUtil() { + // Utility class - prevent instantiation + } + + // ==================== Schema Operations ==================== + + /** + * Gets the serialized Arrow schema from Thrift metadata. + * + *

If the metadata contains a pre-serialized Arrow schema, it is returned directly. Otherwise, + * the Hive schema is converted to Arrow format and serialized. + * + * @param metadata The Thrift result set metadata + * @return The serialized Arrow schema bytes + * @throws DatabricksParsingException if schema conversion or serialization fails + */ + public static byte[] getSerializedSchema(TGetResultSetMetadataResp metadata) + throws DatabricksParsingException { + if (metadata == null) { + LOGGER.debug("Metadata is null, returning empty schema"); + return new byte[0]; + } + + // Use pre-serialized Arrow schema if available + if (metadata.getArrowSchema() != null) { + return metadata.getArrowSchema(); + } + + // Convert Hive schema to Arrow and serialize + Schema arrowSchema = hiveSchemaToArrowSchema(metadata.getSchema()); + try { + return SchemaUtility.serialize(arrowSchema); + } catch (IOException e) { + LOGGER.error(e, "Failed to serialize arrow schema"); + throw new DatabricksParsingException( + "Failed to serialize Arrow schema: " + e.getMessage(), + e, + DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR); + } + } + + /** + * Converts a Hive TTableSchema to an Arrow Schema. + * + * @param hiveSchema The Hive table schema from Thrift + * @return The equivalent Arrow schema + * @throws DatabricksParsingException if conversion fails + */ + public static Schema hiveSchemaToArrowSchema(TTableSchema hiveSchema) + throws DatabricksParsingException { + List fields = new ArrayList<>(); + if (hiveSchema == null) { + LOGGER.debug("Hive schema is null, returning empty Arrow schema"); + return new Schema(fields); + } + + try { + LOGGER.debug( + "Converting Hive schema to Arrow schema with {} columns", hiveSchema.getColumnsSize()); + for (TColumnDesc columnDesc : hiveSchema.getColumns()) { + fields.add(columnDescToArrowField(columnDesc)); + } + } catch (SQLException e) { + LOGGER.error("Failed to convert Hive schema to Arrow: {}", e.getMessage(), e); + throw new DatabricksParsingException( + "Failed to convert Hive schema to Arrow: " + e.getMessage(), + e, + DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR); + } + return new Schema(fields); + } + + /** + * Creates an Arrow Field from a Thrift column descriptor. + * + * @param columnDesc The Thrift column descriptor + * @return The equivalent Arrow field + * @throws SQLException if type mapping fails + */ + public static Field columnDescToArrowField(TColumnDesc columnDesc) throws SQLException { + TPrimitiveTypeEntry primitiveTypeEntry = getTPrimitiveTypeOrDefault(columnDesc.getTypeDesc()); + ArrowType arrowType = mapThriftToArrowType(primitiveTypeEntry.getType()); + FieldType fieldType = new FieldType(true, arrowType, null); + return new Field(columnDesc.getColumnName(), fieldType, null); + } + + // ==================== Arrow Stream Operations ==================== + + /** + * Creates a ByteArrayInputStream containing Arrow IPC data from the response. + * + *

This method combines the cached schema with decompressed Arrow batches to create a complete + * Arrow IPC stream that can be parsed by Arrow readers. + * + * @param cachedSchema The serialized Arrow schema bytes (should be cached from first response) + * @param response The Thrift fetch response containing Arrow batches + * @param callerClass The calling class for logging context + * @return ByteArrayInputStream containing the Arrow IPC data + * @throws DatabricksParsingException if processing fails + */ + public static ByteArrayInputStream createArrowByteStream( + byte[] cachedSchema, TFetchResultsResp response, Class callerClass) + throws DatabricksParsingException { + String context = callerClass.getSimpleName(); + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + CompressionCodec compressionCodec = + CompressionCodec.getCompressionMapping(response.getResultSetMetadata()); + + try { + // Write schema if available + if (cachedSchema != null && cachedSchema.length > 0) { + baos.write(cachedSchema); + } + + // Write arrow batches + List arrowBatches = + response.getResults() != null ? response.getResults().getArrowBatches() : null; + writeArrowBatchesToStream(compressionCodec, arrowBatches, baos, context); + + return new ByteArrayInputStream(baos.toByteArray()); + } catch (DatabricksSQLException | IOException e) { + LOGGER.error("{}: Failed to create Arrow byte stream: {}", context, e.getMessage(), e); + throw new DatabricksParsingException( + "Failed to create Arrow byte stream: " + e.getMessage(), + e, + DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR); + } + } + + /** + * Writes decompressed Arrow batches to the output stream. + * + * @param compressionCodec The compression codec used for the batches + * @param arrowBatchList The list of Arrow batches to write + * @param baos The output stream to write to + * @param context Context string for logging (typically caller class simple name) + * @throws DatabricksSQLException if decompression fails + * @throws IOException if writing fails + */ + static void writeArrowBatchesToStream( + CompressionCodec compressionCodec, + List arrowBatchList, + ByteArrayOutputStream baos, + String context) + throws DatabricksSQLException, IOException { + if (arrowBatchList == null) { + return; + } + for (TSparkArrowBatch arrowBatch : arrowBatchList) { + byte[] decompressedBytes = + decompress( + arrowBatch.getBatch(), + compressionCodec, + String.format( + "%s Arrow batch [%d rows] with decompression: [%s]", + context, arrowBatch.getRowCount(), compressionCodec)); + baos.write(decompressedBytes); + } + } + + /** + * Gets the total row count from all Arrow batches in the response. + * + * @param response The Thrift fetch response + * @return The total number of rows across all batches + */ + public static long getTotalRowsInResponse(TFetchResultsResp response) { + long totalRows = 0; + if (response.getResults() != null && response.getResults().getArrowBatches() != null) { + for (TSparkArrowBatch arrowBatch : response.getResults().getArrowBatches()) { + totalRows += arrowBatch.getRowCount(); + } + } + return totalRows; + } + + // ==================== Column Info Operations ==================== + + /** + * Extracts column information from Thrift result set metadata. + * + *

Converts each column descriptor in the Thrift schema to a {@link ColumnInfo} object. + * + * @param resultManifest The Thrift result set metadata containing schema information + * @return A list of ColumnInfo objects, empty list if schema is null + */ + public static List getColumnInfoList(TGetResultSetMetadataResp resultManifest) + throws DatabricksSQLException { + List columnInfos = new ArrayList<>(); + List arrowMetadataList = DatabricksThriftUtil.getArrowMetadata(resultManifest); + if (resultManifest.getSchema() == null) { + return columnInfos; + } + List columns = resultManifest.getSchema().getColumns(); + for (int columnIndex = 0; columnIndex < columns.size(); columnIndex++) { + TColumnDesc tColumnDesc = columns.get(columnIndex); + String columnArrowMetadata = + arrowMetadataList != null && columnIndex < arrowMetadataList.size() + ? arrowMetadataList.get(columnIndex) + : null; + columnInfos.add(getColumnInfoFromTColumnDesc(tColumnDesc, columnArrowMetadata)); + } + return columnInfos; + } +} diff --git a/src/main/java/com/databricks/jdbc/common/util/DatabricksThriftUtil.java b/src/main/java/com/databricks/jdbc/common/util/DatabricksThriftUtil.java index 04d4a506dd..53a7e383bb 100644 --- a/src/main/java/com/databricks/jdbc/common/util/DatabricksThriftUtil.java +++ b/src/main/java/com/databricks/jdbc/common/util/DatabricksThriftUtil.java @@ -1,5 +1,7 @@ package com.databricks.jdbc.common.util; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.ARROW_METADATA_KEY; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.QUERY_EXECUTION_TIMEOUT_SQLSTATE; import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_RESULT_ROW_LIMIT; import static com.databricks.jdbc.common.util.DatabricksTypeUtil.*; import static com.databricks.jdbc.model.client.thrift.generated.TTypeId.*; @@ -11,6 +13,7 @@ import com.databricks.jdbc.dbclient.impl.common.StatementId; import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.exception.DatabricksTimeoutException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.client.thrift.generated.*; @@ -20,9 +23,15 @@ import com.databricks.jdbc.model.core.StatementStatus; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; import com.databricks.sdk.service.sql.StatementState; +import java.io.IOException; import java.nio.ByteBuffer; +import java.sql.SQLException; import java.time.Instant; import java.util.*; +import java.util.stream.Collectors; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.SchemaUtility; public class DatabricksThriftUtil { @@ -80,13 +89,12 @@ public static ExternalLink createExternalLink(TSparkArrowResultLink chunkInfo, l .setRowCount(chunkInfo.getRowCount()); } - public static void verifySuccessStatus(TStatus status, String errorContext) - throws DatabricksHttpException { + public static void verifySuccessStatus(TStatus status, String errorContext) throws SQLException { verifySuccessStatus(status, errorContext, null); } public static void verifySuccessStatus(TStatus status, String errorContext, String statementId) - throws DatabricksHttpException { + throws SQLException { if (!SUCCESS_STATUS_LIST.contains(status.getStatusCode())) { String errorMessage = statementId != null @@ -94,7 +102,14 @@ public static void verifySuccessStatus(TStatus status, String errorContext, Stri "Error thrift response received [%s] for statementId [%s]", errorContext, statementId) : String.format("Error thrift response received [%s]", errorContext); - throw new DatabricksHttpException(errorMessage, status.getSqlState()); + + String sqlState = status.getSqlState(); + if (sqlState != null && QUERY_EXECUTION_TIMEOUT_SQLSTATE.equals(sqlState)) { + throw new DatabricksTimeoutException( + errorMessage, sqlState, null, DatabricksDriverErrorCode.OPERATION_TIMEOUT_ERROR); + } + + throw new DatabricksHttpException(errorMessage, sqlState); } } @@ -210,16 +225,30 @@ public static String getTypeTextFromTypeDesc(TTypeDesc typeDesc) { return primitiveTypeEntry.getType().name().replace("_TYPE", ""); } - public static ColumnInfo getColumnInfoFromTColumnDesc(TColumnDesc columnDesc) { + public static ColumnInfo getColumnInfoFromTColumnDesc( + TColumnDesc columnDesc, String arrowMetadata) { TPrimitiveTypeEntry primitiveTypeEntry = getTPrimitiveTypeOrDefault(columnDesc.getTypeDesc()); ColumnInfoTypeName columnInfoTypeName = T_TYPE_ID_COLUMN_INFO_TYPE_NAME_MAP.get(primitiveTypeEntry.getType()); + + String typeText = getTypeTextFromTypeDesc(columnDesc.getTypeDesc()); + + if (arrowMetadata != null && isComplexType(arrowMetadata)) { + typeText = arrowMetadata; + if (arrowMetadata.startsWith(GEOMETRY)) { + columnInfoTypeName = ColumnInfoTypeName.GEOMETRY; + } else if (arrowMetadata.startsWith(GEOGRAPHY)) { + columnInfoTypeName = ColumnInfoTypeName.GEOGRAPHY; + } + } + ColumnInfo columnInfo = new ColumnInfo() .setName(columnDesc.getColumnName()) .setPosition((long) columnDesc.getPosition()) .setTypeName(columnInfoTypeName) - .setTypeText(getTypeTextFromTypeDesc(columnDesc.getTypeDesc())); + .setTypeText(typeText); + if (primitiveTypeEntry.isSetTypeQualifiers()) { TTypeQualifiers typeQualifiers = primitiveTypeEntry.getTypeQualifiers(); String scaleQualifierKey = TCLIServiceConstants.SCALE, @@ -296,7 +325,7 @@ public static List> convertColumnarToRowBased( TFetchResultsResp resultsResp, IDatabricksStatementInternal parentStatement, IDatabricksSession session) - throws DatabricksSQLException { + throws SQLException { int statementMaxRows = parentStatement != null ? parentStatement.getMaxRows() : DEFAULT_RESULT_ROW_LIMIT; boolean hasRowLimit = statementMaxRows > 0; @@ -345,8 +374,7 @@ public static long getRowCount(TRowSet resultData) throws DatabricksSQLException } public static void checkDirectResultsForErrorStatus( - TSparkDirectResults directResults, String context, String statementId) - throws DatabricksHttpException { + TSparkDirectResults directResults, String context, String statementId) throws SQLException { if (directResults.isSetOperationStatus()) { LOGGER.debug( "direct result operation status being verified for success response for statementId {}", @@ -371,4 +399,31 @@ public static void checkDirectResultsForErrorStatus( verifySuccessStatus(directResults.getResultSet().getStatus(), context, statementId); } } + + /** + * Deserializes the Arrow schema from TGetResultSetMetadataResp. + * + * @param metadata the TGetResultSetMetadataResp containing the binary Arrow schema + * @return the deserialized Arrow Schema + */ + public static List getArrowMetadata(TGetResultSetMetadataResp metadata) + throws DatabricksSQLException { + if (metadata == null + || metadata.getArrowSchema() == null + || metadata.getArrowSchema().length == 0) { + return null; + } + byte[] arrowSchemaBytes = metadata.getArrowSchema(); + try { + Schema arrowSchema = SchemaUtility.deserialize(arrowSchemaBytes, null); + return arrowSchema.getFields().stream() + .map(Field::getMetadata) + .map(e -> e.get(ARROW_METADATA_KEY)) + .collect(Collectors.toList()); + } catch (IOException e) { + String errorMessage = "Failed to deserialize Arrow schema: " + e.getMessage(); + LOGGER.error(errorMessage, e); + throw new DatabricksSQLException(errorMessage, e, DatabricksDriverErrorCode.RESULT_SET_ERROR); + } + } } diff --git a/src/main/java/com/databricks/jdbc/common/util/DatabricksTypeUtil.java b/src/main/java/com/databricks/jdbc/common/util/DatabricksTypeUtil.java index a5612eb987..1a65a53ef6 100644 --- a/src/main/java/com/databricks/jdbc/common/util/DatabricksTypeUtil.java +++ b/src/main/java/com/databricks/jdbc/common/util/DatabricksTypeUtil.java @@ -555,4 +555,20 @@ public static String getDecimalTypeString(BigDecimal bd) { } return DECIMAL + "(" + precision + "," + scale + ")"; } + + /** + * Checks if the given type name represents a complex type (ARRAY, MAP, STRUCT, GEOMETRY, or + * GEOGRAPHY). + * + * @param typeName The type name to check + * @return true if the type name starts with ARRAY, MAP, STRUCT, GEOMETRY, or GEOGRAPHY, false + * otherwise + */ + public static boolean isComplexType(String typeName) { + return typeName.startsWith(ARRAY) + || typeName.startsWith(MAP) + || typeName.startsWith(STRUCT) + || typeName.startsWith(GEOMETRY) + || typeName.startsWith(GEOGRAPHY); + } } diff --git a/src/main/java/com/databricks/jdbc/common/util/DriverUtil.java b/src/main/java/com/databricks/jdbc/common/util/DriverUtil.java index 505c07b5db..581aa773a7 100644 --- a/src/main/java/com/databricks/jdbc/common/util/DriverUtil.java +++ b/src/main/java/com/databricks/jdbc/common/util/DriverUtil.java @@ -20,7 +20,7 @@ public class DriverUtil { private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(DriverUtil.class); - private static final String DRIVER_VERSION = "3.0.7"; + private static final String DRIVER_VERSION = "3.1.1"; private static final String DRIVER_NAME = "oss-jdbc"; private static final String JDBC_VERSION = "4.3"; diff --git a/src/main/java/com/databricks/jdbc/common/util/JsonUtil.java b/src/main/java/com/databricks/jdbc/common/util/JsonUtil.java index d6fc89058b..a532b8b43b 100644 --- a/src/main/java/com/databricks/jdbc/common/util/JsonUtil.java +++ b/src/main/java/com/databricks/jdbc/common/util/JsonUtil.java @@ -1,13 +1,21 @@ package com.databricks.jdbc.common.util; +import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.databind.ObjectMapper; public class JsonUtil { // Thread-safe singleton instance private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final ObjectMapper TELEMETRY_MAPPER = + new ObjectMapper().setSerializationInclusion(JsonInclude.Include.NON_NULL); - // Use the shared instance for your operations + // Use the shared instance for driver operations public static ObjectMapper getMapper() { return MAPPER; } + + // Use the shared instance for telemetry operations + public static ObjectMapper getTelemetryMapper() { + return TELEMETRY_MAPPER; + } } diff --git a/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java b/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java index a377e85fa9..14ac19f05f 100644 --- a/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java +++ b/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java @@ -17,33 +17,96 @@ public class UserAgentManager { public static final String USER_AGENT_THRIFT_CLIENT = "THttpClient"; private static final String VERSION_FILLER = "version"; + /** + * Parse custom user agent string into name and version components. + * + * @param customerUserAgent The custom user agent string (may be URL encoded) + * @return String array [name, version] or null if parsing fails + */ + private static String[] parseCustomerUserAgent(String customerUserAgent) { + try { + String decodedUA = URLDecoder.decode(customerUserAgent, StandardCharsets.UTF_8); + int i = decodedUA.indexOf('/'); + String customerName = (i < 0) ? decodedUA : decodedUA.substring(0, i); + String customerVersion = (i < 0) ? VERSION_FILLER : decodedUA.substring(i + 1); + return new String[] {customerName, customerVersion}; + } catch (Exception e) { + LOGGER.debug("Failed to parse customer userAgent entry {}, Error {}", customerUserAgent, e); + return null; + } + } + /** * Set the user agent for the Databricks JDBC driver. * * @param connectionContext The connection context. */ public static void setUserAgent(IDatabricksConnectionContext connectionContext) { - // Set the base product and client info + // Set the base product UserAgent.withProduct(DEFAULT_USER_AGENT, DriverUtil.getDriverVersion()); + + // Set client info (this may trigger getClientType which fetches feature flags) UserAgent.withOtherInfo(CLIENT_USER_AGENT_PREFIX, connectionContext.getClientUserAgent()); - if (connectionContext.getCustomerUserAgent() == null) { - return; + + // Set custom user agent (maintains proper order: base -> client type -> custom) + if (connectionContext.getCustomerUserAgent() != null) { + String[] parsed = parseCustomerUserAgent(connectionContext.getCustomerUserAgent()); + if (parsed != null) { + try { + UserAgent.withOtherInfo(parsed[0], UserAgent.sanitize(parsed[1])); + } catch (IllegalArgumentException e) { + LOGGER.debug( + "Failed to set user agent for customer userAgent entry {}, Error {}", + connectionContext.getCustomerUserAgent(), + e); + } + } } - try { - String decodedUA = - URLDecoder.decode( + } + + /** + * Build user agent string for connector service requests (without client type to avoid circular + * dependency). This is used specifically for feature flags requests since client type is not yet + * determined. + * + * @param connectionContext The connection context. + * @return User agent string with format: "DatabricksJDBCDriverOSS/version databricks-jdbc-http + * jvm/version os/name [CustomApp/version]" + */ + public static String buildUserAgentForConnectorService( + IDatabricksConnectionContext connectionContext) { + StringBuilder userAgent = new StringBuilder(); + + // Base product: DatabricksJDBCDriverOSS/version + userAgent.append(DEFAULT_USER_AGENT).append("/").append(DriverUtil.getDriverVersion()); + + // JDBC HTTP identifier + userAgent.append(" ").append(JDBC_HTTP_USER_AGENT); + + // JVM version + userAgent + .append(" jvm/") + .append(System.getProperty("java.version", "unknown").replace(" ", "_")); + + // OS name + userAgent.append(" os/").append(System.getProperty("os.name", "unknown").replace(" ", "_")); + + // Custom user agent (if provided) + if (connectionContext.getCustomerUserAgent() != null) { + String[] parsed = parseCustomerUserAgent(connectionContext.getCustomerUserAgent()); + if (parsed != null) { + try { + userAgent.append(" ").append(parsed[0]).append("/").append(UserAgent.sanitize(parsed[1])); + } catch (IllegalArgumentException e) { + LOGGER.debug( + "Failed to include customer userAgent entry {} in connector service UA, Error {}", connectionContext.getCustomerUserAgent(), - StandardCharsets.UTF_8); // This is for encoded userAgentString - int i = decodedUA.indexOf('/'); - String customerName = (i < 0) ? decodedUA : decodedUA.substring(0, i); - String customerVersion = (i < 0) ? VERSION_FILLER : decodedUA.substring(i + 1); - UserAgent.withOtherInfo(customerName, UserAgent.sanitize(customerVersion)); - } catch (Exception e) { - LOGGER.debug( - "Failed to set user agent for customer userAgent entry {}, Error {}", - connectionContext.getCustomerUserAgent(), - e); + e); + } + } } + + return userAgent.toString(); } /** Gets the user agent string for Databricks Driver HTTP Client. */ diff --git a/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java b/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java index 47560015b4..e9af6827bf 100644 --- a/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java +++ b/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java @@ -96,20 +96,24 @@ public static void throwErrorIfNull(String field, Object value) throw new DatabricksValidationException(errorMessage); } - public static void checkHTTPError(HttpResponse response) - throws DatabricksHttpException, IOException { + public static String checkHTTPErrorWithoutThrowingError(HttpResponse response) { int statusCode = response.getStatusLine().getStatusCode(); - String statusLine = response.getStatusLine().toString(); if (statusCode >= 200 && statusCode < 300) { - return; + return EMPTY_STRING; } - String errorReason = - String.format("HTTP request failed by code: %d, status line: %s.", statusCode, statusLine); + String statusLine = response.getStatusLine().toString(); + StringBuilder errorBuilder = new StringBuilder(); + errorBuilder.append( + String.format("HTTP request failed by code: %d, status line: %s.", statusCode, statusLine)); if (response.containsHeader(THRIFT_ERROR_MESSAGE_HEADER)) { - errorReason += + errorBuilder.append( String.format( " Thrift Header : %s", - response.getFirstHeader(THRIFT_ERROR_MESSAGE_HEADER).getValue()); + response.getFirstHeader(THRIFT_ERROR_MESSAGE_HEADER).getValue())); + } + if (response.containsHeader(REQUEST_ID_HEADER)) { + String requestId = response.getFirstHeader(REQUEST_ID_HEADER).getValue(); + errorBuilder.append(String.format(" Request ID: %s.", requestId)); } if (response.getEntity() != null) { try { @@ -117,13 +121,21 @@ public static void checkHTTPError(HttpResponse response) JsonUtil.getMapper().readTree(EntityUtils.toString(response.getEntity())); JsonNode errorNode = jsonNode.path("message"); if (errorNode.isTextual()) { - errorReason += String.format(" Error message: %s", errorNode.textValue()); + errorBuilder.append(String.format(" Error message: %s", errorNode.textValue())); } } catch (Exception e) { LOGGER.warn("Unable to parse JSON from response entity", e); } } + return errorBuilder.toString(); + } + public static void checkHTTPError(HttpResponse response) + throws DatabricksHttpException, IOException { + String errorReason = checkHTTPErrorWithoutThrowingError(response); + if (errorReason.equals(EMPTY_STRING)) { + return; + } LOGGER.error(errorReason); throw new DatabricksHttpException(errorReason, DEFAULT_HTTP_EXCEPTION_SQLSTATE); } diff --git a/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java b/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java index dad8e39c53..d0206de85d 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java @@ -34,7 +34,7 @@ ImmutableSessionInfo createSession( String catalog, String schema, Map sessionConf) - throws DatabricksSQLException; + throws SQLException; /** * Deletes a session for given session-Id @@ -42,7 +42,7 @@ ImmutableSessionInfo createSession( * @param sessionInfo for which the session should be deleted */ @DatabricksMetricsTimed - void deleteSession(ImmutableSessionInfo sessionInfo) throws DatabricksSQLException; + void deleteSession(ImmutableSessionInfo sessionInfo) throws SQLException; /** * Executes a statement in Databricks server @@ -132,8 +132,7 @@ DatabricksResultSet getStatementResult( * @param chunkStartRowOffset the row offset where the chunk starts in the result set */ ChunkLinkFetchResult getResultChunks( - StatementId statementId, long chunkIndex, long chunkStartRowOffset) - throws DatabricksSQLException; + StatementId statementId, long chunkIndex, long chunkStartRowOffset) throws SQLException; /** * Fetches the result data for given chunk index and statement-Id. @@ -155,7 +154,7 @@ ResultData getResultChunksData(StatementId statementId, long chunkIndex) void resetAccessToken(String newAccessToken); TFetchResultsResp getMoreResults(IDatabricksStatementInternal parentStatement) - throws DatabricksSQLException; + throws SQLException; /** Retrieves underlying DatabricksConfig */ DatabricksConfig getDatabricksConfig(); diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/MetadataResultSetBuilder.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/MetadataResultSetBuilder.java index e77ba07d01..7ef3c938b6 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/MetadataResultSetBuilder.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/MetadataResultSetBuilder.java @@ -8,13 +8,17 @@ import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty; import static com.databricks.jdbc.dbclient.impl.common.CommandConstants.*; import static com.databricks.jdbc.dbclient.impl.common.TypeValConstants.*; +import static java.sql.DatabaseMetaData.*; import com.databricks.jdbc.api.impl.DatabricksResultSet; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; +import com.databricks.jdbc.api.internal.IDatabricksSession; import com.databricks.jdbc.common.CommandName; import com.databricks.jdbc.common.Nullable; import com.databricks.jdbc.common.StatementType; import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.log.JdbcLogger; +import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.core.ColumnMetadata; import com.databricks.jdbc.model.core.ResultColumn; import com.databricks.jdbc.model.core.StatementStatus; @@ -29,16 +33,472 @@ import java.util.stream.Collectors; public class MetadataResultSetBuilder { + private static final JdbcLogger LOGGER = + JdbcLoggerFactory.getLogger(MetadataResultSetBuilder.class); private static final IDatabricksResultSetAdapter defaultAdapter = new DefaultDatabricksResultSetAdapter(); private static final IDatabricksResultSetAdapter importedKeysAdapter = new ImportedKeysDatabricksResultSetAdapter(); + + // Static data for TYPE_INFO metadata - JDBC type information constants + private static final Object[][] TYPE_INFO_DATA = + new Object[][] { + { + "TINYINT", + Types.TINYINT, + 3, + null, + null, + null, + typeNullable, + false, + typePredBasic, + false, + false, + null, + "TINYINT", + 0, + 0, + Types.TINYINT, + null, + 10 + }, + { + "BIGINT", + Types.BIGINT, + 19, + null, + null, + null, + typeNullable, + false, + typePredBasic, + false, + false, + null, + "BIGINT", + 0, + 0, + Types.BIGINT, + null, + 10 + }, + { + "BINARY", + Types.BINARY, + 32767, + "0x", + null, + "LENGTH", + typeNullable, + false, + typePredNone, + null, + false, + null, + "BINARY", + null, + null, + Types.BINARY, + null, + null + }, + { + "CHAR", + Types.CHAR, + 255, + "'", + "'", + "LENGTH", + typeNullable, + true, + typeSearchable, + null, + false, + null, + "CHAR", + null, + null, + Types.CHAR, + null, + null + }, + { + "DECIMAL", + Types.DECIMAL, + 38, + null, + null, + null, + typeNullable, + false, + typePredBasic, + false, + false, + null, + "DECIMAL", + 0, + 0, + Types.DECIMAL, + null, + 10 + }, + { + "INT", + Types.INTEGER, + 10, + null, + null, + null, + typeNullable, + false, + typePredBasic, + false, + false, + null, + "INT", + 0, + 0, + Types.INTEGER, + null, + 10 + }, + { + "SMALLINT", + Types.SMALLINT, + 5, + null, + null, + null, + typeNullable, + false, + typePredBasic, + false, + false, + null, + "SMALLINT", + 0, + 0, + Types.SMALLINT, + null, + 10 + }, + { + "FLOAT", + Types.FLOAT, + 7, + null, + null, + null, + typeNullable, + false, + typePredBasic, + false, + false, + null, + "FLOAT", + null, + null, + Types.FLOAT, + null, + 2 + }, + { + "DOUBLE", + Types.DOUBLE, + 15, + null, + null, + null, + typeNullable, + false, + typePredBasic, + false, + false, + null, + "DOUBLE", + null, + null, + Types.DOUBLE, + null, + 2 + }, + { + "ARRAY", + Types.VARCHAR, + 32767, + "'", + "'", + "Type", + typeNullable, + false, + typeSearchable, + null, + false, + null, + "ARRAY", + null, + null, + Types.VARCHAR, + null, + null + }, + { + "MAP", + Types.VARCHAR, + 32767, + "'", + "'", + "Key,Value", + typeNullable, + false, + typeSearchable, + null, + false, + null, + "MAP", + null, + null, + Types.VARCHAR, + null, + null + }, + { + "STRING", + Types.VARCHAR, + 510, + "'", + "'", + "max length", + typeNullable, + true, + typeSearchable, + null, + false, + null, + "STRING", + null, + null, + Types.VARCHAR, + null, + null + }, + { + "STRUCT", + Types.VARCHAR, + 32767, + "'", + "'", + "Column Type, ...", + typeNullable, + false, + typeSearchable, + null, + false, + null, + "STRUCT", + null, + null, + Types.VARCHAR, + null, + null + }, + { + "VARCHAR", + Types.VARCHAR, + 510, + "'", + "'", + "max length", + typeNullable, + true, + typeSearchable, + null, + false, + null, + "VARCHAR", + null, + null, + Types.VARCHAR, + null, + null + }, + { + "VARIANT", + Types.VARCHAR, + 32767, + "'", + "'", + "max length", + typeNullable, + false, + typeSearchable, + null, + false, + null, + "VARIANT", + null, + null, + Types.VARCHAR, + null, + null + }, + { + "BOOLEAN", + Types.BOOLEAN, + 1, + null, + null, + null, + typeNullable, + false, + typePredBasic, + null, + false, + null, + "BOOLEAN", + null, + null, + Types.BOOLEAN, + null, + null + }, + { + "DATE", + Types.DATE, + 10, + "'", + "'", + null, + typeNullable, + false, + typeSearchable, + null, + false, + null, + "DATE", + null, + null, + Types.DATE, + 1, + null + }, + { + "TIMESTAMP", + Types.TIMESTAMP, + 29, + "'", + "'", + null, + typeNullable, + false, + typeSearchable, + null, + false, + null, + "TIMESTAMP", + 0, + 0, + Types.TIMESTAMP, + 3, + null + }, + { + "TIMESTAMP_NTZ", + Types.TIMESTAMP, + 29, + "'", + "'", + null, + typeNullable, + false, + typeSearchable, + null, + false, + null, + "TIMESTAMP_NTZ", + 0, + 0, + Types.TIMESTAMP, + 3, + null + }, + { + "INTERVAL", + Types.VARCHAR, + 40, + "'", + "'", + "Qualifier", + typeNullable, + false, + typeSearchable, + null, + false, + null, + "INTERVAL", + 0, + 6, + Types.VARCHAR, + null, + null + } + }; + + // Static data for CLIENT_INFO_PROPERTIES metadata + private static final Object[][] CLIENT_INFO_PROPERTIES_DATA = + new Object[][] { + { + "APPLICATIONNAME", + 25, + null, + "The name of the application currently utilizing the connection." + }, + { + "CLIENTHOSTNAME", + 25, + null, + "The hostname of the computer the application using the connection is running on." + }, + { + "CLIENTUSER", + 25, + null, + "The name of the user that the application using the connection is performing work for." + } + }; + private final IDatabricksConnectionContext ctx; public MetadataResultSetBuilder(IDatabricksConnectionContext ctx) { this.ctx = ctx; } + public boolean shouldAllowCatalogAccess( + String catalog, String currentCatalog, IDatabricksSession session) throws SQLException { + if (ctx == null || ctx.getEnableMultipleCatalogSupport()) { + return true; + } + + if (catalog == null) { + return true; + } + + if (currentCatalog == null) { + currentCatalog = session.getCurrentCatalog(); + } + + if (currentCatalog != null && currentCatalog.equals(catalog)) { + return true; + } + + LOGGER.debug( + "Catalog access denied for catalog '{}' when enableMultipleCatalogSupport=false. Current catalog is '{}'", + catalog, + currentCatalog); + return false; + } + public DatabricksResultSet getFunctionsResult(DatabricksResultSet resultSet, String catalog) throws SQLException { List> rows = getRowsForFunctions(resultSet, FUNCTION_COLUMNS, catalog); @@ -1051,4 +1511,41 @@ public DatabricksResultSet getFunctionsResult(String catalog, List> GET_FUNCTIONS_STATEMENT_ID, CommandName.LIST_FUNCTIONS); } + + /** + * Creates a new TYPE_INFO ResultSet instance. + * + *

This method creates a fresh ResultSet instance each time to ensure proper cursor state and + * avoid issues with reusing closed ResultSets. See issue #1178. + * + * @return a new DatabricksResultSet with TYPE_INFO data + */ + public DatabricksResultSet getTypeInfoResult() { + // Convert static data to List> + // InlineJsonResult will make the defensive copy, so we just create cheap views here + List> rows = + Arrays.stream(TYPE_INFO_DATA).map(Arrays::asList).collect(Collectors.toList()); + return getResultSetWithGivenRowsAndColumns( + TYPE_INFO_COLUMNS, rows, "typeinfo-metadata", CommandName.LIST_TYPE_INFO); + } + + /** + * Creates a new CLIENT_INFO_PROPERTIES ResultSet instance. + * + *

This method creates a fresh ResultSet instance each time to ensure proper cursor state and + * avoid issues with reusing closed ResultSets. See issue #1178. + * + * @return a new DatabricksResultSet with CLIENT_INFO_PROPERTIES data + */ + public DatabricksResultSet getClientInfoPropertiesResult() { + // Convert static data to List> + // InlineJsonResult will make the defensive copy, so we just create cheap views here + List> rows = + Arrays.stream(CLIENT_INFO_PROPERTIES_DATA).map(Arrays::asList).collect(Collectors.toList()); + return getResultSetWithGivenRowsAndColumns( + CLIENT_INFO_PROPERTIES_COLUMNS, + rows, + "client-info-properties-metadata", + CommandName.GET_CLIENT_INFO_PROPERTIES); + } } diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksEmptyMetadataClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksEmptyMetadataClient.java index 8b7dc5c2da..048142dff8 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksEmptyMetadataClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksEmptyMetadataClient.java @@ -1,7 +1,5 @@ package com.databricks.jdbc.dbclient.impl.sqlexec; -import static com.databricks.jdbc.dbclient.impl.sqlexec.ResultConstants.TYPE_INFO_RESULT; - import com.databricks.jdbc.api.impl.DatabricksResultSet; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.api.internal.IDatabricksSession; @@ -26,7 +24,7 @@ public DatabricksEmptyMetadataClient(IDatabricksConnectionContext ctx) { @Override public DatabricksResultSet listTypeInfo(IDatabricksSession session) throws SQLException { LOGGER.debug("public ResultSet getTypeInfo()"); - return TYPE_INFO_RESULT; + return metadataResultSetBuilder.getTypeInfoResult(); } @Override diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataSdkClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataSdkClient.java index 1054fa97dc..adcfe64216 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataSdkClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataSdkClient.java @@ -1,13 +1,10 @@ package com.databricks.jdbc.dbclient.impl.sqlexec; import static com.databricks.jdbc.common.MetadataResultConstants.*; -import static com.databricks.jdbc.dbclient.impl.common.CommandConstants.GET_TABLES_STATEMENT_ID; import static com.databricks.jdbc.dbclient.impl.common.CommandConstants.METADATA_STATEMENT_ID; -import static com.databricks.jdbc.dbclient.impl.sqlexec.ResultConstants.TYPE_INFO_RESULT; import com.databricks.jdbc.api.impl.DatabricksResultSet; import com.databricks.jdbc.api.internal.IDatabricksSession; -import com.databricks.jdbc.common.MetadataResultConstants; import com.databricks.jdbc.common.StatementType; import com.databricks.jdbc.common.util.JdbcThreadUtils; import com.databricks.jdbc.common.util.WildcardUtil; @@ -46,7 +43,7 @@ public DatabricksMetadataSdkClient(IDatabricksClient sdkClient) { @Override public DatabricksResultSet listTypeInfo(IDatabricksSession session) { LOGGER.debug("public ResultSet getTypeInfo()"); - return TYPE_INFO_RESULT; + return metadataResultSetBuilder.getTypeInfoResult(); } @Override @@ -54,8 +51,11 @@ public DatabricksResultSet listCatalogs(IDatabricksSession session) throws SQLEx // If multiple catalog support is disabled, return only the current catalog if (isMultipleCatalogSupportDisabled()) { String currentCatalog = session.getCurrentCatalog(); - if (currentCatalog == null) { - currentCatalog = ""; + if (currentCatalog == null || currentCatalog.isEmpty()) { + currentCatalog = "spark"; + LOGGER.debug( + "Current catalog is null or empty when multiple catalog support is disabled. Using default catalog: {}", + currentCatalog); } String SQL = String.format("SELECT '%s' AS catalog", currentCatalog); LOGGER.debug("SQL command to fetch catalogs: {}", SQL); @@ -71,7 +71,13 @@ public DatabricksResultSet listCatalogs(IDatabricksSession session) throws SQLEx @Override public DatabricksResultSet listSchemas( IDatabricksSession session, String catalog, String schemaNamePattern) throws SQLException { - catalog = autoFillCatalog(catalog, session); + // Only fetch currentCatalog if multiple catalog support is disabled + String currentCatalog = isMultipleCatalogSupportDisabled() ? session.getCurrentCatalog() : null; + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, currentCatalog, session)) { + return metadataResultSetBuilder.getSchemasResult(new ArrayList<>()); + } + + catalog = autoFillCatalog(catalog, currentCatalog); // Return empty result set if catalog is an empty string if (catalog != null && catalog.isEmpty()) { @@ -111,11 +117,19 @@ public DatabricksResultSet listTables( String tableNamePattern, String[] tableTypes) throws SQLException { - catalog = autoFillCatalog(catalog, session); String[] validatedTableTypes = Optional.ofNullable(tableTypes) .filter(types -> types.length > 0) .orElse(DEFAULT_TABLE_TYPES); + + // Only fetch currentCatalog if multiple catalog support is disabled + String currentCatalog = isMultipleCatalogSupportDisabled() ? session.getCurrentCatalog() : null; + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, currentCatalog, session)) { + return metadataResultSetBuilder.getTablesResult( + catalog, validatedTableTypes, new ArrayList<>()); + } + + catalog = autoFillCatalog(catalog, currentCatalog); CommandBuilder commandBuilder = new CommandBuilder(catalog, session) .setSchemaPattern(schemaNamePattern) @@ -132,11 +146,8 @@ public DatabricksResultSet listTables( // Gracefully handles the case where an older DBSQL version doesn't support all catalogs in // the SHOW TABLES command. LOGGER.debug("SQL command failed with syntax error. Returning empty result set."); - return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns( - MetadataResultConstants.TABLE_COLUMNS, - new ArrayList<>(), - GET_TABLES_STATEMENT_ID, - com.databricks.jdbc.common.CommandName.LIST_TABLES); + return metadataResultSetBuilder.getTablesResult( + catalog, validatedTableTypes, new ArrayList<>()); } else { throw e; } @@ -157,7 +168,13 @@ public DatabricksResultSet listColumns( String tableNamePattern, String columnNamePattern) throws SQLException { - catalog = autoFillCatalog(catalog, session); + // Only fetch currentCatalog if multiple catalog support is disabled + String currentCatalog = isMultipleCatalogSupportDisabled() ? session.getCurrentCatalog() : null; + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, currentCatalog, session)) { + return metadataResultSetBuilder.getColumnsResult(new ArrayList<>()); + } + + catalog = autoFillCatalog(catalog, currentCatalog); // Fetch columns from all catalogs if catalog is null if (catalog == null) { @@ -183,7 +200,13 @@ public DatabricksResultSet listFunctions( String schemaNamePattern, String functionNamePattern) throws SQLException { - catalog = autoFillCatalog(catalog, session); + // Only fetch currentCatalog if multiple catalog support is disabled + String currentCatalog = isMultipleCatalogSupportDisabled() ? session.getCurrentCatalog() : null; + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, currentCatalog, session)) { + return metadataResultSetBuilder.getFunctionsResult(catalog, new ArrayList<>()); + } + + catalog = autoFillCatalog(catalog, currentCatalog); // Fetch current catalog if catalog is null if (catalog == null) { @@ -212,7 +235,13 @@ public DatabricksResultSet listFunctions( @Override public DatabricksResultSet listPrimaryKeys( IDatabricksSession session, String catalog, String schema, String table) throws SQLException { - catalog = autoFillCatalog(catalog, session); + // Only fetch currentCatalog if multiple catalog support is disabled + String currentCatalog = isMultipleCatalogSupportDisabled() ? session.getCurrentCatalog() : null; + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, currentCatalog, session)) { + return metadataResultSetBuilder.getPrimaryKeysResult(new ArrayList<>()); + } + + catalog = autoFillCatalog(catalog, currentCatalog); // Return empty result set if catalog, schema, or table is null if (catalog == null || schema == null || table == null) { @@ -222,7 +251,7 @@ public DatabricksResultSet listPrimaryKeys( schema, table); return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns( - MetadataResultConstants.PRIMARY_KEYS_COLUMNS, + PRIMARY_KEYS_COLUMNS, new ArrayList<>(), METADATA_STATEMENT_ID, com.databricks.jdbc.common.CommandName.LIST_PRIMARY_KEYS); @@ -239,7 +268,14 @@ public DatabricksResultSet listPrimaryKeys( public DatabricksResultSet listImportedKeys( IDatabricksSession session, String catalog, String schema, String table) throws SQLException { LOGGER.debug("public ResultSet listImportedKeys() using SDK"); - catalog = autoFillCatalog(catalog, session); + + // Only fetch currentCatalog if multiple catalog support is disabled + String currentCatalog = isMultipleCatalogSupportDisabled() ? session.getCurrentCatalog() : null; + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, currentCatalog, session)) { + return metadataResultSetBuilder.getImportedKeys(new ArrayList<>()); + } + + catalog = autoFillCatalog(catalog, currentCatalog); // Return empty result set if catalog, schema, or table is null if (catalog == null || schema == null || table == null) { @@ -249,7 +285,7 @@ public DatabricksResultSet listImportedKeys( schema, table); return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns( - MetadataResultConstants.IMPORTED_KEYS_COLUMNS, + IMPORTED_KEYS_COLUMNS, new ArrayList<>(), METADATA_STATEMENT_ID, com.databricks.jdbc.common.CommandName.GET_IMPORTED_KEYS); @@ -265,11 +301,7 @@ public DatabricksResultSet listImportedKeys( // This is a workaround for the issue where the SQL command fails with "syntax error at or // near "foreign"" LOGGER.debug("SQL command failed with syntax error. Returning empty result set."); - return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns( - MetadataResultConstants.IMPORTED_KEYS_COLUMNS, - new ArrayList<>(), - METADATA_STATEMENT_ID, - com.databricks.jdbc.common.CommandName.GET_IMPORTED_KEYS); + return metadataResultSetBuilder.getImportedKeys(new ArrayList<>()); } else { throw e; } @@ -280,13 +312,17 @@ public DatabricksResultSet listImportedKeys( public DatabricksResultSet listExportedKeys( IDatabricksSession session, String catalog, String schema, String table) throws SQLException { LOGGER.debug("public ResultSet listExportedKeys() using SDK"); - catalog = autoFillCatalog(catalog, session); + + // Only fetch currentCatalog if multiple catalog support is disabled + String currentCatalog = isMultipleCatalogSupportDisabled() ? session.getCurrentCatalog() : null; + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, currentCatalog, session)) { + return metadataResultSetBuilder.getExportedKeys(new ArrayList<>()); + } + + catalog = autoFillCatalog(catalog, currentCatalog); + // Exported keys not tracked in DBSQL. Returning empty result set - return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns( - MetadataResultConstants.EXPORTED_KEYS_COLUMNS, - new ArrayList<>(), - METADATA_STATEMENT_ID, - com.databricks.jdbc.common.CommandName.GET_EXPORTED_KEYS); + return metadataResultSetBuilder.getExportedKeys(new ArrayList<>()); } @Override @@ -300,6 +336,15 @@ public DatabricksResultSet listCrossReferences( String foreignTable) throws SQLException { LOGGER.debug("public ResultSet listCrossReferences() using SDK"); + + // Only fetch currentCatalog if multiple catalog support is disabled + String currentCatalog = isMultipleCatalogSupportDisabled() ? session.getCurrentCatalog() : null; + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(parentCatalog, currentCatalog, session) + || !metadataResultSetBuilder.shouldAllowCatalogAccess( + foreignCatalog, currentCatalog, session)) { + return metadataResultSetBuilder.getCrossRefsResult(new ArrayList<>()); + } + CommandBuilder commandBuilder = new CommandBuilder(foreignCatalog, session).setSchema(foreignSchema).setTable(foreignTable); String SQL = commandBuilder.getSQLString(CommandName.LIST_FOREIGN_KEYS); @@ -312,11 +357,7 @@ public DatabricksResultSet listCrossReferences( // near "foreign"" // This is a known issue in Databricks for older DBSQL versions LOGGER.debug("SQL command failed with syntax error. Returning empty result set."); - return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns( - MetadataResultConstants.CROSS_REFERENCE_COLUMNS, - new ArrayList<>(), - METADATA_STATEMENT_ID, - com.databricks.jdbc.common.CommandName.GET_CROSS_REFERENCE); + return metadataResultSetBuilder.getCrossRefsResult(new ArrayList<>()); } else { LOGGER.error( e, @@ -333,10 +374,20 @@ private boolean isMultipleCatalogSupportDisabled() { && !sdkClient.getConnectionContext().getEnableMultipleCatalogSupport(); } - private String autoFillCatalog(String catalog, IDatabricksSession session) throws SQLException { - if (isMultipleCatalogSupportDisabled()) { - String currentCatalog = session.getCurrentCatalog(); - return (currentCatalog != null && !currentCatalog.isEmpty()) ? currentCatalog : ""; + /** + * Auto-fills the catalog parameter if multiple catalog support is disabled and catalog is null. + * + * @param catalog the catalog parameter to auto-fill + * @param currentCatalog the current catalog from the session + * @return the auto-filled catalog or the original catalog if no auto-fill is needed + */ + private String autoFillCatalog(String catalog, String currentCatalog) { + if (isMultipleCatalogSupportDisabled() && catalog == null) { + String result = + (currentCatalog != null && !currentCatalog.isEmpty()) ? currentCatalog : "spark"; + LOGGER.debug( + "Auto-filling null catalog with '{}' when multiple catalog support is disabled", result); + return result; } return catalog; } @@ -446,7 +497,7 @@ private DatabricksResultSet fetchColumnsAcrossCatalogs( // Convert combined data into a result set return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns( - MetadataResultConstants.COLUMN_COLUMNS, + COLUMN_COLUMNS, columnRows, METADATA_STATEMENT_ID, com.databricks.jdbc.common.CommandName.LIST_COLUMNS); 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 0f3715d82e..dd39117402 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 @@ -366,7 +366,7 @@ public DatabricksResultSet getStatementResult( StatementId typedStatementId, IDatabricksSession session, IDatabricksStatementInternal parentStatement) - throws DatabricksSQLException { + throws SQLException { DatabricksThreadContextHolder.setStatementId(typedStatementId); DatabricksThreadContextHolder.setSessionId(session.getSessionId()); String statementId = typedStatementId.toSQLExecStatementId(); diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/ResultConstants.java b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/ResultConstants.java deleted file mode 100644 index 850a85af44..0000000000 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/ResultConstants.java +++ /dev/null @@ -1,513 +0,0 @@ -package com.databricks.jdbc.dbclient.impl.sqlexec; - -import static java.sql.DatabaseMetaData.*; - -import com.databricks.jdbc.api.impl.DatabricksResultSet; -import com.databricks.jdbc.common.StatementType; -import com.databricks.jdbc.dbclient.impl.common.StatementId; -import com.databricks.jdbc.model.core.StatementStatus; -import com.databricks.sdk.service.sql.StatementState; -import java.sql.Types; -import java.util.Arrays; - -public class ResultConstants { - public static final DatabricksResultSet TYPE_INFO_RESULT = - new DatabricksResultSet( - new StatementStatus().setState(StatementState.SUCCEEDED), - new StatementId("typeinfo-metadata"), - Arrays.asList( - "TYPE_NAME", - "DATA_TYPE", - "PRECISION", - "LITERAL_PREFIX", - "LITERAL_SUFFIX", - "CREATE_PARAMS", - "NULLABLE", - "CASE_SENSITIVE", - "SEARCHABLE", - "UNSIGNED_ATTRIBUTE", - "FIXED_PREC_SCALE", - "AUTO_INCREMENT", - "LOCAL_TYPE_NAME", - "MINIMUM_SCALE", - "MAXIMUM_SCALE", - "SQL_DATA_TYPE", - "SQL_DATETIME_SUB", - "NUM_PREC_RADIX"), - Arrays.asList( - "VARCHAR", - "INTEGER", - "INTEGER", - "VARCHAR", - "VARCHAR", - "VARCHAR", - "SMALLINT", - "BIT", - "SMALLINT", - "BIT", - "BIT", - "BIT", - "VARCHAR", - "SMALLINT", - "SMALLINT", - "INTEGER", - "INTEGER", - "INTEGER"), - new int[] { - Types.VARCHAR, - Types.INTEGER, - Types.INTEGER, - Types.VARCHAR, - Types.VARCHAR, - Types.VARCHAR, - Types.SMALLINT, - Types.BIT, - Types.SMALLINT, - Types.BIT, - Types.BIT, - Types.BIT, - Types.VARCHAR, - Types.SMALLINT, - Types.SMALLINT, - Types.INTEGER, - Types.INTEGER, - Types.INTEGER - }, - new int[] {128, 10, 10, 128, 128, 128, 5, 1, 5, 1, 1, 1, 128, 5, 5, 10, 10, 10}, - new int[] {0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1}, - new Object[][] { - { - "TINYINT", - Types.TINYINT, - 3, - null, - null, - null, - typeNullable, - false, - typePredBasic, - false, - false, - null, - "TINYINT", - 0, - 0, - Types.TINYINT, - null, - 10 - }, - { - "BIGINT", - Types.BIGINT, - 19, - null, - null, - null, - typeNullable, - false, - typePredBasic, - false, - false, - null, - "BIGINT", - 0, - 0, - Types.BIGINT, - null, - 10 - }, - { - "BINARY", - Types.BINARY, - 32767, - "0x", - null, - "LENGTH", - typeNullable, - false, - typePredNone, - null, - false, - null, - "BINARY", - null, - null, - Types.BINARY, - null, - null - }, - { - "CHAR", - Types.CHAR, - 255, - "'", - "'", - "LENGTH", - typeNullable, - true, - typeSearchable, - null, - false, - null, - "CHAR", - null, - null, - Types.CHAR, - null, - null - }, - { - "DECIMAL", - Types.DECIMAL, - 38, - null, - null, - null, - typeNullable, - false, - typePredBasic, - false, - false, - null, - "DECIMAL", - 0, - 0, - Types.DECIMAL, - null, - 10 - }, - { - "INT", - Types.INTEGER, - 10, - null, - null, - null, - typeNullable, - false, - typePredBasic, - false, - false, - null, - "INT", - 0, - 0, - Types.INTEGER, - null, - 10 - }, - { - "SMALLINT", - Types.SMALLINT, - 5, - null, - null, - null, - typeNullable, - false, - typePredBasic, - false, - false, - null, - "SMALLINT", - 0, - 0, - Types.SMALLINT, - null, - 10 - }, - { - "FLOAT", - Types.FLOAT, - 7, - null, - null, - null, - typeNullable, - false, - typePredBasic, - false, - false, - null, - "FLOAT", - null, - null, - Types.FLOAT, - null, - 2 - }, - { - "DOUBLE", - Types.DOUBLE, - 15, - null, - null, - null, - typeNullable, - false, - typePredBasic, - false, - false, - null, - "DOUBLE", - null, - null, - Types.DOUBLE, - null, - 2 - }, - { - "ARRAY", - Types.VARCHAR, - 32767, - "'", - "'", - "Type", - typeNullable, - false, - typeSearchable, - null, - false, - null, - "ARRAY", - null, - null, - Types.VARCHAR, - null, - null - }, - { - "MAP", - Types.VARCHAR, - 32767, - "'", - "'", - "Key,Value", - typeNullable, - false, - typeSearchable, - null, - false, - null, - "MAP", - null, - null, - Types.VARCHAR, - null, - null - }, - { - "STRING", - Types.VARCHAR, - 510, - "'", - "'", - "max length", - typeNullable, - true, - typeSearchable, - null, - false, - null, - "STRING", - null, - null, - Types.VARCHAR, - null, - null - }, - { - "STRUCT", - Types.VARCHAR, - 32767, - "'", - "'", - "Column Type, ...", - typeNullable, - false, - typeSearchable, - null, - false, - null, - "STRUCT", - null, - null, - Types.VARCHAR, - null, - null - }, - { - "VARCHAR", - Types.VARCHAR, - 510, - "'", - "'", - "max length", - typeNullable, - true, - typeSearchable, - null, - false, - null, - "VARCHAR", - null, - null, - Types.VARCHAR, - null, - null - }, - { - "VARIANT", - Types.VARCHAR, - 32767, - "'", - "'", - "max length", - typeNullable, - false, - typeSearchable, - null, - false, - null, - "VARIANT", - null, - null, - Types.VARCHAR, - null, - null - }, - { - "BOOLEAN", - Types.BOOLEAN, - 1, - null, - null, - null, - typeNullable, - false, - typePredBasic, - null, - false, - null, - "BOOLEAN", - null, - null, - Types.BOOLEAN, - null, - null - }, - { - "DATE", - Types.DATE, - 10, - "'", - "'", - null, - typeNullable, - false, - typeSearchable, - null, - false, - null, - "DATE", - null, - null, - Types.DATE, - 1, - null - }, - { - "TIMESTAMP", - Types.TIMESTAMP, - 29, - "'", - "'", - null, - typeNullable, - false, - typeSearchable, - null, - false, - null, - "TIMESTAMP", - 0, - 0, - Types.TIMESTAMP, - 3, - null - }, - { - "TIMESTAMP_NTZ", - Types.TIMESTAMP, - 29, - "'", - "'", - null, - typeNullable, - false, - typeSearchable, - null, - false, - null, - "TIMESTAMP_NTZ", - 0, - 0, - Types.TIMESTAMP, - 3, - null - }, - { - "INTERVAL", - Types.VARCHAR, - 40, - "'", - "'", - "Qualifier", - typeNullable, - false, - typeSearchable, - null, - false, - null, - "INTERVAL", - 0, - 6, - Types.VARCHAR, - null, - null - } - }, - StatementType.METADATA); - - /** The static result set for the client info properties metadata. */ - public static final DatabricksResultSet CLIENT_INFO_PROPERTIES_RESULT = - new DatabricksResultSet( - new StatementStatus().setState(StatementState.SUCCEEDED), - new StatementId("client-info-properties-metadata"), - Arrays.asList("NAME", "MAX_LEN", "DEFAULT_VALUE", "DESCRIPTION"), - Arrays.asList("VARCHAR", "INTEGER", "VARCHAR", "VARCHAR"), - new int[] {Types.VARCHAR, Types.INTEGER, Types.VARCHAR, Types.VARCHAR}, - new int[] {128, 10, 128, 128}, - new int[] {0, 0, 1, 1}, - new Object[][] { - { - "APPLICATIONNAME", - 25, - null, - "The name of the application currently utilizing the connection." - }, - { - "CLIENTHOSTNAME", - 25, - null, - "The hostname of the computer the application using the connection is running on." - }, - { - "CLIENTUSER", - 25, - null, - "The name of the user that the application using the connection is performing work for." - } - }, - StatementType.METADATA); -} diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java index 6b5e65d597..14af79bcb9 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java @@ -22,7 +22,7 @@ import com.databricks.jdbc.model.client.thrift.generated.*; import com.databricks.jdbc.model.core.StatementStatus; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; -import com.databricks.jdbc.telemetry.latency.TelemetryCollector; +import com.databricks.jdbc.telemetry.TelemetryHelper; import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.service.sql.StatementState; import java.sql.SQLException; @@ -136,8 +136,7 @@ TBase getThriftResponse(TBase request) throws DatabricksSQLException { * @return TFetchResultsResp containing the results * @throws DatabricksHttpException if fetch fails */ - TFetchResultsResp getResultSetResp(TOperationHandle operationHandle) - throws DatabricksHttpException { + TFetchResultsResp getResultSetResp(TOperationHandle operationHandle) throws SQLException { TFetchResultsReq req = createFetchResultsReqWithDefaults(operationHandle); return executeFetchRequest(req); } @@ -151,7 +150,7 @@ TFetchResultsResp getResultSetResp(TOperationHandle operationHandle) * @throws DatabricksHttpException if fetch fails */ TFetchResultsResp getResultSetResp(TOperationHandle operationHandle, long startRowOffset) - throws DatabricksHttpException { + throws SQLException { TFetchResultsReq req = createFetchResultsReqWithDefaults(operationHandle); req.setStartRowOffset(startRowOffset); return executeFetchRequest(req); @@ -184,7 +183,7 @@ TCloseOperationResp closeOperation(TCloseOperationReq req) throws DatabricksHttp } TFetchResultsResp getMoreResults(IDatabricksStatementInternal parentStatement) - throws DatabricksSQLException { + throws SQLException { TFetchResultsReq req = createFetchResultsReqWithDefaults(getOperationHandle(parentStatement.getStatementId())); setFetchMetadata(req); @@ -502,8 +501,7 @@ void updateConfig(DatabricksConfig newConfig) { this.databricksConfig = newConfig; } - private TFetchResultsResp executeFetchRequest(TFetchResultsReq request) - throws DatabricksHttpException { + private TFetchResultsResp executeFetchRequest(TFetchResultsReq request) throws SQLException { TFetchResultsResp response; try { response = getThriftClient().FetchResults(request); @@ -554,7 +552,7 @@ private void setFetchMetadata(TFetchResultsReq request) { * @throws DatabricksHttpException if the fetch fails */ TFetchResultsResp fetchResultsWithAbsoluteOffset( - TOperationHandle operationHandle, long startRowOffset) throws DatabricksHttpException { + TOperationHandle operationHandle, long startRowOffset) throws SQLException { String statementId = StatementId.loggableStatementId(operationHandle); LOGGER.debug( "Fetching results with FETCH_ABSOLUTE at offset {} for statement {}", @@ -849,8 +847,8 @@ private TGetOperationStatusResp getOperationStatus( "Statement [{}] Thrift operation status latency: {}ms", statementId, operationStatusLatencyMillis); - TelemetryCollector.getInstance() - .recordGetOperationStatus(statementId.toSQLExecStatementId(), operationStatusLatencyMillis); + TelemetryHelper.recordGetOperationStatus( + connectionContext, statementId.toSQLExecStatementId(), operationStatusLatencyMillis); return operationStatus; } } 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 4f78b2f2c1..04c37f4ed0 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 @@ -7,7 +7,6 @@ import static com.databricks.jdbc.common.util.DatabricksTypeUtil.DECIMAL; import static com.databricks.jdbc.common.util.DatabricksTypeUtil.getDecimalTypeString; import static com.databricks.jdbc.dbclient.impl.sqlexec.CommandName.LIST_FUNCTIONS; -import static com.databricks.jdbc.dbclient.impl.sqlexec.ResultConstants.TYPE_INFO_RESULT; import com.databricks.jdbc.api.impl.*; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; @@ -93,7 +92,7 @@ public ImmutableSessionInfo createSession( String catalog, String schema, Map sessionConf) - throws DatabricksSQLException { + throws SQLException { LOGGER.debug( String.format( "public Session createSession(Compute cluster = {%s}, String catalog = {%s}, String schema = {%s}, Map sessionConf = {%s})", @@ -131,7 +130,7 @@ public ImmutableSessionInfo createSession( } @Override - public void deleteSession(ImmutableSessionInfo sessionInfo) throws DatabricksSQLException { + public void deleteSession(ImmutableSessionInfo sessionInfo) throws SQLException { LOGGER.debug( String.format( "public void deleteSession(Session session = {%s}))", sessionInfo.toString())); @@ -232,7 +231,8 @@ private TExecuteStatementReq getRequest( request.setCanDecompressLZ4Result(true); } if (ProtocolFeatureUtil.supportsCloudFetch(serverProtocolVersion)) { - request.setCanDownloadResult(true); + // Use EnableQueryResultDownload param to control CloudFetch vs inline Arrow + request.setCanDownloadResult(this.connectionContext.isCloudFetchEnabled()); } if (ProtocolFeatureUtil.supportsAdvancedArrowTypes(serverProtocolVersion)) { arrowNativeTypes @@ -302,8 +302,7 @@ public DatabricksResultSet getStatementResult( @Override public ChunkLinkFetchResult getResultChunks( - StatementId statementId, long chunkIndex, long chunkStartRowOffset) - throws DatabricksSQLException { + StatementId statementId, long chunkIndex, long chunkStartRowOffset) throws SQLException { // Thrift uses rowOffset with FETCH_ABSOLUTE; chunkIndex is used for link metadata LOGGER.debug( "getResultChunks(statementId={}, chunkIndex={}, rowOffset={}) using Thrift client", @@ -386,7 +385,7 @@ public ResultData getResultChunksData(StatementId statementId, long chunkIndex) @Override public DatabricksResultSet listTypeInfo(IDatabricksSession session) { LOGGER.debug("public ResultSet getTypeInfo()"); - return TYPE_INFO_RESULT; + return metadataResultSetBuilder.getTypeInfoResult(); } @Override @@ -399,8 +398,11 @@ public DatabricksResultSet listCatalogs(IDatabricksSession session) throws SQLEx // If multiple catalog support is disabled, return only the current catalog if (!isMultipleCatalogSupportEnabled()) { String currentCatalog = session.getCurrentCatalog(); - if (currentCatalog == null || currentCatalog.isEmpty()) { - currentCatalog = ""; + if (currentCatalog == null) { + currentCatalog = "spark"; + LOGGER.debug( + "Current catalog is null when multiple catalog support is disabled. Using default catalog: {}", + currentCatalog); } List> singleCatalogRows = new ArrayList<>(); List catalogRow = new ArrayList<>(); @@ -428,6 +430,11 @@ public DatabricksResultSet listSchemas( "Fetching schemas using Thrift client. Session {%s}, catalog {%s}, schemaNamePattern {%s}", session.toString(), catalog, schemaNamePattern); LOGGER.debug(context); + + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, null, session)) { + return metadataResultSetBuilder.getSchemasResult(new ArrayList<>()); + } + DatabricksThreadContextHolder.setSessionId(session.getSessionId()); TGetSchemasReq request = new TGetSchemasReq() @@ -457,6 +464,11 @@ public DatabricksResultSet listTables( "Fetching tables using Thrift client. Session {%s}, catalog {%s}, schemaNamePattern {%s}, tableNamePattern {%s}", session.toString(), catalog, schemaNamePattern, tableNamePattern); LOGGER.debug(context); + + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, null, session)) { + return metadataResultSetBuilder.getTablesResult(catalog, tableTypes, new ArrayList<>()); + } + DatabricksThreadContextHolder.setSessionId(session.getSessionId()); TGetTablesReq request = new TGetTablesReq() @@ -491,12 +503,17 @@ public DatabricksResultSet listColumns( String schemaNamePattern, String tableNamePattern, String columnNamePattern) - throws DatabricksSQLException { + throws SQLException { String context = String.format( "Fetching columns using Thrift client. Session {%s}, catalog {%s}, schemaNamePattern {%s}, tableNamePattern {%s}, columnNamePattern {%s}", session.toString(), catalog, schemaNamePattern, tableNamePattern, columnNamePattern); LOGGER.debug(context); + + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, null, session)) { + return metadataResultSetBuilder.getColumnsResult(new ArrayList<>()); + } + DatabricksThreadContextHolder.setSessionId(session.getSessionId()); TGetColumnsReq request = new TGetColumnsReq() @@ -526,6 +543,11 @@ public DatabricksResultSet listFunctions( session.toString(), catalog, schemaNamePattern, functionNamePattern); DatabricksThreadContextHolder.setSessionId(session.getSessionId()); LOGGER.debug(context); + + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, null, session)) { + return metadataResultSetBuilder.getFunctionsResult(catalog, new ArrayList<>()); + } + if (connectionContext.enableShowCommandsForGetFunctions()) { // Return empty result set if catalog is null for SQL command path if (catalog == null) { @@ -578,6 +600,11 @@ public DatabricksResultSet listPrimaryKeys( "Fetching primary keys using Thrift client. session {%s}, catalog {%s}, schema {%s}, table {%s}", session.toString(), catalog, schema, table); LOGGER.debug(context); + + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, null, session)) { + return metadataResultSetBuilder.getPrimaryKeysResult(new ArrayList<>()); + } + DatabricksThreadContextHolder.setSessionId(session.getSessionId()); TGetPrimaryKeysReq request = new TGetPrimaryKeysReq() @@ -601,6 +628,11 @@ public DatabricksResultSet listImportedKeys( "Fetching imported keys using Thrift client for session {%s}, catalog {%s}, schema {%s}, table {%s}", session.toString(), catalog, schema, table); LOGGER.debug(context); + + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, null, session)) { + return metadataResultSetBuilder.getImportedKeys(new ArrayList<>()); + } + DatabricksThreadContextHolder.setSessionId(session.getSessionId()); // GetImportedKeys is implemented using GetCrossReferences // When only foreign table name is provided, we get imported keys @@ -625,6 +657,11 @@ public DatabricksResultSet listExportedKeys( "Fetching exported keys using Thrift client for session {%s}, catalog {%s}, schema {%s}, table {%s}", session.toString(), catalog, schema, table); LOGGER.debug(context); + + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, null, session)) { + return metadataResultSetBuilder.getExportedKeys(new ArrayList<>()); + } + // GetImportedKeys is implemented using GetCrossReferences // When only parent table name is provided, we get exported keys TGetCrossReferenceReq request = @@ -661,6 +698,12 @@ public DatabricksResultSet listCrossReferences( foreignSchema, foreignTable); LOGGER.debug(context); + + if (!metadataResultSetBuilder.shouldAllowCatalogAccess(parentCatalog, null, session) + || !metadataResultSetBuilder.shouldAllowCatalogAccess(foreignCatalog, null, session)) { + return metadataResultSetBuilder.getCrossRefsResult(new ArrayList<>()); + } + TGetCrossReferenceReq request = new TGetCrossReferenceReq() .setSessionHandle(Objects.requireNonNull(session.getSessionInfo()).sessionHandle()) @@ -679,7 +722,7 @@ public DatabricksResultSet listCrossReferences( } public TFetchResultsResp getMoreResults(IDatabricksStatementInternal parentStatement) - throws DatabricksSQLException { + throws SQLException { return thriftAccessor.getMoreResults(parentStatement); } diff --git a/src/main/java/com/databricks/jdbc/exception/DatabricksTimeoutException.java b/src/main/java/com/databricks/jdbc/exception/DatabricksTimeoutException.java index 4e60f110dc..66f58cc02d 100644 --- a/src/main/java/com/databricks/jdbc/exception/DatabricksTimeoutException.java +++ b/src/main/java/com/databricks/jdbc/exception/DatabricksTimeoutException.java @@ -18,4 +18,14 @@ public DatabricksTimeoutException( reason, TelemetryLogLevel.ERROR); } + + public DatabricksTimeoutException( + String reason, String sqlState, Throwable cause, DatabricksDriverErrorCode internalError) { + super(reason, sqlState, cause); + TelemetryHelper.exportFailureLog( + DatabricksThreadContextHolder.getConnectionContext(), + internalError.name(), + reason, + TelemetryLogLevel.ERROR); + } } diff --git a/src/main/java/com/databricks/jdbc/model/telemetry/DriverConnectionParameters.java b/src/main/java/com/databricks/jdbc/model/telemetry/DriverConnectionParameters.java index ecdc927423..f3394df738 100644 --- a/src/main/java/com/databricks/jdbc/model/telemetry/DriverConnectionParameters.java +++ b/src/main/java/com/databricks/jdbc/model/telemetry/DriverConnectionParameters.java @@ -139,6 +139,9 @@ public class DriverConnectionParameters { @JsonProperty("enable_metric_view_metadata") boolean enableMetricViewMetadata; + @JsonProperty("query_tags") + String queryTags; + public DriverConnectionParameters setHttpPath(String httpPath) { this.httpPath = httpPath; return this; @@ -364,6 +367,11 @@ public DriverConnectionParameters setEnableMetricViewMetadata(boolean enableMetr return this; } + public DriverConnectionParameters setQueryTags(String queryTags) { + this.queryTags = queryTags; + return this; + } + @Override public String toString() { return new ToStringer(DriverConnectionParameters.class) @@ -411,6 +419,7 @@ public String toString() { .add("rowsFetchedPerBlock", rowsFetchedPerBlock) .add("asyncPollIntervalMillis", asyncPollIntervalMillis) .add("enableMetricViewMetadata", enableMetricViewMetadata) + .add("queryTags", queryTags) .toString(); } } diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryClient.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryClient.java index dfa8aa3e98..e1a645c4ab 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryClient.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryClient.java @@ -5,12 +5,12 @@ import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.telemetry.TelemetryFrontendLog; import com.databricks.jdbc.telemetry.latency.TelemetryCollector; +import com.databricks.jdbc.telemetry.latency.TelemetryCollectorManager; import com.databricks.sdk.core.DatabricksConfig; import java.util.LinkedList; import java.util.List; import java.util.concurrent.*; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.atomic.AtomicInteger; public class TelemetryClient implements ITelemetryClient { private static final int MINIMUM_TELEMETRY_FLUSH_MILLISECONDS = 1000; @@ -26,31 +26,27 @@ public class TelemetryClient implements ITelemetryClient { private ScheduledFuture flushTask; private final int flushIntervalMillis; - private static ThreadFactory createSchedulerThreadFactory() { - return new ThreadFactory() { - private final AtomicInteger threadNumber = new AtomicInteger(1); - - @Override - public Thread newThread(Runnable r) { - Thread thread = new Thread(r, "Telemetry-Scheduler-" + threadNumber.getAndIncrement()); - thread.setDaemon(true); - return thread; - } - }; - } - public TelemetryClient( IDatabricksConnectionContext connectionContext, ExecutorService executorService, + ScheduledExecutorService scheduledExecutorService, DatabricksConfig config) { this.eventsBatch = new LinkedList<>(); this.eventsBatchSize = connectionContext.getTelemetryBatchSize(); this.context = connectionContext; this.databricksConfig = config; this.executorService = executorService; - this.scheduledExecutorService = - Executors.newSingleThreadScheduledExecutor(createSchedulerThreadFactory()); - this.flushIntervalMillis = context.getTelemetryFlushIntervalInMilliseconds(); + /* + * The scheduledExecutorService is shared across all telemetry clients and only schedules + * periodic flush checks. The actual flush work (network I/O) is submitted asynchronously + * to the executorService (10-thread pool), so a slow flush on one statement does not block + * flushes for other statements as long as worker threads are available in the pool. + */ + this.scheduledExecutorService = scheduledExecutorService; + this.flushIntervalMillis = + Math.max( + context.getTelemetryFlushIntervalInMilliseconds(), + MINIMUM_TELEMETRY_FLUSH_MILLISECONDS); // To avoid illegalArgument exception in any case this.lastFlushedTime = System.currentTimeMillis(); this.telemetryPushClient = TelemetryClientFactory.getTelemetryPushClient( @@ -59,14 +55,15 @@ public TelemetryClient( } public TelemetryClient( - IDatabricksConnectionContext connectionContext, ExecutorService executorService) { + IDatabricksConnectionContext connectionContext, + ExecutorService executorService, + ScheduledExecutorService scheduledExecutorService) { this.eventsBatch = new LinkedList<>(); eventsBatchSize = connectionContext.getTelemetryBatchSize(); this.context = connectionContext; this.databricksConfig = null; this.executorService = executorService; - this.scheduledExecutorService = - Executors.newSingleThreadScheduledExecutor(createSchedulerThreadFactory()); + this.scheduledExecutorService = scheduledExecutorService; this.flushIntervalMillis = Math.max( context.getTelemetryFlushIntervalInMilliseconds(), @@ -108,8 +105,10 @@ public void exportEvent(TelemetryFrontendLog event) { @Override public void close() { - // Export any pending latency telemetry before flushing - TelemetryCollector.getInstance().exportAllPendingTelemetryDetails(); + // Export any pending latency telemetry before flushing for this connection + TelemetryCollector collector = + TelemetryCollectorManager.getInstance().getOrCreateCollector(context); + collector.exportAllPendingTelemetryDetails(); try { // Synchronously flush the remaining events and wait for the task to complete @@ -129,22 +128,13 @@ public void close() { flushTask.cancel(false); } - // Shut down the scheduler. - // The executorService is assumed to be a shared resource and is not shut down here. - scheduledExecutorService.shutdown(); - try { - if (!scheduledExecutorService.awaitTermination(5, TimeUnit.SECONDS)) { - scheduledExecutorService.shutdownNow(); - } - } catch (InterruptedException ie) { - LOGGER.trace("Interrupted while waiting for flush to finish. Error: {}", ie); - Thread.currentThread().interrupt(); - scheduledExecutorService.shutdownNow(); - } + // Note: Both executorService and scheduledExecutorService are shared resources + // managed by TelemetryClientFactory and should not be shut down here. } /** - * Submits a flush task to the executor service. + * Submits a flush task to the executor service. Non-blocking: uses a shared thread pool (10 + * threads) so slow flushes don't block other statements. * * @param forceFlush - Flushes the eventsBatch for all size variations if forceFlush, otherwise * only flushes if eventsBatch size has breached diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java index de8de5eab4..9fda43caca 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java @@ -7,13 +7,18 @@ import com.databricks.jdbc.exception.DatabricksParsingException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; +import com.databricks.jdbc.telemetry.latency.TelemetryCollector; +import com.databricks.jdbc.telemetry.latency.TelemetryCollectorManager; import com.databricks.sdk.core.DatabricksConfig; import com.google.common.annotations.VisibleForTesting; import java.util.Map; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; public class TelemetryClientFactory { @@ -30,6 +35,7 @@ public class TelemetryClientFactory { final Map noauthTelemetryClientHolders = new ConcurrentHashMap<>(); private final ExecutorService telemetryExecutorService; + private ScheduledExecutorService sharedSchedulerService; private static ThreadFactory createThreadFactory() { return new ThreadFactory() { @@ -45,8 +51,23 @@ public Thread newThread(Runnable r) { }; } + private static ThreadFactory createSchedulerThreadFactory() { + return new ThreadFactory() { + private final AtomicInteger threadNumber = new AtomicInteger(1); + + @Override + public Thread newThread(Runnable r) { + Thread thread = new Thread(r, "Telemetry-Scheduler-" + threadNumber.getAndIncrement()); + thread.setDaemon(true); + return thread; + } + }; + } + private TelemetryClientFactory() { telemetryExecutorService = Executors.newFixedThreadPool(10, createThreadFactory()); + sharedSchedulerService = + Executors.newSingleThreadScheduledExecutor(createSchedulerThreadFactory()); } public static TelemetryClientFactory getInstance() { @@ -69,14 +90,19 @@ public ITelemetryClient getTelemetryClient(IDatabricksConnectionContext connecti try { return new TelemetryClientHolder( new TelemetryClient( - connectionContext, getTelemetryExecutorService(), databricksConfig), - 1); + connectionContext, + getTelemetryExecutorService(), + getSharedSchedulerService(), + databricksConfig), + connectionContext.getConnectionUuid()); } catch (Exception e) { // Validation or other errors during client creation - fail silently + LOGGER.trace("Skipping telemetry, client creation failed {}", e); return null; } } - existing.refCount.incrementAndGet(); + // Track this unique connection + existing.connectionUuids.add(connectionContext.getConnectionUuid()); return existing; }); return holder != null ? holder.client : NoopTelemetryClient.getInstance(); @@ -90,47 +116,74 @@ connectionContext, getTelemetryExecutorService(), databricksConfig), if (existing == null) { try { return new TelemetryClientHolder( - new TelemetryClient(connectionContext, getTelemetryExecutorService()), 1); + new TelemetryClient( + connectionContext, + getTelemetryExecutorService(), + getSharedSchedulerService()), + connectionContext.getConnectionUuid()); } catch (Exception e) { // Validation or other errors during client creation - fail silently - LOGGER.trace("Skipping telemetry, client creation failed {}", e); + LOGGER.trace("Skipping no-auth telemetry, client creation failed {}", e); return null; } } - existing.refCount.incrementAndGet(); + // Track this unique connection + existing.connectionUuids.add(connectionContext.getConnectionUuid()); return existing; }); return holder != null ? holder.client : NoopTelemetryClient.getInstance(); } + /** + * Closes telemetry client for a connection. Thread-safe: computeIfPresent ensures atomic locking, + * preventing race conditions between connection removal and addition. + */ public void closeTelemetryClient(IDatabricksConnectionContext connectionContext) { String key = TelemetryHelper.keyOf(connectionContext); + String connectionUuid = connectionContext.getConnectionUuid(); + // Atomically remove connection and close client if no connections remain for this key telemetryClientHolders.computeIfPresent( key, (k, holder) -> { - if (holder.refCount.get() <= 1) { + holder.connectionUuids.remove(connectionUuid); + if (holder.connectionUuids.isEmpty()) { closeTelemetryClient(holder.client, "telemetry client"); return null; } - holder.refCount.decrementAndGet(); return holder; }); + // Atomically remove connection and close client if no connections remain for this key noauthTelemetryClientHolders.computeIfPresent( key, (k, holder) -> { - if (holder.refCount.get() <= 1) { + holder.connectionUuids.remove(connectionUuid); + if (holder.connectionUuids.isEmpty()) { closeTelemetryClient(holder.client, "unauthenticated telemetry client"); return null; } - holder.refCount.decrementAndGet(); return holder; }); + + // Export and remove the TelemetryCollector for this connection + TelemetryCollector collector = + TelemetryCollectorManager.getInstance().removeCollector(connectionContext); + if (collector != null) { + // Export any remaining telemetry before removing + collector.exportAllPendingTelemetryDetails(); + } + + // Clean up cached connection parameters to prevent memory leaks + TelemetryHelper.removeConnectionParameters(connectionContext.getConnectionUuid()); } public ExecutorService getTelemetryExecutorService() { return telemetryExecutorService; } + public ScheduledExecutorService getSharedSchedulerService() { + return sharedSchedulerService; + } + static ITelemetryPushClient getTelemetryPushClient( Boolean isAuthenticated, IDatabricksConnectionContext connectionContext, @@ -156,13 +209,31 @@ static ITelemetryPushClient getTelemetryPushClient( @VisibleForTesting public void reset() { - // Close all existing clients + // Close all existing clients (cancels their scheduled tasks) telemetryClientHolders.values().forEach(holder -> holder.client.close()); noauthTelemetryClientHolders.values().forEach(holder -> holder.client.close()); // Clear the maps telemetryClientHolders.clear(); noauthTelemetryClientHolders.clear(); + + // Clear cached connection parameters + TelemetryHelper.clearConnectionParameterCache(); + + // Shutdown shared scheduler service (test cleanup only) + sharedSchedulerService.shutdown(); + try { + if (!sharedSchedulerService.awaitTermination(5, TimeUnit.SECONDS)) { + sharedSchedulerService.shutdownNow(); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + sharedSchedulerService.shutdownNow(); + } + + // Recreate the scheduler for subsequent tests + sharedSchedulerService = + Executors.newSingleThreadScheduledExecutor(createSchedulerThreadFactory()); } private void closeTelemetryClient(ITelemetryClient client, String clientType) { @@ -177,11 +248,12 @@ private void closeTelemetryClient(ITelemetryClient client, String clientType) { private static final class TelemetryClientHolder { final TelemetryClient client; - final AtomicInteger refCount; + final Set connectionUuids; // Track unique connections - TelemetryClientHolder(TelemetryClient client, int initialCount) { + TelemetryClientHolder(TelemetryClient client, String connectionUuid) { this.client = client; - this.refCount = new AtomicInteger(initialCount); + this.connectionUuids = ConcurrentHashMap.newKeySet(); + this.connectionUuids.add(connectionUuid); } } diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java index 6d6cfd5f15..2dd206f2b6 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java @@ -1,8 +1,11 @@ package com.databricks.jdbc.telemetry; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.QUERY_TAGS; import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty; +import com.databricks.jdbc.api.impl.DatabricksConnection; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; +import com.databricks.jdbc.api.internal.IDatabricksStatementInternal; import com.databricks.jdbc.common.DatabricksClientConfiguratorManager; import com.databricks.jdbc.common.TelemetryLogLevel; import com.databricks.jdbc.common.safe.DatabricksDriverFeatureFlagsContextFactory; @@ -15,12 +18,15 @@ import com.databricks.jdbc.exception.DatabricksValidationException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; +import com.databricks.jdbc.model.client.thrift.generated.TSparkRowSetType; import com.databricks.jdbc.model.telemetry.*; import com.databricks.jdbc.model.telemetry.latency.OperationType; import com.databricks.jdbc.telemetry.latency.TelemetryCollector; +import com.databricks.jdbc.telemetry.latency.TelemetryCollectorManager; import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.core.ProxyConfig; import com.databricks.sdk.core.UserAgent; +import com.databricks.sdk.service.sql.Format; import com.google.common.annotations.VisibleForTesting; import java.nio.charset.Charset; import java.time.Instant; @@ -147,7 +153,9 @@ public static void exportFailureLog( if (statementId == null) { telemetryDetails = new StatementTelemetryDetails(null); } else { - telemetryDetails = TelemetryCollector.getInstance().getOrCreateTelemetryDetails(statementId); + TelemetryCollector collector = + TelemetryCollectorManager.getInstance().getOrCreateCollector(connectionContext); + telemetryDetails = collector.getOrCreateTelemetryDetails(statementId); } exportTelemetryEvent(connectionContext, telemetryDetails, errorInfo, chunkIndex, logLevel); } @@ -223,7 +231,8 @@ private static DriverConnectionParameters buildDriverConnectionParameters( .setAsyncPollIntervalMillis(connectionContext.getAsyncExecPollInterval()) .setEnableTokenCache(connectionContext.isTokenCacheEnabled()) .setHttpPath(connectionContext.getHttpPath()) - .setEnableMetricViewMetadata(connectionContext.getEnableMetricViewMetadata()); + .setEnableMetricViewMetadata(connectionContext.getEnableMetricViewMetadata()) + .setQueryTags(connectionContext.getSessionConfigs().get(QUERY_TAGS)); } catch (DatabricksValidationException e) { // If configuration validation fails, return null to skip telemetry export // This prevents invalid configuration from breaking telemetry @@ -392,4 +401,181 @@ public static String keyOf(IDatabricksConnectionContext context) { } return context.getHost(); } + + /** + * Removes cached connection parameters for the given connection UUID. Should be called when a + * connection is closed to prevent memory leaks. + * + * @param connectionUuid The connection UUID whose cached parameters should be removed + */ + public static void removeConnectionParameters(String connectionUuid) { + if (connectionUuid != null) { + connectionParameterCache.remove(connectionUuid); + } + } + + /** Clears all cached connection parameters. This should only be used for testing purposes. */ + @VisibleForTesting + static void clearConnectionParameterCache() { + connectionParameterCache.clear(); + } + + // Simplified telemetry recording methods that handle getting the collector internally + + /** + * Records the total chunks for a statement. Silently ignores errors. + * + * @param connectionContext The connection context + * @param statementId The statement ID + * @param chunkCount The total number of chunks + */ + public static void recordTotalChunks( + IDatabricksConnectionContext connectionContext, StatementId statementId, long chunkCount) { + try { + if (connectionContext != null) { + TelemetryCollectorManager.getInstance() + .getOrCreateCollector(connectionContext) + .recordTotalChunks(statementId, chunkCount); + } + } catch (Exception e) { + LOGGER.trace("Error recording total chunks telemetry: {}", e.getMessage()); + } + } + + /** + * Sets the result format for a statement (SQL Execution API). Silently ignores errors. + * + * @param connectionContext The connection context + * @param statementId The statement ID + * @param format The result format + */ + public static void setResultFormat( + IDatabricksConnectionContext connectionContext, StatementId statementId, Format format) { + try { + if (connectionContext != null) { + TelemetryCollectorManager.getInstance() + .getOrCreateCollector(connectionContext) + .setResultFormat(statementId, format); + } + } catch (Exception e) { + LOGGER.trace("Error setting result format telemetry: {}", e.getMessage()); + } + } + + /** + * Sets the result format for a statement (Thrift API). Silently ignores errors. + * + * @param connectionContext The connection context + * @param parentStatement The parent statement + * @param format The result format + */ + public static void setResultFormat( + IDatabricksConnectionContext connectionContext, + IDatabricksStatementInternal parentStatement, + TSparkRowSetType format) { + try { + if (connectionContext != null) { + TelemetryCollectorManager.getInstance() + .getOrCreateCollector(connectionContext) + .setResultFormat(parentStatement, format); + } + } catch (Exception e) { + LOGGER.trace("Error setting result format telemetry: {}", e.getMessage()); + } + } + + /** + * Records a result set iteration. Silently ignores errors. + * + * @param connectionContext The connection context + * @param statementId The statement ID + * @param chunkCount The total number of chunks + * @param hasNext Whether there are more rows + */ + public static void recordResultSetIteration( + IDatabricksConnectionContext connectionContext, + String statementId, + long chunkCount, + boolean hasNext) { + try { + if (connectionContext != null) { + TelemetryCollectorManager.getInstance() + .getOrCreateCollector(connectionContext) + .recordResultSetIteration(statementId, chunkCount, hasNext); + } + } catch (Exception e) { + LOGGER.trace("Error recording result set iteration telemetry: {}", e.getMessage()); + } + } + + /** + * Records a result set iteration from a parent statement. Extracts connection context safely and + * silently ignores all errors. + * + * @param parentStatement The parent statement (can be null) + * @param statementId The statement ID + * @param chunkCount The total number of chunks (can be null) + * @param hasNext Whether there are more rows + */ + public static void recordResultSetIteration( + IDatabricksStatementInternal parentStatement, + StatementId statementId, + Long chunkCount, + boolean hasNext) { + try { + if (parentStatement != null && chunkCount != null) { + IDatabricksConnectionContext connectionContext = + ((DatabricksConnection) parentStatement.getStatement().getConnection()) + .getConnectionContext(); + recordResultSetIteration( + connectionContext, statementId.toSQLExecStatementId(), chunkCount, hasNext); + } + } catch (Exception e) { + LOGGER.trace("Error getting connection context for telemetry: {}", e.getMessage()); + } + } + + /** + * Records get operation status latency. Silently ignores errors. + * + * @param connectionContext The connection context + * @param statementId The statement ID + * @param latencyMillis The operation latency in milliseconds + */ + public static void recordGetOperationStatus( + IDatabricksConnectionContext connectionContext, String statementId, long latencyMillis) { + try { + if (connectionContext != null) { + TelemetryCollectorManager.getInstance() + .getOrCreateCollector(connectionContext) + .recordGetOperationStatus(statementId, latencyMillis); + } + } catch (Exception e) { + LOGGER.trace("Error recording get operation status telemetry: {}", e.getMessage()); + } + } + + /** + * Records chunk download latency. Silently ignores errors. + * + * @param connectionContext The connection context + * @param statementId The statement ID + * @param chunkIndex The chunk index + * @param latencyMillis The download latency in milliseconds + */ + public static void recordChunkDownloadLatency( + IDatabricksConnectionContext connectionContext, + String statementId, + long chunkIndex, + long latencyMillis) { + try { + if (connectionContext != null) { + TelemetryCollectorManager.getInstance() + .getOrCreateCollector(connectionContext) + .recordChunkDownloadLatency(statementId, chunkIndex, latencyMillis); + } + } catch (Exception e) { + LOGGER.trace("Error recording chunk download latency telemetry: {}", e.getMessage()); + } + } } diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryPushClient.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryPushClient.java index 96e3ded8ee..f871819282 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryPushClient.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryPushClient.java @@ -1,5 +1,7 @@ package com.databricks.jdbc.telemetry; +import static com.databricks.jdbc.common.util.JsonUtil.getTelemetryMapper; + import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.common.DatabricksJdbcConstants; import com.databricks.jdbc.common.util.HttpUtil; @@ -12,8 +14,6 @@ import com.databricks.jdbc.model.telemetry.TelemetryRequest; import com.databricks.jdbc.model.telemetry.TelemetryResponse; import com.databricks.sdk.core.DatabricksConfig; -import com.fasterxml.jackson.annotation.JsonInclude; -import com.fasterxml.jackson.databind.ObjectMapper; import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.Map; @@ -31,8 +31,6 @@ public class TelemetryPushClient implements ITelemetryPushClient { private final boolean isAuthenticated; private final IDatabricksConnectionContext connectionContext; private final DatabricksConfig databricksConfig; - private final ObjectMapper objectMapper = - new ObjectMapper().setSerializationInclusion(JsonInclude.Include.NON_NULL); public TelemetryPushClient( boolean isAuthenticated, @@ -54,7 +52,7 @@ public void pushEvent(TelemetryRequest request) throws Exception { String uri = new URIBuilder(connectionContext.getHostUrl()).setPath(path).toString(); HttpPost post = new HttpPost(uri); post.setEntity( - new StringEntity(objectMapper.writeValueAsString(request), StandardCharsets.UTF_8)); + new StringEntity(getTelemetryMapper().writeValueAsString(request), StandardCharsets.UTF_8)); DatabricksJdbcConstants.JSON_HTTP_HEADERS.forEach(post::addHeader); Map authHeaders = isAuthenticated ? databricksConfig.authenticate() : Collections.emptyMap(); @@ -72,8 +70,8 @@ public void pushEvent(TelemetryRequest request) throws Exception { } } TelemetryResponse telResponse = - objectMapper.readValue( - EntityUtils.toString(response.getEntity()), TelemetryResponse.class); + getTelemetryMapper() + .readValue(EntityUtils.toString(response.getEntity()), TelemetryResponse.class); LOGGER.trace( "Pushed Telemetry logs with request-Id {} with events {} with error count {}", response.getFirstHeader(REQUEST_ID_HEADER), @@ -93,7 +91,7 @@ public void pushEvent(TelemetryRequest request) throws Exception { LOGGER.debug( "Failed to push telemetry logs with error: {}, request: {}", e.getMessage(), - objectMapper.writeValueAsString(request)); + getTelemetryMapper().writeValueAsString(request)); if (connectionContext.isTelemetryCircuitBreakerEnabled()) { throw new DatabricksTelemetryException("Exception while pushing telemetry logs", e); } diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryPushTask.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryPushTask.java index c1cbf69e1b..c30dc68571 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryPushTask.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryPushTask.java @@ -1,12 +1,12 @@ package com.databricks.jdbc.telemetry; +import static com.databricks.jdbc.common.util.JsonUtil.getTelemetryMapper; + import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.telemetry.TelemetryFrontendLog; import com.databricks.jdbc.model.telemetry.TelemetryRequest; -import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.core.JsonProcessingException; -import com.fasterxml.jackson.databind.ObjectMapper; import java.util.List; import java.util.Objects; import java.util.stream.Collectors; @@ -17,8 +17,6 @@ class TelemetryPushTask implements Runnable { private final List queueToBePushed; private final ITelemetryPushClient telemetryPushClient; - private final ObjectMapper objectMapper = - new ObjectMapper().setSerializationInclusion(JsonInclude.Include.NON_NULL); TelemetryPushTask( List eventsQueue, ITelemetryPushClient telemetryPushClient) { @@ -41,7 +39,7 @@ public void run() { .map( event -> { try { - return objectMapper.writeValueAsString(event); + return getTelemetryMapper().writeValueAsString(event); } catch (JsonProcessingException e) { LOGGER.trace( "Failed to serialize Telemetry event {} with error: {}", event, e); diff --git a/src/main/java/com/databricks/jdbc/telemetry/latency/DatabricksMetricsTimedProcessor.java b/src/main/java/com/databricks/jdbc/telemetry/latency/DatabricksMetricsTimedProcessor.java index a4d97344c6..c7a0d570df 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/latency/DatabricksMetricsTimedProcessor.java +++ b/src/main/java/com/databricks/jdbc/telemetry/latency/DatabricksMetricsTimedProcessor.java @@ -1,5 +1,6 @@ package com.databricks.jdbc.telemetry.latency; +import com.databricks.jdbc.common.util.DatabricksThreadContextHolder; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import java.lang.reflect.InvocationHandler; @@ -74,8 +75,12 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl argsStr, executionTimeMillis); try { - TelemetryCollector.getInstance() - .recordOperationLatency(executionTimeMillis, methodName); + TelemetryCollector collector = + TelemetryCollectorManager.getInstance() + .getCollectorSafely(DatabricksThreadContextHolder::getConnectionContext); + if (collector != null) { + collector.recordOperationLatency(executionTimeMillis, methodName); + } } catch (Exception e) { LOGGER.trace( "Failed to export latency metrics for method {}: {}", methodName, e.getMessage()); diff --git a/src/main/java/com/databricks/jdbc/telemetry/latency/TelemetryCollector.java b/src/main/java/com/databricks/jdbc/telemetry/latency/TelemetryCollector.java index fd40a75d1f..3ea6365509 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/latency/TelemetryCollector.java +++ b/src/main/java/com/databricks/jdbc/telemetry/latency/TelemetryCollector.java @@ -25,19 +25,15 @@ public class TelemetryCollector { private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(TelemetryCollector.class); - // Singleton instance for global access - private static final TelemetryCollector INSTANCE = new TelemetryCollector(); - // Per-statement latency tracking using StatementLatencyDetails private final ConcurrentHashMap statementTrackers = new ConcurrentHashMap<>(); - private TelemetryCollector() { - // Private constructor for singleton - } - - public static TelemetryCollector getInstance() { - return INSTANCE; + /** + * Package-private constructor - instances should only be created via TelemetryCollectorManager + */ + TelemetryCollector() { + // Constructor for per-connection instances } /** diff --git a/src/main/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorManager.java b/src/main/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorManager.java new file mode 100644 index 0000000000..b0cf2946ec --- /dev/null +++ b/src/main/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorManager.java @@ -0,0 +1,88 @@ +package com.databricks.jdbc.telemetry.latency; + +import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Supplier; + +/** + * Manages TelemetryCollector instances per connection. Each connection gets its own + * TelemetryCollector instance to ensure telemetry data is properly isolated. + */ +public class TelemetryCollectorManager { + private static final TelemetryCollectorManager INSTANCE = new TelemetryCollectorManager(); + private static final String DEFAULT_CONNECTION = "unknown-connection"; + + private final ConcurrentHashMap collectors = + new ConcurrentHashMap<>(); + + private TelemetryCollectorManager() { + // Private constructor for singleton + } + + public static TelemetryCollectorManager getInstance() { + return INSTANCE; + } + + /** + * Gets or creates a TelemetryCollector for the given connection context. Null checks are retained + * due to ThreadLocal usage patterns that may result in null context. + * + * @param context the connection context + * @return the TelemetryCollector instance for this connection + */ + public TelemetryCollector getOrCreateCollector(IDatabricksConnectionContext context) { + String key = + (context != null && context.getConnectionUuid() != null) + ? context.getConnectionUuid() + : DEFAULT_CONNECTION; + return collectors.computeIfAbsent(key, k -> new TelemetryCollector()); + } + + /** + * Removes and returns the TelemetryCollector for the given connection. Should be called when a + * connection is closed. + * + * @param connectionUuid the connection UUID + * @return the removed TelemetryCollector, or null if not found + */ + public TelemetryCollector removeCollector(String connectionUuid) { + String key = connectionUuid != null ? connectionUuid : DEFAULT_CONNECTION; + return collectors.remove(key); + } + + /** + * Removes and returns the TelemetryCollector for the given connection context. + * + * @param context the connection context + * @return the removed TelemetryCollector, or null if not found + */ + public TelemetryCollector removeCollector(IDatabricksConnectionContext context) { + if (context == null) { + return removeCollector(DEFAULT_CONNECTION); + } + return removeCollector(context.getConnectionUuid()); + } + + /** + * Safely gets a TelemetryCollector from connection context, catching and ignoring any errors. + * This is useful for code that doesn't want telemetry failures to affect main logic. + * + * @param connectionContextSupplier a function that provides the connection context + * @return the TelemetryCollector, or null if an error occurred + */ + public TelemetryCollector getCollectorSafely( + Supplier connectionContextSupplier) { + try { + IDatabricksConnectionContext context = connectionContextSupplier.get(); + return getOrCreateCollector(context); + } catch (Exception e) { + // Silently ignore errors when retrieving telemetry collector + return null; + } + } + + /** Clears all collectors. For testing purposes only. */ + public void clear() { + collectors.clear(); + } +} diff --git a/src/test/java/com/databricks/client/jdbc/DataSourceTest.java b/src/test/java/com/databricks/client/jdbc/DataSourceTest.java index 40a5ac1f9a..24ebd6dace 100644 --- a/src/test/java/com/databricks/client/jdbc/DataSourceTest.java +++ b/src/test/java/com/databricks/client/jdbc/DataSourceTest.java @@ -4,7 +4,6 @@ import static org.junit.jupiter.api.Assertions.*; import com.databricks.jdbc.api.impl.DatabricksConnection; -import com.databricks.jdbc.exception.DatabricksSQLException; import java.sql.Connection; import java.sql.SQLException; import java.sql.SQLFeatureNotSupportedException; @@ -35,7 +34,7 @@ public void testGetUrl() { } @Test - public void testGetConnection() throws DatabricksSQLException { + public void testGetConnection() throws SQLException { Properties properties = new Properties(); properties.setProperty(AUTH_MECH.getParamName(), "3"); diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java index 4a84b077bc..09bf1b4f4a 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java @@ -1210,4 +1210,119 @@ public void testEnableTokenFederation() throws DatabricksSQLException { DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, props); assertFalse(ctx.isTokenFederationEnabled()); } + + @Test + public void testIsCloudFetchEnabled() throws DatabricksSQLException { + // Test default value (should be enabled by default) + DatabricksConnectionContext ctx = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, properties); + assertTrue(ctx.isCloudFetchEnabled()); // Default should be true + + // Test via URL parameter - enabled + String urlWithCloudFetchEnabled = + "jdbc:databricks://sample-host.18.azuredatabricks.net:9999/default;httpPath=/sql/1.0/warehouses/999999999;EnableQueryResultDownload=1"; + ctx = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(urlWithCloudFetchEnabled, properties); + assertTrue(ctx.isCloudFetchEnabled()); + + // Test via URL parameter - disabled + String urlWithCloudFetchDisabled = + "jdbc:databricks://sample-host.18.azuredatabricks.net:9999/default;httpPath=/sql/1.0/warehouses/999999999;EnableQueryResultDownload=0"; + ctx = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(urlWithCloudFetchDisabled, properties); + assertFalse(ctx.isCloudFetchEnabled()); + + // Test via Properties - enabled + Properties props = new Properties(); + props.setProperty("password", "passwd"); + props.setProperty("EnableQueryResultDownload", "1"); + ctx = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, props); + assertTrue(ctx.isCloudFetchEnabled()); + + // Test via Properties - disabled + props = new Properties(); + props.setProperty("password", "passwd"); + props.setProperty("EnableQueryResultDownload", "0"); + ctx = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, props); + assertFalse(ctx.isCloudFetchEnabled()); + } + + // ===== Inline Streaming Configuration Tests ===== + + @Test + public void testIsInlineStreamingEnabledDefault() throws DatabricksSQLException { + // Test default value (should be enabled by default - "1") + DatabricksConnectionContext ctx = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, properties); + assertTrue(ctx.isInlineStreamingEnabled()); // Default should be true + } + + @Test + public void testIsInlineStreamingEnabledDisabled() throws DatabricksSQLException { + // Test via URL parameter - disabled + String urlWithStreamingDisabled = + "jdbc:databricks://sample-host.18.azuredatabricks.net:9999/default;httpPath=/sql/1.0/warehouses/999999999;EnableInlineStreaming=0"; + DatabricksConnectionContext ctx = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(urlWithStreamingDisabled, properties); + assertFalse(ctx.isInlineStreamingEnabled()); + + // Test via Properties - disabled + Properties props = new Properties(); + props.setProperty("password", "passwd"); + props.setProperty("EnableInlineStreaming", "0"); + ctx = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, props); + assertFalse(ctx.isInlineStreamingEnabled()); + } + + @Test + public void testGetThriftMaxBatchesInMemoryDefault() throws DatabricksSQLException { + // Test default value (should be 3) + DatabricksConnectionContext ctx = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, properties); + assertEquals(3, ctx.getThriftMaxBatchesInMemory()); + } + + @Test + public void testGetThriftMaxBatchesInMemoryCustom() throws DatabricksSQLException { + // Test custom value via URL + String urlWithCustomBatches = + "jdbc:databricks://sample-host.18.azuredatabricks.net:9999/default;httpPath=/sql/1.0/warehouses/999999999;ThriftMaxBatchesInMemory=5"; + DatabricksConnectionContext ctx = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(urlWithCustomBatches, properties); + assertEquals(5, ctx.getThriftMaxBatchesInMemory()); + + // Test custom value via Properties + Properties props = new Properties(); + props.setProperty("password", "passwd"); + props.setProperty("ThriftMaxBatchesInMemory", "10"); + ctx = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, props); + assertEquals(10, ctx.getThriftMaxBatchesInMemory()); + } + + @Test + public void testGetThriftMaxBatchesInMemoryInvalidFallback() throws DatabricksSQLException { + // Test invalid value falls back to default (3) + Properties props = new Properties(); + props.setProperty("password", "passwd"); + props.setProperty("ThriftMaxBatchesInMemory", "invalid"); + DatabricksConnectionContext ctx = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, props); + assertEquals(3, ctx.getThriftMaxBatchesInMemory()); // Should fall back to default + } } diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionTest.java index 89cf93ab61..c632e20d23 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionTest.java @@ -61,6 +61,10 @@ public class DatabricksConnectionTest { .collect(Collectors.joining(";"))); private static final String IGNORE_TRANSACTIONS_JDBC_URL = "jdbc:databricks://sample-host.18.azuredatabricks.net:4423/default;transportMode=http;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/99999999;IgnoreTransactions=1"; + private static final String TRANSACTIONS_ENABLED_JDBC_URL = + String.format( + "jdbc:databricks://sample-host.18.azuredatabricks.net:4423/%s;transportMode=http;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/99999999;ConnCatalog=%s;ConnSchema=%s;logLevel=FATAL;IgnoreTransactions=0", + SCHEMA, CATALOG, SCHEMA); private static final ImmutableSessionInfo IMMUTABLE_SESSION_INFO = ImmutableSessionInfo.builder().computeResource(warehouse).sessionId(SESSION_ID).build(); @Mock DatabricksSdkClient databricksClient; @@ -69,11 +73,14 @@ public class DatabricksConnectionTest { private static DatabricksConnection connection; private static IDatabricksConnectionContext connectionContext; + private static IDatabricksConnectionContext transactionsEnabledContext; @BeforeAll static void setup() throws DatabricksSQLException { connectionContext = DatabricksConnectionContext.parse(CATALOG_SCHEMA_JDBC_URL, new Properties()); + transactionsEnabledContext = + DatabricksConnectionContext.parse(TRANSACTIONS_ENABLED_JDBC_URL, new Properties()); } @Test @@ -311,7 +318,7 @@ public void testGetUCVolumeClient() throws SQLException { } @Test - public void testStatement() throws DatabricksSQLException { + public void testStatement() throws SQLException { when(databricksClient.createSession( new Warehouse(WAREHOUSE_ID), CATALOG, SCHEMA, new HashMap<>())) .thenReturn(IMMUTABLE_SESSION_INFO); @@ -441,9 +448,9 @@ void testUnsupportedOperations() throws SQLException { ResultSet.CONCUR_READ_ONLY, ResultSet.CLOSE_CURSORS_AT_COMMIT); }); - assertThrows( - DatabricksSQLFeatureNotImplementedException.class, () -> connection.setSavepoint("1")); - assertThrows(DatabricksSQLFeatureNotImplementedException.class, connection::setSavepoint); + // With default IgnoreTransactions=1, savepoint methods return null (no-op) + assertNull(connection.setSavepoint("1")); + assertNull(connection.setSavepoint()); assertThrows(DatabricksSQLFeatureNotImplementedException.class, connection::createClob); assertThrows(DatabricksSQLFeatureNotImplementedException.class, connection::createBlob); assertThrows(DatabricksSQLFeatureNotImplementedException.class, connection::createNClob); @@ -454,10 +461,9 @@ void testUnsupportedOperations() throws SQLException { assertThrows( DatabricksSQLFeatureNotSupportedException.class, () -> connection.prepareStatement(SQL, new String[0])); - assertThrows( - DatabricksSQLFeatureNotImplementedException.class, () -> connection.rollback(null)); - assertThrows( - DatabricksSQLFeatureNotImplementedException.class, () -> connection.releaseSavepoint(null)); + // With default IgnoreTransactions=1, these are no-ops + assertDoesNotThrow(() -> connection.rollback(null)); + assertDoesNotThrow(() -> connection.releaseSavepoint(null)); assertThrows( DatabricksSQLFeatureNotSupportedException.class, () -> connection.setNetworkTimeout(null, 1)); @@ -504,7 +510,7 @@ void testCommonMethods() throws SQLException { } @Test - void testTranslationIsolation() throws DatabricksSQLException { + void testTranslationIsolation() throws SQLException { connection = new DatabricksConnection(connectionContext, databricksClient); connection.open(); assertDoesNotThrow( @@ -516,7 +522,7 @@ void testTranslationIsolation() throws DatabricksSQLException { } @Test - void testReadOnlyAndAbort() throws DatabricksSQLException { + void testReadOnlyAndAbort() throws SQLException { connection = new DatabricksConnection(connectionContext, databricksClient); connection.open(); assertDoesNotThrow(() -> connection.setReadOnly(false)); @@ -646,7 +652,7 @@ public void testSetAutoCommitToFalse() throws SQLException { when(databricksClient.createSession( new Warehouse(WAREHOUSE_ID), CATALOG, SCHEMA, new HashMap<>())) .thenReturn(IMMUTABLE_SESSION_INFO); - connection = new DatabricksConnection(connectionContext, databricksClient); + connection = new DatabricksConnection(transactionsEnabledContext, databricksClient); connection.open(); DatabricksConnection spyConnection = spy(connection); @@ -671,7 +677,7 @@ public void testSetAutoCommitToTrue() throws SQLException { when(databricksClient.createSession( new Warehouse(WAREHOUSE_ID), CATALOG, SCHEMA, new HashMap<>())) .thenReturn(IMMUTABLE_SESSION_INFO); - connection = new DatabricksConnection(connectionContext, databricksClient); + connection = new DatabricksConnection(transactionsEnabledContext, databricksClient); connection.open(); DatabricksConnection spyConnection = spy(connection); @@ -697,7 +703,7 @@ public void testSetAutoCommitWithServerError() throws SQLException { when(databricksClient.createSession( new Warehouse(WAREHOUSE_ID), CATALOG, SCHEMA, new HashMap<>())) .thenReturn(IMMUTABLE_SESSION_INFO); - connection = new DatabricksConnection(connectionContext, databricksClient); + connection = new DatabricksConnection(transactionsEnabledContext, databricksClient); connection.open(); DatabricksConnection spyConnection = spy(connection); @@ -744,7 +750,7 @@ public void testCommitSuccess() throws SQLException { when(databricksClient.createSession( new Warehouse(WAREHOUSE_ID), CATALOG, SCHEMA, new HashMap<>())) .thenReturn(IMMUTABLE_SESSION_INFO); - connection = new DatabricksConnection(connectionContext, databricksClient); + connection = new DatabricksConnection(transactionsEnabledContext, databricksClient); connection.open(); DatabricksConnection spyConnection = spy(connection); @@ -785,7 +791,7 @@ public void testCommitWithServerError() throws SQLException { when(databricksClient.createSession( new Warehouse(WAREHOUSE_ID), CATALOG, SCHEMA, new HashMap<>())) .thenReturn(IMMUTABLE_SESSION_INFO); - connection = new DatabricksConnection(connectionContext, databricksClient); + connection = new DatabricksConnection(transactionsEnabledContext, databricksClient); connection.open(); DatabricksConnection spyConnection = spy(connection); @@ -812,18 +818,47 @@ public void testRollbackSuccess() throws SQLException { when(databricksClient.createSession( new Warehouse(WAREHOUSE_ID), CATALOG, SCHEMA, new HashMap<>())) .thenReturn(IMMUTABLE_SESSION_INFO); - connection = new DatabricksConnection(connectionContext, databricksClient); + connection = new DatabricksConnection(transactionsEnabledContext, databricksClient); connection.open(); DatabricksConnection spyConnection = spy(connection); DatabricksStatement mockStatement = mock(DatabricksStatement.class); doReturn(mockStatement).when(spyConnection).createStatement(); + when(mockStatement.execute("SET AUTOCOMMIT = FALSE")).thenReturn(true); when(mockStatement.execute("ROLLBACK")).thenReturn(true); + // Must set autocommit to false first — rollback is not valid in auto-commit mode + spyConnection.setAutoCommit(false); spyConnection.rollback(); verify(mockStatement).execute("ROLLBACK"); - verify(mockStatement).close(); + verify(mockStatement, atLeast(2)).close(); + + spyConnection.close(); + } + + @Test + public void testRollbackThrowsWhenAutoCommitEnabled() throws SQLException { + when(databricksClient.createSession( + new Warehouse(WAREHOUSE_ID), CATALOG, SCHEMA, new HashMap<>())) + .thenReturn(IMMUTABLE_SESSION_INFO); + connection = new DatabricksConnection(transactionsEnabledContext, databricksClient); + connection.open(); + + DatabricksConnection spyConnection = spy(connection); + + // Auto-commit is true by default — rollback should throw + assertTrue(spyConnection.getAutoCommit()); + + SQLException thrown = assertThrows(SQLException.class, spyConnection::rollback); + + assertTrue( + thrown.getMessage().contains("auto-commit"), + "Exception should indicate rollback is not valid in auto-commit mode. Got: " + + thrown.getMessage()); + + // Connection should still be usable + assertFalse(spyConnection.isClosed()); spyConnection.close(); } @@ -853,12 +888,16 @@ public void testRollbackWithServerError() throws SQLException { when(databricksClient.createSession( new Warehouse(WAREHOUSE_ID), CATALOG, SCHEMA, new HashMap<>())) .thenReturn(IMMUTABLE_SESSION_INFO); - connection = new DatabricksConnection(connectionContext, databricksClient); + connection = new DatabricksConnection(transactionsEnabledContext, databricksClient); connection.open(); DatabricksConnection spyConnection = spy(connection); DatabricksStatement mockStatement = mock(DatabricksStatement.class); doReturn(mockStatement).when(spyConnection).createStatement(); + when(mockStatement.execute("SET AUTOCOMMIT = FALSE")).thenReturn(true); + + // Must set autocommit to false first — rollback is not valid in auto-commit mode + spyConnection.setAutoCommit(false); SQLException serverError = new SQLException("Unexpected rollback error", "HY000", 99999); when(mockStatement.execute("ROLLBACK")).thenThrow(serverError); @@ -869,7 +908,7 @@ public void testRollbackWithServerError() throws SQLException { assertEquals("HY000", thrown.getSQLState()); assertEquals(99999, thrown.getErrorCode()); // Vendor code preserved - verify(mockStatement).close(); + verify(mockStatement, atLeast(2)).close(); spyConnection.close(); } @@ -879,7 +918,7 @@ public void testTransactionMethodsPreserveExceptionChain() throws SQLException { when(databricksClient.createSession( new Warehouse(WAREHOUSE_ID), CATALOG, SCHEMA, new HashMap<>())) .thenReturn(IMMUTABLE_SESSION_INFO); - connection = new DatabricksConnection(connectionContext, databricksClient); + connection = new DatabricksConnection(transactionsEnabledContext, databricksClient); connection.open(); DatabricksConnection spyConnection = spy(connection); @@ -980,6 +1019,68 @@ public void testGetAutoCommitWithFetchFromServerEnabled_ReturnsFalse() throws SQ spyConnection.close(); } + @Test + public void testGetAutoCommitWithFetchFromServerEnabled_ReturnsNumeric1() throws SQLException { + // Server may return "1" instead of "true" for autocommit enabled + String urlWithFetch = CATALOG_SCHEMA_JDBC_URL + ";FetchAutoCommitFromServer=1"; + IDatabricksConnectionContext contextWithFetch = + DatabricksConnectionContext.parse(urlWithFetch, new Properties()); + + when(databricksClient.createSession( + new Warehouse(WAREHOUSE_ID), CATALOG, SCHEMA, new HashMap<>())) + .thenReturn(IMMUTABLE_SESSION_INFO); + connection = new DatabricksConnection(contextWithFetch, databricksClient); + connection.open(); + + DatabricksConnection spyConnection = spy(connection); + DatabricksStatement mockStatement = mock(DatabricksStatement.class); + ResultSet mockResultSet = mock(ResultSet.class); + + doReturn(mockStatement).when(spyConnection).createStatement(); + when(mockStatement.executeQuery("SET AUTOCOMMIT")).thenReturn(mockResultSet); + when(mockResultSet.next()).thenReturn(true); + when(mockResultSet.getString(1)).thenReturn("1"); // Server returns "1" + + boolean result = spyConnection.getAutoCommit(); + + verify(mockStatement).executeQuery("SET AUTOCOMMIT"); + assertTrue(result); + assertTrue(spyConnection.getSession().getAutoCommit()); + + spyConnection.close(); + } + + @Test + public void testGetAutoCommitWithFetchFromServerEnabled_ReturnsNumeric0() throws SQLException { + // Server may return "0" instead of "false" for autocommit disabled + String urlWithFetch = CATALOG_SCHEMA_JDBC_URL + ";FetchAutoCommitFromServer=1"; + IDatabricksConnectionContext contextWithFetch = + DatabricksConnectionContext.parse(urlWithFetch, new Properties()); + + when(databricksClient.createSession( + new Warehouse(WAREHOUSE_ID), CATALOG, SCHEMA, new HashMap<>())) + .thenReturn(IMMUTABLE_SESSION_INFO); + connection = new DatabricksConnection(contextWithFetch, databricksClient); + connection.open(); + + DatabricksConnection spyConnection = spy(connection); + DatabricksStatement mockStatement = mock(DatabricksStatement.class); + ResultSet mockResultSet = mock(ResultSet.class); + + doReturn(mockStatement).when(spyConnection).createStatement(); + when(mockStatement.executeQuery("SET AUTOCOMMIT")).thenReturn(mockResultSet); + when(mockResultSet.next()).thenReturn(true); + when(mockResultSet.getString(1)).thenReturn("0"); // Server returns "0" + + boolean result = spyConnection.getAutoCommit(); + + verify(mockStatement).executeQuery("SET AUTOCOMMIT"); + assertFalse(result); + assertFalse(spyConnection.getSession().getAutoCommit()); + + spyConnection.close(); + } + @Test public void testGetAutoCommitWithFetchFromServerEnabled_ServerQueryFails() throws SQLException { // Create connection context with FetchAutoCommitFromServer=1 @@ -1059,8 +1160,9 @@ public void testSetAutoCommitSkipsServerCallWhenAlreadyFalse() throws SQLExcepti @Test public void testSetAutoCommitDoesNotSkipWithFetchFromServerEnabled() throws SQLException { - // Create connection with FetchAutoCommitFromServer=1 - String urlWithFetch = CATALOG_SCHEMA_JDBC_URL + ";FetchAutoCommitFromServer=1"; + // Create connection with FetchAutoCommitFromServer=1 and IgnoreTransactions=0 + String urlWithFetch = + CATALOG_SCHEMA_JDBC_URL + ";FetchAutoCommitFromServer=1;IgnoreTransactions=0"; IDatabricksConnectionContext contextWithFetch = DatabricksConnectionContext.parse(urlWithFetch, new Properties()); diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java index a2bf406add..4ea945272a 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java @@ -30,6 +30,8 @@ public class DatabricksDatabaseMetaDataTest { private IDatabricksSession session; private DatabricksDatabaseMetaData metaData; private IDatabricksMetadataClient metadataClient; + private com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder + metadataResultSetBuilder; @BeforeEach public void setup() throws SQLException { @@ -49,7 +51,13 @@ public void setup() throws SQLException { .thenReturn(DatabricksConnectionContext.parse(WAREHOUSE_JDBC_URL, new Properties())); when(metadataClient.listCatalogs(any())).thenReturn(Mockito.mock(DatabricksResultSet.class)); when(metadataClient.listTableTypes(any())).thenReturn(Mockito.mock(DatabricksResultSet.class)); - when(metadataClient.listTypeInfo(any())).thenReturn(Mockito.mock(DatabricksResultSet.class)); + + // Create a real MetadataResultSetBuilder for TYPE_INFO and CLIENT_INFO_PROPERTIES + metadataResultSetBuilder = + new com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder( + session.getConnectionContext()); + when(metadataClient.listTypeInfo(any())) + .thenAnswer(invocation -> metadataResultSetBuilder.getTypeInfoResult()); when(metadataClient.listFunctions(any(), any(), any(), any())) .thenReturn(Mockito.mock(DatabricksResultSet.class)); when(metadataClient.listColumns(any(), any(), any(), any(), any())) @@ -438,7 +446,20 @@ public void getDefaultTransactionIsolation_returnsExpectedIsolationLevel() throw } @Test - public void supportsTransactions_returnsFalse() throws Exception { + public void supportsTransactions_returnsFalseByDefault() throws Exception { + // Default IgnoreTransactions=1 means transactions are ignored, so supportsTransactions returns + // false + boolean supportsTransactions = metaData.supportsTransactions(); + assertFalse(supportsTransactions); + } + + @Test + public void supportsTransactions_returnsTrueWhenTransactionsEnabled() throws Exception { + // When IgnoreTransactions=0, transactions are enabled, so supportsTransactions returns true + String urlWithTransactionsEnabled = WAREHOUSE_JDBC_URL + ";IgnoreTransactions=0"; + when(session.getConnectionContext()) + .thenReturn( + DatabricksConnectionContext.parse(urlWithTransactionsEnabled, new Properties())); boolean supportsTransactions = metaData.supportsTransactions(); assertTrue(supportsTransactions); } @@ -813,7 +834,7 @@ public void testGetDriverName() throws SQLException { @Test public void testGetDriverVersion() throws SQLException { String result = metaData.getDriverVersion(); - assertEquals("3.0.7", result); + assertEquals("3.1.1", result); } @Test @@ -825,7 +846,7 @@ public void testGetDriverMajorVersion() { @Test public void testGetDriverMinorVersion() { int result = metaData.getDriverMinorVersion(); - assertEquals(0, result); + assertEquals(1, result); } @Test @@ -1144,6 +1165,101 @@ public void testGetClientInfoProperties() throws SQLException { assertFalse(resultSet.next()); } + @Test + public void testGetTypeInfoMultipleCalls() throws SQLException { + // This test verifies issue #1178 fix: ResultSets should be usable multiple times per session + + // First call to getTypeInfo() + ResultSet resultSet1 = metaData.getTypeInfo(); + assertNotNull(resultSet1); + + // Consume all rows from first ResultSet + int rowCount1 = 0; + while (resultSet1.next()) { + rowCount1++; + } + assertTrue(rowCount1 > 0, "TYPE_INFO should have at least one row"); + + // Close first ResultSet + resultSet1.close(); + assertTrue(resultSet1.isClosed(), "ResultSet should be closed"); + + // Second call to getTypeInfo() should return a NEW, usable ResultSet + ResultSet resultSet2 = metaData.getTypeInfo(); + assertNotNull(resultSet2); + assertFalse(resultSet2.isClosed(), "Second ResultSet should not be closed"); + + // Should be able to iterate through second ResultSet + int rowCount2 = 0; + while (resultSet2.next()) { + rowCount2++; + } + assertEquals(rowCount1, rowCount2, "Both calls should return the same number of rows"); + + // Third call should also work + ResultSet resultSet3 = metaData.getTypeInfo(); + assertNotNull(resultSet3); + assertFalse(resultSet3.isClosed(), "Third ResultSet should not be closed"); + assertTrue(resultSet3.next(), "Third ResultSet should have data"); + } + + @Test + public void testGetClientInfoPropertiesMultipleCalls() throws SQLException { + // This test verifies issue #1178 fix: ResultSets should be usable multiple times per session + + // First call to getClientInfoProperties() + ResultSet resultSet1 = metaData.getClientInfoProperties(); + assertNotNull(resultSet1); + + // Consume all rows from first ResultSet + int rowCount1 = 0; + while (resultSet1.next()) { + rowCount1++; + } + assertEquals(3, rowCount1, "CLIENT_INFO_PROPERTIES should have exactly 3 rows"); + + // Close first ResultSet + resultSet1.close(); + assertTrue(resultSet1.isClosed(), "ResultSet should be closed"); + + // Second call to getClientInfoProperties() should return a NEW, usable ResultSet + ResultSet resultSet2 = metaData.getClientInfoProperties(); + assertNotNull(resultSet2); + assertFalse(resultSet2.isClosed(), "Second ResultSet should not be closed"); + + // Should be able to iterate through second ResultSet + int rowCount2 = 0; + while (resultSet2.next()) { + rowCount2++; + } + assertEquals(rowCount1, rowCount2, "Both calls should return the same number of rows"); + + // Verify non-nullable columns (per JDBC spec) - get metadata from second ResultSet + ResultSetMetaData metaData = resultSet2.getMetaData(); + assertEquals( + ResultSetMetaData.columnNoNulls, + metaData.isNullable(1), + "NAME column should be non-nullable"); + assertEquals( + ResultSetMetaData.columnNoNulls, + metaData.isNullable(2), + "MAX_LEN column should be non-nullable"); + assertEquals( + ResultSetMetaData.columnNullable, + metaData.isNullable(3), + "DEFAULT_VALUE column should be nullable"); + assertEquals( + ResultSetMetaData.columnNullable, + metaData.isNullable(4), + "DESCRIPTION column should be nullable"); + + // Third call to verify data integrity + ResultSet resultSet3 = this.metaData.getClientInfoProperties(); + assertNotNull(resultSet3); + assertTrue(resultSet3.next(), "Third ResultSet should have data"); + assertEquals("APPLICATIONNAME", resultSet3.getString(1), "First row should be APPLICATIONNAME"); + } + @Test public void testGetCrossReference() throws SQLException { ResultSet resultSet = diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java index 2fa344b3b3..dd54bb0abd 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java @@ -21,6 +21,7 @@ import com.databricks.jdbc.exception.DatabricksTemporaryRedirectException; import com.databricks.jdbc.model.client.thrift.generated.TSessionHandle; import com.databricks.jdbc.telemetry.latency.DatabricksMetricsTimedProcessor; +import java.sql.SQLException; import java.util.Properties; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -43,12 +44,12 @@ public class DatabricksSessionTest { "jdbc:databricks://sample-host.cloud.databricks.com:443/default;transportMode=http;ssl=1;httpPath=sql/protocolv1/o/9999999999999999/9999999999999999999;AuthMech=3;conncatalog=field_demos;connschema=ossjdbc"; private static IDatabricksConnectionContext connectionContext; - static void setupWarehouse(boolean useThrift) throws DatabricksSQLException { + static void setupWarehouse(boolean useThrift) throws SQLException { String url = useThrift ? WAREHOUSE_JDBC_URL : WAREHOUSE_JDBC_URL_WITH_SEA; connectionContext = DatabricksConnectionContext.parse(url, new Properties()); } - private void setupCluster() throws DatabricksSQLException { + private void setupCluster() throws SQLException { connectionContext = DatabricksConnectionContext.parse(VALID_CLUSTER_URL, new Properties()); ImmutableSessionInfo sessionInfo = ImmutableSessionInfo.builder() @@ -60,7 +61,7 @@ private void setupCluster() throws DatabricksSQLException { } @Test - public void testOpenAndCloseSession() throws DatabricksSQLException { + public void testOpenAndCloseSession() throws SQLException { setupWarehouse(true /* useThrift */); ImmutableSessionInfo sessionInfo = ImmutableSessionInfo.builder() @@ -82,7 +83,7 @@ public void testOpenAndCloseSession() throws DatabricksSQLException { } @Test - public void testOpenRedirectedThriftSession() throws DatabricksSQLException { + public void testOpenRedirectedThriftSession() throws SQLException { setupWarehouse(false /* useThrift */); ImmutableSessionInfo sessionInfo = ImmutableSessionInfo.builder() @@ -119,7 +120,7 @@ public void testOpenRedirectedThriftSession() throws DatabricksSQLException { } @Test - public void testOpenAndCloseSessionUsingThrift() throws DatabricksSQLException { + public void testOpenAndCloseSessionUsingThrift() throws SQLException { setupWarehouse(true /* useThrift */); ImmutableSessionInfo sessionInfo = ImmutableSessionInfo.builder() @@ -143,7 +144,7 @@ public void testOpenAndCloseSessionUsingThrift() throws DatabricksSQLException { } @Test - public void testOpenAndCloseSessionForAllPurposeCluster() throws DatabricksSQLException { + public void testOpenAndCloseSessionForAllPurposeCluster() throws SQLException { setupCluster(); DatabricksSession session = new DatabricksSession(connectionContext, thriftClient); assertFalse(session.isOpen()); @@ -158,7 +159,7 @@ public void testOpenAndCloseSessionForAllPurposeCluster() throws DatabricksSQLEx } @Test - public void testCloseIgnoresInvalidSession() throws DatabricksSQLException { + public void testCloseIgnoresInvalidSession() throws SQLException { setupWarehouse(true /* useThrift */); ImmutableSessionInfo sessionInfo = @@ -185,7 +186,7 @@ public void testCloseIgnoresInvalidSession() throws DatabricksSQLException { } @Test - public void testSessionConstructorForWarehouse() throws DatabricksSQLException { + public void testSessionConstructorForWarehouse() throws SQLException { DatabricksSession session = new DatabricksSession( DatabricksConnectionContext.parse(WAREHOUSE_JDBC_URL, new Properties())); @@ -202,7 +203,7 @@ public void testOpenSession_invalidWarehouseUrl() { } @Test - public void testCatalogAndSchema() throws DatabricksSQLException { + public void testCatalogAndSchema() throws SQLException { setupWarehouse(false /* useThrift */); DatabricksSession session = new DatabricksSession(connectionContext); session.setCatalog(NEW_CATALOG); @@ -213,7 +214,7 @@ public void testCatalogAndSchema() throws DatabricksSQLException { } @Test - public void testSessionToString() throws DatabricksSQLException { + public void testSessionToString() throws SQLException { setupWarehouse(false /* useThrift */); DatabricksSession session = new DatabricksSession(connectionContext); assertEquals( @@ -222,7 +223,7 @@ public void testSessionToString() throws DatabricksSQLException { } @Test - public void testSetClientInfoProperty() throws DatabricksSQLException { + public void testSetClientInfoProperty() throws SQLException { DatabricksSession session = new DatabricksSession( DatabricksConnectionContext.parse(VALID_CLUSTER_URL, new Properties()), sdkClient); @@ -231,7 +232,7 @@ public void testSetClientInfoProperty() throws DatabricksSQLException { } @Test - public void testGetClientInfoProperty_ApplicationName() throws DatabricksSQLException { + public void testGetClientInfoProperty_ApplicationName() throws SQLException { DatabricksSession session = new DatabricksSession( DatabricksConnectionContext.parse(VALID_CLUSTER_URL, new Properties()), sdkClient); @@ -241,7 +242,7 @@ public void testGetClientInfoProperty_ApplicationName() throws DatabricksSQLExce } @Test - public void testSetClientInfoProperty_AuthAccessToken() throws DatabricksSQLException { + public void testSetClientInfoProperty_AuthAccessToken() throws SQLException { DatabricksSession session = new DatabricksSession( DatabricksConnectionContext.parse(VALID_CLUSTER_URL, new Properties()), sdkClient); @@ -251,7 +252,7 @@ public void testSetClientInfoProperty_AuthAccessToken() throws DatabricksSQLExce } @Test - public void testSetClientInfoProperty_AuthAccessTokenThrift() throws DatabricksSQLException { + public void testSetClientInfoProperty_AuthAccessTokenThrift() throws SQLException { DatabricksSession session = new DatabricksSession( DatabricksConnectionContext.parse(VALID_CLUSTER_URL, new Properties()), thriftClient); @@ -263,7 +264,7 @@ public void testSetClientInfoProperty_AuthAccessTokenThrift() throws DatabricksS // ===== Lazy Client Type Integration Tests ===== @Test - public void testSessionOpensWithLazyClientType() throws DatabricksSQLException { + public void testSessionOpensWithLazyClientType() throws SQLException { // Create connection context with conditions for SEA client String url = "jdbc:databricks://sample-host.18.azuredatabricks.net:9999/default;AuthMech=3;" @@ -279,7 +280,7 @@ public void testSessionOpensWithLazyClientType() throws DatabricksSQLException { } @Test - public void testSessionOpensWithLazyClientTypeForCluster() throws DatabricksSQLException { + public void testSessionOpensWithLazyClientTypeForCluster() throws SQLException { // Create connection context for all-purpose cluster String url = "jdbc:databricks://sample-host.cloud.databricks.com:9999/default;AuthMech=3;" diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java index e2e065746a..c1fa8195a9 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java @@ -19,7 +19,10 @@ import com.databricks.sdk.service.sql.StatementState; import java.io.InputStream; import java.sql.*; +import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Properties; import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeUnit; @@ -503,196 +506,211 @@ public void testIsSimpleIdentifier() throws Exception { @Test public void testShouldReturnResultSet_SelectQuery() { String query = "-- comment\nSELECT * FROM table;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_ShowQuery() { String query = "SHOW TABLES;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_DescribeQuery() { String query = "DESCRIBE table;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_ExplainQuery() { String query = "EXPLAIN SELECT * FROM table;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_WithQuery() { String query = "WITH cte AS (SELECT * FROM table) SELECT * FROM cte;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_SetQuery() { String query = "SET @var = (SELECT COUNT(*) FROM table);"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_MapQuery() { String query = "MAP table USING some_mapping;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_FromQuery() { String query = "SELECT * FROM table;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_ValuesQuery() { String query = "VALUES (1, 2, 3);"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_UnionQuery() { String query = "SELECT * FROM table1 UNION SELECT * FROM table2;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_IntersectQuery() { String query = "SELECT * FROM table1 INTERSECT SELECT * FROM table2;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_ExceptQuery() { String query = "SELECT * FROM table1 EXCEPT SELECT * FROM table2;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_DeclareQuery() { String query = "DECLARE @var INT;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_PutQuery() { String query = "PUT some_data INTO table;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_GetQuery() { String query = "GET some_data FROM table;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_RemoveQuery() { String query = "REMOVE some_data FROM table;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_ListQuery() { String query = "LIST TABLES;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_UpdateQuery() { String query = "UPDATE table SET column = value;"; - assertFalse(DatabricksStatement.shouldReturnResultSet(query)); + // Without prefix configuration, UPDATE should not return result set + assertFalse(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); + // With UPDATE prefix, it should return result set + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Arrays.asList("UPDATE"))); } @Test public void testShouldReturnResultSet_DeleteQuery() { String query = "DELETE FROM table WHERE condition;"; - assertFalse(DatabricksStatement.shouldReturnResultSet(query)); + // Without prefix configuration, DELETE should not return result set + assertFalse(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); + // With DELETE prefix, it should return result set + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Arrays.asList("DELETE"))); } @Test public void testShouldReturnResultSet_SingleLineCommentAtStart() { String query = "-- This is a comment\nSELECT * FROM table;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_SingleLineCommentAtEnd() { String query = "SELECT * FROM table; -- This is a comment"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_SingleLineCommentInMiddle() { String query = "SELECT * FROM table -- This is a comment\nWHERE id = 1;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_MultiLineCommentAtStart() { String query = "/* This is a comment */ SELECT * FROM table;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_MultiLineCommentAtEnd() { String query = "SELECT * FROM table; /* This is a comment */"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_MultiLineCommentInMiddle() { String query = "SELECT * FROM table /* This is a comment */ WHERE id = 1;"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_MultipleSingleLineComments() { String query = "-- Comment 1\nSELECT * FROM table; -- Comment 2\n-- Comment 3"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_MultipleMultiLineComments() { String query = "/* Comment 1 */ SELECT * FROM table; /* Comment 2 */ /* Comment 3 */"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_SingleAndMultiLineComments() { String query = "-- Single-line comment\nSELECT * FROM table; /* Multi-line comment */"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_CommentSurroundingQuery() { String query = "-- Single-line comment\n/* Multi-line comment */ SELECT * FROM table; /* Another comment */ -- End comment"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); } @Test public void testShouldReturnResultSet_CallStatement() { String query = "-- Single-line comment\n/* Multi-line comment */ CALL send_notifications(12); /* Another comment */ -- End comment"; - assertTrue(DatabricksStatement.shouldReturnResultSet(query)); - assertTrue(DatabricksStatement.shouldReturnResultSet("CALL send_notifications(12);")); + assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList())); + assertTrue( + DatabricksStatement.shouldReturnResultSet( + "CALL send_notifications(12);", Collections.emptyList())); } @Test public void testShouldReturnResultSet_StartWithBegin() { - assertTrue(DatabricksStatement.shouldReturnResultSet("BEGIN")); - assertTrue(DatabricksStatement.shouldReturnResultSet(" begin ")); - assertTrue(DatabricksStatement.shouldReturnResultSet("BEGIN; WORK; END")); - assertTrue(DatabricksStatement.shouldReturnResultSet("BEGIN; SELECT 1")); + assertTrue(DatabricksStatement.shouldReturnResultSet("BEGIN", Collections.emptyList())); + assertTrue(DatabricksStatement.shouldReturnResultSet(" begin ", Collections.emptyList())); + assertTrue( + DatabricksStatement.shouldReturnResultSet("BEGIN; WORK; END", Collections.emptyList())); + assertTrue( + DatabricksStatement.shouldReturnResultSet("BEGIN; SELECT 1", Collections.emptyList())); // Not supporting for transaction statements - assertFalse(DatabricksStatement.shouldReturnResultSet("BEGIN TRANSACTION")); - assertFalse(DatabricksStatement.shouldReturnResultSet(" BeGiN TRANSACTION")); - assertFalse(DatabricksStatement.shouldReturnResultSet("BEGIN transaction ")); + assertFalse( + DatabricksStatement.shouldReturnResultSet("BEGIN TRANSACTION", Collections.emptyList())); + assertFalse( + DatabricksStatement.shouldReturnResultSet( + " BeGiN TRANSACTION", Collections.emptyList())); + assertFalse( + DatabricksStatement.shouldReturnResultSet( + "BEGIN transaction ", Collections.emptyList())); } @Test @@ -1051,4 +1069,128 @@ public void testJdbcStopCondition() throws Exception { // After loop, should be able to call getLargeUpdateCount() without exception assertEquals(-1, statement.getLargeUpdateCount()); } + + @Test + public void testShouldReturnResultSet_WithNonRowcountQueryPrefixes_Insert() { + List prefixes = Arrays.asList("INSERT", "UPDATE", "DELETE"); + String query = "INSERT INTO table VALUES (1, 2, 3)"; + assertTrue(DatabricksStatement.shouldReturnResultSet(query, prefixes)); + } + + @Test + public void testShouldReturnResultSet_WithNonRowcountQueryPrefixes_Update() { + List prefixes = Arrays.asList("INSERT", "UPDATE", "DELETE"); + String query = "UPDATE table SET col = 1 WHERE id = 2"; + assertTrue(DatabricksStatement.shouldReturnResultSet(query, prefixes)); + } + + @Test + public void testShouldReturnResultSet_WithNonRowcountQueryPrefixes_Delete() { + List prefixes = Arrays.asList("INSERT", "UPDATE", "DELETE"); + String query = "DELETE FROM table WHERE id = 1"; + assertTrue(DatabricksStatement.shouldReturnResultSet(query, prefixes)); + } + + @Test + public void testShouldReturnResultSet_WithNonRowcountQueryPrefixes_Merge() { + List prefixes = Arrays.asList("MERGE"); + String query = "MERGE INTO target USING source ON target.id = source.id"; + assertTrue(DatabricksStatement.shouldReturnResultSet(query, prefixes)); + } + + @Test + public void testShouldReturnResultSet_WithNonRowcountQueryPrefixes_CaseInsensitive() { + List prefixes = Arrays.asList("INSERT", "UPDATE"); + String query = "insert into table values (1, 2, 3)"; + assertTrue(DatabricksStatement.shouldReturnResultSet(query, prefixes)); + } + + @Test + public void testShouldReturnResultSet_WithNonRowcountQueryPrefixes_WithComments() { + List prefixes = Arrays.asList("INSERT", "UPDATE", "DELETE"); + String query = "-- Comment\n/* Block comment */ INSERT INTO table VALUES (1, 2, 3)"; + assertTrue(DatabricksStatement.shouldReturnResultSet(query, prefixes)); + } + + @Test + public void testShouldReturnResultSet_WithNonRowcountQueryPrefixes_NoMatch() { + List prefixes = Arrays.asList("INSERT", "UPDATE", "DELETE"); + String query = "CREATE TABLE test (id INT)"; + assertFalse(DatabricksStatement.shouldReturnResultSet(query, prefixes)); + } + + @Test + public void testShouldReturnResultSet_WithNonRowcountQueryPrefixes_EmptyList() { + List prefixes = Collections.emptyList(); + String query = "INSERT INTO table VALUES (1, 2, 3)"; + assertFalse(DatabricksStatement.shouldReturnResultSet(query, prefixes)); + } + + @Test + public void testShouldReturnResultSet_WithNonRowcountQueryPrefixes_SelectStillWorks() { + List prefixes = Arrays.asList("INSERT", "UPDATE", "DELETE"); + String query = "SELECT * FROM table"; + assertTrue(DatabricksStatement.shouldReturnResultSet(query, prefixes)); + } + + @Test + public void testExecuteInsertWithNonRowcountQueryPrefixes() throws Exception { + // Create connection with NonRowcountQueryPrefixes=INSERT + String jdbcUrlWithPrefix = JDBC_URL + "NonRowcountQueryPrefixes=INSERT"; + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(jdbcUrlWithPrefix, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksStatement statement = new DatabricksStatement(connection); + + String insertStatement = "INSERT INTO table VALUES (1, 2, 3)"; + + when(client.executeStatement( + eq(insertStatement), + eq(new Warehouse(WAREHOUSE_ID)), + eq(new HashMap<>()), + eq(StatementType.SQL), + any(IDatabricksSession.class), + eq(statement))) + .thenReturn(resultSet); + + // Execute INSERT statement + boolean hasResultSet = statement.execute(insertStatement); + + // With NonRowcountQueryPrefixes=INSERT, execute() should return true (has result set) + assertTrue(hasResultSet, "INSERT with NonRowcountQueryPrefixes=INSERT should return true"); + assertNotNull(statement.getResultSet(), "ResultSet should be available"); + assertEquals( + -1, statement.getUpdateCount(), "Update count should be -1 when result set is returned"); + + statement.close(); + } + + @Test + public void testExecuteInsertWithoutNonRowcountQueryPrefixes() throws Exception { + // Create connection without NonRowcountQueryPrefixes + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksStatement statement = new DatabricksStatement(connection); + + String insertStatement = "INSERT INTO table VALUES (1, 2, 3)"; + + when(client.executeStatement( + eq(insertStatement), + eq(new Warehouse(WAREHOUSE_ID)), + eq(new HashMap<>()), + eq(StatementType.SQL), + any(IDatabricksSession.class), + eq(statement))) + .thenReturn(resultSet); + + // Execute INSERT statement + boolean hasResultSet = statement.execute(insertStatement); + + // Without NonRowcountQueryPrefixes, execute() should return false (update count, not result + // set) + assertFalse(hasResultSet, "INSERT without NonRowcountQueryPrefixes should return false"); + + statement.close(); + } } diff --git a/src/test/java/com/databricks/jdbc/api/impl/ExecutionResultFactoryTest.java b/src/test/java/com/databricks/jdbc/api/impl/ExecutionResultFactoryTest.java index 2efb1e33a6..fce9805a1a 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/ExecutionResultFactoryTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/ExecutionResultFactoryTest.java @@ -1,22 +1,33 @@ package com.databricks.jdbc.api.impl; -import static com.databricks.jdbc.TestConstants.ARROW_BATCH_LIST; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.Mockito.when; import com.databricks.jdbc.api.impl.arrow.ArrowStreamResult; +import com.databricks.jdbc.api.impl.arrow.LazyThriftInlineArrowResult; +import com.databricks.jdbc.api.impl.arrow.StreamingInlineArrowResult; +import com.databricks.jdbc.api.impl.thrift.StreamingColumnarResult; import com.databricks.jdbc.api.impl.volume.VolumeOperationResult; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.api.internal.IDatabricksStatementInternal; import com.databricks.jdbc.dbclient.impl.common.StatementId; -import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.exception.DatabricksSQLFeatureNotSupportedException; import com.databricks.jdbc.model.client.thrift.generated.*; import com.databricks.jdbc.model.core.ResultData; import com.databricks.jdbc.model.core.ResultManifest; import com.databricks.jdbc.model.core.ResultSchema; import com.databricks.sdk.service.sql.Format; +import java.io.ByteArrayOutputStream; import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowStreamWriter; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -34,7 +45,7 @@ public class ExecutionResultFactoryTest { @Mock TFetchResultsResp fetchResultsResp; @Test - public void testGetResultSet_jsonInline() throws DatabricksSQLException { + public void testGetResultSet_jsonInline() throws SQLException { ResultManifest manifest = new ResultManifest(); manifest.setFormat(Format.JSON_ARRAY); ResultData data = new ResultData(); @@ -45,7 +56,7 @@ public void testGetResultSet_jsonInline() throws DatabricksSQLException { } @Test - public void testGetResultSet_externalLink() throws DatabricksSQLException { + public void testGetResultSet_externalLink() throws SQLException { when(connectionContext.getConnectionUuid()).thenReturn("sample-uuid"); when(connectionContext.getHttpMaxConnectionsPerRoute()).thenReturn(100); when(session.getConnectionContext()).thenReturn(connectionContext); @@ -63,7 +74,7 @@ public void testGetResultSet_externalLink() throws DatabricksSQLException { } @Test - public void testGetResultSet_volumeOperation() throws DatabricksSQLException { + public void testGetResultSet_volumeOperation() throws SQLException { when(connectionContext.getConnectionUuid()).thenReturn("sample-uuid"); when(session.getConnectionContext()).thenReturn(connectionContext); ResultData data = new ResultData(); @@ -99,6 +110,8 @@ public void testGetResultSet_volumeOperationThriftResp() throws Exception { public void testGetResultSet_thriftColumnar() throws SQLException { when(resultSetMetadataResp.getResultFormat()).thenReturn(TSparkRowSetType.COLUMN_BASED_SET); when(fetchResultsResp.getResultSetMetadata()).thenReturn(resultSetMetadataResp); + when(session.getConnectionContext()).thenReturn(connectionContext); + when(connectionContext.isInlineStreamingEnabled()).thenReturn(false); IExecutionResult result = ExecutionResultFactory.getResultSet(fetchResultsResp, session, parentStatement); assertInstanceOf(LazyThriftResult.class, result); @@ -128,14 +141,136 @@ public void testGetResultSet_thriftURL() throws SQLException { @Test public void testGetResultSet_thriftInlineArrow() throws SQLException { - when(connectionContext.getConnectionUuid()).thenReturn("sample-uuid"); when(resultSetMetadataResp.getResultFormat()).thenReturn(TSparkRowSetType.ARROW_BASED_SET); when(fetchResultsResp.getResultSetMetadata()).thenReturn(resultSetMetadataResp); when(fetchResultsResp.getResults()).thenReturn(tRowSet); when(session.getConnectionContext()).thenReturn(connectionContext); - when(tRowSet.getArrowBatches()).thenReturn(ARROW_BATCH_LIST); + when(connectionContext.isInlineStreamingEnabled()).thenReturn(false); IExecutionResult result = ExecutionResultFactory.getResultSet(fetchResultsResp, session, parentStatement); - assertInstanceOf(ArrowStreamResult.class, result); + assertInstanceOf(LazyThriftInlineArrowResult.class, result); + } + + @Test + public void testGetResultSet_thriftColumnarWithStreamingEnabled() throws SQLException { + // Create real TFetchResultsResp with valid columnar data and hasMoreRows=false + // so the prefetch thread exits immediately + TFetchResultsResp realFetchResp = createColumnarFetchResponse(); + + when(session.getConnectionContext()).thenReturn(connectionContext); + when(connectionContext.isInlineStreamingEnabled()).thenReturn(true); + when(connectionContext.getThriftMaxBatchesInMemory()).thenReturn(3); + + IExecutionResult result = + ExecutionResultFactory.getResultSet(realFetchResp, session, parentStatement); + + assertInstanceOf(StreamingColumnarResult.class, result); + result.close(); // Clean up prefetch thread + } + + @Test + public void testGetResultSet_thriftInlineArrowWithStreamingEnabled() throws SQLException { + // Create real TFetchResultsResp with valid Arrow data and hasMoreRows=false + // so the prefetch thread exits immediately + TFetchResultsResp realFetchResp = createArrowFetchResponse(); + + when(session.getConnectionContext()).thenReturn(connectionContext); + when(connectionContext.isInlineStreamingEnabled()).thenReturn(true); + when(connectionContext.getThriftMaxBatchesInMemory()).thenReturn(3); + when(parentStatement.getStatementId()).thenReturn(STATEMENT_ID); + + IExecutionResult result = + ExecutionResultFactory.getResultSet(realFetchResp, session, parentStatement); + + assertInstanceOf(StreamingInlineArrowResult.class, result); + result.close(); // Clean up prefetch thread + } + + // ==================== Helper Methods ==================== + + /** Creates a valid TFetchResultsResp with columnar data for streaming tests. */ + private TFetchResultsResp createColumnarFetchResponse() { + TFetchResultsResp response = new TFetchResultsResp(); + response.hasMoreRows = false; // Prefetch thread will exit immediately + + // Create metadata + TGetResultSetMetadataResp metadata = new TGetResultSetMetadataResp(); + metadata.setResultFormat(TSparkRowSetType.COLUMN_BASED_SET); + response.setResultSetMetadata(metadata); + + // Create columnar data with one row + TRowSet rowSet = new TRowSet(); + TColumn column = new TColumn(); + TStringColumn stringCol = new TStringColumn(); + stringCol.setValues(Arrays.asList("test_value")); + column.setStringVal(stringCol); + rowSet.setColumns(Collections.singletonList(column)); + response.setResults(rowSet); + + return response; + } + + /** Creates a valid TFetchResultsResp with Arrow data for streaming tests. */ + private TFetchResultsResp createArrowFetchResponse() { + TFetchResultsResp response = new TFetchResultsResp(); + response.hasMoreRows = false; // Prefetch thread will exit immediately + + // Create Arrow batch data + byte[] arrowData = createArrowBytes(); + + // Create metadata with schema + TGetResultSetMetadataResp metadata = new TGetResultSetMetadataResp(); + metadata.setResultFormat(TSparkRowSetType.ARROW_BASED_SET); + metadata.setArrowSchema(new byte[0]); // Empty - schema is in the IPC stream + metadata.setSchema(createTableSchema()); + response.setResultSetMetadata(metadata); + + // Create Arrow batch + TSparkArrowBatch arrowBatch = new TSparkArrowBatch(); + arrowBatch.setRowCount(1); + arrowBatch.setBatch(arrowData); + + TRowSet rowSet = new TRowSet(); + rowSet.setArrowBatches(Collections.singletonList(arrowBatch)); + response.setResults(rowSet); + + return response; + } + + /** Creates valid Arrow IPC format bytes with one row. */ + private byte[] createArrowBytes() { + try (BufferAllocator allocator = new RootAllocator(); + ByteArrayOutputStream out = new ByteArrayOutputStream(); + IntVector intVector = new IntVector("col_0", allocator)) { + + intVector.allocateNew(1); + intVector.set(0, 42); + intVector.setValueCount(1); + + try (VectorSchemaRoot root = VectorSchemaRoot.of(intVector); + ArrowStreamWriter writer = new ArrowStreamWriter(root, null, out)) { + writer.start(); + root.setRowCount(1); + writer.writeBatch(); + writer.end(); + } + + return out.toByteArray(); + } catch (Exception e) { + throw new RuntimeException("Failed to create Arrow data", e); + } + } + + /** Creates a TTableSchema for Arrow result metadata. */ + private TTableSchema createTableSchema() { + List columns = new ArrayList<>(); + TPrimitiveTypeEntry primitiveType = new TPrimitiveTypeEntry().setType(TTypeId.INT_TYPE); + TTypeEntry typeEntry = new TTypeEntry(); + typeEntry.setPrimitiveEntry(primitiveType); + TTypeDesc typeDesc = new TTypeDesc().setTypes(Collections.singletonList(typeEntry)); + TColumnDesc columnDesc = + new TColumnDesc().setColumnName("col_0").setTypeDesc(typeDesc).setPosition(0); + columns.add(columnDesc); + return new TTableSchema().setColumns(columns); } } diff --git a/src/test/java/com/databricks/jdbc/api/impl/LazyThriftResultTest.java b/src/test/java/com/databricks/jdbc/api/impl/LazyThriftResultTest.java index 5fee41b547..523a0fae6b 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/LazyThriftResultTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/LazyThriftResultTest.java @@ -10,6 +10,7 @@ import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.model.client.thrift.generated.*; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.sql.SQLException; import java.util.*; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -25,14 +26,14 @@ public class LazyThriftResultTest { @Mock private IDatabricksClient databricksClient; @BeforeEach - void setUp() throws DatabricksSQLException { + void setUp() throws SQLException { // Lenient stubbing to avoid unnecessary strictness lenient().when(session.getDatabricksClient()).thenReturn(databricksClient); lenient().when(statement.getMaxRows()).thenReturn(0); // No limit by default } @Test - void testConstructorWithEmptyResult() throws DatabricksSQLException { + void testConstructorWithEmptyResult() throws SQLException { TFetchResultsResp emptyResponse = createEmptyResponse(false); LazyThriftResult result = new LazyThriftResult(emptyResponse, statement, session); @@ -46,7 +47,7 @@ void testConstructorWithEmptyResult() throws DatabricksSQLException { } @Test - void testBasicIteration() throws DatabricksSQLException { + void testBasicIteration() throws SQLException { TFetchResultsResp response = createResponseWithStringData( Arrays.asList("row1_col1", "row1_col2"), @@ -81,7 +82,7 @@ void testBasicIteration() throws DatabricksSQLException { } @Test - void testLazyFetching() throws DatabricksSQLException { + void testLazyFetching() throws SQLException { // Setup first batch TFetchResultsResp firstBatch = createResponseWithStringData( @@ -126,7 +127,7 @@ void testLazyFetching() throws DatabricksSQLException { } @Test - void testEmptyBatchHandling() throws DatabricksSQLException { + void testEmptyBatchHandling() throws SQLException { // Setup first batch (empty) TFetchResultsResp emptyBatch = createEmptyResponse(true); // hasMoreRows = true @@ -151,7 +152,7 @@ void testEmptyBatchHandling() throws DatabricksSQLException { } @Test - void testMultipleEmptyBatches() throws DatabricksSQLException { + void testMultipleEmptyBatches() throws SQLException { // Setup first batch (empty) TFetchResultsResp emptyBatch1 = createEmptyResponse(true); @@ -174,7 +175,7 @@ void testMultipleEmptyBatches() throws DatabricksSQLException { } @Test - void testMaxRowsLimit() throws DatabricksSQLException { + void testMaxRowsLimit() throws SQLException { when(statement.getMaxRows()).thenReturn(2); // Limit to 2 rows TFetchResultsResp response = @@ -202,7 +203,7 @@ void testMaxRowsLimit() throws DatabricksSQLException { } @Test - void testMaxRowsLimitAcrossBatches() throws DatabricksSQLException { + void testMaxRowsLimitAcrossBatches() throws SQLException { when(statement.getMaxRows()).thenReturn(3); // Limit to 3 rows // First batch has 2 rows @@ -238,7 +239,7 @@ void testMaxRowsLimitAcrossBatches() throws DatabricksSQLException { } @Test - void testErrorHandlingDuringFetch() throws DatabricksSQLException { + void testErrorHandlingDuringFetch() throws SQLException { TFetchResultsResp firstBatch = createResponseWithStringData( Arrays.asList("row1_col1", "row1_col2"), true); // hasMoreRows = true @@ -259,7 +260,7 @@ void testErrorHandlingDuringFetch() throws DatabricksSQLException { } @Test - void testClosedResultAccess() throws DatabricksSQLException { + void testClosedResultAccess() throws SQLException { TFetchResultsResp response = createResponseWithStringData(Arrays.asList("row1_col1", "row1_col2"), false); @@ -274,7 +275,7 @@ void testClosedResultAccess() throws DatabricksSQLException { } @Test - void testInvalidColumnIndex() throws DatabricksSQLException { + void testInvalidColumnIndex() throws SQLException { TFetchResultsResp response = createResponseWithStringData( Arrays.asList("col1", "col2"), // 2 columns @@ -293,7 +294,7 @@ void testInvalidColumnIndex() throws DatabricksSQLException { } @Test - void testAccessBeforeFirstRow() throws DatabricksSQLException { + void testAccessBeforeFirstRow() throws SQLException { TFetchResultsResp response = createResponseWithStringData(Arrays.asList("row1_col1", "row1_col2"), false); @@ -306,7 +307,7 @@ void testAccessBeforeFirstRow() throws DatabricksSQLException { } @Test - void testNullHandling() throws DatabricksSQLException { + void testNullHandling() throws SQLException { TFetchResultsResp response = createResponseWithNulls(); LazyThriftResult result = new LazyThriftResult(response, statement, session); @@ -317,7 +318,7 @@ void testNullHandling() throws DatabricksSQLException { } @Test - void testChunkCount() throws DatabricksSQLException { + void testChunkCount() throws SQLException { TFetchResultsResp response = createEmptyResponse(false); LazyThriftResult result = new LazyThriftResult(response, statement, session); diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunkStatusTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunkStatusTest.java index 4394b9f9e2..bfc8977220 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunkStatusTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunkStatusTest.java @@ -9,7 +9,7 @@ import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksParsingException; import com.databricks.jdbc.model.core.ExternalLink; -import com.databricks.jdbc.telemetry.latency.TelemetryCollector; +import com.databricks.jdbc.telemetry.latency.TelemetryCollectorManager; import java.io.ByteArrayInputStream; import java.io.FilterInputStream; import java.io.IOException; @@ -39,7 +39,7 @@ public class ArrowResultChunkStatusTest { @AfterEach void clearTelemetry() { - TelemetryCollector.getInstance().exportAllPendingTelemetryDetails(); + TelemetryCollectorManager.getInstance().clear(); } @Test diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunkTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunkTest.java index 9d50b498ab..38db00f44b 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunkTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunkTest.java @@ -16,7 +16,9 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.Random; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.*; @@ -261,4 +263,56 @@ public void testEmptyRecordBatches() throws DatabricksSQLException { 10, iterator.getColumnObjectAtCurrentRow(0, ColumnInfoTypeName.INT, "INT", intColumnInfo)); assertFalse(iterator.hasNextRow()); } + + @Test + public void testMetadataExtractionWithZeroRows() throws Exception { + // Arrange - Create schema with Arrow metadata + // This test verifies that VectorSchemaRoot metadata is available even when there are 0 rows + Map metadata1 = new HashMap<>(); + metadata1.put("Spark:DataType:SqlName", "ARRAY"); + FieldType fieldType1 = new FieldType(false, Types.MinorType.INT.getType(), null, metadata1); + + Map metadata2 = new HashMap<>(); + metadata2.put("Spark:DataType:SqlName", "MAP"); + FieldType fieldType2 = new FieldType(false, Types.MinorType.INT.getType(), null, metadata2); + + List fieldList = new ArrayList<>(); + fieldList.add(new Field("col1", fieldType1, null)); + fieldList.add(new Field("col2", fieldType2, null)); + Schema schema = new Schema(fieldList); + + // Create Arrow file with 0 rows + Object[][] emptyData = new Object[2][0]; // 2 columns, 0 rows + File arrowFile = + createTestArrowFile( + "TestZeroRowsMetadata", schema, emptyData, new RootAllocator(Integer.MAX_VALUE)); + + // Create chunk info for 0 rows + BaseChunkInfo chunkInfo = + new BaseChunkInfo().setChunkIndex(0L).setByteCount(200L).setRowOffset(0L).setRowCount(0L); + + ArrowResultChunk arrowResultChunk = + ArrowResultChunk.builder() + .withStatementId(TEST_STATEMENT_ID) + .withChunkInfo(chunkInfo) + .withChunkStatus(ChunkStatus.PROCESSING_SUCCEEDED) + .build(); + + // Act + arrowResultChunk.initializeData(new FileInputStream(arrowFile)); + + // Assert - Metadata should be available even with 0 rows + List metadata = arrowResultChunk.getArrowMetadata(); + assertNotNull(metadata, "Metadata should not be null even with 0 rows"); + assertEquals(2, metadata.size(), "Should have metadata for 2 columns"); + assertEquals("ARRAY", metadata.get(0), "First column metadata should be ARRAY"); + assertEquals( + "MAP", + metadata.get(1), + "Second column metadata should be MAP"); + + // Cleanup + arrowResultChunk.releaseChunk(); + arrowFile.delete(); + } } diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResultTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResultTest.java index 6e32861fec..3ede51f430 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResultTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResultTest.java @@ -133,25 +133,6 @@ public void testIteration() throws Exception { assertFalse(result.next()); } - @Test - public void testInlineArrow() throws DatabricksSQLException { - IDatabricksConnectionContext connectionContext = - DatabricksConnectionContextFactory.create(JDBC_URL, new Properties()); - when(session.getConnectionContext()).thenReturn(connectionContext); - when(metadataResp.getSchema()).thenReturn(TEST_TABLE_SCHEMA); - when(fetchResultsResp.getResults()).thenReturn(resultData); - when(fetchResultsResp.getResultSetMetadata()).thenReturn(metadataResp); - ArrowStreamResult result = - new ArrowStreamResult(fetchResultsResp, true, parentStatement, session); - assertEquals(-1, result.getCurrentRow()); - assertTrue(result.hasNext()); - assertFalse(result.next()); - assertEquals(0, result.getCurrentRow()); - assertFalse(result.hasNext()); - assertDoesNotThrow(result::close); - assertFalse(result.hasNext()); - } - @Test public void testCloudFetchArrow() throws Exception { IDatabricksConnectionContext connectionContext = @@ -164,7 +145,7 @@ public void testCloudFetchArrow() throws Exception { when(fetchResultsResp.getResultSetMetadata()).thenReturn(metadataResp); when(parentStatement.getStatementId()).thenReturn(STATEMENT_ID); ArrowStreamResult result = - new ArrowStreamResult(fetchResultsResp, false, parentStatement, session, mockHttpClient); + new ArrowStreamResult(fetchResultsResp, parentStatement, session, mockHttpClient); assertEquals(-1, result.getCurrentRow()); assertTrue(result.hasNext()); assertDoesNotThrow(result::close); @@ -586,7 +567,7 @@ public void testStreamingChunkProviderEnabledForThriftResult() throws Exception when(mockHttpClient.execute(isA(HttpUriRequest.class), eq(true))).thenReturn(httpResponse); ArrowStreamResult result = - new ArrowStreamResult(fetchResultsResp, false, parentStatement, session, mockHttpClient); + new ArrowStreamResult(fetchResultsResp, parentStatement, session, mockHttpClient); // Verify result was created successfully with StreamingChunkProvider for Thrift assertNotNull(result); diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadServiceTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadServiceTest.java index 81f61ac4cc..9651790c78 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadServiceTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadServiceTest.java @@ -14,6 +14,7 @@ import com.databricks.jdbc.model.core.ChunkLinkFetchResult; import com.databricks.jdbc.model.core.ExternalLink; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.sql.SQLException; import java.time.Instant; import java.time.temporal.ChronoUnit; import java.util.*; @@ -53,7 +54,7 @@ void setUp() { @Test void testGetLinkForChunk_Success() - throws DatabricksSQLException, InterruptedException, ExecutionException, TimeoutException { + throws SQLException, InterruptedException, ExecutionException, TimeoutException { when(mockSession.getDatabricksClient()).thenReturn(mockClient); // Mock the response to link requests @@ -105,7 +106,7 @@ void testGetLinkForChunk_InvalidIndex() throws ExecutionException, InterruptedEx @Test void testGetLinkForChunk_ClientError() - throws DatabricksSQLException, ExecutionException, InterruptedException { + throws SQLException, ExecutionException, InterruptedException { long chunkIndex = 1L; DatabricksSQLException expectedError = new DatabricksSQLException("Test error", DatabricksDriverErrorCode.INVALID_STATE); @@ -127,7 +128,7 @@ void testGetLinkForChunk_ClientError() } @Test - void testAutoTriggerForSEAClient() throws DatabricksSQLException, InterruptedException { + void testAutoTriggerForSEAClient() throws SQLException, InterruptedException { when(mockSession.getDatabricksClient()).thenReturn(mockClient); // Mock the response to link requests when(mockClient.getResultChunks(eq(mockStatementId), eq(1L), anyLong())) @@ -149,7 +150,7 @@ void testAutoTriggerForSEAClient() throws DatabricksSQLException, InterruptedExc @Test void testHandleExpiredLinks() - throws DatabricksSQLException, ExecutionException, InterruptedException, TimeoutException { + throws SQLException, ExecutionException, InterruptedException, TimeoutException { when(mockSession.getConnectionContext().getClientType()).thenReturn(DatabricksClientType.SEA); // Create an expired link for chunk index 1 ExternalLink expiredLinkForChunkIndex_1 = @@ -189,7 +190,7 @@ void testHandleExpiredLinks() @Test void testBatchDownloadChaining() - throws DatabricksSQLException, ExecutionException, InterruptedException, TimeoutException { + throws SQLException, ExecutionException, InterruptedException, TimeoutException { // Use a far future date to ensure links are never considered expired String farFutureExpiration = Instant.now().plus(10, ChronoUnit.MINUTES).toString(); diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkFetcherTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkFetcherTest.java index 3a71bd1f60..2990ec820d 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkFetcherTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkFetcherTest.java @@ -9,6 +9,7 @@ import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.model.core.ChunkLinkFetchResult; import com.databricks.jdbc.model.core.ExternalLink; +import java.sql.SQLException; import java.util.Collections; import java.util.List; import org.junit.jupiter.api.BeforeEach; @@ -64,7 +65,7 @@ void setUp() { @Test @DisplayName("fetchLinks should delegate to client with chunkIndex") - void testFetchLinksDelegatesToClient() throws DatabricksSQLException { + void testFetchLinksDelegatesToClient() throws SQLException { ExternalLink link = createLink(5, 500, 100); ChunkLinkFetchResult expectedResult = ChunkLinkFetchResult.of(List.of(link), true, 6, 600); @@ -78,7 +79,7 @@ void testFetchLinksDelegatesToClient() throws DatabricksSQLException { @Test @DisplayName("refetchLink should return matching link from result") - void testRefetchLinkReturnsMatchingLink() throws DatabricksSQLException { + void testRefetchLinkReturnsMatchingLink() throws SQLException { ExternalLink link5 = createLink(5, 500, 100); ExternalLink link6 = createLink(6, 600, 100); ChunkLinkFetchResult result = ChunkLinkFetchResult.of(List.of(link5, link6), true, 7, 700); @@ -92,7 +93,7 @@ void testRefetchLinkReturnsMatchingLink() throws DatabricksSQLException { @Test @DisplayName("refetchLink should throw when result is end of stream") - void testRefetchLinkThrowsOnEndOfStream() throws DatabricksSQLException { + void testRefetchLinkThrowsOnEndOfStream() throws SQLException { when(mockClient.getResultChunks(STATEMENT_ID, 5L, 500L)) .thenReturn(ChunkLinkFetchResult.endOfStream()); @@ -104,7 +105,7 @@ void testRefetchLinkThrowsOnEndOfStream() throws DatabricksSQLException { @Test @DisplayName("refetchLink should throw when no matching link found") - void testRefetchLinkThrowsWhenNoMatch() throws DatabricksSQLException { + void testRefetchLinkThrowsWhenNoMatch() throws SQLException { // Return links for chunk 6 and 7, but we're looking for chunk 5 ExternalLink link6 = createLink(6, 600, 100); ExternalLink link7 = createLink(7, 700, 100); @@ -140,7 +141,7 @@ void setUp() { @Test @DisplayName("fetchLinks should delegate to client with rowOffset") - void testFetchLinksDelegatesToClient() throws DatabricksSQLException { + void testFetchLinksDelegatesToClient() throws SQLException { ExternalLink link = createLink(5, 500, 100); ChunkLinkFetchResult expectedResult = ChunkLinkFetchResult.of(List.of(link), true, 6, 600); @@ -154,7 +155,7 @@ void testFetchLinksDelegatesToClient() throws DatabricksSQLException { @Test @DisplayName("refetchLink should return matching link from result") - void testRefetchLinkReturnsMatchingLink() throws DatabricksSQLException { + void testRefetchLinkReturnsMatchingLink() throws SQLException { ExternalLink link5 = createLink(5, 500, 100); ExternalLink link6 = createLink(6, 600, 100); ChunkLinkFetchResult result = ChunkLinkFetchResult.of(List.of(link5, link6), true, 7, 700); @@ -168,7 +169,7 @@ void testRefetchLinkReturnsMatchingLink() throws DatabricksSQLException { @Test @DisplayName("refetchLink should retry when hasMore=true but no links") - void testRefetchLinkRetriesWhenHasMoreButNoLinks() throws DatabricksSQLException { + void testRefetchLinkRetriesWhenHasMoreButNoLinks() throws SQLException { // First two calls return empty with hasMore=true, third call returns the link ChunkLinkFetchResult emptyWithMore = ChunkLinkFetchResult.of(Collections.emptyList(), true, 5, 500); @@ -188,7 +189,7 @@ void testRefetchLinkRetriesWhenHasMoreButNoLinks() throws DatabricksSQLException @Test @DisplayName("refetchLink should throw when hasMore=false and no links") - void testRefetchLinkThrowsWhenNoMoreAndNoLinks() throws DatabricksSQLException { + void testRefetchLinkThrowsWhenNoMoreAndNoLinks() throws SQLException { ChunkLinkFetchResult emptyNoMore = ChunkLinkFetchResult.of(Collections.emptyList(), false, -1, 500); @@ -204,7 +205,7 @@ void testRefetchLinkThrowsWhenNoMoreAndNoLinks() throws DatabricksSQLException { @Test @DisplayName("refetchLink should throw when no matching link found") - void testRefetchLinkThrowsWhenNoMatch() throws DatabricksSQLException { + void testRefetchLinkThrowsWhenNoMatch() throws SQLException { // Return links for chunk 6 and 7, but we're looking for chunk 5 ExternalLink link6 = createLink(6, 600, 100); ExternalLink link7 = createLink(7, 700, 100); diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/InlineChunkProviderTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/InlineChunkProviderTest.java index 86be512d4d..8392daf683 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/arrow/InlineChunkProviderTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/InlineChunkProviderTest.java @@ -1,27 +1,17 @@ package com.databricks.jdbc.api.impl.arrow; -import static com.databricks.jdbc.TestConstants.ARROW_BATCH_LIST; -import static com.databricks.jdbc.TestConstants.TEST_TABLE_SCHEMA; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import com.databricks.jdbc.api.internal.IDatabricksSession; -import com.databricks.jdbc.api.internal.IDatabricksStatementInternal; import com.databricks.jdbc.common.CompressionCodec; -import com.databricks.jdbc.exception.DatabricksParsingException; import com.databricks.jdbc.exception.DatabricksSQLException; -import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp; -import com.databricks.jdbc.model.client.thrift.generated.TGetResultSetMetadataResp; -import com.databricks.jdbc.model.client.thrift.generated.TRowSet; -import com.databricks.jdbc.model.client.thrift.generated.TSparkArrowBatch; import com.databricks.jdbc.model.core.ColumnInfo; import com.databricks.jdbc.model.core.ColumnInfoTypeName; import com.databricks.jdbc.model.core.ResultData; import com.databricks.jdbc.model.core.ResultManifest; import java.io.ByteArrayOutputStream; import java.io.IOException; -import java.util.Collections; import net.jpountz.lz4.LZ4FrameOutputStream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; @@ -37,41 +27,9 @@ public class InlineChunkProviderTest { private static final long TOTAL_ROWS = 2L; - @Mock TGetResultSetMetadataResp metadata; - @Mock TFetchResultsResp fetchResultsResp; - @Mock IDatabricksStatementInternal parentStatement; - @Mock IDatabricksSession session; @Mock private ResultData mockResultData; @Mock private ResultManifest mockResultManifest; - @Test - void testInitialisation() throws DatabricksParsingException { - when(fetchResultsResp.getResultSetMetadata()).thenReturn(metadata); - when(metadata.getArrowSchema()).thenReturn(null); - when(metadata.getSchema()).thenReturn(TEST_TABLE_SCHEMA); - when(fetchResultsResp.getResults()).thenReturn(new TRowSet().setArrowBatches(ARROW_BATCH_LIST)); - when(metadata.isSetLz4Compressed()).thenReturn(false); - InlineChunkProvider inlineChunkProvider = - new InlineChunkProvider(fetchResultsResp, parentStatement, session); - assertTrue(inlineChunkProvider.hasNextChunk()); - assertTrue(inlineChunkProvider.next()); - assertFalse(inlineChunkProvider.next()); - } - - @Test - void handleErrorTest() throws DatabricksParsingException { - TSparkArrowBatch arrowBatch = - new TSparkArrowBatch().setRowCount(0).setBatch(new byte[] {65, 66, 67}); - when(fetchResultsResp.getResultSetMetadata()).thenReturn(metadata); - when(fetchResultsResp.getResults()) - .thenReturn(new TRowSet().setArrowBatches(Collections.singletonList(arrowBatch))); - InlineChunkProvider inlineChunkProvider = - new InlineChunkProvider(fetchResultsResp, parentStatement, session); - assertThrows( - DatabricksParsingException.class, - () -> inlineChunkProvider.handleError(new RuntimeException())); - } - @Test void testConstructorSuccessfulCreation() throws DatabricksSQLException, IOException { // Create valid Arrow data with two rows and one column: [1, 2] diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResultTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResultTest.java new file mode 100644 index 0000000000..3cac2c52c0 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResultTest.java @@ -0,0 +1,429 @@ +package com.databricks.jdbc.api.impl.arrow; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.databricks.jdbc.api.impl.DatabricksConnectionContextFactory; +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.dbclient.IDatabricksClient; +import com.databricks.jdbc.dbclient.impl.common.StatementId; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.model.client.thrift.generated.*; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.io.ByteArrayOutputStream; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Properties; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowStreamWriter; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +public class LazyThriftInlineArrowResultTest { + + private static final String JDBC_URL = + "jdbc:databricks://sample-host.18.azuredatabricks.net:9999/default;transportMode=http;ssl=1;" + + "AuthMech=3;httpPath=/sql/1.0/warehouses/99999999;"; + + @Mock private IDatabricksSession session; + @Mock private IDatabricksStatementInternal statement; + @Mock private IDatabricksClient databricksClient; + private IDatabricksConnectionContext connectionContext; + private static final StatementId STATEMENT_ID = new StatementId("test_statement_id"); + private static final TTableSchema TWO_COLUMN_SCHEMA = + createTableSchema(TTypeId.INT_TYPE, TTypeId.STRING_TYPE); + + @BeforeEach + void setUp() throws Exception { + connectionContext = DatabricksConnectionContextFactory.create(JDBC_URL, new Properties()); + } + + /** + * Creates valid Arrow IPC format bytes with two columns (int and string). This creates a complete + * Arrow IPC stream (with embedded schema) that can be parsed by ArrowStreamReader. + * + *

Int column values: batch_index * 100 + row_index (e.g., 0, 1, 2... for batch 0) String + * column values: "row_{batch_index}_{row_index}" (e.g., "row_0_0", "row_0_1"...) + * + * @param batchCount Number of record batches to create + * @param rowsPerBatch Number of rows in each batch + * @return Valid Arrow IPC bytes + */ + private static byte[] createValidArrowData(int batchCount, int rowsPerBatch) { + try (BufferAllocator allocator = new RootAllocator(); + ByteArrayOutputStream out = new ByteArrayOutputStream()) { + + try (IntVector intVector = new IntVector("int_column", allocator); + VarCharVector stringVector = new VarCharVector("string_column", allocator)) { + + intVector.allocateNew(rowsPerBatch); + stringVector.allocateNew(rowsPerBatch); + + try (VectorSchemaRoot root = VectorSchemaRoot.of(intVector, stringVector); + ArrowStreamWriter writer = new ArrowStreamWriter(root, null, out)) { + writer.start(); + + for (int batch = 0; batch < batchCount; batch++) { + for (int i = 0; i < rowsPerBatch; i++) { + intVector.set(i, batch * 100 + i); + stringVector.setSafe(i, ("row_" + batch + "_" + i).getBytes()); + } + intVector.setValueCount(rowsPerBatch); + stringVector.setValueCount(rowsPerBatch); + root.setRowCount(rowsPerBatch); + writer.writeBatch(); + } + + writer.end(); + } + } + + return out.toByteArray(); + } catch (Exception e) { + throw new RuntimeException("Failed to create test Arrow data", e); + } + } + + /** + * Creates a TTableSchema with the specified column types. The schema structure matches what + * hiveSchemaToArrowSchema expects. + */ + private static TTableSchema createTableSchema(TTypeId... types) { + List columns = new ArrayList<>(); + for (int i = 0; i < types.length; i++) { + TPrimitiveTypeEntry primitiveType = new TPrimitiveTypeEntry().setType(types[i]); + TTypeEntry typeEntry = new TTypeEntry(); + typeEntry.setPrimitiveEntry(primitiveType); + TTypeDesc typeDesc = new TTypeDesc().setTypes(Collections.singletonList(typeEntry)); + TColumnDesc columnDesc = + new TColumnDesc().setColumnName("col_" + i).setTypeDesc(typeDesc).setPosition(i); + columns.add(columnDesc); + } + return new TTableSchema().setColumns(columns); + } + + /** + * Creates a TFetchResultsResp with valid Arrow data. Uses empty arrowSchema bytes so that the + * complete Arrow IPC stream (with embedded schema) is used as-is. + * + * @param arrowData Valid Arrow IPC bytes (should include schema) + * @param rowCount Row count for the batch + * @param hasMoreRows Whether there are more rows to fetch + * @return Configured TFetchResultsResp + */ + private TFetchResultsResp createFetchResultsResp( + byte[] arrowData, int rowCount, boolean hasMoreRows) { + TSparkArrowBatch arrowBatch = new TSparkArrowBatch().setRowCount(rowCount).setBatch(arrowData); + TRowSet rowSet = new TRowSet().setArrowBatches(Collections.singletonList(arrowBatch)); + + // Use empty arrowSchema - this causes getSerializedSchema to return empty bytes, + // which effectively means nothing is prepended to our complete Arrow IPC stream + TGetResultSetMetadataResp metadata = + new TGetResultSetMetadataResp().setSchema(TWO_COLUMN_SCHEMA).setArrowSchema(new byte[0]); + + TFetchResultsResp response = + new TFetchResultsResp().setResultSetMetadata(metadata).setResults(rowSet); + response.hasMoreRows = hasMoreRows; + + return response; + } + + @Test + void testEmptyResultSet() throws SQLException { + byte[] arrowData = createValidArrowData(1, 0); + TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, 0, false); + + when(statement.getMaxRows()).thenReturn(0); + when(statement.getStatementId()).thenReturn(STATEMENT_ID); + + LazyThriftInlineArrowResult result = + new LazyThriftInlineArrowResult(initialResponse, statement, session); + + // Verify all initial state for empty result + assertEquals(-1, result.getCurrentRow()); + assertEquals(0, result.getRowCount()); + assertEquals(0, result.getTotalRowsFetched()); + assertEquals(0, result.getChunkCount()); + assertFalse(result.hasNext()); + assertFalse(result.next()); + assertTrue(result.isCompletelyFetched()); + } + + @Test + void testGetObjectThrowsWhenClosed() throws SQLException { + byte[] arrowData = createValidArrowData(1, 1); + TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, 1, false); + + when(statement.getMaxRows()).thenReturn(0); + when(statement.getStatementId()).thenReturn(STATEMENT_ID); + + LazyThriftInlineArrowResult result = + new LazyThriftInlineArrowResult(initialResponse, statement, session); + result.close(); + + // Verify close behavior + assertFalse(result.hasNext()); + assertFalse(result.next()); + assertEquals(0, result.getRowCount()); + + DatabricksSQLException exception = + assertThrows(DatabricksSQLException.class, () -> result.getObject(0)); + assertEquals("Result is already closed", exception.getMessage()); + assertEquals(DatabricksDriverErrorCode.STATEMENT_CLOSED.name(), exception.getSQLState()); + } + + @Test + void testGetObjectThrowsWhenBeforeFirstRow() throws SQLException { + byte[] arrowData = createValidArrowData(1, 1); + TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, 1, false); + + when(statement.getMaxRows()).thenReturn(0); + when(statement.getStatementId()).thenReturn(STATEMENT_ID); + + LazyThriftInlineArrowResult result = + new LazyThriftInlineArrowResult(initialResponse, statement, session); + + DatabricksSQLException exception = + assertThrows(DatabricksSQLException.class, () -> result.getObject(0)); + assertEquals("Cursor is before first row", exception.getMessage()); + assertEquals(DatabricksDriverErrorCode.INVALID_STATE.name(), exception.getSQLState()); + } + + @Test + void testIsCompletelyFetchedWithMoreRows() throws SQLException { + byte[] arrowData = createValidArrowData(1, 0); + TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, 0, true); + + when(statement.getMaxRows()).thenReturn(0); + when(statement.getStatementId()).thenReturn(STATEMENT_ID); + + LazyThriftInlineArrowResult result = + new LazyThriftInlineArrowResult(initialResponse, statement, session); + + assertFalse(result.isCompletelyFetched()); + assertTrue(result.hasNext()); // hasNext is true because hasMoreRows is true + } + + @Test + void testIterateThroughRowsWithValidArrowData() throws SQLException { + int rowCount = 5; + byte[] arrowData = createValidArrowData(1, rowCount); + TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, rowCount, false); + + when(statement.getMaxRows()).thenReturn(0); + when(statement.getStatementId()).thenReturn(STATEMENT_ID); + + LazyThriftInlineArrowResult result = + new LazyThriftInlineArrowResult(initialResponse, statement, session); + + // Verify initial state + assertEquals(-1, result.getCurrentRow()); + assertTrue(result.hasNext()); + + // Iterate through all rows + for (int i = 0; i < rowCount; i++) { + assertTrue(result.hasNext(), "Should have next at row " + i); + assertTrue(result.next(), "next() should return true at row " + i); + assertEquals(i, result.getCurrentRow()); + } + + // After all rows + assertFalse(result.hasNext()); + assertFalse(result.next()); + assertEquals(rowCount, result.getTotalRowsFetched()); + } + + @Test + void testGetObjectReturnsCorrectIntegerValue() throws SQLException { + int rowCount = 3; + byte[] arrowData = createValidArrowData(1, rowCount); + TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, rowCount, false); + + when(statement.getMaxRows()).thenReturn(0); + when(statement.getStatementId()).thenReturn(STATEMENT_ID); + when(session.getConnectionContext()).thenReturn(connectionContext); + + LazyThriftInlineArrowResult result = + new LazyThriftInlineArrowResult(initialResponse, statement, session); + + // Move to first row and get value + assertTrue(result.next()); + Object value = result.getObject(0); + assertNotNull(value); + assertInstanceOf(Integer.class, value); + assertEquals(0, value); // First row: batch_0 * 100 + row_0 = 0 + + // Move to second row + assertTrue(result.next()); + value = result.getObject(0); + assertEquals(1, value); // Second row: batch_0 * 100 + row_1 = 1 + + // Move to third row + assertTrue(result.next()); + value = result.getObject(0); + assertEquals(2, value); // Third row: batch_0 * 100 + row_2 = 2 + } + + @Test + void testGetObjectWithTwoColumns() throws SQLException { + int rowCount = 2; + byte[] arrowData = createValidArrowData(1, rowCount); + TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, rowCount, false); + + when(statement.getMaxRows()).thenReturn(0); + when(statement.getStatementId()).thenReturn(STATEMENT_ID); + when(session.getConnectionContext()).thenReturn(connectionContext); + + LazyThriftInlineArrowResult result = + new LazyThriftInlineArrowResult(initialResponse, statement, session); + + assertTrue(result.next()); + + // Get integer column + Object intValue = result.getObject(0); + assertNotNull(intValue); + assertInstanceOf(Integer.class, intValue); + assertEquals(0, intValue); + + // Get string column + Object stringValue = result.getObject(1); + assertNotNull(stringValue); + assertEquals("row_0_0", stringValue.toString()); + } + + @Test + void testGetObjectThrowsForColumnIndexOutOfBounds() throws SQLException { + int rowCount = 1; + byte[] arrowData = createValidArrowData(1, rowCount); + TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, rowCount, false); + + when(statement.getMaxRows()).thenReturn(0); + when(statement.getStatementId()).thenReturn(STATEMENT_ID); + // Note: session.getConnectionContext() is not stubbed here because the column index + // validation happens before the connection context is accessed + + LazyThriftInlineArrowResult result = + new LazyThriftInlineArrowResult(initialResponse, statement, session); + + assertTrue(result.next()); + + // Test negative index + DatabricksSQLException negativeException = + assertThrows(DatabricksSQLException.class, () -> result.getObject(-1)); + assertTrue(negativeException.getMessage().contains("Column index out of bounds")); + assertEquals(DatabricksDriverErrorCode.INVALID_STATE.name(), negativeException.getSQLState()); + + // Test index beyond column count (we have 2 columns: 0 and 1) + DatabricksSQLException beyondException = + assertThrows(DatabricksSQLException.class, () -> result.getObject(2)); + assertTrue(beyondException.getMessage().contains("Column index out of bounds")); + assertEquals(DatabricksDriverErrorCode.INVALID_STATE.name(), beyondException.getSQLState()); + } + + @Test + void testMaxRowsLimitEnforced() throws SQLException { + int totalRows = 10; + int maxRows = 3; + byte[] arrowData = createValidArrowData(1, totalRows); + TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, totalRows, false); + + when(statement.getMaxRows()).thenReturn(maxRows); + when(statement.getStatementId()).thenReturn(STATEMENT_ID); + + LazyThriftInlineArrowResult result = + new LazyThriftInlineArrowResult(initialResponse, statement, session); + + // Should only be able to iterate up to maxRows + int rowsRetrieved = 0; + while (result.next()) { + rowsRetrieved++; + } + + assertEquals(maxRows, rowsRetrieved); + assertFalse(result.hasNext()); + assertEquals(maxRows - 1, result.getCurrentRow()); // 0-based index + } + + @Test + void testGetArrowMetadataReturnsMetadata() throws SQLException { + int rowCount = 1; + byte[] arrowData = createValidArrowData(1, rowCount); + TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, rowCount, false); + + when(statement.getMaxRows()).thenReturn(0); + when(statement.getStatementId()).thenReturn(STATEMENT_ID); + + LazyThriftInlineArrowResult result = + new LazyThriftInlineArrowResult(initialResponse, statement, session); + + List metadata = result.getArrowMetadata(); + assertNotNull(metadata); + // The metadata list should have one entry per column (2 columns: int and string) + assertEquals(2, metadata.size()); + } + + @Test + void testFetchNextChunkFromServer() throws SQLException { + int rowsPerChunk = 2; + byte[] arrowData1 = createValidArrowData(1, rowsPerChunk); + byte[] arrowData2 = createValidArrowData(1, rowsPerChunk); + + // First chunk with hasMoreRows = true + TFetchResultsResp initialResponse = createFetchResultsResp(arrowData1, rowsPerChunk, true); + + // Second chunk with hasMoreRows = false + TFetchResultsResp secondResponse = createFetchResultsResp(arrowData2, rowsPerChunk, false); + + when(statement.getMaxRows()).thenReturn(0); + when(statement.getStatementId()).thenReturn(STATEMENT_ID); + when(session.getDatabricksClient()).thenReturn(databricksClient); + when(databricksClient.getMoreResults(statement)).thenReturn(secondResponse); + + LazyThriftInlineArrowResult result = + new LazyThriftInlineArrowResult(initialResponse, statement, session); + + // Iterate through first chunk + assertTrue(result.next()); + assertTrue(result.next()); + assertFalse(result.isCompletelyFetched()); // Still has more rows + + // This should trigger fetch of next chunk + assertTrue(result.next()); + assertTrue(result.next()); + + // After all rows + assertFalse(result.next()); + assertTrue(result.isCompletelyFetched()); + assertEquals(rowsPerChunk * 2, result.getTotalRowsFetched()); + + // Verify that getMoreResults was called + verify(databricksClient).getMoreResults(statement); + } + + @Test + void testGetRowCountReturnsCurrentChunkRowCount() throws SQLException { + int rowCount = 5; + byte[] arrowData = createValidArrowData(1, rowCount); + TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, rowCount, false); + + when(statement.getMaxRows()).thenReturn(0); + when(statement.getStatementId()).thenReturn(STATEMENT_ID); + + LazyThriftInlineArrowResult result = + new LazyThriftInlineArrowResult(initialResponse, statement, session); + + assertEquals(rowCount, result.getRowCount()); + } +} diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkProviderTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkProviderTest.java index 31b59fc365..39cfb98545 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkProviderTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkProviderTest.java @@ -300,6 +300,7 @@ private StreamingChunkProvider createProvider( linkPrefetchWindow, CHUNK_READY_TIMEOUT_SECONDS, CLOUD_FETCH_SPEED_THRESHOLD, + null, // connectionContext - not needed for tests initialLinks); } diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResultTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResultTest.java new file mode 100644 index 0000000000..0280faa923 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResultTest.java @@ -0,0 +1,739 @@ +package com.databricks.jdbc.api.impl.arrow; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.databricks.jdbc.api.impl.DatabricksConnectionContextFactory; +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.dbclient.IDatabricksClient; +import com.databricks.jdbc.dbclient.impl.common.StatementId; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.model.client.thrift.generated.*; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.io.ByteArrayOutputStream; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Properties; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowStreamWriter; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +/** + * Unit tests for StreamingInlineArrowResult. Verifies functional equivalence with + * LazyThriftInlineArrowResult from the consumer's perspective, along with streaming-specific + * behaviors. + */ +@ExtendWith(MockitoExtension.class) +public class StreamingInlineArrowResultTest { + + private static final String JDBC_URL = + "jdbc:databricks://sample-host.18.azuredatabricks.net:9999/default;transportMode=http;ssl=1;" + + "AuthMech=3;httpPath=/sql/1.0/warehouses/99999999;"; + + @Mock private IDatabricksSession session; + @Mock private IDatabricksStatementInternal statement; + @Mock private IDatabricksClient databricksClient; + + private IDatabricksConnectionContext connectionContext; + private static final StatementId STATEMENT_ID = new StatementId("test_statement_id"); + private static final TTableSchema TWO_COLUMN_SCHEMA = + createTableSchema(TTypeId.INT_TYPE, TTypeId.STRING_TYPE); + + @BeforeEach + void setUp() throws Exception { + connectionContext = DatabricksConnectionContextFactory.create(JDBC_URL, new Properties()); + lenient().when(session.getDatabricksClient()).thenReturn(databricksClient); + lenient().when(session.getConnectionContext()).thenReturn(connectionContext); + lenient().when(statement.getMaxRows()).thenReturn(0); + lenient().when(statement.getStatementId()).thenReturn(STATEMENT_ID); + } + + @Test + void testBasicIteration() throws SQLException { + int rowCount = 5; + byte[] arrowData = createValidArrowData(1, rowCount); + TFetchResultsResp response = createFetchResultsResp(arrowData, rowCount, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + try { + // Initial state + assertEquals(-1, result.getCurrentRow()); + assertTrue(result.hasNext()); + + // Iterate through all rows + for (int i = 0; i < rowCount; i++) { + assertTrue(result.hasNext(), "Should have next at row " + i); + assertTrue(result.next(), "next() should return true at row " + i); + assertEquals(i, result.getCurrentRow()); + } + + // End + assertFalse(result.hasNext()); + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testGetObjectReturnsCorrectValues() throws SQLException { + int rowCount = 3; + byte[] arrowData = createValidArrowData(1, rowCount); + TFetchResultsResp response = createFetchResultsResp(arrowData, rowCount, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + try { + // Row 0 + assertTrue(result.next()); + Object value = result.getObject(0); + assertNotNull(value); + assertInstanceOf(Integer.class, value); + assertEquals(0, value); + + // Row 1 + assertTrue(result.next()); + value = result.getObject(0); + assertEquals(1, value); + + // Row 2 + assertTrue(result.next()); + value = result.getObject(0); + assertEquals(2, value); + } finally { + result.close(); + } + } + + @Test + void testMultiColumnAccess() throws SQLException { + int rowCount = 2; + byte[] arrowData = createValidArrowData(1, rowCount); + TFetchResultsResp response = createFetchResultsResp(arrowData, rowCount, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + try { + assertTrue(result.next()); + + // Int column + Object intValue = result.getObject(0); + assertNotNull(intValue); + assertInstanceOf(Integer.class, intValue); + assertEquals(0, intValue); + + // String column + Object stringValue = result.getObject(1); + assertNotNull(stringValue); + assertEquals("row_0_0", stringValue.toString()); + } finally { + result.close(); + } + } + + @Test + void testMaxRowsLimit() throws SQLException { + int totalRows = 10; + int maxRows = 3; + when(statement.getMaxRows()).thenReturn(maxRows); + + byte[] arrowData = createValidArrowData(1, totalRows); + TFetchResultsResp response = createFetchResultsResp(arrowData, totalRows, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + try { + int rowsRetrieved = 0; + while (result.next()) { + rowsRetrieved++; + } + + assertEquals(maxRows, rowsRetrieved); + assertFalse(result.hasNext()); + assertEquals(maxRows - 1, result.getCurrentRow()); + } finally { + result.close(); + } + } + + @Test + void testColumnIndexBounds() throws SQLException { + byte[] arrowData = createValidArrowData(1, 1); + TFetchResultsResp response = createFetchResultsResp(arrowData, 1, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + try { + assertTrue(result.next()); + + // Negative index + DatabricksSQLException negativeException = + assertThrows(DatabricksSQLException.class, () -> result.getObject(-1)); + assertTrue(negativeException.getMessage().contains("Column index out of bounds")); + + // Beyond column count + DatabricksSQLException beyondException = + assertThrows(DatabricksSQLException.class, () -> result.getObject(2)); + assertTrue(beyondException.getMessage().contains("Column index out of bounds")); + } finally { + result.close(); + } + } + + @Test + void testAccessAfterClose() throws SQLException { + byte[] arrowData = createValidArrowData(1, 1); + TFetchResultsResp response = createFetchResultsResp(arrowData, 1, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + result.close(); + + assertFalse(result.hasNext()); + assertFalse(result.next()); + assertThrows(DatabricksSQLException.class, () -> result.getObject(0)); + } + + @Test + void testAccessBeforeFirstRow() throws SQLException { + byte[] arrowData = createValidArrowData(1, 1); + TFetchResultsResp response = createFetchResultsResp(arrowData, 1, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + try { + DatabricksSQLException exception = + assertThrows(DatabricksSQLException.class, () -> result.getObject(0)); + assertTrue(exception.getMessage().contains("before first row")); + } finally { + result.close(); + } + } + + @Test + void testSingleRowResult() throws SQLException { + // Test single row instead of empty - streaming has different empty handling + byte[] arrowData = createValidArrowData(1, 1); + TFetchResultsResp response = createFetchResultsResp(arrowData, 1, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + try { + assertEquals(-1, result.getCurrentRow()); + assertEquals(1, result.getRowCount()); + assertTrue(result.hasNext()); + assertTrue(result.next()); + assertEquals(0, result.getCurrentRow()); + assertFalse(result.hasNext()); + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testMultiBatchFetching() throws SQLException, InterruptedException { + int rowsPerChunk = 2; + byte[] arrowData1 = createValidArrowData(1, rowsPerChunk); + byte[] arrowData2 = createValidArrowData(1, rowsPerChunk); + + TFetchResultsResp firstResponse = createFetchResultsResp(arrowData1, rowsPerChunk, true); + TFetchResultsResp secondResponse = createFetchResultsResp(arrowData2, rowsPerChunk, false); + + when(databricksClient.getMoreResults(statement)).thenReturn(secondResponse); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(firstResponse, statement, session); + + try { + // Give prefetch thread time to start + Thread.sleep(100); + + // Iterate through first batch + assertTrue(result.next()); + assertTrue(result.next()); + + // Move to second batch + assertTrue(result.next()); + assertTrue(result.next()); + + // End + assertFalse(result.next()); + + verify(databricksClient, atLeastOnce()).getMoreResults(statement); + } finally { + result.close(); + } + } + + @Test + void testErrorDuringFetch() throws Exception { + byte[] arrowData = createValidArrowData(1, 2); + TFetchResultsResp firstResponse = createFetchResultsResp(arrowData, 2, true); + + DatabricksSQLException expectedException = + new DatabricksSQLException("Network error", DatabricksDriverErrorCode.CONNECTION_ERROR); + when(databricksClient.getMoreResults(statement)).thenThrow(expectedException); + + // The prefetch thread starts during construction and may fail before or after + // the constructor completes. We need to handle both cases. + DatabricksSQLException caughtException = null; + StreamingInlineArrowResult result = null; + + try { + result = new StreamingInlineArrowResult(firstResponse, statement, session); + + // If construction succeeded, consume rows until we hit the error + while (result.next()) { + // Keep iterating - error will be thrown when we need the next batch + } + } catch (DatabricksSQLException e) { + // Error thrown during construction or iteration + caughtException = e; + } finally { + if (result != null) { + result.close(); + } + } + + // Verify we caught the expected error + assertNotNull(caughtException, "Expected DatabricksSQLException to be thrown"); + assertTrue( + caughtException.getMessage().contains("Prefetch failed") + || caughtException.getMessage().contains("Network error"), + "Exception message should contain error details: " + caughtException.getMessage()); + } + + @Test + void testMaxRowsLimitAcrossBatches() throws SQLException, InterruptedException { + // MaxRows limit of 3, spanning across 2 batches (2 rows each) + int rowsPerChunk = 2; + when(statement.getMaxRows()).thenReturn(3); + + byte[] arrowData1 = createValidArrowData(1, rowsPerChunk); + byte[] arrowData2 = createValidArrowData(1, rowsPerChunk); + + TFetchResultsResp firstResponse = createFetchResultsResp(arrowData1, rowsPerChunk, true); + TFetchResultsResp secondResponse = createFetchResultsResp(arrowData2, rowsPerChunk, false); + + when(databricksClient.getMoreResults(statement)).thenReturn(secondResponse); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(firstResponse, statement, session); + + try { + // Give prefetch thread time to start + Thread.sleep(100); + + // Consume first batch (rows 0 and 1) + assertTrue(result.next()); + assertEquals(0, result.getCurrentRow()); + assertTrue(result.next()); + assertEquals(1, result.getCurrentRow()); + + // Get one row from second batch (row 2) + assertTrue(result.next()); + assertEquals(2, result.getCurrentRow()); + + // Should stop at maxRows=3 + assertFalse(result.hasNext()); + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testGetArrowMetadata() throws SQLException { + byte[] arrowData = createValidArrowData(1, 2); + TFetchResultsResp response = createFetchResultsResp(arrowData, 2, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + try { + assertTrue(result.next()); + List metadata = result.getArrowMetadata(); + assertNotNull(metadata); + // Two columns: int and string + assertEquals(2, metadata.size()); + } finally { + result.close(); + } + } + + @Test + void testGetTotalRowsFetched() throws SQLException, InterruptedException { + byte[] arrowData1 = createValidArrowData(1, 3); + byte[] arrowData2 = createValidArrowData(1, 2); + TFetchResultsResp firstBatch = createFetchResultsResp(arrowData1, 3, true); + TFetchResultsResp secondBatch = createFetchResultsResp(arrowData2, 2, false); + + when(databricksClient.getMoreResults(statement)).thenReturn(secondBatch); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(firstBatch, statement, session); + + try { + // Initial batch has at least 3 rows + assertTrue(result.getTotalRowsFetched() >= 3); + + // Give prefetch thread time to fetch second batch + Thread.sleep(100); + + // After prefetch completes, should have 5 total rows fetched + assertEquals(5, result.getTotalRowsFetched()); + } finally { + result.close(); + } + } + + @Test + void testIsCompletelyFetched() throws SQLException, InterruptedException { + byte[] arrowData1 = createValidArrowData(1, 2); + byte[] arrowData2 = createValidArrowData(1, 2); + TFetchResultsResp firstBatch = createFetchResultsResp(arrowData1, 2, true); + TFetchResultsResp secondBatch = createFetchResultsResp(arrowData2, 2, false); + + when(databricksClient.getMoreResults(statement)).thenReturn(secondBatch); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(firstBatch, statement, session); + + try { + // Give prefetch thread time to fetch second batch + Thread.sleep(100); + + // After fetching final batch (hasMoreRows=false), should be completely fetched + assertTrue(result.isCompletelyFetched()); + } finally { + result.close(); + } + } + + @Test + void testBatchesInMemoryTracking() throws SQLException { + byte[] arrowData = createValidArrowData(1, 2); + TFetchResultsResp response = createFetchResultsResp(arrowData, 2, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + try { + // Should have at least initial batch + assertTrue(result.getBatchesInMemory() > 0); + } finally { + result.close(); + } + } + + @Test + void testChunkCount() throws SQLException { + byte[] arrowData = createValidArrowData(1, 2); + TFetchResultsResp response = createFetchResultsResp(arrowData, 2, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + try { + // Streaming results always return 0 for chunk count (not applicable) + assertEquals(0, result.getChunkCount()); + } finally { + result.close(); + } + } + + @Test + void testGetRowCount() throws SQLException { + byte[] arrowData = createValidArrowData(1, 5); + TFetchResultsResp response = createFetchResultsResp(arrowData, 5, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + try { + assertTrue(result.next()); + assertEquals(5, result.getRowCount()); + } finally { + result.close(); + } + } + + @Test + void testInitializationWithEmptyInitialBatch() throws SQLException, InterruptedException { + // Initial batch is EMPTY but hasMoreRows=true + TFetchResultsResp emptyInitial = createEmptyArrowResponse(true); + + // Second batch has actual data + byte[] arrowData = createValidArrowData(1, 3); + TFetchResultsResp dataBatch = createFetchResultsResp(arrowData, 3, false); + + when(databricksClient.getMoreResults(statement)).thenReturn(dataBatch); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(emptyInitial, statement, session); + + try { + // Give prefetch time to fetch the data batch + Thread.sleep(100); + + // Should have skipped empty batch and positioned on data batch + assertTrue(result.hasNext()); + assertTrue(result.next()); + assertEquals(0, result.getObject(0)); // First int value + + assertTrue(result.next()); + assertTrue(result.next()); + assertFalse(result.next()); // Only 3 rows + } finally { + result.close(); + } + } + + @Test + void testDoubleClose() throws SQLException { + byte[] arrowData = createValidArrowData(1, 2); + TFetchResultsResp response = createFetchResultsResp(arrowData, 2, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + // First close + result.close(); + // Second close should not throw + assertDoesNotThrow(result::close); + } + + @Test + void testExhaustAllRowsInSingleBatch() throws SQLException { + // Single batch with 3 rows, no more data + byte[] arrowData = createValidArrowData(1, 3); + TFetchResultsResp response = createFetchResultsResp(arrowData, 3, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + try { + // Consume all 3 rows + assertTrue(result.next()); + assertEquals(0, result.getObject(0)); // First row + assertTrue(result.next()); + assertEquals(1, result.getObject(0)); // Second row + assertTrue(result.next()); + assertEquals(2, result.getObject(0)); // Third row + + // No more rows - should return false + assertFalse(result.next()); + // Calling again should still return false + assertFalse(result.next()); + + // hasNext should also be false + assertFalse(result.hasNext()); + } finally { + result.close(); + } + } + + @Test + void testExhaustAllRowsAcrossMultipleBatches() throws SQLException, InterruptedException { + // First batch with 2 rows + byte[] arrowData1 = createValidArrowData(1, 2); + TFetchResultsResp firstBatch = createFetchResultsResp(arrowData1, 2, true); + + // Second batch with 1 row - final batch + byte[] arrowData2 = createValidArrowData(1, 1); + TFetchResultsResp secondBatch = createFetchResultsResp(arrowData2, 1, false); + + when(databricksClient.getMoreResults(statement)).thenReturn(secondBatch); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(firstBatch, statement, session); + + try { + // Give prefetch time + Thread.sleep(100); + + // Consume batch 1 (2 rows) + assertTrue(result.next()); + assertEquals(0, result.getObject(0)); + assertTrue(result.next()); + assertEquals(1, result.getObject(0)); + + // Consume batch 2 (1 row) + assertTrue(result.next()); + assertEquals(0, result.getObject(0)); // First row of second batch + + // No more rows - triggers end of stream detection + assertFalse(result.next()); + assertFalse(result.hasNext()); + + // Multiple calls should consistently return false + assertFalse(result.next()); + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testBatchTransitionAtExactBoundary() throws SQLException, InterruptedException { + // First batch with exactly 1 row + byte[] arrowData1 = createValidArrowData(1, 1); + TFetchResultsResp firstBatch = createFetchResultsResp(arrowData1, 1, true); + + // Second batch with 1 row + byte[] arrowData2 = createValidArrowData(1, 1); + TFetchResultsResp secondBatch = createFetchResultsResp(arrowData2, 1, false); + + when(databricksClient.getMoreResults(statement)).thenReturn(secondBatch); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(firstBatch, statement, session); + + try { + // Give prefetch time + Thread.sleep(100); + + // Row 1 - from batch 1 + assertTrue(result.next()); + assertEquals(0, result.getObject(0)); + assertEquals(0, result.getCurrentRow()); + + // Row 2 - requires batch transition + assertTrue(result.next()); + assertEquals(0, result.getObject(0)); // First row of second batch + assertEquals(1, result.getCurrentRow()); + + // End - no more batches + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testNextAfterEndReturnsConsistentFalse() throws SQLException { + byte[] arrowData = createValidArrowData(1, 1); + TFetchResultsResp response = createFetchResultsResp(arrowData, 1, false); + + StreamingInlineArrowResult result = + new StreamingInlineArrowResult(response, statement, session); + + try { + // Consume the only row + assertTrue(result.next()); + + // Multiple calls to next() after exhaustion should all return false + assertFalse(result.next()); + assertFalse(result.next()); + assertFalse(result.next()); + + // hasReachedEnd should be set + assertFalse(result.hasNext()); + } finally { + result.close(); + } + } + + // ==================== Helper Methods ==================== + + private TFetchResultsResp createEmptyArrowResponse(boolean hasMoreRows) { + // Create valid Arrow data with schema and 1 batch of 0 rows + // The Arrow data includes its own schema via ArrowStreamWriter + byte[] emptyArrowData = createValidArrowData(1, 0); + TSparkArrowBatch arrowBatch = new TSparkArrowBatch().setRowCount(0).setBatch(emptyArrowData); + TRowSet rowSet = new TRowSet().setArrowBatches(Collections.singletonList(arrowBatch)); + + // Use empty arrowSchema so cachedSchema won't be prepended (Arrow data has its own schema) + TGetResultSetMetadataResp metadata = + new TGetResultSetMetadataResp().setSchema(TWO_COLUMN_SCHEMA).setArrowSchema(new byte[0]); + TFetchResultsResp response = + new TFetchResultsResp().setResultSetMetadata(metadata).setResults(rowSet); + response.hasMoreRows = hasMoreRows; + return response; + } + + private static byte[] createValidArrowData(int batchCount, int rowsPerBatch) { + try (BufferAllocator allocator = new RootAllocator(); + ByteArrayOutputStream out = new ByteArrayOutputStream()) { + + try (IntVector intVector = new IntVector("int_column", allocator); + VarCharVector stringVector = new VarCharVector("string_column", allocator)) { + + intVector.allocateNew(rowsPerBatch); + stringVector.allocateNew(rowsPerBatch); + + try (VectorSchemaRoot root = VectorSchemaRoot.of(intVector, stringVector); + ArrowStreamWriter writer = new ArrowStreamWriter(root, null, out)) { + writer.start(); + + for (int batch = 0; batch < batchCount; batch++) { + for (int i = 0; i < rowsPerBatch; i++) { + intVector.set(i, batch * 100 + i); + stringVector.setSafe(i, ("row_" + batch + "_" + i).getBytes()); + } + intVector.setValueCount(rowsPerBatch); + stringVector.setValueCount(rowsPerBatch); + root.setRowCount(rowsPerBatch); + writer.writeBatch(); + } + + writer.end(); + } + } + + return out.toByteArray(); + } catch (Exception e) { + throw new RuntimeException("Failed to create test Arrow data", e); + } + } + + private static TTableSchema createTableSchema(TTypeId... types) { + List columns = new ArrayList<>(); + for (int i = 0; i < types.length; i++) { + TPrimitiveTypeEntry primitiveType = new TPrimitiveTypeEntry().setType(types[i]); + TTypeEntry typeEntry = new TTypeEntry(); + typeEntry.setPrimitiveEntry(primitiveType); + TTypeDesc typeDesc = new TTypeDesc().setTypes(Collections.singletonList(typeEntry)); + TColumnDesc columnDesc = + new TColumnDesc().setColumnName("col_" + i).setTypeDesc(typeDesc).setPosition(i); + columns.add(columnDesc); + } + return new TTableSchema().setColumns(columns); + } + + private TFetchResultsResp createFetchResultsResp( + byte[] arrowData, int rowCount, boolean hasMoreRows) { + TSparkArrowBatch arrowBatch = new TSparkArrowBatch().setRowCount(rowCount).setBatch(arrowData); + TRowSet rowSet = new TRowSet().setArrowBatches(Collections.singletonList(arrowBatch)); + + TGetResultSetMetadataResp metadata = + new TGetResultSetMetadataResp().setSchema(TWO_COLUMN_SCHEMA).setArrowSchema(new byte[0]); + + TFetchResultsResp response = + new TFetchResultsResp().setResultSetMetadata(metadata).setResults(rowSet); + response.hasMoreRows = hasMoreRows; + + return response; + } +} diff --git a/src/test/java/com/databricks/jdbc/api/impl/streaming/ColumnarResponseProcessorTest.java b/src/test/java/com/databricks/jdbc/api/impl/streaming/ColumnarResponseProcessorTest.java new file mode 100644 index 0000000000..c7b9090080 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/api/impl/streaming/ColumnarResponseProcessorTest.java @@ -0,0 +1,142 @@ +package com.databricks.jdbc.api.impl.streaming; + +import static org.junit.jupiter.api.Assertions.*; + +import com.databricks.jdbc.api.impl.ColumnarRowView; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.model.client.thrift.generated.*; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Unit tests for ColumnarResponseProcessor. */ +public class ColumnarResponseProcessorTest { + + private ColumnarResponseProcessor processor; + + @BeforeEach + void setUp() { + processor = new ColumnarResponseProcessor(); + } + + @Test + void testProcessInitialResponse() throws DatabricksSQLException { + TFetchResultsResp response = createResponseWithData(2, true); + + StreamingBatch batch = processor.processInitialResponse(response); + + assertNotNull(batch); + assertEquals(0, batch.getBatchIndex()); + assertEquals(0, batch.getRowOffset()); + assertEquals(2, batch.getRowCount()); + assertTrue(batch.hasMoreRows()); + assertTrue(batch.isReady()); + assertNotNull(batch.getData()); + } + + @Test + void testProcessResponse() throws DatabricksSQLException { + TFetchResultsResp response = createResponseWithData(5, false); + + StreamingBatch batch = processor.processResponse(response, 3, 100); + + assertEquals(3, batch.getBatchIndex()); + assertEquals(100, batch.getRowOffset()); + assertEquals(5, batch.getRowCount()); + assertFalse(batch.hasMoreRows()); + } + + @Test + void testProcessEmptyResponse() throws DatabricksSQLException { + TFetchResultsResp response = createEmptyResponse(true); + + StreamingBatch batch = processor.processInitialResponse(response); + + assertNotNull(batch); + assertEquals(0, batch.getRowCount()); + assertTrue(batch.hasMoreRows()); + assertTrue(batch.isReady()); + } + + @Test + void testDataAccessible() throws DatabricksSQLException { + TFetchResultsResp response = createResponseWithData(2, false); + + StreamingBatch batch = processor.processInitialResponse(response); + ColumnarRowView view = batch.getData(); + + assertNotNull(view); + assertEquals(2, view.getRowCount()); + assertEquals(1, view.getColumnCount()); + assertEquals("value_0", view.getValue(0, 0)); + assertEquals("value_1", view.getValue(1, 0)); + } + + @Test + void testReleaseActionIsNoOp() throws DatabricksSQLException { + TFetchResultsResp response = createResponseWithData(1, false); + StreamingBatch batch = processor.processInitialResponse(response); + + ColumnarRowView viewBeforeRelease = batch.getData(); + assertNotNull(viewBeforeRelease); + + // Release should work without throwing (no-op for columnar) + batch.release(); + + assertEquals(StreamingBatch.Status.RELEASED, batch.getStatus()); + assertNull(batch.getData()); // Data reference is cleared + } + + @Test + void testGetReleaseAction() { + // Should return a no-op consumer (doesn't throw) + var releaseAction = processor.getReleaseAction(); + assertNotNull(releaseAction); + + // Should not throw when called with any value + assertDoesNotThrow(() -> releaseAction.accept(null)); + + TFetchResultsResp response = createResponseWithData(1, false); + try { + ColumnarRowView view = processor.processInitialResponse(response).getData(); + assertDoesNotThrow(() -> releaseAction.accept(view)); + } catch (DatabricksSQLException e) { + fail("Should not throw: " + e.getMessage()); + } + } + + // ==================== Helper Methods ==================== + + private TFetchResultsResp createEmptyResponse(boolean hasMoreRows) { + TFetchResultsResp response = new TFetchResultsResp(); + response.hasMoreRows = hasMoreRows; + TRowSet rowSet = new TRowSet(); + rowSet.setColumns(Collections.emptyList()); + response.setResults(rowSet); + return response; + } + + private TFetchResultsResp createResponseWithData(int rowCount, boolean hasMoreRows) { + TFetchResultsResp response = new TFetchResultsResp(); + response.hasMoreRows = hasMoreRows; + + TRowSet rowSet = new TRowSet(); + List columns = new ArrayList<>(); + + TColumn column = new TColumn(); + TStringColumn stringCol = new TStringColumn(); + List values = new ArrayList<>(); + for (int i = 0; i < rowCount; i++) { + values.add("value_" + i); + } + stringCol.setValues(values); + column.setStringVal(stringCol); + columns.add(column); + + rowSet.setColumns(columns); + response.setResults(rowSet); + return response; + } +} diff --git a/src/test/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessorTest.java b/src/test/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessorTest.java new file mode 100644 index 0000000000..f46bb83c62 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessorTest.java @@ -0,0 +1,230 @@ +package com.databricks.jdbc.api.impl.streaming; + +import static org.junit.jupiter.api.Assertions.*; + +import com.databricks.jdbc.api.impl.arrow.ArrowResultChunk; +import com.databricks.jdbc.dbclient.impl.common.StatementId; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.model.client.thrift.generated.*; +import java.io.ByteArrayOutputStream; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowStreamWriter; +import org.junit.jupiter.api.Test; + +/** Unit tests for InlineArrowResponseProcessor. Tests schema caching and Arrow chunk creation. */ +public class InlineArrowResponseProcessorTest { + + private static final StatementId STATEMENT_ID = new StatementId("test_statement"); + + @Test + void testProcessInitialResponseCachesSchema() throws DatabricksSQLException { + InlineArrowResponseProcessor processor = new InlineArrowResponseProcessor(STATEMENT_ID); + + TFetchResultsResp initialResponse = createArrowFetchResponse(2, true); + StreamingBatch batch = processor.processInitialResponse(initialResponse); + + assertNotNull(batch); + assertEquals(0, batch.getBatchIndex()); + assertEquals(0, batch.getRowOffset()); + assertEquals(2, batch.getRowCount()); + assertTrue(batch.hasMoreRows()); + assertTrue(batch.isReady()); + + ArrowResultChunk chunk = batch.getData(); + assertNotNull(chunk); + + // Clean up native memory + batch.release(); + } + + @Test + void testProcessSubsequentResponseUsesSchemaCache() throws DatabricksSQLException { + InlineArrowResponseProcessor processor = new InlineArrowResponseProcessor(STATEMENT_ID); + + // First, process initial response to cache schema + TFetchResultsResp initialResponse = createArrowFetchResponse(2, true); + StreamingBatch batch1 = processor.processInitialResponse(initialResponse); + + // Now process a subsequent response (simulating second batch) + // The schema should be cached from the initial response + TFetchResultsResp subsequentResponse = createArrowFetchResponseNoSchema(3, false); + StreamingBatch batch2 = processor.processResponse(subsequentResponse, 1, 100); + + assertNotNull(batch2); + assertEquals(1, batch2.getBatchIndex()); + assertEquals(100, batch2.getRowOffset()); + assertEquals(3, batch2.getRowCount()); + assertFalse(batch2.hasMoreRows()); + + // Clean up native memory + batch1.release(); + batch2.release(); + } + + @Test + void testReleaseActionReleasesNativeMemory() throws DatabricksSQLException { + InlineArrowResponseProcessor processor = new InlineArrowResponseProcessor(STATEMENT_ID); + + TFetchResultsResp response = createArrowFetchResponse(1, false); + StreamingBatch batch = processor.processInitialResponse(response); + + ArrowResultChunk chunk = batch.getData(); + assertNotNull(chunk); + + // Release should clean up native memory without throwing + batch.release(); + + assertEquals(StreamingBatch.Status.RELEASED, batch.getStatus()); + assertNull(batch.getData()); + } + + @Test + void testGetReleaseAction() throws DatabricksSQLException { + InlineArrowResponseProcessor processor = new InlineArrowResponseProcessor(STATEMENT_ID); + + var releaseAction = processor.getReleaseAction(); + assertNotNull(releaseAction); + + // Should handle null gracefully + assertDoesNotThrow(() -> releaseAction.accept(null)); + + // Create a real chunk and verify release works + TFetchResultsResp response = createArrowFetchResponse(1, false); + StreamingBatch batch = processor.processInitialResponse(response); + ArrowResultChunk chunk = batch.getData(); + + // Release should work without throwing + assertDoesNotThrow(() -> releaseAction.accept(chunk)); + } + + @Test + void testProcessInitialResponseWithNullStatementId() throws DatabricksSQLException { + // Processor should work even with null statement ID + InlineArrowResponseProcessor processor = new InlineArrowResponseProcessor(null); + + TFetchResultsResp response = createArrowFetchResponse(1, false); + StreamingBatch batch = processor.processInitialResponse(response); + + assertNotNull(batch); + assertTrue(batch.isReady()); + + batch.release(); + } + + @Test + void testBatchMetadataCorrectlyPropagated() throws DatabricksSQLException { + InlineArrowResponseProcessor processor = new InlineArrowResponseProcessor(STATEMENT_ID); + + // Test with hasMoreRows = true + TFetchResultsResp responseWithMore = createArrowFetchResponse(5, true); + StreamingBatch batch1 = processor.processInitialResponse(responseWithMore); + assertEquals(5, batch1.getRowCount()); + assertTrue(batch1.hasMoreRows()); + + // Test with hasMoreRows = false + TFetchResultsResp responseNoMore = createArrowFetchResponseNoSchema(3, false); + StreamingBatch batch2 = processor.processResponse(responseNoMore, 1, 5); + assertEquals(3, batch2.getRowCount()); + assertFalse(batch2.hasMoreRows()); + + batch1.release(); + batch2.release(); + } + + // ==================== Helper Methods ==================== + + /** Creates an Arrow response with full schema (for initial response). */ + private TFetchResultsResp createArrowFetchResponse(int rowCount, boolean hasMoreRows) { + TFetchResultsResp response = new TFetchResultsResp(); + response.hasMoreRows = hasMoreRows; + + byte[] arrowData = createArrowBytes(rowCount); + + TGetResultSetMetadataResp metadata = new TGetResultSetMetadataResp(); + metadata.setResultFormat(TSparkRowSetType.ARROW_BASED_SET); + metadata.setArrowSchema(new byte[0]); // Empty - schema is in the IPC stream + metadata.setSchema(createTableSchema()); + response.setResultSetMetadata(metadata); + + TSparkArrowBatch arrowBatch = new TSparkArrowBatch(); + arrowBatch.setRowCount(rowCount); + arrowBatch.setBatch(arrowData); + + TRowSet rowSet = new TRowSet(); + rowSet.setArrowBatches(Collections.singletonList(arrowBatch)); + response.setResults(rowSet); + + return response; + } + + /** + * Creates an Arrow response for subsequent batches (minimal metadata, schema cached from + * initial). + */ + private TFetchResultsResp createArrowFetchResponseNoSchema(int rowCount, boolean hasMoreRows) { + TFetchResultsResp response = new TFetchResultsResp(); + response.hasMoreRows = hasMoreRows; + + byte[] arrowData = createArrowBytes(rowCount); + + // Minimal metadata required for processing (processor uses cached schema) + TGetResultSetMetadataResp metadata = new TGetResultSetMetadataResp(); + metadata.setResultFormat(TSparkRowSetType.ARROW_BASED_SET); + response.setResultSetMetadata(metadata); + + TSparkArrowBatch arrowBatch = new TSparkArrowBatch(); + arrowBatch.setRowCount(rowCount); + arrowBatch.setBatch(arrowData); + + TRowSet rowSet = new TRowSet(); + rowSet.setArrowBatches(Collections.singletonList(arrowBatch)); + response.setResults(rowSet); + + return response; + } + + /** Creates valid Arrow IPC format bytes with the specified row count. */ + private byte[] createArrowBytes(int rowCount) { + try (BufferAllocator allocator = new RootAllocator(); + ByteArrayOutputStream out = new ByteArrayOutputStream(); + IntVector intVector = new IntVector("col_0", allocator)) { + + intVector.allocateNew(rowCount); + for (int i = 0; i < rowCount; i++) { + intVector.set(i, i * 100); + } + intVector.setValueCount(rowCount); + + try (VectorSchemaRoot root = VectorSchemaRoot.of(intVector); + ArrowStreamWriter writer = new ArrowStreamWriter(root, null, out)) { + writer.start(); + root.setRowCount(rowCount); + writer.writeBatch(); + writer.end(); + } + + return out.toByteArray(); + } catch (Exception e) { + throw new RuntimeException("Failed to create Arrow data", e); + } + } + + /** Creates a TTableSchema with a single INT column. */ + private TTableSchema createTableSchema() { + List columns = new ArrayList<>(); + TPrimitiveTypeEntry primitiveType = new TPrimitiveTypeEntry().setType(TTypeId.INT_TYPE); + TTypeEntry typeEntry = new TTypeEntry(); + typeEntry.setPrimitiveEntry(primitiveType); + TTypeDesc typeDesc = new TTypeDesc().setTypes(Collections.singletonList(typeEntry)); + TColumnDesc columnDesc = + new TColumnDesc().setColumnName("col_0").setTypeDesc(typeDesc).setPosition(0); + columns.add(columnDesc); + return new TTableSchema().setColumns(columns); + } +} diff --git a/src/test/java/com/databricks/jdbc/api/impl/streaming/StreamingBatchTest.java b/src/test/java/com/databricks/jdbc/api/impl/streaming/StreamingBatchTest.java new file mode 100644 index 0000000000..4741cc7877 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/api/impl/streaming/StreamingBatchTest.java @@ -0,0 +1,159 @@ +package com.databricks.jdbc.api.impl.streaming; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.jupiter.api.Test; + +/** Unit tests for StreamingBatch state management and lifecycle. */ +public class StreamingBatchTest { + + @Test + void testInitialState() { + StreamingBatch batch = new StreamingBatch<>(0, 0, s -> {}); + + assertEquals(StreamingBatch.Status.PENDING, batch.getStatus()); + assertEquals(0, batch.getBatchIndex()); + assertEquals(0, batch.getRowOffset()); + assertNull(batch.getData()); + assertFalse(batch.isReady()); + } + + @Test + void testSetFetching() { + StreamingBatch batch = new StreamingBatch<>(1, 100, s -> {}); + + batch.setFetching(); + + assertEquals(StreamingBatch.Status.FETCHING, batch.getStatus()); + assertFalse(batch.isReady()); + } + + @Test + void testSetData() { + StreamingBatch batch = new StreamingBatch<>(2, 200, s -> {}); + + batch.setData("test_data", 50, true); + + assertEquals(StreamingBatch.Status.READY, batch.getStatus()); + assertTrue(batch.isReady()); + assertEquals("test_data", batch.getData()); + assertEquals(50, batch.getRowCount()); + assertTrue(batch.hasMoreRows()); + } + + @Test + void testSetError() { + StreamingBatch batch = new StreamingBatch<>(3, 300, s -> {}); + Exception testError = new RuntimeException("Test error"); + + batch.setError(testError); + + assertEquals(StreamingBatch.Status.ERROR, batch.getStatus()); + assertFalse(batch.isReady()); + assertEquals(testError, batch.getError()); + } + + @Test + void testRelease() { + AtomicBoolean releaseActionCalled = new AtomicBoolean(false); + StreamingBatch batch = new StreamingBatch<>(4, 400, s -> releaseActionCalled.set(true)); + + batch.setData("to_be_released", 10, false); + batch.release(); + + assertEquals(StreamingBatch.Status.RELEASED, batch.getStatus()); + assertNull(batch.getData()); + assertTrue(releaseActionCalled.get()); + } + + @Test + void testReleaseWithNullData() { + AtomicBoolean releaseActionCalled = new AtomicBoolean(false); + StreamingBatch batch = new StreamingBatch<>(5, 500, s -> releaseActionCalled.set(true)); + + // Release without setting data + batch.release(); + + assertEquals(StreamingBatch.Status.RELEASED, batch.getStatus()); + assertFalse(releaseActionCalled.get()); // Should not call release action for null data + } + + @Test + void testWaitUntilReadyCompletesWhenDataSet() throws Exception { + StreamingBatch batch = new StreamingBatch<>(6, 600, s -> {}); + + // Set data in another thread + Thread setter = + new Thread( + () -> { + try { + Thread.sleep(50); + batch.setData("async_data", 5, false); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + setter.start(); + + // Wait should complete successfully + batch.waitUntilReady(5); + + assertTrue(batch.isReady()); + assertEquals("async_data", batch.getData()); + } + + @Test + void testWaitUntilReadyThrowsOnError() { + StreamingBatch batch = new StreamingBatch<>(7, 700, s -> {}); + RuntimeException testError = new RuntimeException("Fetch failed"); + + // Set error in another thread + Thread errorSetter = + new Thread( + () -> { + try { + Thread.sleep(50); + batch.setError(testError); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + errorSetter.start(); + + // Wait should throw ExecutionException + assertThrows(ExecutionException.class, () -> batch.waitUntilReady(5)); + } + + @Test + void testWaitUntilReadyTimeout() { + StreamingBatch batch = new StreamingBatch<>(8, 800, s -> {}); + + // Never set data - should timeout + assertThrows(TimeoutException.class, () -> batch.waitUntilReady(1)); + } + + @Test + void testStateTransitionFullLifecycle() { + AtomicBoolean released = new AtomicBoolean(false); + StreamingBatch batch = new StreamingBatch<>(9, 900, s -> released.set(true)); + + // PENDING + assertEquals(StreamingBatch.Status.PENDING, batch.getStatus()); + + // PENDING -> FETCHING + batch.setFetching(); + assertEquals(StreamingBatch.Status.FETCHING, batch.getStatus()); + + // FETCHING -> READY + batch.setData("lifecycle_data", 25, false); + assertEquals(StreamingBatch.Status.READY, batch.getStatus()); + + // READY -> RELEASED + batch.release(); + assertEquals(StreamingBatch.Status.RELEASED, batch.getStatus()); + assertTrue(released.get()); + } +} diff --git a/src/test/java/com/databricks/jdbc/api/impl/streaming/ThriftStreamingProviderTest.java b/src/test/java/com/databricks/jdbc/api/impl/streaming/ThriftStreamingProviderTest.java new file mode 100644 index 0000000000..88e4b571f3 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/api/impl/streaming/ThriftStreamingProviderTest.java @@ -0,0 +1,572 @@ +package com.databricks.jdbc.api.impl.streaming; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.databricks.jdbc.api.impl.ColumnarRowView; +import com.databricks.jdbc.api.impl.thrift.ThriftBatchFetcher; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.model.client.thrift.generated.*; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.sql.SQLException; +import java.util.*; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +/** + * Tests for streaming-specific behavior that doesn't exist in the lazy implementations. These tests + * cover the sliding window, batch release, and prefetch error handling. + */ +@ExtendWith(MockitoExtension.class) +public class ThriftStreamingProviderTest { + + @Mock private ThriftBatchFetcher batchFetcher; + + @Test + void testSingleBatchNoMoreRows() throws SQLException { + TFetchResultsResp response = createResponseWithStringData(2, false); + + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, response, 3, 30); + + try { + // Move to first batch + assertTrue(provider.hasNextBatch()); + assertTrue(provider.nextBatch()); + + // Initial batch has data + StreamingBatch batch = provider.getCurrentBatch(); + assertNotNull(batch); + assertEquals(2, batch.getRowCount()); + assertFalse(batch.hasMoreRows()); + + // No more batches since hasMoreRows was false + assertFalse(provider.hasNextBatch()); + assertFalse(provider.nextBatch()); + + // Should be at end of stream + assertTrue(provider.isEndOfStreamReached()); + } finally { + provider.close(); + } + } + + @Test + void testSlidingWindowBoundsMemory() throws SQLException, InterruptedException { + int maxBatchesInMemory = 3; + + // Initial batch + TFetchResultsResp initialResponse = createResponseWithStringData(2, true); + + // Subsequent batches + TFetchResultsResp batch2 = createResponseWithStringData(2, true); + TFetchResultsResp batch3 = createResponseWithStringData(2, true); + TFetchResultsResp batch4 = createResponseWithStringData(2, true); + TFetchResultsResp batch5 = createResponseWithStringData(2, false); + + when(batchFetcher.fetchNextBatch()) + .thenReturn(batch2) + .thenReturn(batch3) + .thenReturn(batch4) + .thenReturn(batch5); + + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, maxBatchesInMemory, 30); + + try { + // Let prefetch thread run a bit + Thread.sleep(100); + + // Provider should never have more than maxBatchesInMemory batches + assertTrue( + provider.getBatchesInMemory() <= maxBatchesInMemory, + "Batches in memory (" + + provider.getBatchesInMemory() + + ") should not exceed max (" + + maxBatchesInMemory + + ")"); + + // Consume batches and verify memory is bounded + provider.nextBatch(); // batch 0 + Thread.sleep(50); + assertTrue(provider.getBatchesInMemory() <= maxBatchesInMemory); + + provider.nextBatch(); // batch 1 (releases batch 0) + Thread.sleep(50); + assertTrue(provider.getBatchesInMemory() <= maxBatchesInMemory); + + provider.nextBatch(); // batch 2 (releases batch 1) + Thread.sleep(50); + assertTrue(provider.getBatchesInMemory() <= maxBatchesInMemory); + + } finally { + provider.close(); + } + } + + @Test + void testEmptyBatchesSkippedByNextBatch() throws SQLException { + // Initial response with empty batch + TFetchResultsResp emptyBatch = createEmptyResponse(true); + + // Second batch also empty + TFetchResultsResp emptyBatch2 = createEmptyResponse(true); + + // Third batch has data + TFetchResultsResp dataBatch = createResponseWithStringData(3, false); + + when(batchFetcher.fetchNextBatch()).thenReturn(emptyBatch2).thenReturn(dataBatch); + + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, emptyBatch, 3, 30); + + try { + // nextBatch should skip empty batches and return when it finds data + assertTrue(provider.nextBatch()); + + StreamingBatch batch = provider.getCurrentBatch(); + assertNotNull(batch); + assertTrue(batch.getRowCount() > 0); + assertEquals(3, batch.getRowCount()); + } finally { + provider.close(); + } + } + + @Test + void testPrefetchErrorPropagated() throws SQLException, InterruptedException { + TFetchResultsResp initialResponse = createResponseWithStringData(2, true); + + DatabricksSQLException fetchError = + new DatabricksSQLException("Network failure", DatabricksDriverErrorCode.CONNECTION_ERROR); + when(batchFetcher.fetchNextBatch()).thenThrow(fetchError); + + // The prefetch thread starts during construction and may fail before or after + // the first nextBatch() call. We need to handle both cases. + DatabricksSQLException caughtException = null; + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, 3, 30); + + try { + // Consume batches until we hit the error + while (provider.nextBatch()) { + // Keep iterating - error will be thrown when we need a batch that failed to prefetch + } + } catch (DatabricksSQLException e) { + caughtException = e; + } finally { + provider.close(); + } + + // Verify we caught the expected error + assertNotNull(caughtException, "Expected DatabricksSQLException to be thrown"); + assertTrue( + caughtException.getMessage().contains("Prefetch failed"), + "Exception should contain 'Prefetch failed': " + caughtException.getMessage()); + } + + @Test + void testCloseStopsPrefetchAndClosesFetcher() throws SQLException, InterruptedException { + TFetchResultsResp initialResponse = createResponseWithStringData(2, true); + TFetchResultsResp batch2 = createResponseWithStringData(2, false); + + when(batchFetcher.fetchNextBatch()).thenReturn(batch2); + + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, 3, 30); + + // Let prefetch run + Thread.sleep(100); + + assertTrue(provider.getBatchesInMemory() > 0); + + provider.close(); + + // Verify fetcher was closed + verify(batchFetcher).close(); + + // After close, hasNextBatch should return false + assertFalse(provider.hasNextBatch()); + } + + @Test + void testNullInputValidation() { + TFetchResultsResp validResponse = createResponseWithStringData(1, false); + + // Null response + assertThrows( + IllegalArgumentException.class, + () -> ThriftStreamingProvider.forColumnar(batchFetcher, null, 3, 30)); + + // Null fetcher + assertThrows( + IllegalArgumentException.class, + () -> ThriftStreamingProvider.forColumnar(null, validResponse, 3, 30)); + } + + @Test + void testWaitForBatchCreationTimeout() throws SQLException { + // Initial batch with hasMoreRows=true so prefetch thread starts + TFetchResultsResp initialResponse = createResponseWithStringData(1, true); + + // Mock fetcher to block indefinitely (simulating slow server) + when(batchFetcher.fetchNextBatch()) + .thenAnswer( + inv -> { + Thread.sleep(10000); // Sleep longer than timeout + return createResponseWithStringData(1, false); + }); + + // Use very short timeout (1 second) to trigger timeout quickly + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, 3, 1); + + try { + // Consume initial batch + provider.nextBatch(); + + // Try to get next batch - should timeout waiting for batch creation + DatabricksSQLException thrown = + assertThrows(DatabricksSQLException.class, provider::nextBatch); + assertTrue( + thrown.getMessage().contains("Timeout") || thrown.getMessage().contains("timeout"), + "Expected timeout error but got: " + thrown.getMessage()); + } finally { + provider.close(); + } + } + + @Test + void testWaitForBatchInterrupted() throws SQLException, InterruptedException { + TFetchResultsResp initialResponse = createResponseWithStringData(1, true); + + // Mock fetcher to block + when(batchFetcher.fetchNextBatch()) + .thenAnswer( + inv -> { + Thread.sleep(10000); + return createResponseWithStringData(1, false); + }); + + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, 3, 30); + + // Start a thread that will try to get next batch + Thread testThread = + new Thread( + () -> { + try { + provider.nextBatch(); // Initial batch + provider.nextBatch(); // Will block waiting for batch 1 + } catch (DatabricksSQLException e) { + // Expected - interrupted + assertTrue( + e.getMessage().contains("Interrupt") || e.getMessage().contains("interrupt")); + } + }); + + try { + testThread.start(); + Thread.sleep(100); // Let the thread start waiting + + // Interrupt the waiting thread + testThread.interrupt(); + testThread.join(2000); + + assertFalse(testThread.isAlive(), "Test thread should have terminated"); + } finally { + provider.close(); + } + } + + @Test + void testGetCurrentBatchWaitsForReady() throws SQLException, InterruptedException { + TFetchResultsResp initialResponse = createResponseWithStringData(2, true); + TFetchResultsResp batch2 = createResponseWithStringData(2, false); + + // Add a small delay to batch fetching + when(batchFetcher.fetchNextBatch()) + .thenAnswer( + inv -> { + Thread.sleep(50); + return batch2; + }); + + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, 3, 30); + + try { + // Move to initial batch + assertTrue(provider.nextBatch()); + StreamingBatch batch = provider.getCurrentBatch(); + assertNotNull(batch); + assertEquals(2, batch.getRowCount()); + + // Move to next batch - should wait for it to be ready + assertTrue(provider.nextBatch()); + batch = provider.getCurrentBatch(); + assertNotNull(batch); + assertEquals(2, batch.getRowCount()); + } finally { + provider.close(); + } + } + + @Test + void testCloseWhileWaiting() throws SQLException, InterruptedException { + TFetchResultsResp initialResponse = createResponseWithStringData(1, true); + + // Mock fetcher to block indefinitely + when(batchFetcher.fetchNextBatch()) + .thenAnswer( + inv -> { + Thread.sleep(10000); + return createResponseWithStringData(1, false); + }); + + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, 3, 30); + + // Start a thread that will wait for next batch + Thread waitingThread = + new Thread( + () -> { + try { + provider.nextBatch(); + provider.nextBatch(); // Will block + } catch (DatabricksSQLException e) { + // Expected after close + } + }); + + try { + waitingThread.start(); + Thread.sleep(100); // Let thread start waiting + + // Close should unblock the waiting thread + provider.close(); + + waitingThread.join(2000); + assertFalse(waitingThread.isAlive(), "Waiting thread should have terminated after close"); + } finally { + if (waitingThread.isAlive()) { + waitingThread.interrupt(); + } + } + } + + @Test + void testGetCurrentBatchBeforeFirstBatch() throws SQLException { + TFetchResultsResp initialResponse = createResponseWithStringData(2, false); + + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, 3, 30); + + try { + // Before calling nextBatch(), currentBatchIndex is -1 + // getCurrentBatch should return null + StreamingBatch batch = provider.getCurrentBatch(); + assertNull(batch, "getCurrentBatch should return null before first nextBatch()"); + } finally { + provider.close(); + } + } + + @Test + void testGetCurrentBatchAfterNextBatch() throws SQLException { + TFetchResultsResp initialResponse = createResponseWithStringData(2, false); + + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, 3, 30); + + try { + // Move to first batch + assertTrue(provider.nextBatch()); + + // getCurrentBatch should return the batch + StreamingBatch batch = provider.getCurrentBatch(); + assertNotNull(batch); + assertEquals(0, batch.getBatchIndex()); + assertEquals(2, batch.getRowCount()); + assertTrue(batch.isReady()); + } finally { + provider.close(); + } + } + + @Test + void testGetCurrentBatchWithPrefetchError() throws SQLException, InterruptedException { + TFetchResultsResp initialResponse = createResponseWithStringData(1, true); + + DatabricksSQLException fetchError = + new DatabricksSQLException("Server error", DatabricksDriverErrorCode.CONNECTION_ERROR); + when(batchFetcher.fetchNextBatch()).thenThrow(fetchError); + + // The prefetch thread starts during construction and may fail at any point. + DatabricksSQLException caughtException = null; + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, 3, 30); + + try { + // Move to initial batch + provider.nextBatch(); + + // Let prefetch thread fail + Thread.sleep(100); + + // Either nextBatch or getCurrentBatch should propagate the prefetch error + while (provider.nextBatch()) { + provider.getCurrentBatch(); + } + } catch (DatabricksSQLException e) { + caughtException = e; + } finally { + provider.close(); + } + + // Verify we caught the expected error + assertNotNull(caughtException, "Expected DatabricksSQLException to be thrown"); + assertTrue( + caughtException.getMessage().contains("Prefetch failed"), + "Exception should contain 'Prefetch failed': " + caughtException.getMessage()); + } + + @Test + void testGetCurrentBatchMultipleCalls() throws SQLException { + TFetchResultsResp initialResponse = createResponseWithStringData(2, false); + + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, 3, 30); + + try { + provider.nextBatch(); + + // Multiple calls to getCurrentBatch should return same batch + StreamingBatch batch1 = provider.getCurrentBatch(); + StreamingBatch batch2 = provider.getCurrentBatch(); + StreamingBatch batch3 = provider.getCurrentBatch(); + + assertSame(batch1, batch2); + assertSame(batch2, batch3); + } finally { + provider.close(); + } + } + + @Test + void testBatchReadyTimeout() throws SQLException, InterruptedException { + TFetchResultsResp initialResponse = createResponseWithStringData(1, true); + + // Create a batch that will never become ready by making fetcher very slow + when(batchFetcher.fetchNextBatch()) + .thenAnswer( + inv -> { + // Simulate a batch that takes too long + Thread.sleep(5000); + return createResponseWithStringData(1, false); + }); + + // Use 1 second timeout + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, 3, 1); + + try { + // Consume initial batch + provider.nextBatch(); + + // Try to move to next batch - should timeout + DatabricksSQLException thrown = + assertThrows(DatabricksSQLException.class, provider::nextBatch); + assertTrue( + thrown.getMessage().toLowerCase().contains("timeout"), + "Expected timeout error: " + thrown.getMessage()); + } finally { + provider.close(); + } + } + + @Test + void testTotalRowsFetched() throws SQLException, InterruptedException { + TFetchResultsResp initialResponse = createResponseWithStringData(3, true); + TFetchResultsResp batch2 = createResponseWithStringData(5, false); + + when(batchFetcher.fetchNextBatch()).thenReturn(batch2); + + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, 3, 30); + + try { + // Initial batch has at least 3 rows (prefetch may have already fetched more) + assertTrue(provider.getTotalRowsFetched() >= 3); + + // Let prefetch run + Thread.sleep(100); + + // After prefetch, should have 8 total rows (3 + 5) + assertEquals(8, provider.getTotalRowsFetched()); + } finally { + provider.close(); + } + } + + @Test + void testBatchesInMemoryCount() throws SQLException, InterruptedException { + TFetchResultsResp initialResponse = createResponseWithStringData(1, true); + TFetchResultsResp batch2 = createResponseWithStringData(1, true); + TFetchResultsResp batch3 = createResponseWithStringData(1, false); + + when(batchFetcher.fetchNextBatch()).thenReturn(batch2).thenReturn(batch3); + + ThriftStreamingProvider provider = + ThriftStreamingProvider.forColumnar(batchFetcher, initialResponse, 3, 30); + + try { + // Initial: at least 1 batch in memory (prefetch may have already fetched more) + assertTrue(provider.getBatchesInMemory() >= 1); + + // Let prefetch run + Thread.sleep(100); + + // After prefetch of 2 more batches (3 total), should be within bounds + assertTrue(provider.getBatchesInMemory() >= 1); + assertTrue(provider.getBatchesInMemory() <= 3); + } finally { + provider.close(); + } + } + + // ==================== Helper Methods ==================== + + private TFetchResultsResp createEmptyResponse(boolean hasMoreRows) { + TFetchResultsResp response = new TFetchResultsResp(); + response.hasMoreRows = hasMoreRows; + TRowSet emptyRowSet = new TRowSet(); + emptyRowSet.setColumns(Collections.emptyList()); + response.setResults(emptyRowSet); + return response; + } + + private TFetchResultsResp createResponseWithStringData(int rowCount, boolean hasMoreRows) { + TFetchResultsResp response = new TFetchResultsResp(); + response.hasMoreRows = hasMoreRows; + + TRowSet rowSet = new TRowSet(); + List columns = new ArrayList<>(); + + TColumn column = new TColumn(); + TStringColumn stringCol = new TStringColumn(); + List values = new ArrayList<>(); + for (int i = 0; i < rowCount; i++) { + values.add("value_" + i); + } + stringCol.setValues(values); + column.setStringVal(stringCol); + columns.add(column); + + rowSet.setColumns(columns); + response.setResults(rowSet); + return response; + } +} diff --git a/src/test/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResultTest.java b/src/test/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResultTest.java new file mode 100644 index 0000000000..f044ec8546 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResultTest.java @@ -0,0 +1,764 @@ +package com.databricks.jdbc.api.impl.thrift; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +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.dbclient.IDatabricksClient; +import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.model.client.thrift.generated.*; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import java.sql.SQLException; +import java.util.*; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +/** + * Unit tests for StreamingColumnarResult. Tests functional parity with LazyThriftResult from the + * consumer's perspective, along with streaming-specific behaviors (such as background prefetch). + */ +@ExtendWith(MockitoExtension.class) +public class StreamingColumnarResultTest { + + @Mock private IDatabricksSession session; + @Mock private IDatabricksStatementInternal statement; + @Mock private IDatabricksClient databricksClient; + @Mock private IDatabricksConnectionContext connectionContext; + + @BeforeEach + void setUp() throws SQLException { + lenient().when(session.getDatabricksClient()).thenReturn(databricksClient); + lenient().when(session.getConnectionContext()).thenReturn(connectionContext); + lenient().when(connectionContext.getThriftMaxBatchesInMemory()).thenReturn(3); + lenient().when(statement.getMaxRows()).thenReturn(0); // No limit by default + } + + @Test + void testBasicIteration() throws SQLException { + TFetchResultsResp response = + createResponseWithStringData( + Arrays.asList("row1_col1", "row1_col2"), + Arrays.asList("row2_col1", "row2_col2"), + false); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + + try { + // Initial state + assertEquals(-1, result.getCurrentRow()); + assertEquals(2, result.getRowCount()); + assertTrue(result.hasNext()); + + // First row + assertTrue(result.next()); + assertEquals(0, result.getCurrentRow()); + assertEquals("row1_col1", result.getObject(0)); + assertEquals("row1_col2", result.getObject(1)); + + // Second row + assertTrue(result.next()); + assertEquals(1, result.getCurrentRow()); + assertEquals("row2_col1", result.getObject(0)); + assertEquals("row2_col2", result.getObject(1)); + + // End + assertFalse(result.hasNext()); + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testMultiBatchFetching() throws SQLException, InterruptedException { + TFetchResultsResp firstBatch = + createResponseWithStringData( + Arrays.asList("row1_col1", "row1_col2"), Arrays.asList("row2_col1", "row2_col2"), true); + + TFetchResultsResp secondBatch = + createResponseWithStringData( + Arrays.asList("row3_col1", "row3_col2"), + Arrays.asList("row4_col1", "row4_col2"), + false); + + when(databricksClient.getMoreResults(statement)).thenReturn(secondBatch); + + StreamingColumnarResult result = new StreamingColumnarResult(firstBatch, statement, session); + + try { + // Give prefetch thread time to start + Thread.sleep(100); + + // Consume first batch + assertTrue(result.next()); + assertEquals("row1_col1", result.getObject(0)); + assertTrue(result.next()); + assertEquals("row2_col1", result.getObject(0)); + + // Move to second batch + assertTrue(result.next()); + assertEquals(2, result.getCurrentRow()); + assertEquals("row3_col1", result.getObject(0)); + + assertTrue(result.next()); + assertEquals("row4_col1", result.getObject(0)); + + // End + assertFalse(result.next()); + + verify(databricksClient, atLeastOnce()).getMoreResults(statement); + } finally { + result.close(); + } + } + + @Test + void testMaxRowsLimit() throws SQLException { + when(statement.getMaxRows()).thenReturn(2); + + // Use hasMoreRows=false so the prefetch thread doesn't try to fetch + TFetchResultsResp response = + createResponseWithStringData( + Arrays.asList("row1_col1", "row1_col2"), + Arrays.asList("row2_col1", "row2_col2"), + Arrays.asList("row3_col1", "row3_col2"), + false); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + + try { + // Should only get 2 rows due to maxRows limit + assertTrue(result.next()); + assertEquals("row1_col1", result.getObject(0)); + + assertTrue(result.next()); + assertEquals("row2_col1", result.getObject(0)); + + // Should stop due to maxRows limit even though there's more data + assertFalse(result.hasNext()); + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testAccessAfterClose() throws SQLException { + TFetchResultsResp response = + createResponseWithStringData(Arrays.asList("row1_col1", "row1_col2"), false); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + result.next(); + result.close(); + + assertThrows(DatabricksSQLException.class, () -> result.getObject(0)); + assertFalse(result.hasNext()); + assertFalse(result.next()); + } + + @Test + void testAccessBeforeFirstRow() throws SQLException { + TFetchResultsResp response = + createResponseWithStringData(Arrays.asList("row1_col1", "row1_col2"), false); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + + try { + DatabricksSQLException thrown = + assertThrows(DatabricksSQLException.class, () -> result.getObject(0)); + assertTrue(thrown.getMessage().contains("before first row")); + } finally { + result.close(); + } + } + + @Test + void testInvalidColumnIndex() throws SQLException { + TFetchResultsResp response = createResponseWithStringData(Arrays.asList("col1", "col2"), false); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + + try { + result.next(); + + // Valid indices + assertDoesNotThrow(() -> result.getObject(0)); + assertDoesNotThrow(() -> result.getObject(1)); + + // Invalid indices + assertThrows(DatabricksSQLException.class, () -> result.getObject(2)); + assertThrows(DatabricksSQLException.class, () -> result.getObject(-1)); + } finally { + result.close(); + } + } + + @Test + void testNullHandling() throws SQLException { + TFetchResultsResp response = createResponseWithNulls(); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + + try { + assertTrue(result.next()); + assertEquals("value1", result.getObject(0)); + assertNull(result.getObject(1)); + } finally { + result.close(); + } + } + + @Test + void testChunkCount() throws SQLException { + TFetchResultsResp response = + createResponseWithStringData(Arrays.asList("row1_col1", "row1_col2"), false); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + + try { + // StreamingColumnarResult doesn't use chunks like Arrow results + assertEquals(0, result.getChunkCount()); + } finally { + result.close(); + } + } + + @Test + void testSingleRowResult() throws SQLException { + // Test single row iteration (empty results have different behavior in streaming due to init) + TFetchResultsResp response = + createResponseWithStringData(Arrays.asList("only_row_col1", "only_row_col2"), false); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + + try { + assertEquals(1, result.getRowCount()); + assertTrue(result.hasNext()); + assertTrue(result.next()); + assertEquals(0, result.getCurrentRow()); + assertEquals("only_row_col1", result.getObject(0)); + + assertFalse(result.hasNext()); + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testErrorDuringFetch() throws SQLException, InterruptedException { + TFetchResultsResp firstBatch = + createResponseWithStringData(Arrays.asList("row1_col1", "row1_col2"), true); + + DatabricksSQLException expectedException = + new DatabricksSQLException("Network error", DatabricksDriverErrorCode.CONNECTION_ERROR); + when(databricksClient.getMoreResults(statement)).thenThrow(expectedException); + + // The prefetch thread starts during construction and may fail at any point. + // The error can surface during construction (if prefetch thread runs fast) or during iteration. + DatabricksSQLException caughtException = null; + StreamingColumnarResult result = null; + + try { + result = new StreamingColumnarResult(firstBatch, statement, session); + // Consume rows until we hit the error + while (result.next()) { + // Keep iterating - error will be thrown when we need a batch that failed to prefetch + } + } catch (DatabricksSQLException e) { + caughtException = e; + } finally { + if (result != null) { + result.close(); + } + } + + // Verify we caught the expected error + assertNotNull(caughtException, "Expected DatabricksSQLException to be thrown"); + assertTrue( + caughtException.getMessage().contains("Prefetch failed") + || caughtException.getMessage().contains("Network error"), + "Exception should contain error details: " + caughtException.getMessage()); + } + + @Test + void testMaxRowsLimitAcrossBatches() throws SQLException, InterruptedException { + // MaxRows limit of 3, spanning across 2 batches (2 rows each) + when(statement.getMaxRows()).thenReturn(3); + + TFetchResultsResp firstBatch = + createResponseWithStringData( + Arrays.asList("row1_col1", "row1_col2"), + Arrays.asList("row2_col1", "row2_col2"), + true); // hasMoreRows = true + + TFetchResultsResp secondBatch = + createResponseWithStringData( + Arrays.asList("row3_col1", "row3_col2"), + Arrays.asList("row4_col1", "row4_col2"), + false); + + when(databricksClient.getMoreResults(statement)).thenReturn(secondBatch); + + StreamingColumnarResult result = new StreamingColumnarResult(firstBatch, statement, session); + + try { + // Give prefetch thread time to start + Thread.sleep(100); + + // Consume first batch (rows 0 and 1) + assertTrue(result.next()); + assertEquals("row1_col1", result.getObject(0)); + assertTrue(result.next()); + assertEquals("row2_col1", result.getObject(0)); + + // Get one row from second batch (row 2) + assertTrue(result.next()); + assertEquals("row3_col1", result.getObject(0)); + + // Should stop at maxRows=3 + assertFalse(result.hasNext()); + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testBatchesInMemoryTracking() throws SQLException { + TFetchResultsResp response = + createResponseWithStringData( + Arrays.asList("row1_col1", "row1_col2"), + Arrays.asList("row2_col1", "row2_col2"), + false); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + + try { + // Should have at least initial batch + assertTrue(result.getBatchesInMemory() > 0); + } finally { + result.close(); + } + // Note: getBatchesInMemory() may return a small number after close due to cleanup timing + // The important thing is that close() was called without exception + } + + @Test + void testGetTotalRowsFetched() throws SQLException, InterruptedException { + TFetchResultsResp firstBatch = + createResponseWithStringData( + Arrays.asList("row1_col1", "row1_col2"), Arrays.asList("row2_col1", "row2_col2"), true); + + TFetchResultsResp secondBatch = + createResponseWithStringData( + Arrays.asList("row3_col1", "row3_col2"), + Arrays.asList("row4_col1", "row4_col2"), + false); + + when(databricksClient.getMoreResults(statement)).thenReturn(secondBatch); + + StreamingColumnarResult result = new StreamingColumnarResult(firstBatch, statement, session); + + try { + // Initial batch has at least 2 rows (prefetch may have already fetched more) + assertTrue(result.getTotalRowsFetched() >= 2); + + // Give prefetch thread time to fetch second batch + Thread.sleep(100); + + // After prefetch completes, should have 4 total rows fetched + assertEquals(4, result.getTotalRowsFetched()); + } finally { + result.close(); + } + } + + @Test + void testIsCompletelyFetched() throws SQLException, InterruptedException { + TFetchResultsResp firstBatch = + createResponseWithStringData( + Arrays.asList("row1_col1", "row1_col2"), true); // hasMoreRows = true + + TFetchResultsResp secondBatch = + createResponseWithStringData( + Arrays.asList("row2_col1", "row2_col2"), false); // hasMoreRows = false + + when(databricksClient.getMoreResults(statement)).thenReturn(secondBatch); + + StreamingColumnarResult result = new StreamingColumnarResult(firstBatch, statement, session); + + try { + // Give prefetch thread time to fetch second batch + Thread.sleep(100); + + // After fetching final batch (hasMoreRows=false), should be completely fetched + assertTrue(result.isCompletelyFetched()); + } finally { + result.close(); + } + } + + @Test + void testInitializationWithEmptyInitialBatch() throws SQLException, InterruptedException { + // Initial batch is EMPTY but hasMoreRows=true + TFetchResultsResp emptyInitial = createEmptyResponse(true); + + // Second batch has actual data + TFetchResultsResp dataBatch = + createResponseWithStringData( + Arrays.asList("row1_col1", "row1_col2"), + Arrays.asList("row2_col1", "row2_col2"), + false); + + when(databricksClient.getMoreResults(statement)).thenReturn(dataBatch); + + StreamingColumnarResult result = new StreamingColumnarResult(emptyInitial, statement, session); + + try { + // Give prefetch time to fetch the data batch + Thread.sleep(100); + + // Should have skipped empty batch and positioned on data batch + assertTrue(result.hasNext()); + assertTrue(result.next()); + assertEquals("row1_col1", result.getObject(0)); + assertEquals("row1_col2", result.getObject(1)); + + assertTrue(result.next()); + assertEquals("row2_col1", result.getObject(0)); + assertEquals("row2_col2", result.getObject(1)); + + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testGetRowCount() throws SQLException { + TFetchResultsResp response = + createResponseWithStringData( + Arrays.asList("row1_col1", "row1_col2"), + Arrays.asList("row2_col1", "row2_col2"), + Arrays.asList("row3_col1", "row3_col2"), + false); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + + try { + assertTrue(result.next()); + assertEquals(3, result.getRowCount()); + } finally { + result.close(); + } + } + + @Test + void testDoubleClose() throws SQLException { + TFetchResultsResp response = + createResponseWithStringData(Arrays.asList("row1_col1", "row1_col2"), false); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + + // First close + result.close(); + // Second close should not throw + assertDoesNotThrow(result::close); + } + + @Test + void testEmptyResultNoMoreRows() throws SQLException { + // Single row result to verify end-of-stream detection works + TFetchResultsResp singleRowResponse = + createResponseWithStringData(Arrays.asList("only_value"), false); + + StreamingColumnarResult result = + new StreamingColumnarResult(singleRowResponse, statement, session); + + try { + // Should have exactly one row + assertTrue(result.hasNext()); + assertTrue(result.next()); + assertEquals("only_value", result.getObject(0)); + + // No more rows + assertFalse(result.hasNext()); + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testHasNextAfterExhausted() throws SQLException { + TFetchResultsResp response = createResponseWithStringData(Arrays.asList("value"), false); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + + try { + // Consume all rows + assertTrue(result.next()); + assertFalse(result.next()); + + // hasNext should consistently return false after exhausted + assertFalse(result.hasNext()); + assertFalse(result.hasNext()); + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testMultipleConsecutiveEmptyBatches() throws SQLException, InterruptedException { + // First batch empty + TFetchResultsResp emptyBatch1 = createEmptyResponse(true); + // Second batch also empty + TFetchResultsResp emptyBatch2 = createEmptyResponse(true); + // Third batch has data + TFetchResultsResp dataBatch = + createResponseWithStringData( + Arrays.asList("row1_col1", "row1_col2"), + Arrays.asList("row2_col1", "row2_col2"), + false); + + when(databricksClient.getMoreResults(statement)).thenReturn(emptyBatch2).thenReturn(dataBatch); + + StreamingColumnarResult result = new StreamingColumnarResult(emptyBatch1, statement, session); + + try { + // Give prefetch time + Thread.sleep(100); + + // Should skip both empty batches and find data + assertTrue(result.hasNext()); + assertTrue(result.next()); + assertEquals("row1_col1", result.getObject(0)); + + assertTrue(result.next()); + assertEquals("row2_col1", result.getObject(0)); + + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testExhaustAllRowsInSingleBatch() throws SQLException { + // Single batch with 3 rows, no more data + TFetchResultsResp response = + createResponseWithStringData( + Arrays.asList("row1"), Arrays.asList("row2"), Arrays.asList("row3"), false); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + + try { + // Consume all 3 rows + assertTrue(result.next()); + assertEquals("row1", result.getObject(0)); + assertTrue(result.next()); + assertEquals("row2", result.getObject(0)); + assertTrue(result.next()); + assertEquals("row3", result.getObject(0)); + + // No more rows - should return false + assertFalse(result.next()); + // Calling again should still return false + assertFalse(result.next()); + + // hasNext should also be false + assertFalse(result.hasNext()); + } finally { + result.close(); + } + } + + @Test + void testExhaustAllRowsAcrossMultipleBatches() throws SQLException, InterruptedException { + // First batch with 2 rows + TFetchResultsResp firstBatch = + createResponseWithStringData( + Arrays.asList("batch1_row1"), Arrays.asList("batch1_row2"), true); + + // Second batch with 1 row - final batch + TFetchResultsResp secondBatch = + createResponseWithStringData(Arrays.asList("batch2_row1"), false); + + when(databricksClient.getMoreResults(statement)).thenReturn(secondBatch); + + StreamingColumnarResult result = new StreamingColumnarResult(firstBatch, statement, session); + + try { + // Give prefetch time + Thread.sleep(100); + + // Consume batch 1 + assertTrue(result.next()); + assertEquals("batch1_row1", result.getObject(0)); + assertTrue(result.next()); + assertEquals("batch1_row2", result.getObject(0)); + + // Consume batch 2 + assertTrue(result.next()); + assertEquals("batch2_row1", result.getObject(0)); + + // No more rows - triggers end of stream detection + assertFalse(result.next()); + assertFalse(result.hasNext()); + + // Multiple calls should consistently return false + assertFalse(result.next()); + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testBatchTransitionAtExactBoundary() throws SQLException, InterruptedException { + // First batch with exactly 1 row + TFetchResultsResp firstBatch = + createResponseWithStringData(Arrays.asList("only_row_in_batch1"), true); + + // Second batch with 1 row + TFetchResultsResp secondBatch = + createResponseWithStringData(Arrays.asList("only_row_in_batch2"), false); + + when(databricksClient.getMoreResults(statement)).thenReturn(secondBatch); + + StreamingColumnarResult result = new StreamingColumnarResult(firstBatch, statement, session); + + try { + // Give prefetch time + Thread.sleep(100); + + // Row 1 - from batch 1 + assertTrue(result.next()); + assertEquals("only_row_in_batch1", result.getObject(0)); + assertEquals(0, result.getCurrentRow()); + + // Row 2 - requires batch transition + assertTrue(result.next()); + assertEquals("only_row_in_batch2", result.getObject(0)); + assertEquals(1, result.getCurrentRow()); + + // End - no more batches + assertFalse(result.next()); + } finally { + result.close(); + } + } + + @Test + void testNextAfterEndReturnsConsistentFalse() throws SQLException { + TFetchResultsResp response = createResponseWithStringData(Arrays.asList("single_row"), false); + + StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); + + try { + // Consume the only row + assertTrue(result.next()); + + // Multiple calls to next() after exhaustion should all return false + assertFalse(result.next()); + assertFalse(result.next()); + assertFalse(result.next()); + + // hasReachedEnd should be set + assertFalse(result.hasNext()); + } finally { + result.close(); + } + } + + // ==================== Helper Methods ==================== + + private TFetchResultsResp createEmptyResponse(boolean hasMoreRows) { + TFetchResultsResp response = new TFetchResultsResp(); + response.hasMoreRows = hasMoreRows; + TRowSet emptyRowSet = new TRowSet(); + emptyRowSet.setColumns(Collections.emptyList()); + response.setResults(emptyRowSet); + return response; + } + + private TFetchResultsResp createResponseWithStringData(List row, boolean hasMoreRows) { + return createResponseWithStringRows(Arrays.asList(row), hasMoreRows); + } + + private TFetchResultsResp createResponseWithStringData( + List row1, List row2, boolean hasMoreRows) { + return createResponseWithStringRows(Arrays.asList(row1, row2), hasMoreRows); + } + + private TFetchResultsResp createResponseWithStringData( + List row1, List row2, List row3, boolean hasMoreRows) { + return createResponseWithStringRows(Arrays.asList(row1, row2, row3), hasMoreRows); + } + + private TFetchResultsResp createResponseWithStringRows( + List> rows, boolean hasMoreRows) { + TFetchResultsResp response = new TFetchResultsResp(); + response.hasMoreRows = hasMoreRows; + + if (rows.isEmpty()) { + return createEmptyResponse(hasMoreRows); + } + + TRowSet rowSet = new TRowSet(); + int numColumns = rows.get(0).size(); + List columns = new ArrayList<>(numColumns); + + for (int col = 0; col < numColumns; col++) { + TColumn column = new TColumn(); + TStringColumn stringCol = new TStringColumn(); + List colValues = new ArrayList<>(); + + for (List row : rows) { + colValues.add(col < row.size() ? row.get(col) : null); + } + + stringCol.setValues(colValues); + column.setStringVal(stringCol); + columns.add(column); + } + + rowSet.setColumns(columns); + response.setResults(rowSet); + return response; + } + + private TFetchResultsResp createResponseWithNulls() { + TFetchResultsResp response = new TFetchResultsResp(); + response.hasMoreRows = false; + + TRowSet rowSet = new TRowSet(); + List columns = new ArrayList<>(); + + // First column - no nulls + TColumn col1 = new TColumn(); + TStringColumn stringCol1 = new TStringColumn(); + stringCol1.setValues(Arrays.asList("value1")); + col1.setStringVal(stringCol1); + columns.add(col1); + + // Second column - with null + TColumn col2 = new TColumn(); + TStringColumn stringCol2 = new TStringColumn(); + stringCol2.setValues(Arrays.asList("placeholder")); + stringCol2.setNulls(new byte[] {0x01}); + col2.setStringVal(stringCol2); + columns.add(col2); + + rowSet.setColumns(columns); + response.setResults(rowSet); + return response; + } +} diff --git a/src/test/java/com/databricks/jdbc/api/impl/thrift/ThriftBatchFetcherImplTest.java b/src/test/java/com/databricks/jdbc/api/impl/thrift/ThriftBatchFetcherImplTest.java new file mode 100644 index 0000000000..43d9558014 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/api/impl/thrift/ThriftBatchFetcherImplTest.java @@ -0,0 +1,59 @@ +package com.databricks.jdbc.api.impl.thrift; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.databricks.jdbc.api.internal.IDatabricksSession; +import com.databricks.jdbc.api.internal.IDatabricksStatementInternal; +import com.databricks.jdbc.dbclient.IDatabricksClient; +import com.databricks.jdbc.exception.DatabricksSQLException; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +/** Unit tests for ThriftBatchFetcherImpl. */ +@ExtendWith(MockitoExtension.class) +public class ThriftBatchFetcherImplTest { + + @Mock private IDatabricksSession session; + @Mock private IDatabricksStatementInternal statement; + @Mock private IDatabricksClient databricksClient; + + @BeforeEach + void setUp() { + lenient().when(session.getDatabricksClient()).thenReturn(databricksClient); + } + + @Test + void testFetchNextBatchWhenClosed() throws DatabricksSQLException { + ThriftBatchFetcherImpl fetcher = new ThriftBatchFetcherImpl(session, statement); + + // Close the fetcher + fetcher.close(); + + // Attempting to fetch should throw exception + DatabricksSQLException thrown = + assertThrows(DatabricksSQLException.class, fetcher::fetchNextBatch); + assertTrue(thrown.getMessage().contains("closed")); + } + + @Test + void testCloseAndIsClosed() { + ThriftBatchFetcherImpl fetcher = new ThriftBatchFetcherImpl(session, statement); + + // Initially not closed + assertFalse(fetcher.isClosed()); + + // Close it + fetcher.close(); + + // Now should be closed + assertTrue(fetcher.isClosed()); + + // Closing again should be safe (idempotent) + assertDoesNotThrow(fetcher::close); + assertTrue(fetcher.isClosed()); + } +} diff --git a/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java b/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java index 833f7c4707..c32fc7bb89 100644 --- a/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java +++ b/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java @@ -37,7 +37,7 @@ class DatabricksDriverFeatureFlagsContextTest { @Mock private ObjectMapper objectMapperMock; private static final String FEATURE_FLAG_NAME = "featureFlagName"; private static final String FEATURE_FLAGS_ENDPOINT = - "https://test-host/api/2.0/connector-service/feature-flags/OSS_JDBC/3.0.7"; + "https://test-host/api/2.0/connector-service/feature-flags/OSS_JDBC/3.1.1"; private DatabricksDriverFeatureFlagsContext context; diff --git a/src/test/java/com/databricks/jdbc/common/util/ArrowUtilTest.java b/src/test/java/com/databricks/jdbc/common/util/ArrowUtilTest.java new file mode 100644 index 0000000000..88fa349abc --- /dev/null +++ b/src/test/java/com/databricks/jdbc/common/util/ArrowUtilTest.java @@ -0,0 +1,177 @@ +package com.databricks.jdbc.common.util; + +import static org.junit.jupiter.api.Assertions.*; + +import com.databricks.jdbc.exception.DatabricksParsingException; +import com.databricks.jdbc.model.client.thrift.generated.*; +import java.io.ByteArrayInputStream; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.Test; + +/** Unit tests for ArrowUtil. */ +public class ArrowUtilTest { + + @Test + void testGetSerializedSchemaWithPreSerializedArrow() throws DatabricksParsingException { + byte[] preSerializedSchema = new byte[] {1, 2, 3, 4, 5}; + + TGetResultSetMetadataResp metadata = new TGetResultSetMetadataResp(); + metadata.setArrowSchema(preSerializedSchema); + + byte[] result = ArrowUtil.getSerializedSchema(metadata); + + assertArrayEquals(preSerializedSchema, result); + } + + @Test + void testGetSerializedSchemaWithNullMetadata() throws DatabricksParsingException { + byte[] result = ArrowUtil.getSerializedSchema(null); + + assertNotNull(result); + assertEquals(0, result.length); + } + + @Test + void testGetSerializedSchemaWithHiveSchemaFallback() throws DatabricksParsingException { + // Create a Hive schema with a simple string column + TTableSchema hiveSchema = createHiveSchema("test_column", TTypeId.STRING_TYPE); + + TGetResultSetMetadataResp metadata = new TGetResultSetMetadataResp(); + metadata.setSchema(hiveSchema); + // Don't set arrowSchema - should trigger fallback + + byte[] result = ArrowUtil.getSerializedSchema(metadata); + + assertNotNull(result); + assertTrue(result.length > 0, "Serialized schema should not be empty"); + } + + @Test + void testHiveSchemaToArrowSchemaWithMultipleColumns() throws DatabricksParsingException { + // Create a Hive schema with multiple columns of different types + TTableSchema hiveSchema = new TTableSchema(); + List columns = new ArrayList<>(); + + columns.add(createColumnDesc("string_col", TTypeId.STRING_TYPE)); + columns.add(createColumnDesc("int_col", TTypeId.INT_TYPE)); + columns.add(createColumnDesc("boolean_col", TTypeId.BOOLEAN_TYPE)); + + hiveSchema.setColumns(columns); + + Schema arrowSchema = ArrowUtil.hiveSchemaToArrowSchema(hiveSchema); + + assertNotNull(arrowSchema); + assertEquals(3, arrowSchema.getFields().size()); + assertEquals("string_col", arrowSchema.getFields().get(0).getName()); + assertEquals("int_col", arrowSchema.getFields().get(1).getName()); + assertEquals("boolean_col", arrowSchema.getFields().get(2).getName()); + } + + @Test + void testHiveSchemaToArrowSchemaWithNullSchema() throws DatabricksParsingException { + Schema arrowSchema = ArrowUtil.hiveSchemaToArrowSchema(null); + + assertNotNull(arrowSchema); + assertEquals(0, arrowSchema.getFields().size()); + } + + @Test + void testColumnDescToArrowField() throws SQLException { + TColumnDesc columnDesc = createColumnDesc("my_column", TTypeId.STRING_TYPE); + + Field field = ArrowUtil.columnDescToArrowField(columnDesc); + + assertNotNull(field); + assertEquals("my_column", field.getName()); + assertTrue(field.isNullable()); + assertNotNull(field.getType()); + } + + @Test + void testCreateArrowByteStream() throws DatabricksParsingException { + byte[] schema = new byte[] {1, 2, 3}; + + TFetchResultsResp response = new TFetchResultsResp(); + TRowSet rowSet = new TRowSet(); + rowSet.setArrowBatches(Collections.emptyList()); + response.setResults(rowSet); + + // Set up metadata with no compression + TGetResultSetMetadataResp metadata = new TGetResultSetMetadataResp(); + response.setResultSetMetadata(metadata); + + ByteArrayInputStream result = ArrowUtil.createArrowByteStream(schema, response, getClass()); + + assertNotNull(result); + // Should at least contain the schema bytes + assertTrue(result.available() >= schema.length); + } + + @Test + void testGetTotalRowsInResponse() { + TFetchResultsResp response = new TFetchResultsResp(); + TRowSet rowSet = new TRowSet(); + + List batches = new ArrayList<>(); + TSparkArrowBatch batch1 = new TSparkArrowBatch(); + batch1.setRowCount(100); + batch1.setBatch(new byte[0]); + batches.add(batch1); + + TSparkArrowBatch batch2 = new TSparkArrowBatch(); + batch2.setRowCount(50); + batch2.setBatch(new byte[0]); + batches.add(batch2); + + rowSet.setArrowBatches(batches); + response.setResults(rowSet); + + long totalRows = ArrowUtil.getTotalRowsInResponse(response); + + assertEquals(150, totalRows); + } + + @Test + void testGetTotalRowsInResponseWithNullResults() { + TFetchResultsResp response = new TFetchResultsResp(); + // Don't set results + + long totalRows = ArrowUtil.getTotalRowsInResponse(response); + + assertEquals(0, totalRows); + } + + // ==================== Helper Methods ==================== + + private TTableSchema createHiveSchema(String columnName, TTypeId typeId) { + TTableSchema schema = new TTableSchema(); + List columns = new ArrayList<>(); + columns.add(createColumnDesc(columnName, typeId)); + schema.setColumns(columns); + return schema; + } + + private TColumnDesc createColumnDesc(String name, TTypeId typeId) { + TColumnDesc columnDesc = new TColumnDesc(); + columnDesc.setColumnName(name); + + TTypeDesc typeDesc = new TTypeDesc(); + List typeEntries = new ArrayList<>(); + + TTypeEntry typeEntry = new TTypeEntry(); + TPrimitiveTypeEntry primitiveType = new TPrimitiveTypeEntry(); + primitiveType.setType(typeId); + typeEntry.setPrimitiveEntry(primitiveType); + typeEntries.add(typeEntry); + + typeDesc.setTypes(typeEntries); + columnDesc.setTypeDesc(typeDesc); + + return columnDesc; + } +} diff --git a/src/test/java/com/databricks/jdbc/common/util/DatabricksThriftUtilTest.java b/src/test/java/com/databricks/jdbc/common/util/DatabricksThriftUtilTest.java index 6554e4cc66..4e53874e62 100644 --- a/src/test/java/com/databricks/jdbc/common/util/DatabricksThriftUtilTest.java +++ b/src/test/java/com/databricks/jdbc/common/util/DatabricksThriftUtilTest.java @@ -1,6 +1,7 @@ package com.databricks.jdbc.common.util; import static com.databricks.jdbc.TestConstants.*; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.QUERY_EXECUTION_TIMEOUT_SQLSTATE; import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_RESULT_ROW_LIMIT; import static com.databricks.jdbc.common.util.DatabricksThriftUtil.checkDirectResultsForErrorStatus; import static org.junit.jupiter.api.Assertions.*; @@ -11,11 +12,14 @@ import com.databricks.jdbc.common.DatabricksJdbcConstants; import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.exception.DatabricksTimeoutException; import com.databricks.jdbc.model.client.thrift.generated.*; import com.databricks.jdbc.model.core.ColumnInfoTypeName; import com.databricks.sdk.service.sql.StatementState; import java.nio.ByteBuffer; +import java.sql.SQLException; import java.util.*; +import java.util.Base64; import java.util.stream.Stream; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -89,6 +93,43 @@ void testVerifySuccessStatus() { assertEquals("42S02", exception.getSQLState()); } + @Test + void testVerifySuccessStatusThrowsTimeoutException() { + // Test that timeout SQL state triggers DatabricksTimeoutException + DatabricksTimeoutException timeoutException = + assertThrows( + DatabricksTimeoutException.class, + () -> + DatabricksThriftUtil.verifySuccessStatus( + new TStatus() + .setStatusCode(TStatusCode.ERROR_STATUS) + .setSqlState(QUERY_EXECUTION_TIMEOUT_SQLSTATE), + "Error while fetching results Request maxRows=100000, maxBytes=404857600. " + + "Response hasMoreRows=false")); + + assertEquals(QUERY_EXECUTION_TIMEOUT_SQLSTATE, timeoutException.getSQLState()); + assertTrue( + timeoutException + .getMessage() + .contains("Error thrift response received [Error while fetching results")); + + // Test with statement ID + timeoutException = + assertThrows( + DatabricksTimeoutException.class, + () -> + DatabricksThriftUtil.verifySuccessStatus( + new TStatus() + .setStatusCode(TStatusCode.ERROR_STATUS) + .setSqlState(QUERY_EXECUTION_TIMEOUT_SQLSTATE), + "Error while fetching results Request maxRows=100000, maxBytes=404857600. " + + "Response hasMoreRows=false", + "f6c9348e-84fb-47e8-8dcd-07455ff39ff5")); + + assertEquals(QUERY_EXECUTION_TIMEOUT_SQLSTATE, timeoutException.getSQLState()); + assertTrue(timeoutException.getMessage().contains("f6c9348e-84fb-47e8-8dcd-07455ff39ff5")); + } + private static Stream resultDataTypes() { // Create a test row set with chunk links TSparkArrowResultLink sampleResultLink = new TSparkArrowResultLink().setRowCount(10); @@ -217,14 +258,14 @@ private static Stream resultDataTypesForGetColumnValue() { @ParameterizedTest @MethodSource("resultDataTypes") - public void testRowCount(TRowSet resultData, int expectedRowCount) throws DatabricksSQLException { + public void testRowCount(TRowSet resultData, int expectedRowCount) throws SQLException { assertEquals(expectedRowCount, DatabricksThriftUtil.getRowCount(resultData)); } @ParameterizedTest @MethodSource("resultDataTypesForGetColumnValue") public void testColumnCount(TRowSet resultData, List> expectedValues) - throws DatabricksSQLException { + throws SQLException { assertEquals(expectedValues, DatabricksThriftUtil.extractRowsFromColumnar(resultData)); } @@ -255,7 +296,7 @@ public void testColumnCount(TGetResultSetMetadataResp resultManifest, int expect } @Test - public void testConvertColumnarToRowBased() throws DatabricksSQLException { + public void testConvertColumnarToRowBased() throws SQLException { when(fetchResultsResp.getResults()).thenReturn(BOOL_ROW_SET); when(parentStatement.getMaxRows()).thenReturn(DEFAULT_RESULT_ROW_LIMIT); List> rowBasedData = @@ -282,7 +323,7 @@ private static TTypeDesc createTypeDesc(TTypeId type) { } @Test - public void testMaxRowsInStatement() throws DatabricksSQLException { + public void testMaxRowsInStatement() throws SQLException { when(fetchResultsResp.getResults()).thenReturn(BOOL_ROW_SET); int maxRows = 2; when(parentStatement.getMaxRows()).thenReturn(maxRows); @@ -308,7 +349,8 @@ public void testMaxRowsInStatement() throws DatabricksSQLException { public void testGetTypeFromTypeDesc(TTypeId type, ColumnInfoTypeName typeName) { TColumnDesc columnDesc = new TColumnDesc().setTypeDesc(createTypeDesc(type)); assertEquals( - typeName, DatabricksThriftUtil.getColumnInfoFromTColumnDesc(columnDesc).getTypeName()); + typeName, + DatabricksThriftUtil.getColumnInfoFromTColumnDesc(columnDesc, null).getTypeName()); } @ParameterizedTest @@ -404,7 +446,7 @@ public void testGetStatementStatusForAsync() throws Exception { } @Test - public void testExtractRowsWithNullsForAllTypes() throws DatabricksSQLException { + public void testExtractRowsWithNullsForAllTypes() throws SQLException { // Create a TRowSet with three columns: STRING, DOUBLE, and BOOLEAN. TRowSet rowSet = new TRowSet(); List columns = new ArrayList<>(); @@ -453,7 +495,7 @@ public void testExtractRowsWithNullsForAllTypes() throws DatabricksSQLException } @Test - public void testExtractRowsWhenNullsArrayIsNull() throws DatabricksSQLException { + public void testExtractRowsWhenNullsArrayIsNull() throws SQLException { // Create a TRowSet with two columns: STRING and DOUBLE, where the nulls arrays are null. TRowSet rowSet = new TRowSet(); List columns = new ArrayList<>(); @@ -490,4 +532,53 @@ public void testExtractRowsWhenNullsArrayIsNull() throws DatabricksSQLException assertEquals(expected.get(i), actual.get(i), "Mismatch in row " + i); } } + + /** + * Real arrow schema captured from Databricks Thrift client execution. + * + *

Query: SELECT ST_POINT(1, 2, 4326) as geom, ST_GeogFromText('POINT(-122.4194 37.7749)') as + * geog_point, array(1, 2, 3) as arr, map('red', 1, 'green', 2) as mp, named_struct('name', + * 'Mumbai', 'population', 20000000) AS city + * + *

Execution: useThriftClient=true, enableArrow=true, cloudFetch=false + * + *

Expected metadata (5 fields): Field[0]: GEOMETRY(4326) Field[1]: GEOGRAPHY(4326) Field[2]: + * ARRAY Field[3]: MAP Field[4]: STRUCT + */ + private static final String REAL_ARROW_SCHEMA_BASE64 = + "/////4AGAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAUAAACABQAAoAQAAIADAADoAQAABAAAAKb6//8UAAAAUAEAAMQBAAAAAAANwAEAAAIAAADYAAAABAAAAFD6//8IAAAArAAAAKEAAAB7InR5cGUiOiJzdHJ1Y3QiLCJmaWVsZHMiOlt7Im5hbWUiOiJuYW1lIiwidHlwZSI6InN0cmluZyIsIm51bGxhYmxlIjpmYWxzZSwibWV0YWRhdGEiOnt9fSx7Im5hbWUiOiJwb3B1bGF0aW9uIiwidHlwZSI6ImludGVnZXIiLCJudWxsYWJsZSI6ZmFsc2UsIm1ldGFkYXRhIjp7fX1dfQAAABcAAABTcGFyazpEYXRhVHlwZTpKc29uVHlwZQAg+///CAAAAEAAAAA3AAAAU1RSVUNUPG5hbWU6IFNUUklORyBOT1QgTlVMTCwgcG9wdWxhdGlvbjogSU5UIE5PVCBOVUxMPgAWAAAAU3Bhcms6RGF0YVR5cGU6U3FsTmFtZQAAAgAAAEQAAAAEAAAACvz//xQAAAAUAAAAFAAAAAAAAAIYAAAAAAAAAAAAAAAg/f//AAAAASAAAAAKAAAAcG9wdWxhdGlvbgAARvz//xQAAAAUAAAAFAAAAAAAAAUQAAAAAAAAAAAAAACk+///BAAAAG5hbWUAAAAAtPv//wQAAABjaXR5AAAAAIb8//8UAAAA3AAAAHwBAAAAAAAReAEAAAIAAACIAAAABAAAADD8//8IAAAAXAAAAFEAAAB7InR5cGUiOiJtYXAiLCJrZXlUeXBlIjoic3RyaW5nIiwidmFsdWVUeXBlIjoiaW50ZWdlciIsInZhbHVlQ29udGFpbnNOdWxsIjpmYWxzZX0AAAAXAAAAU3Bhcms6RGF0YVR5cGU6SnNvblR5cGUAsPz//wgAAAAcAAAAEAAAAE1BUDxTVFJJTkcsIElOVD4AAAAAFgAAAFNwYXJrOkRhdGFUeXBlOlNxbE5hbWUAAAEAAAAEAAAAcv3//xQAAAAUAAAAgAAAAAAAAA18AAAAAAAAAAIAAABAAAAABAAAAJr9//8UAAAAFAAAABQAAAAAAAACGAAAAAAAAAAAAAAAsP7//wAAAAEgAAAABQAAAHZhbHVlAAAA0v3//xQAAAAUAAAAFAAAAAAAAAUQAAAAAAAAAAAAAAAw/f//AwAAAGtleQA8/f//BwAAAGVudHJpZXMATP3//wIAAABtcAAAGv7//xQAAADAAAAABAEAAAAAAAwAAQAAAgAAAHQAAAAEAAAAxP3//wgAAABIAAAAPQAAAHsidHlwZSI6ImFycmF5IiwiZWxlbWVudFR5cGUiOiJpbnRlZ2VyIiwiY29udGFpbnNOdWxsIjpmYWxzZX0AAAAXAAAAU3Bhcms6RGF0YVR5cGU6SnNvblR5cGUAMP7//wgAAAAUAAAACgAAAEFSUkFZPElOVD4AABYAAABTcGFyazpEYXRhVHlwZTpTcWxOYW1lAAABAAAABAAAAOr+//8UAAAAFAAAABwAAAAAAAACIAAAAAAAAAAAAAAACAAMAAgABwAIAAAAAAAAASAAAAAHAAAAZWxlbWVudABo/v//AwAAAGFycgA2////FAAAAKgAAACoAAAAAAAABaQAAAACAAAAWAAAAAQAAADg/v//CAAAACwAAAAhAAAAImdlb2dyYXBoeShPR0M6Q1JTODQsIFNQSEVSSUNBTCkiAAAAFwAAAFNwYXJrOkRhdGFUeXBlOkpzb25UeXBlADD///8IAAAAGAAAAA8AAABHRU9HUkFQSFkoNDMyNikAFgAAAFNwYXJrOkRhdGFUeXBlOlNxbE5hbWUAAAAAAAAo////CgAAAGdlb2dfcG9pbnQAAAAAEgAYABQAAAATAAwAAAAIAAQAEgAAABQAAACkAAAAqAAAAAAAAAWkAAAAAgAAAFQAAAAEAAAAvP///wgAAAAgAAAAFQAAACJnZW9tZXRyeShPR0M6Q1JTODQpIgAAABcAAABTcGFyazpEYXRhVHlwZTpKc29uVHlwZQAIAAwACAAEAAgAAAAIAAAAGAAAAA4AAABHRU9NRVRSWSg0MzI2KQAAFgAAAFNwYXJrOkRhdGFUeXBlOlNxbE5hbWUAAAAAAAAEAAQABAAAAAQAAABnZW9tAAAAAA=="; + + /** Helper method to create TGetResultSetMetadataResp with arrow schema bytes from base64. */ + private TGetResultSetMetadataResp createMetadataWithArrowSchema(String base64ArrowSchema) { + byte[] arrowSchemaBytes = Base64.getDecoder().decode(base64ArrowSchema); + return new TGetResultSetMetadataResp().setArrowSchema(arrowSchemaBytes); + } + + @Test + public void testGetArrowMetadataWithRealData() throws SQLException { + TGetResultSetMetadataResp metadata = createMetadataWithArrowSchema(REAL_ARROW_SCHEMA_BASE64); + List result = DatabricksThriftUtil.getArrowMetadata(metadata); + + List expected = + List.of( + "GEOMETRY(4326)", + "GEOGRAPHY(4326)", + "ARRAY", + "MAP", + "STRUCT"); + + assertEquals(expected, result); + } + + @Test + public void testGetArrowMetadataWithNullMetadata() throws SQLException { + assertNull(DatabricksThriftUtil.getArrowMetadata(null)); + } + + @Test + public void testGetArrowMetadataWithNullArrowSchema() throws SQLException { + TGetResultSetMetadataResp metadata = new TGetResultSetMetadataResp(); + assertNull(DatabricksThriftUtil.getArrowMetadata(metadata)); + } } diff --git a/src/test/java/com/databricks/jdbc/common/util/DriverUtilTest.java b/src/test/java/com/databricks/jdbc/common/util/DriverUtilTest.java index 6688a23cc2..28a83fb970 100644 --- a/src/test/java/com/databricks/jdbc/common/util/DriverUtilTest.java +++ b/src/test/java/com/databricks/jdbc/common/util/DriverUtilTest.java @@ -14,14 +14,14 @@ public class DriverUtilTest { public void testGetDriverVersion() { String version = DriverUtil.getDriverVersion(); assertNotNull(version); - assertEquals("3.0.7", version); + assertEquals("3.1.1", version); } @Test public void testGetDriverVersionWithoutOSSSuffix() { String version = DriverUtil.getDriverVersionWithoutOSSSuffix(); assertNotNull(version); - assertEquals("3.0.7", version); + assertEquals("3.1.1", version); } @Test @@ -40,7 +40,7 @@ public void testGetDriverMajorVersion() { @Test public void testGetDriverMinorVersion() { int minorVersion = DriverUtil.getDriverMinorVersion(); - assertEquals(0, minorVersion); + assertEquals(1, minorVersion); } @Test diff --git a/src/test/java/com/databricks/jdbc/common/util/UserAgentManagerTest.java b/src/test/java/com/databricks/jdbc/common/util/UserAgentManagerTest.java index 605a1c0ff4..176cfc74f0 100644 --- a/src/test/java/com/databricks/jdbc/common/util/UserAgentManagerTest.java +++ b/src/test/java/com/databricks/jdbc/common/util/UserAgentManagerTest.java @@ -86,7 +86,7 @@ void testUserAgentSetsClientCorrectly() throws DatabricksSQLException { DatabricksConnectionContextFactory.create(CLUSTER_JDBC_URL, new Properties()); UserAgentManager.setUserAgent(connectionContext); String userAgent = getUserAgentString(); - assertTrue(userAgent.contains("DatabricksJDBCDriverOSS/3.0.7")); + assertTrue(userAgent.contains("DatabricksJDBCDriverOSS/")); assertTrue(userAgent.contains(" Java/THttpClient")); assertTrue(userAgent.contains(" MyApp/version")); assertTrue(userAgent.contains(" databricks-jdbc-http ")); @@ -97,7 +97,7 @@ void testUserAgentSetsClientCorrectly() throws DatabricksSQLException { DatabricksConnectionContextFactory.create(WAREHOUSE_JDBC_URL, new Properties()); UserAgentManager.setUserAgent(connectionContext); userAgent = getUserAgentString(); - assertTrue(userAgent.contains("DatabricksJDBCDriverOSS/3.0.7")); + assertTrue(userAgent.contains("DatabricksJDBCDriverOSS/")); assertTrue(userAgent.contains(" Java/THttpClient")); assertTrue(userAgent.contains(" MyApp/version")); assertTrue(userAgent.contains(" databricks-jdbc-http ")); @@ -108,12 +108,60 @@ void testUserAgentSetsClientCorrectly() throws DatabricksSQLException { DatabricksConnectionContextFactory.create(WAREHOUSE_JDBC_URL_WITH_SEA, new Properties()); UserAgentManager.setUserAgent(connectionContext); userAgent = getUserAgentString(); - assertTrue(userAgent.contains("DatabricksJDBCDriverOSS/3.0.7")); + assertTrue(userAgent.contains("DatabricksJDBCDriverOSS/")); assertTrue(userAgent.contains(" Java/SQLExecHttpClient")); assertTrue(userAgent.contains(" databricks-jdbc-http ")); assertFalse(userAgent.contains("databricks-sdk-java")); } + @Test + void testCustomUserAgentIncludedBeforeClientTypeEvaluation() throws DatabricksSQLException { + // This test verifies that custom user agent is included in connector service requests + // and maintains proper order: base -> client type -> custom + + // Create connection context with custom user agent entry (use URL without existing UA) + String jdbcUrlWithCustomUA = + WAREHOUSE_JDBC_URL_WITH_THRIFT + ";useragententry=CustomTestApp/2.0.0"; + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContextFactory.create(jdbcUrlWithCustomUA, new Properties()); + + // Call setUserAgent + UserAgentManager.setUserAgent(connectionContext); + + // Get the final user agent string + String userAgent = getUserAgentString(); + + // Verify custom user agent is included + assertTrue( + userAgent.contains("CustomTestApp/2.0.0"), + "Custom user agent should be included: " + userAgent); + + // Verify driver version is present + assertTrue( + userAgent.contains("DatabricksJDBCDriverOSS/"), + "Driver version should be present: " + userAgent); + + // Verify client type is included (proves getClientType was called) + assertTrue( + userAgent.contains("Java/THttpClient") || userAgent.contains("Java/SQLExecHttpClient"), + "Client type should be included: " + userAgent); + + // Verify JDBC HTTP client identifier + assertTrue( + userAgent.contains("databricks-jdbc-http"), + "JDBC HTTP client identifier should be present: " + userAgent); + + // Verify correct order: client type should come BEFORE custom user agent + int clientTypeIndex = + userAgent.contains("Java/THttpClient") + ? userAgent.indexOf("Java/THttpClient") + : userAgent.indexOf("Java/SQLExecHttpClient"); + int customUAIndex = userAgent.indexOf("CustomTestApp/2.0.0"); + assertTrue( + clientTypeIndex < customUAIndex, + "Client type should appear before custom user agent. Order: " + userAgent); + } + @Test void testUserAgentSetsCustomerInput() throws DatabricksSQLException { IDatabricksConnectionContext connectionContext = diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/MetadataResultSetBuilderTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/MetadataResultSetBuilderTest.java index df6ddd9e57..e4a6dfbdd0 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/MetadataResultSetBuilderTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/MetadataResultSetBuilderTest.java @@ -2,12 +2,11 @@ import static com.databricks.jdbc.common.MetadataResultConstants.*; import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; -import static org.mockito.Mockito.withSettings; +import static org.mockito.Mockito.*; import com.databricks.jdbc.api.impl.DatabricksResultSet; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; +import com.databricks.jdbc.api.internal.IDatabricksSession; import com.databricks.jdbc.common.util.DatabricksThreadContextHolder; import com.databricks.jdbc.model.core.ResultColumn; import java.sql.ResultSet; @@ -729,4 +728,42 @@ void testGetTableTypesResultWithMetricViewDisabled() throws SQLException { assertTrue(tableTypes.contains("VIEW")); assertFalse(tableTypes.contains("METRIC_VIEW")); } + + @Test + void shouldAllowCatalogAccessWhenMultipleCatalogSupportEnabled() throws SQLException { + when(connectionContext.getEnableMultipleCatalogSupport()).thenReturn(true); + IDatabricksSession session = mock(IDatabricksSession.class); + + assertTrue(metadataResultSetBuilder.shouldAllowCatalogAccess("cat_a", "cat_b", session)); + verify(session, never()).getCurrentCatalog(); + } + + @Test + void shouldAllowCatalogAccessWhenCatalogIsNull() throws SQLException { + when(connectionContext.getEnableMultipleCatalogSupport()).thenReturn(false); + IDatabricksSession session = mock(IDatabricksSession.class); + + assertTrue(metadataResultSetBuilder.shouldAllowCatalogAccess(null, "cat_b", session)); + verify(session, never()).getCurrentCatalog(); + } + + @Test + void shouldAllowCatalogAccessWhenCurrentCatalogMatchesSession() throws SQLException { + when(connectionContext.getEnableMultipleCatalogSupport()).thenReturn(false); + IDatabricksSession session = mock(IDatabricksSession.class); + when(session.getCurrentCatalog()).thenReturn("main"); + + assertTrue(metadataResultSetBuilder.shouldAllowCatalogAccess("main", null, session)); + verify(session).getCurrentCatalog(); + } + + @Test + void shouldDenyCatalogAccessWhenCatalogDiffers() throws SQLException { + when(connectionContext.getEnableMultipleCatalogSupport()).thenReturn(false); + IDatabricksSession session = mock(IDatabricksSession.class); + when(session.getCurrentCatalog()).thenReturn("main"); + + assertFalse(metadataResultSetBuilder.shouldAllowCatalogAccess("other", null, session)); + verify(session).getCurrentCatalog(); + } } diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksEmptyMetadataClientTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksEmptyMetadataClientTest.java index dd2aca1ad6..1152900d27 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksEmptyMetadataClientTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksEmptyMetadataClientTest.java @@ -115,4 +115,39 @@ void testListPrimaryKeys() throws SQLException { assertEquals(resultSet.getMetaData().getColumnName(2), "TABLE_SCHEM"); assertEquals(resultSet.getMetaData().getColumnName(3), "TABLE_NAME"); } + + @Test + void testListImportedKeys() throws SQLException { + ResultSet resultSet = mockClient.listImportedKeys(session, "catalog", "schema", "table"); + assertNotNull(resultSet); + assertFalse(resultSet.next()); // empty result set + assertEquals(14, resultSet.getMetaData().getColumnCount()); + assertEquals("PKTABLE_CAT", resultSet.getMetaData().getColumnName(1)); + } + + @Test + void testListExportedKeys() throws SQLException { + ResultSet resultSet = mockClient.listExportedKeys(session, "catalog", "schema", "table"); + assertNotNull(resultSet); + assertFalse(resultSet.next()); // empty result set + assertEquals(14, resultSet.getMetaData().getColumnCount()); + assertEquals("PKTABLE_CAT", resultSet.getMetaData().getColumnName(1)); + } + + @Test + void testListCrossReferences() throws SQLException { + ResultSet resultSet = + mockClient.listCrossReferences( + session, + "parentCatalog", + "parentSchema", + "parentTable", + "foreignCatalog", + "foreignSchema", + "foreignTable"); + assertNotNull(resultSet); + assertFalse(resultSet.next()); // empty result set + assertEquals(14, resultSet.getMetaData().getColumnCount()); + assertEquals("PKTABLE_CAT", resultSet.getMetaData().getColumnName(1)); + } } diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataSdkClientTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataSdkClientTest.java index 40453fc7b6..71a7298e23 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataSdkClientTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataSdkClientTest.java @@ -5,8 +5,12 @@ import static com.databricks.jdbc.dbclient.impl.common.CommandConstants.*; import static com.databricks.jdbc.dbclient.impl.common.ImportedKeysDatabricksResultSetAdapter.*; import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.databricks.jdbc.api.impl.DatabricksResultSet; @@ -195,6 +199,39 @@ void testListCatalogs() throws SQLException { assertEquals(((DatabricksResultSetMetaData) actualResult.getMetaData()).getTotalRows(), 2); } + @Test + void listTablesReturnsEmptyWhenCatalogAccessDenied() throws SQLException { + IDatabricksConnectionContext connectionContext = mock(IDatabricksConnectionContext.class); + when(connectionContext.getEnableMultipleCatalogSupport()).thenReturn(false); + when(mockClient.getConnectionContext()).thenReturn(connectionContext); + when(session.getCurrentCatalog()).thenReturn("main"); + + DatabricksMetadataSdkClient metadataClient = new DatabricksMetadataSdkClient(mockClient); + + DatabricksResultSet result = + metadataClient.listTables(session, "other", null, null, new String[] {"TABLE"}); + + assertNotNull(result); + assertFalse(result.next(), "Expected no rows when catalog access is denied"); + verify(mockClient, never()).executeStatement(anyString(), any(), any(), any(), any(), any()); + } + + @Test + void listSchemasReturnsEmptyWhenCatalogIsEmptyString() throws SQLException { + IDatabricksConnectionContext connectionContext = mock(IDatabricksConnectionContext.class); + when(connectionContext.getEnableMultipleCatalogSupport()).thenReturn(false); + when(mockClient.getConnectionContext()).thenReturn(connectionContext); + when(session.getCurrentCatalog()).thenReturn("main"); + + DatabricksMetadataSdkClient metadataClient = new DatabricksMetadataSdkClient(mockClient); + + DatabricksResultSet result = metadataClient.listSchemas(session, "", null); + + assertNotNull(result); + assertFalse(result.next(), "Expected no rows when catalog is empty string"); + verify(mockClient, never()).executeStatement(anyString(), any(), any(), any(), any(), any()); + } + @Test void testListTableTypes() throws SQLException { // Mock the connection context for the table types test @@ -672,7 +709,9 @@ void testListCrossReferences_throwsParseSyntaxError() throws Exception { new DatabricksSQLException( "syntax error at or near \"foreign\"", PARSE_SYNTAX_ERROR_SQL_STATE); when(session.getComputeResource()).thenReturn(WAREHOUSE_COMPUTE); - when(mockClient.getConnectionContext()).thenReturn(mock(IDatabricksConnectionContext.class)); + IDatabricksConnectionContext mockContext = mock(IDatabricksConnectionContext.class); + when(mockContext.getEnableMultipleCatalogSupport()).thenReturn(true); + when(mockClient.getConnectionContext()).thenReturn(mockContext); DatabricksMetadataSdkClient metadataClient = new DatabricksMetadataSdkClient(mockClient); when(mockClient.executeStatement( "SHOW FOREIGN KEYS IN CATALOG `catalog1` IN SCHEMA `testSchema` IN TABLE `testTable`", @@ -1079,7 +1118,9 @@ void testListCrossReferences_handlesNullSqlStateWithoutNPE() throws Exception { "syntax error at or near \"foreign\"", (String) null); // null SQL state when(session.getComputeResource()).thenReturn(WAREHOUSE_COMPUTE); - when(mockClient.getConnectionContext()).thenReturn(mock(IDatabricksConnectionContext.class)); + IDatabricksConnectionContext mockContext = mock(IDatabricksConnectionContext.class); + when(mockContext.getEnableMultipleCatalogSupport()).thenReturn(true); + when(mockClient.getConnectionContext()).thenReturn(mockContext); DatabricksMetadataSdkClient metadataClient = new DatabricksMetadataSdkClient(mockClient); when(mockClient.executeStatement( @@ -1104,4 +1145,113 @@ void testListCrossReferences_handlesNullSqlStateWithoutNPE() throws Exception { TEST_SCHEMA, TEST_TABLE)); } + + @Test + void testListCatalogsWithMultipleCatalogSupportDisabled() throws SQLException { + when(session.getComputeResource()).thenReturn(mockedComputeResource); + when(session.getCurrentCatalog()).thenReturn("my_catalog"); + IDatabricksConnectionContext mockContext = mock(IDatabricksConnectionContext.class); + when(mockContext.getEnableMultipleCatalogSupport()).thenReturn(false); + when(mockClient.getConnectionContext()).thenReturn(mockContext); + + DatabricksMetadataSdkClient metadataClient = new DatabricksMetadataSdkClient(mockClient); + + String expectedSQL = "SELECT 'my_catalog' AS catalog"; + when(mockClient.executeStatement( + expectedSQL, + mockedComputeResource, + new HashMap<>(), + StatementType.METADATA, + session, + null)) + .thenReturn(mockedCatalogResultSet); + + when(mockedCatalogResultSet.next()).thenReturn(true, false); + when(mockedCatalogResultSet.getObject("catalog")).thenReturn("my_catalog"); + doReturn(1).when(mockedMetaData).getColumnCount(); + doReturn("catalog").when(mockedMetaData).getColumnName(1); + doReturn(255).when(mockedMetaData).getPrecision(1); + doReturn(0).when(mockedMetaData).getScale(1); + when(mockedCatalogResultSet.getMetaData()).thenReturn(mockedMetaData); + + DatabricksResultSet actualResult = metadataClient.listCatalogs(session); + + assertEquals(StatementState.SUCCEEDED, actualResult.getStatementStatus().getState()); + assertEquals(GET_CATALOGS_STATEMENT_ID, actualResult.getStatementId()); + assertEquals(1, ((DatabricksResultSetMetaData) actualResult.getMetaData()).getTotalRows()); + } + + @Test + void testListTablesWithMultipleCatalogSupportDisabledAndDifferentCatalog() throws SQLException { + when(session.getCurrentCatalog()).thenReturn("current_catalog"); + IDatabricksConnectionContext mockContext = mock(IDatabricksConnectionContext.class); + when(mockContext.getEnableMultipleCatalogSupport()).thenReturn(false); + when(mockClient.getConnectionContext()).thenReturn(mockContext); + + DatabricksMetadataSdkClient metadataClient = new DatabricksMetadataSdkClient(mockClient); + + // Try to access a different catalog - should return empty + DatabricksResultSet actualResult = + metadataClient.listTables(session, "different_catalog", null, null, null); + + assertEquals(StatementState.SUCCEEDED, actualResult.getStatementStatus().getState()); + assertEquals(GET_TABLES_STATEMENT_ID, actualResult.getStatementId()); + // Should return empty result set + assertEquals(0, ((DatabricksResultSetMetaData) actualResult.getMetaData()).getTotalRows()); + } + + @Test + void testListPrimaryKeysWithMultipleCatalogSupportDisabledAndDifferentCatalog() + throws SQLException { + when(session.getCurrentCatalog()).thenReturn("current_catalog"); + IDatabricksConnectionContext mockContext = mock(IDatabricksConnectionContext.class); + when(mockContext.getEnableMultipleCatalogSupport()).thenReturn(false); + when(mockClient.getConnectionContext()).thenReturn(mockContext); + + DatabricksMetadataSdkClient metadataClient = new DatabricksMetadataSdkClient(mockClient); + + // Try to access a different catalog - should return empty + DatabricksResultSet actualResult = + metadataClient.listPrimaryKeys(session, "different_catalog", "schema", "table"); + + assertEquals(StatementState.SUCCEEDED, actualResult.getStatementStatus().getState()); + assertEquals(METADATA_STATEMENT_ID, actualResult.getStatementId()); + // Should return empty result set + assertEquals(0, ((DatabricksResultSetMetaData) actualResult.getMetaData()).getTotalRows()); + } + + @Test + void testNoUnnecessaryGetCurrentCatalogCallWhenSupportEnabled() throws SQLException { + when(session.getComputeResource()).thenReturn(mockedComputeResource); + IDatabricksConnectionContext mockContext = mock(IDatabricksConnectionContext.class); + when(mockContext.getEnableMultipleCatalogSupport()).thenReturn(true); + when(mockClient.getConnectionContext()).thenReturn(mockContext); + + DatabricksMetadataSdkClient metadataClient = new DatabricksMetadataSdkClient(mockClient); + + String expectedSQL = "SHOW SCHEMAS IN `test_catalog`"; + when(mockClient.executeStatement( + expectedSQL, + mockedComputeResource, + new HashMap<>(), + StatementType.METADATA, + session, + null)) + .thenReturn(mockedResultSet); + + when(mockedResultSet.next()).thenReturn(true, false); + when(mockedResultSet.getObject("databaseName")).thenReturn("schema1"); + doReturn(2).when(mockedMetaData).getColumnCount(); + doReturn(SCHEMA_COLUMN.getResultSetColumnName()).when(mockedMetaData).getColumnName(1); + doReturn(CATALOG_COLUMN.getResultSetColumnName()).when(mockedMetaData).getColumnName(2); + when(mockedResultSet.getMetaData()).thenReturn(mockedMetaData); + when(mockedResultSet.findColumn(CATALOG_RESULT_COLUMN.getResultSetColumnName())) + .thenThrow(DatabricksSQLException.class); + + DatabricksResultSet actualResult = metadataClient.listSchemas(session, "test_catalog", null); + + assertEquals(StatementState.SUCCEEDED, actualResult.getStatementStatus().getState()); + // Verify getCurrentCatalog was NEVER called when support is enabled + verify(session, never()).getCurrentCatalog(); + } } diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java index ca55cc5708..f67e0e497f 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java @@ -118,7 +118,7 @@ void setup(Boolean directResultsEnabled) } @Test - void testOpenSession() throws TException, DatabricksSQLException, DatabricksValidationException { + void testOpenSession() throws TException, SQLException, DatabricksValidationException { setup(true); TOpenSessionReq request = new TOpenSessionReq(); TOpenSessionResp response = new TOpenSessionResp(); @@ -127,7 +127,7 @@ void testOpenSession() throws TException, DatabricksSQLException, DatabricksVali } @Test - void testCloseSession() throws TException, DatabricksSQLException, DatabricksValidationException { + void testCloseSession() throws TException, SQLException, DatabricksValidationException { setup(true); TCloseSessionReq request = new TCloseSessionReq(); TCloseSessionResp response = new TCloseSessionResp(); @@ -323,8 +323,7 @@ void testExecuteThrowsSQLExceptionWithSqlState() } @Test - void testCancelOperation() - throws TException, DatabricksSQLException, DatabricksValidationException { + void testCancelOperation() throws TException, SQLException, DatabricksValidationException { setup(true); TCancelOperationReq request = new TCancelOperationReq() @@ -340,8 +339,7 @@ void testCancelOperation() } @Test - void testCloseOperation() - throws TException, DatabricksSQLException, DatabricksValidationException { + void testCloseOperation() throws TException, SQLException, DatabricksValidationException { setup(true); TCloseOperationReq request = new TCloseOperationReq() @@ -387,8 +385,7 @@ void testCloseOperation_error() } @Test - void testIncludeResultSetMetadataNotSetForOldProtocol() - throws TException, DatabricksSQLException { + void testIncludeResultSetMetadataNotSetForOldProtocol() throws TException, SQLException { TOperationHandle operationHandle = new TOperationHandle() .setOperationId(handleIdentifier) @@ -453,8 +450,7 @@ void testGetStatementResult_pending() throws Exception { } @Test - void testListPrimaryKeys() - throws TException, DatabricksSQLException, DatabricksValidationException { + void testListPrimaryKeys() throws TException, SQLException, DatabricksValidationException { setup(false); TGetPrimaryKeysReq request = new TGetPrimaryKeysReq(); TGetPrimaryKeysResp tGetPrimaryKeysResp = @@ -471,7 +467,7 @@ void testListPrimaryKeys() @Test void testListPrimaryKeysWithDirectResults() - throws TException, DatabricksSQLException, DatabricksValidationException { + throws TException, SQLException, DatabricksValidationException { setup(true); TGetPrimaryKeysReq request = new TGetPrimaryKeysReq(); TGetPrimaryKeysResp tGetPrimaryKeysResp = @@ -485,8 +481,7 @@ void testListPrimaryKeysWithDirectResults() } @Test - void testListFunctions() - throws TException, DatabricksSQLException, DatabricksValidationException { + void testListFunctions() throws TException, SQLException, DatabricksValidationException { setup(false); TGetFunctionsReq request = new TGetFunctionsReq(); TGetFunctionsResp tGetFunctionsResp = @@ -503,7 +498,7 @@ void testListFunctions() @Test void testListFunctionsWithDirectResults() - throws TException, DatabricksSQLException, DatabricksValidationException { + throws TException, SQLException, DatabricksValidationException { setup(true); TGetFunctionsReq request = new TGetFunctionsReq(); TGetFunctionsResp tGetFunctionsResp = @@ -517,7 +512,7 @@ void testListFunctionsWithDirectResults() } @Test - void testListSchemas() throws TException, DatabricksSQLException, DatabricksValidationException { + void testListSchemas() throws TException, SQLException, DatabricksValidationException { setup(false); TGetSchemasReq request = new TGetSchemasReq(); TGetSchemasResp tGetSchemasResp = @@ -534,7 +529,7 @@ void testListSchemas() throws TException, DatabricksSQLException, DatabricksVali @Test void testListSchemasWithDirectResults() - throws TException, DatabricksSQLException, DatabricksValidationException { + throws TException, SQLException, DatabricksValidationException { setup(true); TGetSchemasReq request = new TGetSchemasReq(); TGetSchemasResp tGetSchemasResp = @@ -548,7 +543,7 @@ void testListSchemasWithDirectResults() } @Test - void testListColumns() throws TException, DatabricksSQLException, DatabricksValidationException { + void testListColumns() throws TException, SQLException, DatabricksValidationException { setup(false); TGetColumnsReq request = new TGetColumnsReq(); TGetColumnsResp tGetColumnsResp = @@ -565,7 +560,7 @@ void testListColumns() throws TException, DatabricksSQLException, DatabricksVali @Test void testListColumnsWithDirectResults() - throws TException, DatabricksSQLException, DatabricksValidationException { + throws TException, SQLException, DatabricksValidationException { setup(true); TGetColumnsReq request = new TGetColumnsReq(); TGetColumnsResp tGetColumnsResp = @@ -579,7 +574,7 @@ void testListColumnsWithDirectResults() } @Test - void testListCatalogs() throws TException, DatabricksSQLException, DatabricksValidationException { + void testListCatalogs() throws TException, SQLException, DatabricksValidationException { setup(true); TGetCatalogsReq request = new TGetCatalogsReq(); TGetCatalogsResp tGetCatalogsResp = @@ -596,7 +591,7 @@ void testListCatalogs() throws TException, DatabricksSQLException, DatabricksVal @Test void testListCatalogsWithDirectResults() - throws TException, DatabricksSQLException, DatabricksValidationException { + throws TException, SQLException, DatabricksValidationException { setup(true); TGetCatalogsReq request = new TGetCatalogsReq(); TGetCatalogsResp tGetCatalogsResp = @@ -610,7 +605,7 @@ void testListCatalogsWithDirectResults() } @Test - void testListTables() throws TException, DatabricksSQLException, DatabricksValidationException { + void testListTables() throws TException, SQLException, DatabricksValidationException { setup(false); TGetTablesReq request = new TGetTablesReq(); TGetTablesResp tGetTablesResp = @@ -627,7 +622,7 @@ void testListTables() throws TException, DatabricksSQLException, DatabricksValid @Test void testListTablesWithDirectResults() - throws TException, DatabricksSQLException, DatabricksValidationException { + throws TException, SQLException, DatabricksValidationException { setup(true); TGetTablesReq request = new TGetTablesReq(); TGetTablesResp tGetTablesResp = @@ -641,8 +636,7 @@ void testListTablesWithDirectResults() } @Test - void testListTableTypes() - throws TException, DatabricksSQLException, DatabricksValidationException { + void testListTableTypes() throws TException, SQLException, DatabricksValidationException { setup(false); TGetTableTypesReq request = new TGetTableTypesReq(); TGetTableTypesResp tGetTableTypesResp = @@ -659,7 +653,7 @@ void testListTableTypes() @Test void testListTableTypesWithDirectResults() - throws TException, DatabricksSQLException, DatabricksValidationException { + throws TException, SQLException, DatabricksValidationException { setup(true); TGetTableTypesReq request = new TGetTableTypesReq(); TGetTableTypesResp tGetTableTypesResp = @@ -673,7 +667,7 @@ void testListTableTypesWithDirectResults() } @Test - void testTypeInfo() throws TException, DatabricksSQLException, DatabricksValidationException { + void testTypeInfo() throws TException, SQLException, DatabricksValidationException { setup(false); TGetTypeInfoReq request = new TGetTypeInfoReq(); TGetTypeInfoResp tGetTypeInfoResp = @@ -690,7 +684,7 @@ void testTypeInfo() throws TException, DatabricksSQLException, DatabricksValidat @Test void testTypeInfoWithDirectResults() - throws TException, DatabricksSQLException, DatabricksValidationException { + throws TException, SQLException, DatabricksValidationException { setup(true); TGetTypeInfoReq request = new TGetTypeInfoReq(); TGetTypeInfoResp tGetTypeInfoResp = 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 e07bd93b8f..cdd383a47d 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 @@ -12,6 +12,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -35,6 +36,7 @@ import java.sql.SQLException; import java.util.*; import java.util.stream.Stream; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.junit.jupiter.params.ParameterizedTest; @@ -62,8 +64,15 @@ public class DatabricksThriftServiceClientTest { @Mock DatabricksConfig databricksConfig; @Mock ResultSetMetaData mockedMetaData; + @BeforeEach + void setUp() { + // Enable multiple catalog support by default for all tests + // Individual tests can override this if needed + lenient().when(connectionContext.getEnableMultipleCatalogSupport()).thenReturn(true); + } + @Test - void testCreateSession() throws DatabricksSQLException { + void testCreateSession() throws SQLException { when(connectionContext.getEnableMultipleCatalogSupport()).thenReturn(true); DatabricksThriftServiceClient client = new DatabricksThriftServiceClient(thriftAccessor, connectionContext); @@ -85,7 +94,7 @@ void testCreateSession() throws DatabricksSQLException { } @Test - void testCreateSessionHandlesProtocolVersion() throws DatabricksSQLException { + void testCreateSessionHandlesProtocolVersion() throws SQLException { when(connectionContext.getEnableMultipleCatalogSupport()).thenReturn(true); DatabricksThriftServiceClient client = new DatabricksThriftServiceClient(thriftAccessor, connectionContext); @@ -136,7 +145,7 @@ void testCreateSessionHandlesProtocolVersion() throws DatabricksSQLException { } @Test - void testCloseSession() throws DatabricksSQLException { + void testCloseSession() throws SQLException { DatabricksThriftServiceClient client = new DatabricksThriftServiceClient(thriftAccessor, connectionContext); TCloseSessionReq closeSessionReq = new TCloseSessionReq().setSessionHandle(SESSION_HANDLE); @@ -164,6 +173,9 @@ private static Stream protocolVersionProvider() { void testGetRequestWithDifferentProtocolVersions(TProtocolVersion protocolVersion) throws SQLException { when(connectionContext.shouldEnableArrow()).thenReturn(true); + // Use lenient() because isCloudFetchEnabled() is only called for protocols that support + // CloudFetch + lenient().when(connectionContext.isCloudFetchEnabled()).thenReturn(true); DatabricksThriftServiceClient client = new DatabricksThriftServiceClient(thriftAccessor, connectionContext); when(session.getSessionInfo()).thenReturn(SESSION_INFO); @@ -243,6 +255,7 @@ void testGetRequestWithDifferentProtocolVersions(TProtocolVersion protocolVersio @Test void testExecute() throws SQLException { when(connectionContext.shouldEnableArrow()).thenReturn(true); + when(connectionContext.isCloudFetchEnabled()).thenReturn(true); DatabricksThriftServiceClient client = new DatabricksThriftServiceClient(thriftAccessor, connectionContext); when(session.getSessionInfo()).thenReturn(SESSION_INFO); @@ -281,9 +294,52 @@ void testExecute() throws SQLException { assertEquals(resultSet, actualResultSet); } + @Test + void testExecuteWithCloudFetchDisabled() throws SQLException { + when(connectionContext.shouldEnableArrow()).thenReturn(true); + when(connectionContext.isCloudFetchEnabled()).thenReturn(false); + DatabricksThriftServiceClient client = + new DatabricksThriftServiceClient(thriftAccessor, connectionContext); + when(session.getSessionInfo()).thenReturn(SESSION_INFO); + when(parentStatement.getStatement()).thenReturn(statement); + when(parentStatement.getMaxRows()).thenReturn(10); + when(statement.getQueryTimeout()).thenReturn(10); + TSparkArrowTypes arrowNativeTypes = + new TSparkArrowTypes() + .setComplexTypesAsArrow(true) + .setIntervalTypesAsArrow(true) + .setNullTypeAsArrow(true) + .setDecimalAsArrow(true) + .setTimestampAsArrow(true); + TExecuteStatementReq executeStatementReq = + new TExecuteStatementReq() + .setStatement(TEST_STRING) + .setSessionHandle(SESSION_HANDLE) + .setCanReadArrowResult(true) + .setQueryTimeout(10) + .setResultRowLimit(10) + .setCanDecompressLZ4Result(true) + .setCanDownloadResult(false) + .setParameters(Collections.emptyList()) + .setRunAsync(true) + .setUseArrowNativeTypes(arrowNativeTypes); + when(thriftAccessor.execute(executeStatementReq, parentStatement, session, StatementType.SQL)) + .thenReturn(resultSet); + DatabricksResultSet actualResultSet = + client.executeStatement( + TEST_STRING, + CLUSTER_COMPUTE, + Collections.emptyMap(), + StatementType.SQL, + session, + parentStatement); + assertEquals(resultSet, actualResultSet); + } + @Test void testExecuteAsync() throws SQLException { when(connectionContext.shouldEnableArrow()).thenReturn(true); + when(connectionContext.isCloudFetchEnabled()).thenReturn(true); DatabricksThriftServiceClient client = new DatabricksThriftServiceClient(thriftAccessor, connectionContext); when(session.getSessionInfo()).thenReturn(SESSION_INFO); @@ -630,6 +686,7 @@ void testListFunctionsWithSQLEnabled() throws SQLException { new DatabricksThriftServiceClient(thriftAccessor, connectionContext); when(connectionContext.enableShowCommandsForGetFunctions()).thenReturn(true); when(connectionContext.shouldEnableArrow()).thenReturn(true); + when(connectionContext.isCloudFetchEnabled()).thenReturn(true); when(session.getSessionInfo()).thenReturn(SESSION_INFO); TSparkArrowTypes arrowNativeTypes = new TSparkArrowTypes() diff --git a/src/test/java/com/databricks/jdbc/integration/e2e/ExplicitTransactionStatementTests.java b/src/test/java/com/databricks/jdbc/integration/e2e/ExplicitTransactionStatementTests.java index 786851e0cb..24bfd34ce9 100644 --- a/src/test/java/com/databricks/jdbc/integration/e2e/ExplicitTransactionStatementTests.java +++ b/src/test/java/com/databricks/jdbc/integration/e2e/ExplicitTransactionStatementTests.java @@ -475,8 +475,11 @@ void testSetAutocommitWithoutValue() throws SQLException { ResultSet rs1 = stmt.executeQuery("SET AUTOCOMMIT"); assertTrue(rs1.next(), "SET AUTOCOMMIT should return a result"); // The result should indicate autocommit is TRUE (default) + // Note: The server may return "1"/"0" instead of "true"/"false" String value1 = rs1.getString(1); - assertTrue(Boolean.parseBoolean(value1), "Should return a value"); + assertTrue( + "true".equalsIgnoreCase(value1) || "1".equals(value1), + "Default autocommit should be true/1. Got: " + value1); rs1.close(); // Set autocommit to FALSE diff --git a/src/test/java/com/databricks/jdbc/integration/e2e/GeospatialTests.java b/src/test/java/com/databricks/jdbc/integration/e2e/GeospatialTests.java new file mode 100644 index 0000000000..8b3bf9e4de --- /dev/null +++ b/src/test/java/com/databricks/jdbc/integration/e2e/GeospatialTests.java @@ -0,0 +1,330 @@ +package com.databricks.jdbc.integration.e2e; + +import static com.databricks.jdbc.integration.IntegrationTestUtil.*; +import static org.junit.jupiter.api.Assertions.*; + +import com.databricks.jdbc.api.IGeography; +import com.databricks.jdbc.api.IGeometry; +import com.databricks.jdbc.api.impl.DatabricksResultSetMetaData; +import java.sql.*; +import java.util.List; +import java.util.stream.Stream; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +/** + * E2E integration tests for geospatial POINT types across all configuration combinations: Protocol + * (Thrift/SEA), Serialization (Arrow/Inline), CloudFetch, GeoSpatial support, and Complex type + * support. + * + *

Geospatial objects (IGeometry/IGeography) are returned only when both EnableGeoSpatialSupport + * and EnableComplexDatatypeSupport are enabled AND not in Thrift+Inline mode. Otherwise, returns as + * STRING. + */ +public class GeospatialTests { + + private Connection connection; + + @AfterEach + void cleanUp() throws SQLException { + if (connection != null) { + connection.close(); + } + } + + private void setupConnection( + int useThrift, int enableArrow, int enableGeoSupport, int enableComplexSupport) + throws SQLException { + connection = + getValidJDBCConnection( + List.of( + List.of("UseThriftClient", String.valueOf(useThrift)), + List.of("EnableArrow", String.valueOf(enableArrow)), + List.of("EnableGeoSpatialSupport", String.valueOf(enableGeoSupport)), + List.of("EnableComplexDatatypeSupport", String.valueOf(enableComplexSupport)))); + } + + private static Stream provideAllConfigurations() { + // Note: CloudFetch requires Arrow, so CloudFetch=1 configurations only appear with Arrow=1 + return Stream.of( + // =================================================================== + // SEA (useThrift=0), Inline (arrow=0), No CloudFetch + // =================================================================== + Arguments.of(0, 0, 0, 0, 0, "SEA|Inline|NoCloudFetch|GeoOff|ComplexOff"), + Arguments.of(0, 0, 0, 0, 1, "SEA|Inline|NoCloudFetch|GeoOff|ComplexOn"), + Arguments.of(0, 0, 0, 1, 0, "SEA|Inline|NoCloudFetch|GeoOn|ComplexOff"), + Arguments.of(0, 0, 0, 1, 1, "SEA|Inline|NoCloudFetch|GeoOn|ComplexOn"), + + // =================================================================== + // SEA (useThrift=0), Arrow (arrow=1), No CloudFetch + // =================================================================== + Arguments.of(0, 1, 0, 0, 0, "SEA|Arrow|NoCloudFetch|GeoOff|ComplexOff"), + Arguments.of(0, 1, 0, 0, 1, "SEA|Arrow|NoCloudFetch|GeoOff|ComplexOn"), + Arguments.of(0, 1, 0, 1, 0, "SEA|Arrow|NoCloudFetch|GeoOn|ComplexOff"), + Arguments.of(0, 1, 0, 1, 1, "SEA|Arrow|NoCloudFetch|GeoOn|ComplexOn"), + + // =================================================================== + // SEA (useThrift=0), Arrow (arrow=1), CloudFetch enabled + // =================================================================== + Arguments.of(0, 1, 1, 0, 0, "SEA|Arrow|CloudFetch|GeoOff|ComplexOff"), + Arguments.of(0, 1, 1, 0, 1, "SEA|Arrow|CloudFetch|GeoOff|ComplexOn"), + Arguments.of(0, 1, 1, 1, 0, "SEA|Arrow|CloudFetch|GeoOn|ComplexOff"), + Arguments.of(0, 1, 1, 1, 1, "SEA|Arrow|CloudFetch|GeoOn|ComplexOn"), + + // =================================================================== + // Thrift (useThrift=1), Inline (arrow=0), No CloudFetch + // =================================================================== + Arguments.of(1, 0, 0, 0, 0, "Thrift|Inline|NoCloudFetch|GeoOff|ComplexOff"), + Arguments.of(1, 0, 0, 0, 1, "Thrift|Inline|NoCloudFetch|GeoOff|ComplexOn"), + Arguments.of(1, 0, 0, 1, 0, "Thrift|Inline|NoCloudFetch|GeoOn|ComplexOff"), + Arguments.of(1, 0, 0, 1, 1, "Thrift|Inline|NoCloudFetch|GeoOn|ComplexOn"), + + // =================================================================== + // Thrift (useThrift=1), Arrow (arrow=1), No CloudFetch + // =================================================================== + Arguments.of(1, 1, 0, 0, 0, "Thrift|Arrow|NoCloudFetch|GeoOff|ComplexOff"), + Arguments.of(1, 1, 0, 0, 1, "Thrift|Arrow|NoCloudFetch|GeoOff|ComplexOn"), + Arguments.of(1, 1, 0, 1, 0, "Thrift|Arrow|NoCloudFetch|GeoOn|ComplexOff"), + Arguments.of(1, 1, 0, 1, 1, "Thrift|Arrow|NoCloudFetch|GeoOn|ComplexOn"), + + // =================================================================== + // Thrift (useThrift=1), Arrow (arrow=1), CloudFetch enabled + // =================================================================== + Arguments.of(1, 1, 1, 0, 0, "Thrift|Arrow|CloudFetch|GeoOff|ComplexOff"), + Arguments.of(1, 1, 1, 0, 1, "Thrift|Arrow|CloudFetch|GeoOff|ComplexOn"), + Arguments.of(1, 1, 1, 1, 0, "Thrift|Arrow|CloudFetch|GeoOn|ComplexOff"), + Arguments.of(1, 1, 1, 1, 1, "Thrift|Arrow|CloudFetch|GeoOn|ComplexOn")); + } + + @ParameterizedTest(name = "{5}") + @MethodSource("provideAllConfigurations") + void testGeospatialPoint( + int useThrift, + int enableArrow, + int cloudFetch, + int enableGeoSupport, + int enableComplexSupport, + String desc) + throws SQLException { + + setupConnection(useThrift, enableArrow, enableGeoSupport, enableComplexSupport); + + // Build SQL query - conditionally add sequence for CloudFetch + String sql = + "SELECT " + + "ST_POINT(1, 2, 4326) as geom_point, " + + "ST_GeogFromText('POINT(-122.4194 37.7749)') as geog_point"; + + if (cloudFetch == 1) { + sql += " FROM explode(sequence(1, 1000000)) AS seq"; + } + + ResultSet rs = executeQuery(connection, sql); + assertNotNull(rs, "ResultSet should not be null for config: " + desc); + + ResultSetMetaData rsm = rs.getMetaData(); + assertEquals(2, rsm.getColumnCount(), "Should have 2 columns: " + desc); + + // Assert CloudFetch usage when configured + if (cloudFetch == 1) { + assertTrue( + ((DatabricksResultSetMetaData) rsm).getIsCloudFetchUsed(), + "CloudFetch should be used with 1M rows: " + desc); + } + + // Validate ONLY first row + assertTrue(rs.next(), "Should have at least one row for config: " + desc); + + // Geospatial objects returned only when: + // 1. Both EnableGeoSpatialSupport=1 AND EnableComplexDatatypeSupport=1 + // 2. NOT in Thrift + Inline mode (Thrift doesn't support geospatial without Arrow) + boolean shouldReturnGeospatialObjects = + enableGeoSupport == 1 && enableComplexSupport == 1 && !(useThrift == 1 && enableArrow == 0); + + if (shouldReturnGeospatialObjects) { + validateGeospatialEnabled(rs, rsm); + } else { + validateGeospatialDisabled(rs, rsm); + } + + rs.close(); + } + + @ParameterizedTest(name = "{5}") + @MethodSource("provideAllConfigurations") + void testGeometryAny( + int useThrift, + int enableArrow, + int cloudFetch, + int enableGeoSupport, + int enableComplexSupport, + String desc) + throws SQLException { + + setupConnection(useThrift, enableArrow, enableGeoSupport, enableComplexSupport); + + // Build SQL query - GEOMETRY(ANY) with mixed SRIDs + String sql; + if (cloudFetch == 1) { + sql = + "SELECT CASE WHEN col % 2 = 1 " + + " THEN ST_GeomFromText('POINT(17 7)', 4326) " + + " ELSE ST_GeomFromText('POINT(5 5)', 0) " + + "END as geom " + + "FROM explode(sequence(1, 1000000)) AS t(col)"; + } else { + sql = + "SELECT * FROM VALUES " + + "(ST_GeomFromText('POINT(17 7)', 4326)), " + + "(ST_GeomFromText('POINT(5 5)', 0)) AS t(geom)"; + } + + ResultSet rs = executeQuery(connection, sql); + assertNotNull(rs); + + ResultSetMetaData rsm = rs.getMetaData(); + assertEquals(1, rsm.getColumnCount()); + + // Assert CloudFetch usage when configured + if (cloudFetch == 1) { + assertTrue(((DatabricksResultSetMetaData) rsm).getIsCloudFetchUsed()); + } + + // Geospatial objects returned only when: + // 1. Both EnableGeoSpatialSupport=1 AND EnableComplexDatatypeSupport=1 + // 2. NOT in Thrift + Inline mode + boolean shouldReturnGeospatialObjects = + enableGeoSupport == 1 && enableComplexSupport == 1 && !(useThrift == 1 && enableArrow == 0); + + if (shouldReturnGeospatialObjects) { + validateGeometryAnyEnabled(rs, rsm); + } else { + validateGeometryAnyDisabled(rs, rsm); + } + + rs.close(); + } + + /** Validates geospatial objects are returned (IGeometry/IGeography with correct metadata). */ + private void validateGeospatialEnabled(ResultSet rs, ResultSetMetaData rsm) throws SQLException { + + // Column 1: GEOMETRY POINT + assertEquals("geom_point", rsm.getColumnName(1)); + assertEquals(Types.OTHER, rsm.getColumnType(1)); + assertEquals("GEOMETRY(4326)", rsm.getColumnTypeName(1)); + assertEquals("com.databricks.jdbc.api.IGeometry", rsm.getColumnClassName(1)); + + Object geomObj = rs.getObject("geom_point"); + assertNotNull(geomObj); + assertInstanceOf(IGeometry.class, geomObj); + + IGeometry geom = (IGeometry) geomObj; + assertEquals(4326, geom.getSRID()); + assertEquals("POINT(1 2)", geom.getWKT()); + + // Column 2: GEOGRAPHY POINT + assertEquals("geog_point", rsm.getColumnName(2)); + assertEquals(Types.OTHER, rsm.getColumnType(2)); + assertEquals("GEOGRAPHY(4326)", rsm.getColumnTypeName(2)); + assertEquals("com.databricks.jdbc.api.IGeography", rsm.getColumnClassName(2)); + + Object geogObj = rs.getObject("geog_point"); + assertNotNull(geogObj); + assertInstanceOf(IGeography.class, geogObj); + + IGeography geog = (IGeography) geogObj; + assertEquals(4326, geog.getSRID()); + assertEquals("POINT(-122.4194 37.7749)", geog.getWKT()); + } + + /** Validates geospatial data is returned as STRING. */ + private void validateGeospatialDisabled(ResultSet rs, ResultSetMetaData rsm) throws SQLException { + + // Column 1: GEOMETRY POINT (as STRING) + assertEquals("geom_point", rsm.getColumnName(1)); + assertEquals(Types.VARCHAR, rsm.getColumnType(1)); + assertEquals("STRING", rsm.getColumnTypeName(1)); + assertEquals("java.lang.String", rsm.getColumnClassName(1)); + + Object geomObj = rs.getObject("geom_point"); + assertNotNull(geomObj); + assertInstanceOf(String.class, geomObj); + + String geomStr = (String) geomObj; + assertEquals("SRID=4326;POINT(1 2)", geomStr); + + // Column 2: GEOGRAPHY POINT (as STRING) + assertEquals("geog_point", rsm.getColumnName(2)); + assertEquals(Types.VARCHAR, rsm.getColumnType(2)); + assertEquals("STRING", rsm.getColumnTypeName(2)); + assertEquals("java.lang.String", rsm.getColumnClassName(2)); + + Object geogObj = rs.getObject("geog_point"); + assertNotNull(geogObj); + assertInstanceOf(String.class, geogObj); + + String geogStr = (String) geogObj; + assertEquals("SRID=4326;POINT(-122.4194 37.7749)", geogStr); + } + + /** Validates GEOMETRY(ANY) objects are returned with mixed SRIDs. */ + private void validateGeometryAnyEnabled(ResultSet rs, ResultSetMetaData rsm) throws SQLException { + + // Column metadata assertions + assertEquals("geom", rsm.getColumnName(1)); + assertEquals(Types.OTHER, rsm.getColumnType(1)); + assertEquals("GEOMETRY(ANY)", rsm.getColumnTypeName(1)); + assertEquals("com.databricks.jdbc.api.IGeometry", rsm.getColumnClassName(1)); + + // Row 1: POINT(17 7) with SRID 4326 + assertTrue(rs.next()); + Object geomObj1 = rs.getObject("geom"); + assertNotNull(geomObj1); + assertInstanceOf(IGeometry.class, geomObj1); + + IGeometry geom1 = (IGeometry) geomObj1; + assertEquals(4326, geom1.getSRID()); + assertEquals("POINT(17 7)", geom1.getWKT()); + + // Row 2: POINT(5 5) with SRID 0 + assertTrue(rs.next()); + Object geomObj2 = rs.getObject("geom"); + assertNotNull(geomObj2); + assertInstanceOf(IGeometry.class, geomObj2); + + IGeometry geom2 = (IGeometry) geomObj2; + assertEquals(0, geom2.getSRID()); + assertEquals("POINT(5 5)", geom2.getWKT()); + } + + /** Validates GEOMETRY(ANY) data is returned as STRING. */ + private void validateGeometryAnyDisabled(ResultSet rs, ResultSetMetaData rsm) + throws SQLException { + + // Column metadata assertions + assertEquals("geom", rsm.getColumnName(1)); + assertEquals(Types.VARCHAR, rsm.getColumnType(1)); + assertEquals("STRING", rsm.getColumnTypeName(1)); + assertEquals("java.lang.String", rsm.getColumnClassName(1)); + + // Row 1: POINT(17 7) with SRID 4326 (as STRING) + assertTrue(rs.next()); + Object geomObj1 = rs.getObject("geom"); + assertNotNull(geomObj1); + assertInstanceOf(String.class, geomObj1); + + String geomStr1 = (String) geomObj1; + assertEquals("SRID=4326;POINT(17 7)", geomStr1); + + // Row 2: POINT(5 5) with SRID 0 (as STRING) + assertTrue(rs.next()); + Object geomObj2 = rs.getObject("geom"); + assertNotNull(geomObj2); + assertInstanceOf(String.class, geomObj2); + + String geomStr2 = (String) geomObj2; + assertEquals("POINT(5 5)", geomStr2); + } +} diff --git a/src/test/java/com/databricks/jdbc/integration/e2e/ThriftCloudFetchTests.java b/src/test/java/com/databricks/jdbc/integration/e2e/ThriftCloudFetchTests.java index 1e166b3eaa..5465bd8d29 100644 --- a/src/test/java/com/databricks/jdbc/integration/e2e/ThriftCloudFetchTests.java +++ b/src/test/java/com/databricks/jdbc/integration/e2e/ThriftCloudFetchTests.java @@ -15,12 +15,12 @@ import com.databricks.jdbc.api.internal.IDatabricksSession; import com.databricks.jdbc.dbclient.IDatabricksClient; import com.databricks.jdbc.dbclient.impl.common.StatementId; -import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.core.ExternalLink; import java.sql.Connection; import java.sql.ResultSet; +import java.sql.SQLException; import java.sql.Statement; import java.util.ArrayList; import java.util.Collection; @@ -138,7 +138,7 @@ private void testRefetchLinks( long chunkStartRowOffset, AbstractRemoteChunkProvider chunkProvider, IDatabricksClient client) - throws DatabricksSQLException { + throws SQLException { // Fetch from the startRowOffset of the target chunk. Collection refetchedLinks = client.getResultChunks(statementId, chunkIndex, chunkStartRowOffset).getChunkLinks(); diff --git a/src/test/java/com/databricks/jdbc/integration/e2e/TransactionTests.java b/src/test/java/com/databricks/jdbc/integration/e2e/TransactionTests.java index 69f4c858ba..c580500dee 100644 --- a/src/test/java/com/databricks/jdbc/integration/e2e/TransactionTests.java +++ b/src/test/java/com/databricks/jdbc/integration/e2e/TransactionTests.java @@ -344,16 +344,25 @@ void testCommitWithoutActiveTransaction() throws SQLException { } @Test - @DisplayName("Should safely handle ROLLBACK without active transaction (no-op)") - void testRollbackWithoutActiveTransactionIsNoOp() throws SQLException { + @DisplayName("Should throw exception when rolling back without active transaction") + void testRollbackWithoutActiveTransactionThrows() throws SQLException { // With autoCommit=true (no active transaction) assertTrue(connection.getAutoCommit()); - // ROLLBACK is more forgiving than COMMIT - it should succeed as a no-op - // when there's no active transaction - assertDoesNotThrow( - () -> connection.rollback(), - "ROLLBACK should be a safe no-op when autocommit=true (no active transaction)"); + // ROLLBACK should throw an exception when there is no active transaction + // (connection is in auto-commit mode) + SQLException exception = + assertThrows( + SQLException.class, + () -> connection.rollback(), + "ROLLBACK should throw exception when autocommit=true (no active transaction)"); + + assertTrue( + exception.getMessage().contains("auto-commit") + || exception.getMessage().contains("rollback") + || exception.getMessage().contains("No active transaction"), + "Exception message should indicate rollback is not valid in auto-commit mode. Got: " + + exception.getMessage()); // Verify connection is still usable assertTrue(connection.getAutoCommit()); diff --git a/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/EnableMultipleCatalogSupportIntegrationTests.java b/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/EnableMultipleCatalogSupportIntegrationTests.java index 169ec8cc45..bed8464d0b 100644 --- a/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/EnableMultipleCatalogSupportIntegrationTests.java +++ b/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/EnableMultipleCatalogSupportIntegrationTests.java @@ -55,8 +55,12 @@ void testGetSchemasWithNullCatalogMultipleCatalogSupportEnabled() throws SQLExce } while (schemas.next()); assertTrue( - schemaCount > 1, - "With enableMultipleCatalogSupport=1, should find schemas from multiple catalogs"); + schemaCount >= 1, + "With enableMultipleCatalogSupport=1, should find at least one schema"); + assertTrue( + distinctCatalogs.size() > 1, + "With enableMultipleCatalogSupport=1, should find schemas from multiple catalogs. Found: " + + distinctCatalogs); } } } @@ -94,7 +98,11 @@ void testGetSchemasWithNullCatalogMultipleCatalogSupportDisabled() throws SQLExc assertTrue( schemaCount >= 1, - "With enableMultipleCatalogSupport=0, should find schemas from single catalog"); + "With enableMultipleCatalogSupport=0, should find at least one schema"); + assertTrue( + distinctCatalogs.size() <= 2, + "With enableMultipleCatalogSupport=0, should find schemas from at most 2 catalogs (current catalog + possibly empty). Found: " + + distinctCatalogs); } } } diff --git a/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/ThriftCloudFetchFakeIntegrationTests.java b/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/ThriftCloudFetchFakeIntegrationTests.java index d138d3e828..8ee4aefe55 100644 --- a/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/ThriftCloudFetchFakeIntegrationTests.java +++ b/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/ThriftCloudFetchFakeIntegrationTests.java @@ -15,11 +15,11 @@ import com.databricks.jdbc.api.internal.IDatabricksSession; import com.databricks.jdbc.dbclient.IDatabricksClient; import com.databricks.jdbc.dbclient.impl.common.StatementId; -import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.integration.fakeservice.AbstractFakeServiceIntegrationTests; import com.databricks.jdbc.model.core.ExternalLink; import java.sql.Connection; import java.sql.ResultSet; +import java.sql.SQLException; import java.sql.Statement; import java.util.ArrayList; import java.util.Collection; @@ -127,7 +127,7 @@ private void testRefetchLinks( long chunkStartRowOffset, AbstractRemoteChunkProvider chunkProvider, IDatabricksClient client) - throws DatabricksSQLException { + throws SQLException { // Fetch from the startRowOffset of the target chunk Collection refetchedLinks = diff --git a/src/test/java/com/databricks/jdbc/telemetry/TelemetryClientFactoryTest.java b/src/test/java/com/databricks/jdbc/telemetry/TelemetryClientFactoryTest.java index 5ab298f232..eaa1edbb8f 100644 --- a/src/test/java/com/databricks/jdbc/telemetry/TelemetryClientFactoryTest.java +++ b/src/test/java/com/databricks/jdbc/telemetry/TelemetryClientFactoryTest.java @@ -13,9 +13,14 @@ import java.util.HashMap; import java.util.Map; import java.util.Properties; +import java.util.stream.Stream; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; import org.mockito.Mock; import org.mockito.MockedStatic; import org.mockito.junit.jupiter.MockitoExtension; @@ -87,6 +92,220 @@ public void testGetNoOpTelemetryClientWhenDatabricksConfigIsNull() throws Except } } + // UUID Tracking Tests + + @ParameterizedTest + @ValueSource(ints = {1, 5}) + void testMultipleConnectionsToSameHostShareClient(int connectionCount) throws Exception { + String host = "shared-host.databricks.net"; + IDatabricksConnectionContext[] contexts = new IDatabricksConnectionContext[connectionCount]; + + try (MockedStatic mockedStatic = mockStatic(TelemetryHelper.class)) { + setupTelemetryHelperMock(mockedStatic); + + // Create multiple connections to same host + for (int i = 0; i < connectionCount; i++) { + contexts[i] = createMockContext("uuid-" + i, host, true); + } + + // Get clients - all should return same instance + ITelemetryClient firstClient = + TelemetryClientFactory.getInstance().getTelemetryClient(contexts[0]); + for (int i = 1; i < connectionCount; i++) { + ITelemetryClient client = + TelemetryClientFactory.getInstance().getTelemetryClient(contexts[i]); + assertSame(firstClient, client, "Connection " + i + " should share same client"); + } + + // Verify single holder with all UUIDs + assertEquals(1, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size()); + + // Close connections one by one - holder should remain until last connection closes + for (int i = 0; i < connectionCount - 1; i++) { + TelemetryClientFactory.getInstance().closeTelemetryClient(contexts[i]); + assertEquals( + 1, + TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size(), + "Client should remain after closing connection " + i); + } + + // Close last connection - client should be removed + TelemetryClientFactory.getInstance().closeTelemetryClient(contexts[connectionCount - 1]); + assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size()); + } + } + + private static Stream provideAuthenticationScenarios() { + return Stream.of(Arguments.of("authenticated", true), Arguments.of("no-auth", false)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("provideAuthenticationScenarios") + void testDifferentHostsGetDifferentClients(String scenario, boolean authenticated) + throws Exception { + try (MockedStatic mockedStatic = mockStatic(TelemetryHelper.class)) { + setupTelemetryHelperMock(mockedStatic); + if (authenticated) { + mockedStatic + .when(() -> TelemetryHelper.getDatabricksConfigSafely(any())) + .thenReturn(databricksConfig); + } + + IDatabricksConnectionContext ctx1 = + createMockContext("uuid-1", "host1.databricks.net", authenticated); + IDatabricksConnectionContext ctx2 = + createMockContext("uuid-2", "host2.databricks.net", authenticated); + + ITelemetryClient client1 = TelemetryClientFactory.getInstance().getTelemetryClient(ctx1); + ITelemetryClient client2 = TelemetryClientFactory.getInstance().getTelemetryClient(ctx2); + + assertNotSame(client1, client2, "Different hosts should get different clients"); + + int holderCount = + authenticated + ? TelemetryClientFactory.getInstance().telemetryClientHolders.size() + : TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size(); + assertEquals(2, holderCount); + + // Close first connection + TelemetryClientFactory.getInstance().closeTelemetryClient(ctx1); + holderCount = + authenticated + ? TelemetryClientFactory.getInstance().telemetryClientHolders.size() + : TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size(); + assertEquals(1, holderCount, "First client should be closed"); + + // Close second connection + TelemetryClientFactory.getInstance().closeTelemetryClient(ctx2); + holderCount = + authenticated + ? TelemetryClientFactory.getInstance().telemetryClientHolders.size() + : TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size(); + assertEquals(0, holderCount, "All clients should be closed"); + } + } + + @Test + void testConnectionParamCacheCleanedOnClose() throws Exception { + String uuid = "test-uuid-123"; + String host = "test-host.databricks.net"; + + try (MockedStatic mockedStatic = mockStatic(TelemetryHelper.class)) { + setupTelemetryHelperMock(mockedStatic); + + IDatabricksConnectionContext ctx = createMockContext(uuid, host, false); + TelemetryClientFactory.getInstance().getTelemetryClient(ctx); + + // Close connection and verify cleanup is called + TelemetryClientFactory.getInstance().closeTelemetryClient(ctx); + + // Verify removeConnectionParameters was called + mockedStatic.verify(() -> TelemetryHelper.removeConnectionParameters(uuid), times(1)); + } + } + + @Test + void testUUIDTrackingWithMixedConnectionLifecycles() throws Exception { + String host = "shared-host.databricks.net"; + + try (MockedStatic mockedStatic = mockStatic(TelemetryHelper.class)) { + setupTelemetryHelperMock(mockedStatic); + + IDatabricksConnectionContext conn1 = createMockContext("uuid-1", host, false); + IDatabricksConnectionContext conn2 = createMockContext("uuid-2", host, false); + IDatabricksConnectionContext conn3 = createMockContext("uuid-3", host, false); + + // Open connections in order: 1, 2, 3 + ITelemetryClient client1 = TelemetryClientFactory.getInstance().getTelemetryClient(conn1); + ITelemetryClient client2 = TelemetryClientFactory.getInstance().getTelemetryClient(conn2); + ITelemetryClient client3 = TelemetryClientFactory.getInstance().getTelemetryClient(conn3); + + // All should share same client + assertSame(client1, client2); + assertSame(client2, client3); + assertEquals(1, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size()); + + // Close middle connection - client should remain + TelemetryClientFactory.getInstance().closeTelemetryClient(conn2); + assertEquals(1, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size()); + + // Close first connection - client should still remain + TelemetryClientFactory.getInstance().closeTelemetryClient(conn1); + assertEquals(1, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size()); + + // Close last connection - NOW client should be removed + TelemetryClientFactory.getInstance().closeTelemetryClient(conn3); + assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size()); + } + } + + @Test + void testDoubleCloseDoesNotCauseErrors() throws Exception { + String host = "test-host.databricks.net"; + + try (MockedStatic mockedStatic = mockStatic(TelemetryHelper.class)) { + setupTelemetryHelperMock(mockedStatic); + + IDatabricksConnectionContext ctx = createMockContext("uuid-1", host, false); + TelemetryClientFactory.getInstance().getTelemetryClient(ctx); + + // Close twice - should not throw + assertDoesNotThrow(() -> TelemetryClientFactory.getInstance().closeTelemetryClient(ctx)); + assertDoesNotThrow(() -> TelemetryClientFactory.getInstance().closeTelemetryClient(ctx)); + + assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size()); + } + } + + @Test + void testSameConnectionCallingGetClientMultipleTimes() throws Exception { + String host = "test-host.databricks.net"; + + try (MockedStatic mockedStatic = mockStatic(TelemetryHelper.class)) { + setupTelemetryHelperMock(mockedStatic); + + IDatabricksConnectionContext conn = createMockContext("uuid-1", host, false); + + // Call getTelemetryClient multiple times with same connection (e.g., retry logic) + TelemetryClientFactory.getInstance().getTelemetryClient(conn); + TelemetryClientFactory.getInstance().getTelemetryClient(conn); + TelemetryClientFactory.getInstance().getTelemetryClient(conn); + + assertEquals(1, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size()); + + // Single close should remove client - UUID set deduplicates, old refCount would leak + TelemetryClientFactory.getInstance().closeTelemetryClient(conn); + assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size()); + } + } + + // Helper methods + + private IDatabricksConnectionContext createMockContext( + String uuid, String host, boolean authenticated) { + IDatabricksConnectionContext ctx = mock(IDatabricksConnectionContext.class); + lenient().when(ctx.getConnectionUuid()).thenReturn(uuid); + lenient().when(ctx.getHost()).thenReturn(host); + lenient().when(ctx.isTelemetryEnabled()).thenReturn(true); + lenient().when(ctx.getTelemetryBatchSize()).thenReturn(10); + lenient().when(ctx.getTelemetryFlushIntervalInMilliseconds()).thenReturn(5000); + if (authenticated) { + lenient().when(clientConfigurator.getDatabricksConfig()).thenReturn(databricksConfig); + } + return ctx; + } + + private void setupTelemetryHelperMock(MockedStatic mockedStatic) { + mockedStatic.when(() -> TelemetryHelper.keyOf(any())).thenCallRealMethod(); + mockedStatic.when(() -> TelemetryHelper.getDatabricksConfigSafely(any())).thenReturn(null); + mockedStatic + .when(() -> TelemetryHelper.removeConnectionParameters(anyString())) + .thenAnswer(invocation -> null); + mockedStatic + .when(() -> TelemetryHelper.isTelemetryAllowedForConnection(any())) + .thenReturn(true); + } + private void setupMocksForTelemetryClient(IDatabricksConnectionContext context) { TelemetryClientFactory.getInstance().closeTelemetryClient(context); TelemetryAuthHelper.setupAuthMocks(context, clientConfigurator); diff --git a/src/test/java/com/databricks/jdbc/telemetry/TelemetryClientTest.java b/src/test/java/com/databricks/jdbc/telemetry/TelemetryClientTest.java index bd72ee0c85..fac0432884 100644 --- a/src/test/java/com/databricks/jdbc/telemetry/TelemetryClientTest.java +++ b/src/test/java/com/databricks/jdbc/telemetry/TelemetryClientTest.java @@ -17,6 +17,8 @@ import java.util.Map; import java.util.Properties; import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import org.apache.http.HttpHeaders; import org.apache.http.StatusLine; @@ -42,6 +44,14 @@ public class TelemetryClientTest { @Mock StatusLine mockStatusLine; @Mock DatabricksConfig databricksConfig; + private ScheduledExecutorService createMockScheduler() { + return mock(ScheduledExecutorService.class); + } + + private ScheduledExecutorService createRealScheduler() { + return Executors.newSingleThreadScheduledExecutor(); + } + @Test public void testExportEvent() throws Exception { try (MockedStatic factoryMocked = @@ -59,7 +69,8 @@ public void testExportEvent() throws Exception { IDatabricksConnectionContext context = DatabricksConnectionContext.parse(JDBC_URL, new Properties()); TelemetryClient client = - new TelemetryClient(context, MoreExecutors.newDirectExecutorService()); + new TelemetryClient( + context, MoreExecutors.newDirectExecutorService(), createMockScheduler()); client.exportEvent( new TelemetryFrontendLog(TelemetryLogLevel.ERROR).setFrontendLogEventId("event1")); @@ -101,7 +112,11 @@ public void testExportEvent_authenticated() throws Exception { IDatabricksConnectionContext context = DatabricksConnectionContext.parse(JDBC_URL, new Properties()); TelemetryClient client = - new TelemetryClient(context, MoreExecutors.newDirectExecutorService(), databricksConfig); + new TelemetryClient( + context, + MoreExecutors.newDirectExecutorService(), + createMockScheduler(), + databricksConfig); client.exportEvent( new TelemetryFrontendLog(TelemetryLogLevel.ERROR).setFrontendLogEventId("event1")); @@ -141,7 +156,8 @@ public void testExportEventDoesNotThrowErrorsInFailures() throws Exception { IDatabricksConnectionContext context = DatabricksConnectionContext.parse(JDBC_URL, new Properties()); TelemetryClient client = - new TelemetryClient(context, MoreExecutors.newDirectExecutorService()); + new TelemetryClient( + context, MoreExecutors.newDirectExecutorService(), createMockScheduler()); client.exportEvent( new TelemetryFrontendLog(TelemetryLogLevel.ERROR).setFrontendLogEventId("event1")); @@ -155,6 +171,7 @@ public void testExportEventDoesNotThrowErrorsInFailures() throws Exception { @Test public void testPeriodicFlushWithAuthenticatedClient() throws Exception { + ScheduledExecutorService scheduler = createRealScheduler(); try (MockedStatic factoryMocked = mockStatic(DatabricksHttpClientFactory.class)) { DatabricksHttpClientFactory mockFactory = mock(DatabricksHttpClientFactory.class); @@ -177,7 +194,8 @@ public void testPeriodicFlushWithAuthenticatedClient() throws Exception { IDatabricksConnectionContext context = DatabricksConnectionContext.parse(jdbcUrlWith2SecondsFlush, new Properties()); TelemetryClient client = - new TelemetryClient(context, MoreExecutors.newDirectExecutorService(), databricksConfig); + new TelemetryClient( + context, MoreExecutors.newDirectExecutorService(), scheduler, databricksConfig); // Add a single event that won't trigger batch flush client.exportEvent( @@ -198,6 +216,11 @@ public void testPeriodicFlushWithAuthenticatedClient() throws Exception { // Close the client to trigger final flush client.close(); assertEquals(0, client.getCurrentSize()); + } finally { + scheduler.shutdown(); + if (!scheduler.awaitTermination(5, TimeUnit.SECONDS)) { + scheduler.shutdownNow(); + } } } @@ -205,6 +228,7 @@ public void testPeriodicFlushWithAuthenticatedClient() throws Exception { public void testTimerResetOnBatchSizeFlush() throws Exception { TelemetryClient client = null; ExecutorService executor = null; + ScheduledExecutorService scheduler = createRealScheduler(); try (MockedStatic factoryMocked = mockStatic(DatabricksHttpClientFactory.class)) { DatabricksHttpClientFactory mockFactory = mock(DatabricksHttpClientFactory.class); @@ -223,7 +247,7 @@ public void testTimerResetOnBatchSizeFlush() throws Exception { IDatabricksConnectionContext context = DatabricksConnectionContext.parse(jdbcUrl, new Properties()); executor = MoreExecutors.newDirectExecutorService(); - client = new TelemetryClient(context, executor); + client = new TelemetryClient(context, executor, scheduler); // Add events to trigger batch size flush client.exportEvent( @@ -265,6 +289,10 @@ public void testTimerResetOnBatchSizeFlush() throws Exception { executor.shutdownNow(); } } + scheduler.shutdown(); + if (!scheduler.awaitTermination(5, TimeUnit.SECONDS)) { + scheduler.shutdownNow(); + } // Verify mocks were properly used verify(mockHttpClient, atLeastOnce()).execute(any()); verify(mockHttpResponse, atLeastOnce()).getStatusLine(); diff --git a/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java b/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java index cad946df2f..e9a3c73a30 100644 --- a/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java +++ b/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java @@ -219,6 +219,13 @@ public void testGetDatabricksConfigSafely_HandlesNullContext() { assertNull(result, "Should return null when context is null"); } + @Test + public void testConnectionParameterCacheCleanup() { + assertDoesNotThrow(() -> TelemetryHelper.removeConnectionParameters("test-uuid")); + assertDoesNotThrow(() -> TelemetryHelper.removeConnectionParameters(null)); + assertDoesNotThrow(() -> TelemetryHelper.clearConnectionParameterCache()); + } + @Test public void testTelemetryNotAllowedUsecase() { // Clear thread context to ensure telemetry is not allowed diff --git a/src/test/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorManagerTest.java b/src/test/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorManagerTest.java new file mode 100644 index 0000000000..c05a25396c --- /dev/null +++ b/src/test/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorManagerTest.java @@ -0,0 +1,256 @@ +package com.databricks.jdbc.telemetry.latency; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; +import com.databricks.jdbc.dbclient.impl.common.StatementId; +import java.util.stream.Stream; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +public class TelemetryCollectorManagerTest { + + private TelemetryCollectorManager manager; + private IDatabricksConnectionContext context1; + private IDatabricksConnectionContext context2; + + @BeforeEach + void setUp() { + manager = TelemetryCollectorManager.getInstance(); + manager.clear(); + + // Create mock contexts with different UUIDs + context1 = mock(IDatabricksConnectionContext.class); + when(context1.getConnectionUuid()).thenReturn("connection-uuid-1"); + + context2 = mock(IDatabricksConnectionContext.class); + when(context2.getConnectionUuid()).thenReturn("connection-uuid-2"); + } + + @AfterEach + void tearDown() { + manager.clear(); + } + + @Test + void testGetOrCreateCollectorCreatesNewInstance() { + TelemetryCollector collector = manager.getOrCreateCollector(context1); + assertNotNull(collector, "Collector should not be null"); + } + + @Test + void testGetOrCreateCollectorReturnsSameInstanceForSameConnection() { + TelemetryCollector collector1 = manager.getOrCreateCollector(context1); + TelemetryCollector collector2 = manager.getOrCreateCollector(context1); + + assertSame( + collector1, collector2, "Should return the same instance for the same connection UUID"); + } + + private static Stream provideDifferentConnectionScenarios() { + IDatabricksConnectionContext ctx1 = mock(IDatabricksConnectionContext.class); + when(ctx1.getConnectionUuid()).thenReturn("uuid-1"); + + IDatabricksConnectionContext ctx2 = mock(IDatabricksConnectionContext.class); + when(ctx2.getConnectionUuid()).thenReturn("uuid-2"); + + IDatabricksConnectionContext ctx3 = mock(IDatabricksConnectionContext.class); + when(ctx3.getConnectionUuid()).thenReturn("uuid-3"); + + return Stream.of( + Arguments.of("different UUIDs", ctx1, ctx2), + Arguments.of("another pair of different UUIDs", ctx1, ctx3), + Arguments.of("third pair of different UUIDs", ctx2, ctx3)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("provideDifferentConnectionScenarios") + void testGetOrCreateCollectorReturnsDifferentInstancesForDifferentConnections( + String scenarioName, + IDatabricksConnectionContext contextA, + IDatabricksConnectionContext contextB) { + TelemetryCollector collector1 = manager.getOrCreateCollector(contextA); + TelemetryCollector collector2 = manager.getOrCreateCollector(contextB); + + assertNotSame( + collector1, collector2, "Should return different instances for scenario: " + scenarioName); + } + + @Test + void testPerConnectionIsolation() { + TelemetryCollector collector1 = manager.getOrCreateCollector(context1); + TelemetryCollector collector2 = manager.getOrCreateCollector(context2); + + // Record telemetry for connection 1 + collector1.recordChunkDownloadLatency("statement-1", 0, 100); + + // Record telemetry for connection 2 + collector2.recordChunkDownloadLatency("statement-2", 0, 200); + + // Verify connection 1 has its telemetry + assertNotNull( + collector1.getOrCreateTelemetryDetails("statement-1"), + "Connection 1 should have statement-1"); + + // Verify connection 2 has its telemetry + assertNotNull( + collector2.getOrCreateTelemetryDetails("statement-2"), + "Connection 2 should have statement-2"); + + // Verify collectors are independent by checking they return different instances + assertNotSame( + collector1, + collector2, + "Different connections should have completely independent collectors"); + } + + private static Stream provideContextsForRemoval() { + IDatabricksConnectionContext normalContext = mock(IDatabricksConnectionContext.class); + when(normalContext.getConnectionUuid()).thenReturn("test-uuid"); + + IDatabricksConnectionContext nullUuidContext = mock(IDatabricksConnectionContext.class); + when(nullUuidContext.getConnectionUuid()).thenReturn(null); + + return Stream.of(normalContext, nullUuidContext, null); + } + + @ParameterizedTest + @MethodSource("provideContextsForRemoval") + void testRemoveCollectorReturnsCollector(IDatabricksConnectionContext context) { + TelemetryCollector collector = manager.getOrCreateCollector(context); + TelemetryCollector removed = manager.removeCollector(context); + + assertSame(collector, removed, "Should return the removed collector"); + } + + @ParameterizedTest + @MethodSource("provideContextsForRemoval") + void testRemoveCollectorRemovesFromManager(IDatabricksConnectionContext context) { + TelemetryCollector original = manager.getOrCreateCollector(context); + manager.removeCollector(context); + + // Getting collector again should create a new instance + TelemetryCollector newCollector = manager.getOrCreateCollector(context); + assertNotNull(newCollector, "Should create a new collector after removal"); + assertNotSame( + original, newCollector, "Should be a different instance after removal and re-creation"); + } + + @Test + void testRemoveCollectorDoesNotAffectOtherConnections() { + TelemetryCollector collector1 = manager.getOrCreateCollector(context1); + TelemetryCollector collector2 = manager.getOrCreateCollector(context2); + + // Remove connection 1 + manager.removeCollector(context1); + + // Connection 2 should still have the same collector + TelemetryCollector stillCollector2 = manager.getOrCreateCollector(context2); + assertSame(collector2, stillCollector2, "Connection 2's collector should be unaffected"); + } + + @ParameterizedTest + @MethodSource("provideContextsForRemoval") + void testRemoveNonExistentCollectorReturnsNull(IDatabricksConnectionContext context) { + TelemetryCollector removed = manager.removeCollector(context); + assertNull(removed, "Should return null for non-existent collector"); + } + + private static Stream provideNullContextScenarios() { + IDatabricksConnectionContext contextWithNullUuid = mock(IDatabricksConnectionContext.class); + when(contextWithNullUuid.getConnectionUuid()).thenReturn(null); + + return Stream.of( + Arguments.of("null context and null context", null, null), + Arguments.of("context with null UUID and null context", contextWithNullUuid, null), + Arguments.of( + "context with null UUID and context with null UUID", + contextWithNullUuid, + contextWithNullUuid)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("provideNullContextScenarios") + void testNullContextUsesDefaultConnection( + String scenarioName, + IDatabricksConnectionContext context1, + IDatabricksConnectionContext context2) { + TelemetryCollector collector1 = manager.getOrCreateCollector(context1); + TelemetryCollector collector2 = manager.getOrCreateCollector(context2); + + assertSame( + collector1, + collector2, + "Both contexts should use the same default collector for scenario: " + scenarioName); + } + + @Test + void testClearRemovesAllCollectors() { + manager.getOrCreateCollector(context1); + manager.getOrCreateCollector(context2); + + manager.clear(); + + // After clear, getting collectors should create new instances + TelemetryCollector newCollector1 = manager.getOrCreateCollector(context1); + TelemetryCollector newCollector2 = manager.getOrCreateCollector(context2); + + assertNotNull(newCollector1); + assertNotNull(newCollector2); + } + + @Test + void testExportAndRemoveOnConnectionClose() { + TelemetryCollector collector = manager.getOrCreateCollector(context1); + + // Record some telemetry + StatementId statementId = new StatementId("test-statement-id"); + collector.recordChunkDownloadLatency(statementId.toSQLExecStatementId(), 0, 100); + collector.recordChunkDownloadLatency(statementId.toSQLExecStatementId(), 1, 150); + + assertNotNull( + collector.getOrCreateTelemetryDetails(statementId.toSQLExecStatementId()), + "Should have telemetry details before removal"); + + // Simulate connection close: export and remove + TelemetryCollector removed = manager.removeCollector(context1); + assertNotNull(removed, "Should return the removed collector"); + removed.exportAllPendingTelemetryDetails(); + + // Verify collector was removed from manager + TelemetryCollector newCollector = manager.getOrCreateCollector(context1); + assertNotSame(collector, newCollector, "After removal, should get a new collector instance"); + } + + @Test + void testMultipleConnectionsWithSameHostAreIsolated() { + // Even if connections are to the same host, they should have separate collectors + // because they have different connection UUIDs + IDatabricksConnectionContext conn1 = mock(IDatabricksConnectionContext.class); + when(conn1.getConnectionUuid()).thenReturn("uuid-1"); + + IDatabricksConnectionContext conn2 = mock(IDatabricksConnectionContext.class); + when(conn2.getConnectionUuid()).thenReturn("uuid-2"); + + TelemetryCollector collector1 = manager.getOrCreateCollector(conn1); + TelemetryCollector collector2 = manager.getOrCreateCollector(conn2); + + assertNotSame( + collector1, + collector2, + "Different connections to same host should have different collectors"); + + // Close connection 1 + manager.removeCollector(conn1); + + // Connection 2 should still have its collector + TelemetryCollector stillCollector2 = manager.getOrCreateCollector(conn2); + assertSame( + collector2, stillCollector2, "Connection 2 should not be affected by closing connection 1"); + } +} diff --git a/src/test/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorTest.java b/src/test/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorTest.java index 3655d53379..31a79efbcf 100644 --- a/src/test/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorTest.java +++ b/src/test/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorTest.java @@ -18,7 +18,7 @@ import org.mockito.MockedStatic; public class TelemetryCollectorTest { - private final TelemetryCollector handler = TelemetryCollector.getInstance(); + private final TelemetryCollector handler = new TelemetryCollector(); private static final String TEST_STATEMENT_ID = "test-statement-id"; @BeforeEach diff --git a/thin_public_pom.xml b/thin_public_pom.xml index cd8cf8206a..84bb8d876f 100644 --- a/thin_public_pom.xml +++ b/thin_public_pom.xml @@ -4,7 +4,7 @@ 4.0.0 com.databricks databricks-jdbc-thin - 3.0.7 + 3.1.1 jar Databricks JDBC Driver Thin Databricks JDBC Driver Thin JAR - requires external dependencies. diff --git a/uber-minimal-pom.xml b/uber-minimal-pom.xml index bfe65bb376..a8fac7ed6b 100644 --- a/uber-minimal-pom.xml +++ b/uber-minimal-pom.xml @@ -5,7 +5,7 @@ com.databricks databricks-jdbc - 3.0.7 + 3.1.1 jar Databricks JDBC Driver Databricks JDBC Driver.