Oracle refactor work (#424)

Aryan-Verma · Aryan-Verma · web-flow · commit 36961f1be859 · 2023-09-12T09:37:51.000+05:30
Co-authored-by: Aryan-Verma &lt;aryanverma@google.com&gt;
diff --git a/database-commons/src/main/java/io/cdap/plugin/db/config/AbstractDBSpecificSourceConfig.java b/database-commons/src/main/java/io/cdap/plugin/db/config/AbstractDBSpecificSourceConfig.java
@@ -61,29 +61,29 @@ public abstract class AbstractDBSpecificSourceConfig extends PluginConfig implem
     "For example, 'SELECT * FROM table WHERE $CONDITIONS'. The '$CONDITIONS' string" +
     "will be replaced by 'splitBy' field limits specified by the bounding query.")
   @Macro
-  private String importQuery;
+  protected String importQuery;
 
   @Nullable
   @Name(BOUNDING_QUERY)
   @Description("Bounding Query should return the min and max of the " +
     "values of the 'splitBy' field. For example, 'SELECT MIN(id),MAX(id) FROM table'. " +
     "This is required unless numSplits is set to one.")
   @Macro
-  private String boundingQuery;
+  protected String boundingQuery;
 
   @Nullable
   @Name(SPLIT_BY)
   @Description("Field Name which will be used to generate splits. This is required unless numSplits is set to one.")
   @Macro
-  private String splitBy;
+  protected String splitBy;
 
   @Nullable
   @Name(NUM_SPLITS)
   @Description("The number of splits to generate. If set to one, the boundingQuery is not needed, " +
     "and no $CONDITIONS string needs to be specified in the importQuery. If not specified, the " +
     "execution framework will pick a value.")
   @Macro
-  private Integer numSplits;
+  protected Integer numSplits;
 
   @Nullable
   @Name(SCHEMA)
@@ -97,7 +97,7 @@ public abstract class AbstractDBSpecificSourceConfig extends PluginConfig implem
   @Macro
   @Description("The number of rows to fetch at a time per split. Larger fetch size can result in faster import, " +
     "with the tradeoff of higher memory usage.")
-  private Integer fetchSize;
+  protected Integer fetchSize;
 
   public String getImportQuery() {
     return cleanQuery(importQuery);
diff --git a/database-commons/src/main/java/io/cdap/plugin/db/source/AbstractDBSource.java b/database-commons/src/main/java/io/cdap/plugin/db/source/AbstractDBSource.java
@@ -133,10 +133,8 @@ public Schema getSchema(Class<? extends Driver> driverClass) throws IllegalAcces
     try {
 
       driverCleanup = loadPluginClassAndGetDriver(driverClass);
-      try (Connection connection = getConnection()) {
-        executeInitQueries(connection, sourceConfig.getInitQueries());
-        String query = sourceConfig.getImportQuery();
-        return loadSchemaFromDB(connection, query);
+      try {
+        return getSchema();
       } finally {
         driverCleanup.destroy();
       }
@@ -146,6 +144,23 @@ public Schema getSchema(Class<? extends Driver> driverClass) throws IllegalAcces
     }
   }
 
+  /**
+   * Returns the schema of the importQuery from the database using the SourceConfig details.
+   *
+   * @return Schema instance
+   * @throws SQLException In case of
+   *                      1. Error while creating connection to the database.
+   *                      2. Error while running any init queries.
+   *                      3. Error while running the import query.
+   */
+  public Schema getSchema() throws SQLException {
+    try (Connection connection = getConnection()) {
+      executeInitQueries(connection, sourceConfig.getInitQueries());
+      String query = sourceConfig.getImportQuery();
+      return loadSchemaFromDB(connection, query);
+    }
+  }
+
   private Schema loadSchemaFromDB(Connection connection, String query) throws SQLException {
     Statement statement = connection.createStatement();
     statement.setMaxRows(1);
@@ -229,22 +244,56 @@ public void prepareRun(BatchSourceContext context) throws Exception {
     sourceConfig.validate(collector);
     collector.getOrThrowException();
 
-    String connectionString = sourceConfig.getConnectionString();
-
     LOG.debug("pluginType = {}; pluginName = {}; connectionString = {}; importQuery = {}; " +
                 "boundingQuery = {};",
-              ConnectionConfig.JDBC_PLUGIN_TYPE, sourceConfig.getJdbcPluginName(),
-              connectionString,
-              sourceConfig.getImportQuery(), sourceConfig.getBoundingQuery());
-    ConnectionConfigAccessor connectionConfigAccessor = new ConnectionConfigAccessor();
+              ConnectionConfig.JDBC_PLUGIN_TYPE,
+              sourceConfig.getJdbcPluginName(),
+              sourceConfig.getConnectionString(),
+              sourceConfig.getImportQuery(),
+              sourceConfig.getBoundingQuery());
 
     // Load the plugin class to make sure it is available.
     Class<? extends Driver> driverClass = context.loadPluginClass(getJDBCPluginId());
+    Schema schemaFromDB = loadSchemaFromDB(driverClass);
+
+    ConnectionConfigAccessor connectionConfigAccessor = getConnectionConfigAccessor(
+                                                            driverClass.getName(),
+                                                            schemaFromDB,
+                                                            collector);
+
+    LineageRecorder lineageRecorder = getLineageRecorder(context);
+    Schema schema = sourceConfig.getSchema() == null ? schemaFromDB : sourceConfig.getSchema();
+    lineageRecorder.createExternalDataset(schema);
+    if (schema != null && schema.getFields() != null) {
+      lineageRecorder.recordRead("Read", "Read from database plugin",
+                                 schema.getFields().stream().map(Schema.Field::getName).collect(Collectors.toList()));
+    }
+    context.setInput(Input.of(sourceConfig.getReferenceName(), new SourceInputFormatProvider(
+      DataDrivenETLDBInputFormat.class, connectionConfigAccessor.getConfiguration())));
+  }
+
+  /**
+   * Returns the ConnectionConfigAccessor object containing the Configuration object for the SourceConfig
+   * and Schema. The configuration is later used by the InputFormat object for split calculation, reader creation.
+   *
+   * @param driverClassName   Class name of the driver in use
+   * @param schemaFromDB      Schema object
+   * @param collector         Failure Collector object
+   * @return                  ConnectionConfigAccessor instance
+   * @throws IOException
+   */
+  public ConnectionConfigAccessor getConnectionConfigAccessor (String driverClassName,
+                                                               Schema schemaFromDB,
+                                                               FailureCollector collector) throws IOException {
+    ConnectionConfigAccessor connectionConfigAccessor = new ConnectionConfigAccessor();
+
     if (sourceConfig.getUser() == null && sourceConfig.getPassword() == null) {
-      DBConfiguration.configureDB(connectionConfigAccessor.getConfiguration(), driverClass.getName(), connectionString);
+      DBConfiguration.configureDB(connectionConfigAccessor.getConfiguration(),
+          driverClassName, sourceConfig.getConnectionString());
     } else {
-      DBConfiguration.configureDB(connectionConfigAccessor.getConfiguration(), driverClass.getName(), connectionString,
-                                  sourceConfig.getUser(), sourceConfig.getPassword());
+      DBConfiguration.configureDB(connectionConfigAccessor.getConfiguration(),
+          driverClassName, sourceConfig.getConnectionString(),
+          sourceConfig.getUser(), sourceConfig.getPassword());
     }
 
     if (sourceConfig.getFetchSize() != null) {
@@ -255,7 +304,6 @@ public void prepareRun(BatchSourceContext context) throws Exception {
                                         sourceConfig.getImportQuery(), sourceConfig.getBoundingQuery(),
                                         false);
 
-
     if (sourceConfig.getTransactionIsolationLevel() != null) {
       connectionConfigAccessor.setTransactionIsolationLevel(sourceConfig.getTransactionIsolationLevel());
     }
@@ -273,7 +321,6 @@ public void prepareRun(BatchSourceContext context) throws Exception {
       connectionConfigAccessor.getConfiguration().setInt(MRJobConfig.NUM_MAPS, sourceConfig.getNumSplits());
     }
 
-    Schema schemaFromDB = loadSchemaFromDB(driverClass);
     if (sourceConfig.getSchema() != null) {
       sourceConfig.validateSchema(schemaFromDB, collector);
       collector.getOrThrowException();
@@ -283,15 +330,7 @@ public void prepareRun(BatchSourceContext context) throws Exception {
       connectionConfigAccessor.setSchema(schemaStr);
     }
 
-    LineageRecorder lineageRecorder = getLineageRecorder(context);
-    Schema schema = sourceConfig.getSchema() == null ? schemaFromDB : sourceConfig.getSchema();
-    lineageRecorder.createExternalDataset(schema);
-    if (schema != null && schema.getFields() != null) {
-      lineageRecorder.recordRead("Read", "Read from database plugin",
-                                 schema.getFields().stream().map(Schema.Field::getName).collect(Collectors.toList()));
-    }
-    context.setInput(Input.of(sourceConfig.getReferenceName(), new SourceInputFormatProvider(
-      DataDrivenETLDBInputFormat.class, connectionConfigAccessor.getConfiguration())));
+    return connectionConfigAccessor;
   }
 
   protected LineageRecorder getLineageRecorder(BatchSourceContext context) {
diff --git a/database-commons/src/main/java/io/cdap/plugin/db/source/DataDrivenETLDBInputFormat.java b/database-commons/src/main/java/io/cdap/plugin/db/source/DataDrivenETLDBInputFormat.java
@@ -129,7 +129,7 @@ public Connection createConnection() {
   }
 
   @Override
-  protected RecordReader createDBRecordReader(DBInputSplit split, Configuration conf) throws IOException {
+  public RecordReader createDBRecordReader(DBInputSplit split, Configuration conf) throws IOException {
     final RecordReader dbRecordReader = super.createDBRecordReader(split, conf);
     return new RecordReader() {
       @Override
diff --git a/oracle-plugin/src/main/java/io/cdap/plugin/oracle/OracleConnectorConfig.java b/oracle-plugin/src/main/java/io/cdap/plugin/oracle/OracleConnectorConfig.java
@@ -37,12 +37,18 @@ public class OracleConnectorConfig extends AbstractDBSpecificConnectorConfig {
   private static final String ROLE_NORMAL = "normal";
 
   public OracleConnectorConfig(String host, int port, String user, String password, String jdbcPluginName,
-                               String connectionArguments) {
-    this(host, port, user, password, jdbcPluginName, connectionArguments, null);
+                               String connectionArguments, String database) {
+    this(host, port, user, password, jdbcPluginName, connectionArguments, null, database);
   }
 
   public OracleConnectorConfig(String host, int port, String user, String password, String jdbcPluginName,
-                               String connectionArguments, String connectionType) {
+                               String connectionArguments, String connectionType, String database) {
+    this(host, port, user, password, jdbcPluginName, connectionArguments, connectionType, database, null);
+  }
+
+  public OracleConnectorConfig(String host, int port, String user, String password, String jdbcPluginName,
+                               String connectionArguments, String connectionType, String database,
+                               String role) {
 
     this.host = host;
     this.port = port;
@@ -51,6 +57,8 @@ public OracleConnectorConfig(String host, int port, String user, String password
     this.jdbcPluginName = jdbcPluginName;
     this.connectionArguments = connectionArguments;
     this.connectionType = connectionType;
+    this.database = database;
+    this.role = role;
   }
 
   @Override
diff --git a/oracle-plugin/src/main/java/io/cdap/plugin/oracle/OracleSource.java b/oracle-plugin/src/main/java/io/cdap/plugin/oracle/OracleSource.java
@@ -113,6 +113,22 @@ public static class OracleSourceConfig extends AbstractDBSpecificSourceConfig {
     @Nullable
     private Integer defaultRowPrefetch;
 
+    public OracleSourceConfig(String host, int port, String user, String password, String jdbcPluginName,
+                              String connectionArguments, String connectionType, String database, String role,
+                              int defaultBatchValue, int defaultRowPrefetch,
+                              String importQuery, Integer numSplits, int fetchSize,
+                              String boundingQuery, String splitBy) {
+      this.connection = new OracleConnectorConfig(host, port, user, password, jdbcPluginName, connectionArguments,
+                                                  connectionType, database, role);
+      this.defaultBatchValue = defaultBatchValue;
+      this.defaultRowPrefetch = defaultRowPrefetch;
+      this.fetchSize = fetchSize;
+      this.importQuery = importQuery;
+      this.numSplits = numSplits;
+      this.boundingQuery = boundingQuery;
+      this.splitBy = splitBy;
+    }
+
     @Override
     public String getConnectionString() {
       if (OracleConstants.TNS_CONNECTION_TYPE.equals(connection.getConnectionType())) {
@@ -140,7 +156,7 @@ protected Map<String, String> getDBSpecificArguments() {
     }
 
     @Override
-    protected OracleConnectorConfig getConnection() {
+    public OracleConnectorConfig getConnection() {
       return connection;
     }
 
diff --git a/oracle-plugin/src/test/java/io/cdap/plugin/oracle/OracleConnectorTest.java b/oracle-plugin/src/test/java/io/cdap/plugin/oracle/OracleConnectorTest.java
@@ -29,7 +29,7 @@ public class OracleConnectorTest extends DBSpecificConnectorBaseTest {
   @Test
   public void test() throws IOException, ClassNotFoundException, InstantiationException, IllegalAccessException {
     test(new OracleConnector(
-           new OracleConnectorConfig(host, port, username, password, JDBC_PLUGIN_NAME, connectionArguments)),
+           new OracleConnectorConfig(host, port, username, password, JDBC_PLUGIN_NAME, connectionArguments, database)),
          JDBC_DRIVER_CLASS_NAME, OracleConstants.PLUGIN_NAME);
 
     test(new OracleConnector(
diff --git a/oracle-plugin/src/test/java/io/cdap/plugin/oracle/OracleFailedConnectionTest.java b/oracle-plugin/src/test/java/io/cdap/plugin/oracle/OracleFailedConnectionTest.java
@@ -28,10 +28,10 @@ public class OracleFailedConnectionTest extends DBSpecificFailedConnectionTest {
   public void test() throws ClassNotFoundException, IOException {
 
     OracleConnector connector = new OracleConnector(
-      new OracleConnectorConfig("localhost", 1521, "username", "password", "jdbc", ""));
+      new OracleConnectorConfig("localhost", 1521, "username", "password", "jdbc", "", "database"));
 
     super.test(JDBC_DRIVER_CLASS_NAME, connector, "Failed to create connection to database via connection string:" +
-                                                    " jdbc:oracle:thin:@localhost:1521:null and arguments: " +
+                                                    " jdbc:oracle:thin:@localhost:1521:database and arguments: " +
                                                     "{user=username, oracle.jdbc.timezoneAsRegion=false, " +
                                                     "internal_logon=normal}. Error: ConnectException: Connection " +
                                                     "refused.");

Original file line number	Diff line number	Diff line change
`@@ -129,7 +129,7 @@ public Connection createConnection() {`
`129`	`129`	`}`
`130`	`130`
`131`	`131`	`@Override`
`132`		`- protected RecordReader createDBRecordReader(DBInputSplit split, Configuration conf) throws IOException {`
	`132`	`+ public RecordReader createDBRecordReader(DBInputSplit split, Configuration conf) throws IOException {`
`133`	`133`	`final RecordReader dbRecordReader = super.createDBRecordReader(split, conf);`
`134`	`134`	`return new RecordReader() {`
`135`	`135`	`@Override`