From c55a964447dfc2f4f77ec4f4a1e2302334dfb1d5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Mon, 17 Feb 2025 14:42:06 +0100 Subject: [PATCH] docs: add test for using FOR UPDATE Spanner now supports FOR UPDATE clauses. This change adds a test to verify that FOR UPDATE clauses can be generated with the Spanner SQLAlchemy provider. See also https://cloud.google.com/spanner/docs/release-notes#January_27_2025 --- .../mockserver_tests/mock_server_test_base.py | 47 +++++++++++++++ test/mockserver_tests/test_basics.py | 35 +++++++++++ .../test_read_only_transaction.py | 54 ++--------------- test/mockserver_tests/test_stale_reads.py | 58 ++----------------- 4 files changed, 92 insertions(+), 102 deletions(-) diff --git a/test/mockserver_tests/mock_server_test_base.py b/test/mockserver_tests/mock_server_test_base.py index fa5ed4bd..a8fea819 100644 --- a/test/mockserver_tests/mock_server_test_base.py +++ b/test/mockserver_tests/mock_server_test_base.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. + from google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode from sqlalchemy import Engine, create_engine from sqlalchemy.testing.plugin.plugin_base import fixtures @@ -79,6 +80,52 @@ def add_single_result( MockServerTestBase.spanner_service.mock_spanner.add_result(sql, result) +def add_singer_query_result(sql: str): + result = result_set.ResultSet( + dict( + metadata=result_set.ResultSetMetadata( + dict( + row_type=spanner_type.StructType( + dict( + fields=[ + spanner_type.StructType.Field( + dict( + name="singers_id", + type=spanner_type.Type( + dict(code=spanner_type.TypeCode.INT64) + ), + ) + ), + spanner_type.StructType.Field( + dict( + name="singers_name", + type=spanner_type.Type( + dict(code=spanner_type.TypeCode.STRING) + ), + ) + ), + ] + ) + ) + ) + ), + ) + ) + result.rows.extend( + [ + ( + "1", + "Jane Doe", + ), + ( + "2", + "John Doe", + ), + ] + ) + add_result(sql, result) + + class MockServerTestBase(fixtures.TestBase): server: grpc.Server = None spanner_service: SpannerServicer = None diff --git a/test/mockserver_tests/test_basics.py b/test/mockserver_tests/test_basics.py index 28ea1b86..f895c9e4 100644 --- a/test/mockserver_tests/test_basics.py +++ b/test/mockserver_tests/test_basics.py @@ -25,7 +25,9 @@ String, func, text, + BigInteger, ) +from sqlalchemy.orm import Session, DeclarativeBase, Mapped, mapped_column from sqlalchemy.testing import eq_, is_instance_of from google.cloud.spanner_v1 import ( FixedSizePool, @@ -41,6 +43,7 @@ add_result, add_single_result, add_update_count, + add_singer_query_result, ) @@ -179,3 +182,35 @@ def test_partitioned_dml(self): ) results = connection.execute(text(sql)).rowcount eq_(100, results) + + def test_select_for_update(self): + class Base(DeclarativeBase): + pass + + class Singer(Base): + __tablename__ = "singers" + id: Mapped[int] = mapped_column(BigInteger, primary_key=True) + name: Mapped[str] = mapped_column(String) + + query = ( + "SELECT singers.id AS singers_id, singers.name AS singers_name\n" + "FROM singers\n" + "WHERE singers.id = @a0\n" + " LIMIT @a1 FOR UPDATE" + ) + add_singer_query_result(query) + update = "UPDATE singers SET name=@a0 WHERE singers.id = @a1" + add_update_count(update, 1) + + engine = create_engine( + "spanner:///projects/p/instances/i/databases/d", + connect_args={"client": self.client, "pool": FixedSizePool(size=10)}, + ) + + with Session(engine) as session: + singer = ( + session.query(Singer).filter(Singer.id == 1).with_for_update().first() + ) + singer.name = "New Name" + session.add(singer) + session.commit() diff --git a/test/mockserver_tests/test_read_only_transaction.py b/test/mockserver_tests/test_read_only_transaction.py index 18abf69f..013c0401 100644 --- a/test/mockserver_tests/test_read_only_transaction.py +++ b/test/mockserver_tests/test_read_only_transaction.py @@ -22,10 +22,10 @@ BeginTransactionRequest, TransactionOptions, ) -from test.mockserver_tests.mock_server_test_base import MockServerTestBase -from test.mockserver_tests.mock_server_test_base import add_result -import google.cloud.spanner_v1.types.type as spanner_type -import google.cloud.spanner_v1.types.result_set as result_set +from test.mockserver_tests.mock_server_test_base import ( + MockServerTestBase, + add_singer_query_result, +) class TestReadOnlyTransaction(MockServerTestBase): @@ -71,49 +71,3 @@ def test_read_only_transaction(self): ), begin_request.options, ) - - -def add_singer_query_result(sql: str): - result = result_set.ResultSet( - dict( - metadata=result_set.ResultSetMetadata( - dict( - row_type=spanner_type.StructType( - dict( - fields=[ - spanner_type.StructType.Field( - dict( - name="singers_id", - type=spanner_type.Type( - dict(code=spanner_type.TypeCode.INT64) - ), - ) - ), - spanner_type.StructType.Field( - dict( - name="singers_name", - type=spanner_type.Type( - dict(code=spanner_type.TypeCode.STRING) - ), - ) - ), - ] - ) - ) - ) - ), - ) - ) - result.rows.extend( - [ - ( - "1", - "Jane Doe", - ), - ( - "2", - "John Doe", - ), - ] - ) - add_result(sql, result) diff --git a/test/mockserver_tests/test_stale_reads.py b/test/mockserver_tests/test_stale_reads.py index 02be066c..d3ac91e8 100644 --- a/test/mockserver_tests/test_stale_reads.py +++ b/test/mockserver_tests/test_stale_reads.py @@ -23,17 +23,17 @@ BeginTransactionRequest, TransactionOptions, ) -from test.mockserver_tests.mock_server_test_base import MockServerTestBase -from test.mockserver_tests.mock_server_test_base import add_result -import google.cloud.spanner_v1.types.type as spanner_type -import google.cloud.spanner_v1.types.result_set as result_set +from test.mockserver_tests.mock_server_test_base import ( + MockServerTestBase, + add_singer_query_result, +) class TestStaleReads(MockServerTestBase): def test_stale_read_multi_use(self): from test.mockserver_tests.stale_read_model import Singer - add_singer_query_result("SELECT singers.id, singers.name \n" + "FROM singers") + add_singer_query_result("SELECT singers.id, singers.name \nFROM singers") engine = create_engine( "spanner:///projects/p/instances/i/databases/d", echo=True, @@ -82,7 +82,7 @@ def test_stale_read_multi_use(self): def test_stale_read_single_use(self): from test.mockserver_tests.stale_read_model import Singer - add_singer_query_result("SELECT singers.id, singers.name\n" + "FROM singers") + add_singer_query_result("SELECT singers.id, singers.name \nFROM singers") engine = create_engine( "spanner:///projects/p/instances/i/databases/d", echo=True, @@ -121,49 +121,3 @@ def test_stale_read_single_use(self): ), execute_request.transaction.single_use, ) - - -def add_singer_query_result(sql: str): - result = result_set.ResultSet( - dict( - metadata=result_set.ResultSetMetadata( - dict( - row_type=spanner_type.StructType( - dict( - fields=[ - spanner_type.StructType.Field( - dict( - name="singers_id", - type=spanner_type.Type( - dict(code=spanner_type.TypeCode.INT64) - ), - ) - ), - spanner_type.StructType.Field( - dict( - name="singers_name", - type=spanner_type.Type( - dict(code=spanner_type.TypeCode.STRING) - ), - ) - ), - ] - ) - ) - ) - ), - ) - ) - result.rows.extend( - [ - ( - "1", - "Jane Doe", - ), - ( - "2", - "John Doe", - ), - ] - ) - add_result(sql, result)