Fixed #24031 -- Added CASE expressions to the ORM.

9 files changed, 234 insertions, 28 deletions

diff --git a/django/db/backends/__init__.py b/django/db/backends/__init__.py
index 1c4904a893..fdb4d35d04 100644
--- a/django/db/backends/__init__.py
+++ b/django/db/backends/__init__.py
@@ -821,6 +821,14 @@ class BaseDatabaseOperations(object):
         """
         return "SELECT cache_key FROM %s ORDER BY cache_key LIMIT 1 OFFSET %%s"
 
+    def unification_cast_sql(self, output_field):
+        """
+        Given a field instance, returns the SQL necessary to cast the result of
+        a union to that type. Note that the resulting string should contain a
+        '%s' placeholder for the expression being cast.
+        """
+        return '%s'
+
     def date_extract_sql(self, lookup_type, field_name):
         """
         Given a lookup_type of 'year', 'month' or 'day', returns the SQL that
diff --git a/django/db/backends/postgresql_psycopg2/operations.py b/django/db/backends/postgresql_psycopg2/operations.py
index 9d5b10d01a..31cbe5919f 100644
--- a/django/db/backends/postgresql_psycopg2/operations.py
+++ b/django/db/backends/postgresql_psycopg2/operations.py
@@ -5,6 +5,18 @@ from django.db.backends import BaseDatabaseOperations
 
 
 class DatabaseOperations(BaseDatabaseOperations):
+    def unification_cast_sql(self, output_field):
+        internal_type = output_field.get_internal_type()
+        if internal_type in ("GenericIPAddressField", "IPAddressField", "TimeField", "UUIDField"):
+            # PostgreSQL will resolve a union as type 'text' if input types are
+            # 'unknown'.
+            # http://www.postgresql.org/docs/9.4/static/typeconv-union-case.html
+            # These fields cannot be implicitly cast back in the default
+            # PostgreSQL configuration so we need to explicitly cast them.
+            # We must also remove components of the type within brackets:
+            # varchar(255) -> varchar.
+            return 'CAST(%%s AS %s)' % output_field.db_type(self.connection).split('(')[0]
+        return '%s'
 
     def date_extract_sql(self, lookup_type, field_name):
         # http://www.postgresql.org/docs/current/static/functions-datetime.html#FUNCTIONS-DATETIME-EXTRACT
diff --git a/django/db/models/__init__.py b/django/db/models/__init__.py
index 6fe8fe8aae..9348529625 100644
--- a/django/db/models/__init__.py
+++ b/django/db/models/__init__.py
@@ -4,7 +4,7 @@ import warnings
 
 from django.core.exceptions import ObjectDoesNotExist, ImproperlyConfigured  # NOQA
 from django.db.models.query import Q, QuerySet, Prefetch  # NOQA
-from django.db.models.expressions import ExpressionNode, F, Value, Func  # NOQA
+from django.db.models.expressions import ExpressionNode, F, Value, Func, Case, When  # NOQA
 from django.db.models.manager import Manager  # NOQA
 from django.db.models.base import Model  # NOQA
 from django.db.models.aggregates import *  # NOQA
diff --git a/django/db/models/aggregates.py b/django/db/models/aggregates.py
index 4b57890cf8..06220123ca 100644
--- a/django/db/models/aggregates.py
+++ b/django/db/models/aggregates.py
@@ -14,8 +14,9 @@ class Aggregate(Func):
     contains_aggregate = True
     name = None
 
-    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False):
+    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False):
         assert len(self.source_expressions) == 1
+        # Aggregates are not allowed in UPDATE queries, so ignore for_save
         c = super(Aggregate, self).resolve_expression(query, allow_joins, reuse, summarize)
         if c.source_expressions[0].contains_aggregate and not summarize:
             name = self.source_expressions[0].name
@@ -101,7 +102,6 @@ class Count(Aggregate):
     def __init__(self, expression, distinct=False, **extra):
         if expression == '*':
             expression = Value(expression)
-            expression._output_field = IntegerField()
         super(Count, self).__init__(
             expression, distinct='DISTINCT ' if distinct else '', output_field=IntegerField(), **extra)
 
diff --git a/django/db/models/expressions.py b/django/db/models/expressions.py
index 0f55545961..97a2a9071d 100644
--- a/django/db/models/expressions.py
+++ b/django/db/models/expressions.py
@@ -6,7 +6,7 @@ from django.core.exceptions import FieldError
 from django.db.backends import utils as backend_utils
 from django.db.models import fields
 from django.db.models.constants import LOOKUP_SEP
-from django.db.models.query_utils import refs_aggregate
+from django.db.models.query_utils import refs_aggregate, Q
 from django.utils import timezone
 from django.utils.functional import cached_property
 
@@ -173,7 +173,7 @@ class BaseExpression(object):
                 return True
         return False
 
-    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False):
+    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False):
         """
         Provides the chance to do any preprocessing or validation before being
         added to the query.
@@ -380,11 +380,11 @@ class Expression(ExpressionNode):
         sql = connection.ops.combine_expression(self.connector, expressions)
         return expression_wrapper % sql, expression_params
 
-    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False):
+    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False):
         c = self.copy()
         c.is_summary = summarize
-        c.lhs = c.lhs.resolve_expression(query, allow_joins, reuse, summarize)
-        c.rhs = c.rhs.resolve_expression(query, allow_joins, reuse, summarize)
+        c.lhs = c.lhs.resolve_expression(query, allow_joins, reuse, summarize, for_save)
+        c.rhs = c.rhs.resolve_expression(query, allow_joins, reuse, summarize, for_save)
         return c
 
 
@@ -426,7 +426,7 @@ class F(CombinableMixin):
         """
         self.name = name
 
-    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False):
+    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False):
         return query.resolve_ref(self.name, allow_joins, reuse, summarize)
 
     def refs_aggregate(self, existing_aggregates):
@@ -465,11 +465,11 @@ class Func(ExpressionNode):
             for arg in expressions
         ]
 
-    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False):
+    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False):
         c = self.copy()
         c.is_summary = summarize
         for pos, arg in enumerate(c.source_expressions):
-            c.source_expressions[pos] = arg.resolve_expression(query, allow_joins, reuse, summarize)
+            c.source_expressions[pos] = arg.resolve_expression(query, allow_joins, reuse, summarize, for_save)
         return c
 
     def as_sql(self, compiler, connection, function=None, template=None):
@@ -511,12 +511,24 @@ class Value(ExpressionNode):
         self.value = value
 
     def as_sql(self, compiler, connection):
-        if self.value is None:
+        val = self.value
+        # check _output_field to avoid triggering an exception
+        if self._output_field is not None:
+            if self.for_save:
+                val = self.output_field.get_db_prep_save(val, connection=connection)
+            else:
+                val = self.output_field.get_db_prep_value(val, connection=connection)
+        if val is None:
             # cx_Oracle does not always convert None to the appropriate
             # NULL type (like in case expressions using numbers), so we
             # use a literal SQL NULL
             return 'NULL', []
-        return '%s', [self.value]
+        return '%s', [val]
+
+    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False):
+        c = super(Value, self).resolve_expression(query, allow_joins, reuse, summarize, for_save)
+        c.for_save = for_save
+        return c
 
     def get_group_by_cols(self):
         return []
@@ -599,6 +611,130 @@ class Ref(ExpressionNode):
         return [self]
 
 
+class When(ExpressionNode):
+    template = 'WHEN %(condition)s THEN %(result)s'
+
+    def __init__(self, condition=None, then=Value(None), **lookups):
+        if lookups and condition is None:
+            condition, lookups = Q(**lookups), None
+        if condition is None or not isinstance(condition, Q) or lookups:
+            raise TypeError("__init__() takes either a Q object or lookups as keyword arguments")
+        super(When, self).__init__(output_field=None)
+        self.condition = condition
+        self.result = self._parse_expression(then)
+
+    def __str__(self):
+        return "WHEN %r THEN %r" % (self.condition, self.result)
+
+    def __repr__(self):
+        return "<%s: %s>" % (self.__class__.__name__, self)
+
+    def get_source_expressions(self):
+        return [self.condition, self.result]
+
+    def set_source_expressions(self, exprs):
+        self.condition, self.result = exprs
+
+    def get_source_fields(self):
+        # We're only interested in the fields of the result expressions.
+        return [self.result._output_field_or_none]
+
+    def _parse_expression(self, expression):
+        return expression if hasattr(expression, 'resolve_expression') else F(expression)
+
+    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False):
+        c = self.copy()
+        c.is_summary = summarize
+        c.condition = c.condition.resolve_expression(query, allow_joins, reuse, summarize, False)
+        c.result = c.result.resolve_expression(query, allow_joins, reuse, summarize, for_save)
+        return c
+
+    def as_sql(self, compiler, connection, template=None):
+        template_params = {}
+        sql_params = []
+        condition_sql, condition_params = compiler.compile(self.condition)
+        template_params['condition'] = condition_sql
+        sql_params.extend(condition_params)
+        result_sql, result_params = compiler.compile(self.result)
+        template_params['result'] = result_sql
+        sql_params.extend(result_params)
+        template = template or self.template
+        return template % template_params, sql_params
+
+    def get_group_by_cols(self):
+        # This is not a complete expression and cannot be used in GROUP BY.
+        cols = []
+        for source in self.get_source_expressions():
+            cols.extend(source.get_group_by_cols())
+        return cols
+
+
+class Case(ExpressionNode):
+    """
+    An SQL searched CASE expression:
+
+        CASE
+            WHEN n > 0
+                THEN 'positive'
+            WHEN n < 0
+                THEN 'negative'
+            ELSE 'zero'
+        END
+    """
+    template = 'CASE %(cases)s ELSE %(default)s END'
+    case_joiner = ' '
+
+    def __init__(self, *cases, **extra):
+        if not all(isinstance(case, When) for case in cases):
+            raise TypeError("Positional arguments must all be When objects.")
+        default = extra.pop('default', Value(None))
+        output_field = extra.pop('output_field', None)
+        super(Case, self).__init__(output_field)
+        self.cases = list(cases)
+        self.default = default if hasattr(default, 'resolve_expression') else F(default)
+
+    def __str__(self):
+        return "CASE %s, ELSE %r" % (', '.join(str(c) for c in self.cases), self.default)
+
+    def __repr__(self):
+        return "<%s: %s>" % (self.__class__.__name__, self)
+
+    def get_source_expressions(self):
+        return self.cases + [self.default]
+
+    def set_source_expressions(self, exprs):
+        self.cases = exprs[:-1]
+        self.default = exprs[-1]
+
+    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False):
+        c = self.copy()
+        c.is_summary = summarize
+        for pos, case in enumerate(c.cases):
+            c.cases[pos] = case.resolve_expression(query, allow_joins, reuse, summarize, for_save)
+        c.default = c.default.resolve_expression(query, allow_joins, reuse, summarize, for_save)
+        return c
+
+    def as_sql(self, compiler, connection, template=None, extra=None):
+        if not self.cases:
+            return compiler.compile(self.default)
+        template_params = dict(extra) if extra else {}
+        case_parts = []
+        sql_params = []
+        for case in self.cases:
+            case_sql, case_params = compiler.compile(case)
+            case_parts.append(case_sql)
+            sql_params.extend(case_params)
+        template_params['cases'] = self.case_joiner.join(case_parts)
+        default_sql, default_params = compiler.compile(self.default)
+        template_params['default'] = default_sql
+        sql_params.extend(default_params)
+        template = template or self.template
+        sql = template % template_params
+        if self._output_field_or_none is not None:
+            sql = connection.ops.unification_cast_sql(self.output_field) % sql
+        return sql, sql_params
+
+
 class Date(ExpressionNode):
     """
     Add a date selection column.
@@ -615,7 +751,7 @@ class Date(ExpressionNode):
     def set_source_expressions(self, exprs):
         self.col, = exprs
 
-    def resolve_expression(self, query, allow_joins, reuse, summarize):
+    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False):
         copy = self.copy()
         copy.col = query.resolve_ref(self.lookup, allow_joins, reuse, summarize)
         field = copy.col.output_field
@@ -664,7 +800,7 @@ class DateTime(ExpressionNode):
     def set_source_expressions(self, exprs):
         self.col, = exprs
 
-    def resolve_expression(self, query, allow_joins, reuse, summarize):
+    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False):
         copy = self.copy()
         copy.col = query.resolve_ref(self.lookup, allow_joins, reuse, summarize)
         field = copy.col.output_field
diff --git a/django/db/models/query_utils.py b/django/db/models/query_utils.py
index 3eae98ee65..9e4ddfec58 100644
--- a/django/db/models/query_utils.py
+++ b/django/db/models/query_utils.py
@@ -86,6 +86,27 @@ class Q(tree.Node):
                 clone.children.append(child)
         return clone
 
+    def resolve_expression(self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False):
+        clause, _ = query._add_q(self, reuse, allow_joins=allow_joins)
+        return clause
+
+    def refs_aggregate(self, existing_aggregates):
+        def _refs_aggregate(obj, existing_aggregates):
+            if not isinstance(obj, tree.Node):
+                aggregate, aggregate_lookups = refs_aggregate(obj[0].split(LOOKUP_SEP), existing_aggregates)
+                if not aggregate and hasattr(obj[1], 'refs_aggregate'):
+                    return obj[1].refs_aggregate(existing_aggregates)
+                return aggregate, aggregate_lookups
+            for c in obj.children:
+                aggregate, aggregate_lookups = _refs_aggregate(c, existing_aggregates)
+                if aggregate:
+                    return aggregate, aggregate_lookups
+            return False, ()
+
+        if not existing_aggregates:
+            return False
+        return _refs_aggregate(self, existing_aggregates)
+
 
 class DeferredAttribute(object):
     """
diff --git a/django/db/models/sql/compiler.py b/django/db/models/sql/compiler.py
index b091571835..ebda3be96f 100644
--- a/django/db/models/sql/compiler.py
+++ b/django/db/models/sql/compiler.py
@@ -998,7 +998,9 @@ class SQLUpdateCompiler(SQLCompiler):
         values, update_params = [], []
         for field, model, val in self.query.values:
             if hasattr(val, 'resolve_expression'):
-                val = val.resolve_expression(self.query, allow_joins=False)
+                val = val.resolve_expression(self.query, allow_joins=False, for_save=True)
+                if val.contains_aggregate:
+                    raise FieldError("Aggregate functions are not allowed in this query")
             elif hasattr(val, 'prepare_database_save'):
                 if field.rel:
                     val = val.prepare_database_save(field)
diff --git a/django/db/models/sql/query.py b/django/db/models/sql/query.py
index 5d4538b2d9..c5e7eab28c 100644
--- a/django/db/models/sql/query.py
+++ b/django/db/models/sql/query.py
@@ -961,7 +961,7 @@ class Query(object):
         self.append_annotation_mask([alias])
         self.annotations[alias] = annotation
 
-    def prepare_lookup_value(self, value, lookups, can_reuse):
+    def prepare_lookup_value(self, value, lookups, can_reuse, allow_joins=True):
         # Default lookup if none given is exact.
         used_joins = []
         if len(lookups) == 0:
@@ -980,7 +980,7 @@ class Query(object):
             value = value()
         elif hasattr(value, 'resolve_expression'):
             pre_joins = self.alias_refcount.copy()
-            value = value.resolve_expression(self, reuse=can_reuse)
+            value = value.resolve_expression(self, reuse=can_reuse, allow_joins=allow_joins)
             used_joins = [k for k, v in self.alias_refcount.items() if v > pre_joins.get(k, 0)]
         # Subqueries need to use a different set of aliases than the
         # outer query. Call bump_prefix to change aliases of the inner
@@ -1095,7 +1095,7 @@ class Query(object):
                 (name, lhs.output_field.__class__.__name__))
 
     def build_filter(self, filter_expr, branch_negated=False, current_negated=False,
-                     can_reuse=None, connector=AND):
+                     can_reuse=None, connector=AND, allow_joins=True):
         """
         Builds a WhereNode for a single filter clause, but doesn't add it
         to this Query. Query.add_q() will then add this filter to the where
@@ -1125,10 +1125,12 @@ class Query(object):
         if not arg:
             raise FieldError("Cannot parse keyword query %r" % arg)
         lookups, parts, reffed_aggregate = self.solve_lookup_type(arg)
+        if not allow_joins and len(parts) > 1:
+            raise FieldError("Joined field references are not permitted in this query")
 
         # Work out the lookup type and remove it from the end of 'parts',
         # if necessary.
-        value, lookups, used_joins = self.prepare_lookup_value(value, lookups, can_reuse)
+        value, lookups, used_joins = self.prepare_lookup_value(value, lookups, can_reuse, allow_joins)
 
         clause = self.where_class()
         if reffed_aggregate:
@@ -1225,11 +1227,11 @@ class Query(object):
         """
         if not self._annotations:
             return False
-        if not isinstance(obj, Node):
-            return (refs_aggregate(obj[0].split(LOOKUP_SEP), self.annotations)[0]
-                    or (hasattr(obj[1], 'refs_aggregate')
-                        and obj[1].refs_aggregate(self.annotations)[0]))
-        return any(self.need_having(c) for c in obj.children)
+        if hasattr(obj, 'refs_aggregate'):
+            return obj.refs_aggregate(self.annotations)[0]
+        return (refs_aggregate(obj[0].split(LOOKUP_SEP), self.annotations)[0]
+                or (hasattr(obj[1], 'refs_aggregate')
+                    and obj[1].refs_aggregate(self.annotations)[0]))
 
     def split_having_parts(self, q_object, negated=False):
         """
@@ -1287,7 +1289,7 @@ class Query(object):
         self.demote_joins(existing_inner)
 
     def _add_q(self, q_object, used_aliases, branch_negated=False,
-               current_negated=False):
+               current_negated=False, allow_joins=True):
         """
         Adds a Q-object to the current filter.
         """
@@ -1301,12 +1303,12 @@ class Query(object):
             if isinstance(child, Node):
                 child_clause, needed_inner = self._add_q(
                     child, used_aliases, branch_negated,
-                    current_negated)
+                    current_negated, allow_joins)
                 joinpromoter.add_votes(needed_inner)
             else:
                 child_clause, needed_inner = self.build_filter(
                     child, can_reuse=used_aliases, branch_negated=branch_negated,
-                    current_negated=current_negated, connector=connector)
+                    current_negated=current_negated, connector=connector, allow_joins=allow_joins)
                 joinpromoter.add_votes(needed_inner)
             target_clause.add(child_clause, connector)
         needed_inner = joinpromoter.update_join_types(self)
diff --git a/django/db/models/sql/where.py b/django/db/models/sql/where.py
index 10445555e0..cbb709dfc8 100644
--- a/django/db/models/sql/where.py
+++ b/django/db/models/sql/where.py
@@ -11,6 +11,7 @@ from django.conf import settings
 from django.db.models.fields import DateTimeField, Field
 from django.db.models.sql.datastructures import EmptyResultSet, Empty
 from django.utils.deprecation import RemovedInDjango19Warning
+from django.utils.functional import cached_property
 from django.utils.six.moves import range
 from django.utils import timezone
 from django.utils import tree
@@ -309,6 +310,30 @@ class WhereNode(tree.Node):
                 clone.children.append(child)
         return clone
 
+    def relabeled_clone(self, change_map):
+        clone = self.clone()
+        clone.relabel_aliases(change_map)
+        return clone
+
+    @cached_property
+    def contains_aggregate(self):
+        def _contains_aggregate(obj):
+            if not isinstance(obj, tree.Node):
+                return getattr(obj.lhs, 'contains_aggregate', False) or getattr(obj.rhs, 'contains_aggregate', False)
+            return any(_contains_aggregate(c) for c in obj.children)
+
+        return _contains_aggregate(self)
+
+    def refs_field(self, aggregate_types, field_types):
+        def _refs_field(obj, aggregate_types, field_types):
+            if not isinstance(obj, tree.Node):
+                if hasattr(obj.rhs, 'refs_field'):
+                    return obj.rhs.refs_field(aggregate_types, field_types)
+                return False
+            return any(_refs_field(c, aggregate_types, field_types) for c in obj.children)
+
+        return _refs_field(self, aggregate_types, field_types)
+
 
 class EmptyWhere(WhereNode):
     def add(self, data, connector):