Update services/report_data_handler.py

services/report_data_handler.py

@@ -73,13 +73,18 @@ def save_report_results(
         logger.info(f"Attempting to save agentic analysis to Bubble for org_urn: {org_urn}")
         response = bulk_upload_to_bubble([payload], BUBBLE_REPORT_TABLE_NAME)
 
-        # Assuming bulk_upload_to_bubble returns a list of IDs on success and None or False on failure
-        if response:
+        # Handle API response which could be a list of dicts (for bulk) or a single dict.
+        if response and isinstance(response, list) and len(response) > 0 and isinstance(response[0], dict) and 'id' in response[0]:
+            record_id = response[0]['id'] # Get the ID from the first dictionary in the list
+            logger.info(f"Successfully saved agentic analysis to Bubble. Record ID: {record_id}")
+            return record_id
+        elif response and isinstance(response, dict) and "id" in response: # Handle non-bulk response
             record_id = response["id"]
             logger.info(f"Successfully saved agentic analysis to Bubble. Record ID: {record_id}")
             return record_id
         else:
-            logger.error(f"Failed to save agentic analysis to Bubble. Response: {response}")
+            # Catches None, False, empty lists, or other unexpected formats.
+            logger.error(f"Failed to save agentic analysis to Bubble. Unexpected API Response: {response}")
             return None
             
     except Exception as e:
@@ -110,24 +115,33 @@ def save_objectives(
                 "timeline": obj.get("objective_timeline"),
                 "owner": obj.get("objective_owner"),
                 "report": report_id,
-                 # "organization_urn": org_urn # Assuming 'report' links to the org
             }
             for obj in objectives_data
         ]
 
-        logger.info(f"objectives data {objectives_data}")
+        logger.info(f"Attempting to save {len(payloads)} objectives for report_id: {report_id}")
+        response_data = bulk_upload_to_bubble(payloads, BUBBLE_OKR_TABLE_NAME)
 
-        logger.info(f"payload {payloads}")
+        # Validate response and extract IDs from the list of dictionaries.
+        if not response_data or not isinstance(response_data, list):
+            logger.error(f"Failed to save objectives. API response was not a list: {response_data}")
+            return None
         
-        logger.info(f"Attempting to save {len(payloads)} objectives for report_id: {report_id}")
-        objective_ids = bulk_upload_to_bubble(payloads, BUBBLE_OKR_TABLE_NAME) # Corrected table name
+        try:
+            # Extract the ID from each dictionary in the response list.
+            extracted_ids = [item['id'] for item in response_data]
+        except (TypeError, KeyError):
+            logger.error(f"Failed to parse IDs from API response. Response format invalid: {response_data}", exc_info=True)
+            return None
 
-        if objective_ids is None:
-            logger.error(f"Failed to save objectives to Bubble for report_id: {report_id}. The upload function returned None.")
+        # Check if we extracted the expected number of IDs
+        if len(extracted_ids) != len(payloads):
+            logger.error(f"Failed to save all objectives for report_id: {report_id}. "
+                         f"Expected {len(payloads)} IDs, but got {len(extracted_ids)} from response: {response_data}")
             return None
 
-        logger.info(f"Successfully saved {len(objective_ids)} objectives.")
-        return objective_ids
+        logger.info(f"Successfully saved {len(extracted_ids)} objectives.")
+        return extracted_ids
         
     except Exception as e:
         logger.exception(f"An unexpected error occurred in save_objectives for report_id {report_id}: {e}")
@@ -153,6 +167,11 @@ def save_key_results(
 
     try:
         for objective_data, parent_objective_id in objectives_with_ids:
+            # Defensive check to ensure the parent_objective_id is a valid-looking string.
+            if not isinstance(parent_objective_id, str) or not parent_objective_id:
+                 logger.error(f"Invalid parent_objective_id found: '{parent_objective_id}'. Skipping KRs for this objective.")
+                 continue # Skip this loop iteration
+
             for kr in objective_data.get("key_results", []):
                 key_results_to_process.append(kr)
                 key_result_payloads.append({
@@ -169,14 +188,26 @@ def save_key_results(
             return []
 
         logger.info(f"Attempting to save {len(key_result_payloads)} key results for org_urn: {org_urn}")
-        key_result_ids = bulk_upload_to_bubble(key_result_payloads, BUBBLE_KEY_RESULTS_TABLE_NAME)
+        response_data = bulk_upload_to_bubble(key_result_payloads, BUBBLE_KEY_RESULTS_TABLE_NAME)
 
-        if key_result_ids is None:
-            logger.error(f"Failed to save key results to Bubble for org_urn: {org_urn}.")
+        # Validate response and extract IDs.
+        if not response_data or not isinstance(response_data, list):
+            logger.error(f"Failed to save key results. API response was not a list: {response_data}")
+            return None
+        
+        try:
+            extracted_ids = [item['id'] for item in response_data]
+        except (TypeError, KeyError):
+            logger.error(f"Failed to parse IDs from key result API response: {response_data}", exc_info=True)
+            return None
+
+        if len(extracted_ids) != len(key_result_payloads):
+            logger.error(f"Failed to save all key results for org_urn: {org_urn}. "
+                         f"Expected {len(key_result_payloads)} IDs, but got {len(extracted_ids)} from response: {response_data}")
             return None
 
-        logger.info(f"Successfully saved {len(key_result_ids)} key results.")
-        return list(zip(key_results_to_process, key_result_ids))
+        logger.info(f"Successfully saved {len(extracted_ids)} key results.")
+        return list(zip(key_results_to_process, extracted_ids))
 
     except Exception as e:
         logger.exception(f"An unexpected error occurred in save_key_results for org_urn {org_urn}: {e}")
@@ -220,14 +251,26 @@ def save_tasks(
             return []
 
         logger.info(f"Attempting to save {len(task_payloads)} tasks for org_urn: {org_urn}")
-        task_ids = bulk_upload_to_bubble(task_payloads, BUBBLE_TASKS_TABLE_NAME)
+        response_data = bulk_upload_to_bubble(task_payloads, BUBBLE_TASKS_TABLE_NAME)
 
-        if task_ids is None:
-            logger.error(f"Failed to save tasks to Bubble for org_urn: {org_urn}.")
+        # Validate response and extract IDs.
+        if not response_data or not isinstance(response_data, list):
+            logger.error(f"Failed to save tasks. API response was not a list: {response_data}")
+            return None
+        
+        try:
+            extracted_ids = [item['id'] for item in response_data]
+        except (TypeError, KeyError):
+            logger.error(f"Failed to parse IDs from task API response: {response_data}", exc_info=True)
+            return None
+
+        if len(extracted_ids) != len(task_payloads):
+            logger.error(f"Failed to save all tasks for org_urn: {org_urn}. "
+                         f"Expected {len(task_payloads)} IDs, but got {len(extracted_ids)} from response: {response_data}")
             return None
 
-        logger.info(f"Successfully saved {len(task_ids)} tasks.")
-        return task_ids
+        logger.info(f"Successfully saved {len(extracted_ids)} tasks.")
+        return extracted_ids
 
     except Exception as e:
         logger.exception(f"An unexpected error occurred in save_tasks for org_urn {org_urn}: {e}")
@@ -260,7 +303,6 @@ def save_actionable_okrs(org_urn: str, actionable_okrs: Dict[str, Any], report_i
             return
 
         # Step 1: Save the top-level objectives
-        # Corrected the argument order from your original code.
         objective_ids = save_objectives(org_urn, report_id, objectives_data)
         if objective_ids is None:
             logger.error("OKR save process aborted due to failure in saving objectives.")
@@ -279,7 +321,6 @@ def save_actionable_okrs(org_urn: str, actionable_okrs: Dict[str, Any], report_i
         task_ids = save_tasks(org_urn, key_results_with_ids)
         if task_ids is None:
             logger.error("Task saving failed, but objectives and key results were saved.")
-            # Decide if you want to consider the whole process a failure.
             # For now, we just log the error and complete.
             return