Centralize snowflake IDs for bulk create

apps/utils/models.py

@@ -31,7 +31,27 @@ class SoftDeletableQuerySetMixin(object):
             return super(SoftDeletableQuerySetMixin, self).delete()
 
 
-class SoftDeletableQuerySet(SoftDeletableQuerySetMixin, QuerySet):
+class SnowflakeQuerySet(QuerySet):
+    """为 BaseModel 的批量创建统一准备雪花主键。"""
+
+    def bulk_create(self, objs, *args, **kwargs):
+        objs = list(objs)
+        for obj in objs:
+            if not obj.pk:
+                obj.pk = str(idWorker.get_id())
+            else:
+                obj.pk = str(obj.pk)
+        return super().bulk_create(objs, *args, **kwargs)
+
+
+class SnowflakeManager(models.Manager.from_queryset(SnowflakeQuerySet)):
+    pass
+
+
+class SoftDeletableQuerySet(
+    SoftDeletableQuerySetMixin,
+    SnowflakeQuerySet,
+):
     pass
 
 
@@ -111,6 +131,8 @@ class BaseModel(models.Model):
     is_deleted = models.BooleanField(
         default=False, verbose_name='删除标记', help_text='删除标记')
 
+    objects = SnowflakeManager()
+
     class Meta:
         abstract = True
 

apps/utils/tests.py

@@ -0,0 +1,116 @@
+from unittest.mock import patch
+
+from django.db.models.query import QuerySet
+from django.test import SimpleTestCase, TestCase
+
+from apps.develop.models import Project
+from apps.mtm.models import Material
+from apps.system.models import User
+from apps.wpmw.models import WprDefect
+
+
+class SnowflakeBulkCreateTests(SimpleTestCase):
+    @patch.object(QuerySet, "bulk_create")
+    @patch(
+        "apps.utils.models.idWorker.get_id",
+        side_effect=[100, 101],
+    )
+    def test_base_model_bulk_create_assigns_missing_ids(
+        self,
+        get_id,
+        bulk_create,
+    ):
+        objs = [
+            WprDefect(defect_id="10", wpr_id="20"),
+            WprDefect(defect_id="11", wpr_id="20"),
+        ]
+
+        WprDefect.objects.bulk_create(objs)
+
+        self.assertEqual([obj.id for obj in objs], ["100", "101"])
+        self.assertEqual(get_id.call_count, 2)
+        bulk_create.assert_called_once()
+
+    @patch.object(QuerySet, "bulk_create")
+    @patch("apps.utils.models.idWorker.get_id")
+    def test_base_model_bulk_create_preserves_existing_ids(
+        self,
+        get_id,
+        bulk_create,
+    ):
+        objs = [
+            WprDefect(id=100, defect_id="10", wpr_id="20"),
+            WprDefect(id="101", defect_id="11", wpr_id="20"),
+        ]
+
+        WprDefect.objects.bulk_create(objs)
+
+        self.assertEqual([obj.id for obj in objs], ["100", "101"])
+        get_id.assert_not_called()
+        bulk_create.assert_called_once()
+
+    @patch.object(QuerySet, "bulk_create")
+    @patch(
+        "apps.utils.models.idWorker.get_id",
+        side_effect=[200, 201],
+    )
+    def test_soft_delete_manager_uses_snowflake_bulk_create(
+        self,
+        get_id,
+        bulk_create,
+    ):
+        objs = [
+            Material(name="物料1"),
+            Material(name="物料2"),
+        ]
+
+        Material.objects.bulk_create(objs)
+
+        self.assertEqual([obj.id for obj in objs], ["200", "201"])
+        self.assertEqual(get_id.call_count, 2)
+        bulk_create.assert_called_once()
+
+    @patch.object(QuerySet, "bulk_create")
+    @patch(
+        "apps.utils.models.idWorker.get_id",
+        side_effect=[300, 301],
+    )
+    def test_custom_user_manager_uses_snowflake_bulk_create(
+        self,
+        get_id,
+        bulk_create,
+    ):
+        objs = [
+            User(username="bulk-user-1"),
+            User(username="bulk-user-2"),
+        ]
+
+        User.objects.bulk_create(objs)
+
+        self.assertEqual([obj.id for obj in objs], ["300", "301"])
+        self.assertEqual(get_id.call_count, 2)
+        bulk_create.assert_called_once()
+
+
+class SnowflakeBulkCreateDatabaseTests(TestCase):
+    @patch(
+        "apps.utils.models.idWorker.get_id",
+        side_effect=[900000000000000001, 900000000000000002],
+    )
+    def test_bulk_create_persists_generated_ids(self, get_id):
+        objs = [
+            Project(name="批量项目1", code="bulk-project-1"),
+            Project(name="批量项目2", code="bulk-project-2"),
+        ]
+
+        Project.objects.bulk_create(objs)
+
+        self.assertEqual(
+            [obj.id for obj in objs],
+            ["900000000000000001", "900000000000000002"],
+        )
+        self.assertEqual(
+            set(Project.objects.values_list("id", flat=True)),
+            {"900000000000000001", "900000000000000002"},
+        )
+        self.assertEqual(get_id.call_count, 2)

apps/wpm/services.py

@@ -21,6 +21,7 @@ from apps.wpmw.models import Wpr, WprDefect
 from ..qm.models import Defect, Ftest
 from django.db.models import Count, Q
 from apps.utils.tasks import ctask_run
+from apps.utils.snowflake import idWorker
 from apps.mtm.models import Process, WmScope
 from django.db.models import F
 
@@ -214,6 +215,7 @@ def restore_mlog_wpr_defects(mlog: Mlog):
             continue
         restored.extend(
             WprDefect(
+                id=idWorker.get_id(),
                 wpr=wpr,
                 defect_id=item["defect_id"],
                 is_main=item.get("is_main", False),

apps/wpm/tests.py

@@ -548,6 +548,30 @@ class WMaterialScopeTests(SimpleTestCase):
 
 
 class WprDefectLifecycleTests(SimpleTestCase):
+    @patch("apps.wpmw.models.WprDefect.objects.bulk_create")
+    @patch("apps.wpmw.models.WprDefect.objects.filter")
+    @patch(
+        "apps.wpmw.models.idWorker.get_id",
+        side_effect=["100", "101"],
+    )
+    def test_replace_defects_assigns_unique_ids_before_bulk_create(
+        self,
+        get_id,
+        defect_filter,
+        bulk_create,
+    ):
+        wpr = Wpr(id="1")
+
+        wpr.replace_defects([
+            {"defect_id": "10", "is_main": True},
+            {"defect_id": "11", "is_main": False},
+        ])
+
+        defect_filter.return_value.delete.assert_called_once_with()
+        created = bulk_create.call_args.args[0]
+        self.assertEqual([str(item.id) for item in created], ["100", "101"])
+        self.assertEqual(get_id.call_count, 2)
+
     @patch("apps.wpmw.models.WprDefect.objects.filter")
     def test_change_inventory_never_changes_individual_defects(self, defect_filter):
         material = Material(id="1", tracking=Material.MA_TRACKING_SINGLE)
@@ -671,8 +695,13 @@ class MlogWprDefectSnapshotTests(SimpleTestCase):
     @patch("apps.wpm.services.WprDefect.objects.bulk_create")
     @patch("apps.wpm.services.WprDefect.objects.filter")
     @patch("apps.wpm.services.Wpr.objects.select_for_update")
+    @patch(
+        "apps.wpm.services.idWorker.get_id",
+        side_effect=["100", "101"],
+    )
     def test_restore_replaces_current_defects_exactly(
         self,
+        get_id,
         select_for_update,
         defect_filter,
         bulk_create,
@@ -698,6 +727,8 @@ class MlogWprDefectSnapshotTests(SimpleTestCase):
             [(str(item.defect_id), item.is_main) for item in created],
             [("10", True), ("11", False)],
         )
+        self.assertEqual([str(item.id) for item in created], ["100", "101"])
+        self.assertEqual(get_id.call_count, 2)
 
 
 @skipUnless(

apps/wpmw/models.py

@@ -5,6 +5,7 @@ from apps.utils.models import BaseModel
 from apps.mtm.models import Material
 from rest_framework.exceptions import ParseError
 from apps.wpm.models import WmStateOption, Mlogbw, Handoverbw
+from apps.utils.snowflake import idWorker
 from apps.utils.tools import update_dict
 from apps.inm.models import MIOItemw
 from django.db.models import F, Value
@@ -130,6 +131,7 @@ class Wpr(BaseModel):
         WprDefect.objects.filter(wpr=self).delete()
         WprDefect.objects.bulk_create([
             WprDefect(
+                id=idWorker.get_id(),
                 wpr=self,
                 defect_id=item["defect_id"],
                 is_main=item.get("is_main", False),