diff --git a/samples/noxfile.py b/samples/noxfile.py index 82019f5b..42d1120a 100644 --- a/samples/noxfile.py +++ b/samples/noxfile.py @@ -62,6 +62,11 @@ def transaction(session): _sample(session) +@nox.session() +def tags(session): + _sample(session) + + @nox.session() def isolation_level(session): _sample(session) diff --git a/samples/tags_sample.py b/samples/tags_sample.py new file mode 100644 index 00000000..a75bddd5 --- /dev/null +++ b/samples/tags_sample.py @@ -0,0 +1,58 @@ +# Copyright 2025 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +import uuid + +from sqlalchemy import create_engine +from sqlalchemy.orm import Session + +from sample_helper import run_sample +from model import Singer + + +# Shows how to transaction tags and statement tags with Spanner and SQLAlchemy. +def tags_sample(): + engine = create_engine( + "spanner:///projects/sample-project/" + "instances/sample-instance/" + "databases/sample-database", + echo=True, + ) + # Set a transaction_tag in the execution options for the session to set + # a transaction tag. + with Session( + engine.execution_options(transaction_tag="my_transaction_tag") + ) as session: + # The transaction that is automatically started by SQLAlchemy will use the + # transaction tag that is specified in the execution options. + + # Execute a query with a request tag. + singer_id = str(uuid.uuid4()) + singer = session.get( + Singer, singer_id, execution_options={"request_tag": "my_tag_1"} + ) + + # Add the singer if it was not found. + if singer is None: + # The session.Add(..) function does not support execution_options, but we can + # set the execution_options on the connection of this session. This will be + # propagated to the next statement that is executed on the connection. + session.connection().execution_options(request_tag="insert_singer") + singer = Singer(id=singer_id, first_name="John", last_name="Doe") + session.add(singer) + session.commit() + + +if __name__ == "__main__": + run_sample(tags_sample) diff --git a/test/mockserver_tests/test_tags.py b/test/mockserver_tests/test_tags.py index c422bc5e..8c157154 100644 --- a/test/mockserver_tests/test_tags.py +++ b/test/mockserver_tests/test_tags.py @@ -65,6 +65,11 @@ def test_transaction_tag(self): from test.mockserver_tests.tags_model import Singer add_singer_query_result("SELECT singers.id, singers.name\n" + "FROM singers") + add_single_singer_query_result( + "SELECT singers.id AS singers_id, singers.name AS singers_name\n" + "FROM singers\n" + "WHERE singers.id = @a0" + ) add_update_count("INSERT INTO singers (id, name) VALUES (@a0, @a1)", 1) engine = create_engine( "spanner:///projects/p/instances/i/databases/d", @@ -75,26 +80,32 @@ def test_transaction_tag(self): engine.execution_options(transaction_tag="my-transaction-tag") ) as session: # Execute a query and an insert statement in a read/write transaction. + session.get(Singer, 1, execution_options={"request_tag": "my-tag-1"}) session.scalars( - select(Singer).execution_options(request_tag="my-tag-1") + select(Singer).execution_options(request_tag="my-tag-2") ).all() + session.connection().execution_options(request_tag="insert-singer") session.add(Singer(id=1, name="Some Singer")) session.commit() # Verify the requests that we got. requests = self.spanner_service.requests - eq_(5, len(requests)) + eq_(6, len(requests)) is_instance_of(requests[0], BatchCreateSessionsRequest) is_instance_of(requests[1], BeginTransactionRequest) is_instance_of(requests[2], ExecuteSqlRequest) is_instance_of(requests[3], ExecuteSqlRequest) - is_instance_of(requests[4], CommitRequest) + is_instance_of(requests[4], ExecuteSqlRequest) + is_instance_of(requests[5], CommitRequest) for request in requests[2:]: eq_("my-transaction-tag", request.request_options.transaction_tag) + eq_("my-tag-1", requests[2].request_options.request_tag) + eq_("my-tag-2", requests[3].request_options.request_tag) + eq_("insert-singer", requests[4].request_options.request_tag) -def add_singer_query_result(sql: str): - result = result_set.ResultSet( +def empty_singer_result_set(): + return result_set.ResultSet( dict( metadata=result_set.ResultSetMetadata( dict( @@ -124,6 +135,10 @@ def add_singer_query_result(sql: str): ), ) ) + + +def add_singer_query_result(sql: str): + result = empty_singer_result_set() result.rows.extend( [ ( @@ -137,3 +152,16 @@ def add_singer_query_result(sql: str): ] ) add_result(sql, result) + + +def add_single_singer_query_result(sql: str): + result = empty_singer_result_set() + result.rows.extend( + [ + ( + "1", + "Jane Doe", + ), + ] + ) + add_result(sql, result)