From c40f02c296fd8c623323d947c2d875af58c42aab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
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?= <koloite@gmail.com>
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?= <koloite@gmail.com>
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 <https://cloud.google.com/spanner/docs/schema-design#primary-key-prevent-hotspots>`__).
-We recommend that you use the Python
-`uuid <https://docs.python.org/3/library/uuid.html>`__ 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 <https://cloud.google.com/spanner/docs/schema-design#primary-key-prevent-hotspots>`__.
+
+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 <https://cloud.google.com/spanner/docs/dml-tasks#use-batch>`__.
+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?= <koloite@gmail.com>
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