[ES-1343645] Fix schema in connection string (#661)

Allow specifying schema in connection string. For example:
- `jdbc:databricks://...:443/<schema>;ConnSchema=<connSchema>;ConnCatalog=<connCatalog>`

`<schema>` will be prioritised over `ConnSchema=<connSchema>`. By default, the schema is `default`.

Author: jayantsing-db

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

@@ -65,7 +65,8 @@ public static IDatabricksConnectionContext parse(String url, Properties properti
     Matcher urlMatcher = JDBC_URL_PATTERN.matcher(url);
     if (urlMatcher.find()) {
       String hostUrlVal = urlMatcher.group(1);
-      String urlMinusHost = urlMatcher.group(2);
+      String schema = urlMatcher.group(2);
+      String urlMinusHost = urlMatcher.group(3);
       String[] hostAndPort = hostUrlVal.split(DatabricksJdbcConstants.PORT_DELIMITER);
       String hostValue = hostAndPort[0];
       int portValue =
@@ -75,12 +76,8 @@ public static IDatabricksConnectionContext parse(String url, Properties properti
 
       ImmutableMap.Builder<String, String> parametersBuilder = ImmutableMap.builder();
       String[] urlParts = urlMinusHost.split(DatabricksJdbcConstants.URL_DELIMITER);
-      String schema = urlParts[0];
-      if (nullOrEmptyString(schema)) {
-        schema = DEFAULT_SCHEMA;
-      }
-      for (int urlPartIndex = 1; urlPartIndex < urlParts.length; urlPartIndex++) {
-        String[] pair = urlParts[urlPartIndex].split(DatabricksJdbcConstants.PAIR_DELIMITER);
+      for (String urlPart : urlParts) {
+        String[] pair = urlPart.split(DatabricksJdbcConstants.PAIR_DELIMITER);
         if (pair.length == 1) {
           pair = new String[] {pair[0], ""};
         }
@@ -359,6 +356,10 @@ public String getCatalog() {
 
   @Override
   public String getSchema() {
+    if (!nullOrEmptyString(schema)) {
+      return schema;
+    }
+
     return getParameter(
         DatabricksJdbcUrlParams.CONN_SCHEMA, getParameter(DatabricksJdbcUrlParams.SCHEMA));
   }

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

@@ -7,7 +7,14 @@
 
 public final class DatabricksJdbcConstants {
   public static final Pattern JDBC_URL_PATTERN =
-      Pattern.compile("jdbc:databricks://([^/;]*)(?::\\d+)?/*(.*)");
+      Pattern.compile(
+          "jdbc:databricks://"
+              + // Protocol prefix
+              "([^/;]+)"
+              + // Host[:Port] (captured)
+              "(?:/([^;]*))?"
+              + // Optional Schema (captured without /)
+              "(?:;(.*))?"); // Optional Property=Value pairs (captured without leading ;)
   public static final Pattern HTTP_WAREHOUSE_PATH_PATTERN = Pattern.compile(".*/warehouses/(.+)");
   public static final Pattern HTTP_ENDPOINT_PATH_PATTERN = Pattern.compile(".*/endpoints/(.+)");
   public static final Pattern HTTP_CLI_PATTERN = Pattern.compile(".*cliservice(.+)");

src/test/java/com/databricks/jdbc/TestConstants.java

@@ -129,7 +129,7 @@ public class TestConstants {
   public static final String GCP_TEST_URL =
       "jdbc:databricks://4371047901336987.7.gcp.databricks.com:443/default;transportMode=http;AuthMech=11;Auth_Flow=1;httpPath=/sql/1.0/warehouses/dd5955aacf3f09e5;[email protected];";
   public static final String VALID_URL_1 =
-      "jdbc:databricks://adb-565757575.18.azuredatabricks.net:4423/default;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/erg6767gg;LogLevel=debug;LogPath=./test1;auth_flow=2";
+      "jdbc:databricks://adb-565757575.18.azuredatabricks.net:4423;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/erg6767gg;LogLevel=debug;LogPath=./test1;auth_flow=2";
   public static final String VALID_URL_2 =
       "jdbc:databricks://adb-565656.azuredatabricks.net/default;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/fgff575757;LogLevel=invalid;EnableQueryResultLZ4Compression=1;UseThriftClient=0";
   public static final String VALID_URL_3 =

src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionTest.java

@@ -43,8 +43,8 @@ public class DatabricksConnectionTest {
       "jdbc:databricks://adb-565757575.18.azuredatabricks.net:4423/default;transportMode=http;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/erg6767gg;UserAgentEntry=MyApp";
   private static final String CATALOG_SCHEMA_JDBC_URL =
       String.format(
-          "jdbc:databricks://adb-565757575.18.azuredatabricks.net:4423/default;transportMode=http;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/erg6767gg;ConnCatalog=%s;ConnSchema=%s;logLevel=FATAL",
-          CATALOG, SCHEMA);
+          "jdbc:databricks://adb-565757575.18.azuredatabricks.net:4423/%s;transportMode=http;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/erg6767gg;ConnCatalog=%s;ConnSchema=%s;logLevel=FATAL",
+          SCHEMA, CATALOG, SCHEMA);
   private static final String SESSION_CONF_JDBC_URL =
       String.format(
           "jdbc:databricks://adb-565757575.18.azuredatabricks.net:4423/default;transportMode=http;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/erg6767gg;%s",
@@ -145,7 +145,7 @@ public void testConfInConnection() throws SQLException {
         SESSION_CONFIGS.entrySet().stream()
             .collect(Collectors.toMap(e -> e.getKey().toLowerCase(), Map.Entry::getValue));
     when(databricksClient.createSession(
-            new Warehouse(WAREHOUSE_ID), null, null, lowercaseSessionConfigs))
+            new Warehouse(WAREHOUSE_ID), null, "default", lowercaseSessionConfigs))
         .thenReturn(IMMUTABLE_SESSION_INFO);
     IDatabricksConnectionContext connectionContext =
         DatabricksConnectionContext.parse(SESSION_CONF_JDBC_URL, new Properties());
@@ -215,7 +215,7 @@ public void testSetClientInfo() throws SQLException {
         DatabricksConnectionContext.parse(JDBC_URL, new Properties());
     ImmutableSessionInfo session =
         ImmutableSessionInfo.builder().computeResource(warehouse).sessionId(SESSION_ID).build();
-    when(databricksClient.createSession(warehouse, null, null, new HashMap<>()))
+    when(databricksClient.createSession(warehouse, null, "default", new HashMap<>()))
         .thenReturn(session);
     DatabricksConnection connection =
         Mockito.spy(new DatabricksConnection(connectionContext, databricksClient));

src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java

@@ -29,7 +29,6 @@ public class DatabricksSessionTest {
   @Mock TSessionHandle tSessionHandle;
   private static final String JDBC_URL_INVALID =
       "jdbc:databricks://adb-565757575.18.azuredatabricks.net:4423/default;transportMode=http;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehou/erg6767gg;";
-  private static final String WAREHOUSE_ID = "erg6767gg";
   private static final String NEW_CATALOG = "new_catalog";
   private static final String NEW_SCHEMA = "new_schema";
   private static final String SESSION_ID = "session_id";
@@ -148,7 +147,7 @@ public void testSessionToString() throws DatabricksSQLException {
     setupWarehouse(false);
     DatabricksSession session = new DatabricksSession(connectionContext);
     assertEquals(
-        "DatabricksSession[compute='SQL Warehouse with warehouse ID {warehouse_id}']",
+        "DatabricksSession[compute='SQL Warehouse with warehouse ID {warehouse_id}', schema='default']",
         session.toString());
   }