merge

Author: madhav-db

.github/workflows/prCheck.yml

@@ -66,13 +66,29 @@ jobs:
       - name: Check Unit Tests
         run: mvn test
 
+      - name: Install xmllint
+        if: runner.os == 'Linux'
+        run: sudo apt-get update && sudo apt-get install -y libxml2-utils
+
       - name: JaCoCo report
         run: mvn --batch-mode --errors jacoco:report --file pom.xml
 
-      - name: Publish test coverage
-        uses: codecov/codecov-action@v4
-        with:
-          file: target/site/jacoco/jacoco.xml
-          token: ${{ secrets.CODECOV_TOKEN }}
-          fail_ci_if_error: true
-          verbose: true
+      - name: Extract codeCov percentage
+        run: |
+          COVERAGE_FILE="target/site/jacoco/jacoco.xml"
+          COVERED=$(xmllint --xpath "string(//report/counter[@type='INSTRUCTION']/@covered)" $COVERAGE_FILE)
+          MISSED=$(xmllint --xpath "string(//report/counter[@type='INSTRUCTION']/@missed)" $COVERAGE_FILE)
+          TOTAL=$(($COVERED + $MISSED))
+          PERCENTAGE=$(echo "scale=2; 100 * $COVERED / $TOTAL" | bc)
+          echo $PERCENTAGE > coverage.txt
+          echo "::set-output name=coverage::$PERCENTAGE"
+      - name: Check coverage percentage
+        run: |
+          BRANCH_COVERAGE=$(cat coverage.txt)
+          echo "Branch Coverage: $BRANCH_COVERAGE%"
+          if (( $(echo "$BRANCH_COVERAGE < 85" | bc -l) )); then
+            echo "Coverage is less than 85%"
+            exit 1
+          else
+            echo "Coverage is equal to or greater than 85%"
+          fi

pom.xml

@@ -45,6 +45,7 @@
     <annotation.version>1.3.5</annotation.version>
     <slt.executor>dbsql</slt.executor>
     <slt.token>dummy-token</slt.token>
+    <wiremock.version>3.5.4</wiremock.version>
   </properties>
   <dependencies>
     <dependency>
@@ -168,6 +169,18 @@
       <artifactId>jakarta.annotation-api</artifactId>
       <version>${annotation.version}</version>
     </dependency>
+    <dependency>
+      <groupId>org.wiremock</groupId>
+      <artifactId>wiremock</artifactId>
+      <version>${wiremock.version}</version>
+      <scope>test</scope>
+      <exclusions>
+        <exclusion>
+          <groupId>commons-fileupload</groupId>
+          <artifactId>commons-fileupload</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
   </dependencies>
 
   <build>

src/main/java/com/databricks/jdbc/client/http/DatabricksHttpClient.java

@@ -1,5 +1,8 @@
 package com.databricks.jdbc.client.http;
 
+import static com.databricks.jdbc.driver.DatabricksJdbcConstants.FAKE_SERVICE_URI_PROP_SUFFIX;
+import static com.databricks.jdbc.driver.DatabricksJdbcConstants.IS_FAKE_SERVICE_TEST_PROP;
+
 import com.databricks.jdbc.client.DatabricksHttpException;
 import com.databricks.jdbc.client.IDatabricksHttpClient;
 import com.databricks.jdbc.driver.IDatabricksConnectionContext;
@@ -21,10 +24,13 @@
 import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpUriRequest;
 import org.apache.http.client.protocol.HttpClientContext;
+import org.apache.http.conn.UnsupportedSchemeException;
+import org.apache.http.conn.routing.HttpRoute;
 import org.apache.http.impl.client.BasicCredentialsProvider;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClientBuilder;
 import org.apache.http.impl.client.ProxyAuthenticationStrategy;
+import org.apache.http.impl.conn.DefaultSchemePortResolver;
 import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
 import org.apache.http.protocol.HttpContext;
 import org.slf4j.Logger;
@@ -45,6 +51,8 @@ public class DatabricksHttpClient implements IDatabricksHttpClient {
   private static final int MAX_RETRY_INTERVAL = 10 * 1000; // 10s
   private static final int DEFAULT_RETRY_COUNT = 5;
   private static final String HTTP_GET = "GET";
+  private static final String SDK_USER_AGENT = "databricks-sdk-java";
+  private static final String JDBC_HTTP_USER_AGENT = "databricks-jdbc-http";
   private static final Set<Integer> RETRYABLE_HTTP_CODES = getRetryableHttpCodes();
   protected static final long DEFAULT_IDLE_CONNECTION_TIMEOUT = 5;
 
@@ -153,7 +161,10 @@ public void process(HttpResponse httpResponse, HttpContext httpContext)
           connectionContext.getUseProxyAuth(),
           connectionContext.getProxyUser(),
           connectionContext.getProxyPassword());
+    } else if (Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP))) {
+      setFakeServiceRouteInHttpClient(builder);
     }
+
     return builder.build();
   }
 
@@ -177,6 +188,27 @@ public static void setProxyDetailsInHttpClient(
     }
   }
 
+  @VisibleForTesting
+  static void setFakeServiceRouteInHttpClient(HttpClientBuilder builder) {
+    builder.setRoutePlanner(
+        (host, request, context) -> {
+          // Get the fake service URI for the target URI and set it as proxy
+          final HttpHost proxy =
+              HttpHost.create(System.getProperty(host.toURI() + FAKE_SERVICE_URI_PROP_SUFFIX));
+          final HttpHost target;
+          try {
+            target =
+                new HttpHost(
+                    host.getHostName(),
+                    DefaultSchemePortResolver.INSTANCE.resolve(host),
+                    host.getSchemeName());
+          } catch (UnsupportedSchemeException e) {
+            throw new HttpException(e.getMessage());
+          }
+          return new HttpRoute(target, null, proxy, false);
+        });
+  }
+
   @VisibleForTesting
   static boolean isRetryAllowed(String method) {
     // For now, allowing retry only for GET which is idempotent
@@ -233,13 +265,28 @@ static String getUserAgent() {
     // Remove the SDK part from user agent
     StringBuilder mergedString = new StringBuilder();
     for (int i = 0; i < parts.length; i++) {
-      if (i != 1) { // Skip the second part
+      if (parts[i].startsWith(SDK_USER_AGENT)) {
+        mergedString.append(JDBC_HTTP_USER_AGENT);
+      } else {
         mergedString.append(parts[i]);
-        if (i != parts.length - 1) {
-          mergedString.append(" "); // Add space between parts
-        }
+      }
+      if (i != parts.length - 1) {
+        mergedString.append(" "); // Add space between parts
       }
     }
     return mergedString.toString();
   }
+
+  /** Reset the instance of the http client. This is used for testing purposes only. */
+  @VisibleForTesting
+  public static synchronized void resetInstance() {
+    if (instance != null) {
+      try {
+        instance.httpClient.close();
+      } catch (IOException e) {
+        LOGGER.error("Caught error while closing http client", e);
+      }
+      instance = null;
+    }
+  }
 }

src/main/java/com/databricks/jdbc/client/impl/helper/CommandConstants.java

@@ -13,7 +13,7 @@ public class CommandConstants {
   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 LIKE_SQL = " LIKE `%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;

src/main/java/com/databricks/jdbc/client/impl/helper/MetadataResultConstants.java

@@ -5,8 +5,12 @@
 import java.util.List;
 
 public class MetadataResultConstants {
+  public static final String[] DEFAULT_TABLE_TYPES = {"TABLE", "VIEW", "SYSTEM TABLE"};
   private static final ResultColumn CATALOG_COLUMN =
       new ResultColumn("TABLE_CAT", "catalogName", Types.VARCHAR);
+
+  private static final ResultColumn CATALOG_COLUMN_FOR_GET_CATALOGS =
+      new ResultColumn("TABLE_CAT", "catalog", Types.VARCHAR);
   private static final ResultColumn TYPE_CATALOG_COLUMN =
       new ResultColumn("TYPE_CAT", "TYPE_CATALOG_COLUMN", Types.VARCHAR);
   private static final ResultColumn TYPE_SCHEMA_COLUMN =
@@ -25,6 +29,10 @@ public class MetadataResultConstants {
       new ResultColumn("TYPE_NAME", "TYPE_NAME", Types.VARCHAR);
   private static final ResultColumn SCHEMA_COLUMN =
       new ResultColumn("TABLE_SCHEM", "namespace", Types.VARCHAR);
+
+  private static final ResultColumn SCHEMA_COLUMN_FOR_GET_SCHEMA =
+      new ResultColumn("TABLE_SCHEM", "databaseName", Types.VARCHAR);
+
   private static final ResultColumn TABLE_NAME_COLUMN =
       new ResultColumn("TABLE_NAME", "tableName", Types.VARCHAR);
   private static final ResultColumn TABLE_TYPE_COLUMN =
@@ -132,8 +140,9 @@ public class MetadataResultConstants {
           ORDINAL_POSITION_COLUMN,
           IS_AUTO_INCREMENT_COLUMN,
           IS_GENERATED_COLUMN);
-  public static List<ResultColumn> CATALOG_COLUMNS = List.of(CATALOG_COLUMN);
-  public static List<ResultColumn> SCHEMA_COLUMNS = List.of(SCHEMA_COLUMN, CATALOG_COLUMN);
+  public static List<ResultColumn> CATALOG_COLUMNS = List.of(CATALOG_COLUMN_FOR_GET_CATALOGS);
+  public static List<ResultColumn> SCHEMA_COLUMNS =
+      List.of(SCHEMA_COLUMN_FOR_GET_SCHEMA, CATALOG_COLUMN);
   public static List<ResultColumn> TABLE_COLUMNS =
       List.of(
           CATALOG_COLUMN,

src/main/java/com/databricks/jdbc/client/impl/helper/MetadataResultSetBuilder.java

@@ -5,6 +5,7 @@
 
 import com.databricks.jdbc.client.StatementType;
 import com.databricks.jdbc.core.DatabricksResultSet;
+import com.databricks.jdbc.core.DatabricksSQLException;
 import com.databricks.sdk.service.sql.StatementState;
 import com.databricks.sdk.service.sql.StatementStatus;
 import java.sql.ResultSet;
@@ -29,13 +30,19 @@ public static DatabricksResultSet getCatalogsResult(ResultSet resultSet) throws
     return buildResultSet(CATALOG_COLUMNS, rows, GET_CATALOGS_STATEMENT_ID);
   }
 
-  public static DatabricksResultSet getSchemasResult(ResultSet resultSet) throws SQLException {
-    List<List<Object>> rows = getRows(resultSet, SCHEMA_COLUMNS);
+  public static DatabricksResultSet getSchemasResult(ResultSet resultSet, String catalog)
+      throws SQLException {
+    List<List<Object>> rows = getRowsForSchemas(resultSet, SCHEMA_COLUMNS, catalog);
     return buildResultSet(SCHEMA_COLUMNS, rows, METADATA_STATEMENT_ID);
   }
 
-  public static DatabricksResultSet getTablesResult(ResultSet resultSet) throws SQLException {
-    List<List<Object>> rows = getRows(resultSet, TABLE_COLUMNS);
+  public static DatabricksResultSet getTablesResult(ResultSet resultSet, String[] tableTypes)
+      throws SQLException {
+    List<String> allowedTableTypes = List.of(tableTypes);
+    List<List<Object>> rows =
+        getRows(resultSet, TABLE_COLUMNS).stream()
+            .filter(row -> allowedTableTypes.contains(row.get(3))) // Filtering based on table type
+            .collect(Collectors.toList());
     return buildResultSet(TABLE_COLUMNS, rows, GET_TABLES_STATEMENT_ID);
   }
 
@@ -62,8 +69,42 @@ private static List<List<Object>> getRows(ResultSet resultSet, List<ResultColumn
     while (resultSet.next()) {
       List<Object> row = new ArrayList<>();
       for (ResultColumn column : columns) {
-        Object object = resultSet.getObject(column.getResultSetColumnName());
-        if (object == null) {
+        Object object;
+        try {
+          object = resultSet.getObject(column.getResultSetColumnName());
+          if (object == null) {
+            object = NULL_STRING;
+          }
+        } catch (DatabricksSQLException e) {
+          // Remove non-relevant columns from the obtained result set
+          object = NULL_STRING;
+        }
+        row.add(object);
+      }
+      rows.add(row);
+    }
+    return rows;
+  }
+
+  private static List<List<Object>> getRowsForSchemas(
+      ResultSet resultSet, List<ResultColumn> columns, String catalog) throws SQLException {
+    // TODO(PECO-1677): Remove this method once the server side ResultSet metadata contains catalogs
+    List<List<Object>> rows = new ArrayList<>();
+    while (resultSet.next()) {
+      List<Object> row = new ArrayList<>();
+      for (ResultColumn column : columns) {
+        if (column.getColumnName().equals("TABLE_CAT")) {
+          row.add(catalog);
+          continue;
+        }
+        Object object;
+        try {
+          object = resultSet.getObject(column.getResultSetColumnName());
+          if (object == null) {
+            object = NULL_STRING;
+          }
+        } catch (DatabricksSQLException e) {
+          // Remove non-relevant columns from the obtained result set
           object = NULL_STRING;
         }
         row.add(object);

src/main/java/com/databricks/jdbc/client/impl/sdk/DatabricksNewMetadataSdkClient.java

@@ -1,5 +1,6 @@
 package com.databricks.jdbc.client.impl.sdk;
 
+import static com.databricks.jdbc.client.impl.helper.MetadataResultConstants.DEFAULT_TABLE_TYPES;
 import static com.databricks.jdbc.client.impl.sdk.ResultConstants.TYPE_INFO_RESULT;
 
 import com.databricks.jdbc.client.DatabricksMetadataClient;
@@ -41,7 +42,8 @@ public DatabricksResultSet listCatalogs(IDatabricksSession session) throws SQLEx
     CommandBuilder commandBuilder = new CommandBuilder(session);
     String SQL = commandBuilder.getSQLString(CommandName.LIST_CATALOGS);
     LOGGER.debug("SQL command to fetch catalogs: {}", SQL);
-    return MetadataResultSetBuilder.getCatalogsResult(getResultSet(SQL, session));
+    return MetadataResultSetBuilder.getCatalogsResult(
+        getResultSet(SQL, session, StatementType.METADATA));
   }
 
   @Override
@@ -51,7 +53,8 @@ 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));
+    return MetadataResultSetBuilder.getSchemasResult(
+        getResultSet(SQL, session, StatementType.METADATA), catalog);
   }
 
   @Override
@@ -62,12 +65,17 @@ public DatabricksResultSet listTables(
       String tableNamePattern,
       String[] tableTypes)
       throws SQLException {
+    tableTypes =
+        Optional.ofNullable(tableTypes)
+            .filter(types -> types.length > 0)
+            .orElse(DEFAULT_TABLE_TYPES);
     CommandBuilder commandBuilder =
         new CommandBuilder(catalog, session)
             .setSchemaPattern(schemaNamePattern)
             .setTablePattern(tableNamePattern);
     String SQL = commandBuilder.getSQLString(CommandName.LIST_TABLES);
-    return MetadataResultSetBuilder.getTablesResult(getResultSet(SQL, session));
+    return MetadataResultSetBuilder.getTablesResult(
+        getResultSet(SQL, session, StatementType.METADATA), tableTypes);
   }
 
   @Override
@@ -90,7 +98,8 @@ public DatabricksResultSet listColumns(
             .setTablePattern(tableNamePattern)
             .setColumnPattern(columnNamePattern);
     String SQL = commandBuilder.getSQLString(CommandName.LIST_COLUMNS);
-    return MetadataResultSetBuilder.getColumnsResult(getResultSet(SQL, session));
+    return MetadataResultSetBuilder.getColumnsResult(
+        getResultSet(SQL, session, StatementType.QUERY));
   }
 
   @Override
@@ -106,7 +115,8 @@ public DatabricksResultSet listFunctions(
             .setFunctionPattern(functionNamePattern);
     String SQL = commandBuilder.getSQLString(CommandName.LIST_FUNCTIONS);
     LOGGER.debug("SQL command to fetch functions: {}", SQL);
-    return MetadataResultSetBuilder.getFunctionsResult(getResultSet(SQL, session));
+    return MetadataResultSetBuilder.getFunctionsResult(
+        getResultSet(SQL, session, StatementType.METADATA));
   }
 
   @Override
@@ -116,15 +126,17 @@ public DatabricksResultSet listPrimaryKeys(
         new CommandBuilder(catalog, session).setSchema(schema).setTable(table);
     String SQL = commandBuilder.getSQLString(CommandName.LIST_PRIMARY_KEYS);
     LOGGER.debug("SQL command to fetch primary keys: {}", SQL);
-    return MetadataResultSetBuilder.getPrimaryKeysResult(getResultSet(SQL, session));
+    return MetadataResultSetBuilder.getPrimaryKeysResult(
+        getResultSet(SQL, session, StatementType.METADATA));
   }
 
-  private ResultSet getResultSet(String SQL, IDatabricksSession session) throws SQLException {
+  private ResultSet getResultSet(
+      String SQL, IDatabricksSession session, StatementType statementType) throws SQLException {
     return sdkClient.executeStatement(
         SQL,
         session.getComputeResource(),
         new HashMap<Integer, ImmutableSqlParameter>(),
-        StatementType.METADATA,
+        statementType,
         session,
         null /* parentStatement */);
   }