Skip to content

Commit 5d6bed7

Browse files
oeyhAlekhya Parisha
oeyh
authored and
Alekhya Parisha
committed
Change slot naming (opensearch-project#5699)
Signed-off-by: Hai Yan <oeyh@amazon.com>
1 parent 376f77c commit 5d6bed7

7 files changed

Lines changed: 97 additions & 16 deletions

File tree

‎data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/RdsService.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import org.opensearch.dataprepper.metrics.PluginMetrics;
1010
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager;
1111
import org.opensearch.dataprepper.model.buffer.Buffer;
12+
import org.opensearch.dataprepper.model.configuration.PipelineDescription;
1213
import org.opensearch.dataprepper.model.event.Event;
1314
import org.opensearch.dataprepper.model.event.EventFactory;
1415
import org.opensearch.dataprepper.model.plugin.PluginConfigObservable;
@@ -64,6 +65,7 @@ public class RdsService {
6465
private final RdsSourceConfig sourceConfig;
6566
private final AcknowledgementSetManager acknowledgementSetManager;
6667
private final PluginConfigObservable pluginConfigObservable;
68+
private final PipelineDescription pipelineDescription;
6769
private ExecutorService executor;
6870
private LeaderScheduler leaderScheduler;
6971
private ExportScheduler exportScheduler;
@@ -77,13 +79,15 @@ public RdsService(final EnhancedSourceCoordinator sourceCoordinator,
7779
final ClientFactory clientFactory,
7880
final PluginMetrics pluginMetrics,
7981
final AcknowledgementSetManager acknowledgementSetManager,
80-
final PluginConfigObservable pluginConfigObservable) {
82+
final PluginConfigObservable pluginConfigObservable,
83+
final PipelineDescription pipelineDescription) {
8184
this.sourceCoordinator = sourceCoordinator;
8285
this.eventFactory = eventFactory;
8386
this.pluginMetrics = pluginMetrics;
8487
this.sourceConfig = sourceConfig;
8588
this.acknowledgementSetManager = acknowledgementSetManager;
8689
this.pluginConfigObservable = pluginConfigObservable;
90+
this.pipelineDescription = pipelineDescription;
8791

8892
rdsClient = clientFactory.buildRdsClient();
8993
s3Client = clientFactory.buildS3Client();
@@ -109,7 +113,7 @@ public void start(Buffer<Record<Event>> buffer) {
109113
DbTableMetadata dbTableMetadata = getDbTableMetadata(dbMetadata, schemaManager);
110114

111115
leaderScheduler = new LeaderScheduler(
112-
sourceCoordinator, sourceConfig, s3PathPrefix, schemaManager, dbTableMetadata);
116+
sourceCoordinator, sourceConfig, s3PathPrefix, schemaManager, dbTableMetadata, pipelineDescription.getPipelineName());
113117
runnableList.add(leaderScheduler);
114118

115119
if (sourceConfig.isExportEnabled()) {

‎data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/RdsSource.java‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin;
1212
import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor;
1313
import org.opensearch.dataprepper.model.buffer.Buffer;
14+
import org.opensearch.dataprepper.model.configuration.PipelineDescription;
1415
import org.opensearch.dataprepper.model.event.Event;
1516
import org.opensearch.dataprepper.model.event.EventFactory;
1617
import org.opensearch.dataprepper.model.plugin.PluginConfigObservable;
@@ -39,6 +40,7 @@ public class RdsSource implements Source<Record<Event>>, UsesEnhancedSourceCoord
3940
private final EventFactory eventFactory;
4041
private final AcknowledgementSetManager acknowledgementSetManager;
4142
private final PluginConfigObservable pluginConfigObservable;
43+
private final PipelineDescription pipelineDescription;
4244
private EnhancedSourceCoordinator sourceCoordinator;
4345
private RdsService rdsService;
4446

@@ -48,12 +50,14 @@ public RdsSource(final PluginMetrics pluginMetrics,
4850
final EventFactory eventFactory,
4951
final AwsCredentialsSupplier awsCredentialsSupplier,
5052
final AcknowledgementSetManager acknowledgementSetManager,
51-
final PluginConfigObservable pluginConfigObservable) {
53+
final PluginConfigObservable pluginConfigObservable,
54+
final PipelineDescription pipelineDescription) {
5255
this.pluginMetrics = pluginMetrics;
5356
this.sourceConfig = sourceConfig;
5457
this.eventFactory = eventFactory;
5558
this.acknowledgementSetManager = acknowledgementSetManager;
5659
this.pluginConfigObservable = pluginConfigObservable;
60+
this.pipelineDescription = pipelineDescription;
5761

5862
clientFactory = new ClientFactory(awsCredentialsSupplier, sourceConfig);
5963
}
@@ -64,7 +68,8 @@ public void start(Buffer<Record<Event>> buffer) {
6468
Objects.requireNonNull(sourceCoordinator);
6569
sourceCoordinator.createPartition(new LeaderPartition());
6670

67-
rdsService = new RdsService(sourceCoordinator, sourceConfig, eventFactory, clientFactory, pluginMetrics, acknowledgementSetManager, pluginConfigObservable);
71+
rdsService = new RdsService(sourceCoordinator, sourceConfig, eventFactory, clientFactory, pluginMetrics,
72+
acknowledgementSetManager, pluginConfigObservable, pipelineDescription);
6873

6974
LOG.info("Start RDS service");
7075
rdsService.start(buffer);

‎data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/leader/LeaderScheduler.java‎

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ public class LeaderScheduler implements Runnable {
4646
private final String s3Prefix;
4747
private final SchemaManager schemaManager;
4848
private final DbTableMetadata dbTableMetadata;
49+
private final String pipelineName;
4950

5051
private LeaderPartition leaderPartition;
5152
private List<String> tableNames;
@@ -56,12 +57,14 @@ public LeaderScheduler(final EnhancedSourceCoordinator sourceCoordinator,
5657
final RdsSourceConfig sourceConfig,
5758
final String s3Prefix,
5859
final SchemaManager schemaManager,
59-
final DbTableMetadata dbTableMetadata) {
60+
final DbTableMetadata dbTableMetadata,
61+
final String pipelineName) {
6062
this.sourceCoordinator = sourceCoordinator;
6163
this.sourceConfig = sourceConfig;
6264
this.s3Prefix = s3Prefix;
6365
this.schemaManager = schemaManager;
6466
this.dbTableMetadata = dbTableMetadata;
67+
this.pipelineName = pipelineName;
6568
tableNames = new ArrayList<>(dbTableMetadata.getTableColumnDataTypeMap().keySet());
6669
}
6770

@@ -180,8 +183,9 @@ private void createStreamPartition(RdsSourceConfig sourceConfig) {
180183
} else {
181184
// Postgres
182185
// Create replication slot, which will mark the starting point for stream
183-
final String publicationName = generatePublicationName();
184-
final String slotName = generateReplicationSlotName();
186+
final String suffix = UUID.randomUUID().toString().substring(0, 8);
187+
final String publicationName = generatePublicationName(suffix);
188+
final String slotName = generateReplicationSlotName(suffix);
185189
((PostgresSchemaManager)schemaManager).createLogicalReplicationSlot(tableNames, publicationName, slotName);
186190
final PostgresStreamState postgresStreamState = new PostgresStreamState();
187191
postgresStreamState.setPublicationName(publicationName);
@@ -199,11 +203,17 @@ private Optional<BinlogCoordinate> getCurrentBinlogPosition() {
199203
return binlogCoordinate;
200204
}
201205

202-
private String generatePublicationName() {
203-
return "data_prepper_publication_" + UUID.randomUUID().toString().substring(0, 8);
206+
private String generatePublicationName(final String suffix) {
207+
return "data_prepper_" + getPipelineName() + "_pub_" + suffix;
204208
}
205209

206-
private String generateReplicationSlotName() {
207-
return "data_prepper_slot_" + UUID.randomUUID().toString().substring(0, 8);
210+
private String generateReplicationSlotName(final String suffix) {
211+
return "data_prepper_" + getPipelineName() + "_slot_" + suffix;
212+
}
213+
214+
private String getPipelineName() {
215+
// Shorten the name (if needed) and replace any invalid characters with underscores
216+
final String shortenedPipelineName = pipelineName.length() <= 16 ? pipelineName : pipelineName.substring(0, 16);
217+
return shortenedPipelineName.replaceAll("[^a-zA-Z0-9_]", "_");
208218
}
209219
}

‎data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/schema/PostgresSchemaManager.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ public void createLogicalReplicationSlot(final List<String> tableNames, final St
6565
try {
6666
PreparedStatement statement = conn.prepareStatement(createPublicationStatement);
6767
statement.executeUpdate();
68+
LOG.info("Publication {} created successfully. ", publicationName);
6869
} catch (Exception e) {
6970
LOG.error("Failed to create publication: {}", e.getMessage());
7071
throw e;

‎data-prepper-plugins/rds-source/src/test/java/org/opensearch/dataprepper/plugins/source/rds/RdsServiceTest.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import org.opensearch.dataprepper.metrics.PluginMetrics;
1818
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager;
1919
import org.opensearch.dataprepper.model.buffer.Buffer;
20+
import org.opensearch.dataprepper.model.configuration.PipelineDescription;
2021
import org.opensearch.dataprepper.model.event.Event;
2122
import org.opensearch.dataprepper.model.event.EventFactory;
2223
import org.opensearch.dataprepper.model.plugin.PluginConfigObservable;
@@ -95,6 +96,9 @@ class RdsServiceTest {
9596
@Mock
9697
private PluginConfigObservable pluginConfigObservable;
9798

99+
@Mock
100+
private PipelineDescription pipelineDescription;
101+
98102
@BeforeEach
99103
void setUp() {
100104
when(clientFactory.buildRdsClient()).thenReturn(rdsClient);
@@ -206,6 +210,7 @@ private void prepareMocks() {
206210
}
207211

208212
private RdsService createObjectUnderTest() {
209-
return new RdsService(sourceCoordinator, sourceConfig, eventFactory, clientFactory, pluginMetrics, acknowledgementSetManager, pluginConfigObservable);
213+
return new RdsService(sourceCoordinator, sourceConfig, eventFactory, clientFactory, pluginMetrics,
214+
acknowledgementSetManager, pluginConfigObservable, pipelineDescription);
210215
}
211216
}

‎data-prepper-plugins/rds-source/src/test/java/org/opensearch/dataprepper/plugins/source/rds/RdsSourceTest.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier;
1414
import org.opensearch.dataprepper.metrics.PluginMetrics;
1515
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager;
16+
import org.opensearch.dataprepper.model.configuration.PipelineDescription;
1617
import org.opensearch.dataprepper.model.event.EventFactory;
1718
import org.opensearch.dataprepper.model.plugin.PluginConfigObservable;
1819
import org.opensearch.dataprepper.plugins.source.rds.configuration.AwsAuthenticationConfig;
@@ -45,6 +46,9 @@ class RdsSourceTest {
4546
@Mock
4647
private PluginConfigObservable pluginConfigObservable;
4748

49+
@Mock
50+
private PipelineDescription pipelineDescription;
51+
4852
@BeforeEach
4953
void setUp() {
5054
when(sourceConfig.getAwsAuthenticationConfig()).thenReturn(awsAuthenticationConfig);
@@ -58,6 +62,7 @@ void test_when_buffer_is_null_then_start_throws_exception() {
5862

5963
private RdsSource createObjectUnderTest() {
6064
return new RdsSource(
61-
pluginMetrics, sourceConfig, eventFactory, awsCredentialsSupplier, acknowledgementSetManager, pluginConfigObservable);
65+
pluginMetrics, sourceConfig, eventFactory, awsCredentialsSupplier, acknowledgementSetManager,
66+
pluginConfigObservable, pipelineDescription);
6267
}
6368
}

‎data-prepper-plugins/rds-source/src/test/java/org/opensearch/dataprepper/plugins/source/rds/leader/LeaderSchedulerTest.java‎

Lines changed: 54 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,10 @@
1212
import org.junit.jupiter.params.ParameterizedTest;
1313
import org.junit.jupiter.params.provider.EnumSource;
1414
import org.mockito.Answers;
15+
import org.mockito.ArgumentCaptor;
1516
import org.mockito.Mock;
1617
import org.mockito.junit.jupiter.MockitoExtension;
18+
import org.opensearch.dataprepper.model.configuration.PipelineDescription;
1719
import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator;
1820
import org.opensearch.dataprepper.plugins.source.rds.RdsSourceConfig;
1921
import org.opensearch.dataprepper.plugins.source.rds.configuration.AwsAuthenticationConfig;
@@ -22,9 +24,12 @@
2224
import org.opensearch.dataprepper.plugins.source.rds.coordination.partition.ExportPartition;
2325
import org.opensearch.dataprepper.plugins.source.rds.coordination.partition.GlobalState;
2426
import org.opensearch.dataprepper.plugins.source.rds.coordination.partition.LeaderPartition;
27+
import org.opensearch.dataprepper.plugins.source.rds.coordination.partition.StreamPartition;
2528
import org.opensearch.dataprepper.plugins.source.rds.coordination.state.LeaderProgressState;
29+
import org.opensearch.dataprepper.plugins.source.rds.coordination.state.PostgresStreamState;
2630
import org.opensearch.dataprepper.plugins.source.rds.model.DbTableMetadata;
27-
import org.opensearch.dataprepper.plugins.source.rds.schema.MySqlSchemaManager;
31+
import org.opensearch.dataprepper.plugins.source.rds.schema.PostgresSchemaManager;
32+
import org.opensearch.dataprepper.plugins.source.rds.schema.SchemaManager;
2833

2934
import java.time.Duration;
3035
import java.util.Optional;
@@ -33,6 +38,10 @@
3338
import java.util.concurrent.Executors;
3439

3540
import static org.awaitility.Awaitility.await;
41+
import static org.hamcrest.CoreMatchers.is;
42+
import static org.hamcrest.CoreMatchers.notNullValue;
43+
import static org.hamcrest.CoreMatchers.startsWith;
44+
import static org.hamcrest.MatcherAssert.assertThat;
3645
import static org.mockito.ArgumentMatchers.any;
3746
import static org.mockito.ArgumentMatchers.eq;
3847
import static org.mockito.Mockito.lenient;
@@ -53,7 +62,7 @@ class LeaderSchedulerTest {
5362
private RdsSourceConfig sourceConfig;
5463

5564
@Mock
56-
private MySqlSchemaManager schemaManager;
65+
private SchemaManager schemaManager;
5766

5867
@Mock
5968
private DbTableMetadata dbTableMetadata;
@@ -64,20 +73,26 @@ class LeaderSchedulerTest {
6473
@Mock
6574
private LeaderProgressState leaderProgressState;
6675

76+
@Mock
77+
private PipelineDescription pipelineDescription;
78+
6779
private String s3Prefix;
80+
private String pipelineName;
6881
private LeaderScheduler leaderScheduler;
6982

7083
@BeforeEach
7184
void setUp() {
7285
s3Prefix = UUID.randomUUID().toString();
7386
leaderScheduler = createObjectUnderTest();
87+
pipelineName = UUID.randomUUID().toString();
7488

7589
AwsAuthenticationConfig awsAuthenticationConfig = mock(AwsAuthenticationConfig.class);
7690
lenient().when(awsAuthenticationConfig.getAwsStsRoleArn()).thenReturn(UUID.randomUUID().toString());
7791
lenient().when(sourceConfig.getAwsAuthenticationConfig()).thenReturn(awsAuthenticationConfig);
7892
ExportConfig exportConfig = mock(ExportConfig.class);
7993
lenient().when(exportConfig.getKmsKeyId()).thenReturn(UUID.randomUUID().toString());
8094
lenient().when(sourceConfig.getExport()).thenReturn(exportConfig);
95+
lenient().when(pipelineDescription.getPipelineName()).thenReturn(pipelineName);
8196
}
8297

8398
@Test
@@ -146,7 +161,43 @@ void test_shutDown() {
146161
executorService.shutdownNow();
147162
}
148163

164+
@Test
165+
void leader_node_performs_init_creates_slot_with_expected_name() throws InterruptedException {
166+
final PostgresSchemaManager postgresSchemaManager = mock(PostgresSchemaManager.class);
167+
final String pipelineName = "simple-pipeline";
168+
169+
when(sourceCoordinator.acquireAvailablePartition(LeaderPartition.PARTITION_TYPE)).thenReturn(Optional.of(leaderPartition));
170+
when(leaderPartition.getProgressState()).thenReturn(Optional.of(leaderProgressState));
171+
when(leaderProgressState.isInitialized()).thenReturn(false);
172+
when(sourceConfig.isStreamEnabled()).thenReturn(true);
173+
when(sourceConfig.getEngine()).thenReturn(EngineType.POSTGRES);
174+
175+
final LeaderScheduler leaderScheduler = new LeaderScheduler(sourceCoordinator, sourceConfig, s3Prefix,
176+
postgresSchemaManager, dbTableMetadata, pipelineName);
177+
final ExecutorService executorService = Executors.newSingleThreadExecutor();
178+
executorService.submit(leaderScheduler);
179+
await().atMost(Duration.ofSeconds(1))
180+
.untilAsserted(() -> verify(sourceCoordinator).acquireAvailablePartition(LeaderPartition.PARTITION_TYPE));
181+
Thread.sleep(100);
182+
executorService.shutdownNow();
183+
184+
ArgumentCaptor<StreamPartition> streamPartitionArgumentCaptor = ArgumentCaptor.forClass(StreamPartition.class);
185+
verify(sourceCoordinator).createPartition(any(GlobalState.class));
186+
verify(sourceCoordinator).createPartition(streamPartitionArgumentCaptor.capture());
187+
verify(sourceCoordinator).saveProgressStateForPartition(eq(leaderPartition), any(Duration.class));
188+
189+
final StreamPartition streamPartition = streamPartitionArgumentCaptor.getValue();
190+
assertThat(streamPartition.getProgressState().get().getPostgresStreamState(), notNullValue());
191+
192+
PostgresStreamState postgresStreamState = streamPartition.getProgressState().get().getPostgresStreamState();
193+
final String publicationName = postgresStreamState.getPublicationName();
194+
final String slotName = postgresStreamState.getReplicationSlotName();
195+
assertThat(publicationName, startsWith("data_prepper_simple_pipeline"));
196+
assertThat(slotName, startsWith("data_prepper_simple_pipeline"));
197+
assertThat(publicationName.substring(publicationName.length() - 8), is(slotName.substring(slotName.length() - 8)));
198+
}
199+
149200
private LeaderScheduler createObjectUnderTest() {
150-
return new LeaderScheduler(sourceCoordinator, sourceConfig, s3Prefix, schemaManager, dbTableMetadata);
201+
return new LeaderScheduler(sourceCoordinator, sourceConfig, s3Prefix, schemaManager, dbTableMetadata, pipelineName);
151202
}
152203
}

0 commit comments

Comments
 (0)