fix: Logger initialization

Fixes issue where logger was initialized before app initialization
vectra-ai-research · Jan 11, 2025 · 7f3448d · 7f3448d
1 parent a389f3b
commit 7f3448d
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 24 deletions.
diff --git a/core/entra/entra_token_manager.py b/core/entra/entra_token_manager.py
@@ -3,13 +3,10 @@
 import time
 import datetime
 import requests
-import logging
 from typing import Optional, Dict, Tuple, Union
-from core.Constants import MSFT_TOKENS_FILE, SERVER_LOG_FILE
+from core.Constants import MSFT_TOKENS_FILE
 from core.entra.token_info import Msft_Token
-
-token_logger = logging.getLogger(__name__)
-logging.basicConfig(filename=SERVER_LOG_FILE, format='%(asctime)s - %(levelname)s - %(message)s', level=logging.WARNING)
+from core.logging.logger import graph_logger
 
 class TokenRefreshError(Exception):
     """Custom exception for token refresh failures"""
@@ -48,7 +45,7 @@ def _token_monitor(self):
             try:
                 self._check_and_refresh_tokens()
             except Exception as e:
-                token_logger.error(f"Error in token monitor: {str(e)}")
+                graph_logger.error(f"Error in token monitor: {str(e)}")
             time.sleep(self.check_interval)
 
     def _check_and_refresh_tokens(self):
@@ -75,16 +72,16 @@ def _check_and_refresh_tokens(self):
                     if refresh_token:
                         tokens_to_refresh.append((access_token, refresh_token, token_info))
                     else:
-                        token_logger.info(f"Token near expiration but no refresh token available: {access_token[:10]}...")
+                        graph_logger.info(f"Token near expiration but no refresh token available: {access_token[:10]}...")
             except Exception as e:
-                token_logger.error(f"Error checking token expiration: {str(e)}")
+                graph_logger.error(f"Error checking token expiration: {str(e)}")
 
         # Refresh collected tokens
         for access_token, refresh_token, token_info in tokens_to_refresh:
             try:
                 self._refresh_token(access_token, refresh_token, token_info)
             except TokenRefreshError as e:
-                token_logger.error(f"Failed to refresh token: {str(e)}")
+                graph_logger.error(f"Failed to refresh token: {str(e)}")
 
     def _refresh_token(self, access_token: str, refresh_token: str, token_info: Dict) -> None:
         """
@@ -124,7 +121,7 @@ def _refresh_token(self, access_token: str, refresh_token: str, token_info: Dict
 
                 # Update token in storage
                 self._update_token(access_token, new_access_token, new_refresh_token)
-                token_logger.info(f"Successfully refreshed token: {access_token[:10]}...")
+                graph_logger.info(f"Successfully refreshed token: {access_token[:10]}...")
             else:
                 error_desc = response.json().get('error_description', 'Unknown error')
                 raise TokenRefreshError(f"Token refresh failed: {error_desc}")

diff --git a/core/entra/graph_request.py b/core/entra/graph_request.py
@@ -3,12 +3,8 @@
 import re
 import time
 from typing import Optional, Dict, Any, Union, List
-import logging
 from .entra_token_manager import EntraTokenManager
-from core.Constants import SERVER_LOG_FILE
-
-logger = logging.getLogger(__name__)
-logging.basicConfig(filename=SERVER_LOG_FILE, format='%(asctime)s - %(levelname)s - %(message)s', level=logging.WARNING)
+from core.logging.logger import graph_logger
 
 class RateLimiter:
     """Handles rate limiting for Graph API requests"""
@@ -37,14 +33,14 @@ def wait_if_needed(self, response: Optional[requests.Response] = None) -> None:
 
             if retry_after is not None:
                 self._throttle_until = current_time + int(retry_after)
-                logger.warning(f"Rate limit hit. Waiting {retry_after} seconds.")
+                graph_logger.warning(f"Rate limit hit. Waiting {retry_after} seconds.")
                 time.sleep(int(retry_after))
                 return
 
         # If in a throttle window, wait
         if self._throttle_until and current_time < self._throttle_until:
             wait_time = self._throttle_until - current_time
-            logger.debug(f"In throttle window. Waiting {wait_time:.2f} seconds.")
+            graph_logger.debug(f"In throttle window. Waiting {wait_time:.2f} seconds.")
             time.sleep(wait_time)
             return
 
@@ -97,15 +93,15 @@ def _make_request(self, method: str, url: str, **kwargs) -> requests.Response:
             self.rate_limiter.wait_if_needed(response)
             if response.status_code == 429:  # Too Many Requests
                 retry_after = int(response.headers.get('Retry-After', 30))
-                logger.warning(f"Rate limit exceeded. Retrying after {retry_after} seconds")
+                graph_logger.warning(f"Rate limit exceeded. Retrying after {retry_after} seconds")
                 time.sleep(retry_after)
                 return self._make_request(method, url, **kwargs)
 
             response.raise_for_status()
             return response
 
         except requests.exceptions.RequestException as e:
-            logger.error(f"Request failed: {str(e)}")
+            graph_logger.error(f"Request failed: {str(e)}")
             return response
 
     def get(self, url: str, params: Optional[Dict] = None, 
@@ -141,7 +137,7 @@ def get(self, url: str, params: Optional[Dict] = None,
                 url = result.get('@odata.nextLink') if pagination else None
 
             except Exception as e:
-                logger.error(f"GET request failed: {str(e)}")
+                graph_logger.error(f"GET request failed: {str(e)}")
                 return e
 
         return graph_results
@@ -154,7 +150,7 @@ def post(self, url: str, data: Dict[str, Any],
             return self._make_request('POST', url, headers=headers, 
                                     data=json.dumps(data))
         except Exception as e:
-            logger.error(f"POST request failed: {str(e)}")
+            graph_logger.error(f"POST request failed: {str(e)}")
             return e
 
     def delete(self, url: str, access_token: Optional[str] = None) -> requests.Response:
@@ -163,7 +159,7 @@ def delete(self, url: str, access_token: Optional[str] = None) -> requests.Respo
         try:
             return self._make_request('DELETE', url, headers=headers)
         except Exception as e:
-            logger.error(f"DELETE request failed: {str(e)}")
+            graph_logger.error(f"DELETE request failed: {str(e)}")
             return e
 
     def patch(self, url: str, data: Dict[str, Any], 
@@ -174,7 +170,7 @@ def patch(self, url: str, data: Dict[str, Any],
             return self._make_request('PATCH', url, headers=headers, 
                                     data=json.dumps(data))
         except Exception as e:
-            logger.error(f"PATCH request failed: {str(e)}")
+            graph_logger.error(f"PATCH request failed: {str(e)}")
             return e
 
     def put(self, url: str, data: Dict[str, Any], 
@@ -186,7 +182,7 @@ def put(self, url: str, data: Dict[str, Any],
                                         data=json.dumps(data))
             return response.json()
         except Exception as e:
-            logger.error(f"PUT request failed: {str(e)}")
+            graph_logger.error(f"PUT request failed: {str(e)}")
             return e
 
     @staticmethod