Temp table delete (#1254)

Author: brian-kim31

app/airflow/dags/SR_processing_dag.py

@@ -7,7 +7,12 @@
     process_and_create_scan_report_entries,
     process_data_dictionary,
 )
-from libs.utils import connect_to_storage, create_task, validate_params_SR_processing
+from libs.SR_processing.db_services import handle_failure_and_cleanup_temp_tables
+from libs.utils import (
+    connect_to_storage,
+    create_task,
+    validate_params_SR_processing,
+)
 
 """
 This DAG automates the process of creating scan report tables, fields and values 
@@ -44,6 +49,7 @@
     catchup=False,
     is_paused_upon_creation=False,
     dagrun_timeout=timedelta(minutes=float(AIRFLOW_DAGRUN_TIMEOUT)),
+    on_failure_callback=handle_failure_and_cleanup_temp_tables,
 )
 
 # TODO: add validate for DD file size: DATA_UPLOAD_MAX_MEMORY_SIZE :(

app/airflow/dags/libs/SR_processing/core.py

@@ -157,6 +157,9 @@ def process_and_create_scan_report_entries(**kwargs) -> None:
             logging.error(
                 f"Error inserting tables into mapping_scanreporttable: {str(e)}"
             )
+            # Clean up temp tables before updating job status to fail
+            if table_pairs:
+                delete_temp_tables(scan_report_id, table_pairs)
             update_job_status(
                 stage=JobStageType.UPLOAD_SCAN_REPORT,
                 status=StageStatusType.FAILED,
@@ -172,6 +175,7 @@ def process_and_create_scan_report_entries(**kwargs) -> None:
 
         except Exception as e:
             logging.error(f"Error creating scan report fields: {str(e)}")
+            delete_temp_tables(scan_report_id, table_pairs)
             update_job_status(
                 stage=JobStageType.UPLOAD_SCAN_REPORT,
                 status=StageStatusType.FAILED,
@@ -194,6 +198,7 @@ def process_and_create_scan_report_entries(**kwargs) -> None:
                 )
         except Exception as e:
             logging.error(f"Error creating scan report values: {str(e)}")
+            delete_temp_tables(scan_report_id, table_pairs)
             update_job_status(
                 stage=JobStageType.UPLOAD_SCAN_REPORT,
                 status=StageStatusType.FAILED,

app/airflow/dags/libs/SR_processing/db_services.py

@@ -1,4 +1,5 @@
 import logging
+import time
 from collections import defaultdict
 from typing import Any, Dict, List, Tuple
 
@@ -9,6 +10,7 @@
     AIRFLOW_DAGRUN_TIMEOUT,
     AIRFLOW_DEBUG_MODE,
     EXECUTE_VALUES_PAGE_SIZE,
+    TEMP_TABLE_CLEANUP_DELAY,
 )
 from libs.utils import update_job_status
 from openpyxl.worksheet.worksheet import Worksheet
@@ -238,3 +240,91 @@ def delete_temp_tables(scan_report_id: int, table_pairs: List[Tuple[str, int]])
     except Exception as e:
         logging.error(f"Error deleting temporary tables: {str(e)}")
         raise e
+
+
+def cleanup_temp_tables_for_scan_report(scan_report_id: int) -> List[Tuple[str, int]]:
+    """
+    Clean up temporary tables for a scan report.
+
+    Deletes temporary tables (temp_data_dictionary and temp_field_values) for all tables
+    associated with the given scan_report_id in mapping_scanreporttable.
+
+    Returns:
+        List of (table_name, table_id) tuples for the tables that were cleaned up.
+    """
+
+    query = """
+        SELECT name, id
+        FROM mapping_scanreporttable
+        WHERE scan_report_id = %(scan_report_id)s
+    """
+    records = pg_hook.get_records(query, parameters={"scan_report_id": scan_report_id})
+    table_pairs = [(record[0], record[1]) for record in records] if records else []
+    if table_pairs:
+        delete_temp_tables(scan_report_id, table_pairs)
+    return table_pairs
+
+
+def handle_failure_and_cleanup_temp_tables(context):
+    """
+    Delete temporary tables when the DAG fails or times out.
+
+    This handles cleanup when failures happen outside the normal pipeline code, like
+    timeouts or external errors. Since the tables are already in the database even if
+    the DAG fails, we query mapping_scanreporttable to find all tables for this
+    scan_report_id, then delete the temp tables (temp_data_dictionary and
+    temp_field_values).
+
+    When a timeout occurs, the task that was running may still be creating temporary tables in the background.
+    This function waits TEMP_TABLE_CLEANUP_DELAY seconds first so any in-flight table creation can finish,
+    then runs a single cleanup pass. Cleanup is not urgent, so waiting once is simpler than cleaning twice.
+
+
+    Args:
+        context: Airflow execution context containing task_instance, dag, dag_run, etc.
+    """
+    try:
+        dag_run = context["dag_run"]
+        dag = context.get("dag")
+        dag_run_conf = dag_run.conf or {}
+        scan_report_id = dag_run_conf.get("scan_report_id")
+
+        if not scan_report_id:
+            logging.warning(
+                "No scan_report_id found in DAG run configuration, skipping temp table cleanup"
+            )
+            return
+
+        # Update job status to FAILED for scan_report_processing DAG
+        if dag and dag.dag_id == "scan_report_processing":
+            try:
+                update_job_status(
+                    stage=JobStageType.UPLOAD_SCAN_REPORT,
+                    status=StageStatusType.FAILED,
+                    scan_report=scan_report_id,
+                    details="Scan report processing DAG timed out or failed.",
+                )
+                logging.info(
+                    "Updated job status to FAILED for scan_report_id=%s",
+                    scan_report_id,
+                )
+            except Exception as e:
+                logging.error("Failed to update job status on failure: %s", str(e))
+
+        # Wait so a timed-out task can finish creating tables, then delete the temporary tables
+        delay = TEMP_TABLE_CLEANUP_DELAY
+        time.sleep(delay)
+
+        table_pairs = cleanup_temp_tables_for_scan_report(scan_report_id)
+        if table_pairs:
+            logging.info(
+                "Deleted temp tables for scan_report_id=%s (n=%d)",
+                scan_report_id,
+                len(table_pairs),
+            )
+        logging.info(
+            "Completed temp table cleanup for scan_report_id=%s", scan_report_id
+        )
+
+    except Exception as e:
+        logging.error("Failed to delete temporary tables on failure: %s", str(e))

app/airflow/dags/libs/settings.py

@@ -17,7 +17,11 @@
 # Timedelta for dagrun_timeout in minutes
 AIRFLOW_DAGRUN_TIMEOUT = os.getenv("AIRFLOW_DAGRUN_TIMEOUT", 60)
 
+# In failure callback, temp table cleanup is performed twice.
+# This delay (seconds) gives the timed-out task time to finish creating tables.
+TEMP_TABLE_CLEANUP_DELAY = int(os.getenv("TEMP_TABLE_CLEANUP_DELAY", "60"))
+
 # Page size for bulk database inserts (execute_values)
-EXECUTE_VALUES_PAGE_SIZE = int(os.getenv("EXECUTE_VALUES_PAGE_SIZE", 10000))
+EXECUTE_VALUES_PAGE_SIZE = int(os.getenv("EXECUTE_VALUES_PAGE_SIZE", 1000000))
 
 AIRFLOW_VAR_JSON_VERSION = os.getenv("AIRFLOW_VAR_JSON_VERSION", "v1")

app/airflow/dags/libs/utils.py

@@ -1,6 +1,7 @@
 import ast
 import json
 import logging
+import time
 from typing import Any, Dict, List, Optional, TypedDict
 
 from airflow.models.connection import Connection
@@ -158,6 +159,7 @@ def update_job_status_on_failure(context):
     Args:
         context: Airflow execution context containing task_instance, dag, dag_run, etc.
     """
+    logging.info("update_job_status_on_failure callback triggered")
     try:
         # Extract information from Airflow context
         task_instance = context["task_instance"]

docker-compose.yml

@@ -200,7 +200,9 @@ services:
       # Search recommendations feature flag - controls whether search recommendations task is included in auto_mapping DAG
       SEARCH_ENABLED: False
       AIRFLOW_DAGRUN_TIMEOUT: 60 # Timeout for each dagrun in minutes
+      TEMP_TABLE_CLEANUP_DELAY: 60 # delay between cleanup attempts on failure
       AIRFLOW_VAR_JSON_VERSION: v2
+      EXECUTE_VALUES_PAGE_SIZE: 1000000 # page size for bulk database inserts (execute_values)
 
 volumes:
   db_data: