Merge branch 'main' into abortstatement

Author: gopalldb

.github/workflows/prCheck.yml

@@ -64,7 +64,7 @@ jobs:
           restore-keys: ${{ runner.os }}-m2
 
       - name: Check Unit Tests
-        run: mvn test
+        run: mvn test -Dtest=!**/integration/**,!**/local/**
 
       - name: Install xmllint
         if: runner.os == 'Linux'

pom.xml

@@ -204,7 +204,7 @@
           <version>${maven-surefire-plugin.version}</version>
           <configuration>
             <excludes>
-              <exclude>**/**IntegrationTests.java</exclude>
+              <exclude>**/integration/**/*.java</exclude>
               <exclude>**/*MetadataBenchmarkingTest.java</exclude>
             </excludes>
             <argLine>

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

@@ -30,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);
   }
 
@@ -80,6 +86,34 @@ private static List<List<Object>> getRows(ResultSet resultSet, List<ResultColumn
     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);
+      }
+      rows.add(row);
+    }
+    return rows;
+  }
+
   private static DatabricksResultSet buildResultSet(
       List<ResultColumn> columns, List<List<Object>> rows, String statementId) {
     return new DatabricksResultSet(

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 */);
   }

src/main/java/com/databricks/jdbc/core/DatabricksResultSet.java

@@ -100,7 +100,9 @@ public DatabricksResultSet(
     this.executionResult =
         ExecutionResultFactory.getResultSet(resultData, resultManifest, statementId, session);
     int rowSize = getRowCount(resultData);
-    this.resultSetMetaData = new DatabricksResultSetMetaData(statementId, resultManifest, rowSize);
+    this.resultSetMetaData =
+        new DatabricksResultSetMetaData(
+            statementId, resultManifest, rowSize, resultData.getResultLinksSize());
     this.statementType = statementType;
     this.updateCount = null;
     this.parentStatement = parentStatement;

src/main/java/com/databricks/jdbc/core/DatabricksResultSetMetaData.java

@@ -30,6 +30,7 @@ public class DatabricksResultSetMetaData implements ResultSetMetaData {
   private final ImmutableList<ImmutableDatabricksColumn> columns;
   private final ImmutableMap<String, Integer> columnNameIndex;
   private final long totalRows;
+  private Long chunkCount;
   private static final String DEFAULT_CATALOGUE_NAME = "Spark";
   private static final String NULL_STRING = "null";
 
@@ -80,10 +81,11 @@ public DatabricksResultSetMetaData(
     this.columns = columnsBuilder.build();
     this.columnNameIndex = ImmutableMap.copyOf(columnNameToIndexMap);
     this.totalRows = resultManifest.getTotalRowCount();
+    this.chunkCount = resultManifest.getTotalChunkCount();
   }
 
   public DatabricksResultSetMetaData(
-      String statementId, TGetResultSetMetadataResp resultManifest, int rows) {
+      String statementId, TGetResultSetMetadataResp resultManifest, int rows, long chunkCount) {
     this.statementId = statementId;
     Map<String, Integer> columnNameToIndexMap = new HashMap<>();
     ImmutableList.Builder<ImmutableDatabricksColumn> columnsBuilder = ImmutableList.builder();
@@ -108,6 +110,7 @@ public DatabricksResultSetMetaData(
     this.columns = columnsBuilder.build();
     this.columnNameIndex = ImmutableMap.copyOf(columnNameToIndexMap);
     this.totalRows = rows;
+    this.chunkCount = chunkCount;
   }
 
   public DatabricksResultSetMetaData(
@@ -283,6 +286,10 @@ public long getTotalRows() {
     return totalRows;
   }
 
+  public Long getChunkCount() {
+    return chunkCount;
+  }
+
   private ImmutableDatabricksColumn.Builder getColumnBuilder() {
     return ImmutableDatabricksColumn.builder()
         .isAutoIncrement(false)

src/test/java/com/databricks/jdbc/client/impl/DatabricksNewMetadataSdkClientTest.java

@@ -202,8 +202,11 @@ void testListTables(
         .thenReturn(mockedResultSet);
     when(mockedResultSet.next()).thenReturn(true, false);
     for (ResultColumn resultColumn : TABLE_COLUMNS) {
-      when(mockedResultSet.getObject(resultColumn.getResultSetColumnName()))
-          .thenReturn(TEST_COLUMN);
+      if (resultColumn == TABLE_COLUMNS.get(3)) {
+        when(mockedResultSet.getObject(resultColumn.getResultSetColumnName())).thenReturn("TABLE");
+      } else
+        when(mockedResultSet.getObject(resultColumn.getResultSetColumnName()))
+            .thenReturn(TEST_COLUMN);
     }
     DatabricksResultSet actualResult =
         metadataClient.listTables(session, catalog, schema, table, null);
@@ -231,7 +234,7 @@ void testListColumns(
             sqlStatement,
             mockedComputeResource,
             new HashMap<Integer, ImmutableSqlParameter>(),
-            StatementType.METADATA,
+            StatementType.QUERY,
             session,
             null))
         .thenReturn(mockedResultSet);
@@ -264,10 +267,7 @@ void testListSchemas(String sqlStatement, String schema, String description) thr
             null))
         .thenReturn(mockedResultSet);
     when(mockedResultSet.next()).thenReturn(true, false);
-    for (ResultColumn resultColumn : SCHEMA_COLUMNS) {
-      when(mockedResultSet.getObject(resultColumn.getResultSetColumnName()))
-          .thenReturn(TEST_COLUMN);
-    }
+    when(mockedResultSet.getObject("databaseName")).thenReturn(TEST_COLUMN);
     DatabricksResultSet actualResult = metadataClient.listSchemas(session, TEST_CATALOG, schema);
     assertEquals(
         actualResult.getStatementStatus().getState(), StatementState.SUCCEEDED, description);

src/test/java/com/databricks/jdbc/client/impl/thrift/commons/DatabricksThriftAccessorTest.java

@@ -18,6 +18,7 @@
 import com.databricks.sdk.core.DatabricksConfig;
 import com.databricks.sdk.service.sql.StatementState;
 import java.sql.SQLException;
+import java.util.ArrayList;
 import java.util.Collections;
 import org.apache.thrift.TException;
 import org.apache.thrift.protocol.TProtocol;
@@ -46,11 +47,16 @@ public class DatabricksThriftAccessorTest {
           .setMaxBytes(DEFAULT_BYTE_LIMIT);
   private static final TGetResultSetMetadataReq resultSetMetadataReq =
       new TGetResultSetMetadataReq().setOperationHandle(tOperationHandle);
+
+  private static final TRowSet rowSet = new TRowSet().setResultLinks(new ArrayList<>(2));
+
   private static final TFetchResultsResp response =
       new TFetchResultsResp()
           .setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS))
           .setResultSetMetadata(
-              new TGetResultSetMetadataResp().setResultFormat(TSparkRowSetType.COLUMN_BASED_SET));
+              new TGetResultSetMetadataResp().setResultFormat(TSparkRowSetType.COLUMN_BASED_SET))
+          .setResults(rowSet);
+
   private static final TSparkDirectResults directResults =
       new TSparkDirectResults()
           .setResultSet(response)

src/test/java/com/databricks/jdbc/core/DatabricksResultSetMetaDataTest.java

@@ -101,7 +101,7 @@ public void testColumnsForVolumeOperation() throws SQLException {
   @Test
   public void testThriftColumns() throws SQLException {
     DatabricksResultSetMetaData metaData =
-        new DatabricksResultSetMetaData(STATEMENT_ID, getThriftResultManifest(), 10);
+        new DatabricksResultSetMetaData(STATEMENT_ID, getThriftResultManifest(), 10, 1);
     Assertions.assertEquals(10, metaData.getTotalRows());
     Assertions.assertEquals(1, metaData.getColumnCount());
     Assertions.assertEquals("testCol", metaData.getColumnName(1));
@@ -111,7 +111,7 @@ public void testThriftColumns() throws SQLException {
   public void testEmptyAndNullThriftColumns() throws SQLException {
     TGetResultSetMetadataResp resultSetMetadataResp = new TGetResultSetMetadataResp();
     DatabricksResultSetMetaData metaData =
-        new DatabricksResultSetMetaData(STATEMENT_ID, resultSetMetadataResp, 0);
+        new DatabricksResultSetMetaData(STATEMENT_ID, resultSetMetadataResp, 0, 1);
     Assertions.assertEquals(0, metaData.getColumnCount());
 
     resultSetMetadataResp.setSchema(new TTableSchema());