创建Django模型时自动触发Citus自定义迁移操作的实现方法

你可以通过自定义Django模型基类+扩展迁移自动生成逻辑的方式实现需求，无需为每个模型硬编码Distribute操作，具体步骤如下：

步骤1：定义通用基础模型 #

让你的所有业务模型继承这个抽象基类，统一读取Meta中的reference标记：

from django.db import models

class CitusModel(models.Model):
    class Meta:
        abstract = True

    @classmethod
    def is_reference_table(cls):
        return getattr(cls._meta, "reference", False)

步骤2：扩展迁移自动生成逻辑 #

在你项目的任意app的management/commands目录下创建makemigrations.py，重写原生的makemigrations命令，自动为带reference标记的模型注入Distribute操作：

from django.core.management.commands.makemigrations import Command as BaseCommand
from django.db.migrations import Migration
from django.apps import apps
from tenant_migrations import Distribute

class Command(BaseCommand):
    def write_migration_files(self, changes):
        # 先执行原生生成逻辑
        super().write_migration_files(changes)
        # 扫描所有带reference标记的模型，追加Distribute操作
        for app_label, migrations in changes.items():
            for migration in migrations:
                if isinstance(migration, Migration):
                    # 收集当前app下所有reference=True的模型名
                    ref_models = [
                        model._meta.object_name
                        for model in apps.get_app_config(app_label).get_models()
                        if getattr(model._meta, "reference", False)
                    ]
                    # 追加Distribute操作到迁移操作列表末尾
                    for model_name in ref_models:
                        migration.operations.append(Distribute(model_name))
                # 重新写入修改后的迁移文件
                self.write_migration_file(app_label, migration)

步骤3：业务模型使用示例 #

后续新增模型只要在Meta中指定reference = True即可，执行makemigrations时会自动生成对应的Distribute操作：

from .base import CitusModel

class Task(CitusModel):
    name = models.CharField(max_length=200)
    # 其他字段定义

    class Meta:
        reference = True

注意事项 #

• 如果是已有项目首次接入，建议先执行一次全量扫描生成所有reference表的Distribute迁移，避免遗漏
• 如果你需要按租户ID分发的分布式表而非引用表，只需要在Meta中新增distribute_by属性，再修改上述自动生成逻辑适配即可
• 操作生成的迁移文件完全符合Django的迁移规范，支持回滚等原生操作

内容的提问来源于stack exchange，提问作者reindeer