|
3 | 3 | import uuid |
4 | 4 |
|
5 | 5 | from django.contrib.contenttypes.models import ContentType |
6 | | -from django.db import models |
| 6 | +from django.db import models, connection |
7 | 7 | from django.db.models import Case, Count, FilteredRelation, Q, Sum, When, F |
8 | 8 | from django.db.utils import IntegrityError, NotSupportedError |
9 | 9 | from django.test import TransactionTestCase |
| 10 | +from django.test.utils import CaptureQueriesContext |
10 | 11 |
|
11 | 12 | from polymorphic import query_translate |
12 | 13 | from polymorphic.managers import PolymorphicManager |
@@ -1234,3 +1235,179 @@ def test_refresh_from_db_fields(self): |
1234 | 1235 | def test_non_polymorphic_parent(self): |
1235 | 1236 | obj = NonPolymorphicParent.objects.create() |
1236 | 1237 | assert obj.delete() |
| 1238 | + |
| 1239 | + def test_iteration(self): |
| 1240 | + Model2A.objects.all().delete() |
| 1241 | + |
| 1242 | + for i in range(250): |
| 1243 | + Model2B.objects.create(field1=f"B1-{i}", field2=f"B2-{i}") |
| 1244 | + for i in range(1000): |
| 1245 | + Model2C.objects.create( |
| 1246 | + field1=f"C1-{i + 250}", field2=f"C2-{i + 250}", field3=f"C3-{i + 250}" |
| 1247 | + ) |
| 1248 | + for i in range(2000): |
| 1249 | + Model2D.objects.create( |
| 1250 | + field1=f"D1-{i + 1250}", |
| 1251 | + field2=f"D2-{i + 1250}", |
| 1252 | + field3=f"D3-{i + 1250}", |
| 1253 | + field4=f"D4-{i + 1250}", |
| 1254 | + ) |
| 1255 | + |
| 1256 | + with CaptureQueriesContext(connection) as base_all: |
| 1257 | + for _ in Model2A.objects.non_polymorphic().all(): |
| 1258 | + pass # Evaluating the queryset |
| 1259 | + |
| 1260 | + len_base_all = len(base_all) |
| 1261 | + assert len_base_all == 1, ( |
| 1262 | + f"Expected 1 queries for chunked iteration over 3250 base objects. {len_base_all}" |
| 1263 | + ) |
| 1264 | + |
| 1265 | + with CaptureQueriesContext(connection) as base_iterator: |
| 1266 | + for _ in Model2A.objects.non_polymorphic().iterator(): |
| 1267 | + pass # Evaluating the queryset |
| 1268 | + |
| 1269 | + len_base_iterator = len(base_iterator) |
| 1270 | + assert len_base_iterator == 1, ( |
| 1271 | + f"Expected 1 queries for chunked iteration over 3250 base objects. {len_base_iterator}" |
| 1272 | + ) |
| 1273 | + |
| 1274 | + with CaptureQueriesContext(connection) as base_chunked: |
| 1275 | + for _ in Model2A.objects.non_polymorphic().iterator(chunk_size=1000): |
| 1276 | + pass # Evaluating the queryset |
| 1277 | + |
| 1278 | + len_base_chunked = len(base_chunked) |
| 1279 | + assert len_base_chunked == 1, ( |
| 1280 | + f"Expected 1 queries for chunked iteration over 3250 base objects. {len_base_chunked}" |
| 1281 | + ) |
| 1282 | + |
| 1283 | + with CaptureQueriesContext(connection) as poly_all: |
| 1284 | + b, c, d = 0, 0, 0 |
| 1285 | + for idx, obj in enumerate(reversed(list(Model2A.objects.order_by("-pk").all()))): |
| 1286 | + if isinstance(obj, Model2D): |
| 1287 | + d += 1 |
| 1288 | + assert obj.field1 == f"D1-{idx}" |
| 1289 | + assert obj.field2 == f"D2-{idx}" |
| 1290 | + assert obj.field3 == f"D3-{idx}" |
| 1291 | + assert obj.field4 == f"D4-{idx}" |
| 1292 | + elif isinstance(obj, Model2C): |
| 1293 | + c += 1 |
| 1294 | + assert obj.field1 == f"C1-{idx}" |
| 1295 | + assert obj.field2 == f"C2-{idx}" |
| 1296 | + assert obj.field3 == f"C3-{idx}" |
| 1297 | + elif isinstance(obj, Model2B): |
| 1298 | + b += 1 |
| 1299 | + assert obj.field1 == f"B1-{idx}" |
| 1300 | + assert obj.field2 == f"B2-{idx}" |
| 1301 | + else: |
| 1302 | + assert False, "Unexpected model type" |
| 1303 | + assert (b, c, d) == (250, 1000, 2000) |
| 1304 | + |
| 1305 | + assert len(poly_all) <= 7, ( |
| 1306 | + f"Expected < 7 queries for chunked iteration over 3250 " |
| 1307 | + f"objects with 3 child models and the default chunk size of 2000, encountered " |
| 1308 | + f"{len(poly_all)}" |
| 1309 | + ) |
| 1310 | + |
| 1311 | + with CaptureQueriesContext(connection) as poly_all: |
| 1312 | + b, c, d = 0, 0, 0 |
| 1313 | + for idx, obj in enumerate(Model2A.objects.order_by("pk").iterator(chunk_size=None)): |
| 1314 | + if isinstance(obj, Model2D): |
| 1315 | + d += 1 |
| 1316 | + assert obj.field1 == f"D1-{idx}" |
| 1317 | + assert obj.field2 == f"D2-{idx}" |
| 1318 | + assert obj.field3 == f"D3-{idx}" |
| 1319 | + assert obj.field4 == f"D4-{idx}" |
| 1320 | + elif isinstance(obj, Model2C): |
| 1321 | + c += 1 |
| 1322 | + assert obj.field1 == f"C1-{idx}" |
| 1323 | + assert obj.field2 == f"C2-{idx}" |
| 1324 | + assert obj.field3 == f"C3-{idx}" |
| 1325 | + elif isinstance(obj, Model2B): |
| 1326 | + b += 1 |
| 1327 | + assert obj.field1 == f"B1-{idx}" |
| 1328 | + assert obj.field2 == f"B2-{idx}" |
| 1329 | + else: |
| 1330 | + assert False, "Unexpected model type" |
| 1331 | + assert (b, c, d) == (250, 1000, 2000) |
| 1332 | + |
| 1333 | + assert len(poly_all) <= 7, ( |
| 1334 | + f"Expected < 7 queries for chunked iteration over 3250 " |
| 1335 | + f"objects with 3 child models and a chunk size of 2000, encountered " |
| 1336 | + f"{len(poly_all)}" |
| 1337 | + ) |
| 1338 | + |
| 1339 | + with CaptureQueriesContext(connection) as poly_iterator: |
| 1340 | + b, c, d = 0, 0, 0 |
| 1341 | + for idx, obj in enumerate(Model2A.objects.order_by("pk").iterator()): |
| 1342 | + if isinstance(obj, Model2D): |
| 1343 | + d += 1 |
| 1344 | + assert obj.field1 == f"D1-{idx}" |
| 1345 | + assert obj.field2 == f"D2-{idx}" |
| 1346 | + assert obj.field3 == f"D3-{idx}" |
| 1347 | + assert obj.field4 == f"D4-{idx}" |
| 1348 | + elif isinstance(obj, Model2C): |
| 1349 | + c += 1 |
| 1350 | + assert obj.field1 == f"C1-{idx}" |
| 1351 | + assert obj.field2 == f"C2-{idx}" |
| 1352 | + assert obj.field3 == f"C3-{idx}" |
| 1353 | + elif isinstance(obj, Model2B): |
| 1354 | + b += 1 |
| 1355 | + assert obj.field1 == f"B1-{idx}" |
| 1356 | + assert obj.field2 == f"B2-{idx}" |
| 1357 | + else: |
| 1358 | + assert False, "Unexpected model type" |
| 1359 | + assert (b, c, d) == (250, 1000, 2000) |
| 1360 | + |
| 1361 | + assert len(poly_iterator) <= 7, ( |
| 1362 | + f"Expected <= 7 queries for chunked iteration over 3250 " |
| 1363 | + f"objects with 3 child models and a default chunk size of 2000, encountered " |
| 1364 | + f"{len(poly_iterator)}" |
| 1365 | + ) |
| 1366 | + |
| 1367 | + with CaptureQueriesContext(connection) as poly_chunked: |
| 1368 | + b, c, d = 0, 0, 0 |
| 1369 | + for idx, obj in enumerate(Model2A.objects.order_by("pk").iterator(chunk_size=4000)): |
| 1370 | + if isinstance(obj, Model2D): |
| 1371 | + d += 1 |
| 1372 | + assert obj.field1 == f"D1-{idx}" |
| 1373 | + assert obj.field2 == f"D2-{idx}" |
| 1374 | + assert obj.field3 == f"D3-{idx}" |
| 1375 | + assert obj.field4 == f"D4-{idx}" |
| 1376 | + elif isinstance(obj, Model2C): |
| 1377 | + c += 1 |
| 1378 | + assert obj.field1 == f"C1-{idx}" |
| 1379 | + assert obj.field2 == f"C2-{idx}" |
| 1380 | + assert obj.field3 == f"C3-{idx}" |
| 1381 | + elif isinstance(obj, Model2B): |
| 1382 | + b += 1 |
| 1383 | + assert obj.field1 == f"B1-{idx}" |
| 1384 | + assert obj.field2 == f"B2-{idx}" |
| 1385 | + else: |
| 1386 | + assert False, "Unexpected model type" |
| 1387 | + assert (b, c, d) == (250, 1000, 2000) |
| 1388 | + |
| 1389 | + assert len(poly_chunked) <= 7, ( |
| 1390 | + f"Expected <= 7 queries for chunked iteration over 3250 objects with 3 child " |
| 1391 | + f"models and a chunk size of 4000, encountered {len(poly_chunked)}" |
| 1392 | + ) |
| 1393 | + |
| 1394 | + if connection.vendor == "postgresql": |
| 1395 | + assert len(poly_chunked) == 4, "On postgres with a 4000 chunk size, expected 4 queries" |
| 1396 | + |
| 1397 | + try: |
| 1398 | + result = Model2A.objects.all().delete() |
| 1399 | + assert result == ( |
| 1400 | + 11500, |
| 1401 | + { |
| 1402 | + "tests.Model2D": 2000, |
| 1403 | + "tests.Model2C": 3000, |
| 1404 | + "tests.Model2A": 3250, |
| 1405 | + "tests.Model2B": 3250, |
| 1406 | + }, |
| 1407 | + ) |
| 1408 | + except AttributeError: |
| 1409 | + if connection.vendor == "oracle": |
| 1410 | + # FIXME |
| 1411 | + # known deletion issue with oracle |
| 1412 | + # https://github.com/jazzband/django-polymorphic/issues/673 |
| 1413 | + pass |
0 commit comments