Ticket #3566: aggregate.2.diff

django/db/models/sql/query.py

@@ -56 +56 @@
         self.start_meta = None
 
         # SQL-related attributes
+        self.aggregates = []
         self.select = []
         self.tables = []    # Aliases in the order they are created.
         self.where = where()
@@ -141 +142 @@
         obj.standard_ordering = self.standard_ordering
         obj.start_meta = self.start_meta
         obj.select = self.select[:]
+        obj.aggregates = self.aggregates[:]
         obj.tables = self.tables[:]
         obj.where = deepcopy(self.where)
         obj.where_class = self.where_class
@@ -174 +176 @@
                     row = self.resolve_columns(row, fields)
                 yield row
 
+    def get_aggregation(self):
+        for field in self.select:
+            self.group_by.append(field)
+        self.select.extend(self.aggregates)
+        self.aggregates = []
+        #print self.as_sql()
+        #print 'after', self.select
+
+        get_name = lambda x : isinstance(x, tuple) and x[1] or x.aliased_name
+
+        print 'final query', self.as_sql() 
+
+        if self.group_by:
+            data = self.execute_sql(MULTI)
+            result = []
+            for rs in data.next():
+                result.append(dict(zip([get_name(i) for i in self.select], rs)))
+        else:
+            data = self.execute_sql(SINGLE)
+            result = dict(zip([get_name(i) for i in self.select], data))
+
+        self.select = []
+        return result
+
     def get_count(self):
         """
         Performs a COUNT() query using the current filter constraints.
@@ -811 +837 @@
             self.fill_related_selections(f.rel.to._meta, alias, cur_depth + 1,
                     used, next, restricted)
 
+    def annotate(self, aggregate_expr, aliased_name, model):
+        field_list = aggregate_expr.split(LOOKUP_SEP)
+        opts = model._meta
+
+        aggregate_func = field_list.pop()
+        
+        if len(field_list) > 1:
+            field, target, opts, join_list, last = self.setup_joins(
+                field_list, opts, self.get_initial_alias(), False)
+            final = len(join_list)
+            penultimate = last.pop()
+            if penultimate == final:
+                penultimate = last.pop()
+            if len(join_list) > 1:
+                extra = join_list[penultimate:]
+                final = penultimate
+                col = self.alias_map[extra[0]][LHS_JOIN_COL]
+            else:
+                col = target.column
+                
+            field_name = field_list.pop()
+            alias = join_list[-1]
+            alias = extra[final]
+        else:
+            field_name = field_list[0]
+            alias = opts.db_table
+          
+
+    def add_aggregate(self, aggregate_expr, aliased_name, model):
+        """
+        Adds a single aggregate expression to the Query
+        """
+        
+        field_list = aggregate_expr.split(LOOKUP_SEP)
+        opts = model._meta
+
+        aggregate_func = field_list.pop()
+        
+        if len(field_list) > 1:
+            field, target, opts, join_list, last = self.setup_joins(
+                field_list, opts, self.get_initial_alias(), False)
+            final = len(join_list)
+            penultimate = last.pop()
+            if penultimate == final:
+                penultimate = last.pop()
+            if len(join_list) > 1:
+                extra = join_list[penultimate:]
+                final = penultimate
+                col = self.alias_map[extra[0]][LHS_JOIN_COL]
+            else:
+                col = target.column
+                
+            field_name = field_list.pop()
+            alias = join_list[-1]
+            alias = extra[final]
+        else:
+            field_name = field_list[0]
+            alias = opts.db_table
+
+        class AggregateNode:
+            def __init__(self, field_name, aggregate_func, aliased_name, alias):
+                self.field_name = field_name
+                self.aggregate_func = aggregate_func
+                self.aliased_name = aliased_name
+                self.alias = alias
+                
+            def as_sql(self, quote_func=None):
+                if not quote_func:
+                    quote_func = lambda x: x
+                return '%s(%s.%s)' % (self.aggregate_func.upper(),
+                                      quote_func(self.alias),
+                                      quote_func(self.field_name))
+
+        self.aggregates.append(AggregateNode(field_name, aggregate_func, aliased_name, alias))
+        
     def add_filter(self, filter_expr, connector=AND, negate=False, trim=False,
             single_filter=False):
         """
@@ -829 +930 @@
         if not parts:
             raise FieldError("Cannot parse keyword query %r" % arg)
 
+        # if arg in (x.aliased_name for x in self.aggregates):
+        #     self.having.append(arg)
+        #     return
+
         # Work out the lookup type and remove it from 'parts', if necessary.
         if len(parts) == 1 or parts[-1] not in self.query_terms:
             lookup_type = 'exact'

django/db/models/query.py

@@ -158 +158 @@
                 setattr(obj, k, row[i])
             yield obj
 
+    def aggregate(self, *args, **kwargs):
+        """
+        Returns the aggregation over the current model as a
+        dictionary.
+
+        When applied to a ValuesQuerySet the results are GROUP BY-ed
+        by the fields specified in the values queryset.
+
+        The kwargs are parsed as expression='alias'.
+
+        If args is present the expression is passed as a kwarg with
+        itself as an alias.
+        """
+        #Bug (or is it?): when doing both an aggregation on a related
+        #field and one on a 'local' field the local one goes
+        #wrong. something similar to: SELECT SUM(a.f1) FROM a INNER JOIN b;
+        #the value gets aggregated more than one time.
+
+        if args:
+            newargs = {}
+            for arg in args:
+                newargs[arg] = arg
+            kwargs.update(newargs)
+            
+        for (aggregate_expr, alias) in kwargs.items():
+            self.query.add_aggregate(aggregate_expr, alias, self.model)
+        return self.query.get_aggregation()
+
     def count(self):
         """
         Performs a SELECT COUNT() and returns the number of records as an
@@ -326 +354 @@
         """
         return self._clone(klass=EmptyQuerySet)
 
+    def annotate(self, *args, **kwargs):
+        # Fix: Values is not working propperly
+        # Suffers from the same bug as aggrgate
+        # To-Do: HAVING
+
+        if args:
+            newargs = {}
+            for arg in args:
+                newargs[arg] = arg
+            kwargs.update(newargs)
+
+        opts = self.model._meta
+        fields = []
+        
+        if isinstance(self, ValuesQuerySet):
+            obj = self._clone()
+        else:
+            fields.extend([f.name for f in opts.fields])
+            obj = self._clone(klass=ValuesQuerySet, setup=True, _fields=fields)
+            
+
+        for (aggregate_expr, alias) in kwargs.items():
+            obj.query.add_aggregate(aggregate_expr, alias, self.model)
+
+        return obj
+
     ##################################################################
     # PUBLIC METHODS THAT ALTER ATTRIBUTES AND RETURN A NEW QUERYSET #
     ##################################################################
@@ -335 +389 @@
         Returns a new QuerySet that is a copy of the current one. This allows a
         QuerySet to proxy for a model manager in some cases.
         """
-        return self._clone()
+        return self._clone()        
 
     def filter(self, *args, **kwargs):
         """
@@ -488 +542 @@
         # names of the model fields to select.
 
     def __iter__(self):
+        if self.query.aggregates:
+            return self.aggregate_iterator()
         return self.iterator()
 
+    def aggregate_iterator(self):
+        #Not lazy.. review
+        for i in self.query.get_aggregation():
+            yield i
+
     def iterator(self):
         self.query.trim_extra_select(self.extra_names)
         names = self.query.extra_select.keys() + self.field_names
         for row in self.query.results_iter():
             yield dict(zip(names, row))
-
+ 
     def _setup_query(self):
         """
         Constructs the field_names list that the values query will be