Update sequences after inserting data in PostgreSQL

[Crocmagnon]

• dbcut version: dbcut, version 0.1.4
• Python version: Python 3.6.10
• Operating System: macOS 10.15.5 (up to date)

Description

dbcut should update the sequences in a PostgreSQL database after the data is inserted. Currently, sequences are not updated and therefore inserts following the data load fail because the PK it tries to insert with already exists.

Workaround

I put together a quick and (very) dirty script, only tested on postgresql 12.3, to reset the sequences to the max value of the index after inserting the data (only on "id" columns)

from sqlalchemy import create_engine
from sqlalchemy.sql import text

engine = create_engine("postgres://user:password@localhost:5432/dbname")
conn = engine.connect()
query = text("SELECT c.relname FROM pg_class c WHERE c.relkind = 'r';")
res = conn.execute(query)
all_tables = map(lambda x: x[0], res.fetchall())
query = text("SELECT c.relname FROM pg_class c WHERE c.relkind = 'S';")
res = conn.execute(query)
for sequence in res.fetchall():
    sequence = sequence[0]
    table = sequence.replace("_id_seq", "")
    if table not in all_tables:
        continue
    query = text("SELECT COALESCE(MAX(id), 0) + 1 as value FROM {};".format(table))
    res = conn.execute(query)
    data = res.fetchall()
    if not data:
        continue
    last_id = data[0][0]
    query = text("ALTER SEQUENCE {} RESTART WITH :last_id;".format(sequence))
    conn.execute(query, sequence=sequence, last_id=last_id)