Ticket #11863: version_2.diff

django/db/models/manager.py

@@ -1 +1 @@
 import copy
-from django.db.models.query import QuerySet, EmptyQuerySet, insert_query
+from django.db.models.query import QuerySet, EmptyQuerySet, insert_query, RawQuerySet
 from django.db.models import signals
 from django.db.models.fields import FieldDoesNotExist
 
@@ -181 +181 @@
     def _update(self, values, **kwargs):
         return self.get_query_set()._update(values, **kwargs)
 
+    def raw(self, query, params=None, *args, **kwargs):
+        return RawQuerySet(model=self.model, query=query, params=params, *args, **kwargs)
+
 class ManagerDescriptor(object):
     # This class ensures managers aren't accessible via model instances.
     # For example, Poll.objects works, but poll_obj.objects raises AttributeError.

django/db/models/query.py

@@ -257 +257 @@
         Returns a dictionary containing the calculations (aggregation)
         over the current queryset
 
-        If args is present the expression is passed as a kwarg ussing
+        If args is present the expression is passed as a kwarg using
         the Aggregate object's default alias.
         """
         for arg in args:
@@ -1075 +1075 @@
         if forced_managed:
             transaction.leave_transaction_management()
 
+class InsuficientFields(Exception):
+    """
+    The query passed to raw doesn't include all of the fields needed to build
+    a model instance.
+    """
+    pass
+
+class RawQuerySet(object):
+    """
+    Provides an iterator which converts the results of raw database
+    queries into annotated model instances.
+    """
+    def __init__(self, query, model=None, query_obj=None, params=None, translations=None):
+        self.model = model
+        self.raw_results = query_obj or sql.RawQuery(sql=query, connection=connection, params=params)
+        self.query = query
+        self.params = params or ()
+        # Figure out the column names
+        self.columns = self.raw_results.get_columns()
+        # Adjust any column names which don't match field names
+        if translations is not None:
+            for translation in translations:
+                try:
+                    index = self.columns.index(translation[0])
+                    self.columns[index] = translation[1]
+                except ValueError:
+                    # Ignore transations for non-existant column names
+                    pass
+
+        # Build a list of column names known by the model
+        self.model_fields = {}
+        for field in model._meta.fields:
+            name, column = field.get_attname_column()
+            self.model_fields[column] = name
+
+    def __iter__(self):
+        for row in self.raw_results:
+            yield self.transform_results(row)
+
+    def __len__(self):
+        return len(self.raw_results)
+
+    def __repr__(self):
+        return self.query % self.params
+
+    def transform_results(self, values):
+        kwargs = {}
+        annotations = ()
+
+        # Associate fields to values
+        for pos, value in enumerate(values):
+            column = self.columns[pos]
+
+            # Separate properties from annotations
+            if column in self.model_fields.keys():
+                kwargs[self.model_fields[column]] = value
+            else:
+                annotations += (column, value),
+
+        if len(kwargs) < len(self.model_fields):
+            missing = [column for column, field in self.model_fields.items() if field not in kwargs.keys()]
+            raise InsuficientFields("They query passed doesn't contain all of the needed fields.  The missing fields are: %s" % ', '.join(missing))
+
+        # Construct model instance
+        instance = self.model(**kwargs)
+
+        # Apply annotations
+        for field, value in annotations:
+            setattr(instance, field, value)
+
+        # make the kwargs and annotation metadata available to the unit tests
+        return instance
 
 def insert_query(model, values, return_id=False, raw_values=False):
     """

django/db/models/sql/query.py

@@ -23 +23 @@
 from datastructures import EmptyResultSet, Empty, MultiJoin
 from constants import *
 
-__all__ = ['Query', 'BaseQuery']
+__all__ = ['Query', 'BaseQuery', 'RawQuery', 'InvalidQuery']
+
+class InvalidQuery(Exception):
+    """
+    The query passed to raw isn't a safe query to use with raw.
+    """
+    pass
+
+class RawQuery(object):
+    """
+    A single raw SQL query
+    """
+
+    def __init__(self, sql, connection, params=None):
+        self.validate_sql(sql)
+        self.params = params or ()
+        self.sql = sql
+        self.connection = connection
+        self.cursor = connection.cursor()
+        self.cursor.execute(sql, self.params)
+        self._cache = None
+        self._count = None
+
+
+    def get_columns(self):
+        return [column_meta[0] for column_meta in self.cursor.description]
+
+    def validate_sql(self, sql):
+        if not sql.lower().strip().startswith('select'):
+            raise InvalidQuery('Raw SQL are limited to SELECT queries.  Use connection.cursor for any other query types.')
+
+    def __len__(self):
+        if self._count is None:
+            self._count = self.cursor.rowcount
+            # Handle sqlite's quirky rowcount behavior
+            if self._count == -1:
+                self._populate_cache()
+                self._count = len(self._cache)
+
+        return self._count
+
+    def __iter__(self):
+        self._populate_cache()
+        for value in self._cache:
+            yield value
+
+
+    def __str__(self):
+        return self.sql % self.params
+
+    def _populate_cache(self):
+        if self._cache is None:
+            self._cache = self.cursor.fetchall()
 
 class BaseQuery(object):
     """

new file tests/modeltests/raw_query/fixtures/initial_data.json

@@ +1 @@
+[
+    {
+        "pk": 1, 
+        "model": "raw_query.author", 
+        "fields": {
+            "dob": "1950-09-20", 
+            "first_name": "Joe", 
+            "last_name": "Smith"
+        }
+    }, 
+    {
+        "pk": 2, 
+        "model": "raw_query.author", 
+        "fields": {
+            "dob": "1920-04-02", 
+            "first_name": "Jill", 
+            "last_name": "Doe"
+        }
+    }, 
+    {
+        "pk": 3, 
+        "model": "raw_query.author", 
+        "fields": {
+            "dob": "1986-01-25", 
+            "first_name": "Bob", 
+            "last_name": "Smith"
+        }
+    }, 
+    {
+        "pk": 4, 
+        "model": "raw_query.author", 
+        "fields": {
+            "dob": "1932-05-10", 
+            "first_name": "Bill", 
+            "last_name": "Jones"
+        }
+    }, 
+    {
+        "pk": 1, 
+        "model": "raw_query.book", 
+        "fields": {
+            "author": 1, 
+            "title": "The awesome book"
+        }
+    }, 
+    {
+        "pk": 2, 
+        "model": "raw_query.book", 
+        "fields": {
+            "author": 1, 
+            "title": "The horrible book"
+        }
+    }, 
+    {
+        "pk": 3, 
+        "model": "raw_query.book", 
+        "fields": {
+            "author": 1, 
+            "title": "Another awesome book"
+        }
+    }, 
+    {
+        "pk": 4, 
+        "model": "raw_query.book", 
+        "fields": {
+            "author": 3, 
+            "title": "Some other book"
+        }
+    }, 
+    {
+        "pk": 1, 
+        "model": "raw_query.coffee", 
+        "fields": {
+            "brand": "dunkin doughnuts"
+        }
+    }, 
+    {
+        "pk": 2, 
+        "model": "raw_query.coffee", 
+        "fields": {
+            "brand": "starbucks"
+        }
+    }, 
+    {
+        "pk": 1, 
+        "model": "raw_query.reviewer", 
+        "fields": {
+            "reviewed": [
+                2, 
+                3, 
+                4
+            ]
+        }
+    }, 
+    {
+        "pk": 2, 
+        "model": "raw_query.reviewer", 
+        "fields": {
+            "reviewed": []
+        }
+    }
+]

new file tests/modeltests/raw_query/models.py

@@ +1 @@
+from django.db import models
+
+class Author(models.Model):
+    first_name = models.CharField(max_length=255)
+    last_name = models.CharField(max_length=255)
+    dob = models.DateField()
+
+    def __init__(self, *args, **kwargs):
+        super(Author, self).__init__(*args, **kwargs)
+
+        # Perform a check for missing params to protect against real fields
+        # being added as annotations.
+        if len(args) == 0:
+            expected_keywords = ['dob', 'first_name', 'last_name']
+            for keyword in expected_keywords:
+                if keyword not in kwargs.keys():
+                    raise Exception('All fields are required.  %s is missing' %
+                                      keyword)
+
+
+
+class Book(models.Model):
+    title = models.CharField(max_length=255)
+    author = models.ForeignKey(Author)
+
+class Coffee(models.Model):
+    brand = models.CharField(max_length=255, db_column="name")
+
+class Reviewer(models.Model):
+    reviewed = models.ManyToManyField(Book)
+ No newline at end of file

new file tests/modeltests/raw_query/tests.py

@@ +1 @@
+from django.test import TestCase
+from datetime import datetime
+from models import Author, Book, Coffee, Reviewer
+
+from django.db.models.query import InsuficientFields
+from django.db.models.sql.query import InvalidQuery
+
+class RawQueryTests(TestCase):
+
+    def assertSuccessfulRawQuery(self, model, query, expected_results, \
+            expected_annotations=(), params=[], translations=None):
+        """
+        Execute the passed query against the passed model and check the output
+        """
+        results = model.objects.raw(query=query, params=params, \
+                            translations=translations)
+        self.assertProcessed(results, expected_results, expected_annotations)
+        self.assertAnnotations(results, expected_annotations)
+
+    def assertProcessed(self, results, orig, expected_annotations=()):
+        """
+        Compare the results of a raw query against expected results
+        """
+        self.assertEqual(len(results), len(orig))
+        for index, item in enumerate(results):
+            orig_item = orig[index]
+            for annotation in expected_annotations:
+                setattr(orig_item, *annotation)
+
+            self.assertEqual(item.id, orig_item.id)
+
+    def assertNoAnnotations(self, results):
+        """
+        Check that the results of a raw query contain no annotations
+        """
+        self.assertAnnotations(results, ())
+
+    def assertAnnotations(self, results, expected_annotations):
+        """
+        Check that the passed raw query results contain the expected
+        annotations
+        """
+        if expected_annotations:
+            for index, result in enumerate(results):
+                annotation, value = expected_annotations[index]
+                self.assertTrue(hasattr(result, annotation))
+                self.assertEqual(getattr(result, annotation), value)
+
+    def testSimpleRawQuery(self):
+        """
+        Basic test of raw query with a simple database query
+        """
+        query = "SELECT * FROM raw_query_author"
+        authors = Author.objects.all()
+        self.assertSuccessfulRawQuery(Author, query, authors)
+
+    def testFkeyRawQuery(self):
+        """
+        Test of a simple raw query against a model containing a foreign key
+        """
+        query = "SELECT * FROM raw_query_book"
+        books = Book.objects.all()
+        self.assertSuccessfulRawQuery(Book, query, books)
+
+    def testDBColumnHandler(self):
+        """
+        Test of a simple raw query against a model containing a field with
+        db_column defined.
+        """
+        query = "SELECT * FROM raw_query_coffee"
+        coffees = Coffee.objects.all()
+        self.assertSuccessfulRawQuery(Coffee, query, coffees)
+
+    def testOrderHandler(self):
+        """
+        Test of raw raw query's tolerance for columns being returned in any
+        order
+        """
+        selects = (
+            ('dob, last_name, first_name, id'),
+            ('last_name, dob, first_name, id'),
+            ('first_name, last_name, dob, id'),
+        )
+
+        for select in selects:
+            query = "SELECT %s FROM raw_query_author" % select
+            authors = Author.objects.all()
+            self.assertSuccessfulRawQuery(Author, query, authors)
+
+    def testTranslations(self):
+        """
+        Test of raw query's optional ability to translate unexpected result
+        column names to specific model fields
+        """
+        query = "SELECT first_name AS first, last_name AS last, dob, id FROM raw_query_author"
+        translations = (
+            ('first', 'first_name'),
+            ('last', 'last_name'),
+        )
+        authors = Author.objects.all()
+        self.assertSuccessfulRawQuery(Author, query, authors, translations=translations)
+
+    def testParams(self):
+        """
+        Test passing optional query parameters
+        """
+        query = "SELECT * FROM raw_query_author WHERE first_name = %s"
+        author = Author.objects.all()[2]
+        params = [author.first_name]
+        results = Author.objects.raw(query=query, params=params)
+        self.assertProcessed(results, [author])
+        self.assertNoAnnotations(results)
+        self.assertEqual(len(results), 1)
+
+    def testManyToMany(self):
+        """
+        Test of a simple raw query against a model containing a m2m field
+        """
+        query = "SELECT * FROM raw_query_reviewer"
+        reviewers = Reviewer.objects.all()
+        self.assertSuccessfulRawQuery(Reviewer, query, reviewers)
+
+    def testExtraConversions(self):
+        """
+        Test to insure that extra translations are ignored.
+        """
+        query = "SELECT * FROM raw_query_author"
+        translations = (('something', 'else'),)
+        authors = Author.objects.all()
+        self.assertSuccessfulRawQuery(Author, query, authors, translations=translations)
+
+    def testInsufficientColumns(self):
+        query = "SELECT first_name, dob FROM raw_query_author"
+        raised = False
+        try:
+            results = Author.objects.raw(query)
+            results_list = list(results)
+        except InsuficientFields:
+            raised = True
+
+        self.assertTrue(raised)
+
+    def testAnnotations(self):
+        query = "SELECT a.*, count(b.id) as book_count FROM raw_query_author a LEFT JOIN raw_query_book b ON a.id = b.author_id GROUP BY a.id, a.first_name, a.last_name, a.dob ORDER BY a.id"
+        expected_annotations = (
+            ('book_count', 3),
+            ('book_count', 0),
+            ('book_count', 1),
+            ('book_count', 0),
+        )
+        authors = Author.objects.all()
+        self.assertSuccessfulRawQuery(Author, query, authors, expected_annotations)
+
+    def testInvalidQuery(self):
+        query = "UPDATE raw_query_author SET first_name='thing' WHERE first_name='Joe'"
+        self.assertRaises(InvalidQuery, Author.objects.raw, query)
+
+    def testWhiteSpaceQuery(self):
+        query = "    SELECT * FROM raw_query_author"
+        authors = Author.objects.all()
+        self.assertSuccessfulRawQuery(Author, query, authors)
+
+    def testMultipleIterations(self):
+        query = "SELECT * FROM raw_query_author"
+        normal_authors = Author.objects.all()
+        raw_authors = Author.objects.raw(query)
+
+        # First Iteration
+        first_iterations = 0
+        for index, raw_author in enumerate(raw_authors):
+            self.assertEqual(normal_authors[index], raw_author)
+            first_iterations += 1
+
+        # Second Iteration
+        second_iterations = 0
+        for index, raw_author in enumerate(raw_authors):
+            self.assertEqual(normal_authors[index], raw_author)
+            second_iterations += 1
+
+        self.assertEqual(first_iterations, second_iterations)
+ No newline at end of file