Merge "Add baked.Result.with_post_criteria method"

2 files changed, 114 insertions, 0 deletions

diff --git a/test/ext/test_baked.py b/test/ext/test_baked.py
index d2fcfbab8..47da6d0ed 100644
--- a/test/ext/test_baked.py
+++ b/test/ext/test_baked.py
@@ -13,6 +13,7 @@ from sqlalchemy.orm import exc as orm_exc
 import itertools
 from sqlalchemy.testing import mock
 from sqlalchemy.testing.assertsql import CompiledSQL
+import contextlib
 
 
 class BakedTest(_fixtures.FixtureTest):
@@ -375,6 +376,71 @@ class LikeQueryTest(BakedTest):
         eq_(len(bq._bakery), 4)
 
 
+class ResultPostCriteriaTest(BakedTest):
+
+    @classmethod
+    def setup_mappers(cls):
+        User = cls.classes.User
+        Address = cls.classes.Address
+        Order = cls.classes.Order
+
+        mapper(User, cls.tables.users, properties={
+            "addresses": relationship(
+                Address, order_by=cls.tables.addresses.c.id),
+            "orders": relationship(
+                Order, order_by=cls.tables.orders.c.id)
+        })
+        mapper(Address, cls.tables.addresses)
+        mapper(Order, cls.tables.orders)
+
+    @contextlib.contextmanager
+    def _fixture(self):
+        from sqlalchemy import event
+        User = self.classes.User
+
+        with testing.db.connect() as conn:
+            @event.listens_for(conn, "before_execute")
+            def before_execute(conn, clauseelement, multiparams, params):
+                assert "yes" in conn._execution_options
+
+            bq = self.bakery(
+                lambda s: s.query(User.id).order_by(User.id))
+
+            sess = Session(conn)
+
+            yield sess, bq
+
+    def test_first(self):
+        with self._fixture() as (sess, bq):
+            result = bq(sess).with_post_criteria(
+                lambda q: q.execution_options(yes=True))
+            eq_(result.first(), (7, ))
+
+    def test_iter(self):
+        with self._fixture() as (sess, bq):
+            result = bq(sess).with_post_criteria(
+                lambda q: q.execution_options(yes=True))
+            eq_(list(result)[0], (7, ))
+
+    def test_spoiled(self):
+        with self._fixture() as (sess, bq):
+
+            result = bq.spoil()(sess).with_post_criteria(
+                lambda q: q.execution_options(yes=True))
+
+            eq_(list(result)[0], (7, ))
+
+    def test_get(self):
+        User = self.classes.User
+        with self._fixture() as (sess, bq):
+            bq = self.bakery(
+                lambda s: s.query(User))
+
+            result = bq(sess).with_post_criteria(
+                lambda q: q.execution_options(yes=True))
+            eq_(result.get(7), User(id=7))
+
+
 class ResultTest(BakedTest):
     __backend__ = True
 
diff --git a/test/ext/test_horizontal_shard.py b/test/ext/test_horizontal_shard.py
index 2a596d8c0..79487b2a7 100644
--- a/test/ext/test_horizontal_shard.py
+++ b/test/ext/test_horizontal_shard.py
@@ -186,6 +186,54 @@ class ShardTest(object):
         eq_(set([c.city for c in asia_and_europe]), set(['Tokyo',
             'London', 'Dublin']))
 
+    def test_get_baked_query(self):
+        sess = self._fixture_data()
+
+        tokyo = sess.query(WeatherLocation).filter_by(city="Tokyo").one()
+        tokyo.city
+        sess.expunge_all()
+
+        from sqlalchemy.ext.baked import BakedQuery
+
+        bakery = BakedQuery.bakery()
+
+        bq = bakery(lambda session: session.query(WeatherLocation))
+        t = bq(sess).get(tokyo.id)
+        eq_(t.city, tokyo.city)
+
+    def test_get_baked_query_shard_id(self):
+        sess = self._fixture_data()
+
+        tokyo = sess.query(WeatherLocation).filter_by(city="Tokyo").one()
+        tokyo.city
+        sess.expunge_all()
+
+        from sqlalchemy.ext.baked import BakedQuery
+
+        bakery = BakedQuery.bakery()
+
+        bq = bakery(lambda session: session.query(WeatherLocation))
+        t = bq(sess).with_post_criteria(
+            lambda q: q.set_shard("asia")).get(tokyo.id)
+        eq_(t.city, tokyo.city)
+
+    def test_filter_baked_query_shard_id(self):
+        sess = self._fixture_data()
+
+        tokyo = sess.query(WeatherLocation).filter_by(city="Tokyo").one()
+        tokyo.city
+        sess.expunge_all()
+
+        from sqlalchemy.ext.baked import BakedQuery
+
+        bakery = BakedQuery.bakery()
+
+        bq = bakery(lambda session: session.query(WeatherLocation)).\
+            with_criteria(lambda q: q.filter_by(id=tokyo.id))
+        t = bq(sess).with_post_criteria(
+            lambda q: q.set_shard("asia")).one()
+        eq_(t.city, tokyo.city)
+
     def test_shard_id_event(self):
         canary = []