diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaData.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaData.java index 12d960df3..54f126a84 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaData.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaData.java @@ -25,10 +25,7 @@ import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Types; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; +import java.util.*; import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Collectors; @@ -44,6 +41,7 @@ public class DatabricksResultSetMetaData implements ResultSetMetaData { private final long totalRows; private Long chunkCount; private final boolean isCloudFetchUsed; + private final boolean truncated; /** * Constructs a {@code DatabricksResultSetMetaData} object for a SEA result set. @@ -144,6 +142,7 @@ public DatabricksResultSetMetaData( this.totalRows = resultManifest.getTotalRowCount(); this.chunkCount = resultManifest.getTotalChunkCount(); this.isCloudFetchUsed = usesExternalLinks; + this.truncated = Objects.requireNonNullElse(resultManifest.getTruncated(), false); } /** @@ -260,6 +259,7 @@ public DatabricksResultSetMetaData( this.totalRows = rows; this.chunkCount = chunkCount; this.isCloudFetchUsed = getIsCloudFetchFromManifest(resultManifest); + this.truncated = false; } /** @@ -309,6 +309,7 @@ public DatabricksResultSetMetaData( this.columnNameIndex = CaseInsensitiveImmutableMap.copyOf(columnNameToIndexMap); this.totalRows = totalRows; this.isCloudFetchUsed = false; + this.truncated = false; } /** @@ -361,6 +362,7 @@ public DatabricksResultSetMetaData( this.columnNameIndex = CaseInsensitiveImmutableMap.copyOf(columnNameToIndexMap); this.totalRows = totalRows; this.isCloudFetchUsed = false; + this.truncated = false; } /** @@ -419,6 +421,7 @@ public DatabricksResultSetMetaData( this.columnNameIndex = CaseInsensitiveImmutableMap.copyOf(columnNameToIndexMap); this.totalRows = totalRows; this.isCloudFetchUsed = false; + this.truncated = false; } /** @@ -496,6 +499,7 @@ public DatabricksResultSetMetaData( this.totalRows = -1; this.columns = columnsBuilder.build(); this.columnNameIndex = CaseInsensitiveImmutableMap.copyOf(columnNameToIndexMap); + this.truncated = false; } @Override @@ -645,6 +649,10 @@ private boolean getIsCloudFetchFromManifest(TGetResultSetMetadataResp resultMani return resultManifest.getResultFormat() == TSparkRowSetType.URL_BASED_SET; } + public boolean getIsTruncated() { + return truncated; + } + public Long getChunkCount() { return chunkCount; } diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaDataTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaDataTest.java index 0485e96a8..230596fa3 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaDataTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaDataTest.java @@ -461,6 +461,7 @@ public void testGetDispositionThrift(TSparkRowSetType resultFormat) { } else { assertFalse(metaData.getIsCloudFetchUsed()); } + assertFalse(metaData.getIsTruncated()); } @Test @@ -476,6 +477,26 @@ public void testCloudFetchUsedSdk() { assertFalse(metaData.getIsCloudFetchUsed()); } + @Test + public void testSdkTruncated() { + ResultManifest resultManifest = getResultManifest(); + resultManifest.setTruncated(null); + + DatabricksResultSetMetaData metaData = + new DatabricksResultSetMetaData(STATEMENT_ID, resultManifest, true, connectionContext); + assertFalse(metaData.getIsTruncated()); + + resultManifest.setTruncated(true); + metaData = + new DatabricksResultSetMetaData(STATEMENT_ID, resultManifest, false, connectionContext); + assertTrue(metaData.getIsTruncated()); + + resultManifest.setTruncated(false); + metaData = + new DatabricksResultSetMetaData(STATEMENT_ID, resultManifest, false, connectionContext); + assertFalse(metaData.getIsTruncated()); + } + @Test public void testSEAInlineComplexType() throws SQLException { ResultManifest resultManifest = new ResultManifest(); diff --git a/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/SqlExecApiHybridResultsIntegrationTests.java b/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/SqlExecApiHybridResultsIntegrationTests.java index f5253eb7e..eb49b5eef 100644 --- a/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/SqlExecApiHybridResultsIntegrationTests.java +++ b/src/test/java/com/databricks/jdbc/integration/fakeservice/tests/SqlExecApiHybridResultsIntegrationTests.java @@ -52,6 +52,7 @@ void testHybridSmallQuery() throws SQLException { assertEquals(maxRows, metaData.getTotalRows()); // For small query, arrow results are received inline in hybrid mode assertFalse(metaData.getIsCloudFetchUsed()); + assertFalse(metaData.getIsTruncated()); // For small query, arrow results are received inline in hybrid mode so no cloud fetch calls // are made @@ -93,6 +94,7 @@ void testHybridLargeQuery() throws SQLException { assertEquals(maxRows, metaData.getTotalRows()); // For large query, arrow results are fetched using cloud fetch assertTrue(metaData.getIsCloudFetchUsed()); + assertFalse(metaData.getIsTruncated()); // The number of cloud fetch calls should be equal to the number of chunks final int cloudFetchCalls =