Initial PR for stream support for Postgres in Rds source #5310

oeyh · 2025-01-06T15:08:25Z

Description

Initial PR for stream support for Postgres in Rds source. The implementation uses the existing structure for MySQL stream support. The correspondence is as follows:

For MySQL	For Postgres	Notes
`BinaryLogClient`/`BinlogClientWrapper`	`LogicalReplicationClient`	Client for connecting to database and read CDC stream.
`MySqlConnectionManager` (used to be `ConnectionManager`)	`PostgresConnectionManager`	Handles connection to database
`MySqlSchemaManager` (used to be `SchemaManager`)	`PostgresSchemaManager`	Handles methods to gather table schemas and other metadata about the database

Major changes:

Adds classes specific to Postgres stream (logical replication) support
Rename classes specific to MySQL (binlog) stream support to be clearer
Refactor some common classes to accommodate both MySQL and Postgres
Adds/Updates unit tests

This is an initial PR, will follow up with more changes to add e2e ack, checkpointing, resync, data type mapping etc. for Postgres.

Testing

Tested Data Prepper locally against Aurora Postgres DB.

Issues Resolved

Contributes to #5309

Check List

New functionality includes testing.
New functionality has a documentation issue. Please link to it in this PR.
- New functionality has javadoc added
Commits are signed with a real name per the DCO

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

Signed-off-by: Hai Yan <[email protected]>

dinujoh · 2025-01-13T17:08:32Z

...ugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/RdsService.java

@@ -121,21 +133,23 @@ public void start(Buffer<Record<Event>> buffer) {
        }

        if (sourceConfig.isStreamEnabled()) {
-            BinlogClientFactory binaryLogClientFactory = new BinlogClientFactory(sourceConfig, rdsClient, dbMetadata);
+            ReplicationLogClientFactory replicationLogClientFactory = new ReplicationLogClientFactory(sourceConfig, rdsClient, dbMetadata);


This class needs some refactor. We should follow single responsibility principle here. This class has dependency on MySQL and Postgres Client. We should move these dependencies to separate class.

I did some refactoring and extracted out the creation of schemaManagers in separate classes. The ReplicationLogClientFactory already hides the creation of MySQL/Postgres clients.

dinujoh · 2025-01-13T17:10:10Z

...va/org/opensearch/dataprepper/plugins/source/rds/coordination/state/StreamProgressState.java

@@ -16,6 +16,12 @@ public class StreamProgressState {
    @JsonProperty("currentPosition")
    private BinlogCoordinate currentPosition;

+    @JsonProperty("currentLsn")


Consider creating new class to separate MySQL vs Postgres properties.

This model is used in serialization/deserialization by the source coordinator, it's a bit more involved to separate into two classes. I added a TODO in code and will address it in followup PRs.

dinujoh · 2025-01-13T17:13:06Z

...c/main/java/org/opensearch/dataprepper/plugins/source/rds/schema/MySqlConnectionManager.java

@@ -10,7 +10,7 @@
 import java.sql.SQLException;
 import java.util.Properties;

-public class ConnectionManager {
+public class MySqlConnectionManager {


Consider creating ConnectionManager interface and implementation for MySql/Postgres ConnectionManager

Created the interface.

dinujoh · 2025-01-13T17:20:15Z

...a/org/opensearch/dataprepper/plugins/source/rds/stream/LogicalReplicationEventProcessor.java

+        // If it's INSERT/UPDATE/DELETE, prepare events
+        // If it's a COMMIT, convert all prepared events and send to buffer
+        char messageType = (char) msg.get();
+        if (messageType == 'B') {


Consider using enums instead of char for message types for type safety, readability and maintainability

dinujoh · 2025-01-13T17:23:01Z

...a/org/opensearch/dataprepper/plugins/source/rds/stream/LogicalReplicationEventProcessor.java

+        pipelineEvents = new ArrayList<>();
+    }
+
+    public void process(ByteBuffer msg) {


Implement the missing functionality for UPDATE/DELETE message types 'K' and 'O' ?

https://www.postgresql.org/docs/current/protocol-logicalrep-message-formats.html

Can ByteBuffer msg be null ? Is there any validation to be done on msg ByteBuffer ?

Implement the missing functionality for UPDATE/DELETE message types 'K' and 'O' ?

Done. I updated the UPDATE message processor to include 'K' (primary keys were updated) and 'O' (old row data is included). For DELETE, it shouldn't matter as long as the row data includes primary keys.

Can ByteBuffer msg be null ? Is there any validation to be done on msg ByteBuffer ?

There's a null check for msg in the main while loop in LogicalReplicationClient.connect().

dinujoh · 2025-01-13T17:26:45Z

...a/org/opensearch/dataprepper/plugins/source/rds/stream/LogicalReplicationEventProcessor.java

+        long epochMicro = msg.getLong();
+
+        if (currentLsn != commitLsn) {
+            // This shouldn't happen


Should we consider throwing an exception instead of just logging ?

Updated to throw an exception.

Signed-off-by: Hai Yan <[email protected]>

chenqi0805 · 2025-01-17T21:12:03Z

...rc/main/java/org/opensearch/dataprepper/plugins/source/rds/datatype/postgres/ColumnType.java

+import java.util.Map;
+
+public enum ColumnType {
+    BOOLEAN(16, "boolean"),


Are those typeId from postgres? Is there way we can use their SDK variables instead of the hardcoded number

Yes, they are from postgres, but I don't think they have those in their libraries. Postgres has a system table pg_type that contains those type information. So an alternative would be to query the pg_type table the type ids and other type information.

chenqi0805 · 2025-01-17T21:46:33Z

...a/org/opensearch/dataprepper/plugins/source/rds/stream/LogicalReplicationEventProcessor.java

+            throw new RuntimeException("Commit LSN does not match current LSN, skipping");
+        }
+
+        writeToBuffer(bufferAccumulator);


QUES: will we flush uncommitted messages if numberOfRecordsToAccumulate or bufferTimeout is reached? If yes, Is it something we want to avoid?

This writeToBuffer() is called when processing the COMMIT message. So we will only write records to buffer when the changes are committed.

Thanks for clarifying. One follow up question, could there be uncommitted operations stay in the bufferAccumulator forever since user has never commit them (i.e. aborted operations)?

The source only writes to bufferAccumulator when commit message is received. Before that, the records stay in a temporary list (pipelineEvents).

dinujoh · 2025-01-21T18:12:52Z

...rc/main/java/org/opensearch/dataprepper/plugins/source/rds/datatype/postgres/ColumnType.java

+    }
+
+    public static ColumnType getByTypeId(int typeId) {
+        return TYPE_ID_MAP.get(typeId);


Add exception handling for typeId

dinujoh · 2025-01-21T18:21:13Z

...rc/main/java/org/opensearch/dataprepper/plugins/source/rds/schema/PostgresSchemaManager.java

+            // Create replication slot
+            PGReplicationConnection replicationConnection = pgConnection.getReplicationAPI();
+            try {
+                replicationConnection.createReplicationSlot()


Check if the replication slot exist and reuse if one exist.

Are we reusing if the replication slot exist ?

dinujoh · 2025-01-21T18:31:27Z

...rc/main/java/org/opensearch/dataprepper/plugins/source/rds/schema/PostgresSchemaManager.java

+            retry++;
+        }
+        LOG.warn("Failed to get primary keys for table {}", table);
+        return List.of();


Exception handling when PrimaryKey is not present vs Internal DB Exception/Errors.

dinujoh · 2025-01-21T18:33:34Z

...main/java/org/opensearch/dataprepper/plugins/source/rds/stream/LogicalReplicationClient.java

+                    .withSlotOption("proto_version", "1")
+                    .withSlotOption("publication_names", "my_publication");


Move the constant to static variable

Use the created publication names.

dinujoh · 2025-01-21T18:38:41Z

...main/java/org/opensearch/dataprepper/plugins/source/rds/stream/LogicalReplicationClient.java

+                        stream.setFlushedLSN(lsn);
+                        stream.setAppliedLSN(lsn);
+                    } catch (Exception e) {
+                        LOG.error("Exception while processing Postgres replication stream. ", e);


Add metrics in this class for exception, error and success cases.

Will follow up on this in the next PR.

dinujoh · 2025-01-21T18:44:06Z

...a/org/opensearch/dataprepper/plugins/source/rds/stream/LogicalReplicationEventProcessor.java

+        final long eventTimestampMillis = currentEventTimestamp;
+
+        char typeId = (char) msg.get();
+        if (typeId == 'N') {


move the typeId to enum

Signed-off-by: Hai Yan <[email protected]>

dinujoh · 2025-01-27T18:14:06Z

...urce/src/main/java/org/opensearch/dataprepper/plugins/source/rds/leader/LeaderScheduler.java

+        return "data_prepper_publication_" + UUID.randomUUID().toString().substring(0, 8);
+    }
+
+    private String generateReplicationSlotName() {
+        return "data_prepper_slot_" + UUID.randomUUID().toString().substring(0, 8);


Should this include pipeline identifier ? Can client view these in server ?

Should this include pipeline identifier ?

Sure. I will try to address this in next PR.

Can client view these in server ?

Which client? If you mean users, yes, they can view the replication slot names if they query for it against the server.

dinujoh · 2025-01-27T18:15:33Z

...main/java/org/opensearch/dataprepper/plugins/source/rds/schema/ConnectionManagerFactory.java

+    }
+
+    public ConnectionManager getConnectionManager() {
+        if (sourceConfig.getEngine() == EngineType.MYSQL) {


sourceConfig.getEngine() is this mandatory field ? We can throw exception is this is not set.

Yes, it's MYSQL by default.

dinujoh · 2025-01-27T18:22:50Z

...ugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/RdsService.java

    }

    private QueryManager getQueryManager(final RdsSourceConfig sourceConfig, final DbMetadata dbMetadata) {
        final String readerEndpoint = dbMetadata.getReaderEndpoint() != null ? dbMetadata.getReaderEndpoint() : dbMetadata.getEndpoint();
        final int readerPort = dbMetadata.getReaderPort() == 0 ? dbMetadata.getPort() : dbMetadata.getReaderPort();
-        final ConnectionManager readerConnectionManager = new ConnectionManager(
+        final MySqlConnectionManager readerConnectionManager = new MySqlConnectionManager(


Do we need connection manager for Postgres ?

This QueryManager is currently specific for MySQL due to the need for Resync'ing changes associated with foreign key cascades. Keeping those as is for now. If we ever need a QueryManager for Postgres, I will refactor these.

dinujoh · 2025-01-27T18:23:17Z

...ugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/RdsService.java

@@ -203,13 +213,11 @@ private String getS3PathPrefix() {
        return s3PathPrefix;
    }

-    private Map<String, Map<String, String>> getColumnDataTypeMap(final SchemaManager schemaManager) {
+    private Map<String, Map<String, String>> getColumnDataTypeMap(final MySqlSchemaManager schemaManager) {


Can this still be interface SchemaManager ?

Yes, it should be possible. Let me make some change after Postgres data type mapping work. It's currently only needed for MySQL.

Signed-off-by: Hai Yan <[email protected]>

oeyh · 2025-01-27T22:38:30Z

Had to remove mocking of ByteBuffer to fix Java 21 build failure due to org.mockito.exceptions.base.MockitoException: Unsupported settings with this type 'java.nio.ByteBuffer' (https://github.com/opensearch-project/data-prepper/actions/runs/12953484201/job/36133256794).

Build is passing now (except a known flaky OTEL tests with Java 11)

oeyh added 5 commits December 28, 2024 23:50

First working version

3767b1f

Signed-off-by: Hai Yan <[email protected]>

More progress and update existing unit tests

5f72e59

Signed-off-by: Hai Yan <[email protected]>

Add unit tests

9381480

Signed-off-by: Hai Yan <[email protected]>

Remove and rename classes

4fbdbdb

Signed-off-by: Hai Yan <[email protected]>

Remove test code

9e70372

Signed-off-by: Hai Yan <[email protected]>

oeyh requested review from sb2k16, chenqi0805, engechas, graytaylor0, dinujoh, kkondaka, KarstenSchnitter and dlvenable as code owners January 6, 2025 15:08

dinujoh reviewed Jan 13, 2025

View reviewed changes

oeyh requested a review from san81 as a code owner January 16, 2025 06:52

Address review comments

81c5883

Signed-off-by: Hai Yan <[email protected]>

oeyh force-pushed the rds-postgres-stream branch from c61da05 to 81c5883 Compare January 16, 2025 18:05

Address minor issues

a3bec6f

Signed-off-by: Hai Yan <[email protected]>

chenqi0805 reviewed Jan 17, 2025

View reviewed changes

chenqi0805 previously approved these changes Jan 17, 2025

View reviewed changes

chenqi0805 reviewed Jan 17, 2025

View reviewed changes

dinujoh reviewed Jan 21, 2025

View reviewed changes

Group MySQL and Postgres stream states

00abf68

Signed-off-by: Hai Yan <[email protected]>

oeyh dismissed chenqi0805’s stale review via f4e1683 January 24, 2025 06:25

oeyh requested a review from srikanthjg as a code owner January 24, 2025 06:25

Address more comments

561d938

Signed-off-by: Hai Yan <[email protected]>

oeyh force-pushed the rds-postgres-stream branch from f4e1683 to 561d938 Compare January 24, 2025 16:23

dinujoh reviewed Jan 27, 2025

View reviewed changes

dinujoh previously approved these changes Jan 27, 2025

View reviewed changes

chenqi0805 previously approved these changes Jan 27, 2025

View reviewed changes

Fix Java21 build

5108fa7

Signed-off-by: Hai Yan <[email protected]>

oeyh dismissed stale reviews from chenqi0805 and dinujoh via 5108fa7 January 27, 2025 21:29

chenqi0805 approved these changes Jan 27, 2025

View reviewed changes

dinujoh approved these changes Jan 27, 2025

View reviewed changes

oeyh merged commit f217b24 into opensearch-project:main Jan 27, 2025
45 of 47 checks passed

oeyh mentioned this pull request Jan 28, 2025

Fix Postgres stream processing #5367

Merged

4 tasks

		.withSlotOption("proto_version", "1")
		.withSlotOption("publication_names", "my_publication");

Initial PR for stream support for Postgres in Rds source #5310

Initial PR for stream support for Postgres in Rds source #5310

Conversation

oeyh commented Jan 6, 2025

Description

Testing

Issues Resolved

Check List

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chenqi0805 Jan 17, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chenqi0805 Jan 27, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

oeyh commented Jan 27, 2025 • edited Loading

chenqi0805 Jan 17, 2025 •

edited

Loading

chenqi0805 Jan 27, 2025 •

edited

Loading

oeyh commented Jan 27, 2025 •

edited

Loading