PLUGIN-1936: Rework based on new changes suggested in the design

Author: sgarg-CS

pom.xml

@@ -99,6 +99,11 @@
       <version>${cdap.version}</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>commons-io</groupId>
+      <artifactId>commons-io</artifactId>
+      <version>2.5</version>
+    </dependency>
     <dependency>
       <groupId>org.apache.hadoop</groupId>
       <artifactId>hadoop-common</artifactId>

src/main/java/io/cdap/plugin/servicenow/ServiceNowBaseConfig.java

@@ -17,12 +17,14 @@
 package io.cdap.plugin.servicenow;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.gson.stream.JsonReader;
 import io.cdap.cdap.api.annotation.Description;
 import io.cdap.cdap.api.annotation.Macro;
 import io.cdap.cdap.api.annotation.Name;
 import io.cdap.cdap.api.plugin.PluginConfig;
 import io.cdap.cdap.etl.api.FailureCollector;
 import io.cdap.plugin.common.ConfigUtil;
+import io.cdap.plugin.servicenow.apiclient.ServiceNowAPIException;
 import io.cdap.plugin.servicenow.apiclient.ServiceNowTableAPIClientImpl;
 import io.cdap.plugin.servicenow.apiclient.ServiceNowTableAPIRequestBuilder;
 import io.cdap.plugin.servicenow.connector.ServiceNowConnectorConfig;
@@ -31,14 +33,21 @@
 import io.cdap.plugin.servicenow.util.SchemaType;
 import io.cdap.plugin.servicenow.util.ServiceNowConstants;
 import io.cdap.plugin.servicenow.util.SourceValueType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.nio.charset.StandardCharsets;
 import javax.annotation.Nullable;
 
 /**
  * ServiceNow Base Config. Contains connection properties and methods.
  */
 public class ServiceNowBaseConfig extends PluginConfig {
 
+  private static final Logger log = LoggerFactory.getLogger(ServiceNowBaseConfig.class);
   @Name(ConfigUtil.NAME_USE_CONNECTION)
   @Nullable
   @Description("Whether to use an existing connection.")
@@ -140,7 +149,7 @@ public void validateTable(String tableName, SourceValueType valueType, FailureCo
       requestBuilder.setResponseHeaders(ServiceNowConstants.HEADER_NAME_TOTAL_COUNT);
 
       apiResponse = serviceNowTableAPIClient.executeGetWithRetries(requestBuilder.build());
-      if (serviceNowTableAPIClient.parseResponseToResultListOfMap(apiResponse.getResponseBody()).isEmpty()) {
+      if (isResultEmpty(apiResponse)) {
         // Removed config property as in case of MultiSource, only first table error was populating.
         collector.addFailure("Table: " + tableName + " is empty.", "");
       }
@@ -152,4 +161,26 @@ public void validateTable(String tableName, SourceValueType valueType, FailureCo
     }
   }
 
+  /**
+   * Check whether the result is empty or not.
+   * @param restAPIResponse
+   * @return true if result is empty
+   * @throws IOException
+   */
+  public boolean isResultEmpty(RestAPIResponse restAPIResponse) throws IOException {
+    JsonReader reader = new JsonReader(new InputStreamReader(restAPIResponse.getBodyAsStream(),
+      StandardCharsets.UTF_8));
+    reader.beginObject();
+    while (reader.hasNext()) {
+      String name = reader.nextName();
+      if (ServiceNowConstants.RESULT.equals(name)) {
+        reader.beginArray();
+        return !reader.hasNext();
+      } else {
+        reader.skipValue();
+      }
+    }
+    return true;
+  }
+
 }

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

@@ -161,9 +161,6 @@ public RestAPIResponse fetchTableRecords(
     requestBuilder.setAuthHeader(accessToken);
     RestAPIResponse apiResponse = executeGetWithRetries(requestBuilder.build());
     return apiResponse;
-    //return parseResponseToResultListOfMap(apiResponse.getResponseBody());
-    // return parseResponseStreamToRecord(apiResponse.getInputStream());
-
   }
 
   private void applyDateRangeToRequest(ServiceNowTableAPIRequestBuilder requestBuilder, String startDate,
@@ -206,38 +203,35 @@ public List<Map<String, String>> parseResponseToResultListOfMap(String responseB
     return GSON.fromJson(ja, type);
   }
 
-  public Map<String, String> parseResponseStreamToRecord(InputStream in) throws ServiceNowAPIException {
-    // List<Map<String, String>> records = new ArrayList<>();
-    // InputStream in = httpResponse.getEntity().getContent();
+  public List<Map<String, String>> parseResponseToResultListOfMap(InputStream in) throws ServiceNowAPIException {
     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<>();
+
+      List<Map<String, String>> records = new ArrayList<>();
       while (jsonReader.hasNext()) {
         String name = jsonReader.nextName();
         if (ServiceNowConstants.RESULT.equals(name) && jsonReader.peek() == JsonToken.BEGIN_ARRAY) {
           jsonReader.beginArray();
           while (jsonReader.hasNext()) {
             jsonReader.beginObject();
             while (jsonReader.hasNext()) {
+              Map<String, String> record = new HashMap<>();
               String field = jsonReader.nextName();
               JsonToken token = jsonReader.peek();
-              // JsonElement resultElement = GSON.fromJson(jsonReader, JsonElement.class);
-              // responseBody = resultElement.toString();
               record.put(field, token == JsonToken.NULL ? null : jsonReader.nextString());
+              records.add(record);
             }
             jsonReader.endObject();
-            // records.add(record);
           }
           jsonReader.endArray();
         } else {
-          // skip other fields (e.g., metadata like result_count)
           jsonReader.skipValue();
         }
       }
       jsonReader.endObject();
-      return record;
+      return records;
     } catch (IOException e) {
       throw new ServiceNowAPIException(e, null);
     }
@@ -325,8 +319,8 @@ public Schema fetchTableSchema(String tableName, FailureCollector collector) {
   }
 
   @VisibleForTesting
-  public MetadataAPISchemaResponse parseSchemaResponse(String responseBody) {
-    return GSON.fromJson(responseBody, MetadataAPISchemaResponse.class);
+  public MetadataAPISchemaResponse parseSchemaResponse(InputStream responseStream) {
+    return GSON.fromJson(createJsonReader(responseStream), MetadataAPISchemaResponse.class);
   }
 
   /**
@@ -400,7 +394,7 @@ public Schema fetchTableSchema(String tableName, String accessToken, SourceValue
   private Schema prepareSchemaWithSchemaAPI(RestAPIResponse restAPIResponse, List<ServiceNowColumn> columns,
                                             String tableName) throws ServiceNowAPIException {
     SchemaAPISchemaResponse schemaAPISchemaResponse =
-      GSON.fromJson(restAPIResponse.getResponseBody(), SchemaAPISchemaResponse.class);
+      GSON.fromJson(createJsonReader(restAPIResponse.getBodyAsStream()), SchemaAPISchemaResponse.class);
 
     if (schemaAPISchemaResponse.getResult() == null || schemaAPISchemaResponse.getResult().isEmpty()) {
       throw new ServiceNowAPIException(
@@ -434,8 +428,7 @@ private Schema prepareSchemaWithSchemaAPI(RestAPIResponse restAPIResponse, List<
   private Schema prepareSchemaWithMetadataAPI(RestAPIResponse restAPIResponse, List<ServiceNowColumn> columns,
                                               String tableName, SourceValueType valueType) throws
     ServiceNowAPIException {
-    MetadataAPISchemaResponse metadataAPISchemaResponse = parseSchemaResponse(restAPIResponse.getResponseBody());
-
+    MetadataAPISchemaResponse metadataAPISchemaResponse = parseSchemaResponse(restAPIResponse.getBodyAsStream());
     if (metadataAPISchemaResponse.getResult() == null || metadataAPISchemaResponse.getResult().getColumns() == null ||
       metadataAPISchemaResponse.getResult().getColumns().isEmpty()) {
       throw new ServiceNowAPIException(
@@ -461,6 +454,11 @@ private Schema prepareSchemaWithMetadataAPI(RestAPIResponse restAPIResponse, Lis
     return SchemaBuilder.constructSchema(tableName, columns);
   }
 
+  public JsonReader createJsonReader(InputStream inputStream) {
+    Objects.requireNonNull(inputStream, "InputStream must not be null");
+    return new JsonReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
+  }
+
   /**
    * Get the total number of records in the table
    *
@@ -551,8 +549,8 @@ public String createRecordInDisplayMode(String tableName, HttpEntity entity) thr
   }
 
   private String getSystemId(RestAPIResponse restAPIResponse) {
-    CreateRecordAPIResponse apiResponse = GSON.fromJson(restAPIResponse.getResponseBody(),
-                                                           CreateRecordAPIResponse.class);
+    CreateRecordAPIResponse apiResponse = GSON.fromJson(
+      new InputStreamReader(restAPIResponse.getBodyAsStream(), StandardCharsets.UTF_8), CreateRecordAPIResponse.class);
     return apiResponse.getResult().get(ServiceNowConstants.SYSTEM_ID).toString();
   }
 
@@ -575,7 +573,8 @@ public Map<String, String> getRecordFromServiceNowTable(String tableName, String
     requestBuilder.setAuthHeader(accessToken);
     restAPIResponse = executeGetWithRetries(requestBuilder.build());
 
-    APIResponse apiResponse = GSON.fromJson(restAPIResponse.getResponseBody(), APIResponse.class);
+    APIResponse apiResponse = GSON.fromJson(
+      new InputStreamReader(restAPIResponse.getBodyAsStream(), StandardCharsets.UTF_8), APIResponse.class);
     return apiResponse.getResult().get(0);
   }
 
@@ -593,8 +592,9 @@ public Map<String, String> getRecordFromServiceNowTable(String tableName, String
    * @throws RuntimeException if the schema response is null or contains no result.
    */
   private Schema prepareStringBasedSchema(RestAPIResponse restAPIResponse, List<ServiceNowColumn> columns,
-                                          String tableName) {
-    List<Map<String, String>> result = parseResponseToResultListOfMap(restAPIResponse.getResponseBody());
+                                          String tableName) throws ServiceNowAPIException {
+    List<Map<String, String>> result = parseResponseToResultListOfMap(restAPIResponse.getBodyAsStream());
+    // List<Map<String, String>> result = parseResponseToResultListOfMap(restAPIResponse.getResponseBody());
     if (result != null && !result.isEmpty()) {
       Map<String, String> firstRecord = result.get(0);
       for (String key : firstRecord.keySet()) {

src/main/java/io/cdap/plugin/servicenow/connector/ServiceNowConnector.java

@@ -135,7 +135,7 @@ private TableList listTables(String accessToken) throws ServiceNowAPIException {
     ServiceNowTableAPIClientImpl serviceNowTableAPIClient = new ServiceNowTableAPIClientImpl(config, true);
     RestAPIResponse apiResponse =
         serviceNowTableAPIClient.executeGetWithRetries(requestBuilder.build());
-    return GSON.fromJson(apiResponse.getResponseBody(), TableList.class);
+    return GSON.fromJson(serviceNowTableAPIClient.createJsonReader(apiResponse.getBodyAsStream()), TableList.class);
   }
 
   public ConnectorSpec generateSpec(ConnectorContext connectorContext, ConnectorSpecRequest connectorSpecRequest) {
@@ -184,7 +184,7 @@ private List<StructuredRecord> getTableData(String tableName, int limit)
     requestBuilder.setResponseHeaders(ServiceNowConstants.HEADER_NAME_TOTAL_COUNT);
     RestAPIResponse apiResponse = serviceNowTableAPIClient.executeGetWithRetries(requestBuilder.build());
     List<Map<String, String>> result = serviceNowTableAPIClient.parseResponseToResultListOfMap
-      (apiResponse.getResponseBody());
+      (apiResponse.getBodyAsStream());
     List<StructuredRecord> recordList = new ArrayList<>();
     Schema schema = getSchema(tableName);
     if (schema != null) {

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

@@ -83,7 +83,7 @@ public RestAPIResponse executeGet(RestAPIRequest request) throws IOException {
       }
     } catch (ConnectTimeoutException | SocketException e) {
       ServiceNowAPIException exception = new ServiceNowAPIException(e, null);
-      return new RestAPIResponse(Collections.emptyMap(), (InputStream) null, exception);
+      return new RestAPIResponse(Collections.emptyMap(), null, exception);
     }
   }
 

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

@@ -20,13 +20,16 @@
 import com.google.gson.JsonObject;
 import io.cdap.plugin.servicenow.apiclient.ServiceNowAPIException;
 import io.cdap.plugin.servicenow.util.ServiceNowConstants;
+import org.apache.commons.io.IOUtils;
+import org.apache.commons.io.input.BoundedInputStream;
 import org.apache.http.Header;
 import org.apache.http.HttpEntity;
 import org.apache.http.HttpResponse;
 import org.apache.http.HttpStatus;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
@@ -44,42 +47,33 @@
  * Pojo class to capture the API response.
  */
 public class RestAPIResponse {
+  private static final Logger LOG = LoggerFactory.getLogger(RestAPIResponse.class);
   private static final Gson GSON = new Gson();
   private static final String HTTP_ERROR_MESSAGE = "Http call to ServiceNow instance returned status code %d.";
   private static final String REST_ERROR_MESSAGE = "Rest Api response has errors. Error message: %s.";
   private static final Set<Integer> SUCCESS_CODES = new HashSet<>(Arrays.asList(HttpStatus.SC_CREATED,
                                                                                 HttpStatus.SC_OK));
+  private static final long MAX_PAGE_BYTES = 50L * 1024 * 1024; // 50 MB (Upper Bound)
   private final Map<String, String> headers;
-  // Deprecated: storing full body as String can cause OOM
-  @Deprecated
-  private String responseBody;
   @Nullable private final ServiceNowAPIException exception;
 
-  // New: store InputStream for streaming consumption
-  private InputStream inputStream;
-
-  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;
-  }
+  // New: store byte array
+  private byte[] responseBody;
 
   public RestAPIResponse(
     Map<String, String> headers,
-    InputStream inputStream,
+    byte[] responseBody,
     @Nullable ServiceNowAPIException exception) {
     this.headers = headers;
-    this.inputStream = inputStream;
+    this.responseBody = responseBody;
     this.exception = exception;
   }
 
   /**
    * Parses HttpResponse into RestAPIResponse object when no errors occur.
+   * The RESTAPIResponse contains the HTTP response body as a stream. This stream is:
+   * single-use, forward-only and owned by the caller. Caller is responsible for consuming and closing it.
+   *
    * Throws a {@link ServiceNowAPIException}.
    *
    * @param httpResponse The HttpResponse object to parse
@@ -100,43 +94,33 @@ public static RestAPIResponse parse(HttpResponse httpResponse, String... headerN
 
     ServiceNowAPIException serviceNowAPIException = validateHttpResponse(httpResponse);
     if (serviceNowAPIException != null) {
-      return new RestAPIResponse(headers, null, null, serviceNowAPIException);
+      return new RestAPIResponse(headers, null, serviceNowAPIException);
     }
-    /*try {
-      responseBody = EntityUtils.toString(httpResponse.getEntity());
-    } catch (IOException e) {
-      return new RestAPIResponse(headers, null, null, new ServiceNowAPIException(e, httpResponse));
-    }*/
     try {
-      return prepareResponseWithBodyAndStream(httpResponse, headers, serviceNowAPIException);
+      return prepareResponseStream(httpResponse, headers, serviceNowAPIException);
     } catch (IOException e) {
-      return new RestAPIResponse(headers, null, null, new ServiceNowAPIException(e, httpResponse));
+      return new RestAPIResponse(headers, null, new ServiceNowAPIException(e, httpResponse));
     }
   }
 
-  public static RestAPIResponse prepareResponseWithBodyAndStream(HttpResponse httpResponse, Map<String, String> headers,
+  public static RestAPIResponse prepareResponseStream(HttpResponse httpResponse, Map<String, String> headers,
       ServiceNowAPIException serviceNowAPIException) throws IOException {
     HttpEntity httpEntity = httpResponse.getEntity();
+    byte[] responseBody = new byte[0];
+    InputStream inputStream;
     if (httpEntity != null) {
-      try (InputStream inputStream = httpEntity.getContent();
-           ByteArrayOutputStream buffer = new ByteArrayOutputStream()) {
-
-        // Copy the InputStream into the ByteArrayOutputStream
-        byte[] data = new byte[8192];
-        int bytesRead;
-        while ((bytesRead = inputStream.read(data)) != -1) {
-          buffer.write(data, 0, bytesRead);
-        }
-        // Convert the buffer to a String for the responseBody
-        String responseBody = buffer.toString(String.valueOf(StandardCharsets.UTF_8));
-        serviceNowAPIException = validateRestApiResponse(httpResponse, responseBody);
-        // Create a new InputStream from the buffer for further processing
-        InputStream reusableStream = new ByteArrayInputStream(buffer.toByteArray());
-        // return new RestAPIResponse(headers, responseBody, serviceNowAPIException);
-        return new RestAPIResponse(headers, responseBody, reusableStream, serviceNowAPIException);
+      inputStream = httpEntity.getContent();
+      BoundedInputStream boundedInputStream = new BoundedInputStream(
+        inputStream, MAX_PAGE_BYTES + 1); // +1 to detect overflow
+      responseBody = IOUtils.toByteArray(boundedInputStream);
+      LOG.info("RAW JSON: {}", new String(responseBody, StandardCharsets.UTF_8));
+      if (responseBody.length > MAX_PAGE_BYTES) {
+        throw new IOException(
+          "ServiceNow page exceeded max allowed size: " + MAX_PAGE_BYTES);
       }
+      return new RestAPIResponse(headers, responseBody, serviceNowAPIException);
     } else {
-      return new RestAPIResponse(headers, null, null, serviceNowAPIException);
+      return new RestAPIResponse(headers, responseBody, serviceNowAPIException);
     }
   }
 
@@ -174,12 +158,16 @@ public Map<String, String> getHeaders() {
   }
 
   @Nullable
-  public String getResponseBody() {
+  public byte[] getResponseBody() {
     return responseBody;
   }
 
-  public InputStream getInputStream() {
-    return inputStream;
+  /**
+   * Returns a fresh InputStream for the response body. Caller must close the stream.
+   * @return InputStream
+   */
+  public InputStream getBodyAsStream() {
+    return responseBody == null ? null : new ByteArrayInputStream(responseBody);
   }
 
   @Nullable