From 9273799c7393350c8206469a0867d52e5e711e82 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 14 Aug 2026 20:52:22 +0000 Subject: [PATCH] Add Thrift-native metadata support Signed-off-by: Vu Anh Phung --- NEXT_CHANGELOG.md | 1 + .../api/impl/DatabricksConnectionContext.java | 5 + .../jdbc/api/impl/DatabricksResultSet.java | 7 + .../IDatabricksConnectionContext.java | 3 + .../jdbc/common/DatabricksJdbcConstants.java | 9 +- .../jdbc/common/DatabricksJdbcUrlParams.java | 4 + .../jdbc/common/MetadataOperationType.java | 26 +- .../jdbc/common/util/SqlStateClassifier.java | 4 +- ...ferenceKeysDatabricksResultSetAdapter.java | 27 +- .../impl/common/MetadataResultSetBuilder.java | 56 +++- .../DatabricksMetadataQueryClient.java | 61 +++- .../impl/sqlexec/DatabricksSdkClient.java | 6 + .../impl/thrift/DatabricksThriftAccessor.java | 4 +- .../jdbc/model/core/ResultManifest.java | 19 +- .../common/MetadataOperationTypeTest.java | 79 +----- .../common/MetadataResultSetBuilderTest.java | 253 ++++++++++++++++- .../DatabricksMetadataQueryClientTest.java | 263 ++++++++++++++---- .../impl/sqlexec/DatabricksSdkClientTest.java | 128 ++++++++- .../thrift/DatabricksThriftAccessorTest.java | 8 +- .../jdbc/model/core/ResultManifestTest.java | 24 ++ 20 files changed, 817 insertions(+), 170 deletions(-) create mode 100644 src/test/java/com/databricks/jdbc/model/core/ResultManifestTest.java diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index e12c8875f4..f82e2f1233 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -3,6 +3,7 @@ ## [Unreleased] ### Added +- Added `EnableThriftNativeMetadata` to request and consume supported Thrift-native SEA metadata results. ### Updated - `DatabaseMetaData.getColumns(...)` with a `null` catalog now issues a single `SHOW COLUMNS IN ALL CATALOGS` statement (consistent with `getSchemas`/`getTables`) instead of enumerating every catalog and issuing a per-catalog `SHOW COLUMNS`. Older DBR versions that do not support the syntax transparently fall back to the previous enumerate-and-fan-out behavior. 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 dfa4b70f9c..c209794b12 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -1511,6 +1511,11 @@ public boolean isSeaSyncMetadataEnabled() { return getParameter(DatabricksJdbcUrlParams.ENABLE_SEA_SYNC_METADATA).equals("1"); } + @Override + public boolean isThriftNativeMetadataEnabled() { + return getParameter(DatabricksJdbcUrlParams.ENABLE_THRIFT_NATIVE_METADATA).equals("1"); + } + @Override public boolean getDisableOauthRefreshToken() { return getParameter(DatabricksJdbcUrlParams.DISABLE_OAUTH_REFRESH_TOKEN, "1").equals("1"); 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 cde481ccf3..a2a3ce4ab1 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -81,6 +81,7 @@ enum ResultSetType { // Set to true when next() returns false for the bounded-SEA path, so that isAfterLast() // returns true only after the cursor has moved PAST the last row (not while ON it). private boolean boundedSeaExhausted = false; + private boolean thriftNativeMetadataResult = false; // Cached telemetry collector resolved once at construction time to avoid // per-row overhead in next(). The connection-to-collector mapping is stable @@ -109,6 +110,8 @@ public DatabricksResultSet( throws SQLException { this.executionStatus = new ExecutionStatus(statementStatus); this.statementId = statementId; + this.thriftNativeMetadataResult = + resultManifest != null && Boolean.TRUE.equals(resultManifest.getIsNativeMetadataResult()); if (resultData != null) { this.executionResult = ExecutionResultFactory.getResultSet( @@ -803,6 +806,10 @@ public ResultSetMetaData getMetaData() throws SQLException { return resultSetMetaData; } + public boolean isThriftNativeMetadataResult() { + return thriftNativeMetadataResult; + } + /** * Checks if the given type name represents a geospatial type (GEOMETRY or GEOGRAPHY). * 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 fb0d745a78..61cb4929f6 100644 --- a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java @@ -443,6 +443,9 @@ default int getHeartbeatIntervalSeconds() { */ boolean isSeaSyncMetadataEnabled(); + /** Returns whether SEA metadata requests should require Thrift-native execution. */ + boolean isThriftNativeMetadataEnabled(); + /** Returns whether OAuth refresh tokens should be disabled (omit offline_access by default). */ boolean getDisableOauthRefreshToken(); diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java index 896c8a0a0f..3e874a2566 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java @@ -107,9 +107,16 @@ public final class DatabricksJdbcConstants { public static final String AAD_CLIENT_ID = "databricks-sql-jdbc"; public static final String GCP_GOOGLE_CREDENTIALS_AUTH_TYPE = "google-credentials"; public static final String GCP_GOOGLE_ID_AUTH_TYPE = "google-id"; - public static final String DEFAULT_HTTP_EXCEPTION_SQLSTATE = "08000"; + + /** SQL state used by Thrift for generic operation errors (SQLSTATE 08000). */ + public static final String OPERATION_ERROR_SQLSTATE = "08000"; + + public static final String DEFAULT_HTTP_EXCEPTION_SQLSTATE = OPERATION_ERROR_SQLSTATE; public static final String QUERY_EXECUTION_TIMEOUT_SQLSTATE = "57KD0"; + /** Standard SQL state for syntax error or access rule violation (SQLSTATE 42000). */ + public static final String SYNTAX_OR_ACCESS_VIOLATION_SQLSTATE = "42000"; + /** Standard SQL state for operation cancelled (SQLSTATE HY008). */ public static final String OPERATION_CANCELLED_SQLSTATE = "HY008"; diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java index 7fea2fe2c3..55fde41a7f 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java @@ -215,6 +215,10 @@ public enum DatabricksJdbcUrlParams { "UseBoundedSeaApi", "Use bounded SEA API for CloudFetch: send row_offset on GetResultData, force StreamingChunkProvider, stop relying on total_chunk_count. Requires server support.", "0"), + ENABLE_THRIFT_NATIVE_METADATA( + "EnableThriftNativeMetadata", + "Request Thrift-native SEA results for catalogs, schemas, tables, columns, functions, primary keys, imported keys, and cross references", + "0"), DISABLE_OAUTH_REFRESH_TOKEN( "DisableOauthRefreshToken", "Disable requesting OAuth refresh tokens (omit offline_access unless explicitly provided)", diff --git a/src/main/java/com/databricks/jdbc/common/MetadataOperationType.java b/src/main/java/com/databricks/jdbc/common/MetadataOperationType.java index fdadfe5ee3..9ba5a278a1 100644 --- a/src/main/java/com/databricks/jdbc/common/MetadataOperationType.java +++ b/src/main/java/com/databricks/jdbc/common/MetadataOperationType.java @@ -5,24 +5,30 @@ * HTTP headers to track which metadata operation is being performed. */ public enum MetadataOperationType { - GET_CATALOGS("GetCatalogs"), - GET_SCHEMAS("GetSchemas"), - GET_TABLES("GetTables"), - GET_COLUMNS("GetColumns"), - GET_FUNCTIONS("GetFunctions"), - GET_PRIMARY_KEYS("GetPrimaryKeys"), - GET_CROSS_REFERENCE("GetCrossReference"), - GET_PROCEDURES("GetProcedures"), - GET_PROCEDURE_COLUMNS("GetProcedureColumns"); + GET_CATALOGS("GetCatalogs", true), + GET_SCHEMAS("GetSchemas", true), + GET_TABLES("GetTables", true), + GET_COLUMNS("GetColumns", true), + GET_FUNCTIONS("GetFunctions", true), + GET_PRIMARY_KEYS("GetPrimaryKeys", true), + GET_CROSS_REFERENCE("GetCrossReference", true), + GET_PROCEDURES("GetProcedures", false), + GET_PROCEDURE_COLUMNS("GetProcedureColumns", false); private final String headerValue; + private final boolean thriftNativeSupported; - MetadataOperationType(String headerValue) { + MetadataOperationType(String headerValue, boolean thriftNativeSupported) { this.headerValue = headerValue; + this.thriftNativeSupported = thriftNativeSupported; } /** Returns the header value to be sent in the HTTP request. */ public String getHeaderValue() { return headerValue; } + + public boolean isThriftNativeSupported() { + return thriftNativeSupported; + } } diff --git a/src/main/java/com/databricks/jdbc/common/util/SqlStateClassifier.java b/src/main/java/com/databricks/jdbc/common/util/SqlStateClassifier.java index 7630128a7f..e4e6361ca1 100644 --- a/src/main/java/com/databricks/jdbc/common/util/SqlStateClassifier.java +++ b/src/main/java/com/databricks/jdbc/common/util/SqlStateClassifier.java @@ -2,6 +2,7 @@ import static com.databricks.jdbc.common.DatabricksJdbcConstants.COMMUNICATION_LINK_FAILURE_SQLSTATE; import static com.databricks.jdbc.common.DatabricksJdbcConstants.SERIALIZATION_FAILURE_SQLSTATE; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.SYNTAX_OR_ACCESS_VIOLATION_SQLSTATE; /** * Reclassifies SQL states for known transient or mis-categorized server errors so callers can @@ -25,9 +26,6 @@ * regress the classifier. */ public final class SqlStateClassifier { - - private static final String SYNTAX_OR_ACCESS_VIOLATION_SQLSTATE = "42000"; - private SqlStateClassifier() {} /** diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/CrossReferenceKeysDatabricksResultSetAdapter.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/CrossReferenceKeysDatabricksResultSetAdapter.java index e47e7f40a4..70d4b9cf84 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/CrossReferenceKeysDatabricksResultSetAdapter.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/CrossReferenceKeysDatabricksResultSetAdapter.java @@ -40,23 +40,22 @@ public boolean includeRow(ResultSet resultSet, List columns) throw final ResultColumn parentNamespaceColumn = mapColumn(PKTABLE_SCHEM); final ResultColumn parentTableNameColumn = mapColumn(PKTABLE_NAME); - boolean isParentCatalogMatching = - resultSet - .getString(parentCatalogNameColumn.getResultSetColumnName()) - .equalsIgnoreCase(targetParentCatalogName); - boolean isParentNamespaceMatching = - resultSet - .getString(parentNamespaceColumn.getResultSetColumnName()) - .equalsIgnoreCase(targetParentNamespaceName); - boolean isParentTableMatching = - resultSet - .getString(parentTableNameColumn.getResultSetColumnName()) - .equalsIgnoreCase(targetParentTableName); - - if (!isParentTableMatching || !isParentCatalogMatching || !isParentNamespaceMatching) { + if (!matchesParent( + resultSet.getString(parentCatalogNameColumn.getResultSetColumnName()), + resultSet.getString(parentNamespaceColumn.getResultSetColumnName()), + resultSet.getString(parentTableNameColumn.getResultSetColumnName()))) { return false; } return super.includeRow(resultSet, columns); } + + boolean matchesParent(String catalog, String schema, String table) { + return catalog != null + && schema != null + && table != null + && catalog.equalsIgnoreCase(targetParentCatalogName) + && schema.equalsIgnoreCase(targetParentNamespaceName) + && table.equalsIgnoreCase(targetParentTableName); + } } 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 187bb3d75e..a37941b79f 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 @@ -502,6 +502,9 @@ public boolean shouldAllowCatalogAccess( public DatabricksResultSet getFunctionsResult(DatabricksResultSet resultSet, String catalog) throws SQLException { + if (resultSet.isThriftNativeMetadataResult()) { + return getFunctionsResult(catalog, copyThriftNativeMetadataRows(resultSet)); + } List> rows = getRowsForFunctions(resultSet, FUNCTION_COLUMNS, catalog); return buildResultSet( FUNCTION_COLUMNS, @@ -550,6 +553,9 @@ public DatabricksResultSet getProcedureColumnsResult(List> rows) { } public DatabricksResultSet getColumnsResult(DatabricksResultSet resultSet) throws SQLException { + if (resultSet.isThriftNativeMetadataResult()) { + return getColumnsResult(copyThriftNativeMetadataRows(resultSet)); + } List> rows = getRows(resultSet, COLUMN_COLUMNS, defaultAdapter); return buildResultSet( COLUMN_COLUMNS, @@ -560,6 +566,9 @@ public DatabricksResultSet getColumnsResult(DatabricksResultSet resultSet) throw } public DatabricksResultSet getCatalogsResult(DatabricksResultSet resultSet) throws SQLException { + if (resultSet.isThriftNativeMetadataResult()) { + return getCatalogsResult(copyThriftNativeMetadataRows(resultSet)); + } List> rows = getRows(resultSet, CATALOG_COLUMNS, defaultAdapter); return buildResultSet( CATALOG_COLUMNS, @@ -571,6 +580,9 @@ public DatabricksResultSet getCatalogsResult(DatabricksResultSet resultSet) thro public DatabricksResultSet getSchemasResult(DatabricksResultSet resultSet, String catalog) throws SQLException { + if (resultSet.isThriftNativeMetadataResult()) { + return getSchemasResult(copyThriftNativeMetadataRows(resultSet)); + } List> rows = getRowsForSchemas( resultSet, SCHEMA_COLUMNS, catalog, new SchemasDatabricksResultSetAdapter()); @@ -582,8 +594,11 @@ public DatabricksResultSet getSchemasResult(DatabricksResultSet resultSet, Strin CommandName.LIST_SCHEMAS); } - public DatabricksResultSet getTablesResult(DatabricksResultSet resultSet, String[] tableTypes) - throws SQLException { + public DatabricksResultSet getTablesResult( + DatabricksResultSet resultSet, String catalog, String[] tableTypes) throws SQLException { + if (resultSet.isThriftNativeMetadataResult()) { + return getTablesResult(catalog, tableTypes, copyThriftNativeMetadataRows(resultSet)); + } List allowedTableTypes = List.of(tableTypes); List> rows = getRows(resultSet, TABLE_COLUMNS, defaultAdapter).stream() @@ -622,6 +637,9 @@ public DatabricksResultSet getTableTypesResult() { public DatabricksResultSet getPrimaryKeysResult(DatabricksResultSet resultSet) throws SQLException { + if (resultSet.isThriftNativeMetadataResult()) { + return getPrimaryKeysResult(copyThriftNativeMetadataRows(resultSet)); + } List> rows = getRows(resultSet, PRIMARY_KEYS_COLUMNS, defaultAdapter); return buildResultSet( PRIMARY_KEYS_COLUMNS, @@ -633,6 +651,9 @@ public DatabricksResultSet getPrimaryKeysResult(DatabricksResultSet resultSet) public DatabricksResultSet getImportedKeysResult(DatabricksResultSet resultSet) throws SQLException { + if (resultSet.isThriftNativeMetadataResult()) { + return getImportedKeys(copyThriftNativeMetadataRows(resultSet)); + } List> rows = getRows(resultSet, IMPORTED_KEYS_COLUMNS, importedKeysAdapter); return buildResultSet( IMPORTED_KEYS_COLUMNS, @@ -651,6 +672,20 @@ public DatabricksResultSet getCrossReferenceKeysResult( final CrossReferenceKeysDatabricksResultSetAdapter crossReferenceKeysResultSetAdapter = new CrossReferenceKeysDatabricksResultSetAdapter( targetParentCatalogName, targetParentNamespaceName, targetParentTableName); + // Cross-reference SQL narrows only the foreign side, so parent filtering remains necessary. + if (resultSet.isThriftNativeMetadataResult()) { + List> rows = copyThriftNativeMetadataRows(resultSet); + int parentCatalogIndex = CROSS_REFERENCE_COLUMNS.indexOf(PKTABLE_CAT); + int parentSchemaIndex = CROSS_REFERENCE_COLUMNS.indexOf(PKTABLE_SCHEM); + int parentTableIndex = CROSS_REFERENCE_COLUMNS.indexOf(PKTABLE_NAME); + rows.removeIf( + row -> + !crossReferenceKeysResultSetAdapter.matchesParent( + (String) row.get(parentCatalogIndex), + (String) row.get(parentSchemaIndex), + (String) row.get(parentTableIndex))); + return getCrossRefsResult(rows); + } List> rows = getRows(resultSet, CROSS_REFERENCE_COLUMNS, crossReferenceKeysResultSetAdapter); @@ -662,6 +697,21 @@ public DatabricksResultSet getCrossReferenceKeysResult( CommandName.GET_CROSS_REFERENCE); } + /** Copies native rows for Thrift normalization and JDBC metadata, not just column ordering. */ + private List> copyThriftNativeMetadataRows(DatabricksResultSet resultSet) + throws SQLException { + List> rows = new ArrayList<>(); + int columnCount = resultSet.getMetaData().getColumnCount(); + while (resultSet.next()) { + List row = new ArrayList<>(columnCount); + for (int columnIndex = 1; columnIndex <= columnCount; columnIndex++) { + row.add(resultSet.getObject(columnIndex)); + } + rows.add(row); + } + return rows; + } + private boolean isTextType(String typeVal) { return (typeVal.contains(TEXT_TYPE) || typeVal.contains(CHAR_TYPE) @@ -1546,7 +1596,7 @@ public DatabricksResultSet getTablesResult( List> updatedRows = new ArrayList<>(); for (List row : rows) { // If the catalog is not null and the catalog does not match, skip the row - if (catalog != null && !row.get(0).toString().equals(catalog)) { + if (catalog != null && !catalog.equals(row.get(0))) { continue; } diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataQueryClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataQueryClient.java index 695faf67d6..0dbbf00024 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataQueryClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataQueryClient.java @@ -1,5 +1,7 @@ package com.databricks.jdbc.dbclient.impl.sqlexec; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.OPERATION_ERROR_SQLSTATE; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.SYNTAX_OR_ACCESS_VIOLATION_SQLSTATE; import static com.databricks.jdbc.common.MetadataResultConstants.*; import static com.databricks.jdbc.dbclient.impl.common.CommandConstants.METADATA_STATEMENT_ID; @@ -69,10 +71,7 @@ public DatabricksResultSet listCatalogs(IDatabricksSession session) throws SQLEx "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); - return metadataResultSetBuilder.getCatalogsResult( - getResultSet(SQL, session, MetadataOperationType.GET_CATALOGS)); + return metadataResultSetBuilder.getCatalogsResult(List.of(List.of(currentCatalog))); } CommandBuilder commandBuilder = new CommandBuilder(session); @@ -145,7 +144,12 @@ public DatabricksResultSet listTables( if (tableTypes != null && tableTypes.length == 0) { return metadataResultSetBuilder.getTablesResult(catalog, tableTypes, new ArrayList<>()); } + // Runtime does not reliably enforce empty or exact type filters. SHOW uses supported defaults; + // native results are post-filtered with the original JDBC types below. String[] validatedTableTypes = tableTypes != null ? tableTypes : DEFAULT_TABLE_TYPES; + // Runtime treats catalog as a pattern and can return temporary views outside it. Preserve the + // original JDBC catalog for exact native-result filtering before resolving catalog below. + String requestedCatalog = catalog; // Only fetch currentCatalog if multiple catalog support is disabled String currentCatalog = isMultipleCatalogSupportDisabled() ? session.getCurrentCatalog() : null; @@ -163,9 +167,15 @@ public DatabricksResultSet listTables( LOGGER.debug("SQL command to fetch tables: {}", SQL); LOGGER.debug(String.format("SQL command to fetch tables: {%s}", SQL)); try { + DatabricksResultSet resultSet = getResultSet(SQL, session, MetadataOperationType.GET_TABLES); + String[] resultTableTypes = + resultSet.isThriftNativeMetadataResult() ? tableTypes : validatedTableTypes; return metadataResultSetBuilder.getTablesResult( - getResultSet(SQL, session, MetadataOperationType.GET_TABLES), validatedTableTypes); + resultSet, requestedCatalog, resultTableTypes); } catch (SQLException e) { + if (isThriftNativeMetadataRequested()) { + throw e; + } if ((PARSE_SYNTAX_ERROR_SQL_STATE.equals(e.getSQLState()) && catalog == null) || isObjectNotFoundException(e) || isEmptyPatternError(schemaNamePattern, tableNamePattern)) { @@ -214,6 +224,9 @@ public DatabricksResultSet listColumns( return metadataResultSetBuilder.getColumnsResult( getResultSet(SQL, session, MetadataOperationType.GET_COLUMNS)); } catch (SQLException e) { + if (isThriftNativeMetadataRequested()) { + throw e; + } if (catalog == null && PARSE_SYNTAX_ERROR_SQL_STATE.equals(e.getSQLState())) { // Fallback for older DBR versions that don't support "SHOW COLUMNS IN ALL CATALOGS": // enumerate the catalogs and issue a per-catalog SHOW COLUMNS. @@ -236,6 +249,11 @@ public DatabricksResultSet listFunctions( String schemaNamePattern, String functionNamePattern) throws SQLException { + // Native GetFunctions ignores catalog filtering and returns an empty FUNCTION_CAT. Preserve + // the original JDBC argument (including null) for Thrift's native-only column override; the + // SHOW path is catalog-aware and uses the resolved catalog below. + String requestedCatalog = catalog; + // Only fetch currentCatalog if multiple catalog support is disabled String currentCatalog = isMultipleCatalogSupportDisabled() ? session.getCurrentCatalog() : null; if (!metadataResultSetBuilder.shouldAllowCatalogAccess(catalog, currentCatalog, session)) { @@ -266,8 +284,10 @@ public DatabricksResultSet listFunctions( String SQL = commandBuilder.getSQLString(CommandName.LIST_FUNCTIONS); LOGGER.debug("SQL command to fetch functions: {}", SQL); try { + DatabricksResultSet resultSet = + getResultSet(SQL, session, MetadataOperationType.GET_FUNCTIONS); return metadataResultSetBuilder.getFunctionsResult( - getResultSet(SQL, session, MetadataOperationType.GET_FUNCTIONS), catalog); + resultSet, resultSet.isThriftNativeMetadataResult() ? requestedCatalog : catalog); } catch (SQLException e) { if (isObjectNotFoundException(e)) { LOGGER.debug("Object not found for getFunctions, returning empty result set."); @@ -401,7 +421,9 @@ public DatabricksResultSet listExportedKeys( if (table == null) { LOGGER.debug("listExportedKeys: table is null, throwing"); throw new DatabricksSQLException( - "Invalid argument: tableName may not be null", DatabricksDriverErrorCode.INVALID_STATE); + "Invalid argument: tableName may not be null", + SYNTAX_OR_ACCESS_VIOLATION_SQLSTATE, + DatabricksDriverErrorCode.EXECUTE_STATEMENT_FAILED); } // Only fetch currentCatalog if multiple catalog support is disabled @@ -488,6 +510,11 @@ private boolean isMultipleCatalogSupportDisabled() { && !queryExecutionClient.getConnectionContext().getEnableMultipleCatalogSupport(); } + private boolean isThriftNativeMetadataRequested() { + IDatabricksConnectionContext connectionContext = queryExecutionClient.getConnectionContext(); + return connectionContext != null && connectionContext.isThriftNativeMetadataEnabled(); + } + /** * Returns true if any of the provided patterns is an empty string. Empty string patterns generate * invalid LIKE '' clauses that cause server errors. Per JDBC spec, empty string means "without a @@ -531,11 +558,19 @@ private String autoFillCatalog(String catalog, String currentCatalog) { */ private String[] resolveKeyBasedParams( String catalog, String schema, String table, IDatabricksSession session) throws SQLException { - if (table == null || table.isEmpty()) { - LOGGER.debug("resolveKeyBasedParams: table is null or empty, throwing"); + if (table == null) { + LOGGER.debug("resolveKeyBasedParams: table is null, throwing"); + throw new DatabricksSQLException( + "Invalid argument: tableName may not be null or empty", + SYNTAX_OR_ACCESS_VIOLATION_SQLSTATE, + DatabricksDriverErrorCode.EXECUTE_STATEMENT_FAILED); + } + if (table.isEmpty()) { + LOGGER.debug("resolveKeyBasedParams: table is empty, throwing"); throw new DatabricksSQLException( "Invalid argument: tableName may not be null or empty", - DatabricksDriverErrorCode.INVALID_STATE); + OPERATION_ERROR_SQLSTATE, + DatabricksDriverErrorCode.EXECUTE_STATEMENT_FAILED); } if (catalog == null) { @@ -549,7 +584,8 @@ private String[] resolveKeyBasedParams( "resolveKeyBasedParams: schema is null with explicit catalog '{}', throwing", catalog); throw new DatabricksSQLException( "Invalid argument: schema may not be null when catalog is specified", - DatabricksDriverErrorCode.INVALID_STATE); + OPERATION_ERROR_SQLSTATE, + DatabricksDriverErrorCode.EXECUTE_STATEMENT_FAILED); } // Safety net: getCurrentCatalogAndSchema() returned null values @@ -560,7 +596,8 @@ private String[] resolveKeyBasedParams( schema); throw new DatabricksSQLException( "Invalid argument: could not resolve catalog or schema", - DatabricksDriverErrorCode.INVALID_STATE); + OPERATION_ERROR_SQLSTATE, + DatabricksDriverErrorCode.EXECUTE_STATEMENT_FAILED); } return new String[] {catalog, schema, table}; 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 a62a4c6318..0d57e0dec0 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 @@ -66,6 +66,8 @@ public class DatabricksSdkClient implements IDatabricksClient { private static final String ASYNC_TIMEOUT_VALUE = "0s"; private static final String HEADER_METADATA_OPERATION_TYPE = "X-Databricks-Metadata-Operation-Type"; + private static final String HEADER_REQUIRE_THRIFT_NATIVE_METADATA = + "X-Databricks-Require-Thrift-Native-Metadata"; private final IDatabricksConnectionContext connectionContext; private final ClientConfigurator clientConfigurator; @@ -220,6 +222,10 @@ public DatabricksResultSet executeStatement( if (metadataOperationType != null) { additionalHeaders.put( HEADER_METADATA_OPERATION_TYPE, metadataOperationType.getHeaderValue()); + if (connectionContext.isThriftNativeMetadataEnabled() + && metadataOperationType.isThriftNativeSupported()) { + additionalHeaders.put(HEADER_REQUIRE_THRIFT_NATIVE_METADATA, "true"); + } } req.withHeaders(getHeaders("executeStatement", statementType, false, additionalHeaders)); response = apiClient.execute(req, ExecuteStatementResponse.class); 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 18c126b2bf..0518c61be9 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 @@ -127,7 +127,9 @@ TBase getThriftResponse(TBase request) throws DatabricksSQLException { request, e.getMessage()); LOGGER.error(e, errorMessage); if (e instanceof SQLException) { - throw new DatabricksSQLException(errorMessage, e, ((SQLException) e).getSQLState()); + SQLException sqlException = (SQLException) e; + throw new DatabricksSQLException( + errorMessage, sqlException.getSQLState(), sqlException.getErrorCode(), sqlException); } else { throw new DatabricksSQLException(errorMessage, e, DatabricksDriverErrorCode.INVALID_STATE); } diff --git a/src/main/java/com/databricks/jdbc/model/core/ResultManifest.java b/src/main/java/com/databricks/jdbc/model/core/ResultManifest.java index 4188c41adc..4e473cab0e 100644 --- a/src/main/java/com/databricks/jdbc/model/core/ResultManifest.java +++ b/src/main/java/com/databricks/jdbc/model/core/ResultManifest.java @@ -41,6 +41,9 @@ public class ResultManifest { @JsonProperty("is_volume_operation") private Boolean isVolumeOperation; + @JsonProperty("is_native_metadata_result") + private Boolean isNativeMetadataResult; + public ResultManifest() {} public ResultManifest setChunks(Collection chunks) { @@ -124,6 +127,15 @@ public Boolean getIsVolumeOperation() { return this.isVolumeOperation; } + public ResultManifest setIsNativeMetadataResult(Boolean isNativeMetadataResult) { + this.isNativeMetadataResult = isNativeMetadataResult; + return this; + } + + public Boolean getIsNativeMetadataResult() { + return this.isNativeMetadataResult; + } + public boolean equals(Object o) { if (this == o) { return true; @@ -136,7 +148,8 @@ public boolean equals(Object o) { && Objects.equals(this.totalChunkCount, that.totalChunkCount) && Objects.equals(this.totalRowCount, that.totalRowCount) && Objects.equals(this.truncated, that.truncated) - && Objects.equals(this.isVolumeOperation, that.isVolumeOperation); + && Objects.equals(this.isVolumeOperation, that.isVolumeOperation) + && Objects.equals(this.isNativeMetadataResult, that.isNativeMetadataResult); } else { return false; } @@ -151,7 +164,8 @@ public int hashCode() { this.totalChunkCount, this.totalRowCount, this.truncated, - this.isVolumeOperation); + this.isVolumeOperation, + this.isNativeMetadataResult); } public String toString() { @@ -164,6 +178,7 @@ public String toString() { .add("totalRowCount", this.totalRowCount) .add("truncated", this.truncated) .add("isVolumeOperation", this.isVolumeOperation) + .add("isNativeMetadataResult", this.isNativeMetadataResult) .toString(); } } diff --git a/src/test/java/com/databricks/jdbc/common/MetadataOperationTypeTest.java b/src/test/java/com/databricks/jdbc/common/MetadataOperationTypeTest.java index 613f1f70af..c93e84b74a 100644 --- a/src/test/java/com/databricks/jdbc/common/MetadataOperationTypeTest.java +++ b/src/test/java/com/databricks/jdbc/common/MetadataOperationTypeTest.java @@ -1,7 +1,6 @@ package com.databricks.jdbc.common; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -12,79 +11,25 @@ public class MetadataOperationTypeTest { @Test void testAllEnumValuesExist() { - // Verify all expected enum values exist assertEquals(9, MetadataOperationType.values().length); - assertNotNull(MetadataOperationType.GET_CATALOGS); - assertNotNull(MetadataOperationType.GET_SCHEMAS); - assertNotNull(MetadataOperationType.GET_TABLES); - assertNotNull(MetadataOperationType.GET_COLUMNS); - assertNotNull(MetadataOperationType.GET_FUNCTIONS); - assertNotNull(MetadataOperationType.GET_PRIMARY_KEYS); - assertNotNull(MetadataOperationType.GET_CROSS_REFERENCE); - assertNotNull(MetadataOperationType.GET_PROCEDURES); - assertNotNull(MetadataOperationType.GET_PROCEDURE_COLUMNS); } @ParameterizedTest @CsvSource({ - "GET_CATALOGS, GetCatalogs", - "GET_SCHEMAS, GetSchemas", - "GET_TABLES, GetTables", - "GET_COLUMNS, GetColumns", - "GET_FUNCTIONS, GetFunctions", - "GET_PRIMARY_KEYS, GetPrimaryKeys", - "GET_CROSS_REFERENCE, GetCrossReference", - "GET_PROCEDURES, GetProcedures", - "GET_PROCEDURE_COLUMNS, GetProcedureColumns" + "GET_CATALOGS, GetCatalogs, true", + "GET_SCHEMAS, GetSchemas, true", + "GET_TABLES, GetTables, true", + "GET_COLUMNS, GetColumns, true", + "GET_FUNCTIONS, GetFunctions, true", + "GET_PRIMARY_KEYS, GetPrimaryKeys, true", + "GET_CROSS_REFERENCE, GetCrossReference, true", + "GET_PROCEDURES, GetProcedures, false", + "GET_PROCEDURE_COLUMNS, GetProcedureColumns, false" }) - void testHeaderValues(String enumName, String expectedHeaderValue) { + void testHeaderAndThriftNativeSupport( + String enumName, String expectedHeaderValue, boolean thriftNativeSupported) { MetadataOperationType operationType = MetadataOperationType.valueOf(enumName); assertEquals(expectedHeaderValue, operationType.getHeaderValue()); - } - - @Test - void testGetCatalogsHeaderValue() { - assertEquals("GetCatalogs", MetadataOperationType.GET_CATALOGS.getHeaderValue()); - } - - @Test - void testGetSchemasHeaderValue() { - assertEquals("GetSchemas", MetadataOperationType.GET_SCHEMAS.getHeaderValue()); - } - - @Test - void testGetTablesHeaderValue() { - assertEquals("GetTables", MetadataOperationType.GET_TABLES.getHeaderValue()); - } - - @Test - void testGetColumnsHeaderValue() { - assertEquals("GetColumns", MetadataOperationType.GET_COLUMNS.getHeaderValue()); - } - - @Test - void testGetFunctionsHeaderValue() { - assertEquals("GetFunctions", MetadataOperationType.GET_FUNCTIONS.getHeaderValue()); - } - - @Test - void testGetPrimaryKeysHeaderValue() { - assertEquals("GetPrimaryKeys", MetadataOperationType.GET_PRIMARY_KEYS.getHeaderValue()); - } - - @Test - void testGetCrossReferenceHeaderValue() { - assertEquals("GetCrossReference", MetadataOperationType.GET_CROSS_REFERENCE.getHeaderValue()); - } - - @Test - void testGetProceduresHeaderValue() { - assertEquals("GetProcedures", MetadataOperationType.GET_PROCEDURES.getHeaderValue()); - } - - @Test - void testGetProcedureColumnsHeaderValue() { - assertEquals( - "GetProcedureColumns", MetadataOperationType.GET_PROCEDURE_COLUMNS.getHeaderValue()); + assertEquals(thriftNativeSupported, operationType.isThriftNativeSupported()); } } 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 db849a10cd..bd42fef538 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 @@ -10,11 +10,13 @@ import com.databricks.jdbc.common.util.DatabricksThreadContextHolder; import com.databricks.jdbc.model.core.ResultColumn; import java.sql.ResultSet; +import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Types; import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Stream; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -42,6 +44,255 @@ void tearDown() { DatabricksThreadContextHolder.clearAllContext(); } + @Test + void testThriftNativeFormattingMatchesRawThriftBuilder() throws SQLException { + assertNativeFormattingMatchesThrift( + resultSet -> metadataResultSetBuilder.getFunctionsResult(resultSet, "catalog"), + rows -> metadataResultSetBuilder.getFunctionsResult("catalog", rows), + FUNCTION_COLUMNS); + assertNativeFormattingMatchesThrift( + metadataResultSetBuilder::getCatalogsResult, + metadataResultSetBuilder::getCatalogsResult, + CATALOG_COLUMNS); + assertNativeFormattingMatchesThrift( + resultSet -> metadataResultSetBuilder.getSchemasResult(resultSet, "catalog"), + metadataResultSetBuilder::getSchemasResult, + SCHEMA_COLUMNS); + assertNativeFormattingMatchesThrift( + resultSet -> + metadataResultSetBuilder.getTablesResult(resultSet, null, new String[] {"TABLE"}), + rows -> metadataResultSetBuilder.getTablesResult(null, new String[] {"TABLE"}, rows), + TABLE_COLUMNS); + assertNativeFormattingMatchesThrift( + metadataResultSetBuilder::getPrimaryKeysResult, + metadataResultSetBuilder::getPrimaryKeysResult, + PRIMARY_KEYS_COLUMNS); + assertNativeFormattingMatchesThrift( + metadataResultSetBuilder::getImportedKeysResult, + metadataResultSetBuilder::getImportedKeys, + IMPORTED_KEYS_COLUMNS); + } + + @Test + void testThriftNativeGetColumnsFormattingMatchesRawThriftBuilder() throws SQLException { + List nativeRow = + Arrays.asList( + "catalog", + "schema", + "table", + "column", + Types.INTEGER, + "INT", + 10, + null, + 0, + 10, + 1, + null, + null, + Types.INTEGER, + null, + null, + 0, + "YES", + null, + null, + null, + null, + "YES"); + DatabricksResultSet actual = + metadataResultSetBuilder.getColumnsResult( + nativeMetadataResult(List.of(new ArrayList<>(nativeRow)))); + DatabricksResultSet expected = + metadataResultSetBuilder.getColumnsResult(List.of(new ArrayList<>(nativeRow))); + + assertResultSetsEqual(expected, actual); + } + + @Test + void testLegacyMetadataResultWithNativeColumnCountIsTransformed() throws SQLException { + DatabricksResultSet resultSet = + mockMetadataResultSetWithColumnNames( + List.of( + "functionName", + "namespace", + "catalogName", + "remarks", + "functionType", + "specificName")); + when(resultSet.next()).thenReturn(false); + + assertNotSame(resultSet, metadataResultSetBuilder.getFunctionsResult(resultSet, "catalog")); + verify(resultSet).next(); + } + + private void assertNativeFormattingMatchesThrift( + MetadataResultCall nativeCall, RawRowsResultCall thriftCall, List columns) + throws SQLException { + List row = metadataRow(columns); + DatabricksResultSet actual = + nativeCall.execute(nativeMetadataResult(List.of(new ArrayList<>(row)))); + DatabricksResultSet expected = thriftCall.execute(List.of(new ArrayList<>(row))); + + assertResultSetsEqual(expected, actual); + } + + private DatabricksResultSet nativeMetadataResult(List> rows) throws SQLException { + DatabricksResultSet resultSet = mock(DatabricksResultSet.class); + ResultSetMetaData metadata = mock(ResultSetMetaData.class); + AtomicInteger rowIndex = new AtomicInteger(-1); + when(resultSet.isThriftNativeMetadataResult()).thenReturn(true); + when(resultSet.getMetaData()).thenReturn(metadata); + when(metadata.getColumnCount()).thenReturn(rows.get(0).size()); + when(resultSet.next()).thenAnswer(ignored -> rowIndex.incrementAndGet() < rows.size()); + when(resultSet.getObject(anyInt())) + .thenAnswer( + invocation -> + rows.get(rowIndex.get()).get(invocation.getArgument(0, Integer.class) - 1)); + return resultSet; + } + + private void assertResultSetsEqual(DatabricksResultSet expected, DatabricksResultSet actual) + throws SQLException { + ResultSetMetaData expectedMetadata = expected.getMetaData(); + ResultSetMetaData actualMetadata = actual.getMetaData(); + int columnCount = expectedMetadata.getColumnCount(); + assertEquals(columnCount, actualMetadata.getColumnCount()); + for (int columnIndex = 1; columnIndex <= columnCount; columnIndex++) { + assertEquals( + expectedMetadata.getColumnName(columnIndex), actualMetadata.getColumnName(columnIndex)); + assertEquals( + expectedMetadata.getColumnType(columnIndex), actualMetadata.getColumnType(columnIndex)); + assertEquals( + expectedMetadata.getColumnTypeName(columnIndex), + actualMetadata.getColumnTypeName(columnIndex)); + assertEquals( + expectedMetadata.getPrecision(columnIndex), actualMetadata.getPrecision(columnIndex)); + assertEquals(expectedMetadata.getScale(columnIndex), actualMetadata.getScale(columnIndex)); + assertEquals( + expectedMetadata.isNullable(columnIndex), actualMetadata.isNullable(columnIndex)); + } + while (expected.next()) { + assertTrue(actual.next()); + for (int columnIndex = 1; columnIndex <= columnCount; columnIndex++) { + assertEquals(expected.getObject(columnIndex), actual.getObject(columnIndex)); + } + } + assertFalse(actual.next()); + } + + private List metadataRow(List columns) { + List row = new ArrayList<>(columns.size()); + for (int columnIndex = 1; columnIndex <= columns.size(); columnIndex++) { + row.add(metadataValue(columns.get(columnIndex - 1), columnIndex)); + } + return row; + } + + private Object metadataValue(ResultColumn column, int columnIndex) { + if (TABLE_TYPE_COLUMN.getColumnName().equals(column.getColumnName())) { + return "TABLE"; + } + switch (column.getColumnTypeInt()) { + case Types.SMALLINT: + return (short) columnIndex; + case Types.INTEGER: + return columnIndex; + case Types.BIT: + return true; + default: + return "value-" + columnIndex; + } + } + + @Test + void testThriftNativeTablesStillApplyTableTypeFilter() throws SQLException { + DatabricksResultSet result = + metadataResultSetBuilder.getTablesResult( + nativeMetadataResult(List.of(metadataRow(TABLE_COLUMNS))), + null, + new String[] {"NONEXISTENT_TYPE"}); + + assertFalse(result.next()); + } + + @Test + void testThriftNativeTablesStillApplyExactCatalogFilter() throws SQLException { + List mismatchedCatalogRow = metadataRow(TABLE_COLUMNS); + mismatchedCatalogRow.set(0, "comparator_tests"); + List nullCatalogRow = new ArrayList<>(mismatchedCatalogRow); + nullCatalogRow.set(0, null); + + DatabricksResultSet result = + metadataResultSetBuilder.getTablesResult( + nativeMetadataResult(List.of(mismatchedCatalogRow, nullCatalogRow)), + "COMPARATOR-TESTS", + new String[] {"TABLE"}); + + assertFalse(result.next()); + } + + @Test + void testThriftNativeCrossReferenceStillFiltersParentTable() throws SQLException { + List matchingRow = metadataRow(CROSS_REFERENCE_COLUMNS); + matchingRow.set(0, "parent-catalog"); + matchingRow.set(1, "parent-schema"); + matchingRow.set(2, "parent-table"); + List nonMatchingRow = new ArrayList<>(matchingRow); + nonMatchingRow.set(2, "other-parent-table"); + + DatabricksResultSet result = + metadataResultSetBuilder.getCrossReferenceKeysResult( + nativeMetadataResult(List.of(matchingRow, nonMatchingRow)), + "PARENT-CATALOG", + "PARENT-SCHEMA", + "PARENT-TABLE"); + + assertTrue(result.next()); + assertEquals("parent-catalog", result.getString("PKTABLE_CAT")); + assertEquals("parent-schema", result.getString("PKTABLE_SCHEM")); + assertEquals("parent-table", result.getString("PKTABLE_NAME")); + assertFalse(result.next()); + } + + @Test + void testThriftNativeFunctionsUseThriftPostProcessing() throws SQLException { + List row = + new ArrayList<>( + Arrays.asList( + "server-catalog", "schema", "function", null, (short) 1, "specific-name")); + DatabricksResultSet actual = + metadataResultSetBuilder.getFunctionsResult( + nativeMetadataResult(List.of(new ArrayList<>(row))), "requested-catalog"); + DatabricksResultSet expected = + metadataResultSetBuilder.getFunctionsResult( + "requested-catalog", List.of(new ArrayList<>(row))); + + assertResultSetsEqual(expected, actual); + } + + private DatabricksResultSet mockMetadataResultSetWithColumnNames(List columnNames) + throws SQLException { + DatabricksResultSet resultSet = mock(DatabricksResultSet.class); + ResultSetMetaData metadata = mock(ResultSetMetaData.class); + when(resultSet.getMetaData()).thenReturn(metadata); + when(metadata.getColumnCount()).thenReturn(columnNames.size()); + for (int i = 0; i < columnNames.size(); i++) { + when(metadata.getColumnName(i + 1)).thenReturn(columnNames.get(i)); + } + return resultSet; + } + + @FunctionalInterface + private interface MetadataResultCall { + DatabricksResultSet execute(DatabricksResultSet resultSet) throws SQLException; + } + + @FunctionalInterface + private interface RawRowsResultCall { + DatabricksResultSet execute(List> rows) throws SQLException; + } + @Test void testGetCode() { assert metadataResultSetBuilder.getCode("STRING") == 12; @@ -529,7 +780,7 @@ void testGetTablesResultSortingInSeaMode() throws SQLException { // Call SEA mode method with both TABLE and VIEW types String[] tableTypes = new String[] {"TABLE", "VIEW"}; - ResultSet resultSet = metadataResultSetBuilder.getTablesResult(mockResultSet, tableTypes); + ResultSet resultSet = metadataResultSetBuilder.getTablesResult(mockResultSet, null, tableTypes); // Verify sorting: TABLE_TYPE first, then TABLE_CAT, TABLE_SCHEM, TABLE_NAME // Expected order after sorting: diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataQueryClientTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataQueryClientTest.java index 6d7bd61596..91300f31e9 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataQueryClientTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataQueryClientTest.java @@ -1,9 +1,12 @@ package com.databricks.jdbc.dbclient.impl.sqlexec; import static com.databricks.jdbc.TestConstants.*; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.OPERATION_ERROR_SQLSTATE; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.SYNTAX_OR_ACCESS_VIOLATION_SQLSTATE; import static com.databricks.jdbc.common.MetadataResultConstants.*; import static com.databricks.jdbc.dbclient.impl.common.CommandConstants.*; import static com.databricks.jdbc.dbclient.impl.common.ImportedKeysDatabricksResultSetAdapter.*; +import static com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode.EXECUTE_STATEMENT_FAILED; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; @@ -219,6 +222,91 @@ void listTablesReturnsEmptyWhenCatalogAccessDenied() throws SQLException { .executeStatement(anyString(), any(), any(), any(), any(), any(), any()); } + @Test + void nativeListTablesAppliesExactRequestedCatalogFilter() throws SQLException { + when(session.getComputeResource()).thenReturn(mockedComputeResource); + when(mockClient.executeStatement( + eq("SHOW TABLES IN CATALOG `COMPARATOR-TESTS`"), + eq(mockedComputeResource), + any(), + eq(StatementType.METADATA), + eq(session), + any(), + eq(MetadataOperationType.GET_TABLES))) + .thenReturn(mockedResultSet); + when(mockedResultSet.isThriftNativeMetadataResult()).thenReturn(true); + when(mockedResultSet.getMetaData()).thenReturn(mockedMetaData); + when(mockedMetaData.getColumnCount()).thenReturn(TABLE_COLUMNS.size()); + when(mockedResultSet.next()).thenReturn(true, false); + when(mockedResultSet.getObject(1)).thenReturn("comparator-tests"); + when(mockedResultSet.getObject(2)).thenReturn(TEST_SCHEMA); + when(mockedResultSet.getObject(3)).thenReturn(TEST_TABLE); + when(mockedResultSet.getObject(4)).thenReturn("TABLE"); + + DatabricksMetadataQueryClient metadataClient = new DatabricksMetadataQueryClient(mockClient); + DatabricksResultSet result = + metadataClient.listTables(session, "COMPARATOR-TESTS", null, null, null); + + assertFalse(result.next()); + } + + @Test + void nativeListTablesWithNullTypesReturnsUnrecognizedTypes() throws SQLException { + when(session.getComputeResource()).thenReturn(mockedComputeResource); + when(mockClient.executeStatement( + eq("SHOW TABLES IN CATALOG `catalog1`"), + eq(mockedComputeResource), + any(), + eq(StatementType.METADATA), + eq(session), + any(), + eq(MetadataOperationType.GET_TABLES))) + .thenReturn(mockedResultSet); + when(mockedResultSet.isThriftNativeMetadataResult()).thenReturn(true); + when(mockedResultSet.getMetaData()).thenReturn(mockedMetaData); + when(mockedMetaData.getColumnCount()).thenReturn(TABLE_COLUMNS.size()); + when(mockedResultSet.next()).thenReturn(true, false); + when(mockedResultSet.getObject(1)).thenReturn(TEST_CATALOG); + when(mockedResultSet.getObject(2)).thenReturn(TEST_SCHEMA); + when(mockedResultSet.getObject(3)).thenReturn(TEST_TABLE); + when(mockedResultSet.getObject(4)).thenReturn("FUTURE_TABLE_TYPE"); + + DatabricksMetadataQueryClient metadataClient = new DatabricksMetadataQueryClient(mockClient); + DatabricksResultSet result = metadataClient.listTables(session, TEST_CATALOG, null, null, null); + + assertTrue(result.next()); + assertEquals("FUTURE_TABLE_TYPE", result.getString("TABLE_TYPE")); + assertFalse(result.next()); + } + + @Test + void nativeListTablesPropagatesServerErrors() throws SQLException { + when(session.getComputeResource()).thenReturn(mockedComputeResource); + IDatabricksConnectionContext connectionContext = mock(IDatabricksConnectionContext.class); + when(connectionContext.getEnableMultipleCatalogSupport()).thenReturn(true); + when(connectionContext.isThriftNativeMetadataEnabled()).thenReturn(true); + when(mockClient.getConnectionContext()).thenReturn(connectionContext); + DatabricksSQLException nativeError = + new DatabricksSQLException("native metadata failure", OBJECT_NOT_FOUND_SQL_STATE); + when(mockClient.executeStatement( + eq("SHOW TABLES IN CATALOG ``"), + eq(mockedComputeResource), + any(), + eq(StatementType.METADATA), + eq(session), + any(), + eq(MetadataOperationType.GET_TABLES))) + .thenThrow(nativeError); + + DatabricksMetadataQueryClient metadataClient = new DatabricksMetadataQueryClient(mockClient); + + assertSame( + nativeError, + assertThrows( + DatabricksSQLException.class, + () -> metadataClient.listTables(session, "", null, null, null))); + } + @Test void listSchemasReturnsEmptyWhenCatalogIsEmptyString() throws SQLException { IDatabricksConnectionContext connectionContext = mock(IDatabricksConnectionContext.class); @@ -371,7 +459,7 @@ void testListColumns( eq(StatementType.METADATA), eq(session), any(), - any(MetadataOperationType.class))) + eq(MetadataOperationType.GET_COLUMNS))) .thenReturn(mockedResultSet); when(mockedResultSet.next()).thenReturn(true, false); @@ -429,6 +517,34 @@ void testListColumns( } } + @Test + void nativeListColumnsPropagatesServerErrorsForEmptyPatterns() throws SQLException { + when(session.getComputeResource()).thenReturn(mockedComputeResource); + IDatabricksConnectionContext connectionContext = mock(IDatabricksConnectionContext.class); + when(connectionContext.getEnableMultipleCatalogSupport()).thenReturn(true); + when(connectionContext.isThriftNativeMetadataEnabled()).thenReturn(true); + when(mockClient.getConnectionContext()).thenReturn(connectionContext); + DatabricksSQLException nativeError = + new DatabricksSQLException("native metadata failure", SYNTAX_OR_ACCESS_VIOLATION_SQLSTATE); + when(mockClient.executeStatement( + eq("SHOW COLUMNS IN CATALOG `catalog1` SCHEMA LIKE ''"), + eq(mockedComputeResource), + any(), + eq(StatementType.METADATA), + eq(session), + any(), + eq(MetadataOperationType.GET_COLUMNS))) + .thenThrow(nativeError); + + DatabricksMetadataQueryClient metadataClient = new DatabricksMetadataQueryClient(mockClient); + + assertSame( + nativeError, + assertThrows( + DatabricksSQLException.class, + () -> metadataClient.listColumns(session, TEST_CATALOG, "", null, null))); + } + private void stubColumnsMetaData() throws SQLException { doReturn(13).when(mockedMetaData).getColumnCount(); doReturn(COL_NAME_COLUMN.getResultSetColumnName()).when(mockedMetaData).getColumnName(1); @@ -1047,58 +1163,117 @@ void testListFunctionsWithNullCatalog() throws SQLException { } @Test - void testKeyBasedOpsThrowForNullTable() { + void nativeListFunctionsPreservesNullRequestedCatalog() throws SQLException { + when(session.getComputeResource()).thenReturn(WAREHOUSE_COMPUTE); + when(session.getCurrentCatalog()).thenReturn("current_catalog"); + IDatabricksConnectionContext mockContext = mock(IDatabricksConnectionContext.class); + when(mockContext.getEnableMultipleCatalogSupport()).thenReturn(true); + when(mockClient.getConnectionContext()).thenReturn(mockContext); + when(mockClient.executeStatement( + eq( + "SHOW FUNCTIONS IN CATALOG `current_catalog` SCHEMA LIKE 'testSchema' LIKE 'functionPattern'"), + eq(WAREHOUSE_COMPUTE), + any(), + eq(StatementType.METADATA), + eq(session), + any(), + eq(MetadataOperationType.GET_FUNCTIONS))) + .thenReturn(mockedResultSet); + when(mockedResultSet.isThriftNativeMetadataResult()).thenReturn(true); + when(mockedResultSet.getMetaData()).thenReturn(mockedMetaData); + when(mockedMetaData.getColumnCount()).thenReturn(FUNCTION_COLUMNS.size()); + when(mockedResultSet.next()).thenReturn(true, false); + when(mockedResultSet.getObject(1)).thenReturn("current_catalog"); + when(mockedResultSet.getObject(2)).thenReturn(TEST_SCHEMA); + when(mockedResultSet.getObject(3)).thenReturn("function"); + when(mockedResultSet.getObject(4)).thenReturn(null); + when(mockedResultSet.getObject(5)).thenReturn((short) 1); + when(mockedResultSet.getObject(6)).thenReturn("function"); + DatabricksMetadataQueryClient metadataClient = new DatabricksMetadataQueryClient(mockClient); + DatabricksResultSet result = + metadataClient.listFunctions(session, null, TEST_SCHEMA, TEST_FUNCTION_PATTERN); - assertThrows( - DatabricksSQLException.class, - () -> metadataClient.listPrimaryKeys(session, TEST_CATALOG, TEST_SCHEMA, null), - "listPrimaryKeys should throw for null table"); + assertTrue(result.next()); + assertNull(result.getString("FUNCTION_CAT")); + assertEquals("function", result.getString("FUNCTION_NAME")); + assertFalse(result.next()); + } - assertThrows( - DatabricksSQLException.class, - () -> metadataClient.listImportedKeys(session, TEST_CATALOG, TEST_SCHEMA, null), - "listImportedKeys should throw for null table"); + @Test + void testKeyBasedOpsThrowForNullTable() { + DatabricksMetadataQueryClient metadataClient = new DatabricksMetadataQueryClient(mockClient); + + DatabricksSQLException primaryKeysError = + assertThrows( + DatabricksSQLException.class, + () -> metadataClient.listPrimaryKeys(session, TEST_CATALOG, TEST_SCHEMA, null), + "listPrimaryKeys should throw for null table"); + assertEquals(EXECUTE_STATEMENT_FAILED.getCode(), primaryKeysError.getErrorCode()); + assertEquals(SYNTAX_OR_ACCESS_VIOLATION_SQLSTATE, primaryKeysError.getSQLState()); + + DatabricksSQLException importedKeysError = + assertThrows( + DatabricksSQLException.class, + () -> metadataClient.listImportedKeys(session, TEST_CATALOG, TEST_SCHEMA, null), + "listImportedKeys should throw for null table"); + assertEquals(EXECUTE_STATEMENT_FAILED.getCode(), importedKeysError.getErrorCode()); + assertEquals(SYNTAX_OR_ACCESS_VIOLATION_SQLSTATE, importedKeysError.getSQLState()); } @Test void testKeyBasedOpsThrowForEmptyTable() { DatabricksMetadataQueryClient metadataClient = new DatabricksMetadataQueryClient(mockClient); - assertThrows( - DatabricksSQLException.class, - () -> metadataClient.listPrimaryKeys(session, TEST_CATALOG, TEST_SCHEMA, ""), - "listPrimaryKeys should throw for empty table"); - - assertThrows( - DatabricksSQLException.class, - () -> metadataClient.listImportedKeys(session, TEST_CATALOG, TEST_SCHEMA, ""), - "listImportedKeys should throw for empty table"); + DatabricksSQLException primaryKeysError = + assertThrows( + DatabricksSQLException.class, + () -> metadataClient.listPrimaryKeys(session, TEST_CATALOG, TEST_SCHEMA, ""), + "listPrimaryKeys should throw for empty table"); + assertEquals(EXECUTE_STATEMENT_FAILED.getCode(), primaryKeysError.getErrorCode()); + assertEquals(OPERATION_ERROR_SQLSTATE, primaryKeysError.getSQLState()); + + DatabricksSQLException importedKeysError = + assertThrows( + DatabricksSQLException.class, + () -> metadataClient.listImportedKeys(session, TEST_CATALOG, TEST_SCHEMA, ""), + "listImportedKeys should throw for empty table"); + assertEquals(EXECUTE_STATEMENT_FAILED.getCode(), importedKeysError.getErrorCode()); + assertEquals(OPERATION_ERROR_SQLSTATE, importedKeysError.getSQLState()); } @Test void testKeyBasedOpsThrowForNullSchemaWithExplicitCatalog() { DatabricksMetadataQueryClient metadataClient = new DatabricksMetadataQueryClient(mockClient); - assertThrows( - DatabricksSQLException.class, - () -> metadataClient.listPrimaryKeys(session, "any_catalog", null, TEST_TABLE), - "listPrimaryKeys should throw for null schema with explicit catalog"); - - assertThrows( - DatabricksSQLException.class, - () -> metadataClient.listImportedKeys(session, "any_catalog", null, TEST_TABLE), - "listImportedKeys should throw for null schema with explicit catalog"); + DatabricksSQLException primaryKeysError = + assertThrows( + DatabricksSQLException.class, + () -> metadataClient.listPrimaryKeys(session, "any_catalog", null, TEST_TABLE), + "listPrimaryKeys should throw for null schema with explicit catalog"); + assertEquals(EXECUTE_STATEMENT_FAILED.getCode(), primaryKeysError.getErrorCode()); + assertEquals(OPERATION_ERROR_SQLSTATE, primaryKeysError.getSQLState()); + + DatabricksSQLException importedKeysError = + assertThrows( + DatabricksSQLException.class, + () -> metadataClient.listImportedKeys(session, "any_catalog", null, TEST_TABLE), + "listImportedKeys should throw for null schema with explicit catalog"); + assertEquals(EXECUTE_STATEMENT_FAILED.getCode(), importedKeysError.getErrorCode()); + assertEquals(OPERATION_ERROR_SQLSTATE, importedKeysError.getSQLState()); } @Test void testExportedKeysThrowsForNullTable() { DatabricksMetadataQueryClient metadataClient = new DatabricksMetadataQueryClient(mockClient); - assertThrows( - DatabricksSQLException.class, - () -> metadataClient.listExportedKeys(session, TEST_CATALOG, TEST_SCHEMA, null), - "listExportedKeys should throw for null table"); + DatabricksSQLException exportedKeysError = + assertThrows( + DatabricksSQLException.class, + () -> metadataClient.listExportedKeys(session, TEST_CATALOG, TEST_SCHEMA, null), + "listExportedKeys should throw for null table"); + assertEquals(EXECUTE_STATEMENT_FAILED.getCode(), exportedKeysError.getErrorCode()); + assertEquals(SYNTAX_OR_ACCESS_VIOLATION_SQLSTATE, exportedKeysError.getSQLState()); } @Test @@ -1362,7 +1537,6 @@ void testListCrossReferences_handlesNullSqlStateWithoutNPE() throws Exception { @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); @@ -1370,30 +1544,15 @@ void testListCatalogsWithMultipleCatalogSupportDisabled() throws SQLException { DatabricksMetadataQueryClient metadataClient = new DatabricksMetadataQueryClient(mockClient); - String expectedSQL = "SELECT 'my_catalog' AS catalog"; - when(mockClient.executeStatement( - eq(expectedSQL), - eq(mockedComputeResource), - any(), - eq(StatementType.METADATA), - eq(session), - any(), - any(MetadataOperationType.class))) - .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()); + assertTrue(actualResult.next()); + assertEquals("my_catalog", actualResult.getString("TABLE_CAT")); + verify(mockClient, never()) + .executeStatement(anyString(), any(), any(), any(), any(), any(), any()); } @Test diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java index 4a6e266295..fdfb3ed81a 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java @@ -28,6 +28,7 @@ import com.databricks.jdbc.model.client.sqlexec.*; import com.databricks.jdbc.model.client.sqlexec.ExecuteStatementRequest; import com.databricks.jdbc.model.client.sqlexec.ExecuteStatementResponse; +import com.databricks.jdbc.model.core.ColumnInfo; import com.databricks.jdbc.model.core.Disposition; import com.databricks.jdbc.model.core.ResultData; import com.databricks.jdbc.model.core.ResultManifest; @@ -45,6 +46,10 @@ import javax.net.ssl.SSLHandshakeException; 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.EnumSource; +import org.junit.jupiter.params.provider.NullSource; +import org.junit.jupiter.params.provider.ValueSource; import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @@ -83,6 +88,15 @@ private void setupSessionMocks() throws IOException { } private void setupClientMocks(boolean includeResults, boolean async) throws IOException { + setupClientMocks(includeResults, async, new ArrayList<>(), null); + } + + private void setupClientMocks( + boolean includeResults, + boolean async, + List manifestColumns, + Boolean isNativeMetadataResult) + throws IOException { List params = new ArrayList<>() { { @@ -120,7 +134,11 @@ private void setupClientMocks(boolean includeResults, boolean async) throws IOEx .setManifest( new ResultManifest() .setFormat(Format.JSON_ARRAY) - .setSchema(new ResultSchema().setColumns(new ArrayList<>()).setColumnCount(0L)) + .setSchema( + new ResultSchema() + .setColumns(manifestColumns) + .setColumnCount((long) manifestColumns.size())) + .setIsNativeMetadataResult(isNativeMetadataResult) .setTotalRowCount(0L)); } @@ -1164,6 +1182,114 @@ public void testMetadataOperationTypeHeaderWithGetTables() throws Exception { eq(ExecuteStatementResponse.class)); } + @ParameterizedTest + @EnumSource( + value = MetadataOperationType.class, + names = { + "GET_CATALOGS", + "GET_SCHEMAS", + "GET_TABLES", + "GET_COLUMNS", + "GET_FUNCTIONS", + "GET_PRIMARY_KEYS", + "GET_CROSS_REFERENCE" + }) + public void testSupportedOperationsRequestThriftNativeMetadata( + MetadataOperationType operationType) throws Exception { + DatabricksResultSet resultSet = + executeMetadataOperation(operationType, true, Boolean.TRUE, new ArrayList<>()); + + assertTrue(resultSet.isThriftNativeMetadataResult()); + verifyNativeMetadataHeader(operationType, true); + } + + @ParameterizedTest + @NullSource + @ValueSource(booleans = false) + public void testNativeMetadataSchemaDoesNotOverrideManifestFlag(Boolean manifestFlag) + throws Exception { + DatabricksResultSet resultSet = + executeMetadataOperation( + MetadataOperationType.GET_COLUMNS, + true, + manifestFlag, + List.of( + new ColumnInfo() + .setName("ORDINAL_POSITION") + .setTypeName(STRING) + .setTypeText("STRING"))); + + assertFalse(resultSet.isThriftNativeMetadataResult()); + } + + @Test + public void testNativeMetadataManifestFlagDoesNotDependOnRequestHeader() throws Exception { + DatabricksResultSet resultSet = + executeMetadataOperation( + MetadataOperationType.GET_COLUMNS, false, Boolean.TRUE, new ArrayList<>()); + + assertTrue(resultSet.isThriftNativeMetadataResult()); + verifyNativeMetadataHeader(MetadataOperationType.GET_COLUMNS, false); + } + + @ParameterizedTest + @EnumSource( + value = MetadataOperationType.class, + names = {"GET_PROCEDURES", "GET_PROCEDURE_COLUMNS"}) + public void testUnsupportedOperationDoesNotRequestThriftNativeMetadata( + MetadataOperationType operationType) throws Exception { + DatabricksResultSet resultSet = + executeMetadataOperation(operationType, true, null, new ArrayList<>()); + + assertFalse(resultSet.isThriftNativeMetadataResult()); + verifyNativeMetadataHeader(operationType, false); + } + + private DatabricksResultSet executeMetadataOperation( + MetadataOperationType operationType, + boolean enableThriftNativeMetadata, + Boolean manifestFlag, + List manifestColumns) + throws Exception { + setupClientMocks(true, false, manifestColumns, manifestFlag); + String jdbcUrl = + enableThriftNativeMetadata ? JDBC_URL + "EnableThriftNativeMetadata=1;" : JDBC_URL; + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(jdbcUrl, new Properties()); + DatabricksSdkClient databricksSdkClient = + new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient); + DatabricksConnection connection = + new DatabricksConnection(connectionContext, databricksSdkClient); + connection.open(); + + return databricksSdkClient.executeStatement( + "metadata query", + warehouse, + new HashMap<>(), + StatementType.METADATA, + connection.getSession(), + new DatabricksStatement(connection), + operationType); + } + + private void verifyNativeMetadataHeader(MetadataOperationType operationType, boolean expected) + throws IOException { + verify(apiClient, atLeastOnce()) + .execute( + argThat( + req -> { + Map headers = req.getHeaders(); + return headers != null + && operationType + .getHeaderValue() + .equals(headers.get("X-Databricks-Metadata-Operation-Type")) + && Objects.equals( + expected ? "true" : null, + headers.get("X-Databricks-Require-Thrift-Native-Metadata")); + }), + eq(ExecuteStatementResponse.class)); + } + @Test public void testExecuteStatementWithClosedStatus() throws Exception { // Set up connection and statement 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 58b65b5ba5..7f58270c15 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 @@ -3,6 +3,7 @@ import static com.databricks.jdbc.common.DatabricksJdbcConstants.QUERY_EXECUTION_TIMEOUT_SQLSTATE; import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_BYTE_LIMIT; import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_ROW_LIMIT_PER_BLOCK; +import static com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode.EXECUTE_STATEMENT_FAILED; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.Mockito.*; @@ -321,7 +322,7 @@ void testExecuteThrowsSQLExceptionWithSqlState() assertEquals("Error executing statement", exception.getMessage()); assertEquals("42000", exception.getSQLState()); - assertEquals(1003, exception.getErrorCode()); // EXECUTE_STATEMENT_FAILED stable code + assertEquals(EXECUTE_STATEMENT_FAILED.getCode(), exception.getErrorCode()); } @Test @@ -542,7 +543,7 @@ void testPollingPath_errorStatusWithNullMessage_includesErrorCode() throws Excep // Verify the enriched message includes errorCode instead of "error: [null]" assertTrue(exception.getMessage().contains("errorCode=502")); assertFalse(exception.getMessage().contains("error: [null]")); - assertEquals(1003, exception.getErrorCode()); // EXECUTE_STATEMENT_FAILED stable code + assertEquals(EXECUTE_STATEMENT_FAILED.getCode(), exception.getErrorCode()); } @Test @@ -1155,6 +1156,7 @@ void testMetadataPollingThrowsOnInvalidHandleStatus() DatabricksSQLException exception = assertThrows(DatabricksSQLException.class, () -> accessor.getThriftResponse(request)); assertTrue(exception.getMessage().contains("INVALID_HANDLE_STATUS")); + assertEquals(EXECUTE_STATEMENT_FAILED.getCode(), exception.getErrorCode()); } @Test @@ -1347,7 +1349,7 @@ void testExecute_remapsConcurrentModificationOnOperationStateBranchToSerializati "40001", e.getSQLState(), "Expected ConcurrentModificationException with 42000 to be remapped to 40001"); - assertEquals(1003, e.getErrorCode()); // EXECUTE_STATEMENT_FAILED stable code + assertEquals(EXECUTE_STATEMENT_FAILED.getCode(), e.getErrorCode()); } private TFetchResultsReq getFetchResultsRequest(boolean includeMetadata) diff --git a/src/test/java/com/databricks/jdbc/model/core/ResultManifestTest.java b/src/test/java/com/databricks/jdbc/model/core/ResultManifestTest.java new file mode 100644 index 0000000000..c39ee07d55 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/model/core/ResultManifestTest.java @@ -0,0 +1,24 @@ +package com.databricks.jdbc.model.core; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Test; + +public class ResultManifestTest { + + @Test + void testNativeMetadataResultJsonRoundTrip() throws Exception { + ObjectMapper mapper = new ObjectMapper(); + ResultManifest manifest = + mapper.readValue("{\"is_native_metadata_result\":true}", ResultManifest.class); + + assertEquals(Boolean.TRUE, manifest.getIsNativeMetadataResult()); + assertEquals( + Boolean.TRUE, + mapper + .readTree(mapper.writeValueAsString(manifest)) + .get("is_native_metadata_result") + .booleanValue()); + } +}