From a466d9da15776eb365e854bbe98ea3decd2cb9de Mon Sep 17 00:00:00 2001
From: vikasrathee-cs <vikas.rathee@cloudsufi.com>
Date: Fri, 8 Nov 2024 17:55:38 +0530
Subject: [PATCH] Added check for sheet files.

---
 .../plugin/google/common/APIRequestRetryer.java   | 15 ++++++++++-----
 .../google/common/GoogleDriveFilteringClient.java |  9 ++++++++-
 .../google/sheets/source/SheetTransformer.java    |  3 ++-
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/src/main/java/io/cdap/plugin/google/common/APIRequestRetryer.java b/src/main/java/io/cdap/plugin/google/common/APIRequestRetryer.java
index 334498f..59d18a2 100644
--- a/src/main/java/io/cdap/plugin/google/common/APIRequestRetryer.java
+++ b/src/main/java/io/cdap/plugin/google/common/APIRequestRetryer.java
@@ -76,7 +76,7 @@ public <V> void onRetry(Attempt<V> attempt) {
             GoogleJsonResponseException e = (GoogleJsonResponseException) exceptionCause;
             LOG.warn(String.format(
               "Error code: '%d', message: '%s'. Attempt: '%d'. Delay since first: '%d'. Description: '%s'.",
-              e.getDetails().getCode(),
+              getCodeFromDetailsOfException(e),
               e.getStatusMessage(),
               attempt.getAttemptNumber(),
               attempt.getDelaySinceFirstAttempt(),
@@ -125,21 +125,26 @@ private static boolean checkHttpResponseException(Throwable t) {
 
   private static boolean isTooManyRequestsError(GoogleJsonResponseException e) {
     List<String> possibleMessages = Arrays.asList(TOO_MANY_REQUESTS_MESSAGE, LIMIT_RATE_EXCEEDED_MESSAGE);
-    return e.getDetails().getCode() == TOO_MANY_REQUESTS_CODE && possibleMessages.contains(e.getStatusMessage());
+    return getCodeFromDetailsOfException(e) == TOO_MANY_REQUESTS_CODE && possibleMessages.contains(
+      e.getStatusMessage());
   }
 
   private static boolean isRateLimitError(GoogleJsonResponseException e) {
-    return e.getDetails().getCode() == LIMIT_RATE_EXCEEDED_CODE
+    return getCodeFromDetailsOfException(e) == LIMIT_RATE_EXCEEDED_CODE
       && (LIMIT_RATE_EXCEEDED_MESSAGE.equals(e.getStatusMessage())
       || e.getDetails().getMessage().contains(LIMIT_RATE_EXCEEDED_MESSAGE));
   }
 
   private static boolean isBackendError(GoogleJsonResponseException e) {
-    return e.getDetails().getCode() == BACKEND_ERROR_CODE;
+    return getCodeFromDetailsOfException(e) == BACKEND_ERROR_CODE;
   }
 
   private static boolean isServiceUnavailableError(GoogleJsonResponseException e) {
-    return e.getDetails().getCode() == SERVICE_UNAVAILABLE_CODE;
+    return getCodeFromDetailsOfException(e) == SERVICE_UNAVAILABLE_CODE;
+  }
+
+  private static Integer getCodeFromDetailsOfException(GoogleJsonResponseException e) {
+    return e.getDetails() != null ? e.getDetails().getCode() : null;
   }
 
   private static boolean isRateLimitError(HttpResponseException e) {
diff --git a/src/main/java/io/cdap/plugin/google/common/GoogleDriveFilteringClient.java b/src/main/java/io/cdap/plugin/google/common/GoogleDriveFilteringClient.java
index a666d3f..180b73a 100644
--- a/src/main/java/io/cdap/plugin/google/common/GoogleDriveFilteringClient.java
+++ b/src/main/java/io/cdap/plugin/google/common/GoogleDriveFilteringClient.java
@@ -72,7 +72,14 @@ public List<File> getFilesSummary(List<ExportedType> exportedTypes, int filesNum
       int retrievedFiles = 0;
       int actualFilesNumber = filesNumber;
       if (IdentifierType.FILE_IDENTIFIER.equals(config.getIdentifierType())) {
-        files.add(service.files().get(config.getFileIdentifier()).setSupportsAllDrives(true).execute());
+        File file = service.files().get(config.getFileIdentifier()).setSupportsAllDrives(true).execute();
+        // size will be 1 in case of Sheets plugin with export type as SPREADSHEETS.
+        if (exportedTypes.size() == 1 && exportedTypes.contains(ExportedType.SPREADSHEETS) && !file.getMimeType()
+          .equalsIgnoreCase(DRIVE_SPREADSHEETS_MIME)) {
+          throw new ExecutionException(
+            String.format("File with id: '%s' is not a Google Sheets File.", config.getFileIdentifier()), null);
+        }
+        files.add(file);
         return files;
       }
       Drive.Files.List request = service.files().list()
diff --git a/src/main/java/io/cdap/plugin/google/sheets/source/SheetTransformer.java b/src/main/java/io/cdap/plugin/google/sheets/source/SheetTransformer.java
index e017988..33bdba9 100644
--- a/src/main/java/io/cdap/plugin/google/sheets/source/SheetTransformer.java
+++ b/src/main/java/io/cdap/plugin/google/sheets/source/SheetTransformer.java
@@ -69,7 +69,8 @@ public static StructuredRecord transform(RowRecord rowRecord, Schema schema, boo
         builder.set(metadataRecordName, rowRecord.getMetadata());
       } else {
         ComplexSingleValueColumn complexSingleValueColumn = rowRecord.getHeaderedCells().get(name);
-        if (complexSingleValueColumn.getData() == null && complexSingleValueColumn.getSubColumns().isEmpty()) {
+        if (complexSingleValueColumn == null || (complexSingleValueColumn.getData() == null
+          && complexSingleValueColumn.getSubColumns().isEmpty())) {
           builder.set(name, null);
         } else {
           processCellData(builder, field, complexSingleValueColumn);