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

Author: xaviedoanhduy

dms_import/hooks.py

@@ -2,6 +2,8 @@
 # License AGPL-3.0 or later (https://www.gnu.org/licenses/agpl).
 
 import logging
+import os
+from collections import defaultdict
 from random import randint
 
 from psycopg2.sql import SQL, Identifier
@@ -37,68 +39,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 +113,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 +148,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 +209,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 +224,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 +256,84 @@ 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):
+# flake8: noqa: C901
+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,8 +351,10 @@ 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):
+    Attachment = env["ir.attachment"]
+    filestore = Attachment._filestore()
+    BATCH_SIZE = models.INSERT_BATCH_SIZE
+    for batch_of_ids in split_every(BATCH_SIZE, all_doc_ids):
         _logger.info("Migrating files batch with %d documents", len(batch_of_ids))
         cr.execute(
             SQL(
@@ -337,6 +367,16 @@ def migrate_documents_files(cr, env, folder_mapping, tag_mapping, category_mappi
         if not documents:
             continue
 
+        attachment_ids = [
+            d["attachment_id"] for d in documents if d.get("attachment_id")
+        ]
+        attachments_data = {
+            att["id"]: att["store_fname"]
+            for att in Attachment.with_context(active_test=False).search_read(
+                domain=[("id", "in", attachment_ids)], fields=["store_fname"]
+            )
+        }
+
         # Batch-fetch tag relations for the current batch
         tags_by_doc = {}
         if table_exists(cr, "document_tag_rel"):
@@ -360,33 +400,50 @@ def migrate_documents_files(cr, env, folder_mapping, tag_mapping, category_mappi
                 )
                 continue
 
+            att_id = doc["attachment_id"]
+            if not att_id:
+                _logger.warning("Skipping doc %s: no attachment_id.", doc["id"])
+                continue
+            store_fname = attachments_data.get(att_id)
+            if not store_fname:
+                _logger.warning(
+                    "Skipping doc %s: attachment %s has no store_fname or does not exist.",
+                    doc["id"],
+                    att_id,
+                )
+                continue
+            if not os.path.exists(os.path.join(filestore, store_fname)):
+                _logger.warning(
+                    "Skipping doc %s: file not found in filestore for attachment %s (%s)",
+                    doc["id"],
+                    att_id,
+                    store_fname,
+                )
+                continue
+
             files_to_create_vals.append(
                 {
                     "name": doc["name"],
                     "directory_id": directory_id,
                     "active": doc["active"],
-                    "attachment_id": doc["attachment_id"],
+                    "attachment_id": att_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 DmsFile.env.norecompute():
+                new_files = DmsFile.create(files_to_create_vals)
 
-            attachments_to_update = {}
+            attachments_to_update = defaultdict(list)
             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 vals.get("attachment_id"):
+                    attachments_to_update[new_file.id].append(vals["attachment_id"])
 
-            if attachments_to_update:
-                for att_id, file_id in attachments_to_update.items():
-                    env["ir.attachment"].browse(att_id).write(
+            for file_id, att_ids in attachments_to_update.items():
+                if att_ids:
+                    Attachment.browse(att_ids).write(
                         {"res_model": "dms.file", "res_id": file_id}
                     )
 
@@ -403,15 +460,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 EE 'documents' to CE '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!")