Skip to content

Commit

Permalink
add test
Browse files Browse the repository at this point in the history
  • Loading branch information
JiamingMai committed Dec 29, 2024
1 parent 4d76f58 commit 3bc249e
Show file tree
Hide file tree
Showing 4 changed files with 81 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -24,19 +24,18 @@
import java.util.List;
import java.util.Optional;

import static io.trino.filesystem.alluxio.AlluxioUtils.convertToLocation;
import static java.util.Objects.requireNonNull;

public class AlluxioFileIterator
implements FileIterator
{
private final Iterator<URIStatus> files;
private final String mountRoot;
private final String basePath;

public AlluxioFileIterator(List<URIStatus> files, String mountRoot)
public AlluxioFileIterator(List<URIStatus> files, String basePath)
{
this.files = requireNonNull(files.iterator(), "files is null");
this.mountRoot = requireNonNull(mountRoot, "mountRoot is null");
this.basePath = requireNonNull(basePath, "basePath is null");
}

@Override
Expand All @@ -54,7 +53,8 @@ public FileEntry next()
return null;
}
URIStatus fileStatus = files.next();
Location location = convertToLocation(fileStatus.getPath(), mountRoot);
String filePath = fileStatus.getPath();
Location location = Location.of(basePath + filePath);
return new FileEntry(
location,
fileStatus.getLength(),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
import java.util.stream.Collectors;

import static io.trino.filesystem.alluxio.AlluxioUtils.convertToAlluxioURI;
import static io.trino.filesystem.alluxio.AlluxioUtils.getAlluxioBase;
import static java.util.Objects.requireNonNull;
import static java.util.UUID.randomUUID;

Expand Down Expand Up @@ -189,20 +190,20 @@ public FileIterator listFiles(Location location)
try {
URIStatus status = alluxioClient.getStatus(convertToAlluxioURI(location, mountRoot));
if (status == null) {
new AlluxioFileIterator(Collections.emptyList(), mountRoot);
new AlluxioFileIterator(Collections.emptyList(), getAlluxioBase(location.toString()));
}
if (!status.isFolder()) {
throw new IOException("Location is not a directory: %s".formatted(location));
}
}
catch (NotFoundRuntimeException | AlluxioException e) {
return new AlluxioFileIterator(Collections.emptyList(), mountRoot);
return new AlluxioFileIterator(Collections.emptyList(), getAlluxioBase(location.toString()));
}

try {
List<URIStatus> filesStatus = alluxioClient.listStatus(convertToAlluxioURI(location, mountRoot),
ListStatusPOptions.newBuilder().setRecursive(true).build());
return new AlluxioFileIterator(filesStatus.stream().filter(status -> !status.isFolder() & status.isCompleted()).toList(), mountRoot);
return new AlluxioFileIterator(filesStatus.stream().filter(status -> !status.isFolder() & status.isCompleted()).toList(), getAlluxioBase(location.toString()));
}
catch (AlluxioException e) {
throw new IOException("Error listFiles %s".formatted(location), e);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,16 @@ public static Location convertToLocation(String path, String mountRoot)
return Location.of(schema + mountRootWithSlash + path);
}

public static String getAlluxioBase(String path)
{
requireNonNull(path, "path is null");
if (!path.startsWith("alluxio://")) {
throw new IllegalArgumentException("path is not an alluxio://");
}
int index = path.indexOf('/', "alluxio://".length());
return path.substring(0, index);
}

public static String simplifyPath(String path)
{
// Use a deque to store the path components
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.trino.filesystem.alluxio;

import alluxio.client.file.URIStatus;
import alluxio.wire.FileInfo;
import io.trino.filesystem.FileEntry;
import org.junit.jupiter.api.Test;

import java.io.IOException;
import java.util.List;

import static org.assertj.core.api.Assertions.assertThat;

final class TestAlluxioFileIterator
{
private final String fileName = "000000_0";
private final String filePath = "/s3a/tables/sales/000000_0";
private final String ufsFilePath = "s3a://test-bucket/tables/sales/000000_0";

@Test
void testNext()
throws IOException
{
String alluxioDirPath = "alluxio://master:19998";
FileInfo fileInfo = new FileInfo();
fileInfo.setName(fileName);
fileInfo.setPath(filePath);
fileInfo.setUfsPath(ufsFilePath);
URIStatus fileStatus = new URIStatus(fileInfo);
AlluxioFileIterator iterator = new AlluxioFileIterator(
List.of(fileStatus),
alluxioDirPath);
FileEntry fileEntry = iterator.next();
assertThat(fileEntry.location().toString())
.isEqualTo(alluxioDirPath + filePath);

alluxioDirPath = "alluxio:/";
fileInfo = new FileInfo();
fileInfo.setName(fileName);
fileInfo.setPath(filePath);
fileInfo.setUfsPath(ufsFilePath);
fileStatus = new URIStatus(fileInfo);
iterator = new AlluxioFileIterator(
List.of(fileStatus),
alluxioDirPath);
fileEntry = iterator.next();
assertThat(fileEntry.location().toString())
.isEqualTo(alluxioDirPath + filePath);
}
}

0 comments on commit 3bc249e

Please sign in to comment.