rest token refresh

add global token cache back

update _is_token_expired to instance method

remove global token cache

clean code

clean code

fix

fix

fix code format

fix delete not handle file not found issue

fix dead lock

add global cache back

fix review comments

fix init read

Author: xiaohongbo

paimon-python/pypaimon/api/rest_util.py

@@ -15,8 +15,7 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-from typing import Dict
-from urllib.parse import unquote
+from typing import Dict, Optional
 
 from pypaimon.common.options import Options
 
@@ -31,6 +30,7 @@ def encode_string(value: str) -> str:
     @staticmethod
     def decode_string(encoded: str) -> str:
         """Decode URL-encoded string"""
+        from urllib.parse import unquote
         return unquote(encoded)
 
     @staticmethod
@@ -43,3 +43,24 @@ def extract_prefix_map(
                 new_key = key[len(prefix):]
                 result[new_key] = str(value)
         return result
+
+    @staticmethod
+    def merge(
+            base_properties: Optional[Dict[str, str]],
+            override_properties: Optional[Dict[str, str]]) -> Dict[str, str]:
+        if override_properties is None:
+            override_properties = {}
+        if base_properties is None:
+            base_properties = {}
+
+        result = {}
+
+        for key, value in base_properties.items():
+            if value is not None and key not in override_properties:
+                result[key] = value
+
+        for key, value in override_properties.items():
+            if value is not None:
+                result[key] = value
+
+        return result

paimon-python/pypaimon/catalog/rest/rest_token_file_io.py

@@ -18,27 +18,38 @@
 import logging
 import threading
 import time
+from pathlib import Path
 from typing import Optional
 
+import pyarrow
+import pyarrow.fs
+from pypaimon.common.options.config import CatalogOptions, OssOptions
 from pyarrow._fs import FileSystem
 
 from pypaimon.api.rest_api import RESTApi
+from pypaimon.common.options import Options
+from pypaimon.api.rest_util import RESTUtil
 from pypaimon.catalog.rest.rest_token import RESTToken
 from pypaimon.common.file_io import FileIO
 from pypaimon.common.identifier import Identifier
-from pypaimon.common.options import Options
-from pypaimon.common.options.config import CatalogOptions, OssOptions
+
+from cachetools import TTLCache
 
 
 class RESTTokenFileIO(FileIO):
+    _FILESYSTEM_CACHE: TTLCache = TTLCache(maxsize=1000, ttl=36000)  # 10 hours TTL
+    _CACHE_LOCK = threading.Lock()
+    _TOKEN_CACHE: dict = {}
+    _TOKEN_LOCKS: dict = {}
+    _TOKEN_LOCKS_LOCK = threading.Lock()
 
     def __init__(self, identifier: Identifier, path: str,
                  catalog_options: Optional[Options] = None):
         self.identifier = identifier
         self.path = path
         self.token: Optional[RESTToken] = None
         self.api_instance: Optional[RESTApi] = None
-        self.lock = threading.Lock()
+        self.lock = threading.RLock()
         self.log = logging.getLogger(__name__)
         super().__init__(path, catalog_options)
 
@@ -53,33 +64,241 @@ def __getstate__(self):
     def __setstate__(self, state):
         self.__dict__.update(state)
         # Recreate lock after deserialization
-        self.lock = threading.Lock()
+        self.lock = threading.RLock()
         # api_instance will be recreated when needed
         self.api_instance = None
 
-    def _initialize_oss_fs(self, path) -> FileSystem:
+    def _initialize_oss_fs(self, path, properties: Optional[Options] = None) -> FileSystem:
+        if properties is not None:
+            with self.lock:
+                original_properties = self.properties
+                self.properties = properties
+                try:
+                    return super()._initialize_oss_fs(path)
+                finally:
+                    self.properties = original_properties
+        
         self.try_to_refresh_token()
-        merged_token = self._merge_token_with_catalog_options(self.token.token)
-        self.properties.data.update(merged_token)
-        return super()._initialize_oss_fs(path)
+        if self.token is None:
+            return super()._initialize_oss_fs(path)
+        
+        merged_properties = RESTUtil.merge(
+            self.properties.to_map() if self.properties else {},
+            self._merge_token_with_catalog_options(self.token.token)
+        )
+        merged_options = Options(merged_properties)
+        
+        with self.lock:
+            original_properties = self.properties
+            self.properties = merged_options
+            try:
+                return super()._initialize_oss_fs(path)
+            finally:
+                self.properties = original_properties
+
+    def _initialize_s3_fs(self, properties: Optional[Options] = None) -> FileSystem:
+        if properties is not None:
+            with self.lock:
+                original_properties = self.properties
+                self.properties = properties
+                try:
+                    return super()._initialize_s3_fs()
+                finally:
+                    self.properties = original_properties
+        
+        self.try_to_refresh_token()
+        if self.token is None:
+            return super()._initialize_s3_fs()
+        
+        merged_properties = RESTUtil.merge(
+            self.properties.to_map() if self.properties else {},
+            self._merge_token_with_catalog_options(self.token.token)
+        )
+        merged_options = Options(merged_properties)
+        
+        with self.lock:
+            original_properties = self.properties
+            self.properties = merged_options
+            try:
+                return super()._initialize_s3_fs()
+            finally:
+                self.properties = original_properties
 
     def _merge_token_with_catalog_options(self, token: dict) -> dict:
         """Merge token with catalog options, DLF OSS endpoint should override the standard OSS endpoint."""
         merged_token = dict(token)
-        dlf_oss_endpoint = self.properties.get(CatalogOptions.DLF_OSS_ENDPOINT)
+        with self.lock:
+            dlf_oss_endpoint = self.properties.get(CatalogOptions.DLF_OSS_ENDPOINT)
         if dlf_oss_endpoint and dlf_oss_endpoint.strip():
             merged_token[OssOptions.OSS_ENDPOINT.key()] = dlf_oss_endpoint
         return merged_token
 
+    def get_merged_properties(self) -> Options:
+        self.try_to_refresh_token()
+        if self.token is None:
+            with self.lock:
+                return self.properties
+
+        with self.lock:
+            properties_map = self.properties.to_map() if self.properties else {}
+        
+        merged_properties = RESTUtil.merge(
+            properties_map,
+            self.token.token
+        )
+        return Options(merged_properties)
+
+    def _get_filesystem(self) -> FileSystem:
+        self.try_to_refresh_token()
+
+        if self.token is None:
+            return self.filesystem
+
+        with self._CACHE_LOCK:
+            filesystem = self._FILESYSTEM_CACHE.get(self.token)
+            if filesystem is not None:
+                return filesystem
+
+            with self.lock:
+                properties_map = self.properties.to_map() if self.properties else {}
+            
+            merged_properties = RESTUtil.merge(
+                properties_map,
+                self.token.token
+            )
+            merged_options = Options(merged_properties)
+
+            scheme, netloc, _ = self.parse_location(self.path)
+            if scheme in {"oss"}:
+                filesystem = self._initialize_oss_fs(self.path, merged_options)
+            elif scheme in {"s3", "s3a", "s3n"}:
+                filesystem = self._initialize_s3_fs(merged_options)
+            else:
+                filesystem = self.filesystem
+
+            self._FILESYSTEM_CACHE[self.token] = filesystem
+            return filesystem
+
+    def new_input_stream(self, path: str):
+        filesystem = self._get_filesystem()
+        path_str = self.to_filesystem_path(path)
+        return filesystem.open_input_file(path_str)
+
     def new_output_stream(self, path: str):
-        # Call parent class method to ensure path conversion and parent directory creation
-        return super().new_output_stream(path)
+        filesystem = self._get_filesystem()
+        path_str = self.to_filesystem_path(path)
+        parent_dir = Path(path_str).parent
+        if str(parent_dir) and not self.exists(str(parent_dir)):
+            self.mkdirs(str(parent_dir))
+        return filesystem.open_output_stream(path_str)
+
+    def get_file_status(self, path: str):
+        filesystem = self._get_filesystem()
+        path_str = self.to_filesystem_path(path)
+        file_infos = filesystem.get_file_info([path_str])
+        return file_infos[0]
+
+    def list_status(self, path: str):
+        filesystem = self._get_filesystem()
+        path_str = self.to_filesystem_path(path)
+        selector = pyarrow.fs.FileSelector(path_str, recursive=False, allow_not_found=True)
+        return filesystem.get_file_info(selector)
+
+    def exists(self, path: str) -> bool:
+        filesystem = self._get_filesystem()
+        path_str = self.to_filesystem_path(path)
+        file_info = filesystem.get_file_info([path_str])[0]
+        return file_info.type != pyarrow.fs.FileType.NotFound
+
+    def delete(self, path: str, recursive: bool = False) -> bool:
+        filesystem = self._get_filesystem()
+        path_str = self.to_filesystem_path(path)
+        file_info = filesystem.get_file_info([path_str])[0]
+        if file_info.type == pyarrow.fs.FileType.NotFound:
+            return False
+        if file_info.type == pyarrow.fs.FileType.Directory:
+            if recursive:
+                filesystem.delete_dir_contents(path_str)
+            else:
+                filesystem.delete_dir(path_str)
+        else:
+            filesystem.delete_file(path_str)
+        return True
+
+    def mkdirs(self, path: str) -> bool:
+        filesystem = self._get_filesystem()
+        path_str = self.to_filesystem_path(path)
+        filesystem.create_dir(path_str, recursive=True)
+        return True
+
+    def rename(self, src: str, dst: str) -> bool:
+        filesystem = self._get_filesystem()
+        dst_str = self.to_filesystem_path(dst)
+        dst_parent = Path(dst_str).parent
+        if str(dst_parent) and not self.exists(str(dst_parent)):
+            self.mkdirs(str(dst_parent))
+
+        src_str = self.to_filesystem_path(src)
+        filesystem.move(src_str, dst_str)
+        return True
+
+    def copy_file(self, source_path: str, target_path: str, overwrite: bool = False):
+        if not overwrite and self.exists(target_path):
+            raise FileExistsError(f"Target file {target_path} already exists and overwrite=False")
+
+        filesystem = self._get_filesystem()
+        source_str = self.to_filesystem_path(source_path)
+        target_str = self.to_filesystem_path(target_path)
+        filesystem.copy_file(source_str, target_str)
 
     def try_to_refresh_token(self):
-        if self.should_refresh():
-            with self.lock:
-                if self.should_refresh():
-                    self.refresh_token()
+        identifier_str = str(self.identifier)
+        
+        # Fast path 1: Check instance token
+        if self.token is not None and not self._is_token_expired(self.token):
+            return
+        
+        # Fast path 2: Check global cache
+        cached_token = self._get_cached_token(identifier_str)
+        if cached_token and not self._is_token_expired(cached_token):
+            self.token = cached_token
+            return
+        
+        # Slow path: Acquire global lock for this identifier
+        global_lock = self._get_global_token_lock()
+        with global_lock:
+            cached_token = self._get_cached_token(identifier_str)
+            if cached_token and not self._is_token_expired(cached_token):
+                self.token = cached_token
+                return
+            
+            token_to_check = cached_token if cached_token else self.token
+            if token_to_check is None or self._is_token_expired(token_to_check):
+                self.refresh_token()
+                self._set_cached_token(identifier_str, self.token)
+            else:
+                self.token = cached_token if cached_token else self.token
+
+    def _get_cached_token(self, identifier_str: str) -> Optional[RESTToken]:
+        with self._TOKEN_LOCKS_LOCK:
+            return self._TOKEN_CACHE.get(identifier_str)
+    
+    def _set_cached_token(self, identifier_str: str, token: RESTToken):
+        with self._TOKEN_LOCKS_LOCK:
+            self._TOKEN_CACHE[identifier_str] = token
+    
+    def _get_global_token_lock(self) -> threading.Lock:
+        identifier_str = str(self.identifier)
+        with self._TOKEN_LOCKS_LOCK:
+            if identifier_str not in self._TOKEN_LOCKS:
+                self._TOKEN_LOCKS[identifier_str] = threading.Lock()
+            return self._TOKEN_LOCKS[identifier_str]
+    
+    def _is_token_expired(self, token: RESTToken) -> bool:
+        if token is None:
+            return True
+        current_time = int(time.time() * 1000)
+        return (token.expire_at_millis - current_time) < RESTApi.TOKEN_EXPIRATION_SAFE_TIME_MILLIS
 
     def should_refresh(self):
         if self.token is None:
@@ -90,13 +309,50 @@ def should_refresh(self):
     def refresh_token(self):
         self.log.info(f"begin refresh data token for identifier [{self.identifier}]")
         if self.api_instance is None:
-            self.api_instance = RESTApi(self.properties, False)
+            if not self.properties:
+                raise RuntimeError(
+                    "Cannot refresh token: properties is None or empty. "
+                    "This may indicate a serialization issue when passing RESTTokenFileIO to Ray workers."
+                )
+
+            if not self.properties.contains(CatalogOptions.URI):
+                available_keys = list(self.properties.data.keys()) if self.properties else []
+                raise RuntimeError(
+                    f"Cannot refresh token: missing required configuration '{CatalogOptions.URI}' in properties. "
+                    "This is required to create RESTApi for token refresh. "
+                    f"Available configuration keys: {available_keys}. "
+                    "This may indicate that catalog options were not properly serialized when passing to Ray workers."
+                )
+
+            uri = self.properties.get(CatalogOptions.URI)
+            if not uri or not uri.strip():
+                raise RuntimeError(
+                    f"Cannot refresh token: '{CatalogOptions.URI}' is empty or whitespace. "
+                    f"Value: '{uri}'. Please ensure the REST catalog URI is properly configured."
+                )
+
+            try:
+                self.api_instance = RESTApi(self.properties, False)
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to create RESTApi for token refresh: {e}. "
+                    "Please check that all required catalog options are properly configured. "
+                    f"Identifier: {self.identifier}"
+                ) from e
+
+        try:
+            response = self.api_instance.load_table_token(self.identifier)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to load table token from REST API: {e}. "
+                f"Identifier: {self.identifier}, URI: {self.properties.get(CatalogOptions.URI)}"
+            ) from e
 
-        response = self.api_instance.load_table_token(self.identifier)
         self.log.info(
             f"end refresh data token for identifier [{self.identifier}] expiresAtMillis [{response.expires_at_millis}]"
         )
-        self.token = RESTToken(response.token, response.expires_at_millis)
+        merged_token = self._merge_token_with_catalog_options(response.token)
+        self.token = RESTToken(merged_token, response.expires_at_millis)
 
     def valid_token(self):
         self.try_to_refresh_token()