|
1 | 1 | from bson import SON, ObjectId |
2 | 2 | from django.db import models |
| 3 | +from django.db.models import Case, F, IntegerField, Value, When |
3 | 4 | from django.test import TestCase |
4 | 5 |
|
5 | 6 | from django_mongodb_backend.test import MongoTestCaseMixin |
@@ -1016,3 +1017,166 @@ def test_or_mixed_local_remote_pushdown(self): |
1016 | 1017 | {"$match": {"$or": [{"title": "B1"}, {"queries__author.name": "Bob"}]}}, |
1017 | 1018 | ], |
1018 | 1019 | ) |
| 1020 | + |
| 1021 | + def test_conditional_expression_not_pushed(self): |
| 1022 | + a1 = Author.objects.create(name="Vicente") |
| 1023 | + a2 = Author.objects.create(name="Carlos") |
| 1024 | + a3 = Author.objects.create(name="Maria") |
| 1025 | + |
| 1026 | + Book.objects.create(title="B1", author=a1, isbn="111") |
| 1027 | + b2 = Book.objects.create(title="B2", author=a2, isbn="222") |
| 1028 | + Book.objects.create(title="B3", author=a3, isbn="333") |
| 1029 | + |
| 1030 | + annotated = Book.objects.annotate( |
| 1031 | + score=Case( |
| 1032 | + When(author__name="Vicente", then=Value(2)), |
| 1033 | + When(author__name="Carlos", then=Value(4)), |
| 1034 | + default=Value(1), |
| 1035 | + output_field=IntegerField(), |
| 1036 | + ) |
| 1037 | + + Value(1) |
| 1038 | + ).filter(score__gt=3) |
| 1039 | + |
| 1040 | + with self.assertNumQueries(1) as ctx: |
| 1041 | + self.assertSequenceEqual(annotated, [b2]) |
| 1042 | + self.assertAggregateQuery( |
| 1043 | + ctx.captured_queries[0]["sql"], |
| 1044 | + "queries__book", |
| 1045 | + [ |
| 1046 | + { |
| 1047 | + "$lookup": { |
| 1048 | + "from": "queries__author", |
| 1049 | + "let": {"parent__field__0": "$author_id"}, |
| 1050 | + "pipeline": [ |
| 1051 | + { |
| 1052 | + "$match": { |
| 1053 | + "$expr": {"$and": [{"$eq": ["$$parent__field__0", "$_id"]}]} |
| 1054 | + } |
| 1055 | + } |
| 1056 | + ], |
| 1057 | + "as": "queries__author", |
| 1058 | + } |
| 1059 | + }, |
| 1060 | + {"$unwind": "$queries__author"}, |
| 1061 | + { |
| 1062 | + "$match": { |
| 1063 | + "$expr": { |
| 1064 | + "$gt": [ |
| 1065 | + { |
| 1066 | + "$add": [ |
| 1067 | + { |
| 1068 | + "$switch": { |
| 1069 | + "branches": [ |
| 1070 | + { |
| 1071 | + "case": { |
| 1072 | + "$eq": [ |
| 1073 | + "$queries__author.name", |
| 1074 | + "Vicente", |
| 1075 | + ] |
| 1076 | + }, |
| 1077 | + "then": {"$literal": 2}, |
| 1078 | + }, |
| 1079 | + { |
| 1080 | + "case": { |
| 1081 | + "$eq": [ |
| 1082 | + "$queries__author.name", |
| 1083 | + "Carlos", |
| 1084 | + ] |
| 1085 | + }, |
| 1086 | + "then": {"$literal": 4}, |
| 1087 | + }, |
| 1088 | + ], |
| 1089 | + "default": {"$literal": 1}, |
| 1090 | + } |
| 1091 | + }, |
| 1092 | + {"$literal": 1}, |
| 1093 | + ] |
| 1094 | + }, |
| 1095 | + 3, |
| 1096 | + ] |
| 1097 | + } |
| 1098 | + } |
| 1099 | + }, |
| 1100 | + { |
| 1101 | + "$project": { |
| 1102 | + "score": { |
| 1103 | + "$add": [ |
| 1104 | + { |
| 1105 | + "$switch": { |
| 1106 | + "branches": [ |
| 1107 | + { |
| 1108 | + "case": { |
| 1109 | + "$eq": ["$queries__author.name", "Vicente"] |
| 1110 | + }, |
| 1111 | + "then": {"$literal": 2}, |
| 1112 | + }, |
| 1113 | + { |
| 1114 | + "case": { |
| 1115 | + "$eq": ["$queries__author.name", "Carlos"] |
| 1116 | + }, |
| 1117 | + "then": {"$literal": 4}, |
| 1118 | + }, |
| 1119 | + ], |
| 1120 | + "default": {"$literal": 1}, |
| 1121 | + } |
| 1122 | + }, |
| 1123 | + {"$literal": 1}, |
| 1124 | + ] |
| 1125 | + }, |
| 1126 | + "_id": 1, |
| 1127 | + "title": 1, |
| 1128 | + "author_id": 1, |
| 1129 | + "isbn": 1, |
| 1130 | + } |
| 1131 | + }, |
| 1132 | + ], |
| 1133 | + ) |
| 1134 | + |
| 1135 | + def test_simple_annotation_pushdown(self): |
| 1136 | + a1 = Author.objects.create(name="Alice") |
| 1137 | + a2 = Author.objects.create(name="Bob") |
| 1138 | + b1 = Book.objects.create(title="B1", author=a1, isbn="111") |
| 1139 | + Book.objects.create(title="B2", author=a2, isbn="222") |
| 1140 | + b3 = Book.objects.create(title="B3", author=a1, isbn="333") |
| 1141 | + qs = Book.objects.annotate(name_length=F("author__name")).filter(name_length="Alice") |
| 1142 | + expected = [b1, b3] |
| 1143 | + with self.assertNumQueries(1) as ctx: |
| 1144 | + self.assertSequenceEqual(qs, expected) |
| 1145 | + self.assertAggregateQuery( |
| 1146 | + ctx.captured_queries[0]["sql"], |
| 1147 | + "queries__book", |
| 1148 | + [ |
| 1149 | + { |
| 1150 | + "$lookup": { |
| 1151 | + "from": "queries__author", |
| 1152 | + "let": {"parent__field__0": "$author_id"}, |
| 1153 | + "pipeline": [ |
| 1154 | + { |
| 1155 | + "$match": { |
| 1156 | + "$and": [ |
| 1157 | + { |
| 1158 | + "$expr": { |
| 1159 | + "$and": [{"$eq": ["$$parent__field__0", "$_id"]}] |
| 1160 | + } |
| 1161 | + }, |
| 1162 | + {"name": "Alice"}, |
| 1163 | + ] |
| 1164 | + } |
| 1165 | + } |
| 1166 | + ], |
| 1167 | + "as": "queries__author", |
| 1168 | + } |
| 1169 | + }, |
| 1170 | + {"$unwind": "$queries__author"}, |
| 1171 | + {"$match": {"queries__author.name": "Alice"}}, |
| 1172 | + { |
| 1173 | + "$project": { |
| 1174 | + "queries__author": {"name_length": "$queries__author.name"}, |
| 1175 | + "_id": 1, |
| 1176 | + "title": 1, |
| 1177 | + "author_id": 1, |
| 1178 | + "isbn": 1, |
| 1179 | + } |
| 1180 | + }, |
| 1181 | + ], |
| 1182 | + ) |
0 commit comments