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.

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.

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.

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

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 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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

oeyh commented Jan 6, 2025

dinujoh Jan 13, 2025

oeyh Jan 16, 2025

dinujoh Jan 13, 2025

oeyh Jan 16, 2025

dinujoh Jan 13, 2025

oeyh Jan 16, 2025

dinujoh Jan 13, 2025

oeyh Jan 16, 2025

dinujoh Jan 13, 2025

dinujoh Jan 13, 2025

oeyh Jan 16, 2025

oeyh Jan 16, 2025

dinujoh Jan 13, 2025

oeyh Jan 16, 2025

chenqi0805 Jan 17, 2025

oeyh Jan 17, 2025

chenqi0805 Jan 17, 2025 •

edited

Loading

oeyh Jan 17, 2025

dinujoh Jan 21, 2025

dinujoh Jan 21, 2025

dinujoh Jan 21, 2025

dinujoh Jan 21, 2025

dinujoh Jan 21, 2025

dinujoh Jan 21, 2025

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

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

Are you sure you want to change the base?

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

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