Skip to content

Commit 2d790ce

Browse files
committed
Fix bug with outdated models
This bug would occur when the database would return a column that does not exists on the model. The fix is to not blindly pass all data returned by the database to the model. We now filter it and make sure that the field exists on the model.
1 parent aaa094d commit 2d790ce

File tree

2 files changed

+51
-5
lines changed

2 files changed

+51
-5
lines changed

psqlextra/manager.py

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import List, Dict
1+
from typing import Dict, List
22

33
import django
44
from django.conf import settings
@@ -63,7 +63,23 @@ def upsert_and_get(self, conflict_target: List, fields: Dict):
6363

6464
compiler = self._build_upsert_compiler(conflict_target, fields)
6565
column_data = compiler.execute_sql(return_id=False)
66-
return self.model(**column_data)
66+
67+
# get a list of columns that are officially part of the model
68+
model_columns = [
69+
field.column
70+
for field in self.model._meta.local_concrete_fields
71+
]
72+
73+
# strip out any columns/fields returned by the db that
74+
# are not present in the model
75+
model_init_fields = {}
76+
for column_name, column_value in column_data.items():
77+
if column_name not in model_columns:
78+
continue
79+
80+
model_init_fields[column_name] = column_value
81+
82+
return self.model(**model_init_fields)
6783

6884
def _build_upsert_compiler(self, conflict_target: List, kwargs):
6985
"""Builds the SQL compiler for a insert/update query.

tests/test_upsert.py

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
1+
import pytest
2+
from django.core.exceptions import SuspiciousOperation
3+
from django.db import connection, models
14
from django.test import TestCase
25

36
from psqlextra import HStoreField
47

58
from .fake_model import get_fake_model
6-
from django.db import models
7-
import pytest
89

9-
from django.core.exceptions import SuspiciousOperation
1010

1111
@pytest.mark.django_db
1212
class UpsertTest(TestCase):
@@ -172,3 +172,33 @@ def test_invalid_conflict_target(self):
172172
title='beer'
173173
)
174174
)
175+
176+
def test_outdated_model(self):
177+
"""Tests whether upsert_and_get properly handles
178+
fields that are in the database but not on the model.
179+
180+
This happens if somebody manually modified the database
181+
to add a column that is not present in the model.
182+
183+
This should be handled properly by ignoring the column
184+
returned by the database.
185+
"""
186+
187+
model = get_fake_model({
188+
'title': models.CharField(max_length=140, unique=True)
189+
})
190+
191+
# manually create the colum that is not on the model
192+
with connection.cursor() as cursor:
193+
cursor.execute((
194+
'ALTER TABLE {table} '
195+
'ADD COLUMN beer character varying(50);'
196+
).format(table=model._meta.db_table))
197+
198+
# without proper handling, this would fail with a TypeError
199+
model.objects.upsert_and_get(
200+
conflict_target=['title'],
201+
fields=dict(
202+
title='beer'
203+
)
204+
)

0 commit comments

Comments
 (0)