Search in sources :

Example 46 with StorageObject

use of com.google.api.services.storage.model.StorageObject in project beam by apache.

the class GcsUtilTest method testRecursiveGlobExpansion.

@Test
public void testRecursiveGlobExpansion() throws IOException {
    GcsOptions pipelineOptions = gcsOptionsWithTestCredential();
    GcsUtil gcsUtil = pipelineOptions.getGcsUtil();
    Storage mockStorage = Mockito.mock(Storage.class);
    gcsUtil.setStorageClient(mockStorage);
    Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class);
    Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class);
    Storage.Objects.List mockStorageList = Mockito.mock(Storage.Objects.List.class);
    Objects modelObjects = new Objects();
    List<StorageObject> items = new ArrayList<>();
    // A directory
    items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/"));
    // Files within the directory
    items.add(new StorageObject().setBucket("testbucket").setName("test/directory/file1.txt"));
    items.add(new StorageObject().setBucket("testbucket").setName("test/directory/file2.txt"));
    items.add(new StorageObject().setBucket("testbucket").setName("test/directory/file3.txt"));
    items.add(new StorageObject().setBucket("testbucket").setName("test/directory/otherfile"));
    items.add(new StorageObject().setBucket("testbucket").setName("test/directory/anotherfile"));
    items.add(new StorageObject().setBucket("testbucket").setName("test/file4.txt"));
    modelObjects.setItems(items);
    when(mockStorage.objects()).thenReturn(mockStorageObjects);
    when(mockStorageObjects.get("testbucket", "test/directory/otherfile")).thenReturn(mockStorageGet);
    when(mockStorageObjects.list("testbucket")).thenReturn(mockStorageList);
    when(mockStorageGet.execute()).thenReturn(new StorageObject().setBucket("testbucket").setName("test/directory/otherfile"));
    when(mockStorageList.execute()).thenReturn(modelObjects);
    {
        GcsPath pattern = GcsPath.fromUri("gs://testbucket/test/**/*.txt");
        List<GcsPath> expectedFiles = ImmutableList.of(GcsPath.fromUri("gs://testbucket/test/directory/file1.txt"), GcsPath.fromUri("gs://testbucket/test/directory/file2.txt"), GcsPath.fromUri("gs://testbucket/test/directory/file3.txt"), GcsPath.fromUri("gs://testbucket/test/file4.txt"));
        assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray()));
    }
}
Also used : Storage(com.google.api.services.storage.Storage) StorageObject(com.google.api.services.storage.model.StorageObject) Objects(com.google.api.services.storage.model.Objects) ArrayList(java.util.ArrayList) GcsPath(org.apache.beam.sdk.util.gcsfs.GcsPath) List(java.util.List) ArrayList(java.util.ArrayList) ImmutableList(com.google.common.collect.ImmutableList) GcsOptions(org.apache.beam.sdk.extensions.gcp.options.GcsOptions) Test(org.junit.Test)

Example 47 with StorageObject

use of com.google.api.services.storage.model.StorageObject in project beam by apache.

the class GcsUtilTest method testGlobExpansion.

@Test
public void testGlobExpansion() throws IOException {
    GcsOptions pipelineOptions = gcsOptionsWithTestCredential();
    GcsUtil gcsUtil = pipelineOptions.getGcsUtil();
    Storage mockStorage = Mockito.mock(Storage.class);
    gcsUtil.setStorageClient(mockStorage);
    Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class);
    Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class);
    Storage.Objects.List mockStorageList = Mockito.mock(Storage.Objects.List.class);
    Objects modelObjects = new Objects();
    List<StorageObject> items = new ArrayList<>();
    // A directory
    items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/"));
    // Files within the directory
    items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/file1name"));
    items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/file2name"));
    items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/file3name"));
    items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/otherfile"));
    items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/anotherfile"));
    modelObjects.setItems(items);
    when(mockStorage.objects()).thenReturn(mockStorageObjects);
    when(mockStorageObjects.get("testbucket", "testdirectory/otherfile")).thenReturn(mockStorageGet);
    when(mockStorageObjects.list("testbucket")).thenReturn(mockStorageList);
    when(mockStorageGet.execute()).thenReturn(new StorageObject().setBucket("testbucket").setName("testdirectory/otherfile"));
    when(mockStorageList.execute()).thenReturn(modelObjects);
    // Test a single file.
    {
        GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/otherfile");
        List<GcsPath> expectedFiles = ImmutableList.of(GcsPath.fromUri("gs://testbucket/testdirectory/otherfile"));
        assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray()));
    }
    // Test patterns.
    {
        GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/file*");
        List<GcsPath> expectedFiles = ImmutableList.of(GcsPath.fromUri("gs://testbucket/testdirectory/file1name"), GcsPath.fromUri("gs://testbucket/testdirectory/file2name"), GcsPath.fromUri("gs://testbucket/testdirectory/file3name"));
        assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray()));
    }
    {
        GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/file[1-3]*");
        List<GcsPath> expectedFiles = ImmutableList.of(GcsPath.fromUri("gs://testbucket/testdirectory/file1name"), GcsPath.fromUri("gs://testbucket/testdirectory/file2name"), GcsPath.fromUri("gs://testbucket/testdirectory/file3name"));
        assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray()));
    }
    {
        GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/file?name");
        List<GcsPath> expectedFiles = ImmutableList.of(GcsPath.fromUri("gs://testbucket/testdirectory/file1name"), GcsPath.fromUri("gs://testbucket/testdirectory/file2name"), GcsPath.fromUri("gs://testbucket/testdirectory/file3name"));
        assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray()));
    }
    {
        GcsPath pattern = GcsPath.fromUri("gs://testbucket/test*ectory/fi*name");
        List<GcsPath> expectedFiles = ImmutableList.of(GcsPath.fromUri("gs://testbucket/testdirectory/file1name"), GcsPath.fromUri("gs://testbucket/testdirectory/file2name"), GcsPath.fromUri("gs://testbucket/testdirectory/file3name"));
        assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray()));
    }
}
Also used : Storage(com.google.api.services.storage.Storage) StorageObject(com.google.api.services.storage.model.StorageObject) Objects(com.google.api.services.storage.model.Objects) ArrayList(java.util.ArrayList) GcsPath(org.apache.beam.sdk.util.gcsfs.GcsPath) List(java.util.List) ArrayList(java.util.ArrayList) ImmutableList(com.google.common.collect.ImmutableList) GcsOptions(org.apache.beam.sdk.extensions.gcp.options.GcsOptions) Test(org.junit.Test)

Example 48 with StorageObject

use of com.google.api.services.storage.model.StorageObject in project beam by apache.

the class GcsUtil method create.

/**
   * Same as {@link GcsUtil#create(GcsPath, String)} but allows overriding
   * {code uploadBufferSizeBytes}.
   */
public WritableByteChannel create(GcsPath path, String type, Integer uploadBufferSizeBytes) throws IOException {
    GoogleCloudStorageWriteChannel channel = new GoogleCloudStorageWriteChannel(executorService, storageClient, new ClientRequestHelper<StorageObject>(), path.getBucket(), path.getObject(), AsyncWriteChannelOptions.newBuilder().build(), new ObjectWriteConditions(), Collections.<String, String>emptyMap(), type);
    if (uploadBufferSizeBytes != null) {
        channel.setUploadBufferSize(uploadBufferSizeBytes);
    }
    channel.initialize();
    return channel;
}
Also used : StorageObject(com.google.api.services.storage.model.StorageObject) ObjectWriteConditions(com.google.cloud.hadoop.gcsio.ObjectWriteConditions) GoogleCloudStorageWriteChannel(com.google.cloud.hadoop.gcsio.GoogleCloudStorageWriteChannel)

Example 49 with StorageObject

use of com.google.api.services.storage.model.StorageObject in project beam by apache.

the class GcsUtil method enqueueCopy.

private void enqueueCopy(final GcsPath from, final GcsPath to, BatchRequest batch) throws IOException {
    Storage.Objects.Copy copyRequest = storageClient.objects().copy(from.getBucket(), from.getObject(), to.getBucket(), to.getObject(), null);
    copyRequest.queue(batch, new JsonBatchCallback<StorageObject>() {

        @Override
        public void onSuccess(StorageObject obj, HttpHeaders responseHeaders) {
            LOG.debug("Successfully copied {} to {}", from, to);
        }

        @Override
        public void onFailure(GoogleJsonError e, HttpHeaders responseHeaders) throws IOException {
            throw new IOException(String.format("Error trying to copy %s to %s: %s", from, to, e));
        }
    });
}
Also used : HttpHeaders(com.google.api.client.http.HttpHeaders) StorageObject(com.google.api.services.storage.model.StorageObject) Objects(com.google.api.services.storage.model.Objects) GoogleJsonError(com.google.api.client.googleapis.json.GoogleJsonError) IOException(java.io.IOException)

Example 50 with StorageObject

use of com.google.api.services.storage.model.StorageObject in project beam by apache.

the class GcsUtil method expand.

/**
   * Expands a pattern into matched paths. The pattern path may contain globs, which are expanded
   * in the result. For patterns that only match a single object, we ensure that the object
   * exists.
   */
public List<GcsPath> expand(GcsPath gcsPattern) throws IOException {
    Pattern p = null;
    String prefix = null;
    if (isWildcard(gcsPattern)) {
        // Part before the first wildcard character.
        prefix = getNonWildcardPrefix(gcsPattern.getObject());
        p = Pattern.compile(wildcardToRegexp(gcsPattern.getObject()));
    } else {
        // Not a wildcard.
        try {
            // Use a get request to fetch the metadata of the object, and ignore the return value.
            // The request has strong global consistency.
            getObject(gcsPattern);
            return ImmutableList.of(gcsPattern);
        } catch (FileNotFoundException e) {
            // If the path was not found, return an empty list.
            return ImmutableList.of();
        }
    }
    LOG.debug("matching files in bucket {}, prefix {} against pattern {}", gcsPattern.getBucket(), prefix, p.toString());
    String pageToken = null;
    List<GcsPath> results = new LinkedList<>();
    do {
        Objects objects = listObjects(gcsPattern.getBucket(), prefix, pageToken);
        if (objects.getItems() == null) {
            break;
        }
        // Filter objects based on the regex.
        for (StorageObject o : objects.getItems()) {
            String name = o.getName();
            // Skip directories, which end with a slash.
            if (p.matcher(name).matches() && !name.endsWith("/")) {
                LOG.debug("Matched object: {}", name);
                results.add(GcsPath.fromObject(o));
            }
        }
        pageToken = objects.getNextPageToken();
    } while (pageToken != null);
    return results;
}
Also used : Pattern(java.util.regex.Pattern) StorageObject(com.google.api.services.storage.model.StorageObject) FileNotFoundException(java.io.FileNotFoundException) GcsPath(org.apache.beam.sdk.util.gcsfs.GcsPath) Objects(com.google.api.services.storage.model.Objects) LinkedList(java.util.LinkedList)

Aggregations

StorageObject (com.google.api.services.storage.model.StorageObject)81 Test (org.junit.Test)33 Objects (com.google.api.services.storage.model.Objects)21 IOException (java.io.IOException)20 ArrayList (java.util.ArrayList)16 Storage (com.google.api.services.storage.Storage)12 List (java.util.List)9 URI (java.net.URI)8 RetryHelperException (com.google.cloud.RetryHelper.RetryHelperException)7 GcsOptions (org.apache.beam.sdk.extensions.gcp.options.GcsOptions)7 GcsPath (org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath)7 GoogleJsonError (com.google.api.client.googleapis.json.GoogleJsonError)6 Nullable (javax.annotation.Nullable)6 File (java.io.File)5 Pattern (java.util.regex.Pattern)5 DateTime (com.google.api.client.util.DateTime)4 BlobSourceOption (com.google.cloud.storage.Storage.BlobSourceOption)4 BlobTargetOption (com.google.cloud.storage.Storage.BlobTargetOption)4 BlobWriteOption (com.google.cloud.storage.Storage.BlobWriteOption)4 BucketSourceOption (com.google.cloud.storage.Storage.BucketSourceOption)4