veekun_pokedex/bin/reset-postgresql-sequences

#!/usr/bin/env python2

import sys

from sqlalchemy.sql import func
from sqlalchemy.types import Integer

from pokedex.db import connect
from pokedex.db.tables import mapped_classes

import optparse

parser = optparse.OptionParser()
parser.add_option('-v', '--verbose', action='store_true', default=False,
    help="verbose output")
parser.add_option('-e', '--engine', dest='engine_uri', default=None,
    help="database engine uri")

options, _ = parser.parse_args(sys.argv[1:])

session = connect(options.engine_uri)

assert session.connection().dialect.name == 'postgresql'

def sequence_columns(table):
    for c in table.__table__.primary_key:
        if c.autoincrement and not c.foreign_keys and \
           isinstance(c.type, Integer):
            yield c

for table in sorted(mapped_classes):
    table_name = table.__tablename__
    for c in sequence_columns(table):
        max_value, = session.query(func.max(c)).one()
        if options.verbose:
            print "%s.%s <- %s" % (table_name, c.name, max_value or 0)
        session.execute(func.setval(
            func.pg_get_serial_sequence(table_name, c.name),
            max_value or 1,
            max_value is not None,
        ))