check data token before call fileIO function

Author: jerry-024

paimon-core/src/test/java/org/apache/paimon/rest/RESTCatalogTest.java

@@ -55,6 +55,7 @@
 import org.junit.jupiter.api.Test;
 
 import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -238,7 +239,9 @@ public void testDataTokenExpired() throws Exception {
         FileStoreTable tableTestWrite = (FileStoreTable) catalog.getTable(identifier);
         restCatalogServer.setFileIO(identifier, tableTestWrite.fileIO());
         List<Integer> data = Lists.newArrayList(12);
-        assertThrows(IOException.class, () -> batchWrite(tableTestWrite, data));
+        Exception exception =
+                assertThrows(UncheckedIOException.class, () -> batchWrite(tableTestWrite, data));
+        assertEquals(RESTTestFileIO.TOKEN_EXPIRED_MSG, exception.getCause().getMessage());
         RESTToken dataToken =
                 new RESTToken(
                         ImmutableMap.of("akId", "akId", "akSecret", UUID.randomUUID().toString()),
@@ -262,21 +265,18 @@ public void testDataTokenUnExistInServer() throws Exception {
         restTokenFileIO.isObjectStore();
         restCatalogServer.removeDataToken(identifier);
         restCatalogServer.setFileIO(identifier, tableTestWrite.fileIO());
-        assertThrows(IOException.class, () -> batchWrite(tableTestWrite, data));
+        Exception exception =
+                assertThrows(UncheckedIOException.class, () -> batchWrite(tableTestWrite, data));
+        assertEquals(RESTTestFileIO.TOKEN_UN_EXIST_MSG, exception.getCause().getMessage());
     }
 
     private void batchWrite(FileStoreTable tableTestWrite, List<Integer> data) throws Exception {
         BatchWriteBuilder writeBuilder = tableTestWrite.newBatchWriteBuilder();
         BatchTableWrite write = writeBuilder.newWrite();
-        data.forEach(
-                i -> {
-                    GenericRow record1 = GenericRow.of(i);
-                    try {
-                        write.write(record1);
-                    } catch (Exception e) {
-                        throw new RuntimeException(e);
-                    }
-                });
+        for (Integer i : data) {
+            GenericRow record = GenericRow.of(i);
+            write.write(record);
+        }
         List<CommitMessage> messages = write.prepareCommit();
         BatchTableCommit commit = writeBuilder.newCommit();
         commit.commit(messages);

paimon-core/src/test/java/org/apache/paimon/rest/RESTTestFileIO.java

@@ -54,60 +54,71 @@ public void configure(CatalogContext context) {
 
     @Override
     public SeekableInputStream newInputStream(Path path) throws IOException {
-        boolean isDataTokenEnabled = options.getOptional(DATA_TOKEN_ENABLED).orElse(false);
-        if (isDataTokenEnabled) {
-            RESTToken token = getToken(path);
-            if (token == null) {
-                throw new IOException(TOKEN_UN_EXIST_MSG);
-            } else if (token.expireAtMillis() < System.currentTimeMillis()) {
-                throw new IOException(TOKEN_EXPIRED_MSG);
-            }
-        }
+        checkDataToken(path);
         return super.newInputStream(path);
     }
 
-    private RESTToken getToken(Path path) {
-        String basePath = options.get("warehouse").replaceAll("rest-test-file-io://", "");
-        String filePath = path.toString().split(":")[1].replaceAll(basePath, "");
-        String[] paths = filePath.split("/");
-        String database = paths[1].replaceAll("\\.db", "");
-        String table = paths[2];
-        return DataTokenStore.getDataToken(
-                options.get("catalog-server-id"), Identifier.create(database, table).getFullName());
-    }
-
     @Override
     public PositionOutputStream newOutputStream(Path path, boolean overwrite) throws IOException {
+        checkDataToken(path);
         return super.newOutputStream(path, overwrite);
     }
 
     @Override
     public FileStatus getFileStatus(Path path) throws IOException {
+        checkDataToken(path);
         return super.getFileStatus(path);
     }
 
     @Override
     public FileStatus[] listStatus(Path path) throws IOException {
+        checkDataToken(path);
         return super.listStatus(path);
     }
 
     @Override
     public boolean exists(Path path) throws IOException {
+        checkDataToken(path);
         return super.exists(path);
     }
 
     @Override
     public boolean delete(Path path, boolean recursive) throws IOException {
+        checkDataToken(path);
         return super.delete(path, recursive);
     }
 
     @Override
     public boolean mkdirs(Path path) throws IOException {
+        checkDataToken(path);
         return super.mkdirs(path);
     }
 
     @Override
     public boolean rename(Path src, Path dst) throws IOException {
+        checkDataToken(src);
         return super.rename(src, dst);
     }
+
+    private void checkDataToken(Path path) throws IOException {
+        boolean isDataTokenEnabled = options.getOptional(DATA_TOKEN_ENABLED).orElse(false);
+        if (isDataTokenEnabled) {
+            RESTToken token = getToken(path);
+            if (token == null) {
+                throw new IOException(TOKEN_UN_EXIST_MSG);
+            } else if (token.expireAtMillis() < System.currentTimeMillis()) {
+                throw new IOException(TOKEN_EXPIRED_MSG);
+            }
+        }
+    }
+
+    private RESTToken getToken(Path path) {
+        String basePath = options.get("warehouse").replaceAll("rest-test-file-io://", "");
+        String filePath = path.toString().split(":")[1].replaceAll(basePath, "");
+        String[] paths = filePath.split("/");
+        String database = paths[1].replaceAll("\\.db", "");
+        String table = paths[2];
+        return DataTokenStore.getDataToken(
+                options.get("catalog-server-id"), Identifier.create(database, table).getFullName());
+    }
 }