Django如何使用snowflake自定义生成主键而不是自动生成主键？

之前ID都是用自增实现的，那现在想用Snowflake算法生成主键，要做什么改动呢？

背景介绍

目前工程框架如下

• Django
• django.db.backends.postgresql_psycopg2

目前models声明和save()方法如下

# models.py
# 使用的是默认主键
class User(models.Model):
    name = models.CharField(max_length=100, verbose_name="名字")

# user的新增操作，默认主键
user = User()
user.name = "123"
user.save()

在user.save()时执行的SQL语句是

INSERT INTO "polls_user" ("name") VALUES ('xiaoming') RETURNING "polls_user"."id"; args=('xiaoming',)

很明显，并未在语句中设置ID主键

查看表结构，发现id自动是使用序列自增来实现的

CREATE TABLE "public"."polls_user" (
  "id" int4 NOT NULL DEFAULT nextval('polls_user_id_seq'::regclass),
  "name" varchar(100) COLLATE "pg_catalog"."default" NOT NULL,
  CONSTRAINT "polls_user_pkey" PRIMARY KEY ("id")
)

那么基于目前的情况，如何改写使用snowflake算法生成ID，而不是自动生成主键呢？

实现方案

方案1 - 手动添加主键

实现方式如下

• （1）首先需要声明id，声明为非AutoField类型

  如果不声明，在使用postgresql的情况下，会默认为主键设置为serial类型，使用序列生成主键

# models.py

class User(models.Model):
    id = models.BigIntegerField(primary_key=True)
    name = models.CharField(max_length=100, verbose_name="名字")

• （2）其次是save( )的时候，需要手动设置主键

# user的新增操作

user = User()
user.id = snowflake.next_id()
user.name = "123"
# 可以不指定force_insert=True，如果不指定，默认会先update看是否有记录，如果没有，再执行insert
user.save(force_insert=True)

改动：（1）所有model增加id设置

（2）所有save()之前设置id=snowflake.next_id()

方案2 - 重写save()方法

实现方式如下

• 单独改写某一个model的save( )方法

# models.py

class User(models.Model):
    id = models.BigIntegerField(primary_key=True)
    name = models.CharField(max_length=100, verbose_name="名字")
    
    def save(self, *args, **kwargs):
        if not self.id:
            self.id = snowflake.next_id()
        super(User, self).save(*args, **kwargs)
        return self.id

• 调用的地方不用修改，不需要设置 user.id = snowflake.next_id()

按照上面的方法的话，每个新的model都是需要重写save( )方法的，显然工作量仍然比较多。

改动：（1）所有model增加id设置

（2）所有model重写save()方法

方案3 - 使用 Django Signals 中的pre_save()

官方文档地址：https://docs.djangoproject.com/en/3.0/topics/signals/

实现方式如下

• 修改models.py文件

使用pre_save( )方法，在调用models.save( )之后立即执行

注意

（1）如auth_开头的表的save()以及 其他django_开头的表都会走此逻辑，需要排除这些model的save

（2）save() 包括update和insert，过滤执行update的model

from django.db.models.signals import pre_save
from django.dispatch import receiver
import logging


@receiver(pre_save)
def pre_save_set_snowflake_id(sender, instance, *args, **kwargs):
    """
    Django Signals, pre_save
    适用所有的model
    If we dont include the sender argument in the decorator,
    like @receiver(pre_save, sender=MyModel), the callback will be called for all models.
    """
    # print(__name__)  # = polls.models
    # print(type(instance)) # = <class 'polls.models.Question'>
    if __name__ in str(type(instance)) and not instance.id:
        # 满足条件（1）在本models.py中声明的model （2）id不为空的model 才会使用snowflake生成id
        # 原因是如果不加条件（1）会有如auth_开头的表以及django_开头的表也会用snowflake生成的id，但是其id长度不够
        instance.id = snowflake.next_id()

        
class User(models.Model):
    id = models.BigIntegerField(primary_key=True)
    name = models.CharField(max_length=100, verbose_name="名字")
            

如上，修改好models.py之后，就可以安心使用save()了，原先的业务逻辑也不用有其他修改

改动（1）所有model增加id设置，统一设置 id = models.BigIntegerField(primary_key=True)

（2）增加pre_save( ) 方法，统一设置 id = snowflake.next_id()

方案4 - 仿照models.UUIDField，写一个models.SnowflakeIDField

尝试中，但是还没有写出来。

思路就是新写一个field，比如类似models.UUIDField，能否写一个models.SnowflakeIDField

如下是UUIDField写法

# models.py
import uuid

class TestUUID(models.Model):
    # 一定要设置default=uuid.uuid4
    id = models.UUIDField(primary_key=True, default=uuid.uuid4)
    name = models.CharField(max_length=100, verbose_name="名字")

# 新增操作

test_uuid = TestUUID()
test_uuid.name = '123123'
test_uuid.save()

-- 执行SQL
(0.005) INSERT INTO "polls_testuuid" ("id", "name") VALUES ('74f8e76c-b80a-4197-a469-81992c46e3c4'::uuid, '123123'); args=(UUID('74f8e76c-b80a-4197-a469-81992c46e3c4'), '123123')

TODO… 那么，SnowflakeIDField如何实现呢？

最终了解下来，发现完全没有必要自己写models.SnowflakeIDField

如下，即可完成…

# models.py
# 添加 default=snowflake_util.next_id 即可
id = models.BigIntegerField(primary_key=True, default=snowflake_util.next_id)

总结

• 选择上面的方案3，使用 Django Signals 中的pre_save()，可以增加对于参数的校验，不过本次不需要
• 选择上面的方案4，最简单