Ticket #11863: version_1.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
 
@@ -178 +178 @@
 
     def _update(self, values, **kwargs):
         return self.get_query_set()._update(values, **kwargs)
+        
+    def raw(self, *args, **kwargs):
+        return RawQuerySet(model=self.model, *args, **kwargs)
 
 class ManagerDescriptor(object):
     # This class ensures managers aren't accessible via model instances.

django/db/models/query.py

@@ -1075 +1075 @@
         if forced_managed:
             transaction.leave_transaction_management()
 
+class InsuficientFieldsException(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 ()
+        self._kwargs = {}
+        self._annotations = ()
+        # 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 = self._kwargs
+        annotations = self._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 InsuficientFieldsException("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
+        self._kwargs = kwargs
+        self._annotations = annotations
+        return instance
 
 def insert_query(model, values, return_id=False, raw_values=False):
     """

django/db/models/sql/query.py

@@ -29 +29 @@
 except NameError:
     from sets import Set as set     # Python 2.3 fallback
 
-__all__ = ['Query', 'BaseQuery']
+__all__ = ['Query', 'BaseQuery', 'RawQuery', 'InvalidQueryException']
+
+class InvalidQueryException(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, params)
+        
+        
+    def get_columns(self):
+        return [column_meta[0] for column_meta in self.cursor.description]
+        
+    def validate_sql(self, sql):
+        if not sql.lower().startswith('select'):
+            raise InvalidQueryException('Raw SQL are limited to SELECT queries.  Use connection.cursor for any other query types.')
+            
+    def __len__(self):
+        return self.cursor.rowcount
+        
+    def __iter__(self):
+           values = self.cursor.fetchone()
+           while values:
+               yield values
+               values = self.cursor.fetchone()
+               
+    def __str__(self):
+        return self.sql % self.params
+    
 
 class BaseQuery(object):
     """

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()
+    
+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 InsuficientFieldsException
+from django.db.models.sql.query import InvalidQueryException
+
+class RawQueryTests(TestCase):
+    
+    def setUp(self):
+        self.authors = []
+        self.books = []
+        self.coffees = []
+        self.reviewers = []
+        
+        self.authors.append(Author.objects.create(
+            first_name="Joe",
+            last_name="Smith",
+            dob=datetime(year=1950, month=9, day=20),
+        ))
+        self.authors.append(Author.objects.create(
+            first_name="Jill",
+            last_name="Doe",
+            dob=datetime(year=1920, month=4, day=2),
+        ))
+        self.authors.append(Author.objects.create(
+            first_name="Bob",
+            last_name="Smith",
+            dob=datetime(year=1986, month=1, day=25),
+        ))
+        self.authors.append(Author.objects.create(
+            first_name="Bill",
+            last_name="Jones",
+            dob=datetime(year=1932, month=5, day=10),
+        ))
+        
+        self.books.append(Book.objects.create(
+            title = 'The awesome book',
+            author = self.authors[0],
+        ))
+
+        self.books.append(Book.objects.create(
+            title = 'The horrible book',
+            author = self.authors[0],
+        ))
+
+        self.books.append(Book.objects.create(
+            title = 'Another awesome book',
+            author = self.authors[0],
+        ))
+        
+        self.books.append(Book.objects.create(
+            title = 'Some other book',
+            author = self.authors[2],
+        ))
+        
+        self.coffees.append(Coffee.objects.create(brand="dunkin doughnuts"))
+        self.coffees.append(Coffee.objects.create(brand="starbucks"))
+        
+        self.reviewers.append(Reviewer.objects.create())
+        self.reviewers.append(Reviewer.objects.create())
+        self.reviewers[0].reviewed.add(self.books[3])
+        self.reviewers[0].reviewed.add(self.books[1])
+        self.reviewers[0].reviewed.add(self.books[2])
+        
+    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
+        """
+        self.assertEqual(results._annotations, expected_annotations)
+        
+    def testSimpleRawQuery(self):
+        """
+        Basic test of raw query with a simple database query
+        """
+        query = "SELECT * FROM raw_query_author"
+        self.assertSuccessfulRawQuery(Author, query, self.authors)
+
+    def testFkeyRawQuery(self):
+        """
+        Test of a simple raw query against a model containing a foreign key
+        """
+        query = "SELECT * FROM raw_query_book"
+        self.assertSuccessfulRawQuery(Book, query, self.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"
+        self.assertSuccessfulRawQuery(Coffee, query, self.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
+            self.assertSuccessfulRawQuery(Author, query, self.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'),
+        )
+        self.assertSuccessfulRawQuery(Author, query, self.authors, translations=translations)
+        
+    def testParams(self):
+        """
+        Test passing optional query parameters
+        """
+        query = "SELECT * FROM raw_query_author WHERE first_name = %s"
+        params = [self.authors[2].first_name]
+        results = Author.objects.raw(query=query, params=params)
+        self.assertProcessed(results, [self.authors[2]])
+        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"
+        self.assertSuccessfulRawQuery(Reviewer, query, self.reviewers)
+        
+    def testExtraConversions(self):
+        """
+        Test to insure that extra tranlations are ignored.
+        """
+        query = "SELECT * FROM raw_query_author"
+        translations = (('something', 'else'),)
+        self.assertSuccessfulRawQuery(Author, query, self.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 InsuficientFieldsException:
+            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),
+        )
+        self.assertSuccessfulRawQuery(Author, query, self.authors, expected_annotations)
+        
+    def testInvalidQuery(self):
+        query = "UPDATE raw_query_author SET first_name='thing' WHERE first_name='Joe'"
+        self.assertRaises(InvalidQueryException, Author.objects.raw, query)
+ No newline at end of file