PLUGIN-1936: Streaming Implementation

Author: sgarg-CS

src/e2e-test/java/io/cdap/plugin/servicenowsink/actions/ServiceNowSinkPropertiesPageActions.java

@@ -48,7 +48,7 @@
 
 public class ServiceNowSinkPropertiesPageActions {
   public static ServiceNowSourceConfig config;
-  private static Map<String, String> responseFromServiceNowTable;
+  private static JsonObject responseFromServiceNowTable;
   private static Gson gson = new Gson();
 
   public static void getRecordFromServiceNowTable(String query, String tableName)
@@ -71,10 +71,8 @@ public static boolean verifyIfRecordCreatedInServiceNowIsCorrect(String query, S
     String bqTable = TestSetupHooks.bqTargetTable;
     getRecordFromServiceNowTable(query, tableName);
 
-    JsonObject jsonObject = new JsonObject();
-    responseFromServiceNowTable.forEach(jsonObject::addProperty);
     List<JsonObject> serviceNowResponse = new ArrayList<>();
-    serviceNowResponse.add(jsonObject);
+    serviceNowResponse.add(responseFromServiceNowTable);
 
     List<JsonObject> bigQueryResponse = new ArrayList<>();
     List<Object> bigQueryRows = new ArrayList<>();

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

@@ -17,6 +17,7 @@
 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;
@@ -31,14 +32,19 @@
 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.InputStreamReader;
 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 +146,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 +158,31 @@ public void validateTable(String tableName, SourceValueType valueType, FailureCo
     }
   }
 
+  /**
+   * Checks if the "result" array in the ServiceNow REST API response is empty.
+   * <p>
+   * Determines if the "result" array in a ServiceNow REST API response is empty by specifically looking for a top-level
+   * key named "result". Once found, it opens the associated array and checks for the presence of a first element.
+   * </p>
+   * 
+   * @param restAPIResponse The response object containing the JSON input stream
+   * @return true, if the "result" array exists and is empty, or if the "result" key is never found;
+   * false, if the array contains at least one element.
+   * @throws IOException If there is an error reading the input stream or parsing the JSON.
+   */
+  public boolean isResultEmpty(RestAPIResponse restAPIResponse) throws IOException {
+    JsonReader reader = new JsonReader(new InputStreamReader(restAPIResponse.getResponseStream()));
+    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

@@ -27,7 +27,7 @@
 import com.google.gson.JsonArray;
 import com.google.gson.JsonObject;
 import com.google.gson.reflect.TypeToken;
-import com.sun.org.apache.xpath.internal.operations.Bool;
+import com.google.gson.stream.JsonReader;
 import io.cdap.cdap.api.data.schema.Schema;
 import io.cdap.cdap.etl.api.FailureCollector;
 import io.cdap.plugin.servicenow.connector.ServiceNowConnectorConfig;
@@ -56,7 +56,10 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
 import java.lang.reflect.Type;
+import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -86,6 +89,7 @@ public class ServiceNowTableAPIClientImpl extends RestAPIClient {
   public static JsonArray serviceNowJsonResultArray;
 
   public ServiceNowTableAPIClientImpl(ServiceNowConnectorConfig conf, Boolean useConnection) {
+    super();
     this.conf = conf;
     this.schemaType = getSchemaTypeBasedOnUseConnection(useConnection);
   }
@@ -135,7 +139,7 @@ public String getAccessTokenRetryableMode() throws ExecutionException, RetryExce
    * @return A list of maps, where each map represents a record from the table.
    * @throws ServiceNowAPIException If an error occurs while fetching the records.
    */
-  public List<Map<String, String>> fetchTableRecords(
+  public RestAPIResponse fetchTableRecords(
     String tableName,
     SourceValueType valueType,
     String filterQuery,
@@ -159,7 +163,7 @@ public List<Map<String, String>> fetchTableRecords(
     String accessToken = getAccessToken();
     requestBuilder.setAuthHeader(accessToken);
     RestAPIResponse apiResponse = executeGetWithRetries(requestBuilder.build());
-    return parseResponseToResultListOfMap(apiResponse.getResponseBody());
+    return apiResponse;
   }
 
   private int getRecordCountFromHeader(RestAPIResponse apiResponse) {
@@ -209,30 +213,24 @@ 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 filterQuery, int offset,
                                                                   int limit) throws ServiceNowAPIException {
-    final List<Map<String, String>> results = new ArrayList<>();
-    Callable<Boolean> fetchRecords = () -> {
-      results.addAll(fetchTableRecords(tableName, valueType, filterQuery, offset, limit));
-      return true;
-    };
+    Callable<RestAPIResponse> fetchRecords = () -> fetchTableRecords(tableName, valueType, filterQuery, offset, limit);
 
-    Retryer<Boolean> retryer = RetryerBuilder.<Boolean>newBuilder()
+    Retryer<RestAPIResponse> retryer = RetryerBuilder.<RestAPIResponse>newBuilder()
       .retryIfException(this::isExceptionRetryable)
       .withWaitStrategy(WaitStrategies.exponentialWait(ServiceNowConstants.WAIT_TIME, TimeUnit.MILLISECONDS))
       .withStopStrategy(StopStrategies.stopAfterAttempt(ServiceNowConstants.MAX_NUMBER_OF_RETRY_ATTEMPTS))
       .build();
 
     try {
-      retryer.call(fetchRecords);
+      return retryer.call(fetchRecords);
     } catch (RetryException | ExecutionException e) {
       throw new ServiceNowAPIException(
           String.format("Data Recovery failed for batch %s to %s.", offset, (offset + limit)),
           e, null, false);
     }
-
-    return results;
   }
 
   /**
@@ -255,8 +253,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);
   }
 
   /**
@@ -331,7 +329,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.getResponseStream()), SchemaAPISchemaResponse.class);
 
     if (schemaAPISchemaResponse.getResult() == null || schemaAPISchemaResponse.getResult().isEmpty()) {
       throw new ServiceNowAPIException(
@@ -364,7 +362,7 @@ private Schema prepareSchemaWithSchemaAPI(RestAPIResponse restAPIResponse, List<
    */
   private Schema prepareSchemaWithMetadataAPI(RestAPIResponse restAPIResponse, List<ServiceNowColumn> columns,
     String tableName, SourceValueType valueType, Boolean legacyMapping) throws ServiceNowAPIException {
-    MetadataAPISchemaResponse metadataAPISchemaResponse = parseSchemaResponse(restAPIResponse.getResponseBody());
+    MetadataAPISchemaResponse metadataAPISchemaResponse = parseSchemaResponse(restAPIResponse.getResponseStream());
 
     if (metadataAPISchemaResponse.getResult() == null || metadataAPISchemaResponse.getResult().getColumns() == null ||
       metadataAPISchemaResponse.getResult().getColumns().isEmpty()) {
@@ -391,6 +389,11 @@ private Schema prepareSchemaWithMetadataAPI(RestAPIResponse restAPIResponse, Lis
     return SchemaBuilder.constructSchema(tableName, columns, legacyMapping);
   }
 
+  public JsonReader createJsonReader(InputStream inputStream) {
+    Objects.requireNonNull(inputStream, "InputStream must not be null");
+    return new JsonReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
+  }
+
   /**
    * Gets the total number of records in the table based on the provided date range.
    *
@@ -451,8 +454,8 @@ public String createRecord(String tableName, HttpEntity entity) throws IOExcepti
       requestBuilder.setContentTypeHeader("application/json");
       requestBuilder.setEntity(entity);
       apiResponse = executePost(requestBuilder.build());
-
       systemID = String.valueOf(getSystemId(apiResponse));
+      apiResponse.close();
     } catch (IOException e) {
       throw new ServiceNowAPIException("Error in creating a new record", e, null, false);
     }
@@ -482,15 +485,17 @@ public String createRecordInDisplayMode(String tableName, HttpEntity entity) thr
       apiResponse = executePost(requestBuilder.build());
 
       systemID = String.valueOf(getSystemId(apiResponse));
+      apiResponse.close();
     } catch (IOException e) {
       throw new ServiceNowAPIException("Error in creating a new record", e, null, false);
     }
     return systemID;
   }
 
   private String getSystemId(RestAPIResponse restAPIResponse) {
-    CreateRecordAPIResponse apiResponse = GSON.fromJson(restAPIResponse.getResponseBody(),
-                                                           CreateRecordAPIResponse.class);
+    CreateRecordAPIResponse apiResponse = GSON.fromJson(
+      new InputStreamReader(restAPIResponse.getResponseStream(), StandardCharsets.UTF_8),
+        CreateRecordAPIResponse.class);
     return apiResponse.getResult().get(ServiceNowConstants.SYSTEM_ID).toString();
   }
 
@@ -501,7 +506,7 @@ private String getSystemId(RestAPIResponse restAPIResponse) {
    * @param tableName The ServiceNow table name
    * @param query The query
    */
-  public Map<String, String> getRecordFromServiceNowTable(String tableName, String query)
+  public JsonObject getRecordFromServiceNowTable(String tableName, String query)
       throws ServiceNowAPIException {
 
     ServiceNowTableAPIRequestBuilder requestBuilder = new ServiceNowTableAPIRequestBuilder(
@@ -513,7 +518,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.getResponseStream(), StandardCharsets.UTF_8), APIResponse.class);
     return apiResponse.getResult().get(0);
   }
 
@@ -530,19 +536,45 @@ 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,
+    private Schema prepareStringBasedSchema(RestAPIResponse restAPIResponse, List<ServiceNowColumn> columns,
                                           String tableName) {
-    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()) {
-        columns.add(new ServiceNowColumn(key, "string"));
+      InputStream in = restAPIResponse.getResponseStream();
+      JsonReader reader = new JsonReader(new InputStreamReader(in, StandardCharsets.UTF_8));
+      JsonObject firstRecord;
+      try {
+        firstRecord = getFirstRecord(reader);
+        restAPIResponse.close();
+      } catch (IOException e) {
+        throw new RuntimeException("Failed to parse schema for table: " + tableName, e);
       }
+      if (firstRecord != null) {
+      firstRecord.entrySet().forEach(entry ->
+        columns.add(new ServiceNowColumn(entry.getKey(), "string"))
+      );
       return SchemaBuilder.constructSchema(tableName, columns, false);
     }
     return null;
   }
 
+  public JsonObject getFirstRecord(JsonReader reader) throws IOException {
+    reader.beginObject();
+    while (reader.hasNext()) {
+      String name = reader.nextName();
+      if (name.equals(ServiceNowConstants.RESULT)) {
+        reader.beginArray();
+        if (reader.hasNext()) {
+          // ONLY parse the first object into memory
+          return GSON.fromJson(reader, JsonObject.class);
+        }
+        reader.endArray();
+      } else {
+        reader.skipValue();
+      }
+    }
+    reader.endObject();
+    return null;
+  }
+
 
   /**
    * This function is being used in end-to-end (e2e) tests to delete a record from

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

@@ -16,6 +16,8 @@
 package io.cdap.plugin.servicenow.connector;
 
 import com.google.gson.Gson;
+import com.google.gson.JsonObject;
+import com.google.gson.stream.JsonReader;
 import io.cdap.cdap.api.annotation.Description;
 import io.cdap.cdap.api.annotation.Name;
 import io.cdap.cdap.api.annotation.Plugin;
@@ -55,6 +57,9 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -135,7 +140,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.getResponseStream()), TableList.class);
   }
 
   public ConnectorSpec generateSpec(ConnectorContext connectorContext, ConnectorSpecRequest connectorSpecRequest) {
@@ -172,7 +177,7 @@ public List<StructuredRecord> sample(ConnectorContext connectorContext, SampleRe
   }
 
   private List<StructuredRecord> getTableData(String tableName, int limit)
-      throws OAuthProblemException, OAuthSystemException, ServiceNowAPIException {
+    throws OAuthProblemException, OAuthSystemException, ServiceNowAPIException, IOException {
     ServiceNowTableAPIRequestBuilder requestBuilder = new ServiceNowTableAPIRequestBuilder(
       config.getRestApiEndpoint(), tableName, false, SchemaType.SCHEMA_API_BASED)
       .setExcludeReferenceLink(true)
@@ -183,22 +188,40 @@ private List<StructuredRecord> getTableData(String tableName, int limit)
     requestBuilder.setAuthHeader(accessToken);
     requestBuilder.setResponseHeaders(ServiceNowConstants.HEADER_NAME_TOTAL_COUNT);
     RestAPIResponse apiResponse = serviceNowTableAPIClient.executeGetWithRetries(requestBuilder.build());
-    List<Map<String, String>> result = serviceNowTableAPIClient.parseResponseToResultListOfMap
-      (apiResponse.getResponseBody());
+    InputStream in = apiResponse.getResponseStream();
+    JsonReader reader = new JsonReader(new InputStreamReader(in, StandardCharsets.UTF_8));
     List<StructuredRecord> recordList = new ArrayList<>();
     Schema schema = getSchema(tableName);
     if (schema != null) {
       List<Schema.Field> tableFields = schema.getFields();
-      for (int i = 0; i < result.size(); i++) {
-        StructuredRecord.Builder recordBuilder = StructuredRecord.builder(schema);
-        for (Schema.Field field : tableFields) {
-          String fieldName = field.getName();
-          ServiceNowRecordConverter.convertToValue(fieldName, field.getSchema(), result.get(i), recordBuilder, false);
+      try {
+        reader.beginObject();
+        while (reader.hasNext()) {
+          String name = reader.nextName();
+          if (name.equals(ServiceNowConstants.RESULT)) {
+            reader.beginArray();
+            while (reader.hasNext()) {
+              JsonObject jsonObject = GSON.fromJson(reader, JsonObject.class);
+              StructuredRecord.Builder recordBuilder = StructuredRecord.builder(schema);
+              for (Schema.Field field : tableFields) {
+                String fieldName = field.getName();
+                ServiceNowRecordConverter.convertToValue(fieldName, field.getSchema(), jsonObject, recordBuilder,
+                  false);
+              }
+              StructuredRecord structuredRecord = recordBuilder.build();
+              recordList.add(structuredRecord);
+            }
+            reader.endArray();
+          } else {
+            reader.skipValue();
+          }
         }
-        StructuredRecord structuredRecord = recordBuilder.build();
-        recordList.add(structuredRecord);
+        reader.endObject();
+      } catch (IOException e) {
+        throw new RuntimeException("Error reading JSON response for table " + tableName, e);
       }
     }
+    apiResponse.close();
     return recordList;
 
   }