googleapis · olavloite · Mar 17, 2025 · Mar 3, 2025 · Mar 3, 2025 · Mar 4, 2025
diff --git 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/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/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

diff --git a/samples/auto_generated_primary_key_sample.py b/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)
diff --git 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
@@ -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/auto_increment_model.py 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)