diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryStoragePageSource.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryStoragePageSource.java index 0ba34641f789..efe85260e976 100644 --- a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryStoragePageSource.java +++ b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryStoragePageSource.java @@ -135,6 +135,9 @@ public boolean isFinished() public Page getNextPage() { checkState(pageBuilder.isEmpty(), "PageBuilder is not empty at the beginning of a new page"); + if (!responses.hasNext()) { + return null; + } ReadRowsResponse response = responses.next(); Iterable records = parse(response); for (GenericRecord record : records) { diff --git a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryConnectorTest.java b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryConnectorTest.java index ab5800d0870e..a1734407c303 100644 --- a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryConnectorTest.java +++ b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryConnectorTest.java @@ -37,6 +37,7 @@ import static io.trino.spi.type.VarcharType.VARCHAR; import static io.trino.testing.MaterializedResult.resultBuilder; import static io.trino.testing.assertions.Assert.assertEquals; +import static io.trino.testing.assertions.Assert.assertEventually; import static io.trino.testing.sql.TestTable.randomTableSuffix; import static java.lang.String.format; import static org.assertj.core.api.Assertions.assertThat; @@ -424,6 +425,38 @@ public void testColumnPositionMismatch() } } + @Test + public void testSelectTableWithRowAccessPolicyFilterAll() + { + String policyName = "test_policy" + randomTableSuffix(); + try (TestTable table = new TestTable(this::onBigQuery, "test.test_row_access_policy", "AS SELECT 1 col")) { + assertQuery("SELECT * FROM " + table.getName(), "VALUES 1"); + + // Use assertEventually because there's delay until new row access policies become effective + onBigQuery("CREATE ROW ACCESS POLICY " + policyName + " ON " + table.getName() + " FILTER USING (true)"); + assertEventually(() -> assertQueryReturnsEmptyResult("SELECT * FROM " + table.getName())); + + onBigQuery("DROP ALL ROW ACCESS POLICIES ON " + table.getName()); + assertEventually(() -> assertQuery("SELECT * FROM " + table.getName(), "VALUES 1")); + } + } + + @Test + public void testSelectTableWithRowAccessPolicyFilterPartialRow() + { + String policyName = "test_policy" + randomTableSuffix(); + try (TestTable table = new TestTable(this::onBigQuery, "test.test_row_access_policy", "AS (SELECT 1 col UNION ALL SELECT 2 col)")) { + assertQuery("SELECT * FROM " + table.getName(), "VALUES (1), (2)"); + + // Use assertEventually because there's delay until new row access policies become effective + onBigQuery("CREATE ROW ACCESS POLICY " + policyName + " ON " + table.getName() + " GRANT TO (\"allAuthenticatedUsers\") FILTER USING (col = 1)"); + assertEventually(() -> assertQuery("SELECT * FROM " + table.getName(), "VALUES 1")); + + onBigQuery("DROP ALL ROW ACCESS POLICIES ON " + table.getName()); + assertEventually(() -> assertQuery("SELECT * FROM " + table.getName(), "VALUES (1), (2)")); + } + } + @Test public void testViewDefinitionSystemTable() {