Merge pull request #7 from Golem-Base/rvdp/unsubscribe

feat: allow to unsubscribe from events

Author: r-vdp

example/golem_base_sdk_example/__init__.py

@@ -59,7 +59,7 @@
 }
 
 
-async def run_example(instance: str) -> None:
+async def run_example(instance: str) -> None:  # noqa: PLR0915
     """Run the example."""
     async with await anyio.open_file(
         BaseDirectory.xdg_config_home + "/golembase/private.key",
@@ -73,26 +73,31 @@ async def run_example(instance: str) -> None:
         private_key=key_bytes,
     )
 
-    await client.watch_logs(
-        lambda create: logger.info(
+    watch_logs_handle = await client.watch_logs(
+        label="first",
+        create_callback=lambda create: logger.info(
             """\n
 Got create event: %s
         """,
             create,
         ),
-        lambda update: logger.info(
+        update_callback=lambda update: logger.info(
             """\n
 Got update event: %s
         """,
             update,
         ),
-        lambda deleted_key: logger.info(
+    )
+
+    await client.watch_logs(
+        label="second",
+        delete_callback=lambda deleted_key: logger.info(
             """\n
 Got delete event: %s
         """,
             deleted_key,
         ),
-        lambda extension: logger.info(
+        extend_callback=lambda extension: logger.info(
             """\n
 Got extend event: %s
         """,
@@ -170,11 +175,11 @@ async def run_example(instance: str) -> None:
         logger.info(
             "block number: %s", await client.http_client().eth.get_block_number()
         )
-        update_receipt = await client.update_entities(
+        [update_receipt] = await client.update_entities(
             [GolemBaseUpdate(entity_key, b"hello", 60, [Annotation("app", "demo")], [])]
         )
         logger.info("receipt: %s", update_receipt)
-        entity_key = update_receipt[0].entity_key
+        entity_key = update_receipt.entity_key
 
         logger.info("entity metadata: %s", await client.get_entity_metadata(entity_key))
 
@@ -210,6 +215,8 @@ async def run_example(instance: str) -> None:
             "All entities: %s",
             await client.get_all_entity_keys(),
         )
+
+        await watch_logs_handle.unsubscribe()
     else:
         logger.warning("Could not connect to the API...")
 

example/nix/packages/golem-base-sdk-example.nix

@@ -27,7 +27,7 @@ let
 
 in
 
-pkgs.python3Packages.buildPythonPackage rec {
+pkgs.python3Packages.buildPythonPackage {
   inherit pname;
   version = "0.0.1";
 
@@ -65,9 +65,8 @@ pkgs.python3Packages.buildPythonPackage rec {
   ];
 
   checkPhase = ''
-    mypy --config ${../../../mypy.ini} ${src}/golem_base_sdk_example
-    ruff check --no-cache ${src}/golem_base_sdk_example
-    PYLINTHOME="$TMPDIR" pylint ${src}/golem_base_sdk_example
+    mypy golem_base_sdk_example
+    ruff check --no-cache golem_base_sdk_example
   '';
 
   meta = with lib; {

golem_base_sdk/__init__.py

@@ -7,13 +7,13 @@
 import typing
 from collections.abc import (
     AsyncGenerator,
+    Callable,
     Coroutine,
     Sequence,
 )
 from typing import (
     Any,
-    Callable,
-    Optional,
+    cast,
 )
 
 from eth_typing import ChecksumAddress, HexStr
@@ -48,6 +48,7 @@
     GolemBaseUpdate,
     QueryEntitiesResult,
     UpdateEntityReturnType,
+    WatchLogsHandle,
 )
 from .utils import parse_legacy_btl_extended_log, rlp_encode_transaction
 
@@ -126,12 +127,16 @@ def __init__(self, rpc_url: str):
     async def get_storage_value(self, entity_key: EntityKey) -> bytes:
         """Get the storage value stored in the given entity."""
         return base64.b64decode(
-            await self.eth.get_storage_value(entity_key.as_hex_string())  # type: ignore
+            await self.eth.get_storage_value(  # type: ignore[attr-defined]
+                entity_key.as_hex_string()
+            )
         )
 
     async def get_entity_metadata(self, entity_key: EntityKey) -> EntityMetadata:
         """Get the metadata of the given entity."""
-        metadata = await self.eth.get_entity_metadata(entity_key.as_hex_string())  # type: ignore
+        metadata = await self.eth.get_entity_metadata(  # type: ignore[attr-defined]
+            entity_key.as_hex_string()
+        )
 
         return EntityMetadata(
             entity_key=entity_key,
@@ -158,20 +163,22 @@ async def get_entities_to_expire_at_block(
         return list(
             map(
                 lambda e: EntityKey(GenericBytes.from_hex_string(e)),
-                await self.eth.get_entities_to_expire_at_block(block_number),  # type: ignore
+                await self.eth.get_entities_to_expire_at_block(  # type: ignore[attr-defined]
+                    block_number
+                ),
             )
         )
 
     async def get_entity_count(self) -> int:
         """Get the total entity count in Golem Base."""
-        return await self.eth.get_entity_count()  # type: ignore
+        return cast(int, await self.eth.get_entity_count())  # type: ignore[attr-defined]
 
     async def get_all_entity_keys(self) -> Sequence[EntityKey]:
         """Get all entity keys in Golem Base."""
         return list(
             map(
                 lambda e: EntityKey(GenericBytes.from_hex_string(e)),
-                await self.eth.get_all_entity_keys(),  # type: ignore
+                await self.eth.get_all_entity_keys(),  # type: ignore[attr-defined]
             )
         )
 
@@ -182,9 +189,7 @@ async def get_entities_of_owner(
         return list(
             map(
                 lambda e: EntityKey(GenericBytes.from_hex_string(e)),
-                # https://github.com/pylint-dev/pylint/issues/3162
-                # pylint: disable=no-member
-                await self.eth.get_entities_of_owner(owner),  # type: ignore
+                await self.eth.get_entities_of_owner(owner),  # type: ignore[attr-defined]
             )
         )
 
@@ -195,7 +200,7 @@ async def query_entities(self, query: str) -> Sequence[QueryEntitiesResult]:
                 lambda result: QueryEntitiesResult(
                     entity_key=result.key, storage_value=base64.b64decode(result.value)
                 ),
-                await self.eth.query_entities(query),  # type: ignore
+                await self.eth.query_entities(query),  # type: ignore[attr-defined]
             )
         )
 
@@ -228,6 +233,24 @@ async def create(
         ws_client = await AsyncWeb3(WebSocketProvider(ws_url))
         return GolemBaseClient(rpc_url, ws_client, private_key)
 
+    async def _start_subscription_loop(self) -> None:
+        """Create a long running task to handle subscriptions."""
+        # The loop will finish when there are no subscriptions left, so this method
+        # gets called every time a subscription is created, and we'll check
+        # whether we need to make a new task or whether one is already running.
+        if not self._background_tasks:
+            # Start the asyncio event loop
+            task = asyncio.create_task(
+                self.ws_client().subscription_manager.handle_subscriptions()
+            )
+            self._background_tasks.add(task)
+
+            def task_done(task: asyncio.Task[None]) -> None:
+                logger.info("Subscription background task done, removing...")
+                self._background_tasks.discard(task)
+
+            task.add_done_callback(task_done)
+
     def __init__(self, rpc_url: str, ws_client: AsyncWeb3, private_key: bytes) -> None:
         """Initialise the GolemBaseClient instance."""
         self._http_client = GolemBaseHttpClient(rpc_url)
@@ -262,7 +285,11 @@ async def inner(show_traceback: bool) -> bool:
         # The method on the provider is usually not called directly, instead you
         # can call the eponymous method on the client, which will delegate to the
         # provider.
-        self.http_client().provider.is_connected = is_connected(self.http_client())  # type: ignore
+        object.__setattr__(
+            self.http_client().provider,
+            "is_connected",
+            is_connected(self.http_client()),
+        )
 
         # Allow caching of certain methods to improve performance
         self.http_client().provider.cache_allowed_requests = True
@@ -303,7 +330,7 @@ def ws_client(self) -> AsyncWeb3:
 
     async def is_connected(self) -> bool:
         """Check whether the client's underlying http client is connected."""
-        return await self.http_client().is_connected()
+        return cast(bool, await self.http_client().is_connected())  # type: ignore[redundant-cast]
 
     async def disconnect(self) -> None:
         """
@@ -318,7 +345,7 @@ async def disconnect(self) -> None:
 
     def get_account_address(self) -> ChecksumAddress:
         """Get the address associated with the private key of this client."""
-        return self.account.address  # type: ignore
+        return cast(ChecksumAddress, self.account.address)
 
     async def get_storage_value(self, entity_key: EntityKey) -> bytes:
         """Get the storage value stored in the given entity."""
@@ -382,10 +409,10 @@ async def extend_entities(
 
     async def send_transaction(
         self,
-        creates: Optional[Sequence[GolemBaseCreate]] = None,
-        updates: Optional[Sequence[GolemBaseUpdate]] = None,
-        deletes: Optional[Sequence[GolemBaseDelete]] = None,
-        extensions: Optional[Sequence[GolemBaseExtend]] = None,
+        creates: Sequence[GolemBaseCreate] | None = None,
+        updates: Sequence[GolemBaseUpdate] | None = None,
+        deletes: Sequence[GolemBaseDelete] | None = None,
+        extensions: Sequence[GolemBaseExtend] | None = None,
     ) -> GolemBaseTransactionReceipt:
         """
         Send a generic transaction to Golem Base.
@@ -529,11 +556,13 @@ async def _send_gb_transaction(
 
     async def watch_logs(
         self,
-        create_callback: Callable[[CreateEntityReturnType], None],
-        update_callback: Callable[[UpdateEntityReturnType], None],
-        delete_callback: Callable[[EntityKey], None],
-        extend_callback: Callable[[ExtendEntityReturnType], None],
-    ) -> None:
+        *,
+        label: str,
+        create_callback: Callable[[CreateEntityReturnType], None] | None = None,
+        update_callback: Callable[[UpdateEntityReturnType], None] | None = None,
+        delete_callback: Callable[[EntityKey], None] | None = None,
+        extend_callback: Callable[[ExtendEntityReturnType], None] | None = None,
+    ) -> WatchLogsHandle:
         """
         Subscribe to events on Golem Base.
 
@@ -550,44 +579,62 @@ async def log_handler(
             logger.debug("New log: %s", log_receipt)
             res = await self._process_golem_base_log_receipt(log_receipt)
 
-            for create in res.creates:
-                create_callback(create)
-            for update in res.updates:
-                update_callback(update)
-            for key in res.deletes:
-                delete_callback(key)
-            for extension in res.extensions:
-                extend_callback(extension)
+            if create_callback:
+                for create in res.creates:
+                    create_callback(create)
+            if update_callback:
+                for update in res.updates:
+                    update_callback(update)
+            if delete_callback:
+                for key in res.deletes:
+                    delete_callback(key)
+            if extend_callback:
+                for extension in res.extensions:
+                    extend_callback(extension)
 
         def create_subscription(topic: HexStr) -> LogsSubscription:
             return LogsSubscription(
-                label=f"Golem Base subscription to topic {topic}",
+                label=f"Golem Base subscription to topic {topic} with label {label}",
                 address=self.golem_base_contract.address,
                 topics=[topic],
                 handler=log_handler,
                 # optional `handler_context` args to help parse a response
                 handler_context={},
             )
 
-        async def handle_subscriptions() -> None:
-            await self._ws_client.subscription_manager.subscribe(
-                list(
-                    map(
-                        lambda event: create_subscription(event.topic),
-                        self.golem_base_contract.all_events(),
-                    )
+        event_names = []
+        if create_callback:
+            event_names.append("GolemBaseStorageEntityCreated")
+        if update_callback:
+            event_names.append("GolemBaseStorageEntityUpdated")
+        if delete_callback:
+            event_names.append("GolemBaseStorageEntityDeleted")
+        if extend_callback:
+            event_names.extend(
+                [
+                    "GolemBaseStorageEntityBTLExtended",
+                    "GolemBaseStorageEntityBTLExptended",
+                    "GolemBaseStorageEntityTTLExptended",
+                ]
+            )
+
+        events = list(
+            map(
+                lambda event_name: create_subscription(
+                    self.golem_base_contract.get_event_by_name(event_name).topic
                 ),
+                event_names,
             )
-            # handle subscriptions via configured handlers:
-            await self.ws_client().subscription_manager.handle_subscriptions()
+        )
+        subscription_ids = await self._ws_client.subscription_manager.subscribe(
+            events,
+        )
+        logger.info("Sub ID: %s", subscription_ids)
 
-        # Create a long running task to handle subscriptions that we can run on
-        # the asyncio event loop
-        task = asyncio.create_task(handle_subscriptions())
-        self._background_tasks.add(task)
+        # Start a subscription loop in case there is none running
+        await self._start_subscription_loop()
 
-        def task_done(task: asyncio.Task[None]) -> None:
-            logger.info("Subscription background task done, removing...")
-            self._background_tasks.discard(task)
+        async def unsubscribe() -> None:
+            await self._ws_client.subscription_manager.unsubscribe(subscription_ids)
 
-        task.add_done_callback(task_done)
+        return WatchLogsHandle(_unsubscribe=unsubscribe)