From caac7158f1d7823dde4ff5034594aa8503a98250 Mon Sep 17 00:00:00 2001 From: Marco Falcioni Date: Sun, 29 Sep 2019 22:16:04 -0700 Subject: [PATCH] Use reflection to get correct PK column name (#17) Fixing [issue of using non-id primary key](https://github.com/M1hacka/django-pg-returning/issues/16) by [Marco Falcioni](https://github.com/marcofalcioni) --- setup.py | 2 +- src/django_pg_returning/manager.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/setup.py b/setup.py index d135a42..5d07898 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,7 @@ setup( name='django-pg-returning', - version='1.2.1', + version='1.2.2', packages=['django_pg_returning'], package_dir={'': 'src'}, url='https://github.com/M1hacka/django-pg-returning', diff --git a/src/django_pg_returning/manager.py b/src/django_pg_returning/manager.py index aec81f2..7060966 100644 --- a/src/django_pg_returning/manager.py +++ b/src/django_pg_returning/manager.py @@ -39,7 +39,7 @@ def _insert(self, objs, fields, return_id=False, raw=False, using=None, ignore_c query.insert_values(fields, objs, raw=raw) self.model._insert_returning_cache = self._execute_sql(query, return_fields, using=using) - return self.model._insert_returning_cache.values_list('id', flat=True) if return_id else None + return self.model._insert_returning_cache.values_list(self.model._meta.pk.column, flat=True) if return_id else None _insert.alters_data = True _insert.queryset_only = False @@ -169,9 +169,9 @@ def bulk_create_returning(self, objs, batch_size=None): # Replace values fetched from returned data if result and result[0].pk: # For django 1.10+ where objects can be matched - values_dict = {item['id']: item for item in self.model._insert_returning_cache.values()} + values_dict = {item[self.model._meta.pk.column]: item for item in self.model._insert_returning_cache.values()} for item in result: - for k, v in values_dict[item.id].items(): + for k, v in values_dict[item.pk].items(): setattr(item, k, v) else: # For django before 1.10 which doesn't fetch primary key