Merge "Add count(), scalar() to baked query"

3 files changed, 80 insertions, 2 deletions

diff --git a/doc/build/changelog/changelog_11.rst b/doc/build/changelog/changelog_11.rst
index 72948aaae..440b8294d 100644
--- a/doc/build/changelog/changelog_11.rst
+++ b/doc/build/changelog/changelog_11.rst
@@ -46,6 +46,13 @@
         2000, as the necessary system view is not available prior to SQL Server
         2005.
 
+    .. change:: 3897
+        :tags: feature, ext
+        :tickets: 3896
+
+        Added :meth:`.baked.Result.scalar` and :meth:`.baked.Result.count`
+        to the "baked" query system.
+
     .. change:: 3893
         :tags: bug, orm
         :tickets: 3893
diff --git a/lib/sqlalchemy/ext/baked.py b/lib/sqlalchemy/ext/baked.py
index c329eb47b..090466618 100644
--- a/lib/sqlalchemy/ext/baked.py
+++ b/lib/sqlalchemy/ext/baked.py
@@ -17,7 +17,7 @@ from ..orm.query import Query
 from ..orm import strategies, attributes, properties, \
     strategy_options, util as orm_util, interfaces
 from .. import log as sqla_log
-from ..sql import util as sql_util
+from ..sql import util as sql_util, func, literal_column
 from ..orm import exc as orm_exc
 from .. import exc as sa_exc
 from .. import util
@@ -254,6 +254,40 @@ class Result(object):
         return context.query.params(self._params).\
             with_session(self.session)._execute_and_instances(context)
 
+    def count(self):
+        """return the 'count'.
+
+        Equivalent to :meth:`.Query.count`.
+
+        Note this uses a subquery to ensure an accurate count regardless
+        of the structure of the original statement.
+
+        .. versionadded:: 1.1.6
+
+        """
+
+        col = func.count(literal_column('*'))
+        bq = self.bq.with_criteria(lambda q: q.from_self(col))
+        return bq.for_session(self.session).scalar()
+
+    def scalar(self):
+        """Return the first element of the first result or None
+        if no rows present.  If multiple rows are returned,
+        raises MultipleResultsFound.
+
+        Equivalent to :meth:`.Query.scalar`.
+
+        .. versionadded:: 1.1.6
+
+        """
+        try:
+            ret = self.one()
+            if not isinstance(ret, tuple):
+                return ret
+            return ret[0]
+        except orm_exc.NoResultFound:
+            return None
+
     def first(self):
         """Return the first row.
 
diff --git a/test/ext/test_baked.py b/test/ext/test_baked.py
index 818019de2..337fd4a12 100644
--- a/test/ext/test_baked.py
+++ b/test/ext/test_baked.py
@@ -6,7 +6,7 @@ from sqlalchemy import testing
 from test.orm import _fixtures
 from sqlalchemy.ext.baked import BakedQuery, baked_lazyload, BakedLazyLoader
 from sqlalchemy.ext import baked
-from sqlalchemy import bindparam, func
+from sqlalchemy import bindparam, func, literal_column
 from sqlalchemy.orm import exc as orm_exc
 import itertools
 from sqlalchemy.testing import mock
@@ -241,6 +241,43 @@ class LikeQueryTest(BakedTest):
             eq_(u2.name, 'ed')
         self.assert_sql_count(testing.db, go, 1)
 
+    def test_scalar(self):
+        User = self.classes.User
+
+        bq = self.bakery(lambda s: s.query(User.id))
+
+        sess = Session()
+
+        bq += lambda q: q.filter(User.id == 7)
+
+        eq_(
+            bq(sess).scalar(), 7
+        )
+
+    def test_count(self):
+        User = self.classes.User
+
+        bq = self.bakery(lambda s: s.query(User))
+
+        sess = Session()
+
+        eq_(
+            bq(sess).count(),
+            4
+        )
+
+        bq += lambda q: q.filter(User.id.in_([8, 9]))
+
+        eq_(
+            bq(sess).count(), 2
+        )
+
+        # original query still works
+        eq_(
+            set([(u.id, u.name) for u in bq(sess).all()]),
+            set([(8, 'ed'), (9, 'fred')])
+        )
+
     def test_get_pk_w_null(self):
         """test the re-implementation of logic to do get with IS NULL."""