From fe2721786cdff9a3bdff59c6f318ac286887595a Mon Sep 17 00:00:00 2001 From: krisztina-zsihovszki Date: Mon, 20 Mar 2023 13:21:31 +0100 Subject: [PATCH] NIFI-11266 PutGoogleDrive, ListGoogleDrive, FetchGoogleDrive can't access a SharedDrive This closes #7058 Reviewed-by: Mark Bathori Signed-off-by: Nandor Soma Abonyi --- .../gcp/drive/FetchGoogleDrive.java | 2 ++ .../processors/gcp/drive/ListGoogleDrive.java | 8 +++-- .../processors/gcp/drive/PutGoogleDrive.java | 8 +++-- .../gcp/drive/AbstractGoogleDriveIT.java | 7 ++++ .../gcp/drive/FetchGoogleDriveTest.java | 3 ++ .../gcp/drive/ListGoogleDriveSimpleTest.java | 13 +++++--- .../drive/ListGoogleDriveTestRunnerTest.java | 12 ++++--- .../gcp/drive/PutGoogleDriveIT.java | 32 ++++++++----------- .../gcp/drive/PutGoogleDriveTest.java | 7 +++- 9 files changed, 59 insertions(+), 33 deletions(-) diff --git a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/drive/FetchGoogleDrive.java b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/drive/FetchGoogleDrive.java index 40ae075de1..2f18170758 100644 --- a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/drive/FetchGoogleDrive.java +++ b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/drive/FetchGoogleDrive.java @@ -168,6 +168,7 @@ public class FetchGoogleDrive extends AbstractProcessor implements GoogleDriveTr try (final InputStream driveFileInputStream = driveService .files() .get(fileId) + .setSupportsAllDrives(true) .executeMediaAsInputStream()) { return session.importFrom(driveFileInputStream, flowFile); @@ -178,6 +179,7 @@ public class FetchGoogleDrive extends AbstractProcessor implements GoogleDriveTr return driveService .files() .get(fileId) + .setSupportsAllDrives(true) .setFields("id, name, createdTime, mimeType, size") .execute(); } diff --git a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/drive/ListGoogleDrive.java b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/drive/ListGoogleDrive.java index 4488f9006a..b7855d1102 100644 --- a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/drive/ListGoogleDrive.java +++ b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/drive/ListGoogleDrive.java @@ -78,8 +78,8 @@ import org.apache.nifi.serialization.record.RecordSchema; @TriggerSerially @Tags({"google", "drive", "storage"}) @CapabilityDescription("Lists concrete files (shortcuts are ignored) in a Google Drive folder. " + - "Each listed file may result in one flowfile, the metadata being written as flowfile attributes. " + - "Or - in case the 'Record Writer' property is set - the entire result is written as records to a single flowfile. " + + "Each listed file may result in one FlowFile, the metadata being written as FlowFile attributes. " + + "Or - in case the 'Record Writer' property is set - the entire result is written as records to a single FlowFile. " + "This Processor is designed to run on Primary Node only in a cluster. If the primary node changes, the new Primary Node will pick up where the " + "previous node left off without duplicating all of the data. " + "Please see Additional Details to set up access to Google Drive.") @@ -266,6 +266,8 @@ public class ListGoogleDrive extends AbstractListProcessor do { FileList result = driveService.files() .list() + .setSupportsAllDrives(true) + .setIncludeItemsFromAllDrives(true) .setQ(queryBuilder.toString()) .setPageToken(pageToken) .setFields("nextPageToken, files(id, name, size, createdTime, modifiedTime, mimeType)") @@ -329,6 +331,8 @@ public class ListGoogleDrive extends AbstractListProcessor do { FileList directoryList = service.files() .list() + .setSupportsAllDrives(true) + .setIncludeItemsFromAllDrives(true) .setQ("'" + folderId + "' in parents " + "and mimeType = 'application/vnd.google-apps.folder'" ) diff --git a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/drive/PutGoogleDrive.java b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/drive/PutGoogleDrive.java index 751a6397fb..f0b8683c6e 100644 --- a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/drive/PutGoogleDrive.java +++ b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/drive/PutGoogleDrive.java @@ -187,7 +187,7 @@ public class PutGoogleDrive extends AbstractProcessor implements GoogleDriveTrai REL_FAILURE ))); - public static final String MULTIPART_UPLOAD_URL = "https://www.googleapis.com/upload/drive/v3/files?uploadType=multipart"; + public static final String MULTIPART_UPLOAD_URL = "https://www.googleapis.com/upload/drive/v3/files?uploadType=multipart&supportsAllDrives=true"; private volatile Drive driveService; @@ -324,10 +324,12 @@ public class PutGoogleDrive extends AbstractProcessor implements GoogleDriveTrai if (fileMetadata.getId() == null) { return driveService.files() .create(fileMetadata, mediaContent) + .setSupportsAllDrives(true) .setFields("id, name, createdTime, mimeType, size"); } else { return driveService.files() .update(fileMetadata.getId(), new File(), mediaContent) + .setSupportsAllDrives(true) .setFields("id, name, createdTime, mimeType, size"); } } @@ -347,7 +349,7 @@ public class PutGoogleDrive extends AbstractProcessor implements GoogleDriveTrai return fileMetadata; } else { throw new ProcessException(format("Upload of File [%s] to Folder [%s] failed, HTTP error code: [%d]", - fileMetadata.getName(), fileMetadata.getId(), response.getStatusCode())); + fileMetadata.getName(), fileMetadata.getParents().stream().findFirst().orElse(""), response.getStatusCode())); } } @@ -368,6 +370,8 @@ public class PutGoogleDrive extends AbstractProcessor implements GoogleDriveTrai private Optional checkFileExistence(String fileName, String parentId) throws IOException { final FileList result = driveService.files() .list() + .setSupportsAllDrives(true) + .setIncludeItemsFromAllDrives(true) .setQ(format("name='%s' and ('%s' in parents)", fileName, parentId)) .setFields("files(name, id)") .execute(); diff --git a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/AbstractGoogleDriveIT.java b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/AbstractGoogleDriveIT.java index d2a16a709a..0ff72a5005 100644 --- a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/AbstractGoogleDriveIT.java +++ b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/AbstractGoogleDriveIT.java @@ -25,6 +25,7 @@ import com.google.api.client.json.gson.GsonFactory; import com.google.api.services.drive.Drive; import com.google.api.services.drive.DriveScopes; import com.google.api.services.drive.model.File; +import org.apache.commons.lang3.StringUtils; import org.apache.nifi.processor.Processor; import org.apache.nifi.processors.gcp.credentials.factory.CredentialPropertyDescriptors; import org.apache.nifi.processors.gcp.credentials.service.GCPCredentialsControllerService; @@ -45,6 +46,7 @@ import java.util.Arrays; * SHARED_FOLDER_ID - The ID of a Folder that is shared with the Service Account. The test will create files and sub-folders within this folder.
*
* Created files and folders are cleaned up, but it's advisable to dedicate a folder for this test so that it can be cleaned up easily should the test fail to do so. + * In case your shared folder is located on a shared drive, give Service Account "Manager" permission on the shared drive to make it capable of deleting the created test files. *

* WARNING: The creation of a file is not a synchronized operation, may need to adjust tests accordingly! */ @@ -52,6 +54,8 @@ public abstract class AbstractGoogleDriveIT expected = Arrays.asList( + List expected = singletonList( new GoogleDriveFileInfo.Builder() .id(id) .fileName(filename) @@ -111,7 +114,7 @@ public class ListGoogleDriveSimpleTest { List actual = testSubject.performListing(mockProcessContext, minTimestamp, null); // THEN - List> propertyProviders = Arrays.asList( + List> propertyProviders = asList( GoogleDriveFileInfo::getId, GoogleDriveFileInfo::getIdentifier, GoogleDriveFileInfo::getName, diff --git a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/ListGoogleDriveTestRunnerTest.java b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/ListGoogleDriveTestRunnerTest.java index da628c2515..f883d2029c 100644 --- a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/ListGoogleDriveTestRunnerTest.java +++ b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/ListGoogleDriveTestRunnerTest.java @@ -17,7 +17,7 @@ package org.apache.nifi.processors.gcp.drive; import static java.lang.String.valueOf; -import static java.util.Arrays.asList; +import static java.util.Collections.singletonList; import static org.mockito.Mockito.RETURNS_DEEP_STUBS; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -50,7 +50,7 @@ public class ListGoogleDriveTestRunnerTest implements OutputChecker { private Drive mockDriverService; - private String folderId = "folderId"; + private final String folderId = "folderId"; @BeforeEach void setUp() throws Exception { @@ -118,7 +118,7 @@ public class ListGoogleDriveTestRunnerTest implements OutputChecker { mockFetchedGoogleDriveFileList(id, filename, size, createdTime, modifiedTime, mimeType); - List expectedContents = asList( + List expectedContents = singletonList( "[" + "{" + "\"drive.id\":\"" + id + "\"," + @@ -143,12 +143,14 @@ public class ListGoogleDriveTestRunnerTest implements OutputChecker { private void mockFetchedGoogleDriveFileList(String id, String filename, Long size, Long createdTime, Long modifiedTime, String mimeType) throws IOException { when(mockDriverService.files() .list() + .setSupportsAllDrives(true) + .setIncludeItemsFromAllDrives(true) .setQ("('" + folderId + "' in parents) and (mimeType != 'application/vnd.google-apps.folder') and (mimeType != 'application/vnd.google-apps.shortcut') and trashed = false") .setPageToken(null) .setFields("nextPageToken, files(id, name, size, createdTime, modifiedTime, mimeType)") .execute() .getFiles() - ).thenReturn(asList( + ).thenReturn(singletonList( createFile( id, filename, @@ -170,7 +172,7 @@ public class ListGoogleDriveTestRunnerTest implements OutputChecker { inputFlowFileAttributes.put(GoogleDriveAttributes.TIMESTAMP, valueOf(expectedTimestamp)); inputFlowFileAttributes.put(GoogleDriveAttributes.MIME_TYPE, mimeType); - HashSet> expectedAttributes = new HashSet<>(asList(inputFlowFileAttributes)); + HashSet> expectedAttributes = new HashSet<>(singletonList(inputFlowFileAttributes)); testRunner.run(); diff --git a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/PutGoogleDriveIT.java b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/PutGoogleDriveIT.java index 5ee3dd550f..459323d84c 100644 --- a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/PutGoogleDriveIT.java +++ b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/PutGoogleDriveIT.java @@ -48,14 +48,11 @@ public class PutGoogleDriveIT extends AbstractGoogleDriveIT impl @Test void testUploadFileToFolderById() { - // GIVEN testRunner.setProperty(FOLDER_ID, mainFolderId); testRunner.setProperty(FILE_NAME, TEST_FILENAME); - // WHEN runWithFileContent(); - // THEN testRunner.assertTransferCount(PutGoogleDrive.REL_SUCCESS, 1); testRunner.assertTransferCount(PutGoogleDrive.REL_FAILURE, 0); final List flowFiles = testRunner.getFlowFilesForRelationship(PutGoogleDrive.REL_SUCCESS); @@ -65,22 +62,17 @@ public class PutGoogleDriveIT extends AbstractGoogleDriveIT impl @Test void testUploadedFileAlreadyExistsFailResolution() { - // GIVEN testRunner.setProperty(FOLDER_ID, mainFolderId); testRunner.setProperty(FILE_NAME, TEST_FILENAME); - // WHEN runWithFileContent(); - // THEN testRunner.assertTransferCount(PutGoogleDrive.REL_SUCCESS, 1); testRunner.assertTransferCount(PutGoogleDrive.REL_FAILURE, 0); testRunner.clearTransferState(); - // WHEN runWithFileContent(); - // THEN testRunner.assertTransferCount(PutGoogleDrive.REL_SUCCESS, 0); testRunner.assertTransferCount(PutGoogleDrive.REL_FAILURE, 1); @@ -88,23 +80,18 @@ public class PutGoogleDriveIT extends AbstractGoogleDriveIT impl @Test void testUploadedFileAlreadyExistsReplaceResolution() { - // GIVEN testRunner.setProperty(FOLDER_ID, mainFolderId); testRunner.setProperty(FILE_NAME, TEST_FILENAME); testRunner.setProperty(PutGoogleDrive.CONFLICT_RESOLUTION, REPLACE.getValue()); - // WHEN runWithFileContent(); - // THEN testRunner.assertTransferCount(PutGoogleDrive.REL_SUCCESS, 1); testRunner.assertTransferCount(PutGoogleDrive.REL_FAILURE, 0); testRunner.clearTransferState(); - // WHEN runWithFileContent("012345678"); - // THEN testRunner.assertTransferCount(PutGoogleDrive.REL_SUCCESS, 1); testRunner.assertTransferCount(PutGoogleDrive.REL_FAILURE, 0); @@ -115,27 +102,36 @@ public class PutGoogleDriveIT extends AbstractGoogleDriveIT impl @Test void testUploadedFileAlreadyExistsIgnoreResolution() { - // GIVEN testRunner.setProperty(FOLDER_ID, mainFolderId); testRunner.setProperty(FILE_NAME, TEST_FILENAME); testRunner.setProperty(PutGoogleDrive.CONFLICT_RESOLUTION, IGNORE.getValue()); - // WHEN runWithFileContent(); - // THEN testRunner.assertTransferCount(PutGoogleDrive.REL_SUCCESS, 1); testRunner.assertTransferCount(PutGoogleDrive.REL_FAILURE, 0); testRunner.clearTransferState(); - // WHEN runWithFileContent(); - // THEN testRunner.assertTransferCount(PutGoogleDrive.REL_SUCCESS, 1); testRunner.assertTransferCount(PutGoogleDrive.REL_FAILURE, 0); } + @Test + void testChunkedUpload() { + testRunner.setProperty(FOLDER_ID, mainFolderId); + testRunner.setProperty(FILE_NAME, TEST_FILENAME); + testRunner.setProperty(PutGoogleDrive.CHUNKED_UPLOAD_SIZE, "256 KB"); + testRunner.setProperty(PutGoogleDrive.CHUNKED_UPLOAD_THRESHOLD, "300 KB"); + + runWithFileContent(LARGE_FILE_CONTENT); + + testRunner.assertTransferCount(PutGoogleDrive.REL_SUCCESS, 1); + testRunner.assertTransferCount(PutGoogleDrive.REL_FAILURE, 0); + } + + private void runWithFileContent() { runWithFileContent(DEFAULT_FILE_CONTENT); } diff --git a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/PutGoogleDriveTest.java b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/PutGoogleDriveTest.java index 1e1905ac62..5c9545d6a2 100644 --- a/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/PutGoogleDriveTest.java +++ b/nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/test/java/org/apache/nifi/processors/gcp/drive/PutGoogleDriveTest.java @@ -192,6 +192,7 @@ public class PutGoogleDriveTest extends AbstractGoogleDriveTest{ private void mockFileUpload(File uploadedFile) throws IOException { when(mockDriverService.files() .create(any(File.class), any(InputStreamContent.class)) + .setSupportsAllDrives(true) .setFields("id, name, createdTime, mimeType, size") .execute()) .thenReturn(uploadedFile); @@ -200,6 +201,7 @@ public class PutGoogleDriveTest extends AbstractGoogleDriveTest{ private void mockFileUpdate(File uploadedFile) throws IOException { when(mockDriverService.files() .update(eq(uploadedFile.getId()), any(File.class), any(InputStreamContent.class)) + .setSupportsAllDrives(true) .setFields("id, name, createdTime, mimeType, size") .execute()) .thenReturn(uploadedFile); @@ -207,13 +209,16 @@ public class PutGoogleDriveTest extends AbstractGoogleDriveTest{ private void mockFileUploadError(Exception exception) throws IOException { when(mockDriverService.files() - .create(any(File.class), any(InputStreamContent.class))) + .create(any(File.class), any(InputStreamContent.class)) + .setSupportsAllDrives(true)) .thenThrow(exception); } private void mockFileExists(List fileList) throws IOException { when(mockDriverService.files() .list() + .setSupportsAllDrives(true) + .setIncludeItemsFromAllDrives(true) .setQ(format("name='%s' and ('%s' in parents)", TEST_FILENAME, SHARED_FOLDER_ID)) .setFields("files(name, id)") .execute())