feat: Generic aiohttp logging (#96)

GuyKh · Apr 10, 2024 · 93c7829 · 93c7829
1 parent d765ed8
commit 93c7829
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 21 deletions.
diff --git a/iec_api/commons.py b/iec_api/commons.py
@@ -7,6 +7,7 @@
 from json import JSONDecodeError
 from typing import Any, Optional
 
+import aiohttp
 import pytz
 from aiohttp import ClientError, ClientResponse, ClientSession, StreamReader
 
@@ -103,7 +104,6 @@ async def send_get_request(
         if not timeout:
             timeout = session.timeout
 
-        logger.debug(f"HTTP GET: {url}")
         resp = await session.get(url=url, headers=headers, timeout=timeout)
         json_resp: dict = await resp.json(content_type=None)
     except TimeoutError as ex:
@@ -113,7 +113,6 @@ async def send_get_request(
     except JSONDecodeError as ex:
         raise IECError(-1, f"Received invalid response from IEC API: {str(ex)}")
 
-    logger.debug(f"HTTP GET Response: {json_resp}")
     if resp.status != http.HTTPStatus.OK:
         parse_error_response(resp, json_resp)
 
@@ -134,9 +133,6 @@ async def send_non_json_get_request(
         if not timeout:
             timeout = session.timeout
 
-        logger.debug(
-            f"HTTP GET: {url}",
-        )
         resp = await session.get(url=url, headers=headers, timeout=timeout)
         resp_content = await resp.text(encoding=encoding)
     except TimeoutError as ex:
@@ -146,8 +142,6 @@ async def send_non_json_get_request(
     except JSONDecodeError as ex:
         raise IECError(-1, f"Received invalid response from IEC API: {str(ex)}")
 
-    logger.debug(f"HTTP GET Response: {resp_content}")
-
     return resp_content
 
 
@@ -166,9 +160,6 @@ async def send_post_request(
         if not timeout:
             headers = session.timeout
 
-        logger.debug(f"HTTP POST: {url}")
-        logger.debug(f"HTTP Content: {data or json_data}")
-
         resp = await session.post(url=url, data=data, json=json_data, headers=headers, timeout=timeout)
 
         json_resp: dict = await resp.json(content_type=None)
@@ -179,8 +170,6 @@ async def send_post_request(
     except JSONDecodeError as ex:
         raise IECError(-1, f"Received invalid response from IEC API: {str(ex)}")
 
-    logger.debug(f"HTTP POST Response: {json_resp}")
-
     if resp.status != http.HTTPStatus.OK:
         parse_error_response(resp, json_resp)
     return json_resp
@@ -201,9 +190,6 @@ async def send_non_json_post_request(
         if not timeout:
             headers = session.timeout
 
-        logger.debug(f"HTTP POST: {url}")
-        logger.debug(f"HTTP Content: {data or json_data}")
-
         resp = await session.post(url=url, data=data, json=json_data, headers=headers, timeout=timeout)
     except TimeoutError as ex:
         raise IECError(-1, f"Failed to communicate with IEC API due to time out: ({str(ex)})")
@@ -231,3 +217,18 @@ def convert_to_tz_aware_datetime(dt: Optional[datetime]) -> Optional[datetime]:
         return TIMEZONE.localize(dt)
     else:
         return dt.replace(tzinfo=pytz.utc)
+
+
+async def on_request_start_debug(session: aiohttp.ClientSession, context, params: aiohttp.TraceRequestStartParams):
+    logger.debug(f"HTTP {params.method}: {params.url}")
+
+
+async def on_request_chunk_sent_debug(
+    session: aiohttp.ClientSession, context, params: aiohttp.TraceRequestChunkSentParams
+):
+    if (params.method == "POST" or params.method == "PUT") and params.chunk:
+        logger.debug(f"HTTP Content {params.method}: {params.chunk}")
+
+
+async def on_request_end_debug(session: aiohttp.ClientSession, context, params: aiohttp.TraceRequestEndParams):
+    logger.debug(f"HTTP {params.method} Response <{params.response.status}>: {await params.response.text()}")
diff --git a/iec_api/data.py b/iec_api/data.py
@@ -85,7 +85,6 @@ async def get_customer(session: ClientSession, token: JWT) -> Optional[Customer]
     # sending get request and saving the response as response object
     response = await commons.send_get_request(session=session, url=GET_CONSUMER_URL, headers=headers)
 
-    logger.debug(f"Response: {response}")
     return Customer.from_dict(response)
 
 

diff --git a/iec_api/iec_client.py b/iec_api/iec_client.py
@@ -8,8 +8,7 @@
 import jwt
 from aiohttp import ClientSession
 
-from iec_api import data, login
-from iec_api.commons import is_valid_israeli_id
+from iec_api import commons, data, login
 from iec_api.models.contract import Contract
 from iec_api.models.customer import Account, Customer
 from iec_api.models.device import Device, Devices
@@ -38,13 +37,22 @@ def __init__(self, user_id: str | int, session: Optional[ClientSession] = None):
         """
 
         self._kwh_tariff: Optional[float] = None
-        self._kwh_tariff_fetch_time: Optional[datetime.datetime] = None
-        if not is_valid_israeli_id(user_id):
+        self._kwh_tariff_fetch_time: Optional[datetime] = None
+        if not commons.is_valid_israeli_id(user_id):
             raise ValueError("User ID must be a valid Israeli ID.")
 
+        # Custom Logger to the session
+        trace_config = aiohttp.TraceConfig()
+        trace_config.on_request_start.append(commons.on_request_start_debug)
+        trace_config.on_request_chunk_sent.append(commons.on_request_chunk_sent_debug)
+        trace_config.on_request_end.append(commons.on_request_end_debug)
+        trace_config.freeze()
+
         if not session:
-            session = aiohttp.ClientSession()
+            session = aiohttp.ClientSession(trace_configs=[trace_config])
             atexit.register(self._shutdown)
+        else:
+            session.trace_configs.append(trace_config)
 
         self._session = session