This is an automated email from the ASF dual-hosted git repository.
scwhittle pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new f7537312886 [Spanner Change Stream] Fix Change Streams PostgreSQL
metadata queries (#40755)
f7537312886 is described below
commit f7537312886b4a109c63e1667f66c8978b4a43e9
Author: Jiang Zhu <[email protected]>
AuthorDate: Tue Oct 27 11:35:47 2026 -0700
[Spanner Change Stream] Fix Change Streams PostgreSQL metadata queries
(#40755)
* Fix Spanner Change Streams PostgreSQL metadata queries
* Add PostgreSQL updateToScheduled tests for empty and multiple tokens
---
.../changestreams/dao/PartitionMetadataDao.java | 55 ++++---
.../dao/PartitionMetadataDaoTest.java | 173 +++++++++++++++++++++
2 files changed, 204 insertions(+), 24 deletions(-)
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java
index a60bd5c2143..7b0825cf615 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java
@@ -50,7 +50,6 @@ import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
-import java.util.stream.Collectors;
import javax.annotation.Nullable;
import
org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata;
import
org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata.State;
@@ -83,13 +82,23 @@ public class PartitionMetadataDao {
* @return true if the table exists, false if the table does not exist.
*/
public boolean tableExists() {
- final String checkTableExistsStmt =
- "SELECT t.table_name FROM information_schema.tables AS t "
- + "WHERE t.table_catalog = '' AND "
- + "t.table_schema = '' AND "
- + "t.table_name = '"
- + metadataTableName
- + "'";
+ final String checkTableExistsStmt;
+ if (this.isPostgres()) {
+ checkTableExistsStmt =
+ "SELECT t.table_name FROM information_schema.tables AS t "
+ + "WHERE t.table_schema = 'public' AND "
+ + "t.table_name = '"
+ + metadataTableName
+ + "'";
+ } else {
+ checkTableExistsStmt =
+ "SELECT t.table_name FROM information_schema.tables AS t "
+ + "WHERE t.table_catalog = '' AND "
+ + "t.table_schema = '' AND "
+ + "t.table_name = '"
+ + metadataTableName
+ + "'";
+ }
try (ResultSet queryResultSet =
databaseClient
.singleUseReadOnlyTransaction()
@@ -681,22 +690,20 @@ public class PartitionMetadataDao {
List<String> composedPartitionTokens, State state) {
Statement statement;
if (this.dialect == Dialect.POSTGRESQL) {
- StringBuilder sqlStringBuilder =
- new StringBuilder("SELECT * FROM \"" + metadataTableName + "\"");
- sqlStringBuilder.append(" WHERE \"");
- sqlStringBuilder.append(COLUMN_STATE + "\" = " + "'" +
state.toString() + "'");
- if (!composedPartitionTokens.isEmpty()) {
- sqlStringBuilder.append(" AND \"");
- sqlStringBuilder.append(COLUMN_PARTITION_TOKEN);
- sqlStringBuilder.append("\"");
- sqlStringBuilder.append(" = ANY (Array[");
- sqlStringBuilder.append(
- composedPartitionTokens.stream()
- .map(s -> "'" + s + "'")
- .collect(Collectors.joining(",")));
- sqlStringBuilder.append("])");
- }
- statement = Statement.newBuilder(sqlStringBuilder.toString()).build();
+ statement =
+ Statement.newBuilder(
+ "SELECT * FROM \""
+ + metadataTableName
+ + "\" WHERE \""
+ + COLUMN_STATE
+ + "\" = $1 AND \""
+ + COLUMN_PARTITION_TOKEN
+ + "\" = ANY ($2)")
+ .bind("p1")
+ .to(state.toString())
+ .bind("p2")
+ .to(Value.stringArray(new
ArrayList<>(composedPartitionTokens)))
+ .build();
} else {
statement =
Statement.newBuilder(
diff --git
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDaoTest.java
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDaoTest.java
index 79fe01a2b7d..3b2b550ab1d 100644
---
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDaoTest.java
+++
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDaoTest.java
@@ -18,8 +18,10 @@
package org.apache.beam.sdk.io.gcp.spanner.changestreams.dao;
import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.mock;
@@ -31,7 +33,9 @@ import com.google.cloud.Timestamp;
import com.google.cloud.spanner.DatabaseClient;
import com.google.cloud.spanner.Dialect;
import com.google.cloud.spanner.Mutation;
+import com.google.cloud.spanner.ReadOnlyTransaction;
import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.Statement;
import com.google.cloud.spanner.Struct;
import com.google.cloud.spanner.TransactionContext;
import com.google.cloud.spanner.TransactionRunner;
@@ -406,4 +410,173 @@ public class PartitionMetadataDaoTest {
assertEquals("tvf",
PartitionMetadataDao.extractTvfName("partitionToken123#tvf"));
assertEquals("", PartitionMetadataDao.extractTvfName("partitionToken123"));
}
+
+ @Test
+ public void testTableExistsGoogleSql() {
+ ReadOnlyTransaction readOnlyTx = mock(ReadOnlyTransaction.class);
+ ResultSet resultSet = mock(ResultSet.class);
+ when(databaseClient.singleUseReadOnlyTransaction()).thenReturn(readOnlyTx);
+ when(readOnlyTx.executeQuery(any(Statement.class),
any())).thenReturn(resultSet);
+ when(resultSet.next()).thenReturn(true);
+
+ assertTrue(partitionMetadataDao.tableExists());
+
+ ArgumentCaptor<Statement> statementCaptor =
ArgumentCaptor.forClass(Statement.class);
+ verify(readOnlyTx).executeQuery(statementCaptor.capture(), any());
+ assertEquals(
+ "SELECT t.table_name FROM information_schema.tables AS t "
+ + "WHERE t.table_catalog = '' AND "
+ + "t.table_schema = '' AND "
+ + "t.table_name = '"
+ + METADATA_TABLE_NAME
+ + "'",
+ statementCaptor.getValue().getSql());
+ }
+
+ @Test
+ public void testTableExistsPostgres() {
+ ReadOnlyTransaction readOnlyTx = mock(ReadOnlyTransaction.class);
+ ResultSet resultSet = mock(ResultSet.class);
+ when(databaseClient.singleUseReadOnlyTransaction()).thenReturn(readOnlyTx);
+ when(readOnlyTx.executeQuery(any(Statement.class),
any())).thenReturn(resultSet);
+ when(resultSet.next()).thenReturn(true).thenReturn(false);
+
+ PartitionMetadataDao pgDao =
+ new PartitionMetadataDao(METADATA_TABLE_NAME, databaseClient,
Dialect.POSTGRESQL);
+ assertTrue(pgDao.tableExists());
+ assertFalse(pgDao.tableExists());
+
+ ArgumentCaptor<Statement> statementCaptor =
ArgumentCaptor.forClass(Statement.class);
+ verify(readOnlyTx, times(2)).executeQuery(statementCaptor.capture(),
any());
+ assertEquals(
+ "SELECT t.table_name FROM information_schema.tables AS t "
+ + "WHERE t.table_schema = 'public' AND "
+ + "t.table_name = '"
+ + METADATA_TABLE_NAME
+ + "'",
+ statementCaptor.getValue().getSql());
+ }
+
+ @Test
+ public void testInTransactionContextUpdateToScheduledPostgres() {
+ PartitionMetadataDao.InTransactionContext pgTransactionContext =
+ new PartitionMetadataDao.InTransactionContext(
+ METADATA_TABLE_NAME, transaction, Dialect.POSTGRESQL);
+ ResultSet resultSet = mock(ResultSet.class);
+ when(transaction.executeQuery(any(), any())).thenReturn(resultSet);
+ when(resultSet.next()).thenReturn(true).thenReturn(false);
+ when(resultSet.getString(any())).thenReturn(PARTITION_TOKEN);
+
+ ArgumentCaptor<Statement> statementCaptor =
ArgumentCaptor.forClass(Statement.class);
+ ArgumentCaptor<ImmutableList<Mutation>> mutations =
+ ArgumentCaptor.forClass(ImmutableList.class);
+ doNothing().when(transaction).buffer(mutations.capture());
+
+
assertNull(pgTransactionContext.updateToScheduled(Collections.singletonList(PARTITION_TOKEN)));
+
+ verify(transaction).executeQuery(statementCaptor.capture(), any());
+ Statement captured = statementCaptor.getValue();
+ assertEquals(
+ "SELECT * FROM \""
+ + METADATA_TABLE_NAME
+ + "\" WHERE \"State\" = $1 AND \"PartitionToken\" = ANY ($2)",
+ captured.getSql());
+ assertEquals(State.CREATED.toString(),
captured.getParameters().get("p1").getString());
+ assertEquals(
+ Collections.singletonList(PARTITION_TOKEN),
+ captured.getParameters().get("p2").getStringArray());
+ assertEquals(1, mutations.getValue().size());
+ }
+
+ @Test
+ public void testInTransactionContextUpdateToScheduledPostgresEmptyTokens() {
+ PartitionMetadataDao.InTransactionContext pgTransactionContext =
+ new PartitionMetadataDao.InTransactionContext(
+ METADATA_TABLE_NAME, transaction, Dialect.POSTGRESQL);
+ ResultSet resultSet = mock(ResultSet.class);
+ when(transaction.executeQuery(any(), any())).thenReturn(resultSet);
+ when(resultSet.next()).thenReturn(false);
+
+ ArgumentCaptor<Statement> statementCaptor =
ArgumentCaptor.forClass(Statement.class);
+
+
assertNull(pgTransactionContext.updateToScheduled(Collections.emptyList()));
+
+ verify(transaction).executeQuery(statementCaptor.capture(), any());
+ Statement captured = statementCaptor.getValue();
+ assertEquals(
+ "SELECT * FROM \""
+ + METADATA_TABLE_NAME
+ + "\" WHERE \"State\" = $1 AND \"PartitionToken\" = ANY ($2)",
+ captured.getSql());
+ assertEquals(State.CREATED.toString(),
captured.getParameters().get("p1").getString());
+ assertEquals(Collections.emptyList(),
captured.getParameters().get("p2").getStringArray());
+ verify(transaction, times(0)).buffer(any(Iterable.class));
+ }
+
+ @Test
+ public void
testInTransactionContextUpdateToScheduledPostgresMultipleTokens() {
+ PartitionMetadataDao.InTransactionContext pgTransactionContext =
+ new PartitionMetadataDao.InTransactionContext(
+ METADATA_TABLE_NAME, transaction, Dialect.POSTGRESQL);
+ ResultSet resultSet = mock(ResultSet.class);
+ when(transaction.executeQuery(any(), any())).thenReturn(resultSet);
+ when(resultSet.next()).thenReturn(true).thenReturn(true).thenReturn(false);
+ when(resultSet.getString(PartitionMetadataAdminDao.COLUMN_PARTITION_TOKEN))
+ .thenReturn("token1")
+ .thenReturn("token3");
+
+ ArgumentCaptor<Statement> statementCaptor =
ArgumentCaptor.forClass(Statement.class);
+ ArgumentCaptor<ImmutableList<Mutation>> mutations =
+ ArgumentCaptor.forClass(ImmutableList.class);
+ doNothing().when(transaction).buffer(mutations.capture());
+
+ List<String> inputTokens = ImmutableList.of("token1", "token2", "token3");
+ assertNull(pgTransactionContext.updateToScheduled(inputTokens));
+
+ verify(transaction).executeQuery(statementCaptor.capture(), any());
+ Statement captured = statementCaptor.getValue();
+ assertEquals(
+ "SELECT * FROM \""
+ + METADATA_TABLE_NAME
+ + "\" WHERE \"State\" = $1 AND \"PartitionToken\" = ANY ($2)",
+ captured.getSql());
+ assertEquals(State.CREATED.toString(),
captured.getParameters().get("p1").getString());
+ assertEquals(inputTokens,
captured.getParameters().get("p2").getStringArray());
+
+ verify(transaction, times(2)).buffer(any(ImmutableList.class));
+ List<ImmutableList<Mutation>> bufferedBatches = mutations.getAllValues();
+ assertEquals(2, bufferedBatches.size());
+ assertEquals(
+ "token1",
+ bufferedBatches
+ .get(0)
+ .get(0)
+ .asMap()
+ .get(PartitionMetadataAdminDao.COLUMN_PARTITION_TOKEN)
+ .getString());
+ assertEquals(
+ State.SCHEDULED.toString(),
+ bufferedBatches
+ .get(0)
+ .get(0)
+ .asMap()
+ .get(PartitionMetadataAdminDao.COLUMN_STATE)
+ .getString());
+ assertEquals(
+ "token3",
+ bufferedBatches
+ .get(1)
+ .get(0)
+ .asMap()
+ .get(PartitionMetadataAdminDao.COLUMN_PARTITION_TOKEN)
+ .getString());
+ assertEquals(
+ State.SCHEDULED.toString(),
+ bufferedBatches
+ .get(1)
+ .get(0)
+ .asMap()
+ .get(PartitionMetadataAdminDao.COLUMN_STATE)
+ .getString());
+ }
}
