Merge pull request #63 from cloudsufi/bugfix/ref-name-fix

[PLUGIN-1752] Changed parent class to ReferenceBatchSource

Author: Pushpender Saini

src/main/java/io/cdap/plugin/MultiTableDBSource.java

@@ -32,6 +32,8 @@
 import io.cdap.cdap.etl.api.batch.BatchSourceContext;
 import io.cdap.plugin.common.Asset;
 import io.cdap.plugin.common.LineageRecorder;
+import io.cdap.plugin.common.ReferenceBatchSource;
+import io.cdap.plugin.common.ReferencePluginConfig;
 import io.cdap.plugin.common.SourceInputFormatProvider;
 import io.cdap.plugin.format.DBTableInfo;
 import io.cdap.plugin.format.MultiSQLStatementInputFormat;
@@ -63,19 +65,21 @@
 @Description("Reads from multiple tables in a relational database. " +
   "Outputs one record for each row in each table, with the table name as a record field. " +
   "Also sets a pipeline argument for each table read, which contains the table schema. ")
-public class MultiTableDBSource extends BatchSource<NullWritable, RecordWrapper, StructuredRecord> {
+public class MultiTableDBSource extends ReferenceBatchSource<NullWritable, RecordWrapper, StructuredRecord> {
   private static final Logger LOG = LoggerFactory.getLogger(MultiTableDBSource.class);
 
   private static final String JDBC_PLUGIN_ID = "jdbc.driver";
 
   private final MultiTableConf conf;
 
   public MultiTableDBSource(MultiTableConf conf) {
+    super(new ReferencePluginConfig(conf.getReferenceName()));
     this.conf = conf;
   }
 
   @Override
   public void configurePipeline(PipelineConfigurer pipelineConfigurer) {
+    super.configurePipeline(pipelineConfigurer);
     Class<? extends Driver> jdbcDriverClass = pipelineConfigurer.usePluginClass("jdbc", conf.getJdbcPluginName(),
                                                                                 JDBC_PLUGIN_ID,
                                                                                 PluginProperties.builder().build());

src/test/java/io/cdap/plugin/MultiTableDBSourceTest.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright © 2024 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package io.cdap.plugin;
+
+
+import io.cdap.cdap.api.data.schema.Schema;
+import io.cdap.cdap.etl.api.validation.ValidationException;
+import io.cdap.cdap.etl.mock.common.MockPipelineConfigurer;
+import io.cdap.plugin.format.MultiTableConf;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+/**
+ * Tests for MultiTableDBSource
+ */
+public class MultiTableDBSourceTest {
+
+  MockPipelineConfigurer mockPipelineConfigurer = new MockPipelineConfigurer(Schema.of(Schema.Type.STRING));
+  MultiTableDBSource multiTableDBSource;
+  String referenceNameWithSpace = "Hello MultiTable";
+
+  @Before
+  public void setUp() {
+    multiTableDBSource = new MultiTableDBSource(new MultiTableConf(referenceNameWithSpace));
+  }
+
+  /**
+   * Plugin should not allow space in reference name for MultiTableDBSource
+   */
+  @Test
+  public void testMultiTableReferenceNameWithSpace() {
+    try {
+      multiTableDBSource.configurePipeline(mockPipelineConfigurer);
+    } catch (ValidationException e) {
+      // expected
+    }
+    Assert.assertEquals(1, mockPipelineConfigurer.getStageConfigurer().getFailureCollector().getValidationFailures()
+      .size());
+  }
+}