From c40f02c296fd8c623323d947c2d875af58c42aab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Mon, 3 Mar 2025 18:03:44 +0100 Subject: [PATCH 1/4] feat: support AUTO_INCREMENT and IDENTITY columns Adds support for IDENTITY and AUTO_INCREMENT columns to the Spanner dialect. These are used by default for primary key generation. By default, IDENTITY columns using a backing bit-reversed sequence are used for primary key generation. The sequence kind to use can be configured by setting the attribute default_sequence_kind on the Spanner dialect. The use of AUTO_INCREMENT columns instead of IDENTITY can be configured by setting the use_auto_increment attribute on the Spanner dialect. --- .../sqlalchemy_spanner/sqlalchemy_spanner.py | 37 +++++- test/mockserver_tests/auto_increment_model.py | 28 ++++ test/mockserver_tests/test_auto_increment.py | 120 ++++++++++++++++++ test/mockserver_tests/test_basics.py | 6 +- test/mockserver_tests/test_json.py | 3 +- test/mockserver_tests/test_quickstart.py | 6 +- 6 files changed, 191 insertions(+), 9 deletions(-) create mode 100644 test/mockserver_tests/auto_increment_model.py create mode 100644 test/mockserver_tests/test_auto_increment.py diff --git a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py index 0c8e0c87..3ab38b52 100644 --- a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -409,11 +409,32 @@ def get_column_specification(self, column, **kwargs): if not column.nullable: colspec += " NOT NULL" + has_identity = ( + column.identity is not None and self.dialect.supports_identity_columns + ) default = self.get_column_default_string(column) - if default is not None: - colspec += " DEFAULT (" + default + ")" - if hasattr(column, "computed") and column.computed is not None: + if ( + column.primary_key + and column is column.table._autoincrement_column + and not has_identity + and default is None + ): + if ( + hasattr(self.dialect, "use_auto_increment") + and self.dialect.use_auto_increment + ): + colspec += " AUTO_INCREMENT" + else: + sequence_kind = getattr( + self.dialect, "default_sequence_kind", "BIT_REVERSED_POSITIVE" + ) + colspec += " GENERATED BY DEFAULT AS IDENTITY (%s)" % sequence_kind + elif has_identity: + colspec += " " + self.process(column.identity) + elif default is not None: + colspec += " DEFAULT (" + default + ")" + elif hasattr(column, "computed") and column.computed is not None: colspec += " " + self.process(column.computed) return colspec @@ -526,6 +547,12 @@ def visit_create_index( return text def get_identity_options(self, identity_options): + text = ["bit_reversed_positive"] + if identity_options.start is not None: + text.append("start counter with %d" % identity_options.start) + return " ".join(text) + + def get_sequence_options(self, identity_options): text = ["sequence_kind = 'bit_reversed_positive'"] if identity_options.start is not None: text.append("start_with_counter = %d" % identity_options.start) @@ -534,7 +561,7 @@ def get_identity_options(self, identity_options): def visit_create_sequence(self, create, prefix=None, **kw): """Builds a ``CREATE SEQUENCE`` statement for the sequence.""" text = "CREATE SEQUENCE %s" % self.preparer.format_sequence(create.element) - options = self.get_identity_options(create.element) + options = self.get_sequence_options(create.element) if options: text += " OPTIONS (" + options + ")" return text @@ -628,11 +655,13 @@ class SpannerDialect(DefaultDialect): supports_default_values = False supports_sequences = True sequences_optional = False + supports_identity_columns = True supports_native_enum = True supports_native_boolean = True supports_native_decimal = True supports_statement_cache = True + postfetch_lastrowid = False insert_returning = True update_returning = True delete_returning = True diff --git a/test/mockserver_tests/auto_increment_model.py b/test/mockserver_tests/auto_increment_model.py new file mode 100644 index 00000000..eb67ab89 --- /dev/null +++ b/test/mockserver_tests/auto_increment_model.py @@ -0,0 +1,28 @@ +# 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. + +from sqlalchemy import String +from sqlalchemy.orm import DeclarativeBase +from sqlalchemy.orm import Mapped +from sqlalchemy.orm import mapped_column + + +class Base(DeclarativeBase): + pass + + +class Singer(Base): + __tablename__ = "singers" + id: Mapped[int] = mapped_column(primary_key=True) + name: Mapped[str] = mapped_column(String) diff --git a/test/mockserver_tests/test_auto_increment.py b/test/mockserver_tests/test_auto_increment.py new file mode 100644 index 00000000..16ced3c7 --- /dev/null +++ b/test/mockserver_tests/test_auto_increment.py @@ -0,0 +1,120 @@ +# Copyright 2024 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. + +from sqlalchemy import create_engine +from sqlalchemy.orm import Session +from sqlalchemy.testing import eq_, is_instance_of +from google.cloud.spanner_v1 import ( + FixedSizePool, + ResultSet, + BatchCreateSessionsRequest, + ExecuteSqlRequest, + CommitRequest, + BeginTransactionRequest, +) +from test.mockserver_tests.mock_server_test_base import ( + MockServerTestBase, + add_result, +) +from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest +import google.cloud.spanner_v1.types.type as spanner_type +import google.cloud.spanner_v1.types.result_set as result_set + + +class TestAutoIncrement(MockServerTestBase): + def test_create_table(self): + from test.mockserver_tests.auto_increment_model import Base + + add_result( + """SELECT true +FROM INFORMATION_SCHEMA.TABLES +WHERE TABLE_SCHEMA="" AND TABLE_NAME="singers" +LIMIT 1 +""", + ResultSet(), + ) + engine = create_engine( + "spanner:///projects/p/instances/i/databases/d", + connect_args={"client": self.client, "pool": FixedSizePool(size=10)}, + ) + Base.metadata.create_all(engine) + requests = self.database_admin_service.requests + eq_(1, len(requests)) + is_instance_of(requests[0], UpdateDatabaseDdlRequest) + eq_(1, len(requests[0].statements)) + eq_( + "CREATE TABLE singers (\n" + "\tid INT64 NOT NULL " + "GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE), \n" + "\tname STRING(MAX) NOT NULL\n" + ") PRIMARY KEY (id)", + requests[0].statements[0], + ) + + def test_insert_row(self): + from test.mockserver_tests.auto_increment_model import Singer + + self.add_insert_result( + "INSERT INTO singers (name) VALUES (@a0) THEN RETURN singers.id" + ) + 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 = Singer(name="Test") + session.add(singer) + # Flush the session to send the insert statement to the database. + session.flush() + eq_(987654321, singer.id) + session.commit() + # Verify the requests that we got. + requests = self.spanner_service.requests + eq_(4, 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], CommitRequest) + + def add_insert_result(self, sql): + result = result_set.ResultSet( + dict( + metadata=result_set.ResultSetMetadata( + dict( + row_type=spanner_type.StructType( + dict( + fields=[ + spanner_type.StructType.Field( + dict( + name="id", + type=spanner_type.Type( + dict(code=spanner_type.TypeCode.INT64) + ), + ) + ) + ] + ) + ) + ) + ), + stats=result_set.ResultSetStats( + dict( + row_count_exact=1, + ) + ), + ) + ) + result.rows.extend([("987654321",)]) + add_result(sql, result) diff --git a/test/mockserver_tests/test_basics.py b/test/mockserver_tests/test_basics.py index 29bffa82..36aee22c 100644 --- a/test/mockserver_tests/test_basics.py +++ b/test/mockserver_tests/test_basics.py @@ -127,7 +127,8 @@ def test_create_table(self): eq_(1, len(requests[0].statements)) eq_( "CREATE TABLE users (\n" - "\tuser_id INT64 NOT NULL, \n" + "\tuser_id INT64 NOT NULL " + "GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE), \n" "\tuser_name STRING(16) NOT NULL\n" ") PRIMARY KEY (user_id)", requests[0].statements[0], @@ -163,7 +164,8 @@ def test_create_multiple_tables(self): for i in range(2): eq_( f"CREATE TABLE table{i} (\n" - "\tid INT64 NOT NULL, \n" + "\tid INT64 NOT NULL " + "GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE), \n" "\tvalue STRING(16) NOT NULL" "\n) PRIMARY KEY (id)", requests[0].statements[i], diff --git a/test/mockserver_tests/test_json.py b/test/mockserver_tests/test_json.py index d38eb704..6395fe3a 100644 --- a/test/mockserver_tests/test_json.py +++ b/test/mockserver_tests/test_json.py @@ -58,7 +58,8 @@ def test_create_table(self): eq_(1, len(requests[0].statements)) eq_( "CREATE TABLE venues (\n" - "\tid INT64 NOT NULL, \n" + "\tid INT64 NOT NULL " + "GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE), \n" "\tname STRING(MAX) NOT NULL, \n" "\tdescription JSON\n" ") PRIMARY KEY (id)", diff --git a/test/mockserver_tests/test_quickstart.py b/test/mockserver_tests/test_quickstart.py index 0b31f9e2..c7db636e 100644 --- a/test/mockserver_tests/test_quickstart.py +++ b/test/mockserver_tests/test_quickstart.py @@ -53,7 +53,8 @@ def test_create_tables(self): eq_(2, len(requests[0].statements)) eq_( "CREATE TABLE user_account (\n" - "\tid INT64 NOT NULL, \n" + "\tid INT64 NOT NULL " + "GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE), \n" "\tname STRING(30) NOT NULL, \n" "\tfullname STRING(MAX)\n" ") PRIMARY KEY (id)", @@ -61,7 +62,8 @@ def test_create_tables(self): ) eq_( "CREATE TABLE address (\n" - "\tid INT64 NOT NULL, \n" + "\tid INT64 NOT NULL " + "GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE), \n" "\temail_address STRING(MAX) NOT NULL, \n" "\tuser_id INT64 NOT NULL, \n" "\tFOREIGN KEY(user_id) REFERENCES user_account (id)\n" From c876e40dae784018c4851c47ae4d4d49b376ea5a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Mon, 3 Mar 2025 19:38:47 +0100 Subject: [PATCH 2/4] test: add system test + fix conformance 1.3 test --- .../sqlalchemy_spanner/sqlalchemy_spanner.py | 4 +- test/mockserver_tests/test_auto_increment.py | 61 ++++++++++++++++++- test/system/test_basics.py | 26 +++++++- 3 files changed, 87 insertions(+), 4 deletions(-) diff --git a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py index 3ab38b52..9670327f 100644 --- a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -410,7 +410,9 @@ def get_column_specification(self, column, **kwargs): colspec += " NOT NULL" has_identity = ( - column.identity is not None and self.dialect.supports_identity_columns + hasattr(column, "identity") + and column.identity is not None + and self.dialect.supports_identity_columns ) default = self.get_column_default_string(column) diff --git a/test/mockserver_tests/test_auto_increment.py b/test/mockserver_tests/test_auto_increment.py index 16ced3c7..353cb797 100644 --- a/test/mockserver_tests/test_auto_increment.py +++ b/test/mockserver_tests/test_auto_increment.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC All rights reserved. +# 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. @@ -62,6 +62,65 @@ def test_create_table(self): requests[0].statements[0], ) + def test_create_auto_increment_table(self): + from test.mockserver_tests.auto_increment_model import Base + + add_result( + """SELECT true +FROM INFORMATION_SCHEMA.TABLES +WHERE TABLE_SCHEMA="" AND TABLE_NAME="singers" +LIMIT 1 +""", + ResultSet(), + ) + engine = create_engine( + "spanner:///projects/p/instances/i/databases/d", + connect_args={"client": self.client, "pool": FixedSizePool(size=10)}, + ) + engine.dialect.use_auto_increment = True + Base.metadata.create_all(engine) + requests = self.database_admin_service.requests + eq_(1, len(requests)) + is_instance_of(requests[0], UpdateDatabaseDdlRequest) + eq_(1, len(requests[0].statements)) + eq_( + "CREATE TABLE singers (\n" + "\tid INT64 NOT NULL AUTO_INCREMENT, \n" + "\tname STRING(MAX) NOT NULL\n" + ") PRIMARY KEY (id)", + requests[0].statements[0], + ) + + def test_create_table_with_specific_sequence_kind(self): + from test.mockserver_tests.auto_increment_model import Base + + add_result( + """SELECT true +FROM INFORMATION_SCHEMA.TABLES +WHERE TABLE_SCHEMA="" AND TABLE_NAME="singers" +LIMIT 1 +""", + ResultSet(), + ) + engine = create_engine( + "spanner:///projects/p/instances/i/databases/d", + connect_args={"client": self.client, "pool": FixedSizePool(size=10)}, + ) + engine.dialect.default_sequence_kind = "non_existing_kind" + Base.metadata.create_all(engine) + requests = self.database_admin_service.requests + eq_(1, len(requests)) + is_instance_of(requests[0], UpdateDatabaseDdlRequest) + eq_(1, len(requests[0].statements)) + eq_( + "CREATE TABLE singers (\n" + "\tid INT64 NOT NULL " + "GENERATED BY DEFAULT AS IDENTITY (non_existing_kind), \n" + "\tname STRING(MAX) NOT NULL\n" + ") PRIMARY KEY (id)", + requests[0].statements[0], + ) + def test_insert_row(self): from test.mockserver_tests.auto_increment_model import Singer diff --git a/test/system/test_basics.py b/test/system/test_basics.py index e3099912..e5411988 100644 --- a/test/system/test_basics.py +++ b/test/system/test_basics.py @@ -24,10 +24,11 @@ MetaData, Boolean, BIGINT, + select, ) from sqlalchemy.orm import Session, DeclarativeBase, Mapped, mapped_column from sqlalchemy.types import REAL -from sqlalchemy.testing import eq_ +from sqlalchemy.testing import eq_, is_true from sqlalchemy.testing.plugin.plugin_base import fixtures @@ -50,6 +51,12 @@ def define_tables(cls, metadata): numbers.c.prime.desc(), spanner_storing=[numbers.c.alternative_name], ) + Table( + "users", + metadata, + Column("ID", Integer, primary_key=True), + Column("name", String(20)), + ) def test_hello_world(self, connection): greeting = connection.execute(text("select 'Hello World'")) @@ -69,7 +76,7 @@ def test_reflect(self, connection): engine = connection.engine meta: MetaData = MetaData() meta.reflect(bind=engine) - eq_(1, len(meta.tables)) + eq_(2, len(meta.tables)) table = meta.tables["numbers"] eq_(5, len(table.columns)) eq_("number", table.columns[0].name) @@ -127,6 +134,11 @@ class Number(Base): prime: Mapped[bool] = mapped_column(Boolean) ln: Mapped[float] = mapped_column(REAL) + class User(Base): + __tablename__ = "users" + ID: Mapped[int] = mapped_column(primary_key=True) + name: Mapped[str] = mapped_column(String(20)) + engine = connection.engine with Session(engine) as session: number = Number( @@ -134,3 +146,13 @@ class Number(Base): ) session.add(number) session.commit() + + with Session(engine) as session: + user = User(name="Test") + session.add(user) + session.commit() + + statement = select(User).filter_by(name="Test") + users = session.scalars(statement).all() + eq_(1, len(users)) + is_true(users[0].ID > 0) From 454900fe7d5129d838326a79652691e4ffb7bdc3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Tue, 4 Mar 2025 13:32:09 +0100 Subject: [PATCH 3/4] docs: add sample and update README --- README.rst | 40 ++++++------- samples/auto_generated_primary_key_sample.py | 59 ++++++++++++++++++++ samples/model.py | 13 ++++- samples/noxfile.py | 5 ++ test/mockserver_tests/test_auto_increment.py | 21 +++++++ 5 files changed, 112 insertions(+), 26 deletions(-) create mode 100644 samples/auto_generated_primary_key_sample.py diff --git a/README.rst b/README.rst index 1ec9dfc9..927848dc 100644 --- a/README.rst +++ b/README.rst @@ -293,29 +293,23 @@ This, however, may require to manually repeat a long list of operations, execute In ``AUTOCOMMIT`` mode automatic transactions retry mechanism is disabled, as every operation is committed just in time, and there is no way an ``Aborted`` exception can happen. -Auto-incremented IDs -~~~~~~~~~~~~~~~~~~~~ - -Cloud Spanner doesn't support autoincremented IDs mechanism due to -performance reasons (`see for more -details `__). -We recommend that you use the Python -`uuid `__ module to -generate primary key fields to avoid creating monotonically increasing -keys. - -Though it's not encouraged to do so, in case you *need* the feature, you -can simulate it manually as follows: - -.. code:: python - - with engine.begin() as connection: - top_id = connection.execute( - select([user.c.user_id]).order_by(user.c.user_id.desc()).limit(1) - ).fetchone() - next_id = top_id[0] + 1 if top_id else 1 - - connection.execute(user.insert(), {"user_id": next_id}) +Auto-increment primary keys +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Spanner uses IDENTITY columns for auto-increment primary key values. +IDENTITY columns use a backing bit-reversed sequence to generate unique +values that are safe to use as primary values in Spanner. These values +work the same as standard auto-increment values, except that they are +not monotonically increasing. This prevents hot-spotting for tables that +receive a large number of writes. + +`See this documentation page for more details `__. + +Auto-generated primary keys must be returned by Spanner after each insert +statement using a ``THEN RETURN`` clause. ``THEN RETURN`` clauses are not +supported with `Batch DML `__. +It is therefore recommended to use for example client-side generated UUIDs +as primary key values instead. Query hints ~~~~~~~~~~~ diff --git a/samples/auto_generated_primary_key_sample.py b/samples/auto_generated_primary_key_sample.py new file mode 100644 index 00000000..3475c488 --- /dev/null +++ b/samples/auto_generated_primary_key_sample.py @@ -0,0 +1,59 @@ +# 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 datetime +import uuid + +from sqlalchemy import create_engine +from sqlalchemy.orm import Session + +from sample_helper import run_sample +from model import Singer, Concert, Venue, TicketSale + + +# Shows how to use an IDENTITY column for primary key generation. IDENTITY +# columns use a backing bit-reversed sequence to generate unique values that are +# safe to use for primary keys in Spanner. +# +# IDENTITY columns are used by default by the Spanner SQLAlchemy dialect for +# standard primary key columns. +# +# id: Mapped[int] = mapped_column(primary_key=True) +# +# This leads to the following table definition: +# +# CREATE TABLE ticket_sales ( +# id INT64 NOT NULL GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE), +# ... +# ) PRIMARY KEY (id) +def auto_generated_primary_key_sample(): + engine = create_engine( + "spanner:///projects/sample-project/" + "instances/sample-instance/" + "databases/sample-database", + echo=True, + ) + with Session(engine) as session: + # Venue automatically generates a primary key value using an IDENTITY + # column. We therefore do not need to specify a primary key value when + # we create an instance of Venue. + venue = Venue(code="CH", name="Concert Hall", active=True) + session.add_all([venue]) + session.commit() + + print("Inserted a venue with ID %d" % venue.id) + + +if __name__ == "__main__": + run_sample(auto_generated_primary_key_sample) diff --git a/samples/model.py b/samples/model.py index 13a4c83d..65fc4a41 100644 --- a/samples/model.py +++ b/samples/model.py @@ -31,8 +31,7 @@ ForeignKeyConstraint, Sequence, TextClause, - func, - FetchedValue, + Index, ) from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship @@ -45,6 +44,10 @@ class Base(DeclarativeBase): # This allows inserts to use Batch DML, as the primary key value does not need # to be returned from Spanner using a THEN RETURN clause. # +# The Venue model uses a standard auto-generated integer primary key. This uses +# an IDENTITY column in Spanner. IDENTITY columns use a backing bit-reversed +# sequence to generate unique values that are safe to use for primary keys. +# # The TicketSale model uses a bit-reversed sequence for primary key generation. # This is achieved by creating a bit-reversed sequence and assigning the id # column of the model a server_default value that gets the next value from that @@ -117,7 +120,11 @@ class Track(Base): class Venue(Base): __tablename__ = "venues" - code: Mapped[str] = mapped_column(String(10), primary_key=True) + __table_args__ = (Index("venues_code_unique", "code", unique=True),) + # Venue uses a standard auto-generated primary key. + # This translates to an IDENTITY column in Spanner. + id: Mapped[int] = mapped_column(primary_key=True) + code: Mapped[str] = mapped_column(String(10)) name: Mapped[str] = mapped_column(String(200), nullable=False) description: Mapped[str] = mapped_column(JSON, nullable=True) active: Mapped[bool] = mapped_column(Boolean, nullable=False) diff --git a/samples/noxfile.py b/samples/noxfile.py index 35b744dc..67c3fae5 100644 --- a/samples/noxfile.py +++ b/samples/noxfile.py @@ -22,6 +22,11 @@ def hello_world(session): _sample(session) +@nox.session() +def auto_generated_primary_key(session): + _sample(session) + + @nox.session() def bit_reversed_sequence(session): _sample(session) diff --git a/test/mockserver_tests/test_auto_increment.py b/test/mockserver_tests/test_auto_increment.py index 353cb797..6bc5e2c0 100644 --- a/test/mockserver_tests/test_auto_increment.py +++ b/test/mockserver_tests/test_auto_increment.py @@ -26,6 +26,7 @@ from test.mockserver_tests.mock_server_test_base import ( MockServerTestBase, add_result, + add_update_count, ) from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest import google.cloud.spanner_v1.types.type as spanner_type @@ -147,6 +148,26 @@ def test_insert_row(self): is_instance_of(requests[2], ExecuteSqlRequest) is_instance_of(requests[3], CommitRequest) + def test_insert_row_with_pk_value(self): + from test.mockserver_tests.auto_increment_model import Singer + + # SQLAlchemy should not use a THEN RETURN clause when a value for the + # primary key has been set on the model. + add_update_count("INSERT INTO singers (id, name) VALUES (@a0, @a1)", 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: + # Manually specify a value for the primary key. + singer = Singer(id=1, name="Test") + session.add(singer) + # Flush the session to send the insert statement to the database. + session.flush() + eq_(1, singer.id) + session.commit() + def add_insert_result(self, sql): result = result_set.ResultSet( dict( From 9e775896f91968164d2090de36a0f544f73dfc43 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Mon, 10 Mar 2025 19:23:42 +0100 Subject: [PATCH 4/4] chore: minor cleanup --- samples/auto_generated_primary_key_sample.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/samples/auto_generated_primary_key_sample.py b/samples/auto_generated_primary_key_sample.py index 3475c488..6c74be6f 100644 --- a/samples/auto_generated_primary_key_sample.py +++ b/samples/auto_generated_primary_key_sample.py @@ -12,14 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -import datetime -import uuid - from sqlalchemy import create_engine from sqlalchemy.orm import Session from sample_helper import run_sample -from model import Singer, Concert, Venue, TicketSale +from model import Venue # Shows how to use an IDENTITY column for primary key generation. IDENTITY @@ -44,6 +41,14 @@ def auto_generated_primary_key_sample(): "databases/sample-database", echo=True, ) + + # Add a line like the following to use AUTO_INCREMENT instead of IDENTITY + # when creating tables in SQLAlchemy. + # https://cloud.google.com/spanner/docs/primary-key-default-value#serial-auto-increment + + # engine.dialect.use_auto_increment = True + # Base.metadata.create_all(engine) + with Session(engine) as session: # Venue automatically generates a primary key value using an IDENTITY # column. We therefore do not need to specify a primary key value when