Django reset auto-increment pk/id field for production

后端 未结 2 1458
孤城傲影
孤城傲影 2021-01-02 15:09

(I\'m new to Django, Python, and Postgresql) I\'ve been adding and deleting data in my development and noticed that the pk keeps adding up and never reset to 1 even if I del

2条回答
  •  清酒与你
    2021-01-02 16:00

    You could generate a command that gathers all of the applications in your system and runs a dynamically generated reset statement for each table, while also dynamically gathering the PK column names for each table you loop over in the event you didn't name them all the same value.

    To run: python manage.py reset_sequences

    import psycopg2
    from django.conf import settings
    from django.core.management.base import BaseCommand
    from django.db import connections
    
    
    def dictfetchall(cursor):
        """Return all rows from a cursor as a dict"""
        columns = [col[0] for col in cursor.description]
        return [
            dict(zip(columns, row))
            for row in cursor.fetchall()
        ]
    
    
    class Command(BaseCommand):
        help = "Resets sequencing errors in Postgres which normally occur due to importing/restoring a DB"
    
        def handle(self, *args, **options):
            # loop over all databases in system to figure out the tables that need to be reset
            for name_to_use_for_connection, connection_settings in settings.DATABASES.items():
                db_name = connection_settings['NAME']
                host = connection_settings['HOST']
                user = connection_settings['USER']
                port = connection_settings['PORT']
                password = connection_settings['PASSWORD']
    
                # connect to this specific DB
                conn_str = f"host={host} port={port} user={user} password={password}"
    
                conn = psycopg2.connect(conn_str)
                conn.autocommit = True
    
                select_all_table_statement = f"""SELECT *
                                        FROM information_schema.tables
                                        WHERE table_schema = 'public'
                                        ORDER BY table_name;
                                    """
                # just a visual representation of where we are
                print('-' * 20, db_name)
                try:
                    not_reset_tables = list()
                    # use the specific name for the DB
                    with connections[name_to_use_for_connection].cursor() as cursor:
                        # using the current db as the cursor connection
                        cursor.execute(select_all_table_statement)
                        rows = dictfetchall(cursor)
                        # will loop over table names in the connected DB
                        for row in rows:
                            find_pk_statement = f"""
                                SELECT k.COLUMN_NAME
                                FROM information_schema.table_constraints t
                                LEFT JOIN information_schema.key_column_usage k
                                USING(constraint_name,table_schema,table_name)
                                WHERE t.constraint_type='PRIMARY KEY'
                                    AND t.table_name='{row['table_name']}';
                            """
                            cursor.execute(find_pk_statement)
                            pk_column_names = dictfetchall(cursor)
                            for pk_dict in pk_column_names:
                                column_name = pk_dict['column_name']
    
                            # time to build the reset sequence command for each table
                            # taken from django: https://docs.djangoproject.com/en/3.0/ref/django-admin/#sqlsequencereset
                            # example: SELECT setval(pg_get_serial_sequence('"[TABLE]"','id'), coalesce(max("id"), 1), max("id") IS NOT null) FROM "[TABLE]";
                            try:
                                reset_statement = f"""SELECT setval(pg_get_serial_sequence('"{row['table_name']}"','{column_name}'), 
                                                        coalesce(max("{column_name}"), 1), max("{column_name}") IS NOT null) FROM "{row['table_name']}" """
                                cursor.execute(reset_statement)
                                return_values = dictfetchall(cursor)
                                # will be 1 row
                                for value in return_values:
                                    print(f"Sequence reset to {value['setval']} for {row['table_name']}")
                            except Exception as ex:
                                # will only fail if PK is not an integer...
                                # currently in my system this is from django.contrib.sessions
                                not_reset_tables.append(f"{row['table_name']} not reset")
    
                except psycopg2.Error as ex:
                    raise SystemExit(f'Error: {ex}')
    
                conn.close()
                print('-' * 5, ' ALL ERRORS ', '-' * 5)
                for item_statement in not_reset_tables:
                    # shows which tables produced errors, so far I have only
                    # seen this with PK's that are not integers because of the MAX() method
                    print(item_statement)
    
                # just a visual representation of where we are
                print('-' * 20, db_name)
    
    
    
    

提交回复
热议问题