query_chain.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173
  1. """
  2. QueryChain is a wrapper for sequence of queries.
  3. Features:
  4. * Easy iteration for sequence of queries
  5. * Limit, offset and count which are applied to all queries in the chain
  6. * Smart __getitem__ support
  7. Initialization
  8. ^^^^^^^^^^^^^^
  9. QueryChain takes iterable of queries as first argument. Additionally limit and
  10. offset parameters can be given
  11. ::
  12. chain = QueryChain([session.query(User), session.query(Article)])
  13. chain = QueryChain(
  14. [session.query(User), session.query(Article)],
  15. limit=4
  16. )
  17. Simple iteration
  18. ^^^^^^^^^^^^^^^^
  19. ::
  20. chain = QueryChain([session.query(User), session.query(Article)])
  21. for obj in chain:
  22. print obj
  23. Limit and offset
  24. ^^^^^^^^^^^^^^^^
  25. Lets say you have 5 blog posts, 5 articles and 5 news items in your
  26. database.
  27. ::
  28. chain = QueryChain(
  29. [
  30. session.query(BlogPost),
  31. session.query(Article),
  32. session.query(NewsItem)
  33. ],
  34. limit=5
  35. )
  36. list(chain) # all blog posts but not articles and news items
  37. chain = chain.offset(4)
  38. list(chain) # last blog post, and first four articles
  39. Just like with original query object the limit and offset can be chained to
  40. return a new QueryChain.
  41. ::
  42. chain = chain.limit(5).offset(7)
  43. Chain slicing
  44. ^^^^^^^^^^^^^
  45. ::
  46. chain = QueryChain(
  47. [
  48. session.query(BlogPost),
  49. session.query(Article),
  50. session.query(NewsItem)
  51. ]
  52. )
  53. chain[3:6] # New QueryChain with offset=3 and limit=6
  54. Count
  55. ^^^^^
  56. Let's assume that there are five blog posts, five articles and five news
  57. items in the database, and you have the following query chain::
  58. chain = QueryChain(
  59. [
  60. session.query(BlogPost),
  61. session.query(Article),
  62. session.query(NewsItem)
  63. ]
  64. )
  65. You can then get the total number rows returned by the query chain
  66. with :meth:`~QueryChain.count`::
  67. >>> chain.count()
  68. 15
  69. """
  70. from copy import copy
  71. class QueryChain(object):
  72. """
  73. QueryChain can be used as a wrapper for sequence of queries.
  74. :param queries: A sequence of SQLAlchemy Query objects
  75. :param limit: Similar to normal query limit this parameter can be used for
  76. limiting the number of results for the whole query chain.
  77. :param offset: Similar to normal query offset this parameter can be used
  78. for offsetting the query chain as a whole.
  79. .. versionadded: 0.26.0
  80. """
  81. def __init__(self, queries, limit=None, offset=None):
  82. self.queries = queries
  83. self._limit = limit
  84. self._offset = offset
  85. def __iter__(self):
  86. consumed = 0
  87. skipped = 0
  88. for query in self.queries:
  89. query_copy = copy(query)
  90. if self._limit:
  91. query = query.limit(self._limit - consumed)
  92. if self._offset:
  93. query = query.offset(self._offset - skipped)
  94. obj_count = 0
  95. for obj in query:
  96. consumed += 1
  97. obj_count += 1
  98. yield obj
  99. if not obj_count:
  100. skipped += query_copy.count()
  101. else:
  102. skipped += obj_count
  103. def limit(self, value):
  104. return self[:value]
  105. def offset(self, value):
  106. return self[value:]
  107. def count(self):
  108. """
  109. Return the total number of rows this QueryChain's queries would return.
  110. """
  111. return sum(q.count() for q in self.queries)
  112. def __getitem__(self, key):
  113. if isinstance(key, slice):
  114. return self.__class__(
  115. queries=self.queries,
  116. limit=key.stop if key.stop is not None else self._limit,
  117. offset=key.start if key.start is not None else self._offset
  118. )
  119. else:
  120. for obj in self[key:1]:
  121. return obj
  122. def __repr__(self):
  123. return '<QueryChain at 0x%x>' % id(self)