Skip to content

Commit 6918163

Browse files
committed
fix reconnection handling of CredentialdException - issue #34
Add masked URI to reconnection logic INFO messages to help distinguish multiple connections.
1 parent 96d95cd commit 6918163

File tree

2 files changed

+48
-8
lines changed

2 files changed

+48
-8
lines changed

aiokdb/client.py

Lines changed: 34 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import logging
44
import struct
55
from typing import Any, Optional, Tuple
6-
from urllib.parse import urlparse
6+
from urllib.parse import urlparse, urlunparse
77

88
from aiokdb import cv, logger
99
from aiokdb.server import (
@@ -73,20 +73,47 @@ async def open_qipc_connection(
7373
return q_reader, q_writer
7474

7575

76-
async def maintain_qipc_connection(uri: Optional[str], context: ClientContext) -> None:
76+
async def maintain_qipc_connection(
77+
uri: str, context: ClientContext, raise_bad_creds: bool = False
78+
) -> None:
79+
masked_uri = mask_uri(uri)
7780
while True:
7881
try:
79-
logging.info("attempting connection")
82+
logging.info(f"attempting connection to {masked_uri}")
8083
qr, qw = await open_qipc_connection(uri=uri, context=context)
8184
await qw.writer.wait_closed()
82-
logging.info("connection closed")
85+
logging.info(f"connection closed for {masked_uri}")
8386
except CredentialsException:
84-
raise
87+
logging.warning(f"CredentialsException reported for {masked_uri}")
88+
if raise_bad_creds:
89+
raise
90+
await asyncio.sleep(10)
8591
except Exception:
86-
logging.exception("caught exception")
92+
logging.exception(f"caught exception {masked_uri}")
8793
await asyncio.sleep(10)
8894

89-
logging.info("retry loop exited?")
95+
logging.info(f"retry loop exited for {masked_uri}")
96+
97+
98+
def mask_uri(uri: str) -> str:
99+
"""
100+
Return a version of the URI safe for logging — any password in the
101+
userinfo part (user:pass@) is replaced with '***'.
102+
"""
103+
try:
104+
parsed = urlparse(uri)
105+
# Rebuild netloc with redacted password if needed
106+
userinfo = parsed.username or ""
107+
if parsed.password:
108+
userinfo += ":***"
109+
netloc = f"{parsed.hostname}"
110+
if userinfo:
111+
netloc = f"{userinfo}@{parsed.hostname}"
112+
if parsed.port:
113+
netloc += f":{parsed.port}"
114+
return urlunparse(parsed._replace(netloc=netloc))
115+
except Exception:
116+
return uri
90117

91118

92119
if __name__ == "__main__":

test/test_client_server.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import pytest
55

66
from aiokdb import KException, KObj, MessageType, TypeEnum, cv, kj, kNil
7-
from aiokdb.client import open_qipc_connection
7+
from aiokdb.client import mask_uri, open_qipc_connection
88
from aiokdb.extras import MagicClientContext, MagicServerContext, _string_to_functional
99
from aiokdb.server import CredentialsException, KdbWriter, ServerContext, start_qserver
1010

@@ -271,3 +271,16 @@ async def writer_available(self, dotzw: KdbWriter) -> None:
271271

272272
server.close()
273273
await server.wait_closed()
274+
275+
276+
def test_safe_uri() -> None:
277+
assert (
278+
mask_uri("https://user:[email protected]/path")
279+
== "https://user:***@example.com/path"
280+
)
281+
assert mask_uri("https://example.com/path") == "https://example.com/path"
282+
assert (
283+
mask_uri("postgres://user:topsecret@localhost:5432/db")
284+
== "postgres://user:***@localhost:5432/db"
285+
)
286+
assert mask_uri("kdb://:[email protected]") == "kdb://:***@eg.com"

0 commit comments

Comments
 (0)