31 import heapq

  1065 

  1066 

  1067 class MultiQuery(datastore.Query):

  1068   """Class representing a GQL query requiring multiple datastore queries.

  1069 

  1070   This class is actually a subclass of datastore.Query as it is intended to act

  1071   like a normal Query object (supporting the same interface).

  1072   """

  1073 

  1074   def __init__(self, bound_queries, orderings):

  1075     self.__bound_queries = bound_queries

  1076     self.__orderings = orderings

  1077 

  1078   def __str__(self):

  1079     res = 'MultiQuery: '

  1080     for query in self.__bound_queries:

  1081       res = '%s %s' % (res, str(query))

  1082     return res

  1083 

  1084   def Get(self, limit, offset=0):

  1085     """Get results of the query with a limit on the number of results.

  1086 

  1087     Args:

  1088       limit: maximum number of values to return.

  1089       offset: offset requested -- if nonzero, this will override the offset in

  1090               the original query

  1091 

  1092     Returns:

  1093       A list of entities with at most "limit" entries (less if the query

  1094       completes before reading limit values).

  1095     """

  1096     count = 1

  1097     result = []

  1098 

  1099     iterator = self.Run()

  1100 

  1101     try:

  1102       for i in xrange(offset):

  1103         val = iterator.next()

  1104     except StopIteration:

  1105       pass

  1106 

  1107     try:

  1108       while count <= limit:

  1109         val = iterator.next()

  1110         result.append(val)

  1111         count += 1

  1112     except StopIteration:

  1113       pass

  1114     return result

  1115 

  1116   class SortOrderEntity(object):

  1117     def __init__(self, entity_iterator, orderings):

  1118       self.__entity_iterator = entity_iterator

  1119       self.__entity = None

  1120       self.__min_max_value_cache = {}

  1121       try:

  1122         self.__entity = entity_iterator.next()

  1123       except StopIteration:

  1124         pass

  1125       else:

  1126         self.__orderings = orderings

  1127 

  1128     def __str__(self):

  1129       return str(self.__entity)

  1130 

  1131     def GetEntity(self):

  1132       return self.__entity

  1133 

  1134     def GetNext(self):

  1135       return MultiQuery.SortOrderEntity(self.__entity_iterator,

  1136                                         self.__orderings)

  1137 

  1138     def CmpProperties(self, that):

  1139       """Compare two entities and return their relative order.

  1140 

  1141       Compares self to that based on the current sort orderings and the

  1142       key orders between them. Returns negative, 0, or positive depending on

  1143       whether self is less, equal to, or greater than that. This

  1144       comparison returns as if all values were to be placed in ascending order

  1145       (highest value last).  Only uses the sort orderings to compare (ignores

  1146        keys).

  1147 

  1148       Args:

  1149         self: SortOrderEntity

  1150         that: SortOrderEntity

  1151 

  1152       Returns:

  1153         Negative if self < that

  1154         Zero if self == that

  1155         Positive if self > that

  1156       """

  1157       if not self.__entity:

  1158         return cmp(self.__entity, that.__entity)

  1159 

  1160       for (identifier, order) in self.__orderings:

  1161         value1 = self.__GetValueForId(self, identifier, order)

  1162         value2 = self.__GetValueForId(that, identifier, order)

  1163 

  1164         result = cmp(value1, value2)

  1165         if order == datastore.Query.DESCENDING:

  1166           result = -result

  1167         if result:

  1168           return result

  1169       return 0

  1170 

  1171     def __GetValueForId(self, sort_order_entity, identifier, sort_order):

  1172       value = sort_order_entity.__entity[identifier]

  1173       entity_key = sort_order_entity.__entity.key()

  1174       if self.__min_max_value_cache.has_key((entity_key, identifier)):

  1175         value = self.__min_max_value_cache[(entity_key, identifier)]

  1176       elif isinstance(value, list):

  1177         if sort_order == datastore.Query.DESCENDING:

  1178           value = min(value)

  1179         else:

  1180           value = max(value)

  1181         self.__min_max_value_cache[(entity_key, identifier)] = value

  1182 

  1183       return value

  1184 

  1185     def __cmp__(self, that):

  1186       """Compare self to that w.r.t. values defined in the sort order.

  1187 

  1188       Compare an entity with another, using sort-order first, then the key

  1189       order to break ties. This can be used in a heap to have faster min-value

  1190       lookup.

  1191 

  1192       Args:

  1193         that: other entity to compare to

  1194       Returns:

  1195         negative: if self is less than that in sort order

  1196         zero: if self is equal to that in sort order

  1197         positive: if self is greater than that in sort order

  1198       """

  1199       property_compare = self.CmpProperties(that)

  1200       if property_compare:

  1201         return property_compare

  1202       else:

  1203         return cmp(self.__entity.key(), that.__entity.key())

  1204 

  1205   def Run(self):

  1206     """Return an iterable output with all results in order."""

  1207     results = []

  1208     count = 1

  1209     for bound_query in self.__bound_queries:

  1210       logging.log(LOG_LEVEL, 'Running query #%i' % count)

  1211       results.append(bound_query.Run())

  1212       count += 1

  1213 

  1214     def IterateResults(results):

  1215       """Iterator function to return all results in sorted order.

  1216 

  1217       Iterate over the array of results, yielding the next element, in

  1218       sorted order. This function is destructive (results will be empty

  1219       when the operation is complete).

  1220 

  1221       Args:

  1222         results: list of result iterators to merge and iterate through

  1223 

  1224       Yields:

  1225         The next result in sorted order.

  1226       """

  1227       result_heap = []

  1228       for result in results:

  1229         heap_value = MultiQuery.SortOrderEntity(result, self.__orderings)

  1230         if heap_value.GetEntity():

  1231           heapq.heappush(result_heap, heap_value)

  1232 

  1233       used_keys = set()

  1234 

  1235       while result_heap:

  1236         top_result = heapq.heappop(result_heap)

  1237 

  1238         results_to_push = []

  1239         if top_result.GetEntity().key() not in used_keys:

  1240           yield top_result.GetEntity()

  1241         else:

  1242           pass

  1243 

  1244         used_keys.add(top_result.GetEntity().key())

  1245 

  1246         results_to_push = []

  1247         while result_heap:

  1248           next = heapq.heappop(result_heap)

  1249           if cmp(top_result, next):

  1250             results_to_push.append(next)

  1251             break

  1252           else:

  1253             results_to_push.append(next.GetNext())

  1254         results_to_push.append(top_result.GetNext())

  1255 

  1256         for popped_result in results_to_push:

  1257           if popped_result.GetEntity():

  1258             heapq.heappush(result_heap, popped_result)

  1259 

  1260     return IterateResults(results)

  1261 

  1262   def Count(self, limit=None):

  1263     """Return the number of matched entities for this query.

  1264 

  1265     Will return the de-duplicated count of results.  Will call the more

  1266     efficient Get() function if a limit is given.

  1267 

  1268     Args:

  1269       limit: maximum number of entries to count (for any result > limit, return

  1270       limit).

  1271     Returns:

  1272       count of the number of entries returned.

  1273     """

  1274     if limit is None:

  1275       count = 0

  1276       for value in self.Run():

  1277         count += 1

  1278       return count

  1279     else:

  1280       return len(self.Get(limit))

  1281