|
1 | 1 | """Convenience test mixin."""
|
2 | 2 | from django.core.management import call_command
|
3 |
| -from django.db import connection |
| 3 | +from django.db import connection, models |
4 | 4 |
|
5 | 5 | from .fixture import Fixture
|
6 | 6 |
|
@@ -46,15 +46,30 @@ def test_db_import(self):
|
46 | 46 |
|
47 | 47 | for model in self.dbdiff_models:
|
48 | 48 | if connection.vendor == 'postgresql':
|
| 49 | + pk_field = None |
| 50 | + for field in model._meta.get_fields(): |
| 51 | + if getattr(field, 'primary_key', False): |
| 52 | + pk_field = field |
| 53 | + break |
| 54 | + |
| 55 | + if not isinstance(pk_field, models.AutoField): |
| 56 | + continue |
| 57 | + |
| 58 | + if not pk_field.db_column: |
| 59 | + continue |
| 60 | + |
49 | 61 | reset = """
|
50 | 62 | SELECT
|
51 | 63 | setval(
|
52 |
| - pg_get_serial_sequence('%s', 'id'), |
53 |
| - coalesce(max(id),0) + 1, |
| 64 | + pg_get_serial_sequence('%(table)s', '%(column)s'), |
| 65 | + coalesce(max(%(column)s),0) + 1, |
54 | 66 | false
|
55 | 67 | )
|
56 |
| - FROM %s |
57 |
| - """ % (model._meta.db_table, model._meta.db_table) |
| 68 | + FROM %(table)s |
| 69 | + """ % dict( |
| 70 | + table=model._meta.db_table, |
| 71 | + column=pk_field.db_column, |
| 72 | + ) |
58 | 73 | else:
|
59 | 74 | raise NotImplemented()
|
60 | 75 | connection.cursor().execute(reset)
|
|
0 commit comments