Add ensure_connected method to client_transport (#86)

Why
===

There's currently no way of making sure we're connected other than
sending an RPC. This poses issues where we have errors propagating much
later than they should be.

What changed
============

Add method to allow us to check if we're connected and error as early as
possible

Author: masad-frost

replit_river/client.py

@@ -42,6 +42,9 @@ async def close(self) -> None:
         await self._transport.close()
         logger.info(f"river client {self._client_id} closed")
 
+    async def ensure_connected(self) -> None:
+        await self._transport.get_or_create_session()
+
     async def send_rpc(
         self,
         service_name: str,
@@ -51,7 +54,7 @@ async def send_rpc(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
     ) -> ResponseType:
-        session = await self._transport._get_or_create_session()
+        session = await self._transport.get_or_create_session()
         return await session.send_rpc(
             service_name,
             procedure_name,
@@ -72,7 +75,7 @@ async def send_upload(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
     ) -> ResponseType:
-        session = await self._transport._get_or_create_session()
+        session = await self._transport.get_or_create_session()
         return await session.send_upload(
             service_name,
             procedure_name,
@@ -93,7 +96,7 @@ async def send_subscription(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
     ) -> AsyncIterator[Union[ResponseType, ErrorType]]:
-        session = await self._transport._get_or_create_session()
+        session = await self._transport.get_or_create_session()
         return session.send_subscription(
             service_name,
             procedure_name,
@@ -114,7 +117,7 @@ async def send_stream(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
     ) -> AsyncIterator[Union[ResponseType, ErrorType]]:
-        session = await self._transport._get_or_create_session()
+        session = await self._transport.get_or_create_session()
         return session.send_stream(
             service_name,
             procedure_name,

replit_river/client_transport.py

@@ -73,6 +73,32 @@ async def close(self) -> None:
         self._rate_limiter.close()
         await self._close_all_sessions()
 
+    async def get_or_create_session(self) -> ClientSession:
+        async with self._create_session_lock:
+            existing_session = await self._get_existing_session()
+            if not existing_session:
+                return await self._create_new_session()
+            is_session_open = await existing_session.is_session_open()
+            if not is_session_open:
+                return await self._create_new_session()
+            is_ws_open = await existing_session.is_websocket_open()
+            if is_ws_open:
+                return existing_session
+            new_ws, _, hs_response = await self._establish_new_connection(
+                existing_session
+            )
+            if hs_response.status.sessionId == existing_session.session_id:
+                logger.info(
+                    "Replacing ws connection in session id %s",
+                    existing_session.session_id,
+                )
+                await existing_session.replace_with_new_websocket(new_ws)
+                return existing_session
+            else:
+                logger.info("Closing stale session %s", existing_session.session_id)
+                await existing_session.close()
+                return await self._create_new_session()
+
     async def _get_existing_session(self) -> Optional[ClientSession]:
         async with self._session_lock:
             if not self._sessions:
@@ -190,33 +216,7 @@ async def _create_new_session(
     async def _retry_connection(self) -> ClientSession:
         if not self._transport_options.transparent_reconnect:
             await self._close_all_sessions()
-        return await self._get_or_create_session()
-
-    async def _get_or_create_session(self) -> ClientSession:
-        async with self._create_session_lock:
-            existing_session = await self._get_existing_session()
-            if not existing_session:
-                return await self._create_new_session()
-            is_session_open = await existing_session.is_session_open()
-            if not is_session_open:
-                return await self._create_new_session()
-            is_ws_open = await existing_session.is_websocket_open()
-            if is_ws_open:
-                return existing_session
-            new_ws, _, hs_response = await self._establish_new_connection(
-                existing_session
-            )
-            if hs_response.status.sessionId == existing_session.session_id:
-                logger.info(
-                    "Replacing ws connection in session id %s",
-                    existing_session.session_id,
-                )
-                await existing_session.replace_with_new_websocket(new_ws)
-                return existing_session
-            else:
-                logger.info("Closing stale session %s", existing_session.session_id)
-                await existing_session.close()
-                return await self._create_new_session()
+        return await self.get_or_create_session()
 
     async def _send_handshake_request(
         self,