feat: support AUTO_INCREMENT and IDENTITY columns

olavloite · olavloite · commit c40f02c296fd · 2025-03-03T18:03:47.000+01:00
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.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -53,15 +53,17 @@ 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)",
             requests[0].statements[0],
         )
         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"
