Added support for enableMultipleCatalogSupport param (#1025)

## Description
<!-- Provide a brief summary of the changes made and the issue they aim
to address.-->
Added support for the connection param of enableMultipleCatalogSupport.
The default value is kept to 1.
The behaviour remains the same when the param is enabled. When the param
is disabled, the behaviour is as follows:
- getCatalogs() returns only the current catalog
- Metadata operations are done on the current catalog
- setCatalog() method does nothing (doesn't execute USE CATALOG)
- The SQL query of USE CATALOG can still be executed.

To maintain the current catalog, we make the query to get the current
catalog everytime when this param is disabled. 
 
## Testing
<!-- Describe how the changes have been tested-->

## Additional Notes to the Reviewer
<!-- Share any additional context or insights that may help the reviewer
understand the changes better. This could include challenges faced,
limitations, or compromises made during the development process.
Also, mention any areas of the code that you would like the reviewer to
focus on specifically. -->
[PECOBLR-938](https://databricks.atlassian.net/browse/PECOBLR-938)


[PECOBLR-938]:
https://databricks.atlassian.net/browse/PECOBLR-938?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

Author: msrathore-db

.github/workflows/runIntegrationTests.yml

@@ -53,4 +53,7 @@ jobs:
           DATABRICKS_JDBC_M2M_PRIVATE_KEY_AUTH_KID: ${{ secrets.DATABRICKS_JDBC_M2M_PRIVATE_KEY_AUTH_KID }}
           DATABRICKS_JDBC_M2M_PRIVATE_KEY_JWT_KEY_PASSPHRASE: ${{ secrets.DATABRICKS_JDBC_M2M_PRIVATE_KEY_JWT_KEY_PASSPHRASE }}
           DATABRICKS_JDBC_M2M_PRIVATE_KEY_CREDENTIALS_HTTP_PATH: ${{ secrets.DATABRICKS_JDBC_M2M_PRIVATE_KEY_CREDENTIALS_HTTP_PATH }}
-          DATABRICKS_JDBC_M2M_PRIVATE_KEY_CREDENTIALS_HOST: ${{ secrets.DATABRICKS_JDBC_M2M_PRIVATE_KEY_CREDENTIALS_HOST }}
+          DATABRICKS_JDBC_M2M_PRIVATE_KEY_CREDENTIALS_HOST: ${{ secrets.DATABRICKS_JDBC_M2M_PRIVATE_KEY_CREDENTIALS_HOST }}
+          DATABRICKS_BENCHFOOD_TOKEN: ${{ secrets.DATABRICKS_BENCHFOOD_TOKEN }}
+          DATABRICKS_BENCHFOOD_HTTP_PATH: ${{ secrets.DATABRICKS_BENCHFOOD_HTTP_PATH }}
+          DATABRICKS_BENCHFOOD_HOST: ${{ secrets.DATABRICKS_BENCHFOOD_HOST }}

NEXT_CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Added
 
+- Added `enableMultipleCatalogSupport` connection parameter to control catalog metadata behavior.
+
 ### Updated
 
 ### Fixed

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

@@ -199,6 +199,11 @@ public boolean isReadOnly() throws SQLException {
 
   @Override
   public void setCatalog(String catalog) throws SQLException {
+    // If enableMultipleCatalogSupport is disabled, setCatalog does nothing
+    if (!connectionContext.getEnableMultipleCatalogSupport()) {
+      LOGGER.debug("setCatalog ignored - enableMultipleCatalogSupport is disabled");
+      return;
+    }
     Statement statement = this.createStatement();
     statement.execute("SET CATALOG " + catalog);
     this.session.setCatalog(catalog);

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

@@ -236,6 +236,10 @@ public boolean getEnableSQLValidationForIsValid() {
         .equals("1");
   }
 
+  public boolean getEnableMultipleCatalogSupport() {
+    return getParameter(DatabricksJdbcUrlParams.ENABLE_MULTIPLE_CATALOG_SUPPORT, "1").equals("1");
+  }
+
   @Override
   public String getHostForOAuth() {
     return this.host;

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

@@ -9,6 +9,7 @@
 import com.databricks.jdbc.common.DatabricksClientType;
 import com.databricks.jdbc.common.DatabricksJdbcUrlParams;
 import com.databricks.jdbc.common.IDatabricksComputeResource;
+import com.databricks.jdbc.common.StatementType;
 import com.databricks.jdbc.dbclient.IDatabricksClient;
 import com.databricks.jdbc.dbclient.IDatabricksMetadataClient;
 import com.databricks.jdbc.dbclient.impl.sqlexec.DatabricksEmptyMetadataClient;
@@ -24,6 +25,7 @@
 import com.databricks.jdbc.telemetry.latency.DatabricksMetricsTimedProcessor;
 import com.databricks.sdk.support.ToStringer;
 import com.google.common.annotations.VisibleForTesting;
+import java.util.HashMap;
 import java.util.Map;
 import javax.annotation.Nullable;
 
@@ -277,6 +279,30 @@ public IDatabricksConnectionContext getConnectionContext() {
     return this.connectionContext;
   }
 
+  @Override
+  public String getCurrentCatalog() throws DatabricksSQLException {
+    try {
+      DatabricksResultSet resultSet =
+          databricksClient.executeStatement(
+              "SELECT CURRENT_CATALOG()",
+              this.computeResource,
+              new HashMap<>(),
+              StatementType.METADATA,
+              this,
+              null);
+
+      if (resultSet.next()) {
+        String currentCatalog = resultSet.getString(1);
+        return currentCatalog;
+      }
+    } catch (Exception e) {
+      LOGGER.warn(
+          "Failed to get current catalog from database, falling back to session catalog: {}",
+          e.getMessage());
+    }
+    return this.catalog;
+  }
+
   @Override
   public void setEmptyMetadataClient() {
     databricksMetadataClient = new DatabricksEmptyMetadataClient(connectionContext);

src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java

@@ -99,6 +99,9 @@ public interface IDatabricksConnectionContext {
   /** Returns the value of the EnableSQLValidationForIsValid connection property. */
   boolean getEnableSQLValidationForIsValid();
 
+  /** Returns the value of the enableMultipleCatalogSupport connection property. */
+  boolean getEnableMultipleCatalogSupport();
+
   String getProxyHost();
 
   int getProxyPort();

src/main/java/com/databricks/jdbc/api/internal/IDatabricksSession.java

@@ -84,6 +84,9 @@ public interface IDatabricksSession {
   /** Returns the associated connection context for the session */
   IDatabricksConnectionContext getConnectionContext();
 
+  /** Gets the current catalog from the database */
+  String getCurrentCatalog() throws DatabricksSQLException;
+
   void setEmptyMetadataClient();
 
   void forceClose();

src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java

@@ -164,8 +164,9 @@ public enum DatabricksJdbcUrlParams {
       "Enable SQL query execution for connection validation in isValid() method",
       "0"),
   IGNORE_TRANSACTIONS("IgnoreTransactions", "Ignore transaction-related method calls", "0"),
-  ENABLE_METRIC_VIEW_METADATA("EnableMetricViewMetadata", "Enable metric view metadata", "0");
-
+  ENABLE_METRIC_VIEW_METADATA("EnableMetricViewMetadata", "Enable metric view metadata", "0"),
+  ENABLE_MULTIPLE_CATALOG_SUPPORT(
+      "enableMultipleCatalogSupport", "Enable multiple catalog support", "1");
   private final String paramName;
   private final String defaultValue;
   private final String description;

src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksMetadataSdkClient.java

@@ -42,6 +42,19 @@ public DatabricksMetadataSdkClient(IDatabricksClient sdkClient) {
     this.metadataResultSetBuilder = new MetadataResultSetBuilder(sdkClient.getConnectionContext());
   }
 
+  private boolean isMultipleCatalogSupportDisabled() {
+    return sdkClient.getConnectionContext() != null
+        && !sdkClient.getConnectionContext().getEnableMultipleCatalogSupport();
+  }
+
+  private String autoFillCatalog(String catalog, IDatabricksSession session) throws SQLException {
+    if (isMultipleCatalogSupportDisabled()) {
+      String currentCatalog = session.getCurrentCatalog();
+      return (currentCatalog != null && !currentCatalog.isEmpty()) ? currentCatalog : "";
+    }
+    return catalog;
+  }
+
   @Override
   public DatabricksResultSet listTypeInfo(IDatabricksSession session) {
     LOGGER.debug("public ResultSet getTypeInfo()");
@@ -50,6 +63,19 @@ public DatabricksResultSet listTypeInfo(IDatabricksSession session) {
 
   @Override
   public DatabricksResultSet listCatalogs(IDatabricksSession session) throws SQLException {
+    // If multiple catalog support is disabled, return only the current catalog
+    if (isMultipleCatalogSupportDisabled()) {
+      String currentCatalog = session.getCurrentCatalog();
+      if (currentCatalog == null || currentCatalog.isEmpty()) {
+        currentCatalog = "";
+      }
+      List<List<Object>> singleCatalogRows = new ArrayList<>();
+      List<Object> catalogRow = new ArrayList<>();
+      catalogRow.add(currentCatalog);
+      singleCatalogRows.add(catalogRow);
+      return metadataResultSetBuilder.getCatalogsResult(singleCatalogRows);
+    }
+
     CommandBuilder commandBuilder = new CommandBuilder(session);
     String SQL = commandBuilder.getSQLString(CommandName.LIST_CATALOGS);
     LOGGER.debug("SQL command to fetch catalogs: {}", SQL);
@@ -59,6 +85,7 @@ public DatabricksResultSet listCatalogs(IDatabricksSession session) throws SQLEx
   @Override
   public DatabricksResultSet listSchemas(
       IDatabricksSession session, String catalog, String schemaNamePattern) throws SQLException {
+    catalog = autoFillCatalog(catalog, session);
     CommandBuilder commandBuilder =
         new CommandBuilder(catalog, session).setSchemaPattern(schemaNamePattern);
     String SQL = commandBuilder.getSQLString(CommandName.LIST_SCHEMAS);
@@ -87,6 +114,7 @@ public DatabricksResultSet listTables(
       String tableNamePattern,
       String[] tableTypes)
       throws SQLException {
+    catalog = autoFillCatalog(catalog, session);
     String[] validatedTableTypes =
         Optional.ofNullable(tableTypes)
             .filter(types -> types.length > 0)
@@ -132,6 +160,7 @@ public DatabricksResultSet listColumns(
       String tableNamePattern,
       String columnNamePattern)
       throws SQLException {
+    catalog = autoFillCatalog(catalog, session);
     CommandBuilder commandBuilder =
         new CommandBuilder(catalog, session)
             .setSchemaPattern(schemaNamePattern)
@@ -149,6 +178,7 @@ public DatabricksResultSet listFunctions(
       String schemaNamePattern,
       String functionNamePattern)
       throws SQLException {
+    catalog = autoFillCatalog(catalog, session);
     CommandBuilder commandBuilder =
         new CommandBuilder(catalog, session)
             .setSchemaPattern(schemaNamePattern)
@@ -161,6 +191,7 @@ public DatabricksResultSet listFunctions(
   @Override
   public DatabricksResultSet listPrimaryKeys(
       IDatabricksSession session, String catalog, String schema, String table) throws SQLException {
+    catalog = autoFillCatalog(catalog, session);
     CommandBuilder commandBuilder =
         new CommandBuilder(catalog, session).setSchema(schema).setTable(table);
     String SQL = commandBuilder.getSQLString(CommandName.LIST_PRIMARY_KEYS);
@@ -172,6 +203,7 @@ public DatabricksResultSet listPrimaryKeys(
   public DatabricksResultSet listImportedKeys(
       IDatabricksSession session, String catalog, String schema, String table) throws SQLException {
     LOGGER.debug("public ResultSet listImportedKeys() using SDK");
+    catalog = autoFillCatalog(catalog, session);
     CommandBuilder commandBuilder =
         new CommandBuilder(catalog, session).setSchema(schema).setTable(table);
     String SQL = commandBuilder.getSQLString(CommandName.LIST_FOREIGN_KEYS);
@@ -195,8 +227,9 @@ public DatabricksResultSet listImportedKeys(
 
   @Override
   public DatabricksResultSet listExportedKeys(
-      IDatabricksSession session, String catalog, String schema, String table) {
+      IDatabricksSession session, String catalog, String schema, String table) throws SQLException {
     LOGGER.debug("public ResultSet listExportedKeys() using SDK");
+    catalog = autoFillCatalog(catalog, session);
     // Exported keys not tracked in DBSQL. Returning empty result set
     return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
         MetadataResultConstants.EXPORTED_KEYS_COLUMNS,

src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java

@@ -66,6 +66,10 @@ void setServerProtocolVersion(TProtocolVersion serverProtocolVersion) {
     this.serverProtocolVersion = serverProtocolVersion;
   }
 
+  private boolean isMultipleCatalogSupportEnabled() {
+    return connectionContext == null || connectionContext.getEnableMultipleCatalogSupport();
+  }
+
   @Override
   public IDatabricksConnectionContext getConnectionContext() {
     return connectionContext;
@@ -91,7 +95,7 @@ public ImmutableSessionInfo createSession(
     TOpenSessionReq openSessionReq =
         new TOpenSessionReq()
             .setConfiguration(sessionConf)
-            .setCanUseMultipleCatalogs(true)
+            .setCanUseMultipleCatalogs(isMultipleCatalogSupportEnabled())
             .setClient_protocol_i64(JDBC_THRIFT_VERSION.getValue());
     if (catalog != null || schema != null) {
       openSessionReq.setInitialNamespace(getNamespace(catalog, schema));
@@ -331,6 +335,20 @@ public DatabricksResultSet listCatalogs(IDatabricksSession session) throws SQLEx
         String.format("Fetching catalogs using Thrift client. Session {%s}", session.toString());
     LOGGER.debug(context);
     DatabricksThreadContextHolder.setSessionId(session.getSessionId());
+
+    // If multiple catalog support is disabled, return only the current catalog
+    if (!isMultipleCatalogSupportEnabled()) {
+      String currentCatalog = session.getCurrentCatalog();
+      if (currentCatalog == null || currentCatalog.isEmpty()) {
+        currentCatalog = "";
+      }
+      List<List<Object>> singleCatalogRows = new ArrayList<>();
+      List<Object> catalogRow = new ArrayList<>();
+      catalogRow.add(currentCatalog);
+      singleCatalogRows.add(catalogRow);
+      return metadataResultSetBuilder.getCatalogsResult(singleCatalogRows);
+    }
+
     TGetCatalogsReq request =
         new TGetCatalogsReq()
             .setSessionHandle(Objects.requireNonNull(session.getSessionInfo()).sessionHandle());