fix: static code analysis

Author: ChingEnLin

backend/routes/audit.py

@@ -2,31 +2,32 @@
 from pydantic import BaseModel
 from typing import List, Dict, Any, Optional
 from services.audit_service import process_audit_question
-from services.azure_auth import exchange_token_obo
-from services.gemini_service import VisualizationConfig, AuditSummaryResponse
+from services.gemini_service import VisualizationConfig
 
 router = APIRouter()
 
+
 class AuditQueryRequest(BaseModel):
     question: str
 
+
 class AuditQueryResponse(BaseModel):
     sql_query: str
     results: List[Dict[str, Any]]
     summary: str
     visualization: Optional[VisualizationConfig] = None
 
+
 @router.post("/query", response_model=AuditQueryResponse)
 def query_audit_log(
-    body: AuditQueryRequest = Body(...), 
-    authorization: str = Header(...)
+    body: AuditQueryRequest = Body(...), authorization: str = Header(...)
 ):
     if not authorization.startswith("Bearer "):
         raise HTTPException(status_code=401, detail="Invalid token format")
-    
+
     # We might want to validate the token here even if we don't use it for the pg connection directly yet
     # user_token = authorization.replace("Bearer ", "")
-    # access_token = exchange_token_obo(user_token) 
-    
+    # access_token = exchange_token_obo(user_token)
+
     response = process_audit_question(body.question)
     return AuditQueryResponse(**response)

backend/services/audit_service.py

@@ -1,7 +1,6 @@
 from typing import Dict, Any
 from services.pg_connection import get_connection
 from services.gemini_service import generate_audit_sql, summarize_audit_results
-import json
 
 
 def execute_audit_query(sql_query: str) -> list:
@@ -15,20 +14,20 @@ def execute_audit_query(sql_query: str) -> list:
         conn = get_connection()
         cur = conn.cursor()
         cur.execute(sql_query)
-        
+
         # Get column names
         columns = [desc[0] for desc in cur.description]
         results = [dict(zip(columns, row)) for row in cur.fetchall()]
-        
+
         # Serialize datetime and json objects
         for row in results:
             for key, value in row.items():
-                if hasattr(value, 'isoformat'):
+                if hasattr(value, "isoformat"):
                     row[key] = value.isoformat()
                 elif isinstance(value, dict):
                     # Ensure dicts (like diff_data) are kept as dicts for the frontend
                     pass
-                    
+
         cur.close()
         conn.close()
         return results
@@ -45,30 +44,30 @@ def process_audit_question(question: str) -> Dict[str, Any]:
     3. Summarize results (via Gemini)
     """
     sql_query = generate_audit_sql(question)
-    
+
     # If the generator returned an error query or invalid SQL, return it
     if "Error:" in sql_query:
         return {
             "sql_query": sql_query,
             "results": [],
-            "summary": "Could not generate a valid query for your request."
+            "summary": "Could not generate a valid query for your request.",
         }
-        
+
     results = execute_audit_query(sql_query)
-    
+
     # If execution failed
     if results and "error" in results[0]:
         return {
             "sql_query": sql_query,
             "results": [],
-            "summary": f"Error executing query: {results[0]['error']}"
+            "summary": f"Error executing query: {results[0]['error']}",
         }
-        
+
     summary_response = summarize_audit_results(question, sql_query, results)
-    
+
     return {
         "sql_query": sql_query,
         "results": results,
         "summary": summary_response.summary,
-        "visualization": summary_response.visualization
+        "visualization": summary_response.visualization,
     }

backend/services/gemini_service.py

@@ -2,19 +2,28 @@
 from google.genai import types
 from models.schemas import GeneratedCode, CollectionContext, DebugSuggestionResponse
 from pydantic import BaseModel, Field
-from typing import Optional, List, Any, Dict
+from typing import Optional, List
+
 
 class VisualizationConfig(BaseModel):
     available: bool = Field(description="Whether a chart is recommended for this data")
-    type: Optional[str] = Field(description="Type of chart: 'bar', 'line', 'pie', 'scatter'")
+    type: Optional[str] = Field(
+        description="Type of chart: 'bar', 'line', 'pie', 'scatter'"
+    )
     x_key: Optional[str] = Field(description="Key for X-axis data")
     y_key: Optional[str] = Field(description="Key for Y-axis data")
     title: Optional[str] = Field(description="Title for the chart")
-    data_keys: Optional[List[str]] = Field(description="Keys to include in the chart data points (e.g. ['count', 'date'])")
+    data_keys: Optional[List[str]] = Field(
+        description="Keys to include in the chart data points (e.g. ['count', 'date'])"
+    )
+
 
 class AuditSummaryResponse(BaseModel):
     summary: str = Field(description="Markdown summary of the results")
-    visualization: VisualizationConfig = Field(description="Configuration for data visualization")
+    visualization: VisualizationConfig = Field(
+        description="Configuration for data visualization"
+    )
+
 
 PROMPT_TEMPLATE_QUERY = """
 You are an assistant that converts user requests into MongoDB query code.
@@ -202,7 +211,9 @@ def generate_audit_sql(user_input: str) -> str:
     return sql
 
 
-def summarize_audit_results(user_input: str, sql_query: str, results: list) -> AuditSummaryResponse:
+def summarize_audit_results(
+    user_input: str, sql_query: str, results: list
+) -> AuditSummaryResponse:
     # Truncate results if too large to avoid token limits
     results_str = str(results)[:10000]
     full_prompt = PROMPT_TEMPLATE_AUDIT_SUMMARY.format(
@@ -215,20 +226,21 @@ def summarize_audit_results(user_input: str, sql_query: str, results: list) -> A
         config=types.GenerateContentConfig(
             response_mime_type="application/json",
             response_schema=AuditSummaryResponse,
-            thinking_config=types.ThinkingConfig(thinking_budget=0)
+            thinking_config=types.ThinkingConfig(thinking_budget=0),
         ),
     )
-    
-    if hasattr(response, 'parsed') and response.parsed:
+
+    if hasattr(response, "parsed") and response.parsed:
         return response.parsed
-    
+
     import json
+
     try:
         data = json.loads(response.text)
         return AuditSummaryResponse(**data)
     except Exception as e:
         print(f"Error parsing Gemini response: {e}")
         return AuditSummaryResponse(
             summary="Could not generate summary due to parsing error.",
-            visualization=VisualizationConfig(available=False)
+            visualization=VisualizationConfig(available=False),
         )

backend/tests/test_audit_service.py

@@ -2,47 +2,51 @@
 from unittest.mock import patch, MagicMock
 from services.audit_service import process_audit_question
 
+
 class TestAuditService(unittest.TestCase):
 
-    @patch('services.audit_service.generate_audit_sql')
-    @patch('services.audit_service.get_connection')
-    @patch('services.audit_service.summarize_audit_results')
-    def test_process_audit_question_success(self, mock_summarize, mock_get_conn, mock_generate_sql):
+    @patch("services.audit_service.generate_audit_sql")
+    @patch("services.audit_service.get_connection")
+    @patch("services.audit_service.summarize_audit_results")
+    def test_process_audit_question_success(
+        self, mock_summarize, mock_get_conn, mock_generate_sql
+    ):
         # Setup mocks
         mock_generate_sql.return_value = "SELECT * FROM write_audit_log LIMIT 5"
-        
+
         mock_conn = MagicMock()
         mock_cursor = MagicMock()
         mock_get_conn.return_value = mock_conn
         mock_conn.cursor.return_value = mock_cursor
-        
+
         # Mock DB results
-        mock_cursor.description = [('user_email',), ('operation',)]
-        mock_cursor.fetchall.return_value = [('test@example.com', 'insert')]
-        
+        mock_cursor.description = [("user_email",), ("operation",)]
+        mock_cursor.fetchall.return_value = [("test@example.com", "insert")]
+
         mock_summarize.return_value = "Summary of results."
-        
+
         # Execute
         result = process_audit_question("Show me inserts")
-        
+
         # Assertions
-        self.assertEqual(result['sql_query'], "SELECT * FROM write_audit_log LIMIT 5")
-        self.assertEqual(len(result['results']), 1)
-        self.assertEqual(result['results'][0]['user_email'], 'test@example.com')
-        self.assertEqual(result['summary'], "Summary of results.")
-        
+        self.assertEqual(result["sql_query"], "SELECT * FROM write_audit_log LIMIT 5")
+        self.assertEqual(len(result["results"]), 1)
+        self.assertEqual(result["results"][0]["user_email"], "test@example.com")
+        self.assertEqual(result["summary"], "Summary of results.")
+
         mock_generate_sql.assert_called_once()
         mock_cursor.execute.assert_called_with("SELECT * FROM write_audit_log LIMIT 5")
         mock_summarize.assert_called_once()
 
-    @patch('services.audit_service.generate_audit_sql')
+    @patch("services.audit_service.generate_audit_sql")
     def test_process_audit_question_invalid_sql(self, mock_generate_sql):
         mock_generate_sql.return_value = "DELETE FROM write_audit_log"
-        
+
         result = process_audit_question("Delete everything")
-        
-        self.assertIn("Error executing query", result['summary'])
-        self.assertIn("Only SELECT", result['summary'])
 
-if __name__ == '__main__':
+        self.assertIn("Error executing query", result["summary"])
+        self.assertIn("Only SELECT", result["summary"])
+
+
+if __name__ == "__main__":
     unittest.main()