|
3 | 3 | import logging |
4 | 4 | import struct |
5 | 5 | from typing import Any, Optional, Tuple |
6 | | -from urllib.parse import urlparse |
| 6 | +from urllib.parse import urlparse, urlunparse |
7 | 7 |
|
8 | 8 | from aiokdb import cv, logger |
9 | 9 | from aiokdb.server import ( |
@@ -73,20 +73,47 @@ async def open_qipc_connection( |
73 | 73 | return q_reader, q_writer |
74 | 74 |
|
75 | 75 |
|
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) |
77 | 80 | while True: |
78 | 81 | try: |
79 | | - logging.info("attempting connection") |
| 82 | + logging.info(f"attempting connection to {masked_uri}") |
80 | 83 | qr, qw = await open_qipc_connection(uri=uri, context=context) |
81 | 84 | await qw.writer.wait_closed() |
82 | | - logging.info("connection closed") |
| 85 | + logging.info(f"connection closed for {masked_uri}") |
83 | 86 | 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) |
85 | 91 | except Exception: |
86 | | - logging.exception("caught exception") |
| 92 | + logging.exception(f"caught exception {masked_uri}") |
87 | 93 | await asyncio.sleep(10) |
88 | 94 |
|
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 |
90 | 117 |
|
91 | 118 |
|
92 | 119 | if __name__ == "__main__": |
|
0 commit comments