Fix secrets rotation handling in Rds source (opensearch-project#5465)

* Fix secrets rotation handling

Signed-off-by: Hai Yan <oeyh@amazon.com>

* Update message types

Signed-off-by: Hai Yan <oeyh@amazon.com>

---------

Signed-off-by: Hai Yan <oeyh@amazon.com>

Author: oeyh

data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/model/MessageType.java

@@ -10,28 +10,29 @@
 
 package org.opensearch.dataprepper.plugins.source.rds.model;
 
+import java.util.Arrays;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 public enum MessageType {
     BEGIN('B'),
-    RELATION('R'),
-    INSERT('I'),
-    UPDATE('U'),
-    DELETE('D'),
     COMMIT('C'),
-    TYPE('Y');
+    DELETE('D'),
+    INSERT('I'),
+    MESSAGE('M'),
+    ORIGIN('O'),
+    RELATION('R'),
+    TRUNCATE('T'),
+    TYPE('Y'),
+    UPDATE('U');
 
     private final char value;
 
-    private static final Map<Character, MessageType> MESSAGE_TYPE_MAP = Map.of(
-            BEGIN.getValue(), BEGIN,
-            RELATION.getValue(), RELATION,
-            INSERT.getValue(), INSERT,
-            UPDATE.getValue(), UPDATE,
-            DELETE.getValue(), DELETE,
-            COMMIT.getValue(), COMMIT,
-            TYPE.getValue(), TYPE
-    );
+    private static final Map<Character, MessageType> MESSAGE_TYPE_MAP = Arrays.stream(values())
+            .collect(Collectors.toMap(
+                    messageType -> messageType.value,
+                    messageType -> messageType
+            ));
 
     MessageType(char value) {
         this.value = value;

data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/stream/LogicalReplicationClient.java

@@ -35,6 +35,7 @@ public class LogicalReplicationClient implements ReplicationLogClient {
     private LogSequenceNumber startLsn;
     private LogicalReplicationEventProcessor eventProcessor;
 
+    private PGReplicationStream stream = null;
     private volatile boolean disconnectRequested = false;
 
     public LogicalReplicationClient(final ConnectionManager connectionManager,
@@ -48,7 +49,6 @@ public LogicalReplicationClient(final ConnectionManager connectionManager,
     @Override
     public void connect() {
         LOG.debug("Start connecting logical replication stream. ");
-        PGReplicationStream stream;
         try (Connection conn = connectionManager.getConnection()) {
             PGConnection pgConnection = conn.unwrap(PGConnection.class);
 
@@ -85,20 +85,17 @@ public void connect() {
                         stream.setAppliedLSN(lsn);
                     } catch (Exception e) {
                         LOG.error("Exception while processing Postgres replication stream. ", e);
-                        stream.close();
-                        LOG.debug("Replication stream closed.");
+                        closeStream();
                         throw e;
                     }
                 }
             }
 
-            stream.close();
-            LOG.debug("Replication stream closed.");
+            closeStream();
 
             disconnectRequested = false;
-            LOG.debug("Replication stream closed successfully.");
         } catch (Exception e) {
-            LOG.error("Exception while creating Postgres replication stream. ", e);
+            LOG.error("Exception while creating or processing Postgres replication stream. ", e);
             throw new RuntimeException(e);
         }
     }
@@ -108,6 +105,8 @@ public void disconnect() {
         disconnectRequested = true;
         LOG.debug("Requested to disconnect logical replication stream.");
 
+        closeStream();
+
         if (eventProcessor != null) {
             eventProcessor.stopCheckpointManager();
             LOG.debug("Stopped checkpoint manager.");
@@ -121,4 +120,15 @@ public void setEventProcessor(LogicalReplicationEventProcessor eventProcessor) {
     public void setStartLsn(LogSequenceNumber startLsn) {
         this.startLsn = startLsn;
     }
+
+    private void closeStream() {
+        if (stream != null && !stream.isClosed()) {
+            try {
+                stream.close();
+                LOG.debug("Replication stream closed.");
+            } catch (Exception e) {
+                LOG.error("Exception while closing replication stream. ", e);
+            }
+        }
+    }
 }

data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/stream/LogicalReplicationEventProcessor.java

@@ -163,7 +163,16 @@ public void process(ByteBuffer msg) {
         // If it's a RELATION, update table metadata map
         // If it's INSERT/UPDATE/DELETE, prepare events
         // If it's a COMMIT, convert all prepared events and send to buffer
-        MessageType messageType = MessageType.from((char) msg.get());
+        MessageType messageType;
+        char typeChar = '\0';
+        try {
+            typeChar = (char) msg.get();
+            messageType = MessageType.from(typeChar);
+        } catch (IllegalArgumentException e) {
+            LOG.warn("Unknown message type {} received from stream. Skipping.", typeChar);
+            return;
+        }
+
         switch (messageType) {
             case BEGIN:
                 handleMessageWithRetries(msg, this::processBeginMessage, messageType);
@@ -187,7 +196,7 @@ public void process(ByteBuffer msg) {
                 handleMessageWithRetries(msg, this::processTypeMessage, messageType);
                 break;
             default:
-                throw new IllegalArgumentException("Replication message type [" + messageType + "] is not supported. ");
+                LOG.debug("Replication message type '{}' is not supported. Skipping.", messageType);
         }
     }
 

data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/stream/ReplicationLogClientFactory.java

@@ -26,9 +26,9 @@
 
 public class ReplicationLogClientFactory {
 
-    private final RdsSourceConfig sourceConfig;
     private final RdsClient rdsClient;
     private final DbMetadata dbMetadata;
+    private RdsSourceConfig sourceConfig;
     private String username;
     private String password;
     private SSLMode sslMode = SSLMode.REQUIRED;
@@ -82,9 +82,10 @@ public void setSSLMode(SSLMode sslMode) {
         this.sslMode = sslMode;
     }
 
-    public void setCredentials(String username, String password) {
-        this.username = username;
-        this.password = password;
+    public void updateCredentials(RdsSourceConfig sourceConfig) {
+        this.sourceConfig = sourceConfig;
+        this.username = sourceConfig.getAuthenticationConfig().getUsername();
+        this.password = sourceConfig.getAuthenticationConfig().getPassword();
     }
 
     private String getDatabaseName(List<String> tableNames) {

data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/stream/StreamWorker.java

@@ -66,11 +66,15 @@ public void processStream(final StreamPartition streamPartition) {
             sourceCoordinator.giveUpPartition(streamPartition);
             throw new RuntimeException(e);
         } finally {
-            try {
-                replicationLogClient.disconnect();
-            } catch (Exception e) {
-                LOG.error("Binary log client failed to disconnect.", e);
-            }
+            shutdown();
+        }
+    }
+
+    public void shutdown() {
+        try {
+            replicationLogClient.disconnect();
+        } catch (Exception e) {
+            LOG.error("Replication log client failed to disconnect.", e);
         }
     }
 

data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/stream/StreamWorkerTaskRefresher.java

@@ -48,6 +48,7 @@ public class StreamWorkerTaskRefresher implements PluginConfigObserver<RdsSource
 
     private ExecutorService executorService;
     private RdsSourceConfig currentSourceConfig;
+    private StreamWorker streamWorker;
 
     public StreamWorkerTaskRefresher(final EnhancedSourceCoordinator sourceCoordinator,
                                      final StreamPartition streamPartition,
@@ -96,10 +97,10 @@ public void update(RdsSourceConfig sourceConfig) {
             LOG.info("Database credentials were updated. Refreshing stream worker...");
             credentialsChangeCounter.increment();
             try {
+                streamWorker.shutdown();
                 executorService.shutdownNow();
                 executorService = executorServiceSupplier.get();
-                replicationLogClientFactory.setCredentials(
-                        sourceConfig.getAuthenticationConfig().getUsername(), sourceConfig.getAuthenticationConfig().getPassword());
+                replicationLogClientFactory.updateCredentials(sourceConfig);
 
                 refreshTask(sourceConfig);
 
@@ -132,7 +133,7 @@ private void refreshTask(RdsSourceConfig sourceConfig) {
                     streamPartition, sourceConfig, buffer, s3Prefix, pluginMetrics, logicalReplicationClient,
                     streamCheckpointer, acknowledgementSetManager));
         }
-        final StreamWorker streamWorker = StreamWorker.create(sourceCoordinator, replicationLogClient, pluginMetrics);
+        streamWorker = StreamWorker.create(sourceCoordinator, replicationLogClient, pluginMetrics);
         executorService.submit(() -> streamWorker.processStream(streamPartition));
     }
 

data-prepper-plugins/rds-source/src/test/java/org/opensearch/dataprepper/plugins/source/rds/stream/LogicalReplicationClientTest.java

@@ -119,6 +119,7 @@ void test_disconnect() throws SQLException, InterruptedException {
         when(logicalStreamBuilder.start()).thenReturn(stream);
         when(stream.readPending()).thenReturn(message).thenReturn(null);
         when(stream.getLastReceiveLSN()).thenReturn(lsn);
+        when(stream.isClosed()).thenReturn(false, true);
 
         final ExecutorService executorService = Executors.newSingleThreadExecutor();
         executorService.submit(() -> logicalReplicationClient.connect());
@@ -155,6 +156,7 @@ void test_connect_disconnect_cycles() throws SQLException, InterruptedException
         when(logicalStreamBuilder.start()).thenReturn(stream);
         when(stream.readPending()).thenReturn(message).thenReturn(null);
         when(stream.getLastReceiveLSN()).thenReturn(lsn);
+        when(stream.isClosed()).thenReturn(false, true);
 
         // First connect
         final ExecutorService executorService = Executors.newSingleThreadExecutor();
@@ -174,6 +176,7 @@ void test_connect_disconnect_cycles() throws SQLException, InterruptedException
 
         // Second connect
         when(stream.readPending()).thenReturn(message).thenReturn(null);
+        when(stream.isClosed()).thenReturn(false, true);
         executorService.submit(() -> logicalReplicationClient.connect());
         await().atMost(Duration.ofSeconds(1))
                 .untilAsserted(() -> verify(eventProcessor, times(2)).process(message));

data-prepper-plugins/rds-source/src/test/java/org/opensearch/dataprepper/plugins/source/rds/stream/LogicalReplicationEventProcessorTest.java

@@ -14,6 +14,8 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.Answers;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
@@ -30,7 +32,7 @@
 import java.util.Random;
 import java.util.UUID;
 
-import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.spy;
@@ -149,13 +151,21 @@ void test_correct_process_method_invoked_for_type_message() {
         verify(objectUnderTest).processTypeMessage(message);
     }
 
+    @ParameterizedTest
+    @ValueSource(chars = {'M', 'O', 'T'})
+    void test_unsupported_message_type_not_throw_exception(char typeChar) {
+        setMessageType(MessageType.from(typeChar));
+
+        assertDoesNotThrow(() -> objectUnderTest.process(message));
+    }
+
     @Test
-    void test_unsupported_message_type_throws_exception() {
+    void test_unknown_message_type_not_throw_exception() {
         message = ByteBuffer.allocate(1);
         message.put((byte) 'A');
         message.flip();
 
-        assertThrows(IllegalArgumentException.class, () -> objectUnderTest.process(message));
+        assertDoesNotThrow(() -> objectUnderTest.process(message));
     }
 
     @Test

data-prepper-plugins/rds-source/src/test/java/org/opensearch/dataprepper/plugins/source/rds/stream/StreamWorkerTest.java

@@ -89,6 +89,12 @@ void test_processStream_without_current_binlog_coordinates() throws IOException
         verify(binlogClientWrapper).connect();
     }
 
+    @Test
+    void test_shutdown() throws IOException {
+        streamWorker.shutdown();
+        verify(binlogClientWrapper).disconnect();
+    }
+
     private StreamWorker createObjectUnderTest() {
         return new StreamWorker(sourceCoordinator, binlogClientWrapper, pluginMetrics);
     }