nk20-scripts/management/commands/syncsql.py

# Copyright (C) 2018-2020 by BDE ENS Paris-Saclay
# SPDX-License-Identifier: GPL-3.0-or-later

from django.core.management.base import BaseCommand
from django.apps import apps
from django.db import connection


from polymorphic.models import PolymorphicModel

NO_SEQ = [
    "Session",
    "Token",
    "WEIRole",  # dirty fix
]

class Command(BaseCommand):
    """
    Command to synchronise primary sequence of postgres after bulk insert of django.
    """

    def add_arguments(self,parser):
        parser.add_argument('apps', type=str,nargs='*',help='applications which table would be resynchronized')
        return parser

    def handle(self, *args, **kwargs):
        app_list = kwargs["apps"]
        if len(app_list):
            model_classes = list()
            for app in app_list:
                model_classes += apps.get_app_config(app).get_models()
        else:
            # no app specified, sync everything
            model_classes = apps.get_models(include_auto_created=True)

        db_names = [
            m._meta.db_table for m in model_classes
            if m.__base__.__base__ is not PolymorphicModel and m.__name__ not in NO_SEQ and m.objects.count() > 1
        ]
        com = "BEGIN;\n"
        for db_name in db_names:
            com += f'SELECT setval(pg_get_serial_sequence(\'"{db_name}"\',\'id\'), coalesce(max("id"), 1),' \
                   f' max("id") IS NOT null) FROM "{db_name}";\n'
        com += "COMMIT;"
        print(com)
        cur = connection.cursor()
        cur.execute(com)
        cur.close()