diff --git a/flake.lock b/flake.lock index 91c90689..317d7b9e 100644 --- a/flake.lock +++ b/flake.lock @@ -2,11 +2,11 @@ "nodes": { "nixpkgs": { "locked": { - "lastModified": 1774106199, - "narHash": "sha256-US5Tda2sKmjrg2lNHQL3jRQ6p96cgfWh3J1QBliQ8Ws=", + "lastModified": 1782723713, + "narHash": "sha256-oPXCU/SSUokcGaJREHibG1CBX3+s/W7orDWQOZDsEeQ=", "owner": "nixos", "repo": "nixpkgs", - "rev": "6c9a78c09ff4d6c21d0319114873508a6ec01655", + "rev": "b5aa0fbd538984f6e3d201be0005b4463d8b09f8", "type": "github" }, "original": { diff --git a/flake.nix b/flake.nix index a24fda76..84b24cc0 100644 --- a/flake.nix +++ b/flake.nix @@ -34,6 +34,7 @@ pkgs.uv pkgs.gnumake pkgs.docker + pkgs.postgresql_15 ]; workspace = uv2nix.lib.workspace.loadWorkspace { workspaceRoot = ./.; }; @@ -47,11 +48,14 @@ pyproject-overlay = pkgs: final: prev: { ruamel-yaml-clib = prev.ruamel-yaml-clib.overrideAttrs (old: { - nativeBuildInputs = old.nativeBuildInputs ++ [ - (final.resolveBuildSystem { + nativeBuildInputs = old.nativeBuildInputs ++ (final.resolveBuildSystem { setuptools = []; - }) - ]; + }); + }); + urwid-readline = prev.urwid-readline.overrideAttrs (old: { + nativeBuildInputs = old.nativeBuildInputs ++ (final.resolveBuildSystem { + setuptools = []; + }); }); pyiceberg = prev.pyiceberg.overrideAttrs (old: { buildInputs = (old.buildInputs or []) ++ [ final.poetry-core ]; @@ -61,12 +65,9 @@ sed -i '1i from Cython.Build import cythonize' setup.py sed -i 's/ext_modules=[pyroaring_module]/ext_modules=[cythonize(pyroaring_module)]/' setup.py ''; - nativeBuildInputs = old.nativeBuildInputs ++ [ - (final.resolveBuildSystem { + nativeBuildInputs = old.nativeBuildInputs ++ [final.cython] ++ (final.resolveBuildSystem { setuptools = []; - }) - final.cython - ]; + }); }); }; diff --git a/src/realtime/infra/supabase/migrations/20240810125543_init.sql b/src/realtime/infra/supabase/migrations/20240810125543_init.sql index 3f9383bd..c4e64e40 100644 --- a/src/realtime/infra/supabase/migrations/20240810125543_init.sql +++ b/src/realtime/infra/supabase/migrations/20240810125543_init.sql @@ -14,6 +14,11 @@ CREATE POLICY "Allow authenticated access" ON "public"."todos" AS permissive USING (TRUE) WITH CHECK (TRUE); +CREATE POLICY "authenticated can receive broadcast" ON "realtime"."messages" +FOR ALL +TO AUTHENTICATED +USING (TRUE); + ALTER publication supabase_realtime ADD TABLE todos; @@ -26,3 +31,27 @@ create table messages ( -- enable realtime on messages table alter publication supabase_realtime add table messages; + +-- Grant API access to tables for anon and authenticated roles. +-- Required since Supabase no longer grants public schema access by default +-- (see https://github.com/orgs/supabase/discussions/45329). +GRANT ALL ON TABLE public.todos TO anon, authenticated; +GRANT ALL ON TABLE public.messages TO anon, authenticated; + +-- Grant sequence usage for tables with generated identity primary keys +GRANT USAGE, SELECT ON SEQUENCE public.messages_id_seq TO anon, authenticated; + +CREATE OR replace FUNCTION public.send_realtime( + topic text, + event text, + payload jsonb, + private boolean default True +) +RETURNS void +LANGUAGE SQL +SECURITY DEFINER +AS $$ + SELECT realtime.send(payload, event, topic, private); +$$; + +GRANT EXECUTE ON FUNCTION public.send_realtime(text,text,jsonb,boolean) TO anon, authenticated, service_role; diff --git a/src/realtime/pyproject.toml b/src/realtime/pyproject.toml index 3a4aef1f..e3c5e7eb 100644 --- a/src/realtime/pyproject.toml +++ b/src/realtime/pyproject.toml @@ -12,11 +12,13 @@ maintainers = [ license = "MIT" readme = "README.md" repository = "https://github.com/supabase/supabase-py" -requires-python = ">=3.9" +requires-python = ">=3.10" dependencies = [ + "supabase_utils==3.0.0a1", # x-release-please-version "websockets >=11,<16", "typing-extensions >=4.14.0", "pydantic (>=2.11.7,<3.0.0)", + "yarl>=1.22.0", ] [project.urls] @@ -32,6 +34,7 @@ tests = [ "pytest-cov >= 6.2.1", "python-dotenv >= 1.1.1", "pytest-asyncio >= 1.0.0", + "pudb>=2025.1.5", ] lints = [ "ruff >= 0.12.1", @@ -68,11 +71,10 @@ ignore = ["F401", "F403", "F841", "E712", "E501", "E402", "UP006", "UP035"] keep-runtime-typing = true [tool.pytest.ini_options] -asyncio_mode = "strict" -asyncio_default_fixture_loop_scope = "function" +asyncio_mode = "auto" [tool.mypy] -python_version = "3.9" +python_version = "3.10" check_untyped_defs = true allow_redefinition = true diff --git a/src/realtime/src/realtime/__init__.py b/src/realtime/src/realtime/__init__.py index a7749e74..58fb1984 100644 --- a/src/realtime/src/realtime/__init__.py +++ b/src/realtime/src/realtime/__init__.py @@ -5,13 +5,9 @@ from realtime.version import __version__ -from ._async.channel import AsyncRealtimeChannel -from ._async.client import AsyncRealtimeClient -from ._async.presence import AsyncRealtimePresence -from ._sync.channel import SyncRealtimeChannel -from ._sync.client import SyncRealtimeClient -from ._sync.presence import SyncRealtimePresence +from .channel import RealtimeChannel, RealtimeChannelOptions +from .client import RealtimeClient, connect_once from .exceptions import * from .message import * -from .transformers import * +from .presence import RealtimePresence from .types import * diff --git a/src/realtime/src/realtime/_async/__init__.py b/src/realtime/src/realtime/_async/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/src/realtime/src/realtime/_async/channel.py b/src/realtime/src/realtime/_async/channel.py deleted file mode 100644 index a844dd0b..00000000 --- a/src/realtime/src/realtime/_async/channel.py +++ /dev/null @@ -1,563 +0,0 @@ -from __future__ import annotations - -import asyncio -import json -import logging -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional - -from typing_extensions import assert_never - -from realtime.types import ( - BroadcastCallback, - BroadcastPayload, - Callback, - ChannelEvents, - ChannelStates, - PostgresChangesCallback, - PostgresChangesData, - PresenceOnJoinCallback, - PresenceOnLeaveCallback, - RealtimeAcknowledgementStatus, - RealtimeChannelBroadcastConfig, - RealtimeChannelConfig, - RealtimeChannelOptions, - RealtimeChannelPresenceConfig, - RealtimePostgresChangesListenEvent, - RealtimePresenceState, - RealtimeSubscribeStates, -) - -from ..message import ( - BroadcastMessage, - ChannelCloseMessage, - ChannelErrorMessage, - HeartbeatMessage, - Message, - PostgresChangesMessage, - PostgresChangesPayload, - PostgresRowChange, - PresenceDiffMessage, - PresenceStateMessage, - ReplyMessage, - ReplyPostgresChanges, - ServerMessage, - SuccessReplyMessage, - SuccessSystemPayload, - SystemMessage, -) -from ..transformers import http_endpoint_url -from .presence import ( - AsyncRealtimePresence, -) -from .push import AsyncPush -from .timer import AsyncTimer - -if TYPE_CHECKING: - from .client import AsyncRealtimeClient - -logger = logging.getLogger(__name__) - - -class AsyncRealtimeChannel: - """ - Channel is an abstraction for a topic subscription on an existing socket connection. - Each Channel has its own topic and a list of event-callbacks that respond to messages. - Should only be instantiated through `AsyncRealtimeClient.channel(topic)`. - """ - - def __init__( - self, - socket: AsyncRealtimeClient, - topic: str, - params: Optional[RealtimeChannelOptions] = None, - ) -> None: - """ - Initialize the Channel object. - - :param socket: RealtimeClient object - :param topic: Topic that it subscribes to on the realtime server - :param params: Optional parameters for connection. - """ - self.socket = socket - self.params: RealtimeChannelOptions = ( - params - if params - else { - "config": { - "broadcast": {"ack": False, "self": False}, - "presence": {"key": "", "enabled": False}, - "private": False, - } - } - ) - self.topic = topic - self._joined_once = False - self.presence: AsyncRealtimePresence = AsyncRealtimePresence() - self.state = ChannelStates.CLOSED - self._push_buffer: list[AsyncPush] = [] - self.timeout = self.socket.timeout - - self.join_push: AsyncPush = AsyncPush(self, ChannelEvents.join, self.params) - self.messages_waiting_for_ack: dict[str, AsyncPush] = {} - self.broadcast_callbacks: list[BroadcastCallback] = [] - self.system_callbacks: list[Callable[[SuccessSystemPayload], None]] = [] - self.postgres_changes_callbacks: list[PostgresChangesCallback] = [] - self.subscribe_callback: Optional[ - Callable[[RealtimeSubscribeStates, Optional[Exception]], None] - ] = None - - self.rejoin_timer = AsyncTimer( - self._rejoin_until_connected, lambda tries: 2**tries - ) - - self.broadcast_endpoint_url = self._broadcast_endpoint_url() - - def on_join_push_ok(payload: ReplyPostgresChanges): - self.state = ChannelStates.JOINED - self.rejoin_timer.reset() - for push in self._push_buffer: - asyncio.create_task(push.send()) - self._push_buffer = [] - - def on_join_push_timeout(): - if not self.is_joining: - return - - logger.error(f"join push timeout for channel {self.topic}") - self.state = ChannelStates.ERRORED - self.rejoin_timer.schedule_timeout() - - self.join_push.receive( - RealtimeAcknowledgementStatus.Ok, on_join_push_ok - ).receive(RealtimeAcknowledgementStatus.Timeout, on_join_push_timeout) - - def on_close(self): - logger.info(f"channel {self.topic} closed") - self.rejoin_timer.reset() - self.state = ChannelStates.CLOSED - self.socket._remove_channel(self) - - def on_error(self, payload: dict[str, Any]): - if self.is_leaving or self.is_closed: - return - - logger.info(f"channel {self.topic} error: {payload}") - self.state = ChannelStates.ERRORED - self.rejoin_timer.schedule_timeout() - - # Properties - @property - def is_closed(self): - return self.state == ChannelStates.CLOSED - - @property - def is_joining(self): - return self.state == ChannelStates.JOINING - - @property - def is_leaving(self): - return self.state == ChannelStates.LEAVING - - @property - def is_errored(self): - return self.state == ChannelStates.ERRORED - - @property - def is_joined(self): - return self.state == ChannelStates.JOINED - - # Core channel methods - async def subscribe( - self, - callback: Optional[ - Callable[[RealtimeSubscribeStates, Optional[Exception]], None] - ] = None, - ) -> AsyncRealtimeChannel: - """ - Subscribe to the channel. Can only be called once per channel instance. - - :param callback: Optional callback function that receives subscription state updates - and any errors that occur during subscription - :return: The Channel instance for method chaining - :raises: Exception if called multiple times on the same channel instance - """ - if not self.socket.is_connected: - await self.socket.connect() - - if self._joined_once: - raise Exception( - "Tried to subscribe multiple times. 'subscribe' can only be called a single time per channel instance" - ) - else: - config: RealtimeChannelConfig = self.params["config"] - broadcast = config.get("broadcast") - presence = config.get("presence") or RealtimeChannelPresenceConfig( - key="", enabled=False - ) - private = config.get("private", False) - - presence_enabled = self.presence._has_callback_attached or presence.get( - "enabled", False - ) - presence["enabled"] = presence_enabled - - config_payload: Dict[str, Any] = { - "config": { - "broadcast": broadcast, - "presence": presence, - "private": private, - "postgres_changes": [ - c.binding_filter for c in self.postgres_changes_callbacks - ], - } - } - - if self.socket.access_token: - config_payload["access_token"] = self.socket.access_token - - self.join_push.update_payload(config_payload) - self._joined_once = True - - def on_join_push_ok(payload: ReplyPostgresChanges): - server_postgres_changes = payload.postgres_changes - - new_postgres_bindings = [] - - if server_postgres_changes: - for i, postgres_callback in enumerate( - self.postgres_changes_callbacks - ): - server_binding: Optional[PostgresRowChange] = ( - server_postgres_changes[i] - if i < len(server_postgres_changes) - else None - ) - logger.debug(f"{server_binding}, {postgres_callback}") - - if ( - server_binding - and server_binding.event == postgres_callback.event - and server_binding.schema_ == postgres_callback.schema - and server_binding.table == postgres_callback.table - and server_binding.filter == postgres_callback.filter - ): - postgres_callback.id = server_binding.id - new_postgres_bindings.append(postgres_callback) - else: - asyncio.create_task(self.unsubscribe()) - callback and callback( - RealtimeSubscribeStates.CHANNEL_ERROR, - Exception( - "mismatch between server and client bindings for postgres changes" - ), - ) - return - - self.postgres_changes_callbacks = new_postgres_bindings - callback and callback(RealtimeSubscribeStates.SUBSCRIBED, None) - - def on_join_push_error(payload: Dict[str, Any]): - callback and callback( - RealtimeSubscribeStates.CHANNEL_ERROR, - Exception(json.dumps(payload)), - ) - - def on_join_push_timeout(*args): - callback and callback(RealtimeSubscribeStates.TIMED_OUT, None) - - self.join_push.receive( - RealtimeAcknowledgementStatus.Ok, on_join_push_ok - ).receive(RealtimeAcknowledgementStatus.Error, on_join_push_error).receive( - RealtimeAcknowledgementStatus.Timeout, on_join_push_timeout - ) - - await self._rejoin() - - return self - - async def unsubscribe(self) -> None: - """ - Unsubscribe from the channel and leave the topic. - Sets channel state to LEAVING and cleans up timers and pushes. - """ - self.state = ChannelStates.LEAVING - - self.rejoin_timer.reset() - self.join_push.destroy() - - def _close(*args) -> None: - logger.info(f"channel {self.topic} leave") - self.on_close() - - leave_push = AsyncPush(self, ChannelEvents.leave, {}) - leave_push.receive(RealtimeAcknowledgementStatus.Ok, _close).receive( - RealtimeAcknowledgementStatus.Error, _close - ) - await leave_push.send() - - async def push( - self, event: str, payload: Dict[str, Any], timeout: Optional[int] = None - ) -> AsyncPush: - """ - Push a message to the channel. - - :param event: The event name to push - :param payload: The payload to send - :param timeout: Optional timeout in milliseconds - :return: AsyncPush instance representing the push operation - :raises: Exception if called before subscribing to the channel - """ - if not self._joined_once: - raise Exception( - f"tried to push '{event}' to '{self.topic}' before joining. Use channel.subscribe() before pushing events" - ) - - timeout = timeout or self.timeout - - push = AsyncPush(self, event, payload, timeout) - if self._can_push(): - await push.send() - assert push.ref is not None, "Sent AsyncPush should have a ref" - else: - push.start_timeout() - self._push_buffer.append(push) - - return push - - async def join(self) -> AsyncRealtimeChannel: - """ - Coroutine that attempts to join Phoenix Realtime server via a certain topic. - - :return: Channel - """ - config = self.params["config"] - broadcast = config.get("broadcast") - presence = config.get("presence") - private = config.get("private", False) - - config_payload: Dict[str, Any] = { - "config": { - "broadcast": broadcast, - "presence": presence, - "private": private, - "postgres_changes": [ - c.binding_filter for c in self.postgres_changes_callbacks - ], - } - } - message = Message( - topic=self.topic, - event=ChannelEvents.join, - payload={"config": config_payload}, - ref=self.socket._make_ref(), - ) - await self.socket.send(message) - return self - - def on_broadcast( - self, event: str, callback: Callable[[BroadcastPayload], None] - ) -> AsyncRealtimeChannel: - """ - Set up a listener for a specific broadcast event. - - :param event: The name of the broadcast event to listen for - :param callback: Function called with the payload when a matching broadcast is received - :return: The Channel instance for method chaining - """ - self.broadcast_callbacks.append( - BroadcastCallback(callback=callback, event=event) - ) - return self - - def on_postgres_changes( - self, - event: RealtimePostgresChangesListenEvent, - callback: Callable[[PostgresChangesPayload], None], - table: Optional[str] = None, - schema: Optional[str] = None, - filter: Optional[str] = None, - ) -> AsyncRealtimeChannel: - """ - Set up a listener for Postgres database changes. - - :param event: The type of database event to listen for (INSERT, UPDATE, DELETE, or *) - :param callback: Function called with the payload when a matching change is detected - :param table: The table name to monitor. Defaults to "*" for all tables - :param schema: The database schema to monitor. Defaults to "public" - :param filter: Optional filter string to apply - :return: The Channel instance for method chaining - """ - callback = PostgresChangesCallback( - callback=callback, event=event, table=table, schema=schema, filter=filter - ) - self.postgres_changes_callbacks.append(callback) - return self - - def on_system( - self, callback: Callable[[SuccessSystemPayload], None] - ) -> AsyncRealtimeChannel: - """ - Set up a listener for system events. - - :param callback: The callback function to execute when a system event is received. - :return: The Channel instance for method chaining. - """ - self.system_callbacks.append(callback) - return self - - # Presence methods - async def track(self, user_status: Dict[str, Any]) -> None: - """ - Track presence status for the current user. - - :param user_status: Dictionary containing the user's presence information - """ - await self.send_presence("track", user_status) - - async def untrack(self) -> None: - """ - Stop tracking presence for the current user. - """ - await self.send_presence("untrack", {}) - - def presence_state(self) -> RealtimePresenceState: - """ - Get the current state of presence on this channel. - - :return: Dictionary mapping presence keys to lists of presence payloads - """ - return self.presence.state - - def on_presence_sync(self, callback: Callable[[], None]) -> AsyncRealtimeChannel: - """ - Register a callback for presence sync events. - - :param callback: The callback function to execute when a presence sync event occurs. - :return: The Channel instance for method chaining. - """ - self.presence.on_sync(callback) - - if self.is_joined: - logger.info( - f"channel {self.topic} resubscribe due to change in presence callbacks on joined channel" - ) - asyncio.create_task(self._resubscribe()) - - return self - - def on_presence_join( - self, callback: PresenceOnJoinCallback - ) -> AsyncRealtimeChannel: - """ - Register a callback for presence join events. - - :param callback: The callback function to execute when a presence join event occurs. - :return: The Channel instance for method chaining. - """ - self.presence.on_join(callback) - if self.is_joined: - logger.info( - f"channel {self.topic} resubscribe due to change in presence callbacks on joined channel" - ) - asyncio.create_task(self._resubscribe()) - - return self - - def on_presence_leave( - self, callback: PresenceOnLeaveCallback - ) -> AsyncRealtimeChannel: - """ - Register a callback for presence leave events. - - :param callback: The callback function to execute when a presence leave event occurs. - :return: The Channel instance for method chaining. - """ - self.presence.on_leave(callback) - if self.is_joined: - logger.info( - f"channel {self.topic} resubscribe due to change in presence callbacks on joined channel" - ) - asyncio.create_task(self._resubscribe()) - return self - - # Broadcast methods - async def send_broadcast(self, event: str, data: Any) -> None: - """ - Send a broadcast message through this channel. - - :param event: The name of the broadcast event - :param data: The payload to broadcast - """ - await self.push( - ChannelEvents.broadcast, - {"type": "broadcast", "event": event, "payload": data}, - ) - - # Internal methods - - async def _resubscribe(self) -> None: - await self.unsubscribe() - await self.subscribe() - - def _broadcast_endpoint_url(self): - return f"{http_endpoint_url(self.socket.http_endpoint)}/api/broadcast" - - async def _rejoin(self) -> None: - if self.is_leaving: - return - logger.debug(f"Rejoining channel after reconnection: {self.topic}") - self.state = ChannelStates.JOINING - await self.join_push.resend() - - def _can_push(self): - return self.socket.is_connected and self._joined_once - - async def send_presence(self, event: str, data: Any) -> None: - await self.push(ChannelEvents.presence, {"event": event, "payload": data}) - - def _handle_message(self, message: ServerMessage): - logger.debug(f"{self.topic} : {message!r}") - if isinstance(message, SystemMessage): - if isinstance(message.payload, SuccessSystemPayload): - for callback in self.system_callbacks: - callback(message.payload) - else: - self.on_error(dict(message.payload)) - elif isinstance(message, ReplyMessage): - reply_payload = message.payload - if message.ref and ( - push := self.messages_waiting_for_ack.pop(message.ref, None) - ): - if reply_payload.status == "ok": - push.trigger( - RealtimeAcknowledgementStatus.Ok, reply_payload.response - ) - else: - push.trigger( - RealtimeAcknowledgementStatus.Error, reply_payload.response - ) - elif isinstance(message, BroadcastMessage): - broadcast_payload = message.payload - for broadcast_callback in self.broadcast_callbacks: - broadcast_callback(broadcast_payload) - elif isinstance(message, PresenceStateMessage): - self.presence._on_state_event(message.payload) - elif isinstance(message, PresenceDiffMessage): - self.presence._on_diff_event(message.payload) - elif isinstance(message, PostgresChangesMessage): - payload = message.payload - for postgres_callback in self.postgres_changes_callbacks: - postgres_callback(payload) - elif isinstance(message, ChannelErrorMessage): - self.on_error(message.payload) - elif isinstance(message, ChannelCloseMessage): - self.on_close() - elif isinstance(message, HeartbeatMessage): # do nothing - return - else: - assert_never(message) - - async def _rejoin_until_connected(self): - self.rejoin_timer.schedule_timeout() - if self.socket.is_connected: - await self._rejoin() diff --git a/src/realtime/src/realtime/_async/client.py b/src/realtime/src/realtime/_async/client.py deleted file mode 100644 index a9b98850..00000000 --- a/src/realtime/src/realtime/_async/client.py +++ /dev/null @@ -1,406 +0,0 @@ -import asyncio -import json -import logging -import re -import sys -from functools import wraps -from typing import Any, Callable, Dict, List, Optional, Union -from urllib.parse import urlencode, urlparse, urlunparse -from warnings import warn - -import websockets -from pydantic import ValidationError -from websockets import connect -from websockets.asyncio.client import ClientConnection - -from ..exceptions import NotConnectedError -from ..message import Message, ServerMessageAdapter -from ..transformers import http_endpoint_url -from ..types import ( - DEFAULT_HEARTBEAT_INTERVAL, - DEFAULT_TIMEOUT, - PHOENIX_CHANNEL, - VSN, - ChannelEvents, - ChannelStates, -) -from ..utils import is_ws_url -from .channel import AsyncRealtimeChannel, RealtimeChannelOptions - - -def deprecated(func: Callable) -> Callable: - @wraps(func) - def wrapper(*args, **kwargs): - logger.warning(f"Warning: {func.__name__} is deprecated.") - return func(*args, **kwargs) - - return wrapper - - -logger = logging.getLogger(__name__) - - -class AsyncRealtimeClient: - def __init__( - self, - url: str, - token: Optional[str] = None, - auto_reconnect: bool = True, - params: Optional[Dict[str, Any]] = None, - hb_interval: int = DEFAULT_HEARTBEAT_INTERVAL, - max_retries: int = 5, - initial_backoff: float = 1.0, - timeout: int = DEFAULT_TIMEOUT, - ) -> None: - """ - Initialize a RealtimeClient instance for WebSocket communication. - - :param url: WebSocket URL of the Realtime server. Starts with `ws://` or `wss://`. - Also accepts default Supabase URL: `http://` or `https://`. - :param token: Authentication token for the WebSocket connection. - :param auto_reconnect: If True, automatically attempt to reconnect on disconnection. Defaults to True. - :param params: Optional parameters for the connection. Defaults to None. - :param hb_interval: Interval (in seconds) for sending heartbeat messages to keep the connection alive. Defaults to 25. - :param max_retries: Maximum number of reconnection attempts. Defaults to 5. - :param initial_backoff: Initial backoff time (in seconds) for reconnection attempts. Defaults to 1.0. - :param timeout: Connection timeout in seconds. Defaults to DEFAULT_TIMEOUT. - """ - if not is_ws_url(url): - raise ValueError("url must be a valid WebSocket URL or HTTP URL string") - if sys.version_info < (3, 10): - warn( - "Python versions below 3.10 are deprecated and will not be supported in future versions. Please upgrade to Python 3.10 or newer.", - DeprecationWarning, - stacklevel=2, - ) - - self.url = f"{re.sub(r'https://', 'wss://', re.sub(r'http://', 'ws://', url, flags=re.IGNORECASE), flags=re.IGNORECASE)}/websocket" - if token: - self.url += f"?apikey={token}" - self.http_endpoint = http_endpoint_url(url) - self.params = params or {} - self.apikey = token - self.access_token = token - self.send_buffer: List[Callable] = [] - self.hb_interval = hb_interval - self._ws_connection: Optional[ClientConnection] = None - self.ref = 0 - self.auto_reconnect = auto_reconnect - self.channels: Dict[str, AsyncRealtimeChannel] = {} - self.max_retries = max_retries - self.initial_backoff = initial_backoff - self.timeout = timeout - self._listen_task: Optional[asyncio.Task] = None - self._heartbeat_task: Optional[asyncio.Task] = None - - @property - def is_connected(self) -> bool: - return self._ws_connection is not None - - async def _listen(self) -> None: - """ - An infinite loop that keeps listening. - :return: None - """ - - if not self._ws_connection: - raise NotConnectedError("_listen") - - try: - async for msg in self._ws_connection: - logger.debug(f"receive: {msg!r}") - - try: - message = ServerMessageAdapter.validate_json(msg) - except ValidationError as e: - logger.error(f"Unrecognized message format {msg!r}\n{e}") - continue - logger.debug(f"parsed message as {message!r}") - if channel := self.channels.get(message.topic): - channel._handle_message(message) - except websockets.exceptions.ConnectionClosedError as e: - await self._on_connect_error(e) - - async def _reconnect(self) -> None: - self._ws_connection = None - - to_rejoin = [ - chan - for chan in self.channels.values() - if chan.state == ChannelStates.JOINED or chan.state == ChannelStates.JOINING - ] - for channel in to_rejoin: - channel.state = ChannelStates.ERRORED - - await self.connect() - - if self.is_connected: - rejoins = [asyncio.Task(chan._rejoin()) for chan in to_rejoin] - await asyncio.wait(rejoins) - - async def connect(self) -> None: - """ - Establishes a WebSocket connection with exponential backoff retry mechanism. - - This method attempts to connect to the WebSocket server. If the connection fails, - it will retry with an exponential backoff strategy up to a maximum number of retries. - - Returns: - None - - Raises: - Exception: If unable to establish a connection after max_retries attempts. - - Note: - - The initial backoff time and maximum retries are set during RealtimeClient initialization. - - The backoff time doubles after each failed attempt, up to a maximum of 60 seconds. - """ - - if self.is_connected: - logger.debug("WebSocket connection already established") - return - - retries = 0 - backoff = self.initial_backoff - - logger.debug(f"Attempting to connect to WebSocket at {self.url}") - - while retries < self.max_retries: - try: - ws = await connect(self.url) - self._ws_connection = ws - logger.debug("WebSocket connection established successfully") - return await self._on_connect() - except Exception as e: - retries += 1 - logger.error(f"Connection attempt failed: {str(e)}") - - if retries >= self.max_retries or not self.auto_reconnect: - logger.error( - f"Connection failed permanently after {retries} attempts. Error: {str(e)}" - ) - raise - else: - wait_time = backoff * (2 ** (retries - 1)) - logger.debug( - f"Retry {retries}/{self.max_retries}: Next attempt in {wait_time:.2f}s (backoff={backoff}s)" - ) - await asyncio.sleep(wait_time) - backoff = min(backoff, 60) - - raise Exception( - f"Failed to establish WebSocket connection after {self.max_retries} attempts" - ) - - @deprecated - async def listen(self): - pass - - async def _on_connect(self) -> None: - if self._listen_task: - self._listen_task.cancel() - self._listen_task = None - - if self._heartbeat_task: - self._heartbeat_task.cancel() - self._heartbeat_task = None - - self._listen_task = asyncio.create_task(self._listen()) - self._heartbeat_task = asyncio.create_task(self._heartbeat()) - await self._flush_send_buffer() - - async def _on_connect_error( - self, e: websockets.exceptions.ConnectionClosedError - ) -> None: - logger.error( - f"WebSocket connection closed with code: {e.code}, reason: {e.reason}" - ) - - if self.auto_reconnect: - logger.debug("Initiating auto-reconnect sequence...") - await self._reconnect() - else: - logger.error("Auto-reconnect disabled, terminating connection") - - async def _flush_send_buffer(self): - if self.is_connected and len(self.send_buffer) > 0: - for callback in self.send_buffer: - await callback() - self.send_buffer = [] - - async def close(self) -> None: - """ - Close the WebSocket connection. - - Returns: - None - - Raises: - NotConnectedError: If the connection is not established when this method is called. - """ - - if self._ws_connection: - await self._ws_connection.close() - - self._ws_connection = None - - if self._listen_task: - self._listen_task.cancel() - self._listen_task = None - - if self._heartbeat_task: - self._heartbeat_task.cancel() - self._heartbeat_task = None - - async def _heartbeat(self) -> None: - if not self._ws_connection: - raise NotConnectedError("_heartbeat") - - while self.is_connected: - try: - data = Message( - topic=PHOENIX_CHANNEL, - event=ChannelEvents.heartbeat, - payload={}, - ref=None, - ) - await self.send(data) - await asyncio.sleep(max(self.hb_interval, 15)) - - except websockets.exceptions.ConnectionClosedError as e: - await self._on_connect_error(e) - except websockets.exceptions.ConnectionClosedOK as e: - pass - - def channel( - self, topic: str, params: Optional[RealtimeChannelOptions] = None - ) -> AsyncRealtimeChannel: - """ - Initialize a channel and create a two-way association with the socket. - - :param topic: The topic to subscribe to - :param params: Optional channel parameters - :return: AsyncRealtimeChannel instance - """ - topic = f"realtime:{topic}" - chan = AsyncRealtimeChannel(self, topic, params) - self.channels[topic] = chan - - return chan - - def get_channels(self) -> List[AsyncRealtimeChannel]: - return list(self.channels.values()) - - def _remove_channel(self, channel: AsyncRealtimeChannel) -> None: - del self.channels[channel.topic] - - async def remove_channel(self, channel: AsyncRealtimeChannel) -> None: - """ - Unsubscribes and removes a channel from the socket - :param channel: Channel to remove - :return: None - """ - if channel.topic in self.channels: - await self.channels[channel.topic].unsubscribe() - self._remove_channel(channel) - - if len(self.channels) == 0: - await self.close() - - async def remove_all_channels(self) -> None: - """ - Unsubscribes and removes all channels from the socket - :return: None - """ - for _, channel in self.channels.items(): - await channel.unsubscribe() - - await self.close() - - async def set_auth(self, token: Optional[str]) -> None: - """ - Set the authentication token for the connection and update all joined channels. - - This method updates the access token for the current connection and sends the new token - to all joined channels. This is useful for refreshing authentication or changing users. - - Args: - token (Optional[str]): The new authentication token. Can be None to remove authentication. - - Returns: - None - """ - self.access_token = token - - for _, channel in self.channels.items(): - if channel._joined_once and channel.is_joined: - await channel.push(ChannelEvents.access_token, {"access_token": token}) - - def _make_ref(self) -> str: - self.ref += 1 - return f"{self.ref}" - - async def send(self, message: Union[Message, Dict[str, Any]]) -> None: - """ - Send a message through the WebSocket connection. - - This method serializes the given message dictionary to JSON, - and sends it through the WebSocket connection. If the connection - is not currently established, the message will be buffered and sent - once the connection is re-established. - - Args: - message (Dict[str, Any]): The message to be sent, as a dictionary. - - Returns: - None - """ - if isinstance(message, Message): - msg = message - else: - logger.warning( - "Warning: calling AsyncRealtimeClient.send with a dictionary is deprecated. Please call it with a Message object instead. This will be a hard error in the future." - ) - msg = Message(**message) - message_str = msg.model_dump_json() - logger.debug(f"send: {message_str}") - - async def send_message(): - if not self._ws_connection: - raise NotConnectedError("_send") - - try: - await self._ws_connection.send(message_str) - except websockets.exceptions.ConnectionClosedError as e: - await self._on_connect_error(e) - except websockets.exceptions.ConnectionClosedOK: - pass - - if self.is_connected: - await send_message() - else: - self.send_buffer.append(send_message) - - async def _leave_open_topic(self, topic: str): - dup_channels = [ - ch - for ch in self.channels.values() - if ch.topic == topic and (ch.is_joined or ch.is_joining) - ] - - for ch in dup_channels: - await ch.unsubscribe() - - def endpoint_url(self) -> str: - parsed_url = urlparse(self.url) - query = urlencode({**self.params, "vsn": VSN}, doseq=True) - return urlunparse( - ( - parsed_url.scheme, - parsed_url.netloc, - parsed_url.path, - parsed_url.params, - query, - parsed_url.fragment, - ) - ) diff --git a/src/realtime/src/realtime/_async/presence.py b/src/realtime/src/realtime/_async/presence.py deleted file mode 100644 index ad376702..00000000 --- a/src/realtime/src/realtime/_async/presence.py +++ /dev/null @@ -1,188 +0,0 @@ -""" -Defines the RealtimePresence class and its dependencies. -""" - -import logging -from typing import Any, Callable, Dict, List, Optional, Union - -from ..types import ( - Presence, - PresenceDiff, - PresenceEvents, - PresenceOnJoinCallback, - PresenceOnLeaveCallback, - PresenceOpts, - RawPresenceDiff, - RawPresenceState, - RealtimePresenceState, -) - -logger = logging.getLogger(__name__) - - -class AsyncRealtimePresence: - @property - def _has_callback_attached(self) -> bool: - return ( - self.on_join_callback is not None - or self.on_leave_callback is not None - or self.on_sync_callback is not None - ) - - def __init__(self): - self.state: RealtimePresenceState = {} - self.on_join_callback: Optional[PresenceOnJoinCallback] = None - self.on_leave_callback: Optional[PresenceOnLeaveCallback] = None - self.on_sync_callback: Optional[Callable[[], None]] = None - - def on_join(self, callback: PresenceOnJoinCallback): - self.on_join_callback = callback - - def on_leave(self, callback: PresenceOnLeaveCallback): - self.on_leave_callback = callback - - def on_sync(self, callback: Callable[[], None]): - self.on_sync_callback = callback - - def _on_state_event(self, payload: RawPresenceState): - state = AsyncRealtimePresence._transform_state(payload) - self.state = self._sync_state(state) - - if self.on_sync_callback: - self.on_sync_callback() - - def _on_diff_event(self, payload: RawPresenceDiff): - joins = AsyncRealtimePresence._transform_state(payload["joins"]) - leaves = AsyncRealtimePresence._transform_state(payload["leaves"]) - self.state = self._sync_diff(joins, leaves) - if self.on_sync_callback: - self.on_sync_callback() - - def _sync_state( - self, - new_state: RealtimePresenceState, - ) -> RealtimePresenceState: - joins = {} - leaves = {k: v for k, v in self.state.items() if k not in new_state} - - for key, value in new_state.items(): - current_presences = self.state.get(key, []) - - if len(current_presences) > 0: - new_presence_refs = {presence["presence_ref"] for presence in value} - cur_presence_refs = { - presence["presence_ref"] for presence in current_presences - } - - joined_presences = [ - p for p in value if p["presence_ref"] not in cur_presence_refs - ] - left_presences = [ - p - for p in current_presences - if p["presence_ref"] not in new_presence_refs - ] - - if joined_presences: - joins[key] = joined_presences - if left_presences: - leaves[key] = left_presences - else: - joins[key] = value - - return self._sync_diff(joins, leaves) - - def _sync_diff( - self, joins: RealtimePresenceState, leaves: RealtimePresenceState - ) -> RealtimePresenceState: - for key, new_presences in joins.items(): - current_presences = self.state.get(key, []) - self.state[key] = new_presences - - if len(current_presences) > 0: - joined_presence_refs = { - presence.get("presence_ref") for presence in new_presences - } - cur_presences = list( - presence - for presence in current_presences - if presence.get("presence_ref") not in joined_presence_refs - ) - self.state[key] = cur_presences + self.state[key] - - if self.on_join_callback: - self.on_join_callback(key, current_presences, new_presences) - - for key, left_presences in leaves.items(): - current_presences = self.state.get(key, []) - - if len(current_presences) == 0: - break - - presence_refs_to_remove = { - presence.get("presence_ref") for presence in left_presences - } - current_presences = [ - presence - for presence in current_presences - if presence.get("presence_ref") not in presence_refs_to_remove - ] - self.state[key] = current_presences - - if self.on_leave_callback: - self.on_leave_callback(key, current_presences, left_presences) - - if len(current_presences) == 0: - del self.state[key] - - return self.state - - @staticmethod - def _transform_state( - state: RawPresenceState, - ) -> RealtimePresenceState: - """ - Transform the raw presence state into a standardized RealtimePresenceState format. - - This method processes the input state, which can be either a RawPresenceState or - an already transformed RealtimePresenceState. It handles the conversion of the - Phoenix channel's presence format to our internal representation. - - Args: - state (Union[RawPresenceState, RealtimePresenceState[T]]): The presence state to transform. - - Returns: - RealtimePresenceState[T]: The transformed presence state. - - Example: - Input (RawPresenceState): - { - "user1": { - "metas": [ - {"phx_ref": "ABC123", "user_id": "user1", "status": "online"}, - {"phx_ref": "DEF456", "phx_ref_prev": "ABC123", "user_id": "user1", "status": "away"} - ] - }, - "user2": [{"user_id": "user2", "status": "offline"}] - } - - Output (RealtimePresenceState): - { - "user1": [ - {"presence_ref": "ABC123", "user_id": "user1", "status": "online"}, - {"presence_ref": "DEF456", "user_id": "user1", "status": "away"} - ], - "user2": [{"user_id": "user2", "status": "offline"}] - } - """ - new_state: RealtimePresenceState = {} - for key, presences in state.items(): - new_state[key] = [] - - for presence in presences["metas"]: - if "phx_ref_prev" in presence: - del presence["phx_ref_prev"] - new_presence: Presence = {"presence_ref": presence.pop("phx_ref")} - new_presence.update(presence) - new_state[key].append(new_presence) - return new_state diff --git a/src/realtime/src/realtime/_async/push.py b/src/realtime/src/realtime/_async/push.py deleted file mode 100644 index d252f0c9..00000000 --- a/src/realtime/src/realtime/_async/push.py +++ /dev/null @@ -1,159 +0,0 @@ -from __future__ import annotations - -import asyncio -import logging -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional - -from typing_extensions import assert_never, overload - -from ..message import Message, ReplyPostgresChanges -from ..types import DEFAULT_TIMEOUT, Callback, RealtimeAcknowledgementStatus, _Hook - -if TYPE_CHECKING: - from .channel import AsyncRealtimeChannel - -logger = logging.getLogger(__name__) - - -class AsyncPush: - def __init__( - self, - channel: AsyncRealtimeChannel, - event: str, - payload: Optional[Mapping[str, Any]] = None, - timeout: int = DEFAULT_TIMEOUT, - ): - self.channel = channel - self.event = event - self.payload = payload or {} - self.timeout = timeout - self.rec_hooks: List[_Hook] = [] - self.ref: Optional[str] = None - self.received_resp: Optional[ - tuple[RealtimeAcknowledgementStatus, Dict[str, Any]] - ] = None - self.sent = False - self.timeout_task: Optional[asyncio.Task] = None - self.ok_callbacks: list[Callback[[ReplyPostgresChanges], None]] = [] - self.error_callbacks: list[Callback[[Dict[str, Any]], None]] = [] - self.timeout_callbacks: list[Callback[[], None]] = [] - - async def resend(self): - self.ref = None - self.received_resp = None - self.sent = False - await self.send() - - async def send(self): - if ( - self.received_resp - and self.received_resp[0] == RealtimeAcknowledgementStatus.Timeout - ): - return - - self.ref = self.channel.socket._make_ref() - self.channel.messages_waiting_for_ack[self.ref] = self - self.start_timeout() - self.sent = True - - message = Message( - topic=self.channel.topic, - event=self.event, - ref=self.ref, - payload=self.payload, - ) - await self.channel.socket.send(message) - - def update_payload(self, payload: Dict[str, Any]): - self.payload = {**self.payload, **payload} - - @overload - def receive( - self, - status: Literal[RealtimeAcknowledgementStatus.Ok], - callback: Callback[[ReplyPostgresChanges], None], - ) -> AsyncPush: ... - @overload - def receive( - self, - status: Literal[RealtimeAcknowledgementStatus.Error], - callback: Callback[[Dict[str, Any]], None], - ) -> AsyncPush: ... - - @overload - def receive( - self, - status: Literal[RealtimeAcknowledgementStatus.Timeout], - callback: Callback[[], None], - ) -> AsyncPush: ... - - def receive(self, status, callback) -> AsyncPush: - if (received := self.received_resp) and received[0] == status: - callback(received[1]) - else: - if status == RealtimeAcknowledgementStatus.Ok: - self.ok_callbacks.append(callback) - elif status == RealtimeAcknowledgementStatus.Error: - self.error_callbacks.append(callback) - elif status == RealtimeAcknowledgementStatus.Timeout: - self.timeout_callbacks.append(callback) - else: - assert_never(status) - return self - - def start_timeout(self): - if self.timeout_task: - return - - async def timeout(self): - await asyncio.sleep(self.timeout) - self.trigger(RealtimeAcknowledgementStatus.Timeout, {}) - if self.ref and self.ref in self.channel.messages_waiting_for_ack: - del self.channel.messages_waiting_for_ack[self.ref] - - self.timeout_task = asyncio.create_task(timeout(self)) - - @overload - def trigger( - self, - status: Literal[RealtimeAcknowledgementStatus.Ok], - response: ReplyPostgresChanges, - ): ... - @overload - def trigger( - self, - status: Literal[RealtimeAcknowledgementStatus.Error], - response: Dict[str, Any], - ): ... - @overload - def trigger( - self, - status: Literal[RealtimeAcknowledgementStatus.Timeout], - response: Dict[str, Any], - ): ... - - def trigger(self, status: RealtimeAcknowledgementStatus, response) -> None: - self.received_resp = (status, response) - if status == RealtimeAcknowledgementStatus.Ok: - self._cancel_timeout() - for ok_callback in self.ok_callbacks: - ok_callback(response) - elif status == RealtimeAcknowledgementStatus.Error: - self._cancel_timeout() - for error_callback in self.error_callbacks: - error_callback(response) - elif status == RealtimeAcknowledgementStatus.Timeout: - for timeout_callback in self.timeout_callbacks: - timeout_callback() - else: - assert_never(status) - - def destroy(self): - self._cancel_timeout() - - def _cancel_timeout(self): - if not self.timeout_task: - return - - self.timeout_task.cancel() - self.timeout_task = None diff --git a/src/realtime/src/realtime/_async/timer.py b/src/realtime/src/realtime/_async/timer.py deleted file mode 100644 index eeee4b01..00000000 --- a/src/realtime/src/realtime/_async/timer.py +++ /dev/null @@ -1,40 +0,0 @@ -import asyncio -import logging -from typing import Callable, Optional - -logger = logging.getLogger(__name__) - - -class AsyncTimer: - def __init__(self, callback: Callable, timer_calc: Callable[[int], float]): - self.callback = callback - self.timer_calc = timer_calc - self.timer: Optional[asyncio.Task] = None - self.tries: int = 0 - - def reset(self): - self.tries = 0 - if self.timer and not self.timer.done(): - self.timer.cancel() - self.timer = None - logger.debug( - "AsyncTimer has been reset and any scheduler tasks have been cancelled" - ) - - def schedule_timeout(self): - if self.timer: - self.timer.cancel() - - self.tries += 1 - delay = self.timer_calc(self.tries + 1) - logger.debug(f"Scheduling callback to run after {delay} seconds.") - self.timer = asyncio.create_task(self._run_timer(delay)) - - async def _run_timer(self, delay: float): - try: - await asyncio.sleep(delay) - await self.callback() - except asyncio.CancelledError: - logger.debug("AsyncTimer task was cancelled.") - except Exception as e: - logger.exception(f"Error in AsyncTimer callback: {e}") diff --git a/src/realtime/src/realtime/_sync/__init__.py b/src/realtime/src/realtime/_sync/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/src/realtime/src/realtime/_sync/channel.py b/src/realtime/src/realtime/_sync/channel.py deleted file mode 100644 index 3512f9b7..00000000 --- a/src/realtime/src/realtime/_sync/channel.py +++ /dev/null @@ -1,30 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Optional - -from realtime.types import RealtimeChannelOptions - -if TYPE_CHECKING: - from .client import SyncRealtimeClient - - -class SyncRealtimeChannel: - """ - `Channel` is an abstraction for a topic listener for an existing socket connection. - Each Channel has its own topic and a list of event-callbacks that responds to messages. - Should only be instantiated through `connection.RealtimeClient().channel(topic)`. - """ - - def __init__( - self, - socket: SyncRealtimeClient, - topic: str, - params: Optional[RealtimeChannelOptions] = None, - ) -> None: - """ - Initialize the Channel object. - - :param socket: RealtimeClient object - :param topic: Topic that it subscribes to on the realtime server - :param params: Optional parameters for connection. - """ diff --git a/src/realtime/src/realtime/_sync/client.py b/src/realtime/src/realtime/_sync/client.py deleted file mode 100644 index 634a5343..00000000 --- a/src/realtime/src/realtime/_sync/client.py +++ /dev/null @@ -1,71 +0,0 @@ -from typing import Any, Dict, List, Optional - -from .channel import RealtimeChannelOptions, SyncRealtimeChannel - -NOT_IMPLEMENTED_MESSAGE = "This feature isn't available in the sync client. You can use the realtime feature in the async client only." - - -class SyncRealtimeClient: - def __init__( - self, - url: str, - token: str, - auto_reconnect: bool = True, - params: Optional[Dict[str, Any]] = None, - hb_interval: int = 30, - max_retries: int = 5, - initial_backoff: float = 1.0, - ) -> None: - """ - Initialize a RealtimeClient instance for WebSocket communication. - - :param url: WebSocket URL of the Realtime server. Starts with `ws://` or `wss://`. - Also accepts default Supabase URL: `http://` or `https://`. - :param token: Authentication token for the WebSocket connection. - :param auto_reconnect: If True, automatically attempt to reconnect on disconnection. Defaults to False. - :param params: Optional parameters for the connection. Defaults to an empty dictionary. - :param hb_interval: Interval (in seconds) for sending heartbeat messages to keep the connection alive. Defaults to 30. - :param max_retries: Maximum number of reconnection attempts. Defaults to 5. - :param initial_backoff: Initial backoff time (in seconds) for reconnection attempts. Defaults to 1.0. - """ - - def channel( - self, topic: str, params: Optional[RealtimeChannelOptions] = None - ) -> SyncRealtimeChannel: - """ - :param topic: Initializes a channel and creates a two-way association with the socket - :return: Channel - """ - raise NotImplementedError(NOT_IMPLEMENTED_MESSAGE) - - def get_channels(self) -> List[SyncRealtimeChannel]: - raise NotImplementedError(NOT_IMPLEMENTED_MESSAGE) - - def remove_channel(self, channel: SyncRealtimeChannel) -> None: - """ - Unsubscribes and removes a channel from the socket - :param channel: Channel to remove - :return: None - """ - raise NotImplementedError(NOT_IMPLEMENTED_MESSAGE) - - def remove_all_channels(self) -> None: - """ - Unsubscribes and removes all channels from the socket - :return: None - """ - raise NotImplementedError(NOT_IMPLEMENTED_MESSAGE) - - def set_auth(self, token: Optional[str]) -> None: - """ - Set the authentication token for the connection and update all joined channels. - - This method updates the access token for the current connection and sends the new token - to all joined channels. This is useful for refreshing authentication or changing users. - - Args: - token (Optional[str]): The new authentication token. Can be None to remove authentication. - - Returns: - None - """ diff --git a/src/realtime/src/realtime/_sync/presence.py b/src/realtime/src/realtime/_sync/presence.py deleted file mode 100644 index cc5840b1..00000000 --- a/src/realtime/src/realtime/_sync/presence.py +++ /dev/null @@ -1,12 +0,0 @@ -""" -Defines the RealtimePresence class and its dependencies. -""" - -from typing import Optional - -from ..types import PresenceOpts - - -class SyncRealtimePresence: - def __init__(self, channel, opts: Optional[PresenceOpts] = None): - pass diff --git a/src/realtime/src/realtime/channel.py b/src/realtime/src/realtime/channel.py new file mode 100644 index 00000000..449bee13 --- /dev/null +++ b/src/realtime/src/realtime/channel.py @@ -0,0 +1,311 @@ +from __future__ import annotations + +import asyncio +import json +import logging +from dataclasses import dataclass, field +from typing import ( + TYPE_CHECKING, + Any, + AsyncIterator, + Callable, + Dict, + List, + Mapping, + Optional, + TypeAlias, +) + +import websockets +from supabase_utils.types import JSON +from typing_extensions import assert_never + +from .message import ( + BroadcastMessage, + ChannelCloseMessage, + ChannelErrorMessage, + HeartbeatMessage, + Message, + PostgresChangesMessage, + PostgresChangesPayload, + PostgresRowChange, + PresenceDiffMessage, + PresenceStateMessage, + ReplyMessage, + ReplyPostgresChanges, + ServerMessage, + ServerMessageAdapter, + SuccessReplyMessage, + SuccessSystemPayload, + SystemMessage, +) +from .presence import PresenceEvent, RealtimePresence +from .types import ( + BroadcastPayload, + Callback, + ChannelEvents, + ChannelStates, + PostgresChangesBinding, + PostgresChangesData, + RealtimeAcknowledgementStatus, + RealtimeChannelBroadcastConfig, + RealtimeChannelConfig, + RealtimeChannelOptionsPayload, + RealtimeChannelPresenceConfig, + RealtimePostgresChangesListenEvent, + RealtimePresenceState, + RealtimeSubscribeStates, + ReplayOption, +) + +if TYPE_CHECKING: + from .client import RealtimeClient + +logger = logging.getLogger(__name__) + + +@dataclass +class RealtimeChannelOptions: + broadcast_ack: bool = False + broadcast_self: bool = False + broadcast_replay_since: int | None = None + broadcast_replay_limit: int | None = None + presence_key: str = "" + presence_enabled: bool = False + is_private: bool = False + token: str | None = None + postgres_changes_bindings: list[PostgresChangesBinding] = field( + default_factory=list + ) + + def broadcast( + self, + ack: bool = False, + listen_self: bool = False, + replay_since: int | None = None, + replay_limit: int | None = None, + ) -> RealtimeChannelOptions: + self.broadcast_ack = ack + self.broadcast_self = listen_self + self.broadcast_replay_since = replay_since + self.broadcast_replay_limit = replay_limit + return self + + def presence(self, key: str = "", enabled: bool = False) -> RealtimeChannelOptions: + self.presence_key = key + self.presence_enabled = enabled + return self + + def private(self, private: bool = False) -> RealtimeChannelOptions: + self.is_private = private + return self + + def postgres_changes( + self, + event: RealtimePostgresChangesListenEvent, + table: str | None = None, + schema: str | None = None, + filter: str | None = None, + id: int | None = None, + ) -> RealtimeChannelOptions: + binding = PostgresChangesBinding( + event=event, table=table, schema=schema, filter=filter, id=id + ) + self.postgres_changes_bindings.append(binding) + return self + + def to_payload(self) -> RealtimeChannelOptionsPayload: + replay = ( + ReplayOption( + since=self.broadcast_replay_since, + limit=self.broadcast_replay_limit, + ) + if self.broadcast_replay_since is not None + else None + ) + payload = RealtimeChannelOptionsPayload( + access_token=self.token, + config=RealtimeChannelConfig( + private=self.is_private, + broadcast=RealtimeChannelBroadcastConfig( + ack=self.broadcast_ack, + self=self.broadcast_self, + replay=replay, + ), + presence=RealtimeChannelPresenceConfig( + key=self.presence_key, + enabled=self.presence_enabled, + ), + postgres_changes=[ + binding.as_dict() for binding in self.postgres_changes_bindings + ], + ), + ) + return payload + + +ChannelMessage: TypeAlias = ( + SystemMessage + | BroadcastMessage + | PostgresChangesMessage + | ChannelErrorMessage + | ChannelCloseMessage + | PresenceEvent +) + + +class RealtimeChannel: + """ + Channel is an abstraction for a topic subscription on an existing socket connection. + Each Channel has its own topic and a list of event-callbacks that respond to messages. + Should only be instantiated through `AsyncRealtimeClient.channel(topic)`. + """ + + def __init__( + self, + socket: RealtimeClient, + topic: str, + params: RealtimeChannelOptions | None = None, + ) -> None: + """ + Initialize the Channel object. + + :param socket: RealtimeClient object + :param topic: Topic that it subscribes to on the realtime server + :param params: Optional parameters for connection. + """ + self.socket = socket + self.params = params or RealtimeChannelOptions() + self.topic = topic + self.presence: RealtimePresence = RealtimePresence() + self.state = ChannelStates.CLOSED + self.joined = False + self.join_ref: str | None = None + self.message_stream: asyncio.Queue[Message] = asyncio.Queue() + self.broadcast_endpoint_url = self._broadcast_endpoint_url() + + async def __aenter__(self) -> RealtimeChannel: + await self.subscribe() + return self + + async def __aexit__(self, *args) -> None: + try: + await self.unsubscribe() + except websockets.ConnectionClosed: + # in the case of disconnect, we cannot guarantee that an unsubscribe message can be sent + pass + + async def subscribe(self) -> ReplyMessage: + if self.socket.last_token is not None: + self.params.token = self.socket.last_token + payload = self.params.to_payload() + message = Message( + topic=self.topic, + event=ChannelEvents.join, + payload=payload.model_dump(exclude_none=True), + ref=self.socket._make_ref(), + join_ref=None, + ) + msg = await self.socket.send(message) + logger.info(f"Subscribe reply: {msg!r}") + if msg.payload.status != "ok": + raise Exception( + f"error while subscribing to channel: {msg.payload.response!r}" + ) + self.join_ref = self.socket._make_ref() + self.joined = True + return msg + + @property + async def messages(self) -> AsyncIterator[ChannelMessage]: + while True: + get_task = asyncio.create_task(self.message_stream.get()) + message = await self.socket.get_message_or_raise(get_task) + server_message = ServerMessageAdapter.validate_python(message) + if isinstance(server_message, PresenceStateMessage): + for presence_event in self.presence._on_state_event( + server_message.payload + ): + yield presence_event + elif isinstance(server_message, PresenceDiffMessage): + for presence_event in self.presence._on_diff_event( + server_message.payload + ): + yield presence_event + else: + yield server_message + + async def unsubscribe(self) -> None: + await self.send_event(ChannelEvents.leave, payload={}) + + # Presence methods + async def track(self, user_status: Dict[str, Any]) -> ReplyMessage: + """ + Track presence status for the current user. + + :param user_status: Dictionary containing the user's presence information + """ + return await self.send_presence("track", user_status) + + async def untrack(self) -> ReplyMessage: + """ + Stop tracking presence for the current user. + """ + return await self.send_presence("untrack", {}) + + def presence_state(self) -> RealtimePresenceState: + """ + Get the current state of presence on this channel. + + :return: Dictionary mapping presence keys to lists of presence payloads + """ + return self.presence.state + + async def send_event( + self, event: ChannelEvents, payload: Mapping[str, JSON] + ) -> ReplyMessage: + message = Message( + topic=self.topic, + event=event, + ref=self.socket._make_ref(), + join_ref=self.join_ref, + payload=payload, + ) + return await self.socket.send(message) + + async def send_event_no_wait( + self, event: ChannelEvents, payload: Mapping[str, JSON] + ) -> None: + message = Message( + topic=self.topic, + event=event, + ref=self.socket._make_ref(), + join_ref=self.join_ref, + payload=payload, + ) + return await self.socket.send_no_wait(message) + + # Broadcast methods + async def send_broadcast(self, event: str, data: JSON) -> ReplyMessage | None: + """ + Send a broadcast message through this channel. + + :param event: The name of the broadcast event + :param data: The payload to broadcast + """ + payload = {"type": "broadcast", "event": event, "payload": data} + if self.params.broadcast_ack: + return await self.send_event(ChannelEvents.broadcast, payload) + else: + await self.send_event_no_wait(ChannelEvents.broadcast, payload) + return None + + # Internal methods + + def _broadcast_endpoint_url(self): + return self.socket.http_endpoint.joinpath("api", "broadcast") + + async def send_presence(self, event: str, data: JSON) -> ReplyMessage: + return await self.send_event( + ChannelEvents.presence, {"event": event, "payload": data} + ) diff --git a/src/realtime/src/realtime/client.py b/src/realtime/src/realtime/client.py new file mode 100644 index 00000000..cab5937a --- /dev/null +++ b/src/realtime/src/realtime/client.py @@ -0,0 +1,344 @@ +import asyncio +import json +import logging +import re +import sys +from functools import wraps +from types import TracebackType +from typing import ( + Any, + AsyncIterator, + Awaitable, + Callable, + Dict, + Generator, + List, + Optional, + Union, +) +from urllib.parse import urlencode, urlparse, urlunparse +from warnings import warn + +import websockets +from pydantic import ValidationError +from websockets.asyncio.client import ClientConnection +from yarl import URL +from yarl._query import get_str_query + +from .channel import RealtimeChannel, RealtimeChannelOptions +from .exceptions import NotConnectedError +from .message import Message, ReplyMessage, ServerMessage +from .types import ( + DEFAULT_HEARTBEAT_INTERVAL, + DEFAULT_TIMEOUT, + PHOENIX_CHANNEL, + VSN, + ChannelEvents, + ChannelStates, +) + +logger = logging.getLogger(__name__) + + +def normalize_url(url: URL, token: str | None) -> URL: + if url.scheme not in {"ws", "wss", "http", "https"}: + raise ValueError("url must be a valid WebSocket URL or HTTP URL string") + url = url.joinpath("websocket") + if token: + url = url.with_query(**url.query, apikey=token) + if url.scheme in {"ws", "wss"}: + return url + if url.scheme == "https": + return url.with_scheme("wss") + return url.with_scheme("ws") + + +class connect_once: + def __init__( + self, + url: str, + token_callback: Callable[[], Awaitable[str | None]] | None = None, + hb_interval: int = DEFAULT_HEARTBEAT_INTERVAL, + hb_timeout: int = 5, + ) -> None: + self.url = URL(url) + self.get_access_token = token_callback + self.listen_task: asyncio.Task | None = None + self.heartbeat_task: asyncio.Task | None = None + self.hb_interval = hb_interval + self.hb_timeout = hb_timeout + + def __await__(self) -> Generator[None, None, "RealtimeClient"]: + return self.__await_impl__().__await__() + + async def __await_impl__(self) -> "RealtimeClient": + if self.get_access_token: + token = await self.get_access_token() + url = normalize_url(self.url, token) + else: + url = normalize_url(self.url, None) + self.ws = await websockets.connect(str(url)) + self.client = RealtimeClient( + self.url, + self.ws, + self.get_access_token, + hb_interval=self.hb_interval, + hb_timeout=self.hb_timeout, + ) + self.listen_task = asyncio.create_task(self.client.listen()) + self.heartbeat_task = asyncio.create_task(self.client.heartbeat()) + return self.client + + async def __aenter__(self) -> "RealtimeClient": + return await self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + assert self.listen_task, "Connection should never exit without listen_task" + self.listen_task.cancel() + assert self.heartbeat_task, "Connection should never exit without hearbeat_task" + self.heartbeat_task.cancel() + await self.client.close() + + +class MaxRetriesExceeded(Exception): + pass + + +def exponential_delay(max_retries: int, retry_count: int) -> float | None: + if retry_count == max_retries: + return None + else: + return float(2**retry_count) + + +async def automatically_reconnect( + url: str, + token_callback: Callable[[], Awaitable[str]], + client_handler: Callable[["RealtimeClient"], Awaitable[None]], + max_retries: int = 3, +) -> None: + retries = 0 + while True: + try: + async with connect_once(url, token_callback) as client: + retries = 0 + await client_handler(client) + break + except websockets.ConnectionClosed: + continue + except (TimeoutError, OSError): + retries += 1 + delay = exponential_delay(max_retries, retries) + if delay is not None: + await asyncio.sleep(delay) + else: + raise MaxRetriesExceeded(retries) + + +class RealtimeClient: + def __init__( + self, + url: URL, + ws_connection: ClientConnection, + token_callback: Callable[[], Awaitable[str | None]] | None = None, + hb_interval: int = DEFAULT_HEARTBEAT_INTERVAL, + hb_timeout: int = 5, + ) -> None: + self.url = url + self.last_token: str | None = None + self.get_access_token = token_callback + self.http_endpoint = url.with_scheme("https").with_path("") + + self._ws_connection: ClientConnection = ws_connection + + self.ref = 0 + self.channels: Dict[str, RealtimeChannel] = {} + self.message_refs: Dict[str, asyncio.Future[Message]] = {} + self._connection_failure: asyncio.Future[Message] = asyncio.Future() + + self.hb_interval = hb_interval + self.heartbeat_timeout = hb_timeout + + async def listen(self) -> None: + """ + An infinite loop that keeps listening. + :return: None + """ + try: + while True: + msg = await self._ws_connection.recv(decode=False) + try: + message = Message.model_validate_json(msg) + except ValidationError as e: + logger.error(f"Unrecognized message format {msg!r}\n{e}") + continue + if message.ref is not None: + self.message_refs[message.ref].set_result(message) + elif channel := self.channels.get(message.topic): + await channel.message_stream.put(message) + else: + logger.info(f"Ignoring message: {message!r}") + except Exception as exc: + if ( + not self._connection_failure.done() + or self._connection_failure.cancelled() + ): + self._connection_failure.set_exception(exc) + + async def close(self) -> None: + """ + Close the WebSocket connection. + + Returns: + None + + Raises: + NotConnectedError: If the connection is not established when this method is called. + """ + + await self._ws_connection.close() + + async def get_message_or_raise(self, fut: asyncio.Future[Message]) -> Message: + done, pending = await asyncio.wait( + (fut, self._connection_failure), return_when=asyncio.FIRST_COMPLETED + ) + if self._connection_failure in done: + fut.cancel() + # raise the exception + await self._connection_failure + return fut.result() + + async def heartbeat(self) -> None: + try: + while True: + data = Message( + topic=PHOENIX_CHANNEL, + event=ChannelEvents.heartbeat, + payload={}, + ref=self._make_ref(), + ) + logger.info("sending heartbeat") + res = await asyncio.wait_for( + self.send(data), timeout=self.heartbeat_timeout + ) + logger.info(f"heartbeat response: {res!r}") + await self.set_auth() + await asyncio.sleep(max(self.hb_interval, 15)) + except asyncio.CancelledError: + return + except (websockets.ConnectionClosed, asyncio.TimeoutError) as exc: + if ( + not self._connection_failure.done() + or self._connection_failure.cancelled() + ): + self._connection_failure.set_exception(exc) + + def channel( + self, topic: str, params: RealtimeChannelOptions | None = None + ) -> RealtimeChannel: + """ + Initialize a channel and create a two-way association with the socket. + + :param topic: The topic to subscribe to + :param params: Optional channel parameters + :return: AsyncRealtimeChannel instance + """ + topic = f"realtime:{topic}" + chan = RealtimeChannel(self, topic, params) + self.channels[topic] = chan + return chan + + def get_channels(self) -> List[RealtimeChannel]: + return list(self.channels.values()) + + def _remove_channel(self, channel: RealtimeChannel) -> None: + del self.channels[channel.topic] + + async def remove_channel(self, channel: RealtimeChannel) -> None: + if chan := self.channels.get(channel.topic): + await chan.unsubscribe() + self._remove_channel(channel) + + async def remove_all_channels(self) -> None: + """ + Unsubscribes and removes all channels from the socket + :return: None + """ + tasks = [] + for channel in self.channels.values(): + tasks.append(asyncio.create_task(channel.unsubscribe())) + await asyncio.wait(tasks) + + async def set_auth(self, token: str | None = None) -> None: + """ + Set the authentication token for the connection and update all joined channels. + + This method updates the access token for the current connection and sends the new token + to all joined channels. This is useful for refreshing authentication or changing users. + + Args: + token (Optional[str]): The new authentication token. Can be None to remove authentication. + + Returns: + None + """ + if token is None: + if self.get_access_token is not None and ( + new_token := await self.get_access_token() + ): + token = new_token + else: + logger.warning( + "tried setting auth but no token callback was registered" + ) + return + if token == self.last_token: + logger.info("generated access token is still the same, skipping") + return + + tasks = [] + for channel in self.channels.values(): + if channel.joined and channel.state == ChannelStates.JOINED: + task = asyncio.create_task( + channel.send_event( + ChannelEvents.access_token, {"access_token": token} + ) + ) + tasks.append(task) + if tasks: + await asyncio.wait(tasks) + self.last_token = token + + def _make_ref(self) -> str: + self.ref += 1 + return f"{self.ref}" + + async def send(self, message: Message) -> ReplyMessage: + assert message.ref, "Cannot wait for reponse on a message without ref" + + message_str = message.model_dump_json() + logger.debug(f"sending: {message_str}") + + receive: asyncio.Future[Message] = asyncio.Future() + self.message_refs[message.ref] = receive + await self._ws_connection.send(message_str) + logger.debug(f"waiting for ref='{message.ref}'") + + response = await self.get_message_or_raise(receive) + del self.message_refs[message.ref] + + return ReplyMessage.model_validate(response) + + async def send_no_wait(self, message: Message) -> None: + message_str = message.model_dump_json() + logger.debug(f"sending: {message_str}") + await self._ws_connection.send(message_str) + + def endpoint_url(self) -> str: + query = self.url.with_query(**self.url.query, vsn=VSN) + return str(query) diff --git a/src/realtime/src/realtime/exceptions.py b/src/realtime/src/realtime/exceptions.py index 39992048..7fdd4705 100644 --- a/src/realtime/src/realtime/exceptions.py +++ b/src/realtime/src/realtime/exceptions.py @@ -18,7 +18,7 @@ class AuthorizationError(Exception): Raised when there is an authorization failure for private channels """ - def __init__(self, message: Optional[str] = None): + def __init__(self, message: str | None = None): self.message: str = message or "Authorization failed for private channel" def __str__(self): diff --git a/src/realtime/src/realtime/message.py b/src/realtime/src/realtime/message.py index eed31962..9f123b63 100644 --- a/src/realtime/src/realtime/message.py +++ b/src/realtime/src/realtime/message.py @@ -1,16 +1,16 @@ -from typing import Any, List, Literal, Mapping, Optional, Union +from typing import Annotated, Any, List, Literal, Mapping, Optional, Union -from pydantic import BaseModel, Field, TypeAdapter +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter +from supabase_utils.types import JSON from typing_extensions import TypeAlias, TypedDict from .types import ( BroadcastPayload, ChannelEvents, - PostgresChangesData, PostgresChangesPayload, RawPresenceDiff, RawPresenceState, - RealtimeChannelOptions, + RealtimeChannelOptionsPayload, RealtimePostgresChangesListenEvent, ) @@ -21,29 +21,29 @@ class Message(BaseModel): """ event: str - payload: Mapping[str, Any] + payload: Mapping[str, JSON] topic: str - ref: Optional[str] = None - join_ref: Optional[str] = None + ref: str | None = None + join_ref: str | None = None class JoinMessage(BaseModel): event: Literal[ChannelEvents.join] topic: str ref: str - payload: RealtimeChannelOptions + payload: RealtimeChannelOptionsPayload class PostgresRowChange(BaseModel): id: int event: RealtimePostgresChangesListenEvent - table: Optional[str] = None - schema_: Optional[str] = Field(alias="schema", default=None) - filter: Optional[str] = None + table: str | None = None + schema_: str | None = Field(alias="schema", default=None) + filter: str | None = None class ReplyPostgresChanges(BaseModel): - postgres_changes: Optional[List[PostgresRowChange]] = None + postgres_changes: List[PostgresRowChange] | None = None class SuccessReplyMessage(BaseModel): @@ -53,14 +53,16 @@ class SuccessReplyMessage(BaseModel): class ErrorReplyMessage(BaseModel): status: Literal["error"] - response: dict[str, Any] # TODO: what goes in here? + response: dict[str, JSON] # TODO: what goes in here? class ReplyMessage(BaseModel): + model_config = ConfigDict(from_attributes=True) + event: Literal[ChannelEvents.reply] topic: str - payload: Union[SuccessReplyMessage, ErrorReplyMessage] - ref: Optional[str] + payload: SuccessReplyMessage | ErrorReplyMessage + ref: str | None class SuccessSystemPayload(BaseModel): @@ -78,9 +80,11 @@ class ErrorSystemPayload(BaseModel): class SystemMessage(BaseModel): + model_config = ConfigDict(from_attributes=True) + event: Literal[ChannelEvents.system] topic: str - payload: Union[SuccessSystemPayload, ErrorSystemPayload] + payload: SuccessSystemPayload | ErrorSystemPayload ref: Literal[None] @@ -107,6 +111,8 @@ class AccessTokenMessage(BaseModel): class PostgresChangesMessage(BaseModel): + model_config = ConfigDict(from_attributes=True) + event: Literal[ChannelEvents.postgres_changes] topic: str payload: PostgresChangesPayload @@ -114,6 +120,8 @@ class PostgresChangesMessage(BaseModel): class BroadcastMessage(BaseModel): + model_config = ConfigDict(from_attributes=True) + event: Literal[ChannelEvents.broadcast] topic: str payload: BroadcastPayload @@ -123,11 +131,13 @@ class BroadcastMessage(BaseModel): class PresenceMessage(BaseModel): event: Literal[ChannelEvents.presence] topic: str - payload: dict[str, Any] + payload: dict[str, JSON] ref: Literal[None] class PresenceStateMessage(BaseModel): + model_config = ConfigDict(from_attributes=True) + event: Literal[ChannelEvents.presence_state] topic: str payload: RawPresenceState @@ -135,6 +145,8 @@ class PresenceStateMessage(BaseModel): class PresenceDiffMessage(BaseModel): + model_config = ConfigDict(from_attributes=True) + event: Literal[ChannelEvents.presence_diff] topic: str payload: RawPresenceDiff @@ -142,35 +154,49 @@ class PresenceDiffMessage(BaseModel): class ChannelErrorMessage(BaseModel): + model_config = ConfigDict(from_attributes=True) + event: Literal[ChannelEvents.error] topic: str - payload: dict[str, Any] - ref: Optional[str] + payload: dict[str, JSON] + ref: str | None class ChannelCloseMessage(BaseModel): + model_config = ConfigDict(from_attributes=True) + event: Literal[ChannelEvents.close] topic: str - payload: dict[str, Any] - ref: Optional[str] - - -ServerMessage: TypeAlias = Union[ - SystemMessage, - ReplyMessage, - HeartbeatMessage, - BroadcastMessage, - PresenceStateMessage, - PresenceDiffMessage, - PostgresChangesMessage, - ChannelErrorMessage, - ChannelCloseMessage, + payload: dict[str, JSON] + ref: str | None + + +class ChannelLeaveMessage(BaseModel): + model_config = ConfigDict(from_attributes=True) + + event: Literal[ChannelEvents.leave] + topic: str + payload: dict[str, JSON] + ref: str + + +ServerMessage: TypeAlias = Annotated[ + SystemMessage + | BroadcastMessage + | PresenceStateMessage + | PresenceDiffMessage + | PostgresChangesMessage + | ChannelErrorMessage + | ChannelCloseMessage, + Field(discriminator="event"), ] ServerMessageAdapter: TypeAdapter[ServerMessage] = TypeAdapter(ServerMessage) -ClientMessage: TypeAlias = Union[ - JoinMessage, - HeartbeatMessage, - BroadcastMessage, - PresenceMessage, - AccessTokenMessage, -] + +ClientMessage: TypeAlias = ( + JoinMessage + | HeartbeatMessage + | BroadcastMessage + | PresenceMessage + | AccessTokenMessage + | ChannelLeaveMessage +) diff --git a/src/realtime/src/realtime/presence.py b/src/realtime/src/realtime/presence.py new file mode 100644 index 00000000..e9545a26 --- /dev/null +++ b/src/realtime/src/realtime/presence.py @@ -0,0 +1,180 @@ +""" +Defines the RealtimePresence class and its dependencies. +""" + +import logging +from dataclasses import dataclass +from typing import Any, Callable, Dict, Generator, List, Optional, Union + +from .types import ( + Presence, + PresenceDiff, + PresenceEvents, + PresenceOpts, + RawPresenceDiff, + RawPresenceState, + RealtimePresenceState, +) + +logger = logging.getLogger(__name__) + + +def _transform_state( + state: RawPresenceState, +) -> RealtimePresenceState: + """ + Transform the raw presence state into a standardized RealtimePresenceState format. + + This method processes the input state, which can be either a RawPresenceState or + an already transformed RealtimePresenceState. It handles the conversion of the + Phoenix channel's presence format to our internal representation. + + Args: + state RawPresenceState: The presence state to transform. + + Returns: + RealtimePresenceState: The transformed presence state. + + Example: + Input (RawPresenceState): + { + "user1": { + "metas": [ + {"phx_ref": "ABC123", "user_id": "user1", "status": "online"}, + {"phx_ref": "DEF456", "phx_ref_prev": "ABC123", "user_id": "user1", "status": "away"} + ] + }, + "user2": [{"user_id": "user2", "status": "offline"}] + } + + Output (RealtimePresenceState): + { + "user1": [ + {"presence_ref": "ABC123", "user_id": "user1", "status": "online"}, + {"presence_ref": "DEF456", "user_id": "user1", "status": "away"} + ], + "user2": [{"user_id": "user2", "status": "offline"}] + } + """ + new_state: RealtimePresenceState = {} + for key, presences in state.items(): + new_state[key] = [] + + for presence in presences["metas"]: + if "phx_ref_prev" in presence: + del presence["phx_ref_prev"] + new_presence: Presence = {"presence_ref": presence.pop("phx_ref")} + new_presence.update(presence) + new_state[key].append(new_presence) + return new_state + + +@dataclass +class PresenceJoin: + key: str + current_presences: List[Presence] + new_presences: List[Presence] + + +@dataclass +class PresenceLeave: + key: str + current_presences: List[Presence] + left_presences: List[Presence] + + +PresenceEvent = PresenceJoin | PresenceLeave + + +class RealtimePresence: + def __init__(self) -> None: + self.state: RealtimePresenceState = {} + + def _on_state_event( + self, payload: RawPresenceState + ) -> Generator[PresenceEvent, None, None]: + state = _transform_state(payload) + self.state = yield from self._sync_state(state) + + def _on_diff_event( + self, payload: RawPresenceDiff + ) -> Generator[PresenceEvent, None, None]: + joins = _transform_state(payload["joins"]) + leaves = _transform_state(payload["leaves"]) + self.state = yield from self._sync_diff(joins, leaves) + + def _sync_state( + self, + new_state: RealtimePresenceState, + ) -> Generator[PresenceEvent, None, RealtimePresenceState]: + joins = {} + leaves = {k: v for k, v in self.state.items() if k not in new_state} + + for key, value in new_state.items(): + current_presences = self.state.get(key, []) + + if len(current_presences) > 0: + new_presence_refs = {presence["presence_ref"] for presence in value} + cur_presence_refs = { + presence["presence_ref"] for presence in current_presences + } + + joined_presences = [ + p for p in value if p["presence_ref"] not in cur_presence_refs + ] + left_presences = [ + p + for p in current_presences + if p["presence_ref"] not in new_presence_refs + ] + + if joined_presences: + joins[key] = joined_presences + if left_presences: + leaves[key] = left_presences + else: + joins[key] = value + + res = yield from self._sync_diff(joins, leaves) + return res + + def _sync_diff( + self, joins: RealtimePresenceState, leaves: RealtimePresenceState + ) -> Generator[PresenceEvent, None, RealtimePresenceState]: + for key, new_presences in joins.items(): + current_presences = self.state.get(key, []) + self.state[key] = new_presences + + if len(current_presences) > 0: + joined_presence_refs = { + presence["presence_ref"] for presence in new_presences + } + cur_presences = list( + presence + for presence in current_presences + if presence["presence_ref"] not in joined_presence_refs + ) + self.state[key].extend(cur_presences) + yield PresenceJoin(key, current_presences, new_presences) + + for key, left_presences in leaves.items(): + current_presences = self.state.get(key, []) + + if len(current_presences) == 0: + break + + presence_refs_to_remove = { + presence["presence_ref"] for presence in left_presences + } + current_presences = [ + presence + for presence in current_presences + if presence["presence_ref"] not in presence_refs_to_remove + ] + self.state[key] = current_presences + + if len(current_presences) == 0: + del self.state[key] + yield PresenceLeave(key, current_presences, left_presences) + + return self.state diff --git a/src/realtime/src/realtime/transformers.py b/src/realtime/src/realtime/transformers.py deleted file mode 100644 index bd541377..00000000 --- a/src/realtime/src/realtime/transformers.py +++ /dev/null @@ -1,9 +0,0 @@ -import re - - -def http_endpoint_url(socket_url: str) -> str: - url = re.sub(r"^ws", "http", socket_url, flags=re.IGNORECASE) - url = re.sub( - r"(\/socket\/websocket|\/socket|\/websocket)\/?$", "", url, flags=re.IGNORECASE - ) - return re.sub(r"\/+$", "", url) diff --git a/src/realtime/src/realtime/types.py b/src/realtime/src/realtime/types.py index 89e7b8fa..de03d002 100644 --- a/src/realtime/src/realtime/types.py +++ b/src/realtime/src/realtime/types.py @@ -96,12 +96,39 @@ class PostgresChangesData(TypedDict): table: str commit_timestamp: str type: RealtimePostgresChangesListenEvent - errors: Optional[str] + errors: str | None columns: List[PostgresChangesColumn] - record: NotRequired[Optional[dict[str, Any]]] + record: NotRequired[dict[str, Any] | None] old_record: NotRequired[dict[str, Any]] # todo: improve this +class PostgresChangesBindingDict(TypedDict): + event: RealtimePostgresChangesListenEvent + table: NotRequired[str] + schema: NotRequired[str] + filter: NotRequired[str] + id: NotRequired[int] + + +@dataclass +class PostgresChangesBinding: + event: RealtimePostgresChangesListenEvent + table: str | None + schema: str | None + filter: str | None + id: int | None = None + + def as_dict(self) -> PostgresChangesBindingDict: + binding: PostgresChangesBindingDict = {"event": self.event} + if self.schema is not None: + binding["schema"] = self.schema + if self.table is not None: + binding["table"] = self.table + if self.filter is not None: + binding["filter"] = self.filter + return binding + + class PostgresChangesPayload(TypedDict): data: PostgresChangesData ids: List[int] @@ -115,54 +142,10 @@ class BroadcastMeta(TypedDict, total=False): class BroadcastPayload(TypedDict): event: str payload: dict[str, Any] + type: NotRequired[str] meta: NotRequired[BroadcastMeta] -@dataclass(frozen=True) -class BroadcastCallback: - callback: Callable[[BroadcastPayload], None] - event: str - - def __call__(self, payload: BroadcastPayload) -> None: - if self.event == payload["event"]: - return self.callback(payload) - - -@dataclass -class PostgresChangesCallback: - callback: Callable[[PostgresChangesPayload], None] - event: RealtimePostgresChangesListenEvent - table: Optional[str] - schema: Optional[str] - filter: Optional[str] - id: Optional[int] = None - - def __call__(self, payload: PostgresChangesPayload) -> None: - event_matches = ( - self.event == payload["data"]["type"] - or self.event == RealtimePostgresChangesListenEvent.All - ) - if self.id and self.id in payload["ids"] and event_matches: - return self.callback(payload) - - @property - def binding_filter(self) -> dict[str, Optional[str]]: - binding: dict[str, Optional[str]] = {"event": self.event} - if self.schema: - binding["schema"] = self.schema - if self.table: - binding["table"] = self.table - if self.filter: - binding["filter"] = self.filter - return binding - - -class _Hook: - def __init__(self, status: str, callback: Callback[[Dict[str, Any]], None]): - self.status = status - self.callback = callback - - @with_config(ConfigDict(extra="allow")) class Presence(TypedDict): presence_ref: str @@ -180,30 +163,32 @@ def __init__(self, events: PresenceEvents): # TypedDicts -class ReplayOption(TypedDict, total=False): +class ReplayOption(BaseModel): since: int - limit: int + limit: int | None -class RealtimeChannelBroadcastConfig(TypedDict, total=False): +class RealtimeChannelBroadcastConfig(BaseModel): ack: bool self: bool - replay: ReplayOption + replay: ReplayOption | None = None -class RealtimeChannelPresenceConfig(TypedDict): +class RealtimeChannelPresenceConfig(BaseModel): key: str enabled: bool -class RealtimeChannelConfig(TypedDict): - broadcast: Optional[RealtimeChannelBroadcastConfig] - presence: Optional[RealtimeChannelPresenceConfig] +class RealtimeChannelConfig(BaseModel): + broadcast: RealtimeChannelBroadcastConfig + presence: RealtimeChannelPresenceConfig + postgres_changes: list[PostgresChangesBindingDict] private: bool -class RealtimeChannelOptions(TypedDict): - config: NotRequired[RealtimeChannelConfig] +class RealtimeChannelOptionsPayload(BaseModel): + config: RealtimeChannelConfig + access_token: str | None = None @with_config(ConfigDict(extra="allow")) @@ -216,9 +201,6 @@ class RawPresenceStateEntry(TypedDict): metas: List[PresenceMeta] -# Custom types -PresenceOnJoinCallback = Callable[[str, List[Any], List[Presence]], None] -PresenceOnLeaveCallback = Callable[[str, List[Any], List[Presence]], None] RealtimePresenceState = Dict[str, List[Presence]] RawPresenceState = Dict[str, RawPresenceStateEntry] diff --git a/src/realtime/src/realtime/utils.py b/src/realtime/src/realtime/utils.py deleted file mode 100644 index 67b22018..00000000 --- a/src/realtime/src/realtime/utils.py +++ /dev/null @@ -1,5 +0,0 @@ -from urllib.parse import urlparse - - -def is_ws_url(url: str) -> bool: - return urlparse(url).scheme in {"wss", "ws", "http", "https"} diff --git a/src/realtime/tests/test_connection.py b/src/realtime/tests/test_connection.py index 1d7a2cff..3e987b45 100644 --- a/src/realtime/tests/test_connection.py +++ b/src/realtime/tests/test_connection.py @@ -1,37 +1,51 @@ import asyncio import datetime +import json import os +from typing import AsyncIterator import aiohttp import pytest from dotenv import load_dotenv from pydantic import BaseModel from websockets import broadcast +from yarl import URL from realtime import ( - AsyncRealtimeChannel, - AsyncRealtimeClient, + RealtimeChannel, + RealtimeClient, RealtimePostgresChangesListenEvent, RealtimeSubscribeStates, ) -from realtime.message import Message +from realtime.channel import RealtimeChannelOptions +from realtime.client import connect_once +from realtime.message import ( + BroadcastMessage, + Message, + PostgresChangesMessage, + SystemMessage, +) from realtime.types import DEFAULT_HEARTBEAT_INTERVAL, DEFAULT_TIMEOUT, ChannelEvents load_dotenv() -URL = os.getenv("SUPABASE_URL") or "http://127.0.0.1:54321" +SUPABASE_URL = os.getenv("SUPABASE_URL") or "http://127.0.0.1:54321" PUBLISHABLE_KEY = ( os.getenv("SUPABASE_PUBLISHABLE_KEY") or "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZS1kZW1vIiwicm9sZSI6ImFub24iLCJleHAiOjE5ODM4MTI5OTZ9.CRXP1A7WOeoJeXxjNni43kdQwgnWNReilDMblYTn_I0" ) +async def get_token() -> str: + return PUBLISHABLE_KEY + + @pytest.fixture -def socket() -> AsyncRealtimeClient: - url = f"{URL}/realtime/v1" - key = PUBLISHABLE_KEY - return AsyncRealtimeClient(url, key) +async def socket() -> AsyncIterator[RealtimeClient]: + url = f"{SUPABASE_URL}/realtime/v1" + async with connect_once(url, get_token) as client: + yield client class SignupMessageResponse(BaseModel): @@ -43,7 +57,7 @@ class SignupMessageResponse(BaseModel): async def access_token() -> str: - url = f"{URL}/auth/v1/signup" + url = f"{SUPABASE_URL}/auth/v1/signup" headers = {"apikey": PUBLISHABLE_KEY, "Content-Type": "application/json"} data = { "email": os.getenv("SUPABASE_TEST_EMAIL") @@ -65,257 +79,161 @@ async def access_token() -> str: ) -def test_init_client(): - client = AsyncRealtimeClient(URL, PUBLISHABLE_KEY) +async def realtime_send( + token: str, topic: str, event: str, payload: dict[str, str] +) -> None: + """ + Trigger a `realtime.send` through PostgREST's rpc endpoint. + This is done through a public function created in the migration, + `send_realtime`, that internally calls the unexposed counterpart. + """ + url = f"{SUPABASE_URL}/rest/v1/rpc/send_realtime" + headers = { + "apikey": PUBLISHABLE_KEY, + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + } + data = { + "payload": payload, + "topic": topic, + "event": event, + "private": True, + } - assert client is not None - assert client.url.startswith("ws://") or client.url.startswith("wss://") - assert "/websocket" in client.url - assert client.url.split("apikey=")[1] == PUBLISHABLE_KEY - assert client.auto_reconnect is True - assert client.params == {} - assert client.hb_interval == DEFAULT_HEARTBEAT_INTERVAL - assert client.max_retries == 5 - assert client.initial_backoff == 1.0 - assert client.timeout == DEFAULT_TIMEOUT + async with aiohttp.ClientSession() as session: + async with session.post(url, headers=headers, json=data) as response: + response.raise_for_status() @pytest.mark.asyncio -async def test_broadcast_events(socket: AsyncRealtimeClient): - await socket.connect() - - channel = socket.channel( - "test-broadcast", - params={ - "config": { - "broadcast": {"self": True, "ack": True}, - "presence": {"enabled": True, "key": ""}, - "private": False, - } - }, - ) - received_events = [] - - semaphore = asyncio.Semaphore(0) - - def broadcast_callback(payload): - received_events.append(payload) - semaphore.release() - - subscribe_event = asyncio.Event() - await channel.on_broadcast("test-event", broadcast_callback).subscribe( - lambda state, _: ( - subscribe_event.set() - if state == RealtimeSubscribeStates.SUBSCRIBED - else None - ) +async def test_broadcast_events(socket: RealtimeClient): + channel_options = ( + RealtimeChannelOptions() + .broadcast(listen_self=True, ack=True) + .presence(enabled=True) ) - - await asyncio.wait_for(subscribe_event.wait(), 5) - - # Send 3 broadcast events - for i in range(3): - await channel.send_broadcast("test-event", {"message": f"Event {i + 1}"}) - await asyncio.wait_for(semaphore.acquire(), 5) - - assert len(received_events) == 3 - assert received_events[0]["payload"]["message"] == "Event 1" - assert received_events[1]["payload"]["message"] == "Event 2" - assert received_events[2]["payload"]["message"] == "Event 3" - - await socket.close() + async with socket.channel( + "test-broadcast", + params=channel_options, + ) as chan: + # Send 3 broadcast events + for i in range(3): + await chan.send_broadcast("test-event", {"message": f"Event {i + 1}"}) + + message_stream = chan.messages + msg1 = await message_stream.__anext__() + assert isinstance(msg1, BroadcastMessage) + msg2 = await message_stream.__anext__() + assert isinstance(msg2, BroadcastMessage) + msg3 = await message_stream.__anext__() + assert isinstance(msg3, BroadcastMessage) + assert msg1.payload["event"] == "test-event" + assert msg2.payload["event"] == "test-event" + assert msg3.payload["event"] == "test-event" + assert msg1.payload["payload"]["message"] == "Event 1" + assert msg2.payload["payload"]["message"] == "Event 2" + assert msg3.payload["payload"]["message"] == "Event 3" @pytest.mark.asyncio -async def test_postgrest_changes(socket: AsyncRealtimeClient): +async def test_postgrest_changes(socket: RealtimeClient): token = await access_token() - - await socket.connect() - await socket.set_auth(token) - channel: AsyncRealtimeChannel = socket.channel("test-postgres-changes") - received_events: dict[str, list[dict]] = { - "all": [], - "insert": [], - "update": [], - "delete": [], - } - - def all_changes_callback(payload): - received_events["all"].append(payload) - - insert_event = asyncio.Event() - - def insert_callback(payload): - received_events["insert"].append(payload) - insert_event.set() - - update_event = asyncio.Event() - - def update_callback(payload): - received_events["update"].append(payload) - update_event.set() - - delete_event = asyncio.Event() - - def delete_callback(payload): - received_events["delete"].append(payload) - delete_event.set() - - subscribed_event = asyncio.Event() - system_event = asyncio.Event() + channel_options = ( + RealtimeChannelOptions() + .postgres_changes(RealtimePostgresChangesListenEvent.All, table="todos") + .postgres_changes(RealtimePostgresChangesListenEvent.Insert, table="todos") + .postgres_changes(RealtimePostgresChangesListenEvent.Update, table="todos") + .postgres_changes(RealtimePostgresChangesListenEvent.Delete, table="todos") + ) - await ( - channel.on_postgres_changes( - RealtimePostgresChangesListenEvent.All, all_changes_callback, table="todos" - ) - .on_postgres_changes( - RealtimePostgresChangesListenEvent.Insert, insert_callback, table="todos" + async with socket.channel("test-postgres-changes", channel_options) as chan: + message_stream = chan.messages + system_msg = await message_stream.__anext__() + assert isinstance(system_msg, SystemMessage) + assert system_msg.payload.status == "ok" + print("received system message") + created_todo_id = await create_todo( + token, {"description": "Test todo", "is_completed": False} ) - .on_postgres_changes( - RealtimePostgresChangesListenEvent.Update, update_callback, table="todos" + print("created todo") + insert = await message_stream.__anext__() + assert isinstance(insert, PostgresChangesMessage) + print("received insert") + await update_todo( + token, + created_todo_id, + {"description": "Updated todo", "is_completed": True}, ) - .on_postgres_changes( - RealtimePostgresChangesListenEvent.Delete, delete_callback, table="todos" - ) - .on_system(lambda _: system_event.set()) - .subscribe( - lambda state, _: ( - subscribed_event.set() - if state == RealtimeSubscribeStates.SUBSCRIBED - else None - ) - ) - ) - - await asyncio.wait_for(system_event.wait(), 10) - - # Wait for the channel to be subscribed - await asyncio.wait_for(subscribed_event.wait(), 10) - - created_todo_id = await create_todo( - token, {"description": "Test todo", "is_completed": False} - ) - await asyncio.wait_for(insert_event.wait(), 10) - - await update_todo( - token, created_todo_id, {"description": "Updated todo", "is_completed": True} - ) - await asyncio.wait_for(update_event.wait(), 10) - - await delete_todo(token, created_todo_id) - await asyncio.wait_for(delete_event.wait(), 10) - - assert len(received_events["all"]) == 3 - - insert = received_events["all"][0] - update = received_events["all"][1] - delete = received_events["all"][2] + update = await message_stream.__anext__() + assert isinstance(update, PostgresChangesMessage) - assert insert["data"]["record"]["id"] == created_todo_id - assert insert["data"]["record"]["description"] == "Test todo" - assert not insert["data"]["record"]["is_completed"] + await delete_todo(token, created_todo_id) + delete = await message_stream.__anext__() + assert isinstance(delete, PostgresChangesMessage) - assert update["data"]["record"]["id"] == created_todo_id - assert update["data"]["record"]["description"] == "Updated todo" - assert update["data"]["record"]["is_completed"] + assert insert.payload["data"]["record"] is not None + assert insert.payload["data"]["record"]["id"] == created_todo_id + assert insert.payload["data"]["record"]["description"] == "Test todo" + assert not insert.payload["data"]["record"]["is_completed"] - assert delete["data"]["old_record"]["id"] == created_todo_id + assert update.payload["data"]["record"] is not None + assert update.payload["data"]["record"]["id"] == created_todo_id + assert update.payload["data"]["record"]["description"] == "Updated todo" + assert update.payload["data"]["record"]["is_completed"] - assert received_events["insert"] == [insert] - assert received_events["update"] == [update] - assert received_events["delete"] == [delete] - - await socket.close() + assert delete.payload["data"]["old_record"]["id"] == created_todo_id @pytest.mark.asyncio -async def test_postgrest_changes_on_different_tables(socket: AsyncRealtimeClient): +async def test_postgrest_changes_on_different_tables(socket: RealtimeClient): token = await access_token() - await socket.connect() - await socket.set_auth(token) - channel: AsyncRealtimeChannel = socket.channel("test-postgres-changes") - received_events: dict[str, list[dict]] = {"all": [], "insert": []} - - def all_changes_callback(payload): - received_events["all"].append(payload) - - insert_event = asyncio.Event() - - def insert_callback(payload): - received_events["insert"].append(payload) - insert_event.set() - - subscribed_event = asyncio.Event() - system_event = asyncio.Event() - - await ( - channel.on_postgres_changes( - RealtimePostgresChangesListenEvent.All, all_changes_callback, table="todos" - ) - .on_postgres_changes( - RealtimePostgresChangesListenEvent.Insert, insert_callback, table="todos" - ) - .on_postgres_changes( - RealtimePostgresChangesListenEvent.Insert, insert_callback, table="messages" + channel_options = ( + RealtimeChannelOptions() + .postgres_changes(RealtimePostgresChangesListenEvent.Insert, table="todos") + .postgres_changes(RealtimePostgresChangesListenEvent.Insert, table="messages") + ) + async with socket.channel( + "test-postgres-changes", params=channel_options + ) as channel: + message_stream = channel.messages + system_msg = await message_stream.__anext__() + assert isinstance(system_msg, SystemMessage) + assert system_msg.payload.status == "ok" + created_todo_id = await create_todo( + token, {"description": "Test todo", "is_completed": False} ) - .on_system(lambda _: system_event.set()) - .subscribe( - lambda state, _: ( - subscribed_event.set() - if state == RealtimeSubscribeStates.SUBSCRIBED - else None - ) + insert_todo = await message_stream.__anext__() + assert isinstance(insert_todo, PostgresChangesMessage) + created_message_id = await create_message( + token, {"title": "Test message", "message": "This is a test message"} ) + insert_message = await message_stream.__anext__() + assert isinstance(insert_message, PostgresChangesMessage) + + assert insert_todo.payload["data"]["record"] is not None + assert insert_todo.payload["data"]["record"]["id"] == created_todo_id + assert insert_todo.payload["data"]["record"]["description"] == "Test todo" + assert insert_todo.payload["data"]["record"]["is_completed"] is False + + assert insert_message.payload["data"]["record"] is not None + assert insert_message.payload["data"]["record"]["id"] == created_message_id + assert insert_message.payload["data"]["record"]["title"] == "Test message" + assert ( + insert_message.payload["data"]["record"]["message"] == "This is a test message" ) - await asyncio.wait_for(system_event.wait(), 10) - - # Wait for the channel to be subscribed - await asyncio.wait_for(subscribed_event.wait(), 10) - - created_todo_id = await create_todo( - token, {"description": "Test todo", "is_completed": False} - ) - await asyncio.wait_for(insert_event.wait(), 10) - insert_event.clear() - - created_message_id = await create_message( - token, {"title": "Test message", "message": "This is a test message"} - ) - - await asyncio.wait_for(insert_event.wait(), 10) - - assert len(received_events["all"]) == 1 - assert len(received_events["insert"]) == 2 - - insert = received_events["all"][0] - message_insert = received_events["insert"][1] - - assert insert["data"]["record"]["id"] == created_todo_id - assert insert["data"]["record"]["description"] == "Test todo" - assert insert["data"]["record"]["is_completed"] is False - - assert received_events["insert"] == [insert, message_insert] - - assert message_insert["data"]["record"]["id"] == created_message_id - assert message_insert["data"]["record"]["title"] == "Test message" - assert message_insert["data"]["record"]["message"] == "This is a test message" - - assert received_events["insert"] == [insert, message_insert] - await socket.close() - class CreateTodoResponse(BaseModel): id: str async def create_todo(access_token: str, todo: dict) -> str: - url = f"{URL}/rest/v1/todos?select=id" + url = f"{SUPABASE_URL}/rest/v1/todos?select=id" headers = { "apikey": PUBLISHABLE_KEY, "Authorization": f"Bearer {access_token}", @@ -337,7 +255,7 @@ async def create_todo(access_token: str, todo: dict) -> str: async def update_todo(access_token: str, id: str, todo: dict): - url = f"{URL}/rest/v1/todos?id=eq.{id}" + url = f"{SUPABASE_URL}/rest/v1/todos?id=eq.{id}" headers = { "apikey": PUBLISHABLE_KEY, "Authorization": f"Bearer {access_token}", @@ -355,7 +273,7 @@ class CreateMsgResponse(BaseModel): async def create_message(access_token: str, message: dict) -> int: - url = f"{URL}/rest/v1/messages?select=id" + url = f"{SUPABASE_URL}/rest/v1/messages?select=id" headers = { "apikey": PUBLISHABLE_KEY, "Authorization": f"Bearer {access_token}", @@ -377,7 +295,7 @@ async def create_message(access_token: str, message: dict) -> int: async def delete_todo(access_token: str, id: str): - url = f"{URL}/rest/v1/todos?id=eq.{id}" + url = f"{SUPABASE_URL}/rest/v1/todos?id=eq.{id}" headers = { "apikey": PUBLISHABLE_KEY, "Authorization": f"Bearer {access_token}", @@ -390,294 +308,38 @@ async def delete_todo(access_token: str, id: str): raise Exception(f"Failed to delete todo. Status: {response.status}") -@pytest.mark.asyncio -async def test_multiple_connect_attempts(socket: AsyncRealtimeClient): - # First connection should succeed - await socket.connect() - assert socket.is_connected - initial_ws = socket._ws_connection - - # Second connection attempt should be a no-op since we're already connected - await socket.connect() - assert socket.is_connected - assert socket._ws_connection == initial_ws # Should be the same connection object - - await socket.close() - assert not socket.is_connected - - # Test connection failure and retry behavior - # Temporarily modify the URL to force a connection failure - original_url = socket.url - socket.url = "ws://invalid-url-that-will-fail:12345/websocket" - socket.max_retries = 2 # Reduce retries for faster test - socket.initial_backoff = 0.1 # Reduce backoff for faster test - - start_time = datetime.datetime.now() - - with pytest.raises(Exception) as exc_info: - await socket.connect() - - end_time = datetime.datetime.now() - duration = (end_time - start_time).total_seconds() - - # Should have tried to connect max_retries times with exponential backoff - # First attempt: immediate - # Second attempt: after 0.1s - # Total time should be at least the sum of backoff times but not much more - assert duration >= 0.1, "Should have waited for backoff between retries" - assert duration < 1.0, "Should not have waited longer than necessary" - - # The error message can vary depending on the system and Python version - # Common messages include DNS resolution errors or connection refused - error_msg = str(exc_info.value) - assert any( - msg in error_msg.lower() - for msg in [ - "temporary failure in name resolution", - "nodename nor servname provided", - "name or service not known", - "connection refused", - "failed to establish", - ] - ), f"Unexpected error message: {error_msg}" - - # Restore original URL and verify we can connect again - socket.url = original_url - await socket.connect() - assert socket.is_connected - - await socket.close() - - -@pytest.mark.asyncio -async def test_send_message_reconnection(socket: AsyncRealtimeClient): - # First establish a connection - await socket.connect() - assert socket.is_connected - - # Create a channel and subscribe to it - channel = socket.channel("test-channel") - subscribe_event = asyncio.Event() - await channel.subscribe( - lambda state, _: ( - subscribe_event.set() - if state == RealtimeSubscribeStates.SUBSCRIBED - else None - ) - ) - await asyncio.wait_for(subscribe_event.wait(), 5) - - # Simulate a connection failure by closing the WebSocket - if socket._ws_connection: - await socket._ws_connection.close() - - # Try to send a message - this should trigger reconnection - message = Message( - topic="test-channel", - event=ChannelEvents.broadcast, - payload={"test": "data"}, - ) - await socket.send(message) - - # Wait for reconnection to complete - await asyncio.sleep(1) # Give some time for reconnection - - # Verify we're connected again - assert socket.is_connected - - # Try sending another message to verify the connection is working - await socket.send(message) - - await socket.close() - - @pytest.mark.asyncio async def test_subscribe_to_private_channel_with_broadcast_replay( - socket: AsyncRealtimeClient, + socket: RealtimeClient, ): """Test that channel subscription sends correct payload with broadcast replay configuration.""" - import json - from unittest.mock import AsyncMock, patch + token = await access_token() + + await socket.set_auth(token) - # Mock the websocket connection - mock_ws = AsyncMock() - socket._ws_connection = mock_ws + messages = 10 - # Connect the socket (this will use our mock) - await socket.connect() + topic = "test-channel" + event = "broadcast" # Calculate replay timestamp ten_mins_ago = datetime.datetime.now() - datetime.timedelta(minutes=10) ten_mins_ago_ms = int(ten_mins_ago.timestamp() * 1000) - # Create channel with broadcast replay configuration - channel: AsyncRealtimeChannel = socket.channel( - "test-private-channel", - params={ - "config": { - "private": True, - "broadcast": {"replay": {"since": ten_mins_ago_ms, "limit": 100}}, - "presence": {"enabled": True, "key": ""}, - } - }, + channel_options = ( + RealtimeChannelOptions() + .private(True) + .broadcast(replay_since=ten_mins_ago_ms, replay_limit=messages) ) - # Mock the subscription callback to be called immediately - callback_called = False - - def mock_callback(state, error): - nonlocal callback_called - callback_called = True - - # Subscribe to the channel - await channel.subscribe(mock_callback) - - # Verify that send was called with the correct payload - assert mock_ws.send.called, "WebSocket send should have been called" - - # Get the sent message - sent_message = mock_ws.send.call_args[0][0] - message_data = json.loads(sent_message) - - # Verify the message structure - assert message_data["topic"] == "realtime:test-private-channel" - assert message_data["event"] == "phx_join" - assert "ref" in message_data - assert "payload" in message_data - - # Verify the payload contains the correct broadcast replay configuration - payload = message_data["payload"] - assert "config" in payload - - config = payload["config"] - assert config["private"] is True - assert "broadcast" in config - - broadcast_config = config["broadcast"] - assert "replay" in broadcast_config - - replay_config = broadcast_config["replay"] - assert replay_config["since"] == ten_mins_ago_ms - assert replay_config["limit"] == 100 - - # Verify postgres_changes array is present (even if empty) - assert "postgres_changes" in config - assert isinstance(config["postgres_changes"], list) - - await socket.close() - - -@pytest.mark.asyncio -async def test_subscribe_to_channel_with_empty_replay_config( - socket: AsyncRealtimeClient, -): - """Test that channel subscription handles empty replay configuration correctly.""" - import json - from unittest.mock import AsyncMock, patch - - # Mock the websocket connection - mock_ws = AsyncMock() - socket._ws_connection = mock_ws - - # Connect the socket - await socket.connect() - - # Create channel with empty replay configuration - channel: AsyncRealtimeChannel = socket.channel( - "test-empty-replay", - params={ - "config": { - "private": False, - "broadcast": {"ack": True, "self": False, "replay": {}}, - "presence": {"enabled": True, "key": ""}, - } - }, - ) - - # Mock the subscription callback - callback_called = False - - def mock_callback(state, error): - nonlocal callback_called - callback_called = True - - # Subscribe to the channel - await channel.subscribe(mock_callback) - - # Verify that send was called - assert mock_ws.send.called, "WebSocket send should have been called" - - # Get the sent message - sent_message = mock_ws.send.call_args[0][0] - message_data = json.loads(sent_message) - - # Verify the payload structure - payload = message_data["payload"] - config = payload["config"] - - assert config["private"] is False - assert "broadcast" in config - - broadcast_config = config["broadcast"] - assert broadcast_config["ack"] is True - assert broadcast_config["self"] is False - assert broadcast_config["replay"] == {} - - await socket.close() - - -@pytest.mark.asyncio -async def test_subscribe_to_channel_without_replay_config(socket: AsyncRealtimeClient): - """Test that channel subscription works without replay configuration.""" - import json - from unittest.mock import AsyncMock, patch - - # Mock the websocket connection - mock_ws = AsyncMock() - socket._ws_connection = mock_ws - - # Connect the socket - await socket.connect() - - # Create channel without replay configuration - channel: AsyncRealtimeChannel = socket.channel( - "test-no-replay", - params={ - "config": { - "private": False, - "broadcast": {"ack": True, "self": True}, - "presence": {"enabled": True, "key": ""}, - } - }, - ) - - # Mock the subscription callback - callback_called = False - - def mock_callback(state, error): - nonlocal callback_called - callback_called = True - - # Subscribe to the channel - await channel.subscribe(mock_callback) - - # Verify that send was called - assert mock_ws.send.called, "WebSocket send should have been called" - - # Get the sent message - sent_message = mock_ws.send.call_args[0][0] - message_data = json.loads(sent_message) - - # Verify the payload structure - payload = message_data["payload"] - config = payload["config"] - - assert config["private"] is False - assert "broadcast" in config - - broadcast_config = config["broadcast"] - assert broadcast_config["ack"] is True - assert broadcast_config["self"] is True - assert "replay" not in broadcast_config - - await socket.close() + for i in range(messages): + await realtime_send(token, topic, event, {"i": str(i)}) + + async with socket.channel(topic, params=channel_options) as chan: + message_stream = chan.messages + for i in range(messages): + broadcast = await message_stream.__anext__() + assert isinstance(broadcast, BroadcastMessage) + assert broadcast.payload["payload"]["i"] == str(i) + assert broadcast.payload["meta"] + assert broadcast.payload["meta"]["replayed"] diff --git a/src/realtime/tests/test_presence.py b/src/realtime/tests/test_presence.py index 0d4c587d..b4ef4a34 100644 --- a/src/realtime/tests/test_presence.py +++ b/src/realtime/tests/test_presence.py @@ -1,12 +1,15 @@ import asyncio import datetime import os -from typing import Dict, List, Tuple +from typing import AsyncIterator, Dict, List, Tuple import pytest from dotenv import load_dotenv -from realtime import AsyncRealtimeChannel, AsyncRealtimeClient, AsyncRealtimePresence +from realtime import RealtimeChannel, RealtimeClient +from realtime.client import connect_once +from realtime.message import PresenceDiffMessage, ServerMessage +from realtime.presence import PresenceJoin, PresenceLeave, _transform_state from realtime.types import ChannelStates, Presence, RawPresenceState load_dotenv() @@ -18,97 +21,59 @@ ) +async def get_token() -> str: + return PUBLISHABLE_KEY + + @pytest.fixture -def socket() -> AsyncRealtimeClient: +async def socket() -> AsyncIterator[RealtimeClient]: url = f"{URL}/realtime/v1" - return AsyncRealtimeClient(url, PUBLISHABLE_KEY) + async with connect_once(url, get_token) as client: + yield client @pytest.mark.asyncio -async def test_presence(socket: AsyncRealtimeClient): - await socket.connect() - - channel: AsyncRealtimeChannel = socket.channel("room") - - join_events: List[Tuple[str, List[Dict], List[Presence]]] = [] - leave_events: List[Tuple[str, List[Dict], List[Presence]]] = [] - - sync_event = asyncio.Event() - - def on_sync(): - sync_event.set() - - def on_join(key: str, current_presences: List[Dict], new_presences: List[Presence]): - join_events.append((key, current_presences, new_presences)) - - def on_leave( - key: str, current_presences: List[Dict], left_presences: List[Presence] - ): - leave_events.append((key, current_presences, left_presences)) - - await ( - channel.on_presence_sync(on_sync) - .on_presence_join(on_join) - .on_presence_leave(on_leave) - .subscribe() - ) - - # Wait for the first sync event, which should be immediate - await asyncio.wait_for(sync_event.wait(), 5) - sync_event.clear() - - # Track first user +async def test_presence(socket: RealtimeClient): user1 = {"user_id": "1", "online_at": datetime.datetime.now().isoformat()} - await channel.track(user1) - - await asyncio.wait_for(sync_event.wait(), 5) - sync_event.clear() - - # Assert first user is in the presence state - presences = [(state, value) for state, value in channel.presence.state.items()] - - assert len(presences) == 1 - assert len(presences[0][1]) == 1 - assert presences[0][1][0]["user_id"] == user1["user_id"] # type: ignore - assert presences[0][1][0]["online_at"] == user1["online_at"] # type: ignore - assert "presence_ref" in presences[0][1][0] - - assert len(join_events) == 1 - assert len(join_events[0][2]) == 1 - assert join_events[0][2][0]["user_id"] == user1["user_id"] # type: ignore - assert join_events[0][2][0]["online_at"] == user1["online_at"] # type: ignore - assert "presence_ref" in join_events[0][2][0] - - # Track second user user2 = {"user_id": "2", "online_at": datetime.datetime.now().isoformat()} - await channel.track(user2) - - await asyncio.wait_for(sync_event.wait(), 5) - sync_event.clear() - - # Assert both users are in the presence state - for key, value in channel.presence.state.items(): - assert len(value) == 1 - assert value[0]["user_id"] in ["1", "2"] # type: ignore - assert "online_at" in value[0] - assert "presence_ref" in value[0] - assert len(join_events) == 2 - assert len(join_events[1][2]) == 1 - assert join_events[1][2][0]["user_id"] == user2["user_id"] # type: ignore - assert join_events[1][2][0]["online_at"] == user2["online_at"] # type: ignore - assert "presence_ref" in join_events[1][2][0] - - # Untrack all users - await channel.untrack() - - await asyncio.wait_for(sync_event.wait(), 5) - # Assert presence state is empty and leave events were triggered - assert channel.presence.state == {} - assert len(leave_events) == 2 - assert leave_events[0] != leave_events[1] - - await socket.close() + async with socket.channel("room") as chan: + res = await chan.track(user1) + messages_stream = chan.messages + fst_join = await messages_stream.__anext__() + if not isinstance(fst_join, PresenceJoin): + raise Exception("unexpected message") + + presences = list(chan.presence.state.items()) + + assert len(presences) == 1 + assert len(presences[0][1]) == 1 + assert presences[0][1][0].get("user_id", "") == user1["user_id"] + assert presences[0][1][0].get("online_at", "") == user1["online_at"] + assert "presence_ref" in presences[0][1][0] + + assert fst_join.new_presences[0].get("user_id", "") == user1["user_id"] + assert fst_join.new_presences[0].get("online_at", "") == user1["online_at"] + assert "presence_ref" in fst_join.new_presences[0] + + await chan.track(user2) + snd_join = await messages_stream.__anext__() + if not isinstance(snd_join, PresenceJoin): + raise Exception("unexpected message") + + assert snd_join.new_presences[0].get("user_id", "") == user2["user_id"] + assert snd_join.new_presences[0].get("online_at", "") == user2["online_at"] + assert "presence_ref" in snd_join.new_presences[0] + + await chan.untrack() + fst_leave = await messages_stream.__anext__() + if not isinstance(fst_leave, PresenceLeave): + raise Exception("unexpected message") + snd_leave = await messages_stream.__anext__() + if not isinstance(snd_leave, PresenceLeave): + raise Exception("unexpected message") + assert chan.presence.state == {} + assert fst_leave != snd_leave def test_transform_state_raw_presence_state() -> None: @@ -137,13 +102,13 @@ def test_transform_state_raw_presence_state() -> None: "user2": [{"presence_ref": "GHI789", "user_id": "user2", "status": "offline"}], } - result = AsyncRealtimePresence._transform_state(raw_state) + result = _transform_state(raw_state) assert result == expected_output def test_transform_state_empty_input() -> None: empty_state: RawPresenceState = {} - result = AsyncRealtimePresence._transform_state(empty_state) + result = _transform_state(empty_state) assert result == {} @@ -172,104 +137,5 @@ def test_transform_state_additional_fields() -> None: ] } - result = AsyncRealtimePresence._transform_state(state_with_additional_fields) + result = _transform_state(state_with_additional_fields) assert result == expected_output - - -def test_presence_has_callback_attached(): - """Test that _has_callback_attached property correctly detects presence callbacks.""" - presence = AsyncRealtimePresence() - - # Initially no callbacks should be attached - assert not presence._has_callback_attached - - # After setting sync callback - presence.on_sync(lambda: None) - assert presence._has_callback_attached - - # Reset and test with join callback - presence = AsyncRealtimePresence() - presence.on_join(lambda key, current, new: None) - assert presence._has_callback_attached - - # Reset and test with leave callback - presence = AsyncRealtimePresence() - presence.on_leave(lambda key, current, left: None) - assert presence._has_callback_attached - - -def test_presence_config_includes_enabled_field() -> None: - """Test that presence config correctly includes enabled flag.""" - from realtime.types import RealtimeChannelPresenceConfig - - # Test creating presence config with enabled field - config: RealtimeChannelPresenceConfig = {"key": "user123", "enabled": True} - assert config["key"] == "user123" - assert config["enabled"] == True - - # Test with enabled False - config_disabled: RealtimeChannelPresenceConfig = {"key": "", "enabled": False} - assert config_disabled["key"] == "" - assert config_disabled["enabled"] == False - - -@pytest.mark.asyncio -async def test_presence_enabled_when_callbacks_attached() -> None: - """Test that presence.enabled is set correctly based on callback attachment.""" - from unittest.mock import AsyncMock, Mock - - socket = AsyncRealtimeClient(f"{URL}/realtime/v1", PUBLISHABLE_KEY) - channel = socket.channel("test") - - # Mock the join_push to capture the payload - mock_join_push = Mock() - mock_join_push.receive = Mock(return_value=mock_join_push) - mock_join_push.update_payload = Mock() - mock_join_push.resend = AsyncMock() - channel.join_push = mock_join_push - - # Mock socket connection by setting _ws_connection - mock_ws = Mock() - socket._ws_connection = mock_ws - socket._leave_open_topic = AsyncMock() # type: ignore - - # Add presence callback before subscription - channel.on_presence_sync(lambda: None) - - await channel.subscribe() - - # Verify that update_payload was called - assert mock_join_push.update_payload.called - - # Get the payload that was passed to update_payload - call_args = mock_join_push.update_payload.call_args - payload = call_args[0][0] - - # Verify presence.enabled is True because callback is attached - assert payload["config"]["presence"]["enabled"] == True - - -@pytest.mark.asyncio -async def test_resubscribe_on_presence_callback_addition() -> None: - """Test that channel resubscribes when presence callbacks are added after joining.""" - import asyncio - from unittest.mock import AsyncMock - - socket = AsyncRealtimeClient(f"{URL}/realtime/v1", PUBLISHABLE_KEY) - channel = socket.channel("test") - - # Mock the channel as joined - channel.state = ChannelStates.JOINED - channel._joined_once = True - - # Mock resubscribe method - channel._resubscribe = AsyncMock() # type: ignore - - # Add presence callbacks after joining - channel.on_presence_sync(lambda: None) - - # Wait a bit for async tasks to complete - await asyncio.sleep(0.1) - - # Verify resubscribe was called - assert channel._resubscribe.call_count == 1 diff --git a/src/realtime/tests/test_reconnection.py b/src/realtime/tests/test_reconnection.py new file mode 100644 index 00000000..80edbacd --- /dev/null +++ b/src/realtime/tests/test_reconnection.py @@ -0,0 +1,195 @@ +import asyncio +import datetime +import os +from typing import AsyncIterator + +import aiohttp +import pytest +from dotenv import load_dotenv +from pydantic import BaseModel, TypeAdapter +from websockets import broadcast +from websockets.asyncio.server import ServerConnection, serve +from websockets.exceptions import ConnectionClosedOK + +from realtime import ( + RealtimeChannel, + RealtimeClient, + RealtimePostgresChangesListenEvent, + RealtimeSubscribeStates, +) +from realtime.channel import RealtimeChannelOptions +from realtime.client import MaxRetriesExceeded, automatically_reconnect, connect_once +from realtime.message import ( + BroadcastMessage, + ChannelLeaveMessage, + ClientMessage, + HeartbeatMessage, + JoinMessage, + Message, + PostgresChangesMessage, + ReplyMessage, + ReplyPostgresChanges, + SuccessReplyMessage, + SystemMessage, +) +from realtime.types import DEFAULT_HEARTBEAT_INTERVAL, DEFAULT_TIMEOUT, ChannelEvents + +PORT = 55555 +URL = "localhost" +PUBLISHABLE_KEY = "my-publishable-key" +MessageParser: TypeAdapter[ClientMessage] = TypeAdapter(ClientMessage) + + +async def get_token() -> str: + return PUBLISHABLE_KEY + + +def reply_ack(topic: str, ref: str) -> str: + reply_msg = ReplyMessage( + event=ChannelEvents.reply, + topic=topic, + ref=ref, + payload=SuccessReplyMessage( + status="ok", response=ReplyPostgresChanges(postgres_changes=[]) + ), + ) + return reply_msg.model_dump_json() + + +class RealtimeServer: + def __init__(self, url: str, port: int): + self.url = url + self.port = port + + async def start(self) -> None: + self.ws_server = await serve(self.handler, self.url, self.port) + await self.ws_server.start_serving() + + async def __aenter__(self) -> "RealtimeServer": + await self.start() + await self.ws_server.__aenter__() + return self + + async def __aexit__(self, type, exc_value, traceback): + await self.ws_server.__aexit__(type, exc_value, traceback) + + async def close(self): + self.ws_server.close() + await self.ws_server.wait_closed() + + async def handler(self, connection: ServerConnection) -> None: + while True: + try: + message = await connection.recv(decode=False) + except ConnectionClosedOK: + break + msg = MessageParser.validate_json(message) + match msg: + case HeartbeatMessage(topic=topic, ref=ref): + await connection.send(reply_ack(topic, ref)) + case JoinMessage(topic=topic, ref=ref): + await connection.send(reply_ack(topic, ref)) + case ChannelLeaveMessage(topic=topic, ref=ref): + await connection.send(reply_ack(topic, ref)) + + +@pytest.fixture +async def server() -> AsyncIterator[RealtimeServer]: + async with RealtimeServer(URL, PORT) as server: + yield server + + +async def test_raises_on_send_broadcast(server: RealtimeServer) -> None: + async with ( + connect_once(f"http://{URL}:{PORT}", get_token) as client, + client.channel("test-channel") as channel, + ): + await server.close() + try: + await channel.send_broadcast("test", {}) + except ConnectionClosedOK: + return + raise Exception("should have raised a connection error, but got nothing") + + +async def test_raises_on_send_presence(server: RealtimeServer) -> None: + async with ( + connect_once(f"http://{URL}:{PORT}", get_token) as client, + client.channel("test-channel") as channel, + ): + await server.close() + try: + await channel.send_presence("presence", {}) + except ConnectionClosedOK: + return + raise Exception("should have raised a connection error, but got nothing") + + +async def test_raises_on_listen_message(server: RealtimeServer) -> None: + async with ( + connect_once(f"http://{URL}:{PORT}", get_token) as client, + client.channel("test-channel") as channel, + ): + await server.close() + try: + async for msg in channel.messages: + pass + except ConnectionClosedOK: + return + raise Exception("should have raised a connection error, but got nothing") + + +async def test_automatic_reconnection(server: RealtimeServer) -> None: + async def restart_server_after_2_seconds(): + await asyncio.sleep(2) + await server.start() + + async def client_callback(client: RealtimeClient) -> None: + return + + await server.close() + asyncio.create_task(restart_server_after_2_seconds()) + await automatically_reconnect( + f"http://{URL}:{PORT}", get_token, client_callback, max_retries=3 + ) + + +async def test_automatic_reconnection_fails_after_max_retries() -> None: + # no server running. + async def client_callback(client: RealtimeClient) -> None: + return + + try: + await automatically_reconnect( + f"http://{URL}:{PORT}", get_token, client_callback, max_retries=3 + ) + except MaxRetriesExceeded as exc: + return + raise Exception("expected max retries exceeded exception but got nothing") + + +async def test_automatic_reconnection_calls_callback_when_connected( + server: RealtimeServer, +) -> None: + called_twice = False + + async def restart_server_after_1_second(): + await asyncio.sleep(1) + await server.start() + + async def client_callback(client: RealtimeClient) -> None: + nonlocal called_twice + if called_twice: + return + called_twice = True + await server.close() + # reschedule server to re-start 1 second later + asyncio.create_task(restart_server_after_1_second()) + # this should raise when the server is offline + async with client.channel("test-channel") as channel: + pass + + await automatically_reconnect( + f"http://{URL}:{PORT}", get_token, client_callback, max_retries=3 + ) + assert called_twice diff --git a/src/realtime/tests/test_timer.py b/src/realtime/tests/test_timer.py deleted file mode 100644 index e683639e..00000000 --- a/src/realtime/tests/test_timer.py +++ /dev/null @@ -1,124 +0,0 @@ -import asyncio - -import pytest - -from realtime._async.timer import AsyncTimer - - -def linear_backoff(tries: int) -> float: - return tries * 0.1 - - -@pytest.mark.asyncio -async def test_timer_initialization(): - async def callback(): - pass - - timer = AsyncTimer(callback, linear_backoff) - assert timer.tries == 0 - assert timer.timer is None - assert timer.callback == callback - assert timer.timer_calc == linear_backoff - - -@pytest.mark.asyncio -async def test_timer_schedule(): - callback_called = False - - async def callback(): - nonlocal callback_called - callback_called = True - - timer = AsyncTimer(callback, linear_backoff) - timer.schedule_timeout() - - assert timer.tries == 1 - assert timer.timer is not None - assert not timer.timer.done() - - # Wait for the timer to complete - await timer.timer - assert callback_called - - -@pytest.mark.asyncio -async def test_timer_reset(): - callback_called = False - - async def callback(): - nonlocal callback_called - callback_called = True - - timer = AsyncTimer(callback, linear_backoff) - timer.schedule_timeout() - - # Reset before the timer completes - timer.reset() - assert timer.tries == 0 - assert timer.timer is None - - # Wait a bit to ensure the original timer doesn't fire - await asyncio.sleep(0.2) - assert not callback_called - - -@pytest.mark.asyncio -async def test_timer_multiple_schedules(): - callback_count = 0 - - async def callback(): - nonlocal callback_count - callback_count += 1 - - timer = AsyncTimer(callback, linear_backoff) - - # Schedule multiple times - timer.schedule_timeout() - timer.schedule_timeout() - timer.schedule_timeout() - - assert timer.tries == 3 - assert timer.timer is not None - - # Wait for the last timer to complete - await timer.timer - assert callback_count == 1 # Only the last schedule should fire - - -@pytest.mark.asyncio -async def test_timer_callback_error(): - error_caught = False - - async def callback(): - raise ValueError("Test error") - - timer = AsyncTimer(callback, linear_backoff) - timer.schedule_timeout() - assert timer.timer is not None - - # Wait for the timer to complete - await timer.timer - # The error should be caught and logged, but not re-raised - assert timer.timer.done() - - -@pytest.mark.asyncio -async def test_timer_cancellation(): - callback_called = False - - async def callback(): - nonlocal callback_called - callback_called = True - - timer = AsyncTimer(callback, linear_backoff) - timer.schedule_timeout() - - # Cancel the timer - assert timer.timer is not None - timer.timer.cancel() - - # Wait a bit to ensure the timer doesn't fire - await asyncio.sleep(0.2) - assert not callback_called - assert timer.timer.done() - assert timer.timer.cancelled() diff --git a/src/supabase/pyproject.toml b/src/supabase/pyproject.toml index ec1d81b2..0da7a181 100644 --- a/src/supabase/pyproject.toml +++ b/src/supabase/pyproject.toml @@ -51,6 +51,7 @@ dev = [ { include-group = "lints" }, ] tests = [ + "storage3[iceberg] == 3.0.0a1", # x-release-please-version "pytest >= 8.4.1", "pytest-cov >= 6.2.1", "pytest-asyncio >=0.24,<1.1", diff --git a/src/supabase/src/supabase/_async/client.py b/src/supabase/src/supabase/_async/client.py index 029a8c84..a06a647d 100644 --- a/src/supabase/src/supabase/_async/client.py +++ b/src/supabase/src/supabase/_async/client.py @@ -1,8 +1,6 @@ -import asyncio import copy -import re from types import TracebackType -from typing import Dict, List, Literal, overload +from typing import Dict, Literal, overload from postgrest import AsyncPostgrestClient from postgrest.request_builder import ( @@ -11,7 +9,7 @@ RPCFilterRequestBuilder, ) from postgrest.types import CountMethod -from realtime import AsyncRealtimeChannel, AsyncRealtimeClient, RealtimeChannelOptions +from realtime import connect_once from storage3 import AsyncStorageClient from supabase_auth import AsyncMemoryStorage, AsyncSupabaseAuthClient from supabase_auth.types import AuthChangeEvent, Session @@ -20,7 +18,6 @@ from yarl import URL from ..lib.client_options import AsyncClientOptions as ClientOptions -from ..types import RealtimeClientOptions # Create an exception class when user does not provide a valid url or key. @@ -58,18 +55,20 @@ def __init__( if not supabase_key: raise SupabaseException("supabase_key is required") + self.supabase_url = ( + URL(supabase_url) if supabase_url.endswith("/") else URL(supabase_url + "/") + ) # Check if the url and key are valid - if not re.match(r"^(https?)://.+", supabase_url): - raise SupabaseException("Invalid URL") + if not ( + self.supabase_url.scheme == "https" or self.supabase_url.scheme == "http" + ): + raise SupabaseException(f"Invalid URL: {self.supabase_url}") if options is None: options = ClientOptions(storage=AsyncMemoryStorage()) self.http_session: AsyncHttpSession = http_session - self.supabase_url = ( - URL(supabase_url) if supabase_url.endswith("/") else URL(supabase_url + "/") - ) self.supabase_key = supabase_key self.options = copy.copy(options) self.options.headers = { @@ -85,20 +84,8 @@ def __init__( self.storage_url = self.supabase_url.joinpath("storage", "v1", "") self.functions_url = self.supabase_url.joinpath("functions", "v1") - # Instantiate clients. - self.auth = self._init_supabase_auth_client( - auth_url=str(self.auth_url), - client_options=self.options, - http_session=self.http_session, - ) - self.realtime = self._init_realtime_client( - realtime_url=self.realtime_url, - supabase_key=self.supabase_key, - options=self.options.realtime if self.options else None, - ) - self._postgrest: AsyncPostgrestClient | None = None - self._storage: AsyncStorageClient | None = None - self._functions: AsyncFunctionsClient | None = None + self.auth_access_token = supabase_key + self.auth.on_auth_state_change(self._listen_to_auth_events) async def __aenter__(self) -> "AsyncClient": @@ -220,64 +207,49 @@ def rpc( return self.postgrest.rpc(fn, params, count=count, head=head, get=get) @property - def postgrest(self) -> AsyncPostgrestClient: - if self._postgrest is None: - self._postgrest = self._init_postgrest_client( - rest_url=str(self.rest_url), - headers=self.options.headers, - schema=self.options.schema, - http_session=self.http_session, - ) + def auth(self) -> AsyncSupabaseAuthClient: + return self._init_supabase_auth_client( + auth_url=str(self.auth_url), + client_options=self.options, + http_session=self.http_session, + ) - return self._postgrest + @property + def postgrest(self) -> AsyncPostgrestClient: + return self._init_postgrest_client( + rest_url=str(self.rest_url), + headers=self.options.headers, + schema=self.options.schema, + http_session=self.http_session, + ) @property def storage(self) -> AsyncStorageClient: - if self._storage is None: - self._storage = self._init_storage_client( - storage_url=str(self.storage_url), - headers=self.options.headers, - http_session=self.http_session, - ) - return self._storage + return self._init_storage_client( + storage_url=str(self.storage_url), + headers=self.options.headers, + http_session=self.http_session, + ) @property def functions(self) -> AsyncFunctionsClient: - if self._functions is None: - self._functions = AsyncFunctionsClient( - url=str(self.functions_url), - headers=self.options.headers, - ) - return self._functions - - def channel( - self, topic: str, params: RealtimeChannelOptions | None = None - ) -> AsyncRealtimeChannel: - """Creates a Realtime channel with Broadcast, Presence, and Postgres Changes.""" - return self.realtime.channel(topic, params or {}) - - def get_channels(self) -> List[AsyncRealtimeChannel]: - """Returns all realtime channels.""" - return self.realtime.get_channels() - - async def remove_channel(self, channel: AsyncRealtimeChannel) -> None: - """Unsubscribes and removes Realtime channel from Realtime client.""" - await self.realtime.remove_channel(channel) + return AsyncFunctionsClient( + url=str(self.functions_url), + headers=self.options.headers, + ) - async def remove_all_channels(self) -> None: - """Unsubscribes and removes all Realtime channels from Realtime client.""" - await self.realtime.remove_all_channels() + async def realtime_get_token_callback(self) -> str | None: + return self.auth_access_token - @staticmethod - def _init_realtime_client( - realtime_url: URL, - supabase_key: str, - options: RealtimeClientOptions | None = None, - ) -> AsyncRealtimeClient: - realtime_options = options or {} + @property + def realtime( + self, + ) -> connect_once: """Private method for creating an instance of the realtime-py client.""" - return AsyncRealtimeClient( - str(realtime_url), token=supabase_key, **realtime_options + return connect_once( + str(self.realtime_url), + token_callback=self.realtime_get_token_callback, + **(self.options.realtime or {}), ) @staticmethod @@ -354,4 +326,4 @@ def _listen_to_auth_events( self.auth.default_headers = self.auth.default_headers.override( "Authorization", auth_header ) - asyncio.create_task(self.realtime.set_auth(access_token)) + self.auth_access_token = access_token diff --git a/src/supabase/src/supabase/types.py b/src/supabase/src/supabase/types.py index 4f774531..93384f4e 100644 --- a/src/supabase/src/supabase/types.py +++ b/src/supabase/src/supabase/types.py @@ -2,7 +2,5 @@ class RealtimeClientOptions(TypedDict, total=False): - auto_reconnect: bool hb_interval: int - max_retries: int - initial_backoff: float + hb_timeout: int diff --git a/src/supabase/tests/_async/test_client.py b/src/supabase/tests/_async/test_client.py index b5c19dbe..ca500d52 100644 --- a/src/supabase/tests/_async/test_client.py +++ b/src/supabase/tests/_async/test_client.py @@ -1,5 +1,4 @@ import os -from unittest.mock import AsyncMock, MagicMock from supabase_auth import AsyncMemoryStorage @@ -54,39 +53,6 @@ async def test_schema_update(create_async_client: AsyncClientCallable) -> None: assert client.schema("new_schema") -async def test_updates_the_authorization_header_on_auth_events( - create_async_client: AsyncClientCallable, -) -> None: - async with create_async_client(url, key) as client: - assert client.options.headers.get("apiKey") == key - assert client.options.headers.get("Authorization") == f"Bearer {key}" - - mock_session = MagicMock(access_token="secretuserjwt") - realtime_mock = AsyncMock() - client.realtime = realtime_mock - - client._listen_to_auth_events("SIGNED_IN", mock_session) - - updated_authorization = f"Bearer {mock_session.access_token}" - - assert client.options.headers.get("apiKey") == key - assert client.options.headers.get("Authorization") == updated_authorization - - assert client.postgrest.default_headers.get("apiKey") == key - assert ( - client.postgrest.default_headers.get("Authorization") - == updated_authorization - ) - - assert client.auth.default_headers.get("apiKey") == key - assert client.auth.default_headers.get("Authorization") == updated_authorization - - assert client.storage.default_headers.get("apiKey") == key - assert ( - client.storage.default_headers.get("Authorization") == updated_authorization - ) - - async def test_supports_setting_a_global_authorization_header( create_async_client: AsyncClientCallable, ) -> None: diff --git a/src/utils/src/supabase_utils/http/io.py b/src/utils/src/supabase_utils/http/io.py index 02810ce9..1f5ca8ed 100644 --- a/src/utils/src/supabase_utils/http/io.py +++ b/src/utils/src/supabase_utils/http/io.py @@ -102,7 +102,6 @@ async def communicate( if request.delay: await asyncio.sleep(request.delay) response = await self.session.send(request) - print(response) http_request = iterator.send(response) except StopIteration: return return_value_iterator.return_value diff --git a/uv.lock b/uv.lock index 0c51e3d6..45653133 100644 --- a/uv.lock +++ b/uv.lock @@ -804,7 +804,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -1909,6 +1909,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5b/5a/bc7b4a4ef808fa59a816c17b20c4bef6884daebbdf627ff2a161da67da19/propcache-0.4.1-py3-none-any.whl", hash = "sha256:af2a6052aeb6cf17d3e46ee169099044fd8224cbaf75c76a2ef596e8163e2237", size = 13305, upload-time = "2025-10-08T19:49:00.792Z" }, ] +[[package]] +name = "pudb" +version = "2025.1.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jedi" }, + { name = "packaging" }, + { name = "pygments" }, + { name = "typing-extensions" }, + { name = "urwid" }, + { name = "urwid-readline" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f5/69/0c52c6058da586e173e48ef8f9c45bc805d9df78d5ad4d69467ff68694b5/pudb-2025.1.5.tar.gz", hash = "sha256:e6dedb81fc3c8cdcd66cf62e86337cb913570ebb2c994c9ab52012d059e086ad", size = 225714, upload-time = "2025-12-06T20:18:13.993Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b6/af/42f4877e71728dd59ed7080205f29eba16406f5783c29cdc850200138c80/pudb-2025.1.5-py3-none-any.whl", hash = "sha256:e40c69f102d534cc7b86020c0217e3d424ee81b14a229d068ac122089c301f59", size = 92638, upload-time = "2025-12-06T20:18:12.664Z" }, +] + [[package]] name = "pycparser" version = "2.23" @@ -2386,13 +2403,16 @@ version = "3.0.0a1" source = { editable = "src/realtime" } dependencies = [ { name = "pydantic" }, + { name = "pytest-asyncio" }, { name = "typing-extensions" }, { name = "websockets" }, + { name = "yarl" }, ] [package.dev-dependencies] dev = [ { name = "aiohttp" }, + { name = "pudb" }, { name = "pylsp-mypy" }, { name = "pytest" }, { name = "pytest-asyncio" }, @@ -2410,6 +2430,7 @@ lints = [ ] tests = [ { name = "aiohttp" }, + { name = "pudb" }, { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-cov" }, @@ -2419,13 +2440,16 @@ tests = [ [package.metadata] requires-dist = [ { name = "pydantic", specifier = ">=2.11.7,<3.0.0" }, + { name = "pytest-asyncio", specifier = ">=0.21.0" }, { name = "typing-extensions", specifier = ">=4.14.0" }, { name = "websockets", specifier = ">=11,<16" }, + { name = "yarl", specifier = ">=1.22.0" }, ] [package.metadata.requires-dev] dev = [ { name = "aiohttp", specifier = ">=3.12.13" }, + { name = "pudb", specifier = ">=2025.1.5" }, { name = "pylsp-mypy", specifier = ">=0.7.0,<0.8.0" }, { name = "pytest", specifier = ">=8.4.1" }, { name = "pytest-asyncio", specifier = ">=1.0.0" }, @@ -2443,6 +2467,7 @@ lints = [ ] tests = [ { name = "aiohttp", specifier = ">=3.12.13" }, + { name = "pudb", specifier = ">=2025.1.5" }, { name = "pytest", specifier = ">=8.4.1" }, { name = "pytest-asyncio", specifier = ">=1.0.0" }, { name = "pytest-cov", specifier = ">=6.2.1" }, @@ -2652,23 +2677,23 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "alabaster", marker = "python_full_version < '3.11'" }, - { name = "babel", marker = "python_full_version < '3.11'" }, - { name = "colorama", marker = "python_full_version < '3.11' and sys_platform == 'win32'" }, - { name = "docutils", marker = "python_full_version < '3.11'" }, - { name = "imagesize", marker = "python_full_version < '3.11'" }, - { name = "jinja2", marker = "python_full_version < '3.11'" }, - { name = "packaging", marker = "python_full_version < '3.11'" }, - { name = "pygments", marker = "python_full_version < '3.11'" }, - { name = "requests", marker = "python_full_version < '3.11'" }, - { name = "snowballstemmer", marker = "python_full_version < '3.11'" }, - { name = "sphinxcontrib-applehelp", marker = "python_full_version < '3.11'" }, - { name = "sphinxcontrib-devhelp", marker = "python_full_version < '3.11'" }, - { name = "sphinxcontrib-htmlhelp", marker = "python_full_version < '3.11'" }, - { name = "sphinxcontrib-jsmath", marker = "python_full_version < '3.11'" }, - { name = "sphinxcontrib-qthelp", marker = "python_full_version < '3.11'" }, - { name = "sphinxcontrib-serializinghtml", marker = "python_full_version < '3.11'" }, - { name = "tomli", marker = "python_full_version < '3.11'" }, + { name = "alabaster" }, + { name = "babel" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "docutils" }, + { name = "imagesize" }, + { name = "jinja2" }, + { name = "packaging" }, + { name = "pygments" }, + { name = "requests" }, + { name = "snowballstemmer" }, + { name = "sphinxcontrib-applehelp" }, + { name = "sphinxcontrib-devhelp" }, + { name = "sphinxcontrib-htmlhelp" }, + { name = "sphinxcontrib-jsmath" }, + { name = "sphinxcontrib-qthelp" }, + { name = "sphinxcontrib-serializinghtml" }, + { name = "tomli" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6f/6d/be0b61178fe2cdcb67e2a92fc9ebb488e3c51c4f74a36a7824c0adf23425/sphinx-8.1.3.tar.gz", hash = "sha256:43c1911eecb0d3e161ad78611bc905d1ad0e523e4ddc202a58a821773dc4c927", size = 8184611, upload-time = "2024-10-13T20:27:13.93Z" } wheels = [ @@ -2684,23 +2709,23 @@ resolution-markers = [ "python_full_version >= '3.11' and python_full_version < '3.13'", ] dependencies = [ - { name = "alabaster", marker = "python_full_version >= '3.11'" }, - { name = "babel", marker = "python_full_version >= '3.11'" }, - { name = "colorama", marker = "python_full_version >= '3.11' and sys_platform == 'win32'" }, - { name = "docutils", marker = "python_full_version >= '3.11'" }, - { name = "imagesize", marker = "python_full_version >= '3.11'" }, - { name = "jinja2", marker = "python_full_version >= '3.11'" }, - { name = "packaging", marker = "python_full_version >= '3.11'" }, - { name = "pygments", marker = "python_full_version >= '3.11'" }, - { name = "requests", marker = "python_full_version >= '3.11'" }, - { name = "roman-numerals-py", marker = "python_full_version >= '3.11'" }, - { name = "snowballstemmer", marker = "python_full_version >= '3.11'" }, - { name = "sphinxcontrib-applehelp", marker = "python_full_version >= '3.11'" }, - { name = "sphinxcontrib-devhelp", marker = "python_full_version >= '3.11'" }, - { name = "sphinxcontrib-htmlhelp", marker = "python_full_version >= '3.11'" }, - { name = "sphinxcontrib-jsmath", marker = "python_full_version >= '3.11'" }, - { name = "sphinxcontrib-qthelp", marker = "python_full_version >= '3.11'" }, - { name = "sphinxcontrib-serializinghtml", marker = "python_full_version >= '3.11'" }, + { name = "alabaster" }, + { name = "babel" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "docutils" }, + { name = "imagesize" }, + { name = "jinja2" }, + { name = "packaging" }, + { name = "pygments" }, + { name = "requests" }, + { name = "roman-numerals-py" }, + { name = "snowballstemmer" }, + { name = "sphinxcontrib-applehelp" }, + { name = "sphinxcontrib-devhelp" }, + { name = "sphinxcontrib-htmlhelp" }, + { name = "sphinxcontrib-jsmath" }, + { name = "sphinxcontrib-qthelp" }, + { name = "sphinxcontrib-serializinghtml" }, ] sdist = { url = "https://files.pythonhosted.org/packages/38/ad/4360e50ed56cb483667b8e6dadf2d3fda62359593faabbe749a27c4eaca6/sphinx-8.2.3.tar.gz", hash = "sha256:398ad29dee7f63a75888314e9424d40f52ce5a6a87ae88e7071e80af296ec348", size = 8321876, upload-time = "2025-03-02T22:31:59.658Z" } wheels = [ @@ -2715,7 +2740,7 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "sphinx", version = "8.1.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "sphinx", version = "8.1.3", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/26/f0/43c6a5ff3e7b08a8c3b32f81b859f1b518ccc31e45f22e2b41ced38be7b9/sphinx_autodoc_typehints-3.0.1.tar.gz", hash = "sha256:b9b40dd15dee54f6f810c924f863f9cf1c54f9f3265c495140ea01be7f44fa55", size = 36282, upload-time = "2025-01-16T18:25:30.958Z" } wheels = [ @@ -2731,7 +2756,7 @@ resolution-markers = [ "python_full_version >= '3.11' and python_full_version < '3.13'", ] dependencies = [ - { name = "sphinx", version = "8.2.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "sphinx", version = "8.2.3", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/04/da/40d1ac3a657d967a8d7024d730eb5e29057a2f998f8c5f3df9c2e33917f0/sphinx_autodoc_typehints-3.5.1.tar.gz", hash = "sha256:6114bc788d7b5118712b4f163e0c693e3828c552296007ff1a60ba1606b04718", size = 37620, upload-time = "2025-10-09T18:38:29.378Z" } wheels = [ @@ -3523,6 +3548,36 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" }, ] +[[package]] +name = "urwid" +version = "4.0.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/98/b8/9ed1c288eb7e9236ee83a3f847d15dfa879841219b9a7d174c6c2ef33f53/urwid-4.0.2.tar.gz", hash = "sha256:6962bd04ab98002326b67a431c59b2fb35e8b5abe2e095feda3ee7d8ea8f1228", size = 861918, upload-time = "2026-06-02T18:20:41.867Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/a0/39d524fb8ed3a9facdd2aa4eeb1a2635b3b8689c300989f8cebb989624ba/urwid-4.0.2-py3-none-any.whl", hash = "sha256:ca5958eca20d55535da37810a2e62cbd81a2ce399ee2e93b04a2718a544029eb", size = 295760, upload-time = "2026-06-02T18:20:40.11Z" }, +] + +[[package]] +name = "urwid-readline" +version = "0.15.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "urwid" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ad/70/be318554495555eba7d8ff6e489f6f74ddb225b24086ba4af62a82e723fd/urwid_readline-0.15.1.tar.gz", hash = "sha256:9301444b86d58f7d26388506b704f142cefd193888488b4070d3a0fdfcfc0f84", size = 9007, upload-time = "2024-09-22T17:51:55.144Z" } + +[[package]] +name = "wcwidth" +version = "0.8.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/49/b4/51fe890511f0f242d07cb1ebe6a5b6db417262b9d2568b460347c57d95cc/wcwidth-0.8.1.tar.gz", hash = "sha256:faf5b4a5366a72dc49cad48cdf21f52bdf63bdda995178e483ba247ff79089b9", size = 1466072, upload-time = "2026-06-08T05:57:23.146Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bd/6e/95b0e537de1f4d4301f76f944642c6da50d1511cc7b3d64dc418a66c7509/wcwidth-0.8.1-py3-none-any.whl", hash = "sha256:f453740b1e4a4f3291faa37944c555d71056c4da08d59809b307ef4feba695c8", size = 323092, upload-time = "2026-06-08T05:57:21.413Z" }, +] + [[package]] name = "webencodings" version = "0.5.1"