From 67a4d3d394e39086928d3edc6be0008ca29feb34 Mon Sep 17 00:00:00 2001 From: Tobie Morgan Hitchcock Date: Sun, 2 Feb 2025 00:32:40 +0000 Subject: [PATCH 1/2] Ensure WebSocket max message size is not limited --- src/surrealdb/__init__.py | 20 ++++++-------------- src/surrealdb/connections/async_http.py | 1 - src/surrealdb/connections/async_ws.py | 18 +++++------------- src/surrealdb/connections/blocking_ws.py | 15 +++++---------- 4 files changed, 16 insertions(+), 38 deletions(-) diff --git a/src/surrealdb/__init__.py b/src/surrealdb/__init__.py index 2e47afcf..d19f704a 100644 --- a/src/surrealdb/__init__.py +++ b/src/surrealdb/__init__.py @@ -28,14 +28,10 @@ def __call__(cls, *args, **kwargs): constructed_url = Url(url) - - # Extract `max_size` with a default if not explicitly provided - max_size = kwargs.get("max_size", 2 ** 20) - if constructed_url.scheme == UrlScheme.HTTP or constructed_url.scheme == UrlScheme.HTTPS: return AsyncHttpSurrealConnection(url=url) elif constructed_url.scheme == UrlScheme.WS or constructed_url.scheme == UrlScheme.WSS: - return AsyncWsSurrealConnection(url=url, max_size=max_size) + return AsyncWsSurrealConnection(url=url) else: raise ValueError(f"Unsupported protocol in URL: {url}. Use 'ws://' or 'http://'.") @@ -54,32 +50,28 @@ def __call__(cls, *args, **kwargs): constructed_url = Url(url) - - # Extract `max_size` with a default if not explicitly provided - max_size = kwargs.get("max_size", 2 ** 20) - if constructed_url.scheme == UrlScheme.HTTP or constructed_url.scheme == UrlScheme.HTTPS: return BlockingHttpSurrealConnection(url=url) elif constructed_url.scheme == UrlScheme.WS or constructed_url.scheme == UrlScheme.WSS: - return BlockingWsSurrealConnection(url=url, max_size=max_size) + return BlockingWsSurrealConnection(url=url) else: raise ValueError(f"Unsupported protocol in URL: {url}. Use 'ws://' or 'http://'.") -def Surreal(url: Optional[str] = None, max_size: int = 2 ** 20) -> Union[BlockingWsSurrealConnection, BlockingHttpSurrealConnection]: +def Surreal(url: Optional[str] = None) -> Union[BlockingWsSurrealConnection, BlockingHttpSurrealConnection]: constructed_url = Url(url) if constructed_url.scheme == UrlScheme.HTTP or constructed_url.scheme == UrlScheme.HTTPS: return BlockingHttpSurrealConnection(url=url) elif constructed_url.scheme == UrlScheme.WS or constructed_url.scheme == UrlScheme.WSS: - return BlockingWsSurrealConnection(url=url, max_size=max_size) + return BlockingWsSurrealConnection(url=url) else: raise ValueError(f"Unsupported protocol in URL: {url}. Use 'ws://' or 'http://'.") -def AsyncSurreal(url: Optional[str] = None, max_size: int = 2 ** 20) -> Union[AsyncWsSurrealConnection, AsyncHttpSurrealConnection]: +def AsyncSurreal(url: Optional[str] = None) -> Union[AsyncWsSurrealConnection, AsyncHttpSurrealConnection]: constructed_url = Url(url) if constructed_url.scheme == UrlScheme.HTTP or constructed_url.scheme == UrlScheme.HTTPS: return AsyncHttpSurrealConnection(url=url) elif constructed_url.scheme == UrlScheme.WS or constructed_url.scheme == UrlScheme.WSS: - return AsyncWsSurrealConnection(url=url, max_size=max_size) + return AsyncWsSurrealConnection(url=url) else: raise ValueError(f"Unsupported protocol in URL: {url}. Use 'ws://' or 'http://'.") diff --git a/src/surrealdb/connections/async_http.py b/src/surrealdb/connections/async_http.py index 5f6ca293..6501f5ad 100644 --- a/src/surrealdb/connections/async_http.py +++ b/src/surrealdb/connections/async_http.py @@ -22,7 +22,6 @@ class AsyncHttpSurrealConnection(AsyncTemplate, UtilsMixin): Attributes: url: The URL of the database to process queries for. - max_size: The maximum size of the connection payload. id: The ID of the connection. """ diff --git a/src/surrealdb/connections/async_ws.py b/src/surrealdb/connections/async_ws.py index 97b062b4..ab99dbda 100644 --- a/src/surrealdb/connections/async_ws.py +++ b/src/surrealdb/connections/async_ws.py @@ -23,35 +23,27 @@ class AsyncWsSurrealConnection(AsyncTemplate, UtilsMixin): """ A single async connection to a SurrealDB instance. To be used once and discarded. - # Notes - A new connection is created for each query. This is because the async websocket connection is - dropped - Attributes: url: The URL of the database to process queries for. user: The username to login on. password: The password to login on. namespace: The namespace that the connection will stick to. database: The database that the connection will stick to. - max_size: The maximum size of the connection. id: The ID of the connection. """ def __init__( self, url: str, - max_size: int = 2 ** 20, ) -> None: """ The constructor for the AsyncSurrealConnection class. :param url: The URL of the database to process queries for. - :param max_size: The maximum size of the connection. """ self.url: Url = Url(url) self.raw_url: str = f"{self.url.raw_url}/rpc" self.host: str = self.url.hostname self.port: int = self.url.port - self.max_size: int = max_size self.id: str = str(uuid.uuid4()) self.token: Optional[str] = None self.socket = None @@ -64,19 +56,18 @@ async def _send(self, message: RequestMessage, process: str, bypass: bool = Fals self.check_response_for_error(response, process) return response - async def connect(self, url: Optional[str] = None, max_size: Optional[int] = None) -> None: + async def connect(self, url: Optional[str] = None) -> None: # overwrite params if passed in if url is not None: self.url = Url(url) self.raw_url: str = f"{self.url.raw_url}/rpc" self.host: str = self.url.hostname self.port: int = self.url.port - if max_size is not None: - self.max_size = max_size if self.socket is None: self.socket = await websockets.connect( self.raw_url, - max_size=self.max_size, + max_size=None, + write_limit=None, subprotocols=[websockets.Subprotocol("cbor")] ) @@ -361,7 +352,8 @@ async def __aenter__(self) -> "AsyncWsSurrealConnection": """ self.socket = await websockets.connect( self.raw_url, - max_size=self.max_size, + max_size=None, + write_limit=None, subprotocols=[websockets.Subprotocol("cbor")] ) return self diff --git a/src/surrealdb/connections/blocking_ws.py b/src/surrealdb/connections/blocking_ws.py index f9a1d824..bf70b133 100644 --- a/src/surrealdb/connections/blocking_ws.py +++ b/src/surrealdb/connections/blocking_ws.py @@ -22,32 +22,25 @@ class BlockingWsSurrealConnection(SyncTemplate, UtilsMixin): """ A single blocking connection to a SurrealDB instance. To be used once and discarded. - # Notes - A new connection is created for each query. This is because the WebSocket connection is - dropped after the query is completed. - Attributes: url: The URL of the database to process queries for. user: The username to login on. password: The password to login on. namespace: The namespace that the connection will stick to. database: The database that the connection will stick to. - max_size: The maximum size of the connection. id: The ID of the connection. """ - def __init__(self, url: str, max_size: int = 2 ** 20) -> None: + def __init__(self, url: str) -> None: """ The constructor for the BlockingWsSurrealConnection class. :param url: (str) the URL of the database to process queries for. - :param max_size: (int) The maximum size of the connection. """ self.url: Url = Url(url) self.raw_url: str = f"{self.url.raw_url}/rpc" self.host: str = self.url.hostname self.port: int = self.url.port - self.max_size: int = max_size self.id: str = str(uuid.uuid4()) self.token: Optional[str] = None self.socket = None @@ -56,7 +49,8 @@ def _send(self, message: RequestMessage, process: str, bypass: bool = False) -> if self.socket is None: self.socket = ws_sync.connect( self.raw_url, - max_size=self.max_size, + max_size=None, + write_limit=None, subprotocols=[websockets.Subprotocol("cbor")], ) self.socket.send(message.WS_CBOR_DESCRIPTOR) @@ -361,7 +355,8 @@ def __enter__(self) -> "BlockingWsSurrealConnection": """ self.socket = ws_sync.connect( self.raw_url, - max_size=self.max_size, + max_size=None, + write_limit=None, subprotocols=[websockets.Subprotocol("cbor")] ) return self From c3f32f63d75c033d36842716a5a5d761fe377e7f Mon Sep 17 00:00:00 2001 From: Tobie Morgan Hitchcock Date: Sun, 2 Feb 2025 00:38:41 +0000 Subject: [PATCH 2/2] Remove incorrect config option --- src/surrealdb/connections/async_ws.py | 2 -- src/surrealdb/connections/blocking_ws.py | 2 -- 2 files changed, 4 deletions(-) diff --git a/src/surrealdb/connections/async_ws.py b/src/surrealdb/connections/async_ws.py index ab99dbda..2cef12b1 100644 --- a/src/surrealdb/connections/async_ws.py +++ b/src/surrealdb/connections/async_ws.py @@ -67,7 +67,6 @@ async def connect(self, url: Optional[str] = None) -> None: self.socket = await websockets.connect( self.raw_url, max_size=None, - write_limit=None, subprotocols=[websockets.Subprotocol("cbor")] ) @@ -353,7 +352,6 @@ async def __aenter__(self) -> "AsyncWsSurrealConnection": self.socket = await websockets.connect( self.raw_url, max_size=None, - write_limit=None, subprotocols=[websockets.Subprotocol("cbor")] ) return self diff --git a/src/surrealdb/connections/blocking_ws.py b/src/surrealdb/connections/blocking_ws.py index bf70b133..a9d0d9ef 100644 --- a/src/surrealdb/connections/blocking_ws.py +++ b/src/surrealdb/connections/blocking_ws.py @@ -50,7 +50,6 @@ def _send(self, message: RequestMessage, process: str, bypass: bool = False) -> self.socket = ws_sync.connect( self.raw_url, max_size=None, - write_limit=None, subprotocols=[websockets.Subprotocol("cbor")], ) self.socket.send(message.WS_CBOR_DESCRIPTOR) @@ -356,7 +355,6 @@ def __enter__(self) -> "BlockingWsSurrealConnection": self.socket = ws_sync.connect( self.raw_url, max_size=None, - write_limit=None, subprotocols=[websockets.Subprotocol("cbor")] ) return self