From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Tue, 27 Jul 2021 15:35:41 +0000 (-0400)
Subject: accommodate plain core textual statements
X-Git-Tag: rel_1_4_23~29
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=95648050cfb2afb5bea8558a6246ed3b5360e7fd;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

accommodate plain core textual statements

Fixed issue where the horizontal sharding extension would not correctly
accommodate for a plain textual SQL statement passed to
:meth:`_orm.Session.execute`.

Fixes: #6816
Change-Id: Ie2b71b06d10793443dbd5e1b271c56cbf9431bb3
---

diff --git a/doc/build/changelog/unreleased_14/6816.rst b/doc/build/changelog/unreleased_14/6816.rst
new file mode 100644
index 0000000000..f778fc11b7
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/6816.rst
@@ -0,0 +1,7 @@
+.. change::
+    :tags: bug, ext
+    :tickets: 6816
+
+    Fixed issue where the horizontal sharding extension would not correctly
+    accommodate for a plain textual SQL statement passed to
+    :meth:`_orm.Session.execute`.
diff --git a/lib/sqlalchemy/ext/horizontal_shard.py b/lib/sqlalchemy/ext/horizontal_shard.py
index 8fb3bf2827..5f13ad2689 100644
--- a/lib/sqlalchemy/ext/horizontal_shard.py
+++ b/lib/sqlalchemy/ext/horizontal_shard.py
@@ -211,9 +211,11 @@ def execute_and_instances(orm_context):
         load_options = active_options = orm_context.load_options
         update_options = None
 
-    else:
+    elif orm_context.is_update or orm_context.is_delete:
         load_options = None
         update_options = active_options = orm_context.update_delete_options
+    else:
+        load_options = update_options = active_options = None
 
     session = orm_context.session
 
@@ -226,7 +228,7 @@ def execute_and_instances(orm_context):
         if orm_context.is_select:
             load_options += {"_refresh_identity_token": shard_id}
             execution_options["_sa_orm_load_options"] = load_options
-        else:
+        elif orm_context.is_update or orm_context.is_delete:
             update_options += {"_refresh_identity_token": shard_id}
             execution_options["_sa_orm_update_options"] = update_options
 
@@ -234,7 +236,7 @@ def execute_and_instances(orm_context):
             bind_arguments=bind_arguments, execution_options=execution_options
         )
 
-    if active_options._refresh_identity_token is not None:
+    if active_options and active_options._refresh_identity_token is not None:
         shard_id = active_options._refresh_identity_token
     elif "_sa_shard_id" in orm_context.execution_options:
         shard_id = orm_context.execution_options["_sa_shard_id"]
diff --git a/test/ext/test_horizontal_shard.py b/test/ext/test_horizontal_shard.py
index 76e94918a8..b05510f7a6 100644
--- a/test/ext/test_horizontal_shard.py
+++ b/test/ext/test_horizontal_shard.py
@@ -15,6 +15,7 @@ from sqlalchemy import sql
 from sqlalchemy import String
 from sqlalchemy import Table
 from sqlalchemy import testing
+from sqlalchemy import text
 from sqlalchemy import update
 from sqlalchemy import util
 from sqlalchemy.ext.horizontal_shard import ShardedSession
@@ -27,6 +28,7 @@ from sqlalchemy.orm import Session
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy.pool import SingletonThreadPool
 from sqlalchemy.sql import operators
+from sqlalchemy.sql import Select
 from sqlalchemy.testing import eq_
 from sqlalchemy.testing import fixtures
 from sqlalchemy.testing import is_
@@ -120,7 +122,7 @@ class ShardTest(object):
                             for value in binary.right.value:
                                 ids.append(shard_lookup[value])
 
-            if query.whereclause is not None:
+            if isinstance(query, Select) and query.whereclause is not None:
                 FindContinent().traverse(query.whereclause)
             if len(ids) == 0:
                 return ["north_america", "asia", "europe", "south_america"]
@@ -677,6 +679,35 @@ class DistinctEngineShardTest(ShardTest, fixtures.TestBase):
         for i in range(1, 5):
             os.remove("shard%d_%s.db" % (i, provision.FOLLOWER_IDENT))
 
+    @testing.combinations((True,), (False,))
+    @testing.uses_deprecated("Using plain strings")
+    def test_plain_core_textual_lookup_w_shard(self, use_legacy_text):
+        sess = self._fixture_data()
+
+        if use_legacy_text:
+            stmt = "SELECT * FROM weather_locations"
+        else:
+            stmt = text("SELECT * FROM weather_locations")
+
+        eq_(
+            sess.execute(stmt, shard_id="asia").fetchall(),
+            [(1, "Asia", "Tokyo")],
+        )
+
+    @testing.combinations((True,), (False,))
+    @testing.uses_deprecated("Using plain strings")
+    def test_plain_core_textual_lookup(self, use_legacy_text):
+        sess = self._fixture_data()
+
+        if use_legacy_text:
+            stmt = "SELECT * FROM weather_locations WHERE id=1"
+        else:
+            stmt = text("SELECT * FROM weather_locations WHERE id=1")
+        eq_(
+            sess.execute(stmt).fetchall(),
+            [(1, "Asia", "Tokyo")],
+        )
+
 
 class AttachedFileShardTest(ShardTest, fixtures.TestBase):
     """Use modern schema conventions along with SQLite ATTACH."""