[IMP] dms_import: improve performance, reduce batch size, and bypass heavy compute fields

Author: xaviedoanhduy

dms_import/hooks.py

@@ -2,6 +2,7 @@
 # License AGPL-3.0 or later (https://www.gnu.org/licenses/agpl).
 
 import logging
+from collections import defaultdict
 from random import randint
 
 from psycopg2.sql import SQL, Identifier
@@ -37,68 +38,48 @@ def _get_or_create_default_storage(env):
     return db_storage
 
 
-def _assign_dms_groups(cr, env, folder_id, new_dir):
-    """Migrate read/write groups from a documents.folder to dms.access.groups."""
-    DmsAccessGroups = env["dms.access.group"]
-    all_new_groups = DmsAccessGroups.browse()
+def _batch_fetch_folder_relations(cr, folder_ids):
+    write_groups = defaultdict(list)
+    read_groups = defaultdict(list)
+    tags_by_folder = defaultdict(list)
 
-    # Handle Write Groups
     if table_exists(cr, "documents_folder_res_groups_rel"):
         cr.execute(
             SQL(
-                """
-                SELECT res_groups_id FROM {}
-                WHERE documents_folder_id = %s
-            """
-            ).format(Identifier("documents_folder_res_groups_rel")),
-            (folder_id,),
+                """SELECT documents_folder_id, res_groups_id
+            FROM documents_folder_res_groups_rel
+            WHERE documents_folder_id = ANY(%s)"""
+            ),
+            (list(folder_ids),),
         )
-        group_ids = [r[0] for r in cr.fetchall()]
-        if group_ids:
-            write_group = DmsAccessGroups.create(
-                {
-                    "name": f"{new_dir.name} Write Group",
-                    "perm_create": True,
-                    "perm_write": True,
-                    "perm_unlink": True,
-                    "group_ids": [Command.set(group_ids)],
-                }
-            )
-            all_new_groups |= write_group
+        for folder_id, group_id in cr.fetchall():
+            write_groups[folder_id].append(group_id)
 
-    # Handle Read Groups
     if table_exists(cr, "documents_folder_read_groups"):
         cr.execute(
             SQL(
-                """
-                SELECT res_groups_id FROM {}
-                WHERE documents_folder_id = %s
-            """
-            ).format(Identifier("documents_folder_read_groups")),
-            (folder_id,),
+                """SELECT documents_folder_id, res_groups_id
+            FROM documents_folder_read_groups
+            WHERE documents_folder_id = ANY(%s)"""
+            ),
+            (list(folder_ids),),
         )
-        read_group_ids = [r[0] for r in cr.fetchall()]
-        if read_group_ids:
-            read_group = DmsAccessGroups.create(
-                {
-                    "name": f"{new_dir.name} Read Group",
-                    "group_ids": [Command.set(read_group_ids)],
-                }
-            )
-            all_new_groups |= read_group
+        for folder_id, group_id in cr.fetchall():
+            read_groups[folder_id].append(group_id)
 
-    # If dont have any group, we need to assign a group to visible folder
-    if not all_new_groups:
-        all_new_groups = DmsAccessGroups.create(
-            {
-                "name": f"{new_dir.name} Default Group",
-                "perm_create": True,
-                "perm_write": True,
-                "perm_unlink": True,
-                "group_ids": [Command.set(env.ref("base.group_user").ids)],
-            }
+    if table_exists(cr, "documents_facet") and table_exists(cr, "documents_tag"):
+        cr.execute(
+            SQL(
+                """SELECT f.folder_id, t.id
+                FROM {} t JOIN {} f ON f.id = t.facet_id
+                WHERE f.folder_id = ANY(%s)"""
+            ).format(Identifier("documents_tag"), Identifier("documents_facet")),
+            (list(folder_ids),),
         )
-    return all_new_groups
+        for folder_id, tag_id in cr.fetchall():
+            tags_by_folder[folder_id].append(tag_id)
+
+    return write_groups, read_groups, tags_by_folder
 
 
 def _create_unique_directory(DmsDirectory, vals):
@@ -131,7 +112,7 @@ def migrate_documents_tags(cr, env, lang):
     tag_mapping = {}
     category_mapping = {}
 
-    # 1. FACETS → CATEGORIES
+    # 1. Facets → categories
     if table_exists(cr, "documents_facet"):
         cr.execute(
             SQL("SELECT id, name->>%s AS name FROM {}").format(
@@ -166,7 +147,7 @@ def migrate_documents_tags(cr, env, lang):
             if norm and norm in existing_categories:
                 category_mapping[f["id"]] = existing_categories[norm].id
 
-    # 2. TAGS
+    # 2. Tags
     existing_tags = {
         (_normalize(tag.name), tag.category_id.id or False): tag
         for tag in DmsTag.search([])
@@ -227,7 +208,7 @@ def migrate_documents_tags(cr, env, lang):
     return tag_mapping, category_mapping
 
 
-def migrate_documents_folders(cr, env, lang, tag_mapping, category_mapping):
+def migrate_documents_folders(cr, env, lang, tag_mapping):
     """Migrate documents.folder to dms.directory, optimized for batch processing."""
     if not table_exists(cr, "documents_folder"):
         _logger.warning("Skipping folder migration: 'documents_folder' not found.")
@@ -242,11 +223,22 @@ def migrate_documents_folders(cr, env, lang, tag_mapping, category_mapping):
         (lang,),
     )
     folders = cr.dictfetchall()
+    if not folders:
+        _logger.info("No folders to migrate.")
+        return {}
 
     DmsDirectory = env["dms.directory"]
+    DmsAccessGroups = env["dms.access.group"]
     folder_mapping = {}
     storage = _get_or_create_default_storage(env)
+    default_user_group_id = env.ref("base.group_user").id
 
+    folder_ids = [f["id"] for f in folders]
+    write_groups, read_groups, tags_by_folder = _batch_fetch_folder_relations(
+        cr, folder_ids
+    )
+    access_groups_to_create_vals = []
+    dirs_to_update_tags = defaultdict(list)
     for folder in folders:
         try:
             parent_id = folder_mapping.get(folder["parent_folder_id"])
@@ -263,49 +255,83 @@ def migrate_documents_folders(cr, env, lang, tag_mapping, category_mapping):
                 "Created dms.directory: '%s' (ID: %s)", new_dir.name, new_dir.id
             )
 
-            # Assign groups
-            dms_groups = _assign_dms_groups(cr, env, folder["id"], new_dir)
-            dms_groups._compute_users()  # trigger to recompute count_users
-            new_dir.group_ids = [Command.set(dms_groups.ids)]
-            _logger.info(
-                "Assigned %s groups to directory '%s'",
-                dms_groups.mapped("name"),
-                new_dir.name,
-            )
+            # Assign groups using pre-fetched data
+            has_groups = False
+            write_group_ids = write_groups.get(folder["id"], [])
+            if write_group_ids:
+                has_groups = True
+                vals = {
+                    "name": f"{new_dir.name} Write Group",
+                    "perm_create": True,
+                    "perm_write": True,
+                    "perm_unlink": True,
+                    "group_ids": [Command.set(write_group_ids)],
+                    "_dir_id": new_dir.id,
+                }
+                access_groups_to_create_vals.append(vals)
+            read_group_ids = read_groups.get(folder["id"], [])
+            if read_group_ids:
+                has_groups = True
+                vals = {
+                    "name": f"{new_dir.name} Read Group",
+                    "group_ids": [Command.set(read_group_ids)],
+                    "_dir_id": new_dir.id,
+                }
+                access_groups_to_create_vals.append(vals)
+            # In case no group, directory will be invisible, so I decide to assign a new group
+            # this group has full permision like no group in document folder
+            if not has_groups:
+                vals = {
+                    "name": f"{new_dir.name} Default Group",
+                    "perm_create": True,
+                    "perm_write": True,
+                    "perm_unlink": True,
+                    "group_ids": [Command.set([default_user_group_id])],
+                    "_dir_id": new_dir.id,
+                }
+                access_groups_to_create_vals.append(vals)
 
             # Assign Tags
-            if table_exists(cr, "documents_facet") and table_exists(
-                cr, "documents_tag"
-            ):
-                cr.execute(
-                    SQL(
-                        """SELECT t.id FROM {} t
-                        JOIN {} f ON f.id = t.facet_id
-                        WHERE f.folder_id = %s"""
-                    ).format(
-                        Identifier("documents_tag"),
-                        Identifier("documents_facet"),
-                    ),
-                    (folder["id"],),
-                )
-                dms_tag_ids = [
-                    tag_mapping.get(r[0])
-                    for r in cr.fetchall()
-                    if tag_mapping.get(r[0])
-                ]
-                if dms_tag_ids:
-                    new_dir.tag_ids = [Command.set(dms_tag_ids)]
+            folder_tag_ids = [
+                tag_mapping[tag_id]
+                for tag_id in tags_by_folder.get(folder["id"], [])
+                if tag_id in tag_mapping
+            ]
+            if folder_tag_ids:
+                dirs_to_update_tags[tuple(sorted(folder_tag_ids))].append(new_dir.id)
 
         except Exception:
             _logger.exception(
                 "Error migrating folder ID %s (%s)", folder["id"], folder["name"]
             )
 
+    for tag_ids_tuple, dir_ids in dirs_to_update_tags.items():
+        DmsDirectory.browse(dir_ids).write(
+            {"tag_ids": [Command.set(list(tag_ids_tuple))]}
+        )
+
+    if access_groups_to_create_vals:
+        new_groups = DmsAccessGroups.create(
+            [
+                {k: v for k, v in vals.items() if k != "_dir_id"}
+                for vals in access_groups_to_create_vals
+            ]
+        )
+        # Recompute all user counts at once
+        new_groups.invalidate_model(fnames=["users"])
+
+        dirs_to_update_groups = defaultdict(list)
+        for vals, group in zip(access_groups_to_create_vals, new_groups):
+            dirs_to_update_groups[vals["_dir_id"]].append(group.id)
+
+        for dir_id, group_ids in dirs_to_update_groups.items():
+            DmsDirectory.browse(dir_id).write({"group_ids": [Command.set(group_ids)]})
+
     _logger.info("Successfully migrated %d folders.", len(folder_mapping))
     return folder_mapping
 
 
-def migrate_documents_files(cr, env, folder_mapping, tag_mapping, category_mapping):
+def migrate_documents_files(cr, env, folder_mapping, tag_mapping):
     """Migrate documents.document to dms.file, optimized for batch processing."""
     if not table_exists(cr, "documents_document") or not folder_mapping:
         _logger.warning("Skipping file migration: table or folder mapping is missing.")
@@ -323,9 +349,15 @@ def migrate_documents_files(cr, env, folder_mapping, tag_mapping, category_mappi
         return
 
     DmsFile = env["dms.file"]
-
-    for batch_of_ids in split_every(models.INSERT_BATCH_SIZE * 10, all_doc_ids):
-        _logger.info("Migrating files batch with %d documents", len(batch_of_ids))
+    BATCH_SIZE = models.INSERT_BATCH_SIZE
+    total_docs = len(all_doc_ids)
+    _logger.info("Processing %d documents in total.", total_docs)
+    for i, batch_of_ids in enumerate(split_every(BATCH_SIZE, all_doc_ids)):
+        _logger.info(
+            "Processing batch %d/%d...",
+            i + 1,
+            (total_docs + BATCH_SIZE - 1) // BATCH_SIZE,
+        )
         cr.execute(
             SQL(
                 """SELECT id, name, folder_id, attachment_id, active
@@ -367,29 +399,38 @@ def migrate_documents_files(cr, env, folder_mapping, tag_mapping, category_mappi
                     "active": doc["active"],
                     "attachment_id": doc["attachment_id"],
                     "tag_ids": [Command.set(tags_by_doc.get(doc["id"], []))],
-                    "_old_attachment_id": doc["attachment_id"],
                 }
             )
 
         if files_to_create_vals:
-            new_files = DmsFile.create(
-                [
-                    {k: v for k, v in vals.items() if k != "_old_attachment_id"}
-                    for vals in files_to_create_vals
-                ]
-            )
+            # avoid _compute_content in dms.file
+            with env.norecompute():
+                new_files = DmsFile.create(files_to_create_vals)
 
-            attachments_to_update = {}
+            # Update attachments in bulk
+            attachment_update_map = {}
             for vals, new_file in zip(files_to_create_vals, new_files):
-                if vals["_old_attachment_id"]:
-                    attachments_to_update[vals["_old_attachment_id"]] = new_file.id
-
-            if attachments_to_update:
-                for att_id, file_id in attachments_to_update.items():
-                    env["ir.attachment"].browse(att_id).write(
-                        {"res_model": "dms.file", "res_id": file_id}
-                    )
+                if vals.get("attachment_id"):
+                    attachment_update_map[vals["attachment_id"]] = new_file.id
+
+            if attachment_update_map:
+                attachment_ids = list(attachment_update_map.keys())
+                case_clauses = []
+                params = []
+                for att_id, file_id in attachment_update_map.items():
+                    case_clauses.append(SQL("WHEN %s THEN %s"))
+                    params.extend([att_id, file_id])
+                params.append(tuple(attachment_ids))
+                query = SQL(
+                    """
+                    UPDATE ir_attachment
+                    SET res_model = 'dms.file',
+                        res_id = CASE id {cases} END
+                    WHERE id IN %s
+                """
+                ).format(cases=SQL(" ").join(case_clauses))
 
+                cr.execute(query, params)
     _logger.info("Files migration finished.")
 
 
@@ -403,15 +444,11 @@ def post_init_hook(cr, registry):
     try:
         lang = get_lang(env).code or env.lang
         _logger.info(
-            "Starting migration from EE 'documents' to OCA 'dms' using lang '%s'", lang
+            "Starting migration from 'documents' to 'dms' using lang '%s'", lang
         )
-
-        tag_mapping, category_mapping = migrate_documents_tags(cr, env, lang)
-        folder_mapping = migrate_documents_folders(
-            cr, env, lang, tag_mapping, category_mapping
-        )
-        migrate_documents_files(cr, env, folder_mapping, tag_mapping, category_mapping)
-
+        tag_mapping, _ = migrate_documents_tags(cr, env, lang)
+        folder_mapping = migrate_documents_folders(cr, env, lang, tag_mapping)
+        migrate_documents_files(cr, env, folder_mapping, tag_mapping)
         _logger.info("Migration from 'documents' to 'dms' completed successfully.")
     except Exception:
         _logger.exception("Migration from 'documents' to 'dms' failed!")