Skip to content

Commit

Permalink
Use reflection to get correct PK column name (#17)
Browse files Browse the repository at this point in the history
Fixing [issue of using non-id primary key](#16) by [Marco Falcioni](https://github.com/marcofalcioni)
  • Loading branch information
marcofalcioni authored and M1ha-Shvn committed Sep 30, 2019
1 parent e2fd90f commit caac715
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 4 deletions.
2 changes: 1 addition & 1 deletion setup.py
Expand Up @@ -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',
Expand Down
6 changes: 3 additions & 3 deletions src/django_pg_returning/manager.py
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit caac715

Please sign in to comment.