Address review comments

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

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

@@ -26,12 +26,13 @@
 import org.opensearch.dataprepper.plugins.source.rds.model.DbMetadata;
 import org.opensearch.dataprepper.plugins.source.rds.model.DbTableMetadata;
 import org.opensearch.dataprepper.plugins.source.rds.resync.ResyncScheduler;
+import org.opensearch.dataprepper.plugins.source.rds.schema.ConnectionManager;
+import org.opensearch.dataprepper.plugins.source.rds.schema.ConnectionManagerFactory;
 import org.opensearch.dataprepper.plugins.source.rds.schema.MySqlConnectionManager;
-import org.opensearch.dataprepper.plugins.source.rds.schema.QueryManager;
 import org.opensearch.dataprepper.plugins.source.rds.schema.MySqlSchemaManager;
+import org.opensearch.dataprepper.plugins.source.rds.schema.QueryManager;
 import org.opensearch.dataprepper.plugins.source.rds.schema.SchemaManager;
-import org.opensearch.dataprepper.plugins.source.rds.schema.PostgresConnectionManager;
-import org.opensearch.dataprepper.plugins.source.rds.schema.PostgresSchemaManager;
+import org.opensearch.dataprepper.plugins.source.rds.schema.SchemaManagerFactory;
 import org.opensearch.dataprepper.plugins.source.rds.stream.ReplicationLogClientFactory;
 import org.opensearch.dataprepper.plugins.source.rds.stream.StreamScheduler;
 import org.opensearch.dataprepper.plugins.source.rds.utils.IdentifierShortener;
@@ -178,29 +179,8 @@ public void shutdown() {
     }
 
     private SchemaManager getSchemaManager(final RdsSourceConfig sourceConfig, final DbMetadata dbMetadata) {
-        // For MySQL
-        if (sourceConfig.getEngine() == EngineType.MYSQL) {
-            final MySqlConnectionManager connectionManager = new MySqlConnectionManager(
-                    dbMetadata.getEndpoint(),
-                    dbMetadata.getPort(),
-                    sourceConfig.getAuthenticationConfig().getUsername(),
-                    sourceConfig.getAuthenticationConfig().getPassword(),
-                    sourceConfig.isTlsEnabled());
-            return new MySqlSchemaManager(connectionManager);
-        }
-        // For Postgres
-        final PostgresConnectionManager connectionManager = new PostgresConnectionManager(
-                dbMetadata.getEndpoint(),
-                dbMetadata.getPort(),
-                sourceConfig.getAuthenticationConfig().getUsername(),
-                sourceConfig.getAuthenticationConfig().getPassword(),
-                sourceConfig.isTlsEnabled(),
-                getDatabaseName(sourceConfig.getTableNames()));
-        return new PostgresSchemaManager(connectionManager);
-    }
-
-    private String getDatabaseName(List<String> tableNames) {
-        return tableNames.get(0).split("\\.")[0];
+        final ConnectionManager connectionManager = new ConnectionManagerFactory(sourceConfig, dbMetadata).getConnectionManager();
+        return new SchemaManagerFactory(connectionManager).getSchemaManager();
     }
 
     private QueryManager getQueryManager(final RdsSourceConfig sourceConfig, final DbMetadata dbMetadata) {

data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/coordination/state/StreamProgressState.java

@@ -10,23 +10,45 @@
 import org.opensearch.dataprepper.plugins.source.rds.model.ForeignKeyRelation;
 
 import java.util.List;
+import java.util.Map;
 
 public class StreamProgressState {
 
+    // TODO: separate MySQL and Postgres properties into different progress state classes
+    // Common
+    @JsonProperty("engineType")
+    private String engineType;
+
+    @JsonProperty("waitForExport")
+    private boolean waitForExport = false;
+
+    /**
+     * Map of table name to primary keys
+     */
+    @JsonProperty("primaryKeyMap")
+    private Map<String, List<String>> primaryKeyMap;
+
+    // For MySQL
     @JsonProperty("currentPosition")
     private BinlogCoordinate currentPosition;
 
+    @JsonProperty("foreignKeyRelations")
+    private List<ForeignKeyRelation> foreignKeyRelations;
+
+    // For Postgres
     @JsonProperty("currentLsn")
     private String currentLsn;
 
     @JsonProperty("replicationSlotName")
     private String replicationSlotName;
 
-    @JsonProperty("waitForExport")
-    private boolean waitForExport = false;
+    public String getEngineType() {
+        return engineType;
+    }
 
-    @JsonProperty("foreignKeyRelations")
-    private List<ForeignKeyRelation> foreignKeyRelations;
+    public void setEngineType(String engineType) {
+        this.engineType = engineType;
+    }
 
     public BinlogCoordinate getCurrentPosition() {
         return currentPosition;
@@ -36,6 +58,14 @@ public String getCurrentLsn() {
         return currentLsn;
     }
 
+    public Map<String, List<String>> getPrimaryKeyMap() {
+        return primaryKeyMap;
+    }
+
+    public void setPrimaryKeyMap(Map<String, List<String>> primaryKeyMap) {
+        this.primaryKeyMap = primaryKeyMap;
+    }
+
     public String getReplicationSlotName() {
         return replicationSlotName;
     }

data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/datatype/postgres/ColumnType.java

@@ -1,3 +1,13 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ */
+
 package org.opensearch.dataprepper.plugins.source.rds.datatype.postgres;
 
 import java.util.HashMap;

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

@@ -19,8 +19,8 @@
 import org.opensearch.dataprepper.plugins.source.rds.model.BinlogCoordinate;
 import org.opensearch.dataprepper.plugins.source.rds.model.DbTableMetadata;
 import org.opensearch.dataprepper.plugins.source.rds.schema.MySqlSchemaManager;
-import org.opensearch.dataprepper.plugins.source.rds.schema.SchemaManager;
 import org.opensearch.dataprepper.plugins.source.rds.schema.PostgresSchemaManager;
+import org.opensearch.dataprepper.plugins.source.rds.schema.SchemaManager;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -156,13 +156,15 @@ private Map<String, List<String>> getPrimaryKeyMap() {
         return sourceConfig.getTableNames().stream()
                 .collect(Collectors.toMap(
                         fullTableName -> fullTableName,
-                        fullTableName -> ((MySqlSchemaManager)schemaManager).getPrimaryKeys(fullTableName.split("\\.")[0], fullTableName.split("\\.")[1])
+                        fullTableName -> schemaManager.getPrimaryKeys(fullTableName)
                 ));
     }
 
     private void createStreamPartition(RdsSourceConfig sourceConfig) {
         final StreamProgressState progressState = new StreamProgressState();
+        progressState.setEngineType(sourceConfig.getEngine().toString());
         progressState.setWaitForExport(sourceConfig.isExportEnabled());
+        progressState.setPrimaryKeyMap(getPrimaryKeyMap());
         if (sourceConfig.getEngine() == EngineType.MYSQL) {
             getCurrentBinlogPosition().ifPresent(progressState::setCurrentPosition);
             progressState.setForeignKeyRelations(((MySqlSchemaManager)schemaManager).getForeignKeyRelations(sourceConfig.getTableNames()));

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

@@ -0,0 +1,30 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ */
+
+package org.opensearch.dataprepper.plugins.source.rds.model;
+
+public enum MessageType {
+    BEGIN('B'),
+    RELATION('R'),
+    INSERT('I'),
+    UPDATE('U'),
+    DELETE('D'),
+    COMMIT('C');
+
+    private final char value;
+
+    MessageType(char value) {
+        this.value = value;
+    }
+
+    public char getValue() {
+        return value;
+    }
+}

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

@@ -0,0 +1,22 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ */
+
+package org.opensearch.dataprepper.plugins.source.rds.schema;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+
+/**
+ * Interface for managing connections to a database.
+ */
+public interface ConnectionManager {
+
+    Connection getConnection() throws SQLException;
+}

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

@@ -0,0 +1,50 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ */
+
+package org.opensearch.dataprepper.plugins.source.rds.schema;
+
+import org.opensearch.dataprepper.plugins.source.rds.RdsSourceConfig;
+import org.opensearch.dataprepper.plugins.source.rds.configuration.EngineType;
+import org.opensearch.dataprepper.plugins.source.rds.model.DbMetadata;
+
+import java.util.List;
+
+public class ConnectionManagerFactory {
+    private final RdsSourceConfig sourceConfig;
+    private final DbMetadata dbMetadata;
+
+    public ConnectionManagerFactory(final RdsSourceConfig sourceConfig, final DbMetadata dbMetadata) {
+        this.sourceConfig = sourceConfig;
+        this.dbMetadata = dbMetadata;
+    }
+
+    public ConnectionManager getConnectionManager() {
+        if (sourceConfig.getEngine() == EngineType.MYSQL) {
+            return new MySqlConnectionManager(
+                    dbMetadata.getEndpoint(),
+                    dbMetadata.getPort(),
+                    sourceConfig.getAuthenticationConfig().getUsername(),
+                    sourceConfig.getAuthenticationConfig().getPassword(),
+                    sourceConfig.isTlsEnabled());
+        }
+
+        return new PostgresConnectionManager(
+                dbMetadata.getEndpoint(),
+                dbMetadata.getPort(),
+                sourceConfig.getAuthenticationConfig().getUsername(),
+                sourceConfig.getAuthenticationConfig().getPassword(),
+                sourceConfig.isTlsEnabled(),
+                getDatabaseName(sourceConfig.getTableNames()));
+    }
+
+    private String getDatabaseName(List<String> tableNames) {
+        return tableNames.get(0).split("\\.")[0];
+    }
+}

data-prepper-plugins/rds-source/src/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 MySqlConnectionManager {
+public class MySqlConnectionManager implements ConnectionManager {
     static final String JDBC_URL_FORMAT = "jdbc:mysql://%s:%d";
     static final String USERNAME_KEY = "user";
     static final String PASSWORD_KEY = "password";
@@ -33,6 +33,7 @@ public MySqlConnectionManager(String hostName, int port, String username, String
         this.requireSSL = requireSSL;
     }
 
+    @Override
     public Connection getConnection() throws SQLException {
         final Properties props = new Properties();
         props.setProperty(USERNAME_KEY, username);

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

@@ -1,6 +1,11 @@
 /*
  * Copyright OpenSearch Contributors
  * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
  */
 
 package org.opensearch.dataprepper.plugins.source.rds.schema;
@@ -39,13 +44,16 @@ public class MySqlSchemaManager implements SchemaManager {
     static final String UPDATE_RULE = "UPDATE_RULE";
     static final String DELETE_RULE = "DELETE_RULE";
     static final String COLUMN_DEF = "COLUMN_DEF";
-    private final MySqlConnectionManager connectionManager;
+    private final ConnectionManager connectionManager;
 
-    public MySqlSchemaManager(MySqlConnectionManager connectionManager) {
+    public MySqlSchemaManager(ConnectionManager connectionManager) {
         this.connectionManager = connectionManager;
     }
 
-    public List<String> getPrimaryKeys(final String database, final String table) {
+    @Override
+    public List<String> getPrimaryKeys(final String fullTableName) {
+        final String database = fullTableName.split("\\.")[0];
+        final String table = fullTableName.split("\\.")[1];
         int retry = 0;
         while (retry <= NUM_OF_RETRIES) {
             final List<String> primaryKeys = new ArrayList<>();

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

@@ -1,3 +1,13 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ */
+
 package org.opensearch.dataprepper.plugins.source.rds.schema;
 
 import org.postgresql.PGProperty;
@@ -9,7 +19,7 @@
 import java.sql.SQLException;
 import java.util.Properties;
 
-public class PostgresConnectionManager {
+public class PostgresConnectionManager implements ConnectionManager {
     private static final Logger LOG = LoggerFactory.getLogger(PostgresConnectionManager.class);
 
     public static final String JDBC_URL_FORMAT = "jdbc:postgresql://%s:%d/%s";
@@ -36,6 +46,7 @@ public PostgresConnectionManager(String endpoint, int port, String username, Str
         this.database = database;
     }
 
+    @Override
     public Connection getConnection() throws SQLException {
         final Properties props = new Properties();
         PGProperty.USER.set(props, username);