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()));
}
}
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()));
}
}
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;
}
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));
}
});
}
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;
}
Aggregations