feat: support AUTO_INCREMENT and IDENTITY columns (#610)

* 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.

* test: add system test + fix conformance 1.3 test

* docs: add sample and update README

* chore: minor cleanup

Author: olavloite

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
 ~~~~~~~~~~~

google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py

@@ -409,11 +409,34 @@ def get_column_specification(self, column, **kwargs):
         if not column.nullable:
             colspec += " NOT NULL"
 
+        has_identity = (
+            hasattr(column, "identity")
+            and 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 +549,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 +563,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 +657,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

samples/auto_generated_primary_key_sample.py

@@ -0,0 +1,64 @@
+# 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 create_engine
+from sqlalchemy.orm import Session
+
+from sample_helper import run_sample
+from model import Venue
+
+
+# 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,
+    )
+
+    # 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
+        # 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)

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)

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)

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)