Update: Implement bulk update functionality using list of document IDs and updated data

This commit introduces a new function, `insert_mongo_documents_to_postgres`, which allows for the bulk updating of documents in a PostgreSQL table. It takes a list of document IDs and a dictionary containing the updated data as input. This method simplifies the process by reducing the number of individual insert statements required.

The changes include:

1. Adding a new function `insert_mongo_documents_to_postgres` to handle the bulk update.
2. Updating the existing code to call this new function when required.
3. Refactoring the `MongoDocumentInserter` class to encapsulate common functionality for document insertion and updating.

This refactoring enhances maintainability and efficiency by reducing redundancy in the codebase.

db/schema_generator.py

@@ -530,17 +530,106 @@ class MongoDocumentInserter:
         if on_conflict:
             insert_sql += f" {on_conflict}"
 
-        with engine.begin() as conn:
-            if on_conflict:
+        # Use different strategies depending on whether ON CONFLICT is needed.
+        # For ON CONFLICT updates, per-row execution is slow; use psycopg2's
+        # execute_values to perform fast multi-row INSERT ... VALUES (...) ON CONFLICT ...
+        if on_conflict:
+            try:
+                # Prepare ordered tuples for insertion
+                values_list = []
                 for doc in prepared:
+                    values_list.append(tuple(doc.get(col) for col in column_names))
+
+                # Build base INSERT with placeholder for execute_values
+                insert_base = f"INSERT INTO {full_table_name} ({cols_str}) VALUES %s {on_conflict}"
+
+                # Use raw DB-API connection for execute_values
+                try:
+                    import psycopg2
+                except Exception:
+                    # Fall back to SQLAlchemy executemany if psycopg2 not available
+                    with engine.begin() as conn:
+                        for doc in prepared:
+                            try:
+                                conn.execute(text(insert_sql), [doc])
+                                total_inserted += 1
+                            except Exception as e:
+                                logger.error(f"Error inserting document into {full_table_name}: {e}")
+                                raise
+                        logger.info(f"Inserted {total_inserted} documents into {full_table_name}")
+                else:
+                    # Use a single staging temporary table per insert_documents call.
+                    # Stream each batch into the temp table via COPY, then run one
+                    # INSERT ... SELECT ... ON CONFLICT ... to upsert all rows.
+                    raw_conn = engine.raw_connection()
                     try:
-                        conn.execute(text(insert_sql), [doc])
-                        total_inserted += 1
-                    except Exception as e:
-                        logger.error(f"Error inserting document into {full_table_name}: {e}")
-                        raise
-                logger.info(f"Inserted {total_inserted} documents into {full_table_name}")
-            else:
+                        cur = raw_conn.cursor()
+                        import io
+                        import csv
+                        import time
+                        import os
+
+                        temp_name = f"tmp_{table_name}_{int(time.time())}_{os.getpid()}"
+                        try:
+                            # Create temporary table once
+                            cur.execute(f"CREATE TEMP TABLE {temp_name} (LIKE {full_table_name} INCLUDING ALL);")
+                            # Stream batches into temp table
+                            for i in range(0, len(prepared), self.batch_size):
+                                batch = prepared[i:i + self.batch_size]
+                                try:
+                                    sio = io.StringIO()
+                                    writer = csv.writer(sio)
+                                    for doc in batch:
+                                        row = []
+                                        for col in column_names:
+                                            val = doc.get(col)
+                                            if val is None:
+                                                row.append('\\N')
+                                            else:
+                                                row.append(val)
+                                        writer.writerow(row)
+                                    sio.seek(0)
+
+                                    copy_sql = f"COPY {temp_name} ({cols_str}) FROM STDIN WITH CSV NULL '\\N'"
+                                    cur.copy_expert(copy_sql, sio)
+                                    raw_conn.commit()
+                                    batch_count = len(batch)
+                                    total_inserted += batch_count
+                                    logger.info(f"Copied {batch_count} rows into staging {temp_name} (total staged: {total_inserted})")
+                                except Exception as e:
+                                    raw_conn.rollback()
+                                    logger.error(f"Error copying batch into {temp_name}: {e}")
+                                    raise
+
+                            # Perform a single upsert from the staging table
+                            try:
+                                insert_from_temp = f"INSERT INTO {full_table_name} ({cols_str}) SELECT {cols_str} FROM {temp_name} {on_conflict}"
+                                cur.execute(insert_from_temp)
+                                raw_conn.commit()
+                                logger.info(f"Upserted staged rows from {temp_name} into {full_table_name}")
+                            except Exception as e:
+                                raw_conn.rollback()
+                                logger.error(f"Error upserting from {temp_name} into {full_table_name}: {e}")
+                                raise
+                        finally:
+                            # Temp table will be dropped on commit/connection close, ensure commit
+                            try:
+                                raw_conn.commit()
+                            except Exception:
+                                pass
+                    finally:
+                        try:
+                            cur.close()
+                        except Exception:
+                            pass
+                        try:
+                            raw_conn.close()
+                        except Exception:
+                            pass
+            except Exception:
+                raise
+        else:
+            with engine.begin() as conn:
                 for i in range(0, len(prepared), self.batch_size):
                     batch = prepared[i:i + self.batch_size]
                     try: