From 02557c859b12433c24c8060ccc6e9eb2a1d67565 Mon Sep 17 00:00:00 2001 From: samikshya-chand_data Date: Tue, 2 Sep 2025 21:03:02 +0530 Subject: [PATCH 1/6] Fix getFunctions in SEA + add a workaround for getFunctions in Thrift --- .../api/impl/DatabricksConnectionContext.java | 5 ++ .../api/impl/DatabricksDatabaseMetaData.java | 36 +++++++++-- .../IDatabricksConnectionContext.java | 2 + .../jdbc/common/DatabricksJdbcUrlParams.java | 4 +- .../impl/common/CommandConstants.java | 3 +- .../dbclient/impl/sqlexec/CommandBuilder.java | 20 +++--- .../thrift/DatabricksThriftServiceClient.java | 2 +- .../impl/DatabricksDatabaseMetaDataTest.java | 63 ++++++++++++++++--- .../impl/sqlexec/CommandBuilderTest.java | 5 +- .../DatabricksMetadataSdkClientTest.java | 9 ++- 10 files changed, 116 insertions(+), 33 deletions(-) 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 9b26b4f4aa..b56a78e6bc 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -958,6 +958,11 @@ public Integer getHttpConnectionRequestTimeout() { return null; } + @Override + public boolean enableShowCommandsForGetFunctions() { + return getParameter(DatabricksJdbcUrlParams.ENABLE_SHOW_COMMAND_FOR_GET_FUNCTIONS).equals("1"); + } + private static boolean nullOrEmptyString(String s) { return s == null || s.isEmpty(); } 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 f9d8d14bf2..8bb0f712c7 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaData.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaData.java @@ -2,6 +2,7 @@ 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.CommandName.LIST_FUNCTIONS; import static com.databricks.jdbc.dbclient.impl.sqlexec.ResultConstants.CLIENT_INFO_PROPERTIES_RESULT; import com.databricks.jdbc.api.impl.converters.ConverterHelper; @@ -9,8 +10,10 @@ import com.databricks.jdbc.api.internal.IDatabricksSession; import com.databricks.jdbc.common.*; import com.databricks.jdbc.common.util.DriverUtil; +import com.databricks.jdbc.common.util.WildcardUtil; import com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder; import com.databricks.jdbc.dbclient.impl.common.StatementId; +import com.databricks.jdbc.dbclient.impl.sqlexec.CommandBuilder; import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; @@ -1534,11 +1537,22 @@ public ResultSet getFunctions(String catalog, String schemaPattern, String funct functionNamePattern)); throwExceptionIfConnectionIsClosed(); try { - return session - .getDatabricksMetadataClient() - .listFunctions(session, catalog, schemaPattern, functionNamePattern); + if (connection.getConnectionContext().enableShowCommandsForGetFunctions()) { + return getFunctionsFromCommand( + catalog, + schemaPattern, + functionNamePattern); // TODO : remove this once thrift implementation is fixed. + } else { + if (WildcardUtil.isNullOrEmpty(functionNamePattern)) { + functionNamePattern = + "%"; // This is because functionName is a required parameter in thrift flow. + } + return session + .getDatabricksMetadataClient() + .listFunctions(session, catalog, schemaPattern, functionNamePattern); + } } catch (Exception e) { - LOGGER.error(e, "Unable to fetch functions, returning empty result set"); + LOGGER.error(e, "Unable to fetch functions with error {}, returning empty result set", e); return metadataResultSetBuilder.getFunctionsResult(catalog, List.of()); } } @@ -1618,4 +1632,18 @@ private void throwExceptionIfConnectionIsClosed() throws SQLException { "Connection closed!", DatabricksDriverErrorCode.CONNECTION_CLOSED); } } + + private ResultSet getFunctionsFromCommand( + String catalog, String schemaPattern, String functionNamePattern) throws SQLException { + Statement internalStatement = connection.createStatement(); + String showFunctionsSqlCommand = + new CommandBuilder(catalog, session) + .setSchemaPattern(schemaPattern) + .setFunctionPattern(functionNamePattern) + .getSQLString(LIST_FUNCTIONS); + DatabricksResultSet rs = + (DatabricksResultSet) internalStatement.executeQuery(showFunctionsSqlCommand); + return new MetadataResultSetBuilder(connection.getConnectionContext()) + .getFunctionsResult(rs, catalog); + } } 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 2da74837d4..c1d04fc676 100644 --- a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java @@ -350,4 +350,6 @@ public interface IDatabricksConnectionContext { /** Returns the HTTP connection request timeout in seconds */ Integer getHttpConnectionRequestTimeout(); + + boolean enableShowCommandsForGetFunctions(); } diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java index 433447a162..55cb4f4d18 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java @@ -153,7 +153,9 @@ public enum DatabricksJdbcUrlParams { HTTP_CONNECTION_REQUEST_TIMEOUT( "HttpConnectionRequestTimeout", "HTTP connection request timeout in seconds"), CLOUD_FETCH_SPEED_THRESHOLD( - "CloudFetchSpeedThreshold", "Minimum expected download speed in MB/s", "0.1"); + "CloudFetchSpeedThreshold", "Minimum expected download speed in MB/s", "0.1"), + ENABLE_SHOW_COMMAND_FOR_GET_FUNCTIONS( + "EnableShowCommandForGetFunctions", "Use SQL command to fetch function list", "0"); private final String paramName; private final String defaultValue; diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/CommandConstants.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/CommandConstants.java index 7abb321f48..b959d2b6be 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/CommandConstants.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/CommandConstants.java @@ -19,11 +19,12 @@ public class CommandConstants { public static final String SHOW_TABLES_SQL = "SHOW TABLES" + IN_CATALOG_SQL; public static final String SHOW_TABLES_IN_ALL_CATALOGS_SQL = "SHOW TABLES" + IN_ALL_CATALOGS_SQL; public static final String SHOW_COLUMNS_SQL = "SHOW COLUMNS" + IN_CATALOG_SQL; - public static final String SHOW_FUNCTIONS_SQL = "SHOW FUNCTIONS" + IN_CATALOG_SQL; + public static final String SHOW_FUNCTIONS_SQL = "SHOW USER FUNCTIONS" + IN_CATALOG_SQL; public static final String SHOW_SCHEMAS_IN_ALL_CATALOGS_SQL = "SHOW SCHEMAS" + IN_ALL_CATALOGS_SQL; public static final String SHOW_PRIMARY_KEYS_SQL = "SHOW KEYS" + IN_CATALOG_SQL + IN_ABSOLUTE_SCHEMA_SQL + IN_ABSOLUTE_TABLE_SQL; public static final String SHOW_FOREIGN_KEYS_SQL = "SHOW FOREIGN KEYS" + IN_CATALOG_SQL + IN_ABSOLUTE_SCHEMA_SQL + IN_ABSOLUTE_TABLE_SQL; + public static final String USE_CATALOG_SQL = "USE CATALOG %s"; } diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/CommandBuilder.java b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/CommandBuilder.java index a1ffba7db8..8645f1726d 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/CommandBuilder.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/CommandBuilder.java @@ -6,10 +6,10 @@ import com.databricks.jdbc.api.internal.IDatabricksSession; import com.databricks.jdbc.common.util.WildcardUtil; -import com.databricks.jdbc.exception.DatabricksSQLFeatureNotSupportedException; +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 java.sql.SQLException; import java.util.Collections; import java.util.HashMap; import java.util.Map; @@ -74,7 +74,7 @@ private String fetchCatalogSQL() { return SHOW_CATALOGS_SQL; } - private String fetchSchemaSQL() throws SQLException { + private String fetchSchemaSQL() { LOGGER.debug( "Building command for fetching schema. Catalog %s, SchemaPattern %s and session context %s", catalogName, schemaPattern, sessionContext); @@ -91,7 +91,7 @@ private String fetchSchemaSQL() throws SQLException { return showSchemasSQL; } - private String fetchTablesSQL() throws SQLException { + private String fetchTablesSQL() { LOGGER.debug( "Building command for fetching tables. Catalog %s, SchemaPattern %s, TablePattern %s and session context %s", catalogName, schemaPattern, tablePattern, sessionContext); @@ -111,7 +111,7 @@ private String fetchTablesSQL() throws SQLException { return showTablesSQL; } - private String fetchColumnsSQL() throws SQLException { + private String fetchColumnsSQL() throws DatabricksSQLException { String contextString = String.format( "Building command for fetching columns. Catalog %s, SchemaPattern %s, TablePattern %s, ColumnPattern %s and session context : %s", @@ -134,7 +134,7 @@ private String fetchColumnsSQL() throws SQLException { return showColumnsSQL; } - private String fetchFunctionsSQL() throws SQLException { + private String fetchFunctionsSQL() throws DatabricksSQLException { String contextString = String.format( "Building command for fetching functions. Catalog %s, SchemaPattern %s, FunctionPattern %s. With session context %s", @@ -156,7 +156,7 @@ private String fetchTableTypesSQL() { return SHOW_TABLE_TYPES_SQL; } - private String fetchPrimaryKeysSQL() throws SQLException { + private String fetchPrimaryKeysSQL() throws DatabricksSQLException { String contextString = String.format( "Building command for fetching primary keys. Catalog %s, Schema %s, Table %s. With session context: %s", @@ -170,7 +170,7 @@ private String fetchPrimaryKeysSQL() throws SQLException { return String.format(SHOW_PRIMARY_KEYS_SQL, catalogName, schemaName, tableName); } - private String fetchForeignKeysSQL() throws SQLException { + private String fetchForeignKeysSQL() throws DatabricksSQLException { String contextString = String.format( "Building command for fetching foreign keys. Catalog %s, Schema %s, Table %s. With session context: %s", @@ -184,7 +184,7 @@ private String fetchForeignKeysSQL() throws SQLException { return String.format(SHOW_FOREIGN_KEYS_SQL, catalogName, schemaName, tableName); } - public String getSQLString(CommandName command) throws SQLException { + public String getSQLString(CommandName command) throws DatabricksSQLException { switch (command) { case LIST_CATALOGS: return fetchCatalogSQL(); @@ -203,7 +203,7 @@ public String getSQLString(CommandName command) throws SQLException { case LIST_FOREIGN_KEYS: return fetchForeignKeysSQL(); } - throw new DatabricksSQLFeatureNotSupportedException( + throw new DatabricksValidationException( String.format("Invalid command issued %s. Context: %s", command, sessionContext)); } } 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 55635d3841..e718eb658c 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 @@ -436,7 +436,7 @@ public DatabricksResultSet listFunctions( String catalog, String schemaNamePattern, String functionNamePattern) - throws DatabricksSQLException { + throws SQLException { String context = String.format( "Fetching functions using Thrift client. Session {%s}, catalog {%s}, schemaNamePattern {%s}, functionNamePattern {%s}.", 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 ea0497478b..0fa254befc 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java @@ -9,31 +9,40 @@ import com.databricks.jdbc.api.internal.IDatabricksSession; import com.databricks.jdbc.common.DatabricksJdbcConstants; import com.databricks.jdbc.dbclient.IDatabricksMetadataClient; +import com.databricks.jdbc.dbclient.impl.thrift.DatabricksThriftServiceClient; import com.databricks.jdbc.exception.DatabricksSQLException; import java.sql.*; 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.mockito.Mock; import org.mockito.Mockito; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.LENIENT) public class DatabricksDatabaseMetaDataTest { + @Mock IDatabricksConnectionInternal connection; + @Mock IDatabricksSession session; + @Mock IDatabricksMetadataClient metadataClient; + @Mock DatabricksConnectionContext connectionContext; + @Mock DatabricksThriftServiceClient client; + @Mock ResultSetMetaData resultSetMetaData; + @Mock DatabricksResultSet resultSet; - private IDatabricksConnectionInternal connection; - private IDatabricksSession session; - private DatabricksDatabaseMetaData metaData; - private IDatabricksMetadataClient metadataClient; + DatabricksDatabaseMetaData metaData; @BeforeEach public void setup() throws SQLException { - connection = Mockito.mock(IDatabricksConnectionInternal.class); - session = Mockito.mock(IDatabricksSession.class); when(connection.getSession()).thenReturn(session); metaData = new DatabricksDatabaseMetaData(connection); - metadataClient = Mockito.mock(IDatabricksMetadataClient.class); when(session.getDatabricksMetadataClient()).thenReturn(metadataClient); when(metadataClient.listTables(any(), any(), any(), any(), any())) .thenReturn(Mockito.mock(DatabricksResultSet.class)); @@ -43,6 +52,8 @@ public void setup() throws SQLException { .thenReturn(Mockito.mock(DatabricksResultSet.class)); when(session.getConnectionContext()) .thenReturn(DatabricksConnectionContext.parse(WAREHOUSE_JDBC_URL, new Properties())); + when(connection.getConnectionContext()) + .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)); @@ -67,7 +78,7 @@ public void getDatabaseProductName_returnsCorrectProductName() throws Exception } @Test - public void getDatabaseProductName_throwsExceptionWhenConnectionIsClosed() throws Exception { + public void getDatabaseProductName_throwsExceptionWhenConnectionIsClosed() throws SQLException { when(connection.getSession().isOpen()).thenReturn(false); try { metaData.getDatabaseProductName(); @@ -798,12 +809,14 @@ public void testGetDriverVersion() throws SQLException { @Test public void testGetDriverMajorVersion() { + // setup() not needed; method does not depend on connection state int result = metaData.getDriverMajorVersion(); assertEquals(1, result); } @Test public void testGetDriverMinorVersion() { + // setup() not needed; method does not depend on connection state int result = metaData.getDriverMinorVersion(); assertEquals(0, result); } @@ -1640,4 +1653,38 @@ private static Stream provideCatalogSchemaEntityPatternParams() { // Test case 9: Special characters in patterns Arguments.of(null, "_test%", "%ENTITY_", "_column%", "Special characters in patterns")); } + + @Test + void testGetTablesFunctionsWithShowCommandEnabled() throws SQLException { + when(connection.getSession()).thenReturn(session); + when(connection.getConnectionContext()).thenReturn(connectionContext); + when(session.getConnectionContext()).thenReturn(connectionContext); + metaData = new DatabricksDatabaseMetaData(connection); + when(session.isOpen()).thenReturn(true); + when(connectionContext.enableShowCommandsForGetFunctions()).thenReturn(true); + when(resultSetMetaData.getColumnCount()).thenReturn(6); + when(resultSetMetaData.getColumnName(1)).thenReturn("functionName"); + when(resultSetMetaData.getColumnName(2)).thenReturn("namespace"); + when(resultSetMetaData.getColumnName(3)).thenReturn("catalogName"); + when(resultSetMetaData.getColumnName(4)).thenReturn("remarks"); + when(resultSetMetaData.getColumnName(5)).thenReturn("functionType"); + when(resultSetMetaData.getColumnName(6)).thenReturn("specificName"); + when(resultSet.getMetaData()).thenReturn(resultSetMetaData); + when(resultSet.next()).thenReturn(true, false); + when(resultSet.getObject("functionName")).thenReturn("my_fn"); + when(resultSet.getObject("namespace")).thenReturn(TEST_SCHEMA); + when(resultSet.getObject("catalogName")).thenReturn(TEST_CATALOG); + when(resultSet.getObject("remarks")).thenReturn("remark"); + when(resultSet.getObject("functionType")).thenReturn(1); + when(resultSet.getObject("specificName")).thenReturn("my_fn"); + Statement statement = Mockito.mock(Statement.class); + when(connection.createStatement()).thenReturn(statement); + when(statement.executeQuery("SHOW USER FUNCTIONS IN CATALOG catalog1 SCHEMA LIKE 'testSchema'")) + .thenReturn(resultSet); + ResultSet out = metaData.getFunctions(TEST_CATALOG, TEST_SCHEMA, null); + assertNotNull(out); + assertTrue(out.next()); + assertEquals(TEST_CATALOG, out.getString("FUNCTION_CAT")); + assertEquals("my_fn", out.getString("FUNCTION_NAME")); + } } diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/CommandBuilderTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/CommandBuilderTest.java index 4130df7543..0de1fe43ab 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/CommandBuilderTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/CommandBuilderTest.java @@ -6,7 +6,7 @@ import com.databricks.jdbc.api.internal.IDatabricksSession; import com.databricks.jdbc.common.util.WildcardUtil; -import com.databricks.jdbc.exception.DatabricksSQLFeatureNotSupportedException; +import com.databricks.jdbc.exception.DatabricksValidationException; import java.sql.SQLException; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; @@ -200,7 +200,6 @@ void shouldThrowExceptionForUnsupportedCommand() { CommandName mockCommand = mock(CommandName.class); - assertThrows( - DatabricksSQLFeatureNotSupportedException.class, () -> builder.getSQLString(mockCommand)); + assertThrows(DatabricksValidationException.class, () -> builder.getSQLString(mockCommand)); } } 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 647f71f93a..1c5810e4ed 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 @@ -82,19 +82,19 @@ private static Stream listSchemasTestParams() { private static Stream listFunctionsTestParams() { return Stream.of( Arguments.of( - "SHOW FUNCTIONS IN CATALOG catalog1 SCHEMA LIKE 'testSchema' LIKE 'functionPattern'", + "SHOW USER FUNCTIONS IN CATALOG catalog1 SCHEMA LIKE 'testSchema' LIKE 'functionPattern'", TEST_CATALOG, TEST_SCHEMA, TEST_FUNCTION_PATTERN, "test for get functions"), Arguments.of( - "SHOW FUNCTIONS IN CATALOG catalog1 LIKE 'functionPattern'", + "SHOW USER FUNCTIONS IN CATALOG catalog1 LIKE 'functionPattern'", TEST_CATALOG, null, TEST_FUNCTION_PATTERN, "test for get functions without schema"), Arguments.of( - "SHOW FUNCTIONS IN CATALOG catalog1 SCHEMA LIKE 'testSchema'", + "SHOW USER FUNCTIONS IN CATALOG catalog1 SCHEMA LIKE 'testSchema'", TEST_CATALOG, TEST_SCHEMA, null, @@ -742,7 +742,7 @@ void testListCrossReferences_notAvailable() throws Exception { @ParameterizedTest @MethodSource("listFunctionsTestParams") - void testTestFunctions( + void testGetFunctions( String sql, String catalog, String schema, String functionPattern, String description) throws SQLException { when(session.getComputeResource()).thenReturn(WAREHOUSE_COMPUTE); @@ -768,7 +768,6 @@ void testTestFunctions( when(mockedResultSet.getMetaData()).thenReturn(mockedMetaData); DatabricksResultSet actualResult = metadataClient.listFunctions(session, catalog, schema, functionPattern); - assertEquals( actualResult.getStatementStatus().getState(), StatementState.SUCCEEDED, description); assertEquals(actualResult.getStatementId(), GET_FUNCTIONS_STATEMENT_ID, description); From d24c61bb2d7af937c39b56e9882b1bdd0b8c7fef Mon Sep 17 00:00:00 2001 From: samikshya-chand_data Date: Tue, 2 Sep 2025 23:51:02 +0530 Subject: [PATCH 2/6] Reset based on spec --- .../jdbc/dbclient/impl/common/CommandConstants.java | 2 +- .../jdbc/api/impl/DatabricksDatabaseMetaDataTest.java | 2 +- .../impl/sqlexec/DatabricksMetadataSdkClientTest.java | 6 +++--- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/CommandConstants.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/CommandConstants.java index b959d2b6be..b292290d42 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/CommandConstants.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/CommandConstants.java @@ -19,7 +19,7 @@ public class CommandConstants { public static final String SHOW_TABLES_SQL = "SHOW TABLES" + IN_CATALOG_SQL; public static final String SHOW_TABLES_IN_ALL_CATALOGS_SQL = "SHOW TABLES" + IN_ALL_CATALOGS_SQL; public static final String SHOW_COLUMNS_SQL = "SHOW COLUMNS" + IN_CATALOG_SQL; - public static final String SHOW_FUNCTIONS_SQL = "SHOW USER FUNCTIONS" + IN_CATALOG_SQL; + public static final String SHOW_FUNCTIONS_SQL = "SHOW FUNCTIONS" + IN_CATALOG_SQL; public static final String SHOW_SCHEMAS_IN_ALL_CATALOGS_SQL = "SHOW SCHEMAS" + IN_ALL_CATALOGS_SQL; public static final String SHOW_PRIMARY_KEYS_SQL = 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 0fa254befc..afc9c12dc8 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java @@ -1679,7 +1679,7 @@ void testGetTablesFunctionsWithShowCommandEnabled() throws SQLException { when(resultSet.getObject("specificName")).thenReturn("my_fn"); Statement statement = Mockito.mock(Statement.class); when(connection.createStatement()).thenReturn(statement); - when(statement.executeQuery("SHOW USER FUNCTIONS IN CATALOG catalog1 SCHEMA LIKE 'testSchema'")) + when(statement.executeQuery("SHOW FUNCTIONS IN CATALOG catalog1 SCHEMA LIKE 'testSchema'")) .thenReturn(resultSet); ResultSet out = metaData.getFunctions(TEST_CATALOG, TEST_SCHEMA, null); assertNotNull(out); 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 1c5810e4ed..f71dab3bdd 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 @@ -82,19 +82,19 @@ private static Stream listSchemasTestParams() { private static Stream listFunctionsTestParams() { return Stream.of( Arguments.of( - "SHOW USER FUNCTIONS IN CATALOG catalog1 SCHEMA LIKE 'testSchema' LIKE 'functionPattern'", + "SHOW FUNCTIONS IN CATALOG catalog1 SCHEMA LIKE 'testSchema' LIKE 'functionPattern'", TEST_CATALOG, TEST_SCHEMA, TEST_FUNCTION_PATTERN, "test for get functions"), Arguments.of( - "SHOW USER FUNCTIONS IN CATALOG catalog1 LIKE 'functionPattern'", + "SHOW FUNCTIONS IN CATALOG catalog1 LIKE 'functionPattern'", TEST_CATALOG, null, TEST_FUNCTION_PATTERN, "test for get functions without schema"), Arguments.of( - "SHOW USER FUNCTIONS IN CATALOG catalog1 SCHEMA LIKE 'testSchema'", + "SHOW FUNCTIONS IN CATALOG catalog1 SCHEMA LIKE 'testSchema'", TEST_CATALOG, TEST_SCHEMA, null, From 4f581c2be29d045b5d1adc0512fc2caac4dc2c12 Mon Sep 17 00:00:00 2001 From: samikshya-chand_data Date: Wed, 10 Sep 2025 13:25:57 +0530 Subject: [PATCH 3/6] address comments --- .../api/impl/DatabricksDatabaseMetaData.java | 46 +++++-------------- .../impl/common/CommandConstants.java | 1 - .../thrift/DatabricksThriftServiceClient.java | 25 +++++++++- .../impl/DatabricksDatabaseMetaDataTest.java | 8 ++-- 4 files changed, 38 insertions(+), 42 deletions(-) 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 8bb0f712c7..2db6a64fad 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.CommandName.LIST_FUNCTIONS; import static com.databricks.jdbc.dbclient.impl.sqlexec.ResultConstants.CLIENT_INFO_PROPERTIES_RESULT; import com.databricks.jdbc.api.impl.converters.ConverterHelper; @@ -13,7 +12,6 @@ import com.databricks.jdbc.common.util.WildcardUtil; import com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder; import com.databricks.jdbc.dbclient.impl.common.StatementId; -import com.databricks.jdbc.dbclient.impl.sqlexec.CommandBuilder; import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; @@ -1530,29 +1528,21 @@ public ResultSet getClientInfoProperties() throws SQLException { public ResultSet getFunctions(String catalog, String schemaPattern, String functionNamePattern) throws SQLException { LOGGER.debug( - String.format( - "public ResultSet getFunctions(String catalog = {}, String schemaPattern = {}, String functionNamePattern = {})", - catalog, - schemaPattern, - functionNamePattern)); + "public ResultSet getFunctions(String catalog = {}, String schemaPattern = {}, String functionNamePattern = {})", + catalog, + schemaPattern, + functionNamePattern); throwExceptionIfConnectionIsClosed(); try { - if (connection.getConnectionContext().enableShowCommandsForGetFunctions()) { - return getFunctionsFromCommand( - catalog, - schemaPattern, - functionNamePattern); // TODO : remove this once thrift implementation is fixed. - } else { - if (WildcardUtil.isNullOrEmpty(functionNamePattern)) { - functionNamePattern = - "%"; // This is because functionName is a required parameter in thrift flow. - } - return session - .getDatabricksMetadataClient() - .listFunctions(session, catalog, schemaPattern, functionNamePattern); + if (WildcardUtil.isNullOrEmpty(functionNamePattern)) { + functionNamePattern = + "%"; // This is because functionName is a required parameter in thrift flow. } + return session + .getDatabricksMetadataClient() + .listFunctions(session, catalog, schemaPattern, functionNamePattern); } catch (Exception e) { - LOGGER.error(e, "Unable to fetch functions with error {}, returning empty result set", e); + LOGGER.error(e, "Unable to fetch functions, returning empty result set"); return metadataResultSetBuilder.getFunctionsResult(catalog, List.of()); } } @@ -1632,18 +1622,4 @@ private void throwExceptionIfConnectionIsClosed() throws SQLException { "Connection closed!", DatabricksDriverErrorCode.CONNECTION_CLOSED); } } - - private ResultSet getFunctionsFromCommand( - String catalog, String schemaPattern, String functionNamePattern) throws SQLException { - Statement internalStatement = connection.createStatement(); - String showFunctionsSqlCommand = - new CommandBuilder(catalog, session) - .setSchemaPattern(schemaPattern) - .setFunctionPattern(functionNamePattern) - .getSQLString(LIST_FUNCTIONS); - DatabricksResultSet rs = - (DatabricksResultSet) internalStatement.executeQuery(showFunctionsSqlCommand); - return new MetadataResultSetBuilder(connection.getConnectionContext()) - .getFunctionsResult(rs, catalog); - } } diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/CommandConstants.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/CommandConstants.java index b292290d42..7abb321f48 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/CommandConstants.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/CommandConstants.java @@ -26,5 +26,4 @@ public class CommandConstants { "SHOW KEYS" + IN_CATALOG_SQL + IN_ABSOLUTE_SCHEMA_SQL + IN_ABSOLUTE_TABLE_SQL; public static final String SHOW_FOREIGN_KEYS_SQL = "SHOW FOREIGN KEYS" + IN_CATALOG_SQL + IN_ABSOLUTE_SCHEMA_SQL + IN_ABSOLUTE_TABLE_SQL; - public static final String USE_CATALOG_SQL = "USE CATALOG %s"; } 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 e718eb658c..64fce9ee7d 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 @@ -4,6 +4,7 @@ import static com.databricks.jdbc.common.util.DatabricksThriftUtil.*; 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.*; @@ -19,6 +20,7 @@ import com.databricks.jdbc.dbclient.IDatabricksMetadataClient; import com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder; import com.databricks.jdbc.dbclient.impl.common.StatementId; +import com.databricks.jdbc.dbclient.impl.sqlexec.CommandBuilder; import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksParsingException; import com.databricks.jdbc.exception.DatabricksSQLException; @@ -144,7 +146,7 @@ public DatabricksResultSet executeStatement( LOGGER.debug( String.format( "public DatabricksResultSet executeStatement(String sql = {%s}, Compute cluster = {%s}, Map parameters = {%s}, StatementType statementType = {%s}, IDatabricksSession session)", - sql, computeResource.toString(), parameters.toString(), statementType)); + sql, computeResource, parameters.toString(), statementType)); DatabricksThreadContextHolder.setStatementType(statementType); @@ -443,6 +445,27 @@ public DatabricksResultSet listFunctions( session.toString(), catalog, schemaNamePattern, functionNamePattern); DatabricksThreadContextHolder.setSessionId(session.getSessionId()); LOGGER.debug(context); + if (connectionContext.enableShowCommandsForGetFunctions()) { + String showFunctionsSqlCommand = + new CommandBuilder(catalog, session) + .setSchemaPattern(schemaNamePattern) + .setFunctionPattern(functionNamePattern) + .getSQLString(LIST_FUNCTIONS); + LOGGER.debug( + "Fetching functions using SQL Command {{}}. Session {{}}", + showFunctionsSqlCommand, + session.toString()); + try (DatabricksResultSet rs = + executeStatement( + showFunctionsSqlCommand, + session.getComputeResource(), + Collections.emptyMap(), + StatementType.METADATA, + session, + null)) { + return metadataResultSetBuilder.getFunctionsResult(rs, catalog); + } + } TGetFunctionsReq request = new TGetFunctionsReq() .setSessionHandle(Objects.requireNonNull(session.getSessionInfo()).sessionHandle()) 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 afc9c12dc8..805729590c 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java @@ -23,11 +23,9 @@ import org.mockito.Mock; import org.mockito.Mockito; import org.mockito.junit.jupiter.MockitoExtension; -import org.mockito.junit.jupiter.MockitoSettings; -import org.mockito.quality.Strictness; @ExtendWith(MockitoExtension.class) -@MockitoSettings(strictness = Strictness.LENIENT) +/*@MockitoSettings(strictness = Strictness.LENIENT)*/ public class DatabricksDatabaseMetaDataTest { @Mock IDatabricksConnectionInternal connection; @Mock IDatabricksSession session; @@ -1656,12 +1654,12 @@ private static Stream provideCatalogSchemaEntityPatternParams() { @Test void testGetTablesFunctionsWithShowCommandEnabled() throws SQLException { + when(connectionContext.enableShowCommandsForGetFunctions()).thenReturn(true); when(connection.getSession()).thenReturn(session); when(connection.getConnectionContext()).thenReturn(connectionContext); when(session.getConnectionContext()).thenReturn(connectionContext); - metaData = new DatabricksDatabaseMetaData(connection); when(session.isOpen()).thenReturn(true); - when(connectionContext.enableShowCommandsForGetFunctions()).thenReturn(true); + metaData = new DatabricksDatabaseMetaData(connection); when(resultSetMetaData.getColumnCount()).thenReturn(6); when(resultSetMetaData.getColumnName(1)).thenReturn("functionName"); when(resultSetMetaData.getColumnName(2)).thenReturn("namespace"); From c00b4b198712288bad4edc22e45e393c5bfbaaa8 Mon Sep 17 00:00:00 2001 From: samikshya-chand_data Date: Wed, 10 Sep 2025 19:11:43 +0530 Subject: [PATCH 4/6] Fix tests --- .../thrift/DatabricksThriftServiceClient.java | 10 +- .../impl/DatabricksDatabaseMetaDataTest.java | 61 ++-------- .../DatabricksThriftServiceClientTest.java | 109 ++++++++++++++++++ 3 files changed, 124 insertions(+), 56 deletions(-) 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 64fce9ee7d..9f2df55888 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 @@ -1,5 +1,6 @@ package com.databricks.jdbc.dbclient.impl.thrift; +import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_STATEMENT_TIMEOUT_SECONDS; import static com.databricks.jdbc.common.EnvironmentVariables.JDBC_THRIFT_VERSION; import static com.databricks.jdbc.common.util.DatabricksThriftUtil.*; import static com.databricks.jdbc.common.util.DatabricksTypeUtil.DECIMAL; @@ -202,11 +203,14 @@ private TExecuteStatementReq getRequest( parameters.values().stream() .map(this::mapToSparkParameterListItem) .collect(Collectors.toList()); - + int timeout = DEFAULT_STATEMENT_TIMEOUT_SECONDS; + if (parentStatement != null && parentStatement.getStatement() != null) { + timeout = parentStatement.getStatement().getQueryTimeout(); + } TExecuteStatementReq request = new TExecuteStatementReq() .setStatement(sql) - .setQueryTimeout(parentStatement.getStatement().getQueryTimeout()) + .setQueryTimeout(timeout) .setSessionHandle(Objects.requireNonNull(session.getSessionInfo()).sessionHandle()) .setCanReadArrowResult(this.connectionContext.shouldEnableArrow()) .setUseArrowNativeTypes(arrowNativeTypes); @@ -230,7 +234,7 @@ private TExecuteStatementReq getRequest( request.setUseArrowNativeTypes(arrowNativeTypes); } - int maxRows = parentStatement.getMaxRows(); + int maxRows = (parentStatement == null) ? 0 : parentStatement.getMaxRows(); if (maxRows > 0) { // set request param only if user has set maxRows. // Similar // behavior 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 805729590c..ea0497478b 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java @@ -9,38 +9,31 @@ import com.databricks.jdbc.api.internal.IDatabricksSession; import com.databricks.jdbc.common.DatabricksJdbcConstants; import com.databricks.jdbc.dbclient.IDatabricksMetadataClient; -import com.databricks.jdbc.dbclient.impl.thrift.DatabricksThriftServiceClient; import com.databricks.jdbc.exception.DatabricksSQLException; import java.sql.*; 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.mockito.Mock; import org.mockito.Mockito; -import org.mockito.junit.jupiter.MockitoExtension; -@ExtendWith(MockitoExtension.class) -/*@MockitoSettings(strictness = Strictness.LENIENT)*/ public class DatabricksDatabaseMetaDataTest { - @Mock IDatabricksConnectionInternal connection; - @Mock IDatabricksSession session; - @Mock IDatabricksMetadataClient metadataClient; - @Mock DatabricksConnectionContext connectionContext; - @Mock DatabricksThriftServiceClient client; - @Mock ResultSetMetaData resultSetMetaData; - @Mock DatabricksResultSet resultSet; - DatabricksDatabaseMetaData metaData; + private IDatabricksConnectionInternal connection; + private IDatabricksSession session; + private DatabricksDatabaseMetaData metaData; + private IDatabricksMetadataClient metadataClient; @BeforeEach public void setup() throws SQLException { + connection = Mockito.mock(IDatabricksConnectionInternal.class); + session = Mockito.mock(IDatabricksSession.class); when(connection.getSession()).thenReturn(session); metaData = new DatabricksDatabaseMetaData(connection); + metadataClient = Mockito.mock(IDatabricksMetadataClient.class); when(session.getDatabricksMetadataClient()).thenReturn(metadataClient); when(metadataClient.listTables(any(), any(), any(), any(), any())) .thenReturn(Mockito.mock(DatabricksResultSet.class)); @@ -50,8 +43,6 @@ public void setup() throws SQLException { .thenReturn(Mockito.mock(DatabricksResultSet.class)); when(session.getConnectionContext()) .thenReturn(DatabricksConnectionContext.parse(WAREHOUSE_JDBC_URL, new Properties())); - when(connection.getConnectionContext()) - .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)); @@ -76,7 +67,7 @@ public void getDatabaseProductName_returnsCorrectProductName() throws Exception } @Test - public void getDatabaseProductName_throwsExceptionWhenConnectionIsClosed() throws SQLException { + public void getDatabaseProductName_throwsExceptionWhenConnectionIsClosed() throws Exception { when(connection.getSession().isOpen()).thenReturn(false); try { metaData.getDatabaseProductName(); @@ -807,14 +798,12 @@ public void testGetDriverVersion() throws SQLException { @Test public void testGetDriverMajorVersion() { - // setup() not needed; method does not depend on connection state int result = metaData.getDriverMajorVersion(); assertEquals(1, result); } @Test public void testGetDriverMinorVersion() { - // setup() not needed; method does not depend on connection state int result = metaData.getDriverMinorVersion(); assertEquals(0, result); } @@ -1651,38 +1640,4 @@ private static Stream provideCatalogSchemaEntityPatternParams() { // Test case 9: Special characters in patterns Arguments.of(null, "_test%", "%ENTITY_", "_column%", "Special characters in patterns")); } - - @Test - void testGetTablesFunctionsWithShowCommandEnabled() throws SQLException { - when(connectionContext.enableShowCommandsForGetFunctions()).thenReturn(true); - when(connection.getSession()).thenReturn(session); - when(connection.getConnectionContext()).thenReturn(connectionContext); - when(session.getConnectionContext()).thenReturn(connectionContext); - when(session.isOpen()).thenReturn(true); - metaData = new DatabricksDatabaseMetaData(connection); - when(resultSetMetaData.getColumnCount()).thenReturn(6); - when(resultSetMetaData.getColumnName(1)).thenReturn("functionName"); - when(resultSetMetaData.getColumnName(2)).thenReturn("namespace"); - when(resultSetMetaData.getColumnName(3)).thenReturn("catalogName"); - when(resultSetMetaData.getColumnName(4)).thenReturn("remarks"); - when(resultSetMetaData.getColumnName(5)).thenReturn("functionType"); - when(resultSetMetaData.getColumnName(6)).thenReturn("specificName"); - when(resultSet.getMetaData()).thenReturn(resultSetMetaData); - when(resultSet.next()).thenReturn(true, false); - when(resultSet.getObject("functionName")).thenReturn("my_fn"); - when(resultSet.getObject("namespace")).thenReturn(TEST_SCHEMA); - when(resultSet.getObject("catalogName")).thenReturn(TEST_CATALOG); - when(resultSet.getObject("remarks")).thenReturn("remark"); - when(resultSet.getObject("functionType")).thenReturn(1); - when(resultSet.getObject("specificName")).thenReturn("my_fn"); - Statement statement = Mockito.mock(Statement.class); - when(connection.createStatement()).thenReturn(statement); - when(statement.executeQuery("SHOW FUNCTIONS IN CATALOG catalog1 SCHEMA LIKE 'testSchema'")) - .thenReturn(resultSet); - ResultSet out = metaData.getFunctions(TEST_CATALOG, TEST_SCHEMA, null); - assertNotNull(out); - assertTrue(out.next()); - assertEquals(TEST_CATALOG, out.getString("FUNCTION_CAT")); - assertEquals("my_fn", out.getString("FUNCTION_NAME")); - } } 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 4fb755bb4e..3966431b7f 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 @@ -27,6 +27,8 @@ import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.service.sql.StatementState; import java.math.BigDecimal; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.util.*; import java.util.stream.Stream; @@ -57,6 +59,7 @@ public class DatabricksThriftServiceClientTest { @Mock IDatabricksStatementInternal parentStatement; @Mock DatabricksStatement statement; @Mock DatabricksConfig databricksConfig; + @Mock ResultSetMetaData mockedMetaData; @Test void testCreateSession() throws DatabricksSQLException { @@ -580,6 +583,112 @@ void testListFunctions() throws SQLException { assertEquals(resultSet.getStatementStatus().getState(), StatementState.SUCCEEDED); } + @Test + void testListFunctionsWithSQLEnabled() throws SQLException { + DatabricksThriftServiceClient client = + new DatabricksThriftServiceClient(thriftAccessor, connectionContext); + when(connectionContext.enableShowCommandsForGetFunctions()).thenReturn(true); + when(connectionContext.shouldEnableArrow()).thenReturn(true); + when(session.getSessionInfo()).thenReturn(SESSION_INFO); + TSparkArrowTypes arrowNativeTypes = + new TSparkArrowTypes() + .setComplexTypesAsArrow(true) + .setIntervalTypesAsArrow(true) + .setNullTypeAsArrow(true) + .setDecimalAsArrow(true) + .setTimestampAsArrow(true); + TExecuteStatementReq executeStatementReq = + new TExecuteStatementReq() + .setStatement("SHOW FUNCTIONS IN CATALOG catalog1 SCHEMA LIKE 'testSchema' LIKE 'test'") + .setSessionHandle(SESSION_HANDLE) + .setCanReadArrowResult(true) + .setCanDecompressLZ4Result(true) + .setCanDownloadResult(true) + .setQueryTimeout(0) + .setParameters(Collections.emptyList()) + .setRunAsync(true) + .setUseArrowNativeTypes(arrowNativeTypes); + when(thriftAccessor.execute(executeStatementReq, null, session, StatementType.METADATA)) + .thenReturn(resultSet); + when(resultSet.getMetaData()).thenReturn(mockedMetaData); + when(mockedMetaData.getColumnCount()).thenReturn(6); + when(mockedMetaData.getColumnName(1)).thenReturn("functionName"); + when(mockedMetaData.getColumnName(2)).thenReturn("namespace"); + when(mockedMetaData.getColumnName(3)).thenReturn("catalogName"); + when(mockedMetaData.getColumnName(4)).thenReturn("remarks"); + when(mockedMetaData.getColumnName(5)).thenReturn("functionType"); + when(mockedMetaData.getColumnName(6)).thenReturn("specificName"); + when(resultSet.next()).thenReturn(true, false); + when(resultSet.getObject("functionName")).thenReturn("my_fn"); + when(resultSet.getObject("namespace")).thenReturn(TEST_SCHEMA); + when(resultSet.getObject("remarks")).thenReturn("remark"); + when(resultSet.getObject("functionType")).thenReturn(1); + when(resultSet.getObject("specificName")).thenReturn("my_fn"); + + ResultSet actualResultSet = + client.listFunctions(session, TEST_CATALOG, TEST_SCHEMA, TEST_STRING); + assertNotNull(actualResultSet); + assertTrue(actualResultSet.next()); + assertEquals(TEST_CATALOG, actualResultSet.getString("FUNCTION_CAT")); + assertEquals("my_fn", actualResultSet.getString("FUNCTION_NAME")); + } + + @Test + void testGetRequest_DefaultTimeoutAndNoRowLimit_WhenParentStatementNull() throws SQLException { + when(connectionContext.shouldEnableArrow()).thenReturn(true); + DatabricksThriftServiceClient client = + new DatabricksThriftServiceClient(thriftAccessor, connectionContext); + when(session.getSessionInfo()).thenReturn(SESSION_INFO); + + when(thriftAccessor.execute( + any(TExecuteStatementReq.class), eq(null), eq(session), eq(StatementType.SQL))) + .thenReturn(resultSet); + + client.executeStatement( + TEST_STRING, CLUSTER_COMPUTE, Collections.emptyMap(), StatementType.SQL, session, null); + + ArgumentCaptor requestCaptor = + ArgumentCaptor.forClass(TExecuteStatementReq.class); + verify(thriftAccessor) + .execute(requestCaptor.capture(), eq(null), eq(session), eq(StatementType.SQL)); + TExecuteStatementReq request = requestCaptor.getValue(); + assertEquals(0, request.getQueryTimeout()); + assertFalse(request.isSetResultRowLimit()); + } + + @Test + void testGetRequest_DefaultTimeout_WhenStatementNull() throws SQLException { + when(connectionContext.shouldEnableArrow()).thenReturn(true); + DatabricksThriftServiceClient client = + new DatabricksThriftServiceClient(thriftAccessor, connectionContext); + when(session.getSessionInfo()).thenReturn(SESSION_INFO); + when(parentStatement.getStatement()).thenReturn(null); + when(parentStatement.getMaxRows()).thenReturn(0); + + when(thriftAccessor.execute( + any(TExecuteStatementReq.class), + eq(parentStatement), + eq(session), + eq(StatementType.SQL))) + .thenReturn(resultSet); + + client.executeStatement( + TEST_STRING, + CLUSTER_COMPUTE, + Collections.emptyMap(), + StatementType.SQL, + session, + parentStatement); + + ArgumentCaptor requestCaptor = + ArgumentCaptor.forClass(TExecuteStatementReq.class); + verify(thriftAccessor) + .execute(requestCaptor.capture(), eq(parentStatement), eq(session), eq(StatementType.SQL)); + TExecuteStatementReq request = requestCaptor.getValue(); + assertEquals(0, request.getQueryTimeout()); + assertFalse(request.isSetResultRowLimit()); + } + @Test void testListPrimaryKeys() throws SQLException { DatabricksThriftServiceClient client = From 3c1102a55009fe50f421c94e84f5ac9c416819ed Mon Sep 17 00:00:00 2001 From: samikshya-chand_data Date: Thu, 11 Sep 2025 01:00:19 +0530 Subject: [PATCH 5/6] fmt --- .../jdbc/api/internal/IDatabricksConnectionContext.java | 1 + 1 file changed, 1 insertion(+) 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 6fa7f2f49f..e8fb6c4811 100644 --- a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java @@ -362,6 +362,7 @@ public interface IDatabricksConnectionContext { Integer getHttpConnectionRequestTimeout(); boolean enableShowCommandsForGetFunctions(); + /** Returns whether batched INSERT optimization is enabled */ boolean isBatchedInsertsEnabled(); } From 02ab084375d97706008cc38195d6f51a88afe54a Mon Sep 17 00:00:00 2001 From: samikshya-chand_data Date: Thu, 11 Sep 2025 01:12:00 +0530 Subject: [PATCH 6/6] fmt --- .../jdbc/api/impl/DatabricksDatabaseMetaData.java | 2 +- .../jdbc/api/impl/arrow/InlineChunkProvider.java | 12 +++++++----- 2 files changed, 8 insertions(+), 6 deletions(-) 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 2db6a64fad..a8e6b67458 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaData.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaData.java @@ -1542,7 +1542,7 @@ public ResultSet getFunctions(String catalog, String schemaPattern, String funct .getDatabricksMetadataClient() .listFunctions(session, catalog, schemaPattern, functionNamePattern); } catch (Exception e) { - LOGGER.error(e, "Unable to fetch functions, returning empty result set"); + LOGGER.error(e, "Unable to fetch functions, returning empty result set {}", e); return metadataResultSetBuilder.getFunctionsResult(catalog, List.of()); } } 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 39150c6da7..e22d974a40 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 @@ -45,11 +45,13 @@ public class InlineChunkProvider implements ChunkProvider { this.currentChunkIndex = -1; this.totalRows = 0; ByteArrayInputStream byteStream = initializeByteStream(resultsResp, session, parentStatement); - arrowResultChunk = - ArrowResultChunk.builder() - .withInputStream(byteStream, totalRows) - .withStatementId(parentStatement.getStatementId()) - .build(); + ArrowResultChunk.Builder builder = + ArrowResultChunk.builder().withInputStream(byteStream, totalRows); + + if (parentStatement != null) { + builder.withStatementId(parentStatement.getStatementId()); + } + arrowResultChunk = builder.build(); } /**