Skip to content

Commit 69c1d24

Browse files
committed
INTPYTHON-348 add support for QuerySet.raw_mql()
- Tim tips and pair w/Jib
1 parent 7a09e2b commit 69c1d24

File tree

1 file changed

+45
-3
lines changed

1 file changed

+45
-3
lines changed

django_mongodb/query.py

Lines changed: 45 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -319,6 +319,36 @@ def raw_mql(self, raw_query, params=(), translations=None, using=None):
319319

320320

321321
class MongoRawQuery(RawQuery):
322+
323+
324+
def __init__(self, sql, using, model, params=()):
325+
self.params = params
326+
self.sql = sql
327+
self.using = using
328+
self.cursor = None
329+
self.model = model
330+
331+
# Mirror some properties of a normal query so that
332+
# the compiler can be used to process results.
333+
self.low_mark, self.high_mark = 0, None # Used for offset/limit
334+
self.extra_select = {}
335+
self.annotation_select = {}
336+
337+
def __iter__(self):
338+
# Always execute a new query for a new iterator.
339+
# This could be optimized with a cache at the expense of RAM.
340+
341+
self.cursor = self._execute_query()
342+
343+
if not connections[self.using].features.can_use_chunked_reads:
344+
# If the database can't use chunked reads we need to make sure we
345+
# evaluate the entire query up front.
346+
result = list(self.cursor)
347+
else:
348+
result = self.cursor
349+
350+
return iter(result)
351+
322352
def _execute_query(self):
323353
connection = connections[self.using]
324354
params_type = self.params_type
@@ -332,8 +362,19 @@ def _execute_query(self):
332362
else:
333363
raise RuntimeError("Unexpected params type: %s" % params_type)
334364

335-
self.cursor = connection.cursor()
336-
self.cursor.execute(self.sql, params)
365+
# self.cursor = connection.cursor()
366+
# self.cursor.execute(self.sql, params)
367+
368+
collection = connection.get_collection(self.model._meta.db_table)
369+
370+
return collection.aggregate(self.sql)
371+
372+
def get_columns(self):
373+
if self.cursor is None:
374+
self._execute_query()
375+
converter = connections[self.using].introspection.identifier_converter
376+
self.cursor.description = ""
377+
return [converter(column_meta[0]) for column_meta in self.cursor.description]
337378

338379

339380
class MongoRawQuerySet(RawQuerySet):
@@ -350,5 +391,6 @@ def __init__(
350391
self.model = model
351392
self._db = using
352393
self._hints = hints or {}
353-
self.query = query or MongoRawQuery(sql=raw_query, using=self.db, params=params)
394+
self.query = query or MongoRawQuery(sql=raw_query, using=self.db, model=self.model, params=params)
354395
self._result_cache = None
396+
self.translations = translations or {}

0 commit comments

Comments
 (0)