Merge pull request #10842 from jumpserver/pr@dev@perf_support_tidb

perf: 修改支持 tidb

apps/authentication/migrations/0007_connectiontoken.py

@@ -1,10 +1,10 @@
 # Generated by Django 3.1.12 on 2022-02-11 06:01
+import uuid
 
 from django.db import migrations, models
 
 
 class Migration(migrations.Migration):
-
     dependencies = [
         ('authentication', '0006_auto_20211227_1059'),
     ]
@@ -13,7 +13,7 @@ class Migration(migrations.Migration):
         migrations.CreateModel(
             name='ConnectionToken',
             fields=[
-                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('id', models.UUIDField(default=uuid.uuid4, primary_key=True, serialize=False),),
                 ('created_by', models.CharField(blank=True, max_length=32, null=True, verbose_name='Created by')),
                 ('updated_by', models.CharField(blank=True, max_length=32, null=True, verbose_name='Updated by')),
                 ('date_created', models.DateTimeField(auto_now_add=True, null=True, verbose_name='Date created')),

apps/authentication/migrations/0011_auto_20220705_1940.py

@@ -1,14 +1,13 @@
 # Generated by Django 3.2.12 on 2022-07-05 11:40
 
-import authentication.models
+import django.db.models.deletion
 from django.conf import settings
 from django.db import migrations, models
-import django.db.models.deletion
+
-import uuid
+import authentication.models
 
 
 class Migration(migrations.Migration):
-
     dependencies = [
         ('applications', '0021_auto_20220629_1826'),
         ('assets', '0091_auto_20220629_1826'),
@@ -20,7 +19,9 @@ class Migration(migrations.Migration):
         migrations.AddField(
             model_name='connectiontoken',
             name='application',
-            field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='connection_tokens', to='applications.application', verbose_name='Application'),
+            field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL,
+                                    related_name='connection_tokens', to='applications.application',
+                                    verbose_name='Application'),
         ),
         migrations.AddField(
             model_name='connectiontoken',
@@ -30,7 +31,8 @@ class Migration(migrations.Migration):
         migrations.AddField(
             model_name='connectiontoken',
             name='asset',
-            field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='connection_tokens', to='assets.asset', verbose_name='Asset'),
+            field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL,
+                                    related_name='connection_tokens', to='assets.asset', verbose_name='Asset'),
         ),
         migrations.AddField(
             model_name='connectiontoken',
@@ -55,7 +57,9 @@ class Migration(migrations.Migration):
         migrations.AddField(
             model_name='connectiontoken',
             name='system_user',
-            field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='connection_tokens', to='assets.systemuser', verbose_name='System user'),
+            field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL,
+                                    related_name='connection_tokens', to='assets.systemuser',
+                                    verbose_name='System user'),
         ),
         migrations.AddField(
             model_name='connectiontoken',
@@ -65,25 +69,24 @@ class Migration(migrations.Migration):
         migrations.AddField(
             model_name='connectiontoken',
             name='type',
-            field=models.CharField(choices=[('asset', 'Asset'), ('application', 'Application')], default='asset', max_length=16, verbose_name='Type'),
+            field=models.CharField(choices=[('asset', 'Asset'), ('application', 'Application')], default='asset',
+                                   max_length=16, verbose_name='Type'),
         ),
         migrations.AddField(
             model_name='connectiontoken',
             name='user',
-            field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='connection_tokens', to=settings.AUTH_USER_MODEL, verbose_name='User'),
+            field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL,
+                                    related_name='connection_tokens', to=settings.AUTH_USER_MODEL, verbose_name='User'),
         ),
         migrations.AddField(
             model_name='connectiontoken',
             name='user_display',
             field=models.CharField(default='', max_length=128, verbose_name='User display'),
         ),
-        migrations.AlterField(
-            model_name='connectiontoken',
-            name='id',
-            field=models.UUIDField(default=uuid.uuid4, primary_key=True, serialize=False),
-        ),
         migrations.AlterModelOptions(
             name='connectiontoken',
-            options={'ordering': ('-date_expired',), 'permissions': [('view_connectiontokensecret', 'Can view connection token secret')], 'verbose_name': 'Connection token'},
+            options={'ordering': ('-date_expired',),
+                     'permissions': [('view_connectiontokensecret', 'Can view connection token secret')],
+                     'verbose_name': 'Connection token'},
         ),
     ]

apps/jumpserver/rewriting/db.py

@@ -15,16 +15,18 @@ def atomic(using=None, savepoint=False):
     return db_atomic(using=using, savepoint=savepoint)
 
 
-class OneToOneField(models.OneToOneField):
+class OneToOneField(models.OneToOneField, ForeignKey):
     def __init__(self, *args, **kwargs):
+        kwargs['unique'] = False
+        if os.getenv('DB_CONSTRAINT', '1') == '0':
             kwargs['db_constraint'] = False
-        super().__init__(*args, **kwargs)
+        ForeignKey.__init__(self, *args, **kwargs)
 
 
 def set_db_constraint():
     if os.getenv('DB_CONSTRAINT', '1') != '0':
         return
-    if sys.argv == 2 and sys.argv[1] == 'makemigrations':
+    if len(sys.argv) == 2 and sys.argv[1] == 'makemigrations':
         return
     print("Set foreignkey db constraint False")
     transaction.atomic = atomic

apps/tickets/migrations/0013_ticket_serial_num.py

@@ -28,7 +28,6 @@ def fill_ticket_serial_number(apps, schema_editor):
 
 
 class Migration(migrations.Migration):
-
     dependencies = [
         ('tickets', '0012_ticketsession'),
     ]
@@ -37,7 +36,11 @@ class Migration(migrations.Migration):
         migrations.AddField(
             model_name='ticket',
             name='serial_num',
-            field=models.CharField(max_length=128, null=True, unique=True, verbose_name='Serial number'),
+            field=models.CharField(max_length=128, null=True, verbose_name='Serial number'),
         ),
         migrations.RunPython(fill_ticket_serial_number),
+        migrations.AlterUniqueTogether(
+            name='ticket',
+            unique_together={('serial_num',)},
+        ),
     ]

apps/tickets/models/ticket/general.py

@@ -13,8 +13,8 @@ from django.utils.translation import ugettext_lazy as _
 from common.db.encoder import ModelJSONFieldEncoder
 from common.db.models import JMSBaseModel
 from common.exceptions import JMSException
-from common.utils.timezone import as_current_tz
 from common.utils import reverse
+from common.utils.timezone import as_current_tz
 from orgs.models import Organization
 from orgs.utils import tmp_to_org
 from tickets.const import (
@@ -296,7 +296,7 @@ class Ticket(StatusMixin, JMSBaseModel):
     )
     comment = models.TextField(default='', blank=True, verbose_name=_('Comment'))
     rel_snapshot = models.JSONField(verbose_name=_('Relation snapshot'), default=dict)
-    serial_num = models.CharField(_('Serial number'), max_length=128, unique=True, null=True)
+    serial_num = models.CharField(_('Serial number'), max_length=128, null=True)
     meta = models.JSONField(encoder=ModelJSONFieldEncoder, default=dict, verbose_name=_("Meta"))
     org_id = models.CharField(
         max_length=36, blank=True, default='', verbose_name=_('Organization'), db_index=True
@@ -305,6 +305,9 @@ class Ticket(StatusMixin, JMSBaseModel):
     class Meta:
         ordering = ('-date_created',)
         verbose_name = _('Ticket')
+        unique_together = (
+            ('serial_num',),
+        )
 
     def __str__(self):
         return '{}({})'.format(self.title, self.applicant)
@@ -329,9 +332,9 @@ class Ticket(StatusMixin, JMSBaseModel):
         # TODO: 与 StatusMixin.process_map 内连表查询有部分重叠 有优化空间 待验证排除是否不影响其它调用
         prefetch_ticket_assignee = Prefetch('ticket_steps__ticket_assignees',
                                             queryset=TicketAssignee.objects.select_related('assignee'), )
-        tickets = cls.objects.prefetch_related(prefetch_ticket_assignee)\
+        tickets = cls.objects.prefetch_related(prefetch_ticket_assignee) \
-            .select_related('applicant')\
+            .select_related('applicant') \
-            .filter(queries)\
+            .filter(queries) \
             .distinct()
         return tickets
 

apps/users/migrations/0034_auto_20210506_1448.py

@@ -4,7 +4,6 @@ from django.db import migrations, models
 
 
 class Migration(migrations.Migration):
-
     dependencies = [
         ('users', '0033_user_need_update_password'),
     ]
@@ -13,11 +12,11 @@ class Migration(migrations.Migration):
         migrations.AddField(
             model_name='user',
             name='dingtalk_id',
-            field=models.CharField(default=None, max_length=128, null=True, unique=True, verbose_name='DingTalk'),
+            field=models.CharField(default=None, max_length=128, null=True, verbose_name='DingTalk'),
         ),
         migrations.AddField(
             model_name='user',
             name='wecom_id',
-            field=models.CharField(default=None, max_length=128, null=True, unique=True, verbose_name='WeCom'),
+            field=models.CharField(default=None, max_length=128, null=True, verbose_name='WeCom'),
         ),
     ]

apps/users/migrations/0036_user_feishu_id.py

@@ -4,7 +4,6 @@ from django.db import migrations, models
 
 
 class Migration(migrations.Migration):
-
     dependencies = [
         ('users', '0035_auto_20210526_1100'),
     ]
@@ -13,6 +12,10 @@ class Migration(migrations.Migration):
         migrations.AddField(
             model_name='user',
             name='feishu_id',
-            field=models.CharField(default=None, max_length=128, null=True, unique=True, verbose_name='FeiShu'),
+            field=models.CharField(default=None, max_length=128, null=True, verbose_name='FeiShu'),
+        ),
+        migrations.AlterUniqueTogether(
+            name='user',
+            unique_together={('feishu_id',), ('wecom_id',), ('dingtalk_id',)},
         ),
     ]

apps/users/models/user.py

@@ -811,9 +811,9 @@ class User(AuthMixin, TokenMixin, RoleMixin, MFAMixin, JSONFilterMixin, Abstract
         default=False, verbose_name=_('Need update password')
     )
     date_updated = models.DateTimeField(auto_now=True, verbose_name=_('Date updated'))
-    wecom_id = models.CharField(null=True, default=None, unique=True, max_length=128, verbose_name=_('WeCom'))
+    wecom_id = models.CharField(null=True, default=None, max_length=128, verbose_name=_('WeCom'))
-    dingtalk_id = models.CharField(null=True, default=None, unique=True, max_length=128, verbose_name=_('DingTalk'))
+    dingtalk_id = models.CharField(null=True, default=None, max_length=128, verbose_name=_('DingTalk'))
-    feishu_id = models.CharField(null=True, default=None, unique=True, max_length=128, verbose_name=_('FeiShu'))
+    feishu_id = models.CharField(null=True, default=None, max_length=128, verbose_name=_('FeiShu'))
 
     DATE_EXPIRED_WARNING_DAYS = 5
 
@@ -945,6 +945,11 @@ class User(AuthMixin, TokenMixin, RoleMixin, MFAMixin, JSONFilterMixin, Abstract
     class Meta:
         ordering = ['username']
         verbose_name = _("User")
+        unique_together = (
+            ('dingtalk_id',),
+            ('wecom_id',),
+            ('feishu_id',),
+        )
         permissions = [
             ('invite_user', _('Can invite user')),
             ('remove_user', _('Can remove user')),