PLUGIN-1936: Refactor code to implement streaming logic

Author: sgarg-CS

src/main/java/io/cdap/plugin/servicenow/apiclient/ServiceNowTableAPIClientImpl.java

@@ -138,7 +138,7 @@ public String getAccessTokenRetryableMode() throws ExecutionException, RetryExce
    * @param limit     The number of records to be fetched
    * @return The list of Map; each Map representing a table row
    */
-  public List<Map<String, String>> fetchTableRecords(
+  public RestAPIResponse fetchTableRecords(
       String tableName,
       SourceValueType valueType,
       String startDate,
@@ -161,8 +161,9 @@ public List<Map<String, String>> fetchTableRecords(
     String accessToken = getAccessToken();
     requestBuilder.setAuthHeader(accessToken);
     RestAPIResponse apiResponse = executeGetWithRetries(requestBuilder.build());
+    return apiResponse;
     //return parseResponseToResultListOfMap(apiResponse.getResponseBody());
-    return parseResponseStreamToResultListOfMap(apiResponse.getInputStream());
+    // return parseResponseStreamToRecord(apiResponse.getInputStream());
 
   }
 
@@ -206,20 +207,20 @@ public List<Map<String, String>> parseResponseToResultListOfMap(String responseB
     return GSON.fromJson(ja, type);
   }
 
-  public List<Map<String, String>> parseResponseStreamToResultListOfMap(InputStream in) throws ServiceNowAPIException {
-    List<Map<String, String>> records = new ArrayList<>();
+  public Map<String, String> parseResponseStreamToRecord(InputStream in) throws ServiceNowAPIException {
+    // List<Map<String, String>> records = new ArrayList<>();
     // InputStream in = httpResponse.getEntity().getContent();
     try (InputStreamReader reader = new InputStreamReader(in, StandardCharsets.UTF_8);
          JsonReader jsonReader = new JsonReader(reader)) {
       jsonReader.setLenient(true);
       jsonReader.beginObject();
+      Map<String, String> record = new HashMap<>();
       while (jsonReader.hasNext()) {
         String name = jsonReader.nextName();
         if (ServiceNowConstants.RESULT.equals(name) && jsonReader.peek() == JsonToken.BEGIN_ARRAY) {
           jsonReader.beginArray();
           while (jsonReader.hasNext()) {
             jsonReader.beginObject();
-            Map<String, String> record = new HashMap<>();
             while (jsonReader.hasNext()) {
               String field = jsonReader.nextName();
               JsonToken token = jsonReader.peek();
@@ -228,7 +229,7 @@ public List<Map<String, String>> parseResponseStreamToResultListOfMap(InputStrea
               record.put(field, token == JsonToken.NULL ? null : jsonReader.nextString());
             }
             jsonReader.endObject();
-            records.add(record);
+            // records.add(record);
           }
           jsonReader.endArray();
         } else {
@@ -237,7 +238,7 @@ public List<Map<String, String>> parseResponseStreamToResultListOfMap(InputStrea
         }
       }
       jsonReader.endObject();
-      return records;
+      return record;
     } catch (IOException e) {
       throw new ServiceNowAPIException(e, null);
     }
@@ -277,12 +278,14 @@ private String getErrorMessage(String responseBody) {
    * @param limit     The number of records to be fetched
    * @return The list of Map; each Map representing a table row
    */
-  public List<Map<String, String>> fetchTableRecordsRetryableMode(String tableName, SourceValueType valueType,
+  public RestAPIResponse fetchTableRecordsRetryableMode(String tableName, SourceValueType valueType,
                                                                   String startDate, String endDate, int offset,
                                                                   int limit) throws ServiceNowAPIException {
-    final List<Map<String, String>> results = new ArrayList<>();
+    //final List<Map<String, String>> results = new ArrayList<>();
+    final RestAPIResponse[] restAPIResponse = new RestAPIResponse[1];
     Callable<Boolean> fetchRecords = () -> {
-      results.addAll(fetchTableRecords(tableName, valueType, startDate, endDate, offset, limit));
+      // results.addAll(fetchTableRecords(tableName, valueType, startDate, endDate, offset, limit));
+      restAPIResponse[0] = fetchTableRecords(tableName, valueType, startDate, endDate, offset, limit);
       return true;
     };
 
@@ -300,7 +303,7 @@ public List<Map<String, String>> fetchTableRecordsRetryableMode(String tableName
           e, null, false);
     }
 
-    return results;
+    return restAPIResponse[0];
   }
 
   /**

src/main/java/io/cdap/plugin/servicenow/restapi/RestAPIResponse.java

@@ -64,9 +64,11 @@ public class RestAPIResponse {
   public RestAPIResponse(
       Map<String, String> headers,
       @Nullable String responseBody,
+      InputStream inputStream,
       @Nullable ServiceNowAPIException exception) {
     this.headers = headers;
     this.responseBody = responseBody;
+    this.inputStream = inputStream;
     this.exception = exception;
   }
 
@@ -101,26 +103,26 @@ public static RestAPIResponse parse(HttpResponse httpResponse, String... headerN
 
     ServiceNowAPIException serviceNowAPIException = validateHttpResponse(httpResponse);
     if (serviceNowAPIException != null) {
-      return new RestAPIResponse(headers, (InputStream) null, serviceNowAPIException);
+      return new RestAPIResponse(headers, null, null, serviceNowAPIException);
     }
 
     String responseBody = null;
     try {
       responseBody = EntityUtils.toString(httpResponse.getEntity());
     } catch (IOException e) {
-      return new RestAPIResponse(headers, (String) null, new ServiceNowAPIException(e, httpResponse));
+      return new RestAPIResponse(headers, null, null, new ServiceNowAPIException(e, httpResponse));
     }
     // Instead of reading the entire entity, store the stream
     HttpEntity httpEntity = httpResponse.getEntity();
     InputStream responseStream;
     try {
       responseStream = (httpEntity != null) ? httpEntity.getContent() : null;
     } catch (IOException e) {
-      return new RestAPIResponse(headers, (InputStream) null, new ServiceNowAPIException(e, httpResponse));
+      return new RestAPIResponse(headers, null, null, new ServiceNowAPIException(e, httpResponse));
     }
     serviceNowAPIException = validateRestApiResponse(httpResponse, responseBody);
     // return new RestAPIResponse(headers, responseBody, serviceNowAPIException);
-    return new RestAPIResponse(headers, responseStream, serviceNowAPIException);
+    return new RestAPIResponse(headers, responseBody, responseStream, serviceNowAPIException);
 
   }
 

src/main/java/io/cdap/plugin/servicenow/source/ServiceNowMultiRecordReader.java

@@ -22,6 +22,7 @@
 import io.cdap.plugin.servicenow.apiclient.ServiceNowAPIException;
 import io.cdap.plugin.servicenow.apiclient.ServiceNowTableAPIClientImpl;
 import io.cdap.plugin.servicenow.connector.ServiceNowRecordConverter;
+import io.cdap.plugin.servicenow.restapi.RestAPIResponse;
 import org.apache.hadoop.mapreduce.InputSplit;
 import org.apache.hadoop.mapreduce.TaskAttemptContext;
 
@@ -91,14 +92,14 @@ public StructuredRecord getCurrentValue() throws IOException {
   }
 
   @VisibleForTesting
-  void fetchData() throws ServiceNowAPIException {
+  RestAPIResponse fetchData() throws ServiceNowAPIException {
     // Get the table data
-    results = restApi.fetchTableRecordsRetryableMode(tableName, multiSourcePluginConf.getValueType(),
-                                                     multiSourcePluginConf.getStartDate(),
-                                                     multiSourcePluginConf.getEndDate(), split.getOffset(),
-                                                     multiSourcePluginConf.getPageSize());
+    RestAPIResponse restAPIResponse = restApi.fetchTableRecordsRetryableMode(tableName,
+      multiSourcePluginConf.getValueType(), multiSourcePluginConf.getStartDate(), multiSourcePluginConf.getEndDate(),
+        split.getOffset(), multiSourcePluginConf.getPageSize());
 
-    iterator = results.iterator();
+    // iterator = results.iterator();
+    return restAPIResponse;
   }
 
   private void fetchSchema(ServiceNowTableAPIClientImpl restApi) {