From 7803d562546ab4d56d57ef61e2ebbb94f1a767d1 Mon Sep 17 00:00:00 2001 From: Hyu Date: Mon, 17 Aug 2026 11:49:26 +0800 Subject: [PATCH] fix(plugin): keep pre-tenancy plugin storage readable (#2446) * fix(plugin): adopt legacy scoped storage rows * fix(plugin): delete adopted legacy storage rows * test(plugin): cover legacy storage deletion * fix(plugin): avoid mutating legacy storage on reads * fix(plugin): make legacy storage adoption atomic * fix(plugin): resolve concurrent legacy adoption * fix(plugin): upsert concurrent storage writes * fix(plugin): retry reads after legacy adoption * fix(plugin): deduplicate migrated storage keys --------- Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com> --- src/langbot/pkg/plugin/handler.py | 130 ++++++++++++++-- .../unit_tests/plugin/test_handler_actions.py | 142 +++++++++++++++++- 2 files changed, 253 insertions(+), 19 deletions(-) diff --git a/src/langbot/pkg/plugin/handler.py b/src/langbot/pkg/plugin/handler.py index 4805f2668..8dae74f6f 100644 --- a/src/langbot/pkg/plugin/handler.py +++ b/src/langbot/pkg/plugin/handler.py @@ -11,6 +11,8 @@ import traceback from dataclasses import dataclass import sqlalchemy +import sqlalchemy.dialects.postgresql +import sqlalchemy.dialects.sqlite from langbot_plugin.runtime.io import handler from langbot_plugin.runtime.io.connection import Connection @@ -431,6 +433,19 @@ class RuntimeConnectionHandler(handler.Handler): return f'{identity.plugin_author}/{identity.plugin_name}' raise ValueError(f'Unsupported binary storage owner_type {owner_type!r}') + @staticmethod + def _legacy_binary_storage_key( + action_context: ActionContext, + *, + owner_type: str, + owner: str, + key: str, + ) -> str: + """Return the pre-tenancy key shape for a row already scoped to this Workspace.""" + + legacy_owner = action_context.workspace_uuid if owner_type == 'workspace' else owner + return f'{owner_type}:{legacy_owner}:{key}' + @classmethod def _binary_storage_key( cls, @@ -896,25 +911,82 @@ class RuntimeConnectionHandler(handler.Handler): .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid) .where(persistence_bstorage.BinaryStorage.unique_key == unique_key) ) + storage = result.first() + if storage is None: + legacy_key = self._legacy_binary_storage_key( + action_context, + owner_type=owner_type, + owner=owner, + key=key, + ) + result = await self.ap.persistence_mgr.execute_async( + sqlalchemy.select(persistence_bstorage.BinaryStorage) + .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid) + .where(persistence_bstorage.BinaryStorage.unique_key == legacy_key) + .where(persistence_bstorage.BinaryStorage.key == key) + .where(persistence_bstorage.BinaryStorage.owner_type == owner_type) + .where(persistence_bstorage.BinaryStorage.owner == owner) + ) + storage = result.first() + if storage is not None: + update_result = await self.ap.persistence_mgr.execute_async( + sqlalchemy.update(persistence_bstorage.BinaryStorage) + .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid) + .where(persistence_bstorage.BinaryStorage.unique_key == legacy_key) + .where(persistence_bstorage.BinaryStorage.key == key) + .where(persistence_bstorage.BinaryStorage.owner_type == owner_type) + .where(persistence_bstorage.BinaryStorage.owner == owner) + .values(unique_key=unique_key, value=value) + ) + if update_result.rowcount: + return handler.ActionResponse.success(data={}) + canonical_update = await self.ap.persistence_mgr.execute_async( + sqlalchemy.update(persistence_bstorage.BinaryStorage) + .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid) + .where(persistence_bstorage.BinaryStorage.unique_key == unique_key) + .where(persistence_bstorage.BinaryStorage.key == key) + .where(persistence_bstorage.BinaryStorage.owner_type == owner_type) + .where(persistence_bstorage.BinaryStorage.owner == owner) + .values(value=value) + ) + if canonical_update.rowcount: + return handler.ActionResponse.success(data={}) + storage = None - if result.first() is not None: + if storage is not None: await self.ap.persistence_mgr.execute_async( sqlalchemy.update(persistence_bstorage.BinaryStorage) .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid) .where(persistence_bstorage.BinaryStorage.unique_key == unique_key) + .where(persistence_bstorage.BinaryStorage.key == key) + .where(persistence_bstorage.BinaryStorage.owner_type == owner_type) + .where(persistence_bstorage.BinaryStorage.owner == owner) .values(value=value) ) - else: - await self.ap.persistence_mgr.execute_async( - sqlalchemy.insert(persistence_bstorage.BinaryStorage).values( - workspace_uuid=action_context.workspace_uuid, - unique_key=unique_key, - key=key, - owner_type=owner_type, - owner=owner, - value=value, - ) + return handler.ActionResponse.success(data={}) + + dialect_name = self.ap.persistence_mgr.get_db_engine().dialect.name + insert = { + 'postgresql': sqlalchemy.dialects.postgresql.insert, + 'sqlite': sqlalchemy.dialects.sqlite.insert, + }.get(dialect_name) + if insert is None: + return handler.ActionResponse.error(message=f'Unsupported storage database dialect: {dialect_name}') + await self.ap.persistence_mgr.execute_async( + insert(persistence_bstorage.BinaryStorage) + .values( + workspace_uuid=action_context.workspace_uuid, + unique_key=unique_key, + key=key, + owner_type=owner_type, + owner=owner, + value=value, ) + .on_conflict_do_update( + index_elements=['workspace_uuid', 'unique_key'], + set_={'value': value}, + ) + ) return handler.ActionResponse.success( data={}, @@ -946,6 +1018,29 @@ class RuntimeConnectionHandler(handler.Handler): ) storage = result.first() + if storage is None: + legacy_key = self._legacy_binary_storage_key( + action_context, + owner_type=owner_type, + owner=owner, + key=key, + ) + result = await self.ap.persistence_mgr.execute_async( + sqlalchemy.select(persistence_bstorage.BinaryStorage) + .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid) + .where(persistence_bstorage.BinaryStorage.unique_key == legacy_key) + .where(persistence_bstorage.BinaryStorage.key == key) + .where(persistence_bstorage.BinaryStorage.owner_type == owner_type) + .where(persistence_bstorage.BinaryStorage.owner == owner) + ) + storage = result.first() + if storage is None: + retry_result = await self.ap.persistence_mgr.execute_async( + sqlalchemy.select(persistence_bstorage.BinaryStorage) + .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid) + .where(persistence_bstorage.BinaryStorage.unique_key == unique_key) + ) + storage = retry_result.first() if storage is None: return handler.ActionResponse.error( message=f'Storage with key {key} not found', @@ -981,10 +1076,19 @@ class RuntimeConnectionHandler(handler.Handler): message=str(e), ) + legacy_key = self._legacy_binary_storage_key( + action_context, + owner_type=owner_type, + owner=owner, + key=key, + ) await self.ap.persistence_mgr.execute_async( sqlalchemy.delete(persistence_bstorage.BinaryStorage) .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid) - .where(persistence_bstorage.BinaryStorage.unique_key == unique_key) + .where(persistence_bstorage.BinaryStorage.unique_key.in_((unique_key, legacy_key))) + .where(persistence_bstorage.BinaryStorage.key == key) + .where(persistence_bstorage.BinaryStorage.owner_type == owner_type) + .where(persistence_bstorage.BinaryStorage.owner == owner) ) return handler.ActionResponse.success( @@ -1012,7 +1116,7 @@ class RuntimeConnectionHandler(handler.Handler): return handler.ActionResponse.success( data={ - 'keys': result.scalars().all(), + 'keys': list(dict.fromkeys(result.scalars().all())), }, ) diff --git a/tests/unit_tests/plugin/test_handler_actions.py b/tests/unit_tests/plugin/test_handler_actions.py index 4089e23be..aa0762471 100644 --- a/tests/unit_tests/plugin/test_handler_actions.py +++ b/tests/unit_tests/plugin/test_handler_actions.py @@ -234,6 +234,7 @@ class TestSetBinaryStorage: }, } mock_app.persistence_mgr = Mock() + mock_app.persistence_mgr.get_db_engine.return_value = SimpleNamespace(dialect=SimpleNamespace(name='sqlite')) mock_app.persistence_mgr.execute_async = AsyncMock(return_value=make_result()) mock_app.logger = Mock() return mock_app @@ -270,8 +271,8 @@ class TestSetBinaryStorage: ) assert response.code == 0 - assert app.persistence_mgr.execute_async.await_count == 2 - insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[1].args[0]) + assert app.persistence_mgr.execute_async.await_count == 3 + insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[2].args[0]) assert insert_params['workspace_uuid'] == 'workspace-a' assert insert_params['unique_key'] == canonical_binary_key( 'plugin', @@ -301,6 +302,69 @@ class TestSetBinaryStorage: assert expected_key in update_params.values() assert update_params['value'] == b'new' + @pytest.mark.asyncio + async def test_adopts_legacy_storage_before_updating(self, app): + """A migrated pre-tenancy row is updated in place rather than duplicated.""" + runtime_handler = make_handler(app) + legacy_storage = SimpleNamespace(unique_key='plugin:test-author/test-plugin:test-key') + adopted = SimpleNamespace(rowcount=1) + app.persistence_mgr.execute_async.side_effect = [ + make_result(), + make_result(legacy_storage), + adopted, + ] + + response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value](self.payload(b'new')) + + assert response.code == 0 + assert app.persistence_mgr.execute_async.await_count == 3 + adoption_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[2].args[0]) + expected_key = canonical_binary_key('plugin', 'test-author/test-plugin', 'test-key') + assert expected_key in adoption_params.values() + assert adoption_params['value'] == b'new' + + @pytest.mark.asyncio + async def test_legacy_adoption_race_updates_winning_canonical_row(self, app): + runtime_handler = make_handler(app) + legacy_storage = SimpleNamespace(unique_key='plugin:test-author/test-plugin:test-key') + lost_race = SimpleNamespace(rowcount=0) + canonical_winner = SimpleNamespace(rowcount=1) + app.persistence_mgr.execute_async.side_effect = [ + make_result(), + make_result(legacy_storage), + lost_race, + canonical_winner, + ] + + response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value](self.payload(b'new')) + + assert response.code == 0 + assert app.persistence_mgr.execute_async.await_count == 4 + winner_update = compiled_params(app.persistence_mgr.execute_async.await_args_list[3].args[0]) + assert canonical_binary_key('plugin', 'test-author/test-plugin', 'test-key') in winner_update.values() + assert winner_update['value'] == b'new' + + @pytest.mark.asyncio + async def test_legacy_adoption_lost_to_delete_inserts_new_value(self, app): + runtime_handler = make_handler(app) + legacy_storage = SimpleNamespace(unique_key='plugin:test-author/test-plugin:test-key') + lost_race = SimpleNamespace(rowcount=0) + app.persistence_mgr.execute_async.side_effect = [ + make_result(), + make_result(legacy_storage), + lost_race, + SimpleNamespace(rowcount=0), + make_result(), + ] + + response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value](self.payload(b'new')) + + assert response.code == 0 + assert app.persistence_mgr.execute_async.await_count == 5 + insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[4].args[0]) + assert insert_params['unique_key'] == canonical_binary_key('plugin', 'test-author/test-plugin', 'test-key') + assert insert_params['value'] == b'new' + @pytest.mark.asyncio async def test_invalid_max_value_bytes_falls_back_to_default_limit(self, app): """Invalid max_value_bytes uses the 10MB default limit.""" @@ -525,6 +589,46 @@ class TestGetBinaryStorage: in statement_params.values() ) + @pytest.mark.asyncio + async def test_reads_legacy_storage_without_mutating_key(self, app): + runtime_handler = make_handler(app) + legacy_storage = SimpleNamespace( + unique_key='plugin:test-author/test-plugin:test-key', + value=b'legacy bytes', + ) + app.persistence_mgr.execute_async.side_effect = [ + make_result(), + make_result(legacy_storage), + ] + + response = await runtime_handler.actions[RuntimeToLangBotAction.GET_BINARY_STORAGE.value]( + {'key': 'test-key', 'owner_type': 'plugin', 'owner': 'ignored'} + ) + + assert response.code == 0 + assert base64.b64decode(response.data['value_base64']) == b'legacy bytes' + assert app.persistence_mgr.execute_async.await_count == 2 + + @pytest.mark.asyncio + async def test_retries_canonical_after_concurrent_legacy_adoption(self, app): + runtime_handler = make_handler(app) + canonical_storage = SimpleNamespace(value=b'adopted bytes') + app.persistence_mgr.execute_async.side_effect = [ + make_result(), + make_result(), + make_result(canonical_storage), + ] + + response = await runtime_handler.actions[RuntimeToLangBotAction.GET_BINARY_STORAGE.value]( + {'key': 'test-key', 'owner_type': 'plugin', 'owner': 'ignored'} + ) + + assert response.code == 0 + assert base64.b64decode(response.data['value_base64']) == b'adopted bytes' + assert app.persistence_mgr.execute_async.await_count == 3 + retry_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[2].args[0]) + assert canonical_binary_key('plugin', 'test-author/test-plugin', 'test-key') in retry_params.values() + @pytest.mark.asyncio async def test_returns_error_when_not_found(self, app): """Missing binary storage rows return an error response.""" @@ -567,21 +671,47 @@ class TestDeleteAndListBinaryStorage: assert response.code == 0 statement_params = compiled_params(app.persistence_mgr.execute_async.await_args.args[0]) - assert 'workspace-a' in statement_params.values() + flat_values = [ + item for value in statement_params.values() for item in (value if isinstance(value, list) else [value]) + ] + assert 'workspace-a' in flat_values assert ( canonical_binary_key( 'plugin', 'test-author/test-plugin', 'test-key', ) - in statement_params.values() + in flat_values ) - assert 'forged-owner' not in statement_params.values() + assert 'forged-owner' not in flat_values + + @pytest.mark.asyncio + async def test_delete_removes_canonical_and_legacy_scoped_keys(self, app): + runtime_handler = make_handler(app) + + response = await runtime_handler.actions[RuntimeToLangBotAction.DELETE_BINARY_STORAGE.value]( + { + 'key': 'test-key', + 'owner_type': 'plugin', + 'owner': 'forged-owner', + } + ) + + assert response.code == 0 + statement_params = compiled_params(app.persistence_mgr.execute_async.await_args.args[0]) + values = [ + item for value in statement_params.values() for item in (value if isinstance(value, list) else [value]) + ] + assert 'workspace-a' in values + assert canonical_binary_key('plugin', 'test-author/test-plugin', 'test-key') in values + assert 'plugin:test-author/test-plugin:test-key' in values + assert 'test-author/test-plugin' in values + assert 'forged-owner' not in values @pytest.mark.asyncio async def test_list_keys_uses_trusted_plugin_owner(self, app): result = Mock() - result.scalars.return_value.all.return_value = ['first', 'second'] + result.scalars.return_value.all.return_value = ['first', 'second', 'first'] app.persistence_mgr.execute_async.return_value = result runtime_handler = make_handler(app)