diff --git a/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py b/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py index 293171285181..861e3d3bb15c 100644 --- a/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py +++ b/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py @@ -326,5 +326,10 @@ def valid_token(self): self.try_to_refresh_token() return self.token + def _existing_valid_token(self) -> Optional[RESTToken]: + """Return this instance's valid token without locks or refresh requests.""" + token = self.token + return token if not self._is_token_expired(token) else None + def close(self): pass diff --git a/paimon-python/pypaimon/read/native_plan.py b/paimon-python/pypaimon/read/native_plan.py index fbad1d8c3760..581acb339734 100644 --- a/paimon-python/pypaimon/read/native_plan.py +++ b/paimon-python/pypaimon/read/native_plan.py @@ -309,7 +309,7 @@ def __getstate__(self): def __setstate__(self, state): self.__init__() - def get(self, response, database, table, options): + def get(self, response, database, table, options, token_loader=None): from pypaimon_rust.datafusion import Table pid = os.getpid() @@ -322,12 +322,36 @@ def get(self, response, database, table, options): key = (response, database, table, tuple(sorted(options.items()))) with state.lock: if state.entry is None or state.entry[0] != key: - native_table = Table.from_rest_response( - response, database=database, table=table, rest_options=options) + reuse_token = getattr(Table, 'from_rest_response_with_token', None) + token = token_loader() if callable(reuse_token) and token_loader else None + if token is None: + native_table = Table.from_rest_response( + response, database=database, table=table, rest_options=options) + else: + native_table = reuse_token( + response, database=database, table=table, rest_options=options, + data_token=dict(token.token), + expires_at_millis=token.expire_at_millis) state.entry = (key, native_table) return state.entry[1] +def _rest_data_token(table): + """Reuse only an existing valid token bound to this exact table.""" + from pypaimon.catalog.rest.rest_token_file_io import RESTTokenFileIO + from pypaimon.filesystem.caching_file_io import CachingFileIO + + file_io = table.file_io + if type(file_io) is CachingFileIO: + file_io = file_io._delegate + if (table.current_branch() != 'main' + or type(file_io) is not RESTTokenFileIO + or file_io.identifier != table.identifier + or file_io.path != table.table_path): + return None + return file_io._existing_valid_token() + + def _native_read_builder(table): """Return a fresh builder with the current schema and shared REST FileIO.""" rest_response = _resolved_rest_table_response(table) @@ -339,7 +363,8 @@ def _native_read_builder(table): rest_response, database=table.identifier.get_database_name(), table=table.identifier.get_object_name(), - options=_catalog_options(table)) + options=_catalog_options(table), + token_loader=lambda: _rest_data_token(table)) rt = rt.copy_with_resolved_schema(_resolved_schema_json(table), branch=table.current_branch()) elif file_io_options is not None: from pypaimon_rust.datafusion import Table diff --git a/paimon-python/pypaimon/tests/native_plan_rest_test.py b/paimon-python/pypaimon/tests/native_plan_rest_test.py index 320bfd61bc00..ef3294c2f3a8 100644 --- a/paimon-python/pypaimon/tests/native_plan_rest_test.py +++ b/paimon-python/pypaimon/tests/native_plan_rest_test.py @@ -149,6 +149,64 @@ def test_resolved_rest_table_keeps_refreshable_file_io(rest_source, rest_catalog refresh.assert_called() +@pytest.mark.skipif(not native_method_available('Table', 'from_rest_response_with_token'), + reason='REST data token reuse binding required') +@pytest.mark.parametrize('local_cache', [False, True]) +def test_native_rest_reuses_python_data_token( + rest_source, rest_catalog, tmp_path, local_cache): + import time + + from pypaimon.catalog.rest.rest_token import RESTToken + from pypaimon.catalog.rest.rest_token_file_io import RESTTokenFileIO + from pypaimon.filesystem.caching_file_io import CachingFileIO + from pypaimon.read.native_plan import _catalog_options, _resolved_rest_table_response + from pypaimon_rust.datafusion import Table as NativeTable + + source, server, _ = rest_source + catalog, _ = rest_catalog + options = dict(catalog.context.options.to_map()) + options['data-token.enabled'] = 'true' + if local_cache: + options.update({'local-cache.enabled': 'true', + 'local-cache.dir': str(tmp_path / 'cache')}) + server.set_table_token( + source.identifier, RESTToken({}, int(time.time() * 1000) + 7_200_000)) + table = CatalogFactory.create(options).get_table(source.identifier) + reused_table = CatalogFactory.create(options).get_table(source.identifier) + uncached_table = CatalogFactory.create(options).get_table(source.identifier) + + def token_file_io(table): + file_io = table.file_io + assert (type(file_io) is CachingFileIO) == local_cache + if type(file_io) is CachingFileIO: + file_io = file_io._delegate + assert type(file_io) is RESTTokenFileIO + return file_io + + # The old bridge fetched a Python token and an independent Rust token. + with patch.object(RESTTokenFileIO, '_TOKEN_CACHE', {}): + with patch.object(server, '_table_token_handle', + wraps=server._table_token_handle) as load: + token_file_io(table).valid_token() + baseline = NativeTable.from_rest_response( + _resolved_rest_table_response(table), database='default', table='t', + rest_options=_catalog_options(table)) + assert baseline.new_read_builder().new_scan().plan().snapshot_id() == 2 + assert load.call_count == 2 + load.reset_mock() + RESTTokenFileIO._TOKEN_CACHE.clear() + + token_file_io(reused_table).token = token_file_io(table).token + + for _ in range(2): + assert reused_table.new_read_builder().new_scan().plan().snapshot_id == 2 + assert load.call_count == 0 + + # No instance token: Rust obtains one instead of making Python refresh. + assert uncached_table.new_read_builder().new_scan().plan().snapshot_id == 2 + assert load.call_count == 1 + + @pytest.mark.parametrize('branch', [None, 'dev']) def test_rest_dotted_database_and_table_keep_identity(rest_catalog, branch): from pypaimon.common.identifier import Identifier diff --git a/paimon-python/pypaimon/tests/native_plan_test.py b/paimon-python/pypaimon/tests/native_plan_test.py index 7f5bd3386f1e..095110cd581a 100644 --- a/paimon-python/pypaimon/tests/native_plan_test.py +++ b/paimon-python/pypaimon/tests/native_plan_test.py @@ -1115,6 +1115,104 @@ def access(): release.set() holder.join(5) + def test_native_rest_token_lookup_after_fork_avoids_inherited_lock(self): + import multiprocessing + from threading import Event, Thread + + from pypaimon.catalog.rest.rest_token_file_io import RESTTokenFileIO + from pypaimon.common.identifier import Identifier + from pypaimon.read.native_plan import _NativeRestTableCache, _rest_data_token + + if 'fork' not in multiprocessing.get_all_start_methods(): + self.skipTest('fork required') + context = multiprocessing.get_context('fork') + identifier = Identifier('default', 't') + path = 'oss://bucket/table' + file_io = RESTTokenFileIO(identifier, path) + table = SimpleNamespace(file_io=file_io, identifier=identifier, + table_path=path, current_branch=lambda: 'main') + fake_df = ModuleType('pypaimon_rust.datafusion') + fake_df.Table = Mock() + fake_df.Table.from_rest_response.return_value = object() + cache = _NativeRestTableCache() + held, release = Event(), Event() + + def hold_token_lock(): + with RESTTokenFileIO._TOKEN_LOCKS_LOCK: + held.set() + release.wait() + + def child(connection): + try: + native_table = cache.get( + 'response', 'default', 't', {}, + token_loader=lambda: _rest_data_token(table)) + connection.send(('ok', native_table is not None)) + except BaseException as exc: + connection.send(('error', repr(exc))) + finally: + connection.close() + + with patch.dict(sys.modules, {'pypaimon_rust.datafusion': fake_df}), \ + patch.object(file_io, '_build_cache_key', return_value='token-key'): + holder = Thread(target=hold_token_lock, daemon=True) + holder.start() + self.assertTrue(held.wait(5)) + receiving, sending = context.Pipe(duplex=False) + process = context.Process(target=child, args=(sending,)) + try: + process.start() + sending.close() + self.assertTrue(receiving.poll(3), 'child deadlocked on inherited token lock') + self.assertEqual(receiving.recv(), ('ok', True)) + process.join(5) + self.assertEqual(process.exitcode, 0) + finally: + if process.is_alive(): + process.terminate() + process.join(5) + receiving.close() + sending.close() + release.set() + holder.join(5) + + def test_native_rest_cache_passes_only_existing_token(self): + import time + + from pypaimon.catalog.rest.rest_token import RESTToken + from pypaimon.catalog.rest.rest_token_file_io import RESTTokenFileIO + from pypaimon.common.identifier import Identifier + from pypaimon.read.native_plan import _NativeRestTableCache, _rest_data_token + + identifier = Identifier('default', 't') + path = 'oss://bucket/table' + file_io = RESTTokenFileIO(identifier, path) + token = RESTToken({'key': 'value'}, int(time.time() * 1000) + 7_200_000) + file_io.token = token + table = SimpleNamespace(file_io=file_io, identifier=identifier, + table_path=path, current_branch=lambda: 'main') + fake_df = ModuleType('pypaimon_rust.datafusion') + fake_df.Table = Mock() + fake_df.Table.from_rest_response_with_token.return_value = object() + + with patch.dict(sys.modules, {'pypaimon_rust.datafusion': fake_df}), \ + patch.object(file_io, 'valid_token', + side_effect=AssertionError('must not refresh')): + cache = _NativeRestTableCache() + cache.get('response', 'default', 't', {}, + token_loader=lambda: _rest_data_token(table)) + fake_df.Table.from_rest_response_with_token.assert_called_once_with( + 'response', database='default', table='t', rest_options={}, + data_token=token.token, expires_at_millis=token.expire_at_millis) + fake_df.Table.from_rest_response.assert_not_called() + + file_io.token = None + cache = _NativeRestTableCache() + cache.get('response', 'default', 't', {}, + token_loader=lambda: _rest_data_token(table)) + fake_df.Table.from_rest_response.assert_called_once_with( + 'response', database='default', table='t', rest_options={}) + def test_native_rest_cache_retries_failed_construction(self): from pypaimon.read.native_plan import _NativeRestTableCache diff --git a/paimon-python/pypaimon/tests/rest/rest_token_file_io_test.py b/paimon-python/pypaimon/tests/rest/rest_token_file_io_test.py index 4e94b21d728e..8255a78ea977 100644 --- a/paimon-python/pypaimon/tests/rest/rest_token_file_io_test.py +++ b/paimon-python/pypaimon/tests/rest/rest_token_file_io_test.py @@ -49,6 +49,22 @@ def tearDown(self): if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) + def test_existing_valid_token_does_not_refresh_or_read_shared_cache(self): + file_io = RESTTokenFileIO(self.identifier, self.warehouse_path) + now = 1700000000 + valid = RESTToken({'key': 'value'}, (now + 7200) * 1000) + expiring = RESTToken({'key': 'value'}, (now + 1800) * 1000) + with patch('pypaimon.catalog.rest.rest_token_file_io.time.time', return_value=now), \ + patch.object(file_io, 'try_to_refresh_token') as refresh, \ + patch.object(file_io, '_get_cached_token') as shared_cache: + self.assertIsNone(file_io._existing_valid_token()) + file_io.token = valid + self.assertIs(file_io._existing_valid_token(), valid) + file_io.token = expiring + self.assertIsNone(file_io._existing_valid_token()) + refresh.assert_not_called() + shared_cache.assert_not_called() + def test_blob_presigned_url_bound_table_root(self): root = "oss://bucket/table-a" file_io = RESTTokenFileIO(self.identifier, root, self.catalog_options)