diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/InitializeDoFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/InitializeDoFn.java index fb0348dac9ac..ab752cda2a94 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/InitializeDoFn.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/InitializeDoFn.java @@ -17,6 +17,8 @@ */ package org.apache.beam.sdk.io.gcp.spanner.changestreams.dofn; +import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.SpannerException; import java.io.Serializable; import java.util.ArrayList; import java.util.List; @@ -31,6 +33,8 @@ import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.SdkHarnessOptions; import org.apache.beam.sdk.transforms.DoFn; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * A DoFn responsible for initializing the change stream Connector. It handles the creation of the @@ -40,6 +44,7 @@ public class InitializeDoFn extends DoFn implements Serializable { private static final long serialVersionUID = -8921188388649003102L; + private static final Logger LOG = LoggerFactory.getLogger(InitializeDoFn.class); private final long heartbeatMillis; @@ -117,10 +122,11 @@ public void processElement(OutputReceiver receiver) { * specified in {@link InitializeDoFn#DEFAULT_HEARTBEAT_MILLIS}. */ private void createFakeParentPartition() { + List parentPartitions = new ArrayList<>(); if (daoFactory.getTvfNameList().isEmpty()) { // For IMMUTABLE_KEY_RANGE change stream or MUTABLE_KEY_RANGE // change stream without the specified tvf name list. - PartitionMetadata parentPartition = + parentPartitions.add( PartitionMetadata.newBuilder() .setPartitionToken(InitialPartition.PARTITION_TOKEN) .setTvfName(ChangeStreamsConstants.DEFAULT_TVF_NAME) @@ -129,13 +135,11 @@ private void createFakeParentPartition() { .setHeartbeatMillis(heartbeatMillis) .setState(State.CREATED) .setWatermark(startTimestamp) - .build(); - daoFactory.getPartitionMetadataDao().insert(parentPartition); + .build()); } else { // For MUTABLE_KEY_RANGE change stream with the specified tvf name list. - List parentPartitions = new ArrayList<>(); for (String tvfName : daoFactory.getTvfNameList()) { - PartitionMetadata parentPartition = + parentPartitions.add( PartitionMetadata.newBuilder() .setPartitionToken(InitialPartition.PARTITION_TOKEN) .setTvfName(tvfName) @@ -144,10 +148,36 @@ private void createFakeParentPartition() { .setHeartbeatMillis(heartbeatMillis) .setState(State.CREATED) .setWatermark(startTimestamp) - .build(); - parentPartitions.add(parentPartition); + .build()); + } + } + try { + daoFactory + .getPartitionMetadataDao() + .runInTransaction( + transaction -> { + List partitionsToInsert = new ArrayList<>(); + for (PartitionMetadata parentPartition : parentPartitions) { + String composedToken = + PartitionMetadataDao.composePartitionTokenWithTvfName( + parentPartition.getPartitionToken(), parentPartition.getTvfName()); + if (transaction.getPartition(composedToken) == null) { + partitionsToInsert.add(parentPartition); + } + } + if (!partitionsToInsert.isEmpty()) { + transaction.insert(partitionsToInsert); + return true; + } + return false; + }, + "InsertFakeParentPartition"); + } catch (SpannerException e) { + if (e.getErrorCode() == ErrorCode.ALREADY_EXISTS) { + LOG.info("Initial parent partition already exists in metadata table, skipping insertion."); + } else { + throw e; } - daoFactory.getPartitionMetadataDao().insert(parentPartitions); } } } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/InitializeDoFnTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/InitializeDoFnTest.java index 6aa3878bef12..611b74d485ca 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/InitializeDoFnTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/InitializeDoFnTest.java @@ -21,12 +21,17 @@ import static org.junit.Assert.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.google.cloud.Timestamp; +import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.SpannerExceptionFactory; import com.google.cloud.spanner.Struct; import java.util.Arrays; import java.util.List; @@ -34,10 +39,12 @@ import org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.DaoFactory; import org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataAdminDao; import org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao; +import org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao.InTransactionContext; import org.apache.beam.sdk.io.gcp.spanner.changestreams.mapper.MapperFactory; import org.apache.beam.sdk.io.gcp.spanner.changestreams.mapper.PartitionMetadataMapper; import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.InitialPartition; import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata; +import org.apache.beam.sdk.io.gcp.spanner.changestreams.util.TestTransactionAnswer; import org.apache.beam.sdk.transforms.DoFn.OutputReceiver; import org.junit.Before; import org.junit.Test; @@ -49,6 +56,7 @@ public class InitializeDoFnTest { private MapperFactory mapperFactory; private OutputReceiver receiver; private PartitionMetadataDao partitionMetadataDao; + private InTransactionContext transaction; private PartitionMetadataAdminDao partitionMetadataAdminDao; private PartitionMetadataMapper partitionMetadataMapper; private InitializeDoFn initializeDoFn; @@ -59,8 +67,11 @@ public void setUp() { mapperFactory = mock(MapperFactory.class); receiver = mock(OutputReceiver.class); partitionMetadataDao = mock(PartitionMetadataDao.class); + transaction = mock(InTransactionContext.class); partitionMetadataAdminDao = mock(PartitionMetadataAdminDao.class); partitionMetadataMapper = mock(PartitionMetadataMapper.class); + when(partitionMetadataDao.runInTransaction(any(), any())) + .thenAnswer(new TestTransactionAnswer(transaction)); initializeDoFn = new InitializeDoFn( daoFactory, @@ -72,29 +83,101 @@ public void setUp() { @Test public void testInitialize() { + String composedToken = + PartitionMetadataDao.composePartitionTokenWithTvfName( + InitialPartition.PARTITION_TOKEN, ChangeStreamsConstants.DEFAULT_TVF_NAME); when(daoFactory.getPartitionMetadataDao()).thenReturn(partitionMetadataDao); when(partitionMetadataDao.tableExists()).thenReturn(false); when(daoFactory.getPartitionMetadataAdminDao()).thenReturn(partitionMetadataAdminDao); doNothing().when(partitionMetadataAdminDao).createPartitionMetadataTable(); - when(partitionMetadataDao.insert(any(PartitionMetadata.class))) - .thenReturn(Timestamp.ofTimeMicroseconds(1L)); - when(partitionMetadataDao.getPartition( - PartitionMetadataDao.composePartitionTokenWithTvfName( - InitialPartition.PARTITION_TOKEN, ChangeStreamsConstants.DEFAULT_TVF_NAME))) - .thenReturn(Struct.newBuilder().build()); + when(transaction.getPartition(composedToken)).thenReturn(null); + ArgumentCaptor> listCaptor = ArgumentCaptor.forClass(List.class); + when(partitionMetadataDao.getPartition(composedToken)).thenReturn(Struct.newBuilder().build()); when(mapperFactory.partitionMetadataMapper()).thenReturn(partitionMetadataMapper); - when(partitionMetadataMapper.from(any())).thenReturn(mock(PartitionMetadata.class)); + PartitionMetadata expectedPartition = mock(PartitionMetadata.class); + when(partitionMetadataMapper.from(any())).thenReturn(expectedPartition); + initializeDoFn.processElement(receiver); + verify(daoFactory, times(2)).getPartitionMetadataDao(); verify(daoFactory, times(1)).getPartitionMetadataAdminDao(); - verify(partitionMetadataDao, times(1)).insert(any(PartitionMetadata.class)); - verify(partitionMetadataDao, times(1)) - .getPartition( - PartitionMetadataDao.composePartitionTokenWithTvfName( - InitialPartition.PARTITION_TOKEN, ChangeStreamsConstants.DEFAULT_TVF_NAME)); + verify(transaction, times(1)).getPartition(composedToken); + verify(transaction, times(1)).insert(listCaptor.capture()); + assertEquals(1, listCaptor.getValue().size()); + assertEquals( + InitialPartition.PARTITION_TOKEN, listCaptor.getValue().get(0).getPartitionToken()); + verify(partitionMetadataDao, times(1)).getPartition(composedToken); verify(partitionMetadataDao, times(1)).tableExists(); verify(mapperFactory, times(1)).partitionMetadataMapper(); verify(partitionMetadataMapper, times(1)).from(any()); + verify(receiver, times(1)).output(expectedPartition); + } + + @Test + public void testInitializeWhenFakeParentPartitionAlreadyExists() { + String composedToken = + PartitionMetadataDao.composePartitionTokenWithTvfName( + InitialPartition.PARTITION_TOKEN, ChangeStreamsConstants.DEFAULT_TVF_NAME); + when(daoFactory.getPartitionMetadataDao()).thenReturn(partitionMetadataDao); + when(partitionMetadataDao.tableExists()).thenReturn(false); + when(daoFactory.getPartitionMetadataAdminDao()).thenReturn(partitionMetadataAdminDao); + doNothing().when(partitionMetadataAdminDao).createPartitionMetadataTable(); + when(transaction.getPartition(composedToken)).thenReturn(Struct.newBuilder().build()); + when(partitionMetadataDao.getPartition(composedToken)).thenReturn(Struct.newBuilder().build()); + when(mapperFactory.partitionMetadataMapper()).thenReturn(partitionMetadataMapper); + PartitionMetadata expectedPartition = mock(PartitionMetadata.class); + when(partitionMetadataMapper.from(any())).thenReturn(expectedPartition); + + initializeDoFn.processElement(receiver); + + verify(transaction, times(1)).getPartition(composedToken); + verify(transaction, never()).insert(any(List.class)); + verify(transaction, never()).insert(any(PartitionMetadata.class)); + verify(receiver, times(1)).output(expectedPartition); + } + + @Test + public void testInitializeIgnoresAlreadyExistsSpannerException() { + String composedToken = + PartitionMetadataDao.composePartitionTokenWithTvfName( + InitialPartition.PARTITION_TOKEN, ChangeStreamsConstants.DEFAULT_TVF_NAME); + when(daoFactory.getPartitionMetadataDao()).thenReturn(partitionMetadataDao); + when(partitionMetadataDao.tableExists()).thenReturn(false); + when(daoFactory.getPartitionMetadataAdminDao()).thenReturn(partitionMetadataAdminDao); + doNothing().when(partitionMetadataAdminDao).createPartitionMetadataTable(); + doThrow( + SpannerExceptionFactory.newSpannerException( + ErrorCode.ALREADY_EXISTS, "Row [Parent0] already exists")) + .when(partitionMetadataDao) + .runInTransaction(any(), any()); + when(partitionMetadataDao.getPartition(composedToken)).thenReturn(Struct.newBuilder().build()); + when(mapperFactory.partitionMetadataMapper()).thenReturn(partitionMetadataMapper); + PartitionMetadata expectedPartition = mock(PartitionMetadata.class); + when(partitionMetadataMapper.from(any())).thenReturn(expectedPartition); + + initializeDoFn.processElement(receiver); + + verify(receiver, times(1)).output(expectedPartition); + } + + @Test + public void testInitializeRethrowsNonAlreadyExistsSpannerException() { + when(daoFactory.getPartitionMetadataDao()).thenReturn(partitionMetadataDao); + when(partitionMetadataDao.tableExists()).thenReturn(false); + when(daoFactory.getPartitionMetadataAdminDao()).thenReturn(partitionMetadataAdminDao); + doNothing().when(partitionMetadataAdminDao).createPartitionMetadataTable(); + doThrow( + SpannerExceptionFactory.newSpannerException( + ErrorCode.PERMISSION_DENIED, "Permission denied")) + .when(partitionMetadataDao) + .runInTransaction(any(), any()); + + try { + initializeDoFn.processElement(receiver); + fail(); + } catch (SpannerException e) { + assertEquals(ErrorCode.PERMISSION_DENIED, e.getErrorCode()); + } } @Test @@ -103,8 +186,6 @@ public void testInitializeWithNoPartition() { when(partitionMetadataDao.tableExists()).thenReturn(false); when(daoFactory.getPartitionMetadataAdminDao()).thenReturn(partitionMetadataAdminDao); doNothing().when(partitionMetadataAdminDao).createPartitionMetadataTable(); - when(partitionMetadataDao.insert(any(PartitionMetadata.class))) - .thenReturn(Timestamp.ofTimeMicroseconds(1L)); when(mapperFactory.partitionMetadataMapper()).thenReturn(partitionMetadataMapper); when(partitionMetadataMapper.from(any())).thenReturn(mock(PartitionMetadata.class)); try { @@ -123,8 +204,6 @@ public void testInitializeWithTvfNameList() { when(daoFactory.getPartitionMetadataAdminDao()).thenReturn(partitionMetadataAdminDao); doNothing().when(partitionMetadataAdminDao).createPartitionMetadataTable(); ArgumentCaptor> listCaptor = ArgumentCaptor.forClass(List.class); - when(partitionMetadataDao.insert(listCaptor.capture())) - .thenReturn(Timestamp.ofTimeMicroseconds(1L)); when(partitionMetadataDao.getPartition( PartitionMetadataDao.composePartitionTokenWithTvfName( InitialPartition.PARTITION_TOKEN, "tvf1"))) @@ -136,11 +215,11 @@ public void testInitializeWithTvfNameList() { verify(daoFactory, times(2)).getPartitionMetadataDao(); verify(daoFactory, times(1)).getPartitionMetadataAdminDao(); + verify(transaction, times(1)).insert(listCaptor.capture()); assertEquals(2, listCaptor.getValue().size()); assertEquals("tvf1", listCaptor.getValue().get(0).getTvfName()); assertEquals("tvf2", listCaptor.getValue().get(1).getTvfName()); - verify(partitionMetadataDao, times(1)).insert(listCaptor.capture()); verify(partitionMetadataDao, times(1)) .getPartition( PartitionMetadataDao.composePartitionTokenWithTvfName( @@ -151,15 +230,39 @@ public void testInitializeWithTvfNameList() { } @Test - public void testInitializeWithNoPartitionWithTvfNameList() { + public void testInitializeWithTvfNameListWhenSomeAlreadyExist() { + String composedToken1 = + PartitionMetadataDao.composePartitionTokenWithTvfName( + InitialPartition.PARTITION_TOKEN, "tvf1"); + String composedToken2 = + PartitionMetadataDao.composePartitionTokenWithTvfName( + InitialPartition.PARTITION_TOKEN, "tvf2"); when(daoFactory.getTvfNameList()).thenReturn(Arrays.asList("tvf1", "tvf2")); when(daoFactory.getPartitionMetadataDao()).thenReturn(partitionMetadataDao); when(partitionMetadataDao.tableExists()).thenReturn(false); when(daoFactory.getPartitionMetadataAdminDao()).thenReturn(partitionMetadataAdminDao); doNothing().when(partitionMetadataAdminDao).createPartitionMetadataTable(); + when(transaction.getPartition(composedToken1)).thenReturn(Struct.newBuilder().build()); + when(transaction.getPartition(composedToken2)).thenReturn(null); ArgumentCaptor> listCaptor = ArgumentCaptor.forClass(List.class); - when(partitionMetadataDao.insert(listCaptor.capture())) - .thenReturn(Timestamp.ofTimeMicroseconds(1L)); + when(partitionMetadataDao.getPartition(composedToken1)).thenReturn(Struct.newBuilder().build()); + when(mapperFactory.partitionMetadataMapper()).thenReturn(partitionMetadataMapper); + when(partitionMetadataMapper.from(any())).thenReturn(mock(PartitionMetadata.class)); + + initializeDoFn.processElement(receiver); + + verify(transaction, times(1)).insert(listCaptor.capture()); + assertEquals(1, listCaptor.getValue().size()); + assertEquals("tvf2", listCaptor.getValue().get(0).getTvfName()); + } + + @Test + public void testInitializeWithNoPartitionWithTvfNameList() { + when(daoFactory.getTvfNameList()).thenReturn(Arrays.asList("tvf1", "tvf2")); + when(daoFactory.getPartitionMetadataDao()).thenReturn(partitionMetadataDao); + when(partitionMetadataDao.tableExists()).thenReturn(false); + when(daoFactory.getPartitionMetadataAdminDao()).thenReturn(partitionMetadataAdminDao); + doNothing().when(partitionMetadataAdminDao).createPartitionMetadataTable(); when(mapperFactory.partitionMetadataMapper()).thenReturn(partitionMetadataMapper); when(partitionMetadataMapper.from(any())).thenReturn(mock(PartitionMetadata.class));