49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
# Copyright (C) 2018-2021 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()
|