diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java index e060766cbd22..342596c746f2 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java @@ -46,6 +46,8 @@ import com.google.cloud.spanner.Options; import com.google.cloud.spanner.Options.RpcPriority; import com.google.cloud.spanner.PartitionOptions; +import com.google.cloud.spanner.ReadOnlyTransaction; +import com.google.cloud.spanner.ResultSet; import com.google.cloud.spanner.Spanner; import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerOptions; @@ -1998,6 +2000,11 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta final MapperFactory mapperFactory = new MapperFactory(changeStreamDatabaseDialect); final ChangeStreamMetrics metrics = new ChangeStreamMetrics(); final RpcPriority rpcPriority = MoreObjects.firstNonNull(getRpcPriority(), RpcPriority.HIGH); + final SpannerAccessor spannerAccessor = + SpannerAccessor.getOrCreate(changeStreamSpannerConfig); + final boolean isMutableChangeStream = + isMutableChangeStream( + spannerAccessor.getDatabaseClient(), changeStreamDatabaseDialect, changeStreamName); final DaoFactory daoFactory = new DaoFactory( changeStreamSpannerConfig, @@ -2007,7 +2014,8 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta rpcPriority, input.getPipeline().getOptions().getJobName(), changeStreamDatabaseDialect, - metadataDatabaseDialect); + metadataDatabaseDialect, + isMutableChangeStream); final ActionFactory actionFactory = new ActionFactory(); final Duration watermarkRefreshRate = @@ -2688,4 +2696,58 @@ static String resolveSpannerProjectId(SpannerConfig config) { ? SpannerOptions.getDefaultProjectId() : config.getProjectId().get(); } + + @VisibleForTesting + static boolean isMutableChangeStream( + DatabaseClient databaseClient, Dialect dialect, String changeStreamName) { + String fetchedPartitionMode = fetchPartitionMode(databaseClient, dialect, changeStreamName); + if (fetchedPartitionMode.isEmpty() + || fetchedPartitionMode.equalsIgnoreCase("IMMUTABLE_KEY_RANGE")) { + return false; + } + return true; + } + + private static String fetchPartitionMode( + DatabaseClient databaseClient, Dialect dialect, String changeStreamName) { + try (ReadOnlyTransaction tx = databaseClient.readOnlyTransaction()) { + Statement statement; + if (dialect == Dialect.POSTGRESQL) { + statement = + Statement.newBuilder( + "select option_value\n" + + "from information_schema.change_stream_options\n" + + "where change_stream_name = $1 and option_name = 'partition_mode'") + .bind("p1") + .to(changeStreamName) + .build(); + } else { + statement = + Statement.newBuilder( + "select option_value\n" + + "from information_schema.change_stream_options\n" + + "where change_stream_name = @changeStreamName and option_name = 'partition_mode'") + .bind("changeStreamName") + .to(changeStreamName) + .build(); + } + ResultSet resultSet = tx.executeQuery(statement); + while (resultSet.next()) { + String value = resultSet.getString(0); + if (value != null) { + return value; + } + } + return ""; + } catch (RuntimeException e) { + // Log the failure (with stack trace) but rethrow so the caller still observes + // the error. + LOG.warn( + "Failed to fetch partition_mode for change stream '{}', dialect={} - will propagate exception", + changeStreamName, + dialect, + e); + throw e; + } + } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamDao.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamDao.java index 3ef9c13f4714..0f32fa46cde7 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamDao.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamDao.java @@ -31,12 +31,12 @@ * as a {@link ResultSet}, which can be consumed until the stream is finished. */ public class ChangeStreamDao { - private final String changeStreamName; private final DatabaseClient databaseClient; private final RpcPriority rpcPriority; private final String jobName; private final Dialect dialect; + private final boolean isMutableChangeStream; /** * Constructs a change stream dao. All the queries performed by this class will be for the given @@ -53,12 +53,14 @@ public class ChangeStreamDao { DatabaseClient databaseClient, RpcPriority rpcPriority, String jobName, - Dialect dialect) { + Dialect dialect, + boolean isMutableChangeStream) { this.changeStreamName = changeStreamName; this.databaseClient = databaseClient; this.rpcPriority = rpcPriority; this.jobName = jobName; this.dialect = dialect; + this.isMutableChangeStream = isMutableChangeStream; } /** @@ -91,8 +93,18 @@ public ChangeStreamResultSet changeStreamQuery( String query = ""; Statement statement; if (this.isPostgres()) { - query = - "SELECT * FROM \"spanner\".\"read_json_" + changeStreamName + "\"($1, $2, $3, $4, null)"; + // Ensure we have determined whether change stream uses mutable key range + if (this.isMutableChangeStream) { + query = + "SELECT * FROM \"spanner\".\"read_proto_bytes_" + + changeStreamName + + "\"($1, $2, $3, $4, null)"; + } else { + query = + "SELECT * FROM \"spanner\".\"read_json_" + + changeStreamName + + "\"($1, $2, $3, $4, null)"; + } statement = Statement.newBuilder(query) .bind("p1") diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/DaoFactory.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/DaoFactory.java index 787abad02e02..67b58bace70f 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/DaoFactory.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/DaoFactory.java @@ -49,6 +49,7 @@ public class DaoFactory implements Serializable { private final String jobName; private final Dialect spannerChangeStreamDatabaseDialect; private final Dialect metadataDatabaseDialect; + private final boolean isMutableChangeStream; /** * Constructs a {@link DaoFactory} with the configuration to be used for the underlying instances. @@ -68,7 +69,8 @@ public DaoFactory( RpcPriority rpcPriority, String jobName, Dialect spannerChangeStreamDatabaseDialect, - Dialect metadataDatabaseDialect) { + Dialect metadataDatabaseDialect, + boolean isMutableChangeStream) { if (metadataSpannerConfig.getInstanceId() == null) { throw new IllegalArgumentException("Metadata instance can not be null"); } @@ -83,6 +85,7 @@ public DaoFactory( this.jobName = jobName; this.spannerChangeStreamDatabaseDialect = spannerChangeStreamDatabaseDialect; this.metadataDatabaseDialect = metadataDatabaseDialect; + this.isMutableChangeStream = isMutableChangeStream; } /** @@ -143,7 +146,8 @@ public synchronized ChangeStreamDao getChangeStreamDao() { spannerAccessor.getDatabaseClient(), rpcPriority, jobName, - this.spannerChangeStreamDatabaseDialect); + this.spannerChangeStreamDatabaseDialect, + this.isMutableChangeStream); } return changeStreamDaoInstance; } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadChangeStreamTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadChangeStreamTest.java index 5fd3548a3004..589d831e1a45 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadChangeStreamTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadChangeStreamTest.java @@ -19,10 +19,23 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import com.google.auth.Credentials; import com.google.cloud.Timestamp; +import com.google.cloud.spanner.DatabaseClient; +import com.google.cloud.spanner.Dialect; import com.google.cloud.spanner.Options.RpcPriority; +import com.google.cloud.spanner.ReadOnlyTransaction; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.Statement; +import java.util.Arrays; +import java.util.Collection; import org.apache.beam.sdk.extensions.gcp.auth.TestCredential; import org.apache.beam.sdk.extensions.gcp.options.GcpOptions; import org.apache.beam.sdk.io.gcp.spanner.changestreams.MetadataSpannerConfigFactory; @@ -30,10 +43,15 @@ import org.junit.Before; import org.junit.Rule; import org.junit.Test; +import org.junit.experimental.runners.Enclosed; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; +import org.mockito.ArgumentCaptor; -@RunWith(JUnit4.class) +@RunWith(Enclosed.class) public class SpannerIOReadChangeStreamTest { private static final String TEST_PROJECT = "my-project"; @@ -44,100 +62,192 @@ public class SpannerIOReadChangeStreamTest { private static final String TEST_METADATA_TABLE = "my-metadata-table"; private static final String TEST_CHANGE_STREAM = "my-change-stream"; - @Rule public final transient TestPipeline testPipeline = TestPipeline.create(); - - private SpannerConfig spannerConfig; - private SpannerIO.ReadChangeStream readChangeStream; - - @Before - public void setUp() throws Exception { - spannerConfig = - SpannerConfig.create() - .withProjectId(TEST_PROJECT) - .withInstanceId(TEST_INSTANCE) - .withDatabaseId(TEST_DATABASE); - - Timestamp startTimestamp = Timestamp.now(); - Timestamp endTimestamp = - Timestamp.ofTimeSecondsAndNanos( - startTimestamp.getSeconds() + 10, startTimestamp.getNanos()); - readChangeStream = - SpannerIO.readChangeStream() - .withSpannerConfig(spannerConfig) - .withChangeStreamName(TEST_CHANGE_STREAM) - .withMetadataInstance(TEST_METADATA_INSTANCE) - .withMetadataDatabase(TEST_METADATA_DATABASE) - .withMetadataTable(TEST_METADATA_TABLE) - .withRpcPriority(RpcPriority.MEDIUM) - .withInclusiveStartAt(startTimestamp) - .withInclusiveEndAt(endTimestamp); - } + /** Basic configuration tests using standard JUnit4. */ + @RunWith(JUnit4.class) + public static class ConfigurationTests { + @Rule public final transient TestPipeline testPipeline = TestPipeline.create(); + private SpannerConfig spannerConfig; + private SpannerIO.ReadChangeStream readChangeStream; + + @Before + public void setUp() { + spannerConfig = + SpannerConfig.create() + .withProjectId(TEST_PROJECT) + .withInstanceId(TEST_INSTANCE) + .withDatabaseId(TEST_DATABASE); + + readChangeStream = + SpannerIO.readChangeStream() + .withSpannerConfig(spannerConfig) + .withChangeStreamName(TEST_CHANGE_STREAM) + .withMetadataInstance(TEST_METADATA_INSTANCE) + .withMetadataDatabase(TEST_METADATA_DATABASE) + .withMetadataTable(TEST_METADATA_TABLE) + .withRpcPriority(RpcPriority.MEDIUM) + .withInclusiveStartAt(Timestamp.now()); + } + + @Test + public void testSetPipelineCredential() { + TestCredential testCredential = new TestCredential(); + // Set the credential in the pipeline options. + testPipeline.getOptions().as(GcpOptions.class).setGcpCredential(testCredential); + SpannerConfig changeStreamSpannerConfig = readChangeStream.buildChangeStreamSpannerConfig(); + SpannerConfig metadataSpannerConfig = + MetadataSpannerConfigFactory.create( + changeStreamSpannerConfig, TEST_METADATA_INSTANCE, TEST_METADATA_DATABASE); + assertNull(changeStreamSpannerConfig.getCredentials()); + assertNull(metadataSpannerConfig.getCredentials()); + + SpannerConfig changeStreamSpannerConfigWithCredential = + SpannerIO.buildSpannerConfigWithCredential( + changeStreamSpannerConfig, testPipeline.getOptions()); + SpannerConfig metadataSpannerConfigWithCredential = + SpannerIO.buildSpannerConfigWithCredential( + metadataSpannerConfig, testPipeline.getOptions()); + assertEquals(testCredential, changeStreamSpannerConfigWithCredential.getCredentials().get()); + assertEquals(testCredential, metadataSpannerConfigWithCredential.getCredentials().get()); + } + + @Test + public void testSetSpannerConfigCredential() { + TestCredential testCredential = new TestCredential(); + // Set the credential in the SpannerConfig. + spannerConfig = spannerConfig.withCredentials(testCredential); + readChangeStream = readChangeStream.withSpannerConfig(spannerConfig); + SpannerConfig changeStreamSpannerConfig = readChangeStream.buildChangeStreamSpannerConfig(); + SpannerConfig metadataSpannerConfig = + MetadataSpannerConfigFactory.create( + changeStreamSpannerConfig, TEST_METADATA_INSTANCE, TEST_METADATA_DATABASE); + assertEquals(testCredential, changeStreamSpannerConfig.getCredentials().get()); + assertEquals(testCredential, metadataSpannerConfig.getCredentials().get()); - @Test - public void testSetPipelineCredential() { - TestCredential testCredential = new TestCredential(); - // Set the credential in the pipeline options. - testPipeline.getOptions().as(GcpOptions.class).setGcpCredential(testCredential); - SpannerConfig changeStreamSpannerConfig = readChangeStream.buildChangeStreamSpannerConfig(); - SpannerConfig metadataSpannerConfig = - MetadataSpannerConfigFactory.create( - changeStreamSpannerConfig, TEST_METADATA_INSTANCE, TEST_METADATA_DATABASE); - assertNull(changeStreamSpannerConfig.getCredentials()); - assertNull(metadataSpannerConfig.getCredentials()); - - SpannerConfig changeStreamSpannerConfigWithCredential = - SpannerIO.buildSpannerConfigWithCredential( - changeStreamSpannerConfig, testPipeline.getOptions()); - SpannerConfig metadataSpannerConfigWithCredential = - SpannerIO.buildSpannerConfigWithCredential( - metadataSpannerConfig, testPipeline.getOptions()); - assertEquals(testCredential, changeStreamSpannerConfigWithCredential.getCredentials().get()); - assertEquals(testCredential, metadataSpannerConfigWithCredential.getCredentials().get()); + SpannerConfig changeStreamSpannerConfigWithCredential = + SpannerIO.buildSpannerConfigWithCredential( + changeStreamSpannerConfig, testPipeline.getOptions()); + SpannerConfig metadataSpannerConfigWithCredential = + SpannerIO.buildSpannerConfigWithCredential( + metadataSpannerConfig, testPipeline.getOptions()); + assertEquals(testCredential, changeStreamSpannerConfigWithCredential.getCredentials().get()); + assertEquals(testCredential, metadataSpannerConfigWithCredential.getCredentials().get()); + } + + @Test + public void testWithDefaultCredential() { + // Get the default credential, without setting any credentials in the pipeline + // options or SpannerConfig. + Credentials defaultCredential = + testPipeline.getOptions().as(GcpOptions.class).getGcpCredential(); + SpannerConfig changeStreamSpannerConfig = readChangeStream.buildChangeStreamSpannerConfig(); + SpannerConfig metadataSpannerConfig = + MetadataSpannerConfigFactory.create( + changeStreamSpannerConfig, TEST_METADATA_INSTANCE, TEST_METADATA_DATABASE); + assertNull(changeStreamSpannerConfig.getCredentials()); + assertNull(metadataSpannerConfig.getCredentials()); + + SpannerConfig changeStreamSpannerConfigWithCredential = + SpannerIO.buildSpannerConfigWithCredential( + changeStreamSpannerConfig, testPipeline.getOptions()); + SpannerConfig metadataSpannerConfigWithCredential = + SpannerIO.buildSpannerConfigWithCredential( + metadataSpannerConfig, testPipeline.getOptions()); + assertEquals( + defaultCredential, changeStreamSpannerConfigWithCredential.getCredentials().get()); + assertEquals(defaultCredential, metadataSpannerConfigWithCredential.getCredentials().get()); + } } - @Test - public void testSetSpannerConfigCredential() { - TestCredential testCredential = new TestCredential(); - // Set the credential in the SpannerConfig. - spannerConfig = spannerConfig.withCredentials(testCredential); - readChangeStream = readChangeStream.withSpannerConfig(spannerConfig); - SpannerConfig changeStreamSpannerConfig = readChangeStream.buildChangeStreamSpannerConfig(); - SpannerConfig metadataSpannerConfig = - MetadataSpannerConfigFactory.create( - changeStreamSpannerConfig, TEST_METADATA_INSTANCE, TEST_METADATA_DATABASE); - assertEquals(testCredential, changeStreamSpannerConfig.getCredentials().get()); - assertEquals(testCredential, metadataSpannerConfig.getCredentials().get()); - - SpannerConfig changeStreamSpannerConfigWithCredential = - SpannerIO.buildSpannerConfigWithCredential( - changeStreamSpannerConfig, testPipeline.getOptions()); - SpannerConfig metadataSpannerConfigWithCredential = - SpannerIO.buildSpannerConfigWithCredential( - metadataSpannerConfig, testPipeline.getOptions()); - assertEquals(testCredential, changeStreamSpannerConfigWithCredential.getCredentials().get()); - assertEquals(testCredential, metadataSpannerConfigWithCredential.getCredentials().get()); + /** Parameterized tests for Dialect and Partition Mode combinations. */ + @RunWith(Parameterized.class) + public static class PartitionModeTests { + + @Parameters(name = "{index}: dialect={0}, mode={1}, expected={2}") + public static Collection data() { + return Arrays.asList( + new Object[][] { + {Dialect.GOOGLE_STANDARD_SQL, "MUTABLE_KEY_RANGE", true}, + {Dialect.GOOGLE_STANDARD_SQL, "IMMUTABLE_KEY_RANGE", false}, + {Dialect.GOOGLE_STANDARD_SQL, "", false}, // Empty string case + {Dialect.POSTGRESQL, "MUTABLE_KEY_RANGE", true}, + {Dialect.POSTGRESQL, "IMMUTABLE_KEY_RANGE", false}, + {Dialect.POSTGRESQL, "", false} + }); + } + + @Parameter(0) + public Dialect dialect; + + @Parameter(1) + public String partitionMode; + + @Parameter(2) + public boolean expected; + + @Test + public void testIsMutableChangeStream() { + DatabaseClient databaseClient = mock(DatabaseClient.class); + ReadOnlyTransaction transaction = mock(ReadOnlyTransaction.class); + ResultSet resultSet = mock(ResultSet.class); + + when(databaseClient.readOnlyTransaction()).thenReturn(transaction); + when(transaction.executeQuery(any(Statement.class))).thenReturn(resultSet); + + // Handle the different return values for the mock ResultSet + if (partitionMode.isEmpty()) { + // If the partition mode is empty (e.g., the option is not set in Spanner), + // simulate an empty result set by having next() return false immediately. + when(resultSet.next()).thenReturn(false); // No row returned + } else { + // If a partition mode exists, simulate a result set containing exactly one row. + // next() returns true for the first call (row exists) and false for the second (end of + // stream). + when(resultSet.next()).thenReturn(true).thenReturn(false); + // When getString(0) is called to retrieve the 'option_value' column from the row, + // return the specified partitionMode string (e.g., "MUTABLE_KEY_RANGE"). + when(resultSet.getString(0)).thenReturn(partitionMode); + } + + boolean actual = SpannerIO.isMutableChangeStream(databaseClient, dialect, TEST_CHANGE_STREAM); + assertEquals(expected, actual); + + // Verify SQL Syntax: Captures the statement to check for dialect-specific + // syntax + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(Statement.class); + verify(transaction).executeQuery(statementCaptor.capture()); + String sql = statementCaptor.getValue().getSql(); + + // Ensure the SQL uses the correct parameter placeholder syntax for the given dialect. + // Different dialects have different requirements for how parameters are bound in queries. + if (dialect == Dialect.POSTGRESQL) { + // PostgreSQL-dialect Spanner databases use positional parameters (e.g., $1, $2) + assertTrue("PostgreSQL SQL should use $1", sql.contains("$1")); + } else { + // Google Standard SQL-dialect Spanner databases use named parameters (e.g., + // @changeStreamName) + assertTrue("GoogleSQL SQL should use @changeStreamName", sql.contains("@changeStreamName")); + } + } } - @Test - public void testWithDefaultCredential() { - // Get the default credential, without setting any credentials in the pipeline options or - // SpannerConfig. - Credentials defaultCredential = - testPipeline.getOptions().as(GcpOptions.class).getGcpCredential(); - SpannerConfig changeStreamSpannerConfig = readChangeStream.buildChangeStreamSpannerConfig(); - SpannerConfig metadataSpannerConfig = - MetadataSpannerConfigFactory.create( - changeStreamSpannerConfig, TEST_METADATA_INSTANCE, TEST_METADATA_DATABASE); - assertNull(changeStreamSpannerConfig.getCredentials()); - assertNull(metadataSpannerConfig.getCredentials()); - - SpannerConfig changeStreamSpannerConfigWithCredential = - SpannerIO.buildSpannerConfigWithCredential( - changeStreamSpannerConfig, testPipeline.getOptions()); - SpannerConfig metadataSpannerConfigWithCredential = - SpannerIO.buildSpannerConfigWithCredential( - metadataSpannerConfig, testPipeline.getOptions()); - assertEquals(defaultCredential, changeStreamSpannerConfigWithCredential.getCredentials().get()); - assertEquals(defaultCredential, metadataSpannerConfigWithCredential.getCredentials().get()); + /** Tests for error handling and exceptions. */ + @RunWith(JUnit4.class) + public static class ErrorTests { + + @Test(expected = RuntimeException.class) + public void testIsMutableChangeStream_PropagatesException() { + DatabaseClient databaseClient = mock(DatabaseClient.class); + ReadOnlyTransaction transaction = mock(ReadOnlyTransaction.class); + + // Mock the transaction creation + when(databaseClient.readOnlyTransaction()).thenReturn(transaction); + + // Simulate a database failure when executing the query + when(transaction.executeQuery(any(Statement.class))) + .thenThrow(new RuntimeException("Database connection failed")); + + // The method should log the error and rethrow the exception + SpannerIO.isMutableChangeStream(databaseClient, Dialect.GOOGLE_STANDARD_SQL, "test-stream"); + } } } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java index 835ca0a0f5a8..c8435bcfdffd 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java @@ -331,6 +331,7 @@ public void testInvalidRecordReceivedWithDefaultSettings() { Timestamp.ofTimeSecondsAndNanos(startTimestamp.getSeconds(), startTimestamp.getNanos() + 1); mockGetDialect(); + mockChangeStreamOptions(); mockTableExists(); mockGetWatermark(startTimestamp); ResultSet getPartitionResultSet = mockGetParentPartition(startTimestamp, endTimestamp); @@ -382,6 +383,38 @@ public void testInvalidRecordReceivedWithDefaultSettings() { } } + private void mockChangeStreamOptions() { + Statement changeStreamOptionsStatement = + Statement.newBuilder( + "select option_value\n" + + "from information_schema.change_stream_options\n" + + "where change_stream_name = @changeStreamName and option_name = 'partition_mode'") + .bind("changeStreamName") + .to(TEST_CHANGE_STREAM) + .build(); + ResultSetMetadata changeStreamOptionsResultSetMetadata = + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setName("option_value") + .setType(Type.newBuilder().setCode(TypeCode.STRING).build()) + .build()) + .build()) + .build(); + ResultSet changeStreamOptionsResultSet = + ResultSet.newBuilder() + .addRows( + ListValue.newBuilder() + .addValues(Value.newBuilder().setStringValue("NEW_VALUES").build()) + .build()) + .setMetadata(changeStreamOptionsResultSetMetadata) + .build(); + mockSpannerService.putPartialStatementResult( + StatementResult.query(changeStreamOptionsStatement, changeStreamOptionsResultSet)); + } + private void mockInvalidChangeStreamRecordReceived(Timestamp now, Timestamp after3Seconds) { Statement changeStreamQueryStatement = Statement.newBuilder( diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamDaoTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamDaoTest.java new file mode 100644 index 000000000000..9bdaa7b9fa5a --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamDaoTest.java @@ -0,0 +1,97 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.gcp.spanner.changestreams.dao; + +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.cloud.spanner.DatabaseClient; +import com.google.cloud.spanner.Dialect; +import com.google.cloud.spanner.Options.RpcPriority; +import com.google.cloud.spanner.ReadOnlyTransaction; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.Statement; +import org.junit.Before; +import org.junit.Test; +import org.mockito.ArgumentCaptor; + +public class ChangeStreamDaoTest { + private DatabaseClient databaseClient; + private RpcPriority rpcPriority; + private static final String CHANGE_STREAM_NAME = "testCS"; + + @Before + public void setUp() { + databaseClient = mock(DatabaseClient.class); + rpcPriority = mock(RpcPriority.class); + } + + // New tests for PostgreSQL branch to verify the chosen TVF in the generated + // SQL. + @Test + public void testChangeStreamQueryPostgresMutable() { + // Arrange: single-use transaction for the actual change stream query + ReadOnlyTransaction singleUseTx = mock(ReadOnlyTransaction.class); + when(databaseClient.singleUse()).thenReturn(singleUseTx); + + ChangeStreamDao changeStreamDao = + new ChangeStreamDao( + CHANGE_STREAM_NAME, databaseClient, rpcPriority, "testjob", Dialect.POSTGRESQL, true); + + // Act: call the method that constructs and executes the statement + changeStreamDao.changeStreamQuery(null, null, null, 0L); + + // Assert: capture the Statement passed to singleUse().executeQuery and verify + // SQL + ArgumentCaptor captor = ArgumentCaptor.forClass(Statement.class); + verify(singleUseTx).executeQuery(captor.capture(), any(), any()); + Statement captured = captor.getValue(); + String sql = captured.getSql(); // adjust if different accessor is used + assertTrue( + "Expected SQL to contain read_proto_bytes_", + sql.contains("read_proto_bytes_" + CHANGE_STREAM_NAME)); + } + + @Test + public void testChangeStreamQueryPostgresImmutable() { + // Arrange: single-use transaction for the actual change stream query + ReadOnlyTransaction singleUseTx = mock(ReadOnlyTransaction.class); + when(databaseClient.singleUse()).thenReturn(singleUseTx); + + ResultSet queryResult = mock(ResultSet.class); + when(singleUseTx.executeQuery(any(), any(), any())).thenReturn(queryResult); + + ChangeStreamDao changeStreamDao = + new ChangeStreamDao( + CHANGE_STREAM_NAME, databaseClient, rpcPriority, "testjob", Dialect.POSTGRESQL, false); + + // Act + changeStreamDao.changeStreamQuery(null, null, null, 0L); + + // Assert + ArgumentCaptor captor = ArgumentCaptor.forClass(Statement.class); + verify(singleUseTx).executeQuery(captor.capture(), any(), any()); + Statement captured = captor.getValue(); + String sql = captured.getSql(); + assertTrue( + "Expected SQL to contain read_json_", sql.contains("read_json_" + CHANGE_STREAM_NAME)); + } +}