Implemented pg_hook.get_cursor to make SR upload fast

Author: brian-kim31

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

@@ -2,6 +2,7 @@
 from collections import defaultdict
 from typing import Any, Dict, List, Tuple
 
+from psycopg2.extras import execute_values
 from airflow.providers.postgres.hooks.postgres import PostgresHook
 from libs.enums import JobStageType, StageStatusType
 from libs.queries import create_fields_query
@@ -105,27 +106,21 @@ def update_temp_data_dictionary_table(
                         }
                     )
 
-        # Insert records into the temporary table
+        # Using execute_values with page_size=5000 instead of insert_rows() (fixed page_size=100) for better performance
         if dictionary_records:
-            pg_hook.insert_rows(
-                table=f"temp_data_dictionary_{scan_report_id}",
-                rows=[
-                    (
-                        d["table_name"],
-                        d["field_name"],
-                        d["value"],
-                        d["value_description"],
-                    )
-                    for d in dictionary_records
-                ],
-                target_fields=[
-                    "table_name",
-                    "field_name",
-                    "value",
-                    "value_description",
-                ],
-                commit_every=10000,  # commit every 50k
-            )
+            rows = [
+                (
+                    d["table_name"],
+                    d["field_name"],
+                    d["value"],
+                    d["value_description"],
+                )
+                for d in dictionary_records
+            ]
+            sql = f"INSERT INTO temp_data_dictionary_{scan_report_id} (table_name, field_name, value, value_description) VALUES %s"
+            with pg_hook.get_cursor() as cur:
+                execute_values(cur, sql, rows, page_size=5000)
+                cur.connection.commit()
 
         logging.info(
             f"Created temporary data dictionary table with {len(dictionary_records)} records"
@@ -197,23 +192,19 @@ def create_temp_field_values_table(
                 )
 
         if field_values_data:
-            pg_hook.insert_rows(
-                table=f"temp_field_values_{table_id}",
-                rows=[
-                    (
-                        d["field_name"],
-                        d["value"],
-                        d["frequency"],
-                    )
-                    for d in field_values_data
-                ],
-                target_fields=[
-                    "field_name",
-                    "value",
-                    "frequency",
-                ],
-                commit_every=10000,  # commit every 50k to reduce transaction overhead
-            )
+            # Using execute_values with page_size=5000 instead of insert_rows() (fixed page_size=100) for better performance
+            rows = [
+                (
+                    d["field_name"],
+                    d["value"],
+                    d["frequency"],
+                )
+                for d in field_values_data
+            ]
+            sql = f"INSERT INTO temp_field_values_{table_id} (field_name, value, frequency) VALUES %s"
+            with pg_hook.get_cursor() as cur:
+                execute_values(cur, sql, rows, page_size=5000)
+                cur.connection.commit()
 
     except Exception as e:
         logging.error(f"Error creating data dictionary table: {str(e)}")