wip: make aiokdb.cli use the reconnecting client, to make it easier to discover any reconnection issues

Author: shuckc

aiokdb/cli.py

@@ -3,17 +3,32 @@
 import logging
 import os
 import traceback
-from typing import Any
+from typing import Any, Optional
 
 from prompt_toolkit.history import FileHistory
 from prompt_toolkit.patch_stdout import patch_stdout
 from prompt_toolkit.shortcuts import PromptSession
 
 from aiokdb import TypeEnum, cv
-from aiokdb.client import open_qipc_connection
+from aiokdb.client import (
+    ClientContext,
+    KdbWriter,
+    maintain_qipc_connection,
+)
 from aiokdb.format import AsciiFormatter
 
 
+class CliClientContext(ClientContext):
+    def __init__(self) -> None:
+        self.writer: Optional[KdbWriter] = None
+
+    async def writer_available(self, dotzw: KdbWriter) -> None:
+        self.writer = dotzw
+
+    def writer_closing(self, dotzw: KdbWriter) -> None:
+        self.writer = None
+
+
 async def main(args: Any) -> None:
     history = FileHistory(os.path.expanduser("~/.aiokdb-cli-history"))
     session: Any = PromptSession("(eval) > ", history=history)
@@ -23,8 +38,12 @@ async def main(args: Any) -> None:
     if password is None:
         password = await session.prompt_async("Password:", is_password=True)
 
-    r, w = await open_qipc_connection(
-        host=args.host, port=args.port, user=args.user, password=password
+    cc = CliClientContext()
+
+    task = asyncio.create_task(
+        maintain_qipc_connection(
+            uri=f"kdb://{args.user}:{password}@{args.host}:{args.port}", context=cc
+        )
     )
 
     # Run echo loop. Read text from stdin, and reply it back.
@@ -33,15 +52,22 @@ async def main(args: Any) -> None:
             inp = await session.prompt_async("q)", is_password=False)
             if inp == "":
                 continue
-            output = await w.sync_req(cv(inp))
-            if output.t != TypeEnum.NIL:
-                print(fmt.format(output))
+            if not cc.writer:
+                print("Writer not connected, wait for re-connect")
+            else:
+                output = await cc.writer.sync_req(cv(inp))
+                if output.t != TypeEnum.NIL:
+                    print(fmt.format(output))
         except KeyboardInterrupt:
             return
         except EOFError:
             break
         except Exception:
             traceback.print_exc(limit=-2)
+
+    task.cancel()
+    await task
+
     return None
 
 

aiokdb/client.py

@@ -100,20 +100,17 @@ def mask_uri(uri: str) -> str:
     Return a version of the URI safe for logging — any password in the
     userinfo part (user:pass@) is replaced with '***'.
     """
-    try:
-        parsed = urlparse(uri)
-        # Rebuild netloc with redacted password if needed
-        userinfo = parsed.username or ""
-        if parsed.password:
-            userinfo += ":***"
-        netloc = f"{parsed.hostname}"
-        if userinfo:
-            netloc = f"{userinfo}@{parsed.hostname}"
-        if parsed.port:
-            netloc += f":{parsed.port}"
-        return urlunparse(parsed._replace(netloc=netloc))
-    except Exception:
-        return uri
+    parsed = urlparse(uri)
+    # Rebuild netloc with redacted password if needed
+    userinfo = parsed.username or ""
+    if parsed.password:
+        userinfo += ":***"
+    netloc = f"{parsed.hostname}"
+    if userinfo:
+        netloc = f"{userinfo}@{parsed.hostname}"
+    if parsed.port:
+        netloc += f":{parsed.port}"
+    return urlunparse(parsed._replace(netloc=netloc))
 
 
 if __name__ == "__main__":

aiokdb/server.py

@@ -7,7 +7,7 @@
 from functools import partial
 from typing import Any, Callable, List, Optional, Tuple
 
-from aiokdb import KException, KObj, MessageType, TypeEnum, b9, d9, krr, logger
+from aiokdb import KException, KObj, MessageType, TypeEnum, b9, d9, kNil, krr, logger
 
 
 class CredentialsException(Exception):
@@ -308,7 +308,14 @@ async def start_qserver(port: int, context: ServerContext) -> Any:
 
 
 async def main(qpassword: str, qport: int) -> None:
-    context = ServerContext(qpassword)
+    class DelayedServerContext(ServerContext):
+        async def on_sync_request(self, cmd: KObj, dotzw: KdbWriter) -> KObj:  # .z.pg
+            logging.info("Got request")
+            await asyncio.sleep(1)
+            logging.info("sending response")
+            return kNil
+
+    context = DelayedServerContext(qpassword)
     server = await start_qserver(qport, context)
     async with server:
         await server.serve_forever()