From 379aa1fac89428650e40ace3a672cdf43cc6aa73 Mon Sep 17 00:00:00 2001 From: Barry Warsaw Date: Sun, 23 Oct 2011 18:53:18 -0400 Subject: Add database specific pre- and post- reset callbacks. Use this to implement primary key reset for PostgreSQL, fixing several tests. --- src/mailman/database/postgresql.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) (limited to 'src/mailman/database/postgresql.py') diff --git a/src/mailman/database/postgresql.py b/src/mailman/database/postgresql.py index b59165243..4e40558a4 100644 --- a/src/mailman/database/postgresql.py +++ b/src/mailman/database/postgresql.py @@ -25,6 +25,7 @@ __all__ = [ ] +from operator import attrgetter from pkg_resources import resource_string from mailman.database.base import StormBaseDatabase @@ -38,10 +39,29 @@ class PostgreSQLDatabase(StormBaseDatabase): """See `BaseDatabase`.""" table_query = ('SELECT table_name FROM information_schema.tables ' "WHERE table_schema = 'public'") - table_names = set(item[0] for item in + table_names = set(item[0] for item in store.execute(table_query)) return 'version' in table_names def _get_schema(self): """See `BaseDatabase`.""" return resource_string('mailman.database.sql', 'postgres.sql') + + def _post_reset(self, store): + """PostgreSQL-specific test suite cleanup. + + Reset the _id_seq.last_value so that primary key ids + restart from zero for new tests. + """ + from mailman.database.model import ModelMeta + classes = sorted(ModelMeta._class_registry, + key=attrgetter('__storm_table__')) + # Recipe adapted from + # http://stackoverflow.com/questions/544791/ + # django-postgresql-how-to-reset-primary-key + for model_class in classes: + store.execute("""\ + SELECT setval('"{0}_id_seq"', coalesce(max("id"), 1), + max("id") IS NOT null) + FROM "{0}"; + """.format(model_class.__storm_table__)) -- cgit v1.2.3-70-g09d2