test_auto_increment.py

# 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 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,
    add_update_count,
)
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_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

        self.add_insert_result("INSERT INTO singers (name) VALUES (@a0) THEN RETURN 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 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(
                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)