diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 035fdbf77e..d0e1d5b441 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -6,7 +6,8 @@ - **Query Tags support**: Added ability to attach key-value tags to SQL queries for analytical purposes that would appear in `system.query.history` table. Example: `jdbc:databricks://host;QUERY_TAGS=team:marketing,dashboard:abc123`. - **SQL Scripting support**: Added support for [SQL Scripting](https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-scripting) -- Added a client property `enableVolumeOperations` to enable GET/PUT/REMOVE volume operations on a stream. For backward compatibility, allowedVolumeIngestionPaths can also be used for REMOVE operation. +- Added a client property `enableVolumeOperations` to enable GET/PUT/REMOVE volume operations on a stream. For backward compatibility, allowedVolumeIngestionPaths can also be used for REMOVE operation. +- Support for fetching schemas across all catalogs (when catalog is specified as null or a wildcard) in `DatabaseMetaData#getSchemas` API in SQL Execution mode. ### Updated - Databricks SDK dependency upgraded to latest version 0.60.0 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 ea4ac8866a..f9d8d14bf2 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaData.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaData.java @@ -9,7 +9,6 @@ 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.JdbcThreadUtils; import com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder; import com.databricks.jdbc.dbclient.impl.common.StatementId; import com.databricks.jdbc.exception.DatabricksSQLException; @@ -20,8 +19,6 @@ import com.databricks.sdk.service.sql.StatementState; import java.sql.*; import java.util.*; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; public class DatabricksDatabaseMetaData implements DatabaseMetaData { @@ -40,9 +37,6 @@ public class DatabricksDatabaseMetaData implements DatabaseMetaData { public static final String SYSTEM_FUNCTIONS = "DATABASE,IFNULL,USER"; public static final String TIME_DATE_FUNCTIONS = "CURDATE,CURRENT_DATE,CURRENT_TIME,CURRENT_TIMESTAMP,CURTIME,DAYNAME,DAYOFMONTH,DAYOFWEEK,DAYOFYEAR,HOUR,MINUTE,MONTH,MONTHNAME,NOW,QUARTER,SECOND,TIMESTAMPADD,TIMESTAMPDIFF,WEEK,YEAR"; - private static final Object THREAD_POOL_LOCK = new Object(); - private static ExecutorService schemasThreadPool = null; - private static final int DEFAULT_MAX_THREADS = 10; private final IDatabricksConnectionInternal connection; private final IDatabricksSession session; private final MetadataResultSetBuilder metadataResultSetBuilder; @@ -1505,48 +1499,6 @@ public ResultSet getSchemas(String catalog, String schemaPattern) throws SQLExce catalog, schemaPattern); throwExceptionIfConnectionIsClosed(); - if (session.getConnectionContext().getClientType() == DatabricksClientType.SEA - && (catalog == null || catalog.equals("*") || catalog.equals("%"))) { - // Fetch catalogs from the metadata client - List catalogList = new ArrayList<>(); - try (ResultSet catalogs = getCatalogs()) { - while (catalogs.next()) { - String c = catalogs.getString(1); - if (c != null && !c.isEmpty()) { - catalogList.add(c); - } - } - } - - // Process catalogs in parallel, gathering schema information - List> schemaRows = - JdbcThreadUtils.parallelFlatMap( - catalogList, - session.getConnectionContext(), - DEFAULT_MAX_THREADS, // Not significant since the executor is provided as a parameter - 90, // 90 seconds timeout - c -> { - List> rows = new ArrayList<>(); - try (ResultSet catalogSchemas = - session.getDatabricksMetadataClient().listSchemas(session, c, schemaPattern)) { - while (catalogSchemas.next()) { - List schemaRow = new ArrayList<>(); - schemaRow.add(catalogSchemas.getString(1)); // TABLE_SCHEM - schemaRow.add(catalogSchemas.getString(2)); // TABLE_CATALOG - rows.add(schemaRow); - } - } catch (SQLException e) { - LOGGER.warn("Error fetching schemas for catalog %s %s", c, e.getMessage()); - } - return rows; - }, - getOrCreateSchemasThreadPool()); - - // Convert combined data into a result set - return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns( - SCHEMA_COLUMNS, schemaRows, METADATA_STATEMENT_ID, CommandName.LIST_SCHEMAS); - } - return session.getDatabricksMetadataClient().listSchemas(session, catalog, schemaPattern); } @@ -1659,23 +1611,6 @@ public boolean isWrapperFor(Class iface) throws SQLException { return iface != null && iface.isAssignableFrom(this.getClass()); } - private static ExecutorService getOrCreateSchemasThreadPool() { - synchronized (THREAD_POOL_LOCK) { - if (schemasThreadPool == null || schemasThreadPool.isShutdown()) { - // Could read max threads from a configuration property - schemasThreadPool = - Executors.newFixedThreadPool( - DEFAULT_MAX_THREADS, - r -> { - Thread t = new Thread(r, "jdbc-schemas-fetcher"); - t.setDaemon(true); - return t; - }); - } - return schemasThreadPool; - } - } - private void throwExceptionIfConnectionIsClosed() throws SQLException { LOGGER.debug("private void throwExceptionIfConnectionIsClosed()"); if (!connection.getSession().isOpen()) { diff --git a/src/main/java/com/databricks/jdbc/common/MetadataResultConstants.java b/src/main/java/com/databricks/jdbc/common/MetadataResultConstants.java index 8654da8293..cd4e2c4a9d 100644 --- a/src/main/java/com/databricks/jdbc/common/MetadataResultConstants.java +++ b/src/main/java/com/databricks/jdbc/common/MetadataResultConstants.java @@ -13,9 +13,9 @@ public class MetadataResultConstants { public static final String[] DEFAULT_TABLE_TYPES = {"TABLE", "VIEW", "SYSTEM TABLE"}; public static final ResultColumn CATALOG_COLUMN = new ResultColumn("TABLE_CAT", "catalogName", Types.VARCHAR); - private static final ResultColumn CATALOG_FULL_COLUMN = + public static final ResultColumn CATALOG_FULL_COLUMN = new ResultColumn("TABLE_CATALOG", "catalogName", Types.VARCHAR); - public static final ResultColumn CATALOG_COLUMN_FOR_GET_CATALOGS = + public static final ResultColumn CATALOG_RESULT_COLUMN = new ResultColumn("TABLE_CAT", "catalog", Types.VARCHAR); public static final ResultColumn TYPE_CATALOG_COLUMN = new ResultColumn("TYPE_CAT", "TYPE_CATALOG_COLUMN", Types.VARCHAR); @@ -216,7 +216,7 @@ public class MetadataResultConstants { IS_AUTO_INCREMENT_COLUMN, IS_GENERATED_COLUMN); - public static List CATALOG_COLUMNS = List.of(CATALOG_COLUMN_FOR_GET_CATALOGS); + public static List CATALOG_COLUMNS = List.of(CATALOG_RESULT_COLUMN); public static List SCHEMA_COLUMNS = List.of(SCHEMA_COLUMN_FOR_GET_SCHEMA, CATALOG_FULL_COLUMN); @@ -491,9 +491,7 @@ public class MetadataResultConstants { MetadataResultConstants.TYPE_NAME_COLUMN, MetadataResultConstants.DATA_TYPE_COLUMN, MetadataResultConstants.PRECISION_COLUMN)); - put( - CommandName.LIST_CATALOGS, - List.of(MetadataResultConstants.CATALOG_COLUMN_FOR_GET_CATALOGS)); + put(CommandName.LIST_CATALOGS, List.of(MetadataResultConstants.CATALOG_RESULT_COLUMN)); put( CommandName.LIST_TABLES, List.of(MetadataResultConstants.TABLE_NAME_COLUMN, TABLE_TYPE_COLUMN)); diff --git a/src/main/java/com/databricks/jdbc/common/util/WildcardUtil.java b/src/main/java/com/databricks/jdbc/common/util/WildcardUtil.java index d62b11eceb..a2e150e278 100644 --- a/src/main/java/com/databricks/jdbc/common/util/WildcardUtil.java +++ b/src/main/java/com/databricks/jdbc/common/util/WildcardUtil.java @@ -21,6 +21,10 @@ public static boolean isNullOrEmpty(String s) { return s == null || s.trim().isEmpty(); } + public static boolean isNullOrWildcard(String s) { + return s == null || isWildcard(s) || s.equals("%"); + } + /** * This function checks if the input string is a wildcard string * 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 08d35facc8..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 @@ -11,14 +11,17 @@ public class CommandConstants { public static final String IN_CATALOG_SQL = " IN CATALOG %s"; public static final String IN_ABSOLUTE_SCHEMA_SQL = " IN SCHEMA %s"; public static final String IN_ABSOLUTE_TABLE_SQL = " IN TABLE %s"; - public static final String SHOW_SCHEMA_IN_CATALOG_SQL = "SHOW SCHEMAS IN %s"; + public static final String IN_ALL_CATALOGS_SQL = " IN ALL CATALOGS"; + public static final String SHOW_SCHEMAS_IN_CATALOG_SQL = "SHOW SCHEMAS IN %s"; public static final String LIKE_SQL = " LIKE '%s'"; public static final String SCHEMA_LIKE_SQL = " SCHEMA" + LIKE_SQL; public static final String TABLE_LIKE_SQL = " TABLE" + LIKE_SQL; 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"; + 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_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 = 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 f27589336b..6db7d66644 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 @@ -68,7 +68,9 @@ public DatabricksResultSet getCatalogsResult(DatabricksResultSet resultSet) thro public DatabricksResultSet getSchemasResult(DatabricksResultSet resultSet, String catalog) throws SQLException { - List> rows = getRowsForSchemas(resultSet, SCHEMA_COLUMNS, catalog); + List> rows = + getRowsForSchemas( + resultSet, SCHEMA_COLUMNS, catalog, new SchemasDatabricksResultSetAdapter()); return buildResultSet( SCHEMA_COLUMNS, rows, @@ -562,19 +564,39 @@ private List> getRowsForFunctions( } private List> getRowsForSchemas( - DatabricksResultSet resultSet, List columns, String catalog) + DatabricksResultSet resultSet, + List columns, + String catalog, + IDatabricksResultSetAdapter adapter) throws SQLException { List> rows = new ArrayList<>(); while (resultSet.next()) { + // Check if this row should be included based on the adapter's filter + if (!adapter.includeRow(resultSet, columns)) { + continue; + } + List row = new ArrayList<>(); for (ResultColumn column : columns) { - if (column.getColumnName().equals("TABLE_CATALOG")) { - row.add(catalog); - continue; + // Map the expected column on client to column in the result set using the adapter + ResultColumn mappedColumn = adapter.mapColumn(column); + + if (mappedColumn + .getResultSetColumnName() + .equals(CATALOG_RESULT_COLUMN.getResultSetColumnName())) { + try { + resultSet.findColumn(mappedColumn.getResultSetColumnName()); + } catch (SQLException e) { + // Result set does not have a catalog column + // Manually add the catalog and move to next column + row.add(catalog); + continue; + } } + Object object; try { - object = resultSet.getObject(column.getResultSetColumnName()); + object = resultSet.getObject(mappedColumn.getResultSetColumnName()); if (object == null) { object = NULL_STRING; } diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/SchemasDatabricksResultSetAdapter.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/SchemasDatabricksResultSetAdapter.java new file mode 100644 index 0000000000..79a0a0b20b --- /dev/null +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/SchemasDatabricksResultSetAdapter.java @@ -0,0 +1,27 @@ +package com.databricks.jdbc.dbclient.impl.common; + +import static com.databricks.jdbc.common.MetadataResultConstants.CATALOG_FULL_COLUMN; +import static com.databricks.jdbc.common.MetadataResultConstants.CATALOG_RESULT_COLUMN; + +import com.databricks.jdbc.model.core.ResultColumn; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.List; + +public class SchemasDatabricksResultSetAdapter implements IDatabricksResultSetAdapter { + @Override + public ResultColumn mapColumn(ResultColumn column) { + String columnName = column.getResultSetColumnName(); + if (columnName.equals(CATALOG_FULL_COLUMN.getResultSetColumnName())) { + // Map CATALOG_FULL_COLUMN to CATALOG_COLUMN_FOR_GET_CATALOGS of result set + return CATALOG_RESULT_COLUMN; + } else { + return column; + } + } + + @Override + public boolean includeRow(ResultSet resultSet, List columns) throws SQLException { + return true; + } +} 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 95590ce7f2..a1ffba7db8 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 @@ -75,27 +75,28 @@ private String fetchCatalogSQL() { } private String fetchSchemaSQL() throws SQLException { - String contextString = - String.format( - "Building command for fetching schema. Catalog %s, SchemaPattern %s and session context %s", - catalogName, schemaPattern, sessionContext); - LOGGER.debug(contextString); - throwErrorIfNull(Collections.singletonMap(CATALOG, catalogName), contextString); - String showSchemaSQL = String.format(SHOW_SCHEMA_IN_CATALOG_SQL, catalogName); + LOGGER.debug( + "Building command for fetching schema. Catalog %s, SchemaPattern %s and session context %s", + catalogName, schemaPattern, sessionContext); + String showSchemasSQL; + if (WildcardUtil.isNullOrWildcard(catalogName)) { + // SHOW SCHEMAS IN ALL CATALOGS + showSchemasSQL = SHOW_SCHEMAS_IN_ALL_CATALOGS_SQL; + } else { + showSchemasSQL = String.format(SHOW_SCHEMAS_IN_CATALOG_SQL, catalogName); + } if (!WildcardUtil.isNullOrEmpty(schemaPattern)) { - showSchemaSQL += String.format(LIKE_SQL, schemaPattern); + showSchemasSQL += String.format(LIKE_SQL, schemaPattern); } - return showSchemaSQL; + return showSchemasSQL; } private String fetchTablesSQL() throws SQLException { - String contextString = - String.format( - "Building command for fetching tables. Catalog %s, SchemaPattern %s, TablePattern %s and session context %s", - catalogName, schemaPattern, tablePattern, sessionContext); - LOGGER.debug(contextString); + LOGGER.debug( + "Building command for fetching tables. Catalog %s, SchemaPattern %s, TablePattern %s and session context %s", + catalogName, schemaPattern, tablePattern, sessionContext); String showTablesSQL; - if (catalogName == null || catalogName.equals("*") || catalogName.equals("%")) { + if (WildcardUtil.isNullOrWildcard(catalogName)) { // SHOW TABLES IN ALL CATALOGS showTablesSQL = SHOW_TABLES_IN_ALL_CATALOGS_SQL; } else { diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataSdkClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataSdkClient.java index 7cbb486b16..a68aabb60d 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataSdkClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataSdkClient.java @@ -1,7 +1,6 @@ package com.databricks.jdbc.dbclient.impl.sqlexec; -import static com.databricks.jdbc.common.MetadataResultConstants.DEFAULT_TABLE_TYPES; -import static com.databricks.jdbc.common.MetadataResultConstants.PARSE_SYNTAX_ERROR_SQL_STATE; +import static com.databricks.jdbc.common.MetadataResultConstants.*; import static com.databricks.jdbc.dbclient.impl.common.CommandConstants.GET_TABLES_STATEMENT_ID; import static com.databricks.jdbc.dbclient.impl.common.CommandConstants.METADATA_STATEMENT_ID; import static com.databricks.jdbc.dbclient.impl.sqlexec.ResultConstants.TYPE_INFO_RESULT; @@ -10,21 +9,31 @@ import com.databricks.jdbc.api.internal.IDatabricksSession; import com.databricks.jdbc.common.MetadataResultConstants; import com.databricks.jdbc.common.StatementType; +import com.databricks.jdbc.common.util.JdbcThreadUtils; +import com.databricks.jdbc.common.util.WildcardUtil; import com.databricks.jdbc.dbclient.IDatabricksClient; import com.databricks.jdbc.dbclient.IDatabricksMetadataClient; import com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; +import java.sql.ResultSet; import java.sql.SQLException; import java.util.ArrayList; import java.util.HashMap; +import java.util.List; import java.util.Optional; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; /** Implementation for {@link IDatabricksMetadataClient} using {@link IDatabricksClient}. */ public class DatabricksMetadataSdkClient implements IDatabricksMetadataClient { private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(DatabricksMetadataSdkClient.class); + private static final int DEFAULT_MAX_THREADS_FETCH_SCHEMAS = 10; + private static final int TASK_TIMEOUT_FETCH_SCHEMAS_SEC = 90; + private static final Object THREAD_POOL_LOCK = new Object(); + private static ExecutorService schemasThreadPool = null; private final IDatabricksClient sdkClient; private final MetadataResultSetBuilder metadataResultSetBuilder; @@ -54,7 +63,20 @@ public DatabricksResultSet listSchemas( new CommandBuilder(catalog, session).setSchemaPattern(schemaNamePattern); String SQL = commandBuilder.getSQLString(CommandName.LIST_SCHEMAS); LOGGER.debug("SQL command to fetch schemas: {}", SQL); - return metadataResultSetBuilder.getSchemasResult(getResultSet(SQL, session), catalog); + try { + return metadataResultSetBuilder.getSchemasResult(getResultSet(SQL, session), catalog); + } catch (SQLException e) { + if (WildcardUtil.isNullOrWildcard(catalog) + && e.getSQLState().equals(PARSE_SYNTAX_ERROR_SQL_STATE)) { + // This is a fallback for the case where the SQL command fails with "syntax error at or near + // "ALL CATALOGS"" + // This is a known issue for older DBR versions + LOGGER.debug("SQL command failed with syntax error. Fetching schemas across all catalogs."); + return fetchSchemasAcrossCatalogs(session, schemaNamePattern); + } else { + throw e; + } + } } @Override @@ -159,7 +181,6 @@ public DatabricksResultSet listImportedKeys( if (e.getSQLState().equals(PARSE_SYNTAX_ERROR_SQL_STATE)) { // This is a workaround for the issue where the SQL command fails with "syntax error at or // near "foreign"" - // This is a known issue in Databricks for older DBSQL versions LOGGER.debug("SQL command failed with syntax error. Returning empty result set."); return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns( MetadataResultConstants.IMPORTED_KEYS_COLUMNS, @@ -233,4 +254,66 @@ private DatabricksResultSet getResultSet(String SQL, IDatabricksSession session) session, null /* parentStatement */); } + + private DatabricksResultSet fetchSchemasAcrossCatalogs( + IDatabricksSession session, String schemaPattern) throws SQLException { + List catalogList = new ArrayList<>(); + try (ResultSet catalogs = session.getDatabricksMetadataClient().listCatalogs(session)) { + while (catalogs.next()) { + String c = catalogs.getString(1); + if (c != null && !c.isEmpty()) { + catalogList.add(c); + } + } + } + + // Process catalogs in parallel, gathering schema information + List> schemaRows = + JdbcThreadUtils.parallelFlatMap( + catalogList, + session.getConnectionContext(), + DEFAULT_MAX_THREADS_FETCH_SCHEMAS, // Not significant since the executor is provided as + // a parameter + TASK_TIMEOUT_FETCH_SCHEMAS_SEC, + c -> { + List> rows = new ArrayList<>(); + try (ResultSet catalogSchemas = + session.getDatabricksMetadataClient().listSchemas(session, c, schemaPattern)) { + while (catalogSchemas.next()) { + List schemaRow = new ArrayList<>(); + schemaRow.add(catalogSchemas.getString(1)); // TABLE_SCHEM + schemaRow.add(catalogSchemas.getString(2)); // TABLE_CATALOG + rows.add(schemaRow); + } + } catch (SQLException e) { + LOGGER.warn("Error fetching schemas for catalog %s %s", c, e.getMessage()); + } + return rows; + }, + getOrCreateSchemasThreadPool()); + + // Convert combined data into a result set + return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns( + SCHEMA_COLUMNS, + schemaRows, + METADATA_STATEMENT_ID, + com.databricks.jdbc.common.CommandName.LIST_SCHEMAS); + } + + public static ExecutorService getOrCreateSchemasThreadPool() { + synchronized (THREAD_POOL_LOCK) { + if (schemasThreadPool == null || schemasThreadPool.isShutdown()) { + // Could read max threads from a configuration property + schemasThreadPool = + Executors.newFixedThreadPool( + DEFAULT_MAX_THREADS_FETCH_SCHEMAS, + r -> { + Thread t = new Thread(r, "jdbc-schemas-fetcher"); + t.setDaemon(true); + return t; + }); + } + return schemasThreadPool; + } + } } 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 e651f14826..ea0497478b 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaDataTest.java @@ -683,14 +683,6 @@ public void testGetSchemas_SqlExec() throws SQLException { DatabricksConnectionContext.parse(WAREHOUSE_JDBC_URL_WITH_SEA, new Properties())); ResultSet resultSet = metaData.getSchemas(); assertNotNull(resultSet); - - // Result set should have 2 columns; TABLE_SCHEM and TABLE_CATALOG - assertEquals(2, resultSet.getMetaData().getColumnCount()); - assertSame("TABLE_SCHEM", resultSet.getMetaData().getColumnName(1)); - assertSame("TABLE_CATALOG", resultSet.getMetaData().getColumnName(2)); - - // For SQL_EXEC execution, result set data should be empty - assertFalse(resultSet.next()); } @Test diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/SchemasDatabricksResultSetAdapterTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/SchemasDatabricksResultSetAdapterTest.java new file mode 100644 index 0000000000..25503c0003 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/SchemasDatabricksResultSetAdapterTest.java @@ -0,0 +1,58 @@ +package com.databricks.jdbc.dbclient.impl.common; + +import static com.databricks.jdbc.common.MetadataResultConstants.CATALOG_FULL_COLUMN; +import static com.databricks.jdbc.common.MetadataResultConstants.CATALOG_RESULT_COLUMN; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.databricks.jdbc.model.core.ResultColumn; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.Arrays; +import java.util.List; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class SchemasDatabricksResultSetAdapterTest { + + private SchemasDatabricksResultSetAdapter adapter; + private ResultSet mockResultSet; + + @BeforeEach + void setUp() { + adapter = new SchemasDatabricksResultSetAdapter(); + mockResultSet = mock(ResultSet.class); + } + + @Test + void mapColumn_shouldMapCatalogFullColumnToCatalogColumnForGetCatalogs() { + ResultColumn inputColumn = mock(ResultColumn.class); + when(inputColumn.getResultSetColumnName()) + .thenReturn(CATALOG_FULL_COLUMN.getResultSetColumnName()); + + ResultColumn result = adapter.mapColumn(inputColumn); + + assertSame(CATALOG_RESULT_COLUMN, result); + } + + @Test + void mapColumn_shouldReturnSameColumnForNonCatalogFullColumn() { + ResultColumn inputColumn = mock(ResultColumn.class); + when(inputColumn.getResultSetColumnName()).thenReturn("some_other_column"); + + ResultColumn result = adapter.mapColumn(inputColumn); + + assertSame(inputColumn, result); + } + + @Test + void includeRow_shouldAlwaysReturnTrue() throws SQLException { + List columns = Arrays.asList(mock(ResultColumn.class), mock(ResultColumn.class)); + + boolean result = adapter.includeRow(mockResultSet, columns); + + assertTrue(result); + } +} 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 8e2dd1a090..4130df7543 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 @@ -147,11 +147,6 @@ void shouldGenerateCorrectSqlForTablesWithWildcardCatalog() throws SQLException CommandBuilder builder1 = new CommandBuilder("*", mockSession); String sql1 = builder1.getSQLString(CommandName.LIST_TABLES); assertEquals(SHOW_TABLES_IN_ALL_CATALOGS_SQL, sql1); - - // Test with '%' wildcard - CommandBuilder builder2 = new CommandBuilder("%", mockSession); - String sql2 = builder2.getSQLString(CommandName.LIST_TABLES); - assertEquals(SHOW_TABLES_IN_ALL_CATALOGS_SQL, sql2); } } 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 8d8d1d5090..647f71f93a 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 @@ -183,9 +183,7 @@ void testListCatalogs() throws SQLException { setupCatalogMocks(); DatabricksMetadataSdkClient metadataClient = new DatabricksMetadataSdkClient(mockClient); doReturn(1).when(mockedMetaData).getColumnCount(); - doReturn(CATALOG_COLUMN_FOR_GET_CATALOGS.getResultSetColumnName()) - .when(mockedMetaData) - .getColumnName(1); + doReturn(CATALOG_RESULT_COLUMN.getResultSetColumnName()).when(mockedMetaData).getColumnName(1); doReturn(255).when(mockedMetaData).getPrecision(1); doReturn(0).when(mockedMetaData).getScale(1); when(mockedCatalogResultSet.getMetaData()).thenReturn(mockedMetaData); @@ -388,7 +386,7 @@ void testListSchemas(String sqlStatement, String schema, String description) thr when(mockClient.executeStatement( sqlStatement, mockedComputeResource, - new HashMap(), + new HashMap<>(), StatementType.METADATA, session, null)) @@ -399,6 +397,8 @@ void testListSchemas(String sqlStatement, String schema, String description) thr doReturn(SCHEMA_COLUMN.getResultSetColumnName()).when(mockedMetaData).getColumnName(1); doReturn(CATALOG_COLUMN.getResultSetColumnName()).when(mockedMetaData).getColumnName(2); when(mockedResultSet.getMetaData()).thenReturn(mockedMetaData); + when(mockedResultSet.findColumn(CATALOG_RESULT_COLUMN.getResultSetColumnName())) + .thenThrow(DatabricksSQLException.class); DatabricksResultSet actualResult = metadataClient.listSchemas(session, TEST_CATALOG, schema); assertEquals( actualResult.getStatementStatus().getState(), StatementState.SUCCEEDED, description); @@ -407,6 +407,40 @@ void testListSchemas(String sqlStatement, String schema, String description) thr ((DatabricksResultSetMetaData) actualResult.getMetaData()).getTotalRows(), 1, description); } + @Test + void testListSchemasNullCatalog() throws SQLException { + when(session.getComputeResource()).thenReturn(mockedComputeResource); + DatabricksMetadataSdkClient metadataClient = new DatabricksMetadataSdkClient(mockClient); + when(mockClient.executeStatement( + "SHOW SCHEMAS IN ALL CATALOGS LIKE 'a*'", + mockedComputeResource, + new HashMap<>(), + StatementType.METADATA, + session, + null)) + .thenReturn(mockedResultSet); + when(mockedResultSet.next()).thenReturn(true, false); + when(mockedResultSet.getObject("databaseName")).thenReturn(TEST_COLUMN); + when(mockedResultSet.getObject("catalog")).thenReturn(TEST_CATALOG); + doReturn(2).when(mockedMetaData).getColumnCount(); + doReturn(SCHEMA_COLUMN.getResultSetColumnName()).when(mockedMetaData).getColumnName(1); + doReturn(CATALOG_COLUMN.getResultSetColumnName()).when(mockedMetaData).getColumnName(2); + when(mockedResultSet.getMetaData()).thenReturn(mockedMetaData); + when(mockedResultSet.findColumn(CATALOG_RESULT_COLUMN.getResultSetColumnName())).thenReturn(2); + DatabricksResultSet actualResult = metadataClient.listSchemas(session, null, "a*"); + assertEquals(actualResult.getStatementStatus().getState(), StatementState.SUCCEEDED); + assertEquals(actualResult.getStatementId(), METADATA_STATEMENT_ID); + assertEquals(((DatabricksResultSetMetaData) actualResult.getMetaData()).getTotalRows(), 1); + + // Check the first row of the result set + assertTrue(actualResult.next()); + assertEquals(TEST_COLUMN, actualResult.getObject(1)); + assertEquals(TEST_CATALOG, actualResult.getObject(2)); + + // Check that the result set is empty after the first row + assertFalse(actualResult.next()); + } + @Test void testListPrimaryKeys() throws SQLException { when(session.getComputeResource()).thenReturn(WAREHOUSE_COMPUTE); @@ -748,9 +782,6 @@ void testThrowsErrorResultInCaseOfNullCatalog() { assertThrows( DatabricksValidationException.class, () -> metadataClient.listColumns(session, null, TEST_SCHEMA, TEST_TABLE, TEST_COLUMN)); - assertThrows( - DatabricksValidationException.class, - () -> metadataClient.listSchemas(session, null, TEST_SCHEMA)); assertThrows( DatabricksValidationException.class, () -> metadataClient.listPrimaryKeys(session, null, TEST_SCHEMA, TEST_TABLE));