changes for closing httpResponse

Author: sgarg-CS

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

@@ -266,8 +266,8 @@ public MetadataAPISchemaResponse parseSchemaResponse(InputStream responseStream)
    * @throws ServiceNowAPIException
    */
   public Schema fetchTableSchema(String tableName, SourceValueType valueType)
-      throws ServiceNowAPIException {
-    return fetchTableSchema(tableName, getAccessToken(), valueType, schemaType, false);
+    throws ServiceNowAPIException {
+    return fetchTableSchema(tableName, getAccessToken(), valueType, schemaType, true);
   }
 
   private SchemaType getSchemaTypeBasedOnUseConnection(Boolean useConnection) {
@@ -291,7 +291,7 @@ private SchemaType getSchemaTypeBasedOnUseConnection(Boolean useConnection) {
    */
   public Schema fetchTableSchema(String tableName, String accessToken, SourceValueType valueType,
                                  SchemaType schemaType, Boolean legacyMapping)
-      throws ServiceNowAPIException {
+    throws ServiceNowAPIException {
     ServiceNowTableAPIRequestBuilder requestBuilder = new ServiceNowTableAPIRequestBuilder(
       this.conf.getRestApiEndpoint(), tableName, true, schemaType)
       .setExcludeReferenceLink(true);
@@ -301,12 +301,16 @@ public Schema fetchTableSchema(String tableName, String accessToken, SourceValue
     restAPIResponse = executeGetWithRetries(requestBuilder.build());
     List<ServiceNowColumn> columns = new ArrayList<>();
 
-    if (schemaType == SchemaType.METADATA_API_BASED) {
-      return prepareSchemaWithMetadataAPI(restAPIResponse, columns, tableName, valueType, legacyMapping);
-    } else if (schemaType == SchemaType.SCHEMA_API_BASED) {
-      return prepareSchemaWithSchemaAPI(restAPIResponse, columns, tableName);
-    } else {
-      return prepareStringBasedSchema(restAPIResponse, columns, tableName);
+    try {
+      if (schemaType == SchemaType.METADATA_API_BASED) {
+        return prepareSchemaWithMetadataAPI(restAPIResponse, columns, tableName, valueType, legacyMapping);
+      } else if (schemaType == SchemaType.SCHEMA_API_BASED) {
+        return prepareSchemaWithSchemaAPI(restAPIResponse, columns, tableName, legacyMapping);
+      } else {
+        return prepareStringBasedSchema(restAPIResponse, columns, tableName, legacyMapping);
+      }
+    } catch (IOException exception) {
+      throw new RuntimeException("Error in fetching schema for table " + tableName, exception);
     }
   }
 
@@ -327,7 +331,7 @@ public Schema fetchTableSchema(String tableName, String accessToken, SourceValue
    * @throws RuntimeException if the schema response is null or contains no result.
    */
   private Schema prepareSchemaWithSchemaAPI(RestAPIResponse restAPIResponse, List<ServiceNowColumn> columns,
-                                            String tableName) throws ServiceNowAPIException {
+    String tableName, Boolean legacyMapping) throws ServiceNowAPIException, IOException {
     SchemaAPISchemaResponse schemaAPISchemaResponse =
       GSON.fromJson(createJsonReader(restAPIResponse.getResponseStream()), SchemaAPISchemaResponse.class);
 
@@ -339,7 +343,7 @@ private Schema prepareSchemaWithSchemaAPI(RestAPIResponse restAPIResponse, List<
     for (SchemaAPISchemaField field : schemaAPISchemaResponse.getResult()) {
       columns.add(new ServiceNowColumn(field.getName(), field.getInternalType()));
     }
-    return SchemaBuilder.constructSchema(tableName, columns, false);
+    return SchemaBuilder.constructSchema(tableName, columns, legacyMapping);
   }
 
   /**
@@ -361,7 +365,7 @@ private Schema prepareSchemaWithSchemaAPI(RestAPIResponse restAPIResponse, List<
    * @throws RuntimeException if the response does not contain valid column information.
    */
   private Schema prepareSchemaWithMetadataAPI(RestAPIResponse restAPIResponse, List<ServiceNowColumn> columns,
-    String tableName, SourceValueType valueType, Boolean legacyMapping) throws ServiceNowAPIException {
+    String tableName, SourceValueType valueType, Boolean legacyMapping) throws ServiceNowAPIException, IOException {
     MetadataAPISchemaResponse metadataAPISchemaResponse = parseSchemaResponse(restAPIResponse.getResponseStream());
 
     if (metadataAPISchemaResponse.getResult() == null || metadataAPISchemaResponse.getResult().getColumns() == null ||
@@ -492,7 +496,7 @@ public String createRecordInDisplayMode(String tableName, HttpEntity entity) thr
     return systemID;
   }
 
-  private String getSystemId(RestAPIResponse restAPIResponse) {
+  private String getSystemId(RestAPIResponse restAPIResponse) throws IOException {
     CreateRecordAPIResponse apiResponse = GSON.fromJson(
       new InputStreamReader(restAPIResponse.getResponseStream(), StandardCharsets.UTF_8),
         CreateRecordAPIResponse.class);
@@ -507,7 +511,7 @@ private String getSystemId(RestAPIResponse restAPIResponse) {
    * @param query The query
    */
   public JsonObject getRecordFromServiceNowTable(String tableName, String query)
-      throws ServiceNowAPIException {
+    throws ServiceNowAPIException, IOException {
 
     ServiceNowTableAPIRequestBuilder requestBuilder = new ServiceNowTableAPIRequestBuilder(
       this.conf.getRestApiEndpoint(), tableName, false, schemaType)
@@ -537,7 +541,7 @@ public JsonObject getRecordFromServiceNowTable(String tableName, String query)
    * @throws RuntimeException if the schema response is null or contains no result.
    */
     private Schema prepareStringBasedSchema(RestAPIResponse restAPIResponse, List<ServiceNowColumn> columns,
-                                          String tableName) {
+                                          String tableName, Boolean legacyMapping) throws IOException {
       InputStream in = restAPIResponse.getResponseStream();
       JsonReader reader = new JsonReader(new InputStreamReader(in, StandardCharsets.UTF_8));
       JsonObject firstRecord;
@@ -551,7 +555,7 @@ private Schema prepareStringBasedSchema(RestAPIResponse restAPIResponse, List<Se
       firstRecord.entrySet().forEach(entry ->
         columns.add(new ServiceNowColumn(entry.getKey(), "string"))
       );
-      return SchemaBuilder.constructSchema(tableName, columns, false);
+      return SchemaBuilder.constructSchema(tableName, columns, legacyMapping);
     }
     return null;
   }

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

@@ -108,7 +108,7 @@ public BrowseDetail browse(ConnectorContext connectorContext, BrowseRequest brow
    * Browse Details for the given AccessToken.
    */
   public BrowseDetail browse(ConnectorContext connectorContext,
-                             String accessToken) throws ServiceNowAPIException {
+                             String accessToken) throws ServiceNowAPIException, IOException {
     int count = 0;
     FailureCollector collector = connectorContext.getFailureCollector();
     config.validateCredentialsFields(collector);
@@ -131,7 +131,7 @@ public BrowseDetail browse(ConnectorContext connectorContext,
   /**
    * @return the list of tables.
    */
-  private TableList listTables(String accessToken) throws ServiceNowAPIException {
+  private TableList listTables(String accessToken) throws ServiceNowAPIException, IOException {
     ServiceNowTableAPIRequestBuilder requestBuilder = new ServiceNowTableAPIRequestBuilder(
       config.getRestApiEndpoint(), OBJECT_TABLE_LIST, false, SchemaType.SCHEMA_API_BASED);
     requestBuilder.setAuthHeader(accessToken);

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

@@ -59,7 +59,7 @@ public abstract class RestAPIClient {
   private static final int DEFAULT_CONNECT_TIMEOUT_MS = 120000;
 
   /* Read Timeout in ms for waiting for data after the connection is established */
-  private static final int DEFAULT_READ_TIMEOUT_MS = 300000;
+  private static final int DEFAULT_READ_TIMEOUT_MS = 90000;
 
   /* Maximum total connections. */
   private static final int MAX_CONNECTIONS = 200;

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

@@ -24,6 +24,7 @@
 import org.apache.http.HttpEntity;
 import org.apache.http.HttpResponse;
 import org.apache.http.HttpStatus;
+import org.apache.http.client.methods.CloseableHttpResponse;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -52,15 +53,15 @@ public class RestAPIResponse {
   private final Map<String, String> headers;
   @Nullable private final ServiceNowAPIException exception;
 
-  // Input stream of the response body.
-  private InputStream responseStream;
+  // Keep reference of HttpResponse for error handling, as it needs to be closed after consuming the response stream.
+  private HttpResponse httpResponse;
 
   public RestAPIResponse(
     Map<String, String> headers,
-    InputStream responseStream,
+    HttpResponse httpResponse,
     @Nullable ServiceNowAPIException exception) {
     this.headers = headers;
-    this.responseStream = responseStream;
+    this.httpResponse = httpResponse;
     this.exception = exception;
   }
 
@@ -89,22 +90,27 @@ public static RestAPIResponse parse(HttpResponse httpResponse, String... headerN
 
     ServiceNowAPIException serviceNowAPIException = validateHttpResponse(httpResponse);
     if (serviceNowAPIException != null) {
-      return new RestAPIResponse(headers, null, serviceNowAPIException);
-    }
-    try {
-      return prepareResponse(httpResponse, headers, serviceNowAPIException);
-    } catch (IOException e) {
-      return new RestAPIResponse(headers, null, new ServiceNowAPIException(e, httpResponse));
+      return new RestAPIResponse(headers, httpResponse, serviceNowAPIException);
     }
+    return new RestAPIResponse(headers, httpResponse, null);
   }
 
   public void close() throws IOException {
-    if (responseStream != null) {
-      responseStream.close();
+    try {
+      InputStream responseStream = getResponseStream();
+      if (responseStream != null) {
+        LOG.info("Closing response stream");
+        responseStream.close();
+      }
+    } finally {
+      if (httpResponse instanceof CloseableHttpResponse) {
+        LOG.info("Closing HttpResponse");
+        ((CloseableHttpResponse) httpResponse).close();
+      }
     }
   }
 
-  public static RestAPIResponse prepareResponse(HttpResponse httpResponse, Map<String, String> headers,
+  /*public static RestAPIResponse prepareResponse(HttpResponse httpResponse, Map<String, String> headers,
       ServiceNowAPIException serviceNowAPIException) throws IOException {
     HttpEntity httpEntity = httpResponse.getEntity();
     InputStream inputStream;
@@ -114,7 +120,7 @@ public static RestAPIResponse prepareResponse(HttpResponse httpResponse, Map<Str
     } else {
       return new RestAPIResponse(headers, null, serviceNowAPIException);
     }
-  }
+  }*/
 
   public static RestAPIResponse parse(HttpResponse httpResponse) throws IOException {
     return parse(httpResponse, new String[0]);
@@ -158,7 +164,11 @@ public boolean hasException() {
     return exception != null;
   }
 
-  public InputStream getResponseStream() {
-    return responseStream;
+  public InputStream getResponseStream() throws IOException {
+    if (httpResponse == null) {
+      return null;
+    }
+    HttpEntity entity = httpResponse.getEntity();
+    return (entity != null) ? entity.getContent() : null;
   }
 }

src/main/java/io/cdap/plugin/servicenow/sink/service/ServiceNowSinkAPIRequestImpl.java

@@ -171,7 +171,7 @@ public void createPostRequest(Map<String, RestRequest> restRequestsMap, String a
     }
   }
 
-  private JsonReader getJsonReader(RestAPIResponse apiResponse) {
+  private JsonReader getJsonReader(RestAPIResponse apiResponse) throws IOException {
     InputStreamReader inputStreamReader = new InputStreamReader(apiResponse.getResponseStream(),
       StandardCharsets.UTF_8);
     JsonReader jsonReader = new JsonReader(inputStreamReader);

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

@@ -51,6 +51,7 @@ public abstract class ServiceNowBaseRecordReader extends RecordReader<NullWritab
   protected JsonObject row;
   protected final Gson gson = new Gson();
   protected JsonReader jsonReader = null;
+  protected RestAPIResponse currentResponse = null;
 
   public ServiceNowBaseRecordReader() {
   }
@@ -93,9 +94,10 @@ public boolean nextKeyValue() throws IOException {
 
   public boolean openNextPage() throws IOException, ServiceNowAPIException {
     closeCurrentPage();
-    RestAPIResponse resp = fetchData();
-    InputStream in = resp.getResponseStream();
+    this.currentResponse = fetchData();
+    InputStream in = this.currentResponse.getResponseStream();
     if (in == null) {
+      closeCurrentPage();
       return false;
     }
     this.jsonReader = new JsonReader(new InputStreamReader(in, StandardCharsets.UTF_8));
@@ -137,7 +139,6 @@ public boolean openNextPage() throws IOException, ServiceNowAPIException {
         jsonReader.endArray();
         // cleanup
         closeCurrentPage();
-        closeRestAPIResponse(resp);
         return false;
       }
     } catch (IOException e) {
@@ -151,6 +152,7 @@ public boolean openNextPage() throws IOException, ServiceNowAPIException {
   }
 
   public void closeCurrentPage() {
+    LOG.info("Closing current page for table {}", tableName);
     if (this.jsonReader != null) {
       try {
         this.jsonReader.close();
@@ -160,14 +162,14 @@ public void closeCurrentPage() {
         this.jsonReader = null;
       }
     }
-  }
 
-  public void closeRestAPIResponse(RestAPIResponse resp) {
-    if (resp != null) {
+    if (this.currentResponse != null) {
       try {
-        resp.close();
+        this.currentResponse.close();
       } catch (IOException e) {
-        LOG.warn("Error closing RestAPIResponse", e);
+        LOG.warn("Error closing RestAPIResponse for table {}", tableName, e);
+      } finally {
+        this.currentResponse = null;
       }
     }
   }
@@ -184,6 +186,11 @@ public float getProgress() throws IOException, InterruptedException {
     return pos / (float) split.getLength();
   }
 
+  public int getPos() {
+    return pos;
+  }
+
   public void close() throws IOException {
+    closeCurrentPage();
   }
 }

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

@@ -87,6 +87,8 @@ public StructuredRecord getCurrentValue() throws IOException {
 
     @Override
     RestAPIResponse fetchData() throws ServiceNowAPIException {
+    LOG.info("Fetching data for table: {}, with offset: {} and page size: {}", tableName, split.getOffset(),
+      pluginConf.getPageSize());
     // Get the table data
     RestAPIResponse restAPIResponse = restApi.fetchTableRecordsRetryableMode(tableName, pluginConf.getValueType(),
      split.getFilterQuery() , split.getOffset(), pluginConf.getPageSize());