chore(django-google-spanner): batch DDL statements during create_test_db

packages/django-google-spanner/django_spanner/creation.py

@@ -48,7 +48,11 @@ def create_test_db(self, *args, **kwargs):
         # by a developer running the tests locally.
         if os.environ.get("RUNNING_SPANNER_BACKEND_TESTS") == "1":
             self.mark_skips()
-        super().create_test_db(*args, **kwargs)
+        self.connection._batch_ddl_in_schema_editor = True
+        try:
+            return super().create_test_db(*args, **kwargs)
+        finally:
+            self.connection._batch_ddl_in_schema_editor = False
 
     def _create_test_db(self, verbosity, autoclobber, keepdb=False):
         """

packages/django-google-spanner/django_spanner/schema.py

@@ -57,6 +57,35 @@ class DatabaseSchemaEditor(BaseDatabaseSchemaEditor):
     # sql_create_inline_fk = "CONSTRAINT FK_%(to_table)s_%(to_column)s_%(from_table)s_%(from_column)s FOREIGN KEY (%(from_column_norm)s) REFERENCES %(to_table_norm)s  (%(to_column_norm)s)"  # noqa
     sql_create_inline_fk = None
 
+    def __enter__(self):
+        super().__enter__()
+        self._restore_autocommit = False
+        if (
+            getattr(self.connection, "_batch_ddl_in_schema_editor", False)
+            and not self.collect_sql
+            and not self.connection.in_atomic_block
+            and self.connection.get_autocommit()
+        ):
+            self.connection.set_autocommit(False)
+            self._restore_autocommit = True
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        clean_exit = False
+        try:
+            res = super().__exit__(exc_type, exc_value, traceback)
+            clean_exit = (exc_type is None) or bool(res)
+            return res
+        finally:
+            if getattr(self, "_restore_autocommit", False):
+                try:
+                    if clean_exit:
+                        self.connection.commit()
+                    else:
+                        self.connection.rollback()
+                finally:
+                    self.connection.set_autocommit(True)
+
     def create_model(self, model):
         """
         Create a table and any accompanying indexes or unique constraints for

packages/django-google-spanner/tests/mockserver_tests/test_basics.py

@@ -16,9 +16,15 @@
 from decimal import Decimal
 from unittest import mock
 
+import django
 import google.cloud.spanner_v1.types.result_set as result_set
 import google.cloud.spanner_v1.types.type as spanner_type
-from django.db import connection, models
+from django.core.management.commands.migrate import Command as MigrateCommand
+from django.db import connection, models, transaction
+from django.db.migrations.migration import Migration
+from django.db.migrations.operations.models import CreateModel
+from django.db.migrations.operations.special import RunPython
+from django.db.migrations.state import ProjectState
 from google.cloud.spanner_v1 import (
     BatchCreateSessionsRequest,
     CommitRequest,
@@ -404,3 +410,178 @@ def test_insert_and_fetch_json_field(self):
         # filter(value=None) request -> also queries JSON null ("null", TypeCode.JSON)
         self.assertEqual(exec_requests[14].param_types["a0"].code, TypeCode.JSON)
         self.assertEqual(exec_requests[14].params["a0"], "null")
+
+    def test_sync_apps_batches_ddl_when_create_test_db_flag_enabled(self):
+        cmd = MigrateCommand()
+        cmd.verbosity = 0
+        with mock.patch.object(
+            connection.introspection, "table_names", return_value=[]
+        ):
+            connection._batch_ddl_in_schema_editor = True
+            try:
+                cmd.sync_apps(connection, ["auth", "contenttypes", "sessions", "sites"])
+            finally:
+                connection._batch_ddl_in_schema_editor = False
+
+        ddl_requests = self.database_admin_service.requests
+        self.assertEqual(len(ddl_requests), 1)
+        self.assertEqual(len(ddl_requests[0].statements), 32)
+        self.assertTrue(transaction.get_autocommit())
+        self.assertTrue(connection.connection.autocommit)
+
+    def test_default_schema_editor_executes_ddl_per_statement(self):
+        class LocalAuthor(models.Model):
+            name = models.CharField(max_length=100, unique=True)
+
+            class Meta:
+                app_label = "tests"
+
+        with connection.schema_editor() as editor:
+            editor.create_model(LocalAuthor)
+
+        # Without _batch_ddl_in_schema_editor, CREATE TABLE and CREATE UNIQUE INDEX
+        # execute as 2 separate UpdateDatabaseDdlRequest calls.
+        ddl_requests = self.database_admin_service.requests
+        self.assertEqual(len(ddl_requests), 2)
+        self.assertTrue(transaction.get_autocommit())
+
+    def test_migration_with_runpython_and_deferred_sql_under_create_test_db_flag(self):
+        add_select1_result()
+
+        def run_query(apps, schema_editor):
+            with schema_editor.connection.cursor() as cursor:
+                cursor.execute("select 1")
+
+        mig = Migration("0001_batch_runpython", "tests")
+        mig.operations = [
+            CreateModel(
+                name="BatchModelOne",
+                fields=[
+                    ("id", models.BigAutoField(primary_key=True)),
+                    ("code", models.CharField(max_length=50, unique=True)),
+                ],
+            ),
+            RunPython(run_query),
+            CreateModel(
+                name="BatchModelTwo",
+                fields=[
+                    ("id", models.BigAutoField(primary_key=True)),
+                    ("title", models.CharField(max_length=50, unique=True)),
+                ],
+            ),
+        ]
+        state = ProjectState()
+
+        connection._batch_ddl_in_schema_editor = True
+        try:
+            with connection.schema_editor(atomic=mig.atomic) as editor:
+                mig.apply(state, editor)
+        finally:
+            connection._batch_ddl_in_schema_editor = False
+
+        ddl_requests = self.database_admin_service.requests
+        self.assertEqual(len(ddl_requests), 2)
+        self.assertEqual(len(ddl_requests[0].statements), 1)
+        self.assertEqual(len(ddl_requests[1].statements), 3)
+        self.assertTrue(
+            any(isinstance(r, CommitRequest) for r in self.spanner_service.requests)
+        )
+        self.assertTrue(transaction.get_autocommit())
+
+    def test_schema_editor_exception_discards_queued_ddl_when_batching(self):
+        class BrokenModel(models.Model):
+            name = models.CharField(max_length=100)
+
+            class Meta:
+                app_label = "tests"
+
+        connection._batch_ddl_in_schema_editor = True
+        try:
+            with self.assertRaises(RuntimeError):
+                with connection.schema_editor() as editor:
+                    editor.create_model(BrokenModel)
+                    raise RuntimeError("abort schema block")
+        finally:
+            connection._batch_ddl_in_schema_editor = False
+
+        self.assertEqual(len(self.database_admin_service.requests), 0)
+        self.assertTrue(transaction.get_autocommit())
+        self.assertTrue(connection.connection.autocommit)
+
+    def test_create_test_db_end_to_end_batches_syncdb_ddls(self):
+        orig_name = connection.settings_dict["NAME"]
+
+        def fake_migrate_syncdb(*args, **kwargs):
+            cmd = MigrateCommand()
+            cmd.verbosity = 0
+            with mock.patch.object(
+                connection.introspection, "table_names", return_value=[]
+            ):
+                cmd.sync_apps(connection, ["sites", "sessions"])
+
+        def fake_call_command(name, *args, **kwargs):
+            if name == "migrate":
+                fake_migrate_syncdb(*args, **kwargs)
+
+        try:
+            with (
+                mock.patch.object(
+                    connection.creation,
+                    "_create_test_db",
+                    return_value="test_django_test_db",
+                ),
+                mock.patch(
+                    "django.core.management.call_command",
+                    side_effect=fake_call_command,
+                ),
+            ):
+                create_kwargs = {"verbosity": 0, "autoclobber": True}
+                if django.VERSION < (6, 0):
+                    create_kwargs["serialize"] = False
+                db_name = connection.creation.create_test_db(**create_kwargs)
+
+            self.assertEqual(db_name, "django_test_db")
+            self.assertFalse(getattr(connection, "_batch_ddl_in_schema_editor", False))
+            ddl_requests = self.database_admin_service.requests
+            self.assertEqual(len(ddl_requests), 1)
+            # sites (table + unique index) + sessions (table + expire_date index) = 4 DDLs
+            self.assertEqual(len(ddl_requests[0].statements), 4)
+            self.assertTrue(transaction.get_autocommit())
+        finally:
+            connection.settings_dict["NAME"] = orig_name
+
+    def test_schema_editor_deferred_sql_exception_discards_queued_ddl(self):
+        connection._batch_ddl_in_schema_editor = True
+        try:
+            with self.assertRaises(ValueError):
+                with connection.schema_editor() as editor:
+                    editor.create_model(Singer)
+                    editor.deferred_sql.append(
+                        "CREATE INDEX idx_s ON tests_singer (first_name); SELECT 1"
+                    )
+        finally:
+            connection._batch_ddl_in_schema_editor = False
+
+        self.assertEqual(len(self.database_admin_service.requests), 0)
+        self.assertTrue(transaction.get_autocommit())
+        self.assertTrue(connection.connection.autocommit)
+
+    def test_schema_editor_server_side_ddl_error_restores_autocommit(self):
+        from google.api_core.exceptions import InvalidArgument
+
+        connection.ensure_connection()
+        connection._batch_ddl_in_schema_editor = True
+        try:
+            with mock.patch.object(
+                connection.connection.database,
+                "update_ddl",
+                side_effect=InvalidArgument("400 Invalid DDL syntax"),
+            ):
+                with self.assertRaises(InvalidArgument):
+                    with connection.schema_editor() as editor:
+                        editor.create_model(Singer)
+        finally:
+            connection._batch_ddl_in_schema_editor = False
+
+        self.assertTrue(transaction.get_autocommit())
+        self.assertTrue(connection.connection.autocommit)

packages/django-google-spanner/tests/unit/django_spanner/test_creation.py

@@ -0,0 +1,66 @@
+# Copyright 2026 Google LLC
+#
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file or at
+# https://developers.google.com/open-source/licenses/bsd
+
+import os
+from unittest import mock
+
+from django.db.backends.base.creation import BaseDatabaseCreation
+
+from django_spanner.creation import DatabaseCreation
+from tests.unit.django_spanner.simple_test import SpannerSimpleTestClass
+
+
+class TestCreation(SpannerSimpleTestClass):
+    def test_create_test_db_enables_ddl_batching_and_returns_db_name(self):
+        creation = DatabaseCreation(self.connection)
+        observed_flags = []
+
+        def fake_super_create_test_db(*args, **kwargs):
+            observed_flags.append(
+                getattr(self.connection, "_batch_ddl_in_schema_editor", False)
+            )
+            return "test_database_id"
+
+        with mock.patch.object(
+            BaseDatabaseCreation,
+            "create_test_db",
+            side_effect=fake_super_create_test_db,
+        ):
+            result = creation.create_test_db(verbosity=0, autoclobber=True)
+
+        self.assertEqual(result, "test_database_id")
+        self.assertEqual(observed_flags, [True])
+        self.assertFalse(self.connection._batch_ddl_in_schema_editor)
+
+    def test_create_test_db_resets_ddl_batching_on_error(self):
+        creation = DatabaseCreation(self.connection)
+
+        with mock.patch.object(
+            BaseDatabaseCreation,
+            "create_test_db",
+            side_effect=RuntimeError("boom"),
+        ):
+            with self.assertRaises(RuntimeError):
+                creation.create_test_db(verbosity=0, autoclobber=True)
+
+        self.assertFalse(self.connection._batch_ddl_in_schema_editor)
+
+    def test_create_test_db_marks_skips_when_env_set(self):
+        creation = DatabaseCreation(self.connection)
+
+        with (
+            mock.patch.dict(os.environ, {"RUNNING_SPANNER_BACKEND_TESTS": "1"}),
+            mock.patch.object(creation, "mark_skips") as mock_mark_skips,
+            mock.patch.object(
+                BaseDatabaseCreation,
+                "create_test_db",
+                return_value="test_database_id",
+            ),
+        ):
+            result = creation.create_test_db(verbosity=0, autoclobber=True)
+
+        self.assertEqual(result, "test_database_id")
+        mock_mark_skips.assert_called_once_with()