Skip to content

[Spanner Change Stream] Make InitializeDoFn fake parent partition creation idempotent #40808

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
jiangzzhu wants to merge 1 commit into from
Closed

[Spanner Change Stream] Make InitializeDoFn fake parent partition creation idempotent #40808

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 38 additions & 8 deletions ...m/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/InitializeDoFn.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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
Expand All @@ -40,6 +44,7 @@
public class InitializeDoFn extends DoFn<byte[], PartitionMetadata> implements Serializable {

private static final long serialVersionUID = -8921188388649003102L;
private static final Logger LOG = LoggerFactory.getLogger(InitializeDoFn.class);

private final long heartbeatMillis;

Expand Down Expand Up @@ -117,10 +122,11 @@ public void processElement(OutputReceiver<PartitionMetadata> receiver) {
* specified in {@link InitializeDoFn#DEFAULT_HEARTBEAT_MILLIS}.
*/
private void createFakeParentPartition() {
List<PartitionMetadata> 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)
Expand All @@ -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<PartitionMetadata> parentPartitions = new ArrayList<>();
for (String tvfName : daoFactory.getTvfNameList()) {
PartitionMetadata parentPartition =
parentPartitions.add(
PartitionMetadata.newBuilder()
.setPartitionToken(InitialPartition.PARTITION_TOKEN)
.setTvfName(tvfName)
Expand All @@ -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<PartitionMetadata> 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);
}
}
}
143 changes: 123 additions & 20 deletions ...c/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/InitializeDoFnTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,23 +21,30 @@
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;
import org.apache.beam.sdk.io.gcp.spanner.changestreams.ChangeStreamsConstants;
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;
Expand All @@ -49,6 +56,7 @@ public class InitializeDoFnTest {
private MapperFactory mapperFactory;
private OutputReceiver<PartitionMetadata> receiver;
private PartitionMetadataDao partitionMetadataDao;
private InTransactionContext transaction;
private PartitionMetadataAdminDao partitionMetadataAdminDao;
private PartitionMetadataMapper partitionMetadataMapper;
private InitializeDoFn initializeDoFn;
Expand All @@ -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,
Expand All @@ -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<List<PartitionMetadata>> 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
Expand All @@ -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 {
Expand All @@ -123,8 +204,6 @@ public void testInitializeWithTvfNameList() {
when(daoFactory.getPartitionMetadataAdminDao()).thenReturn(partitionMetadataAdminDao);
doNothing().when(partitionMetadataAdminDao).createPartitionMetadataTable();
ArgumentCaptor<List<PartitionMetadata>> listCaptor = ArgumentCaptor.forClass(List.class);
when(partitionMetadataDao.insert(listCaptor.capture()))
.thenReturn(Timestamp.ofTimeMicroseconds(1L));
when(partitionMetadataDao.getPartition(
PartitionMetadataDao.composePartitionTokenWithTvfName(
InitialPartition.PARTITION_TOKEN, "tvf1")))
Expand All @@ -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(
Expand All @@ -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<List<PartitionMetadata>> 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));

Expand Down
Loading