From 53c5a4a231b2196562d0117a10f1cef1866467c2 Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Fri, 13 Nov 2026 08:01:59 +0000 Subject: [PATCH 1/2] chore(django-google-spanner): batch DDL statements during create_test_db Scope autocommit-disabled DDL batching in DatabaseSchemaEditor to DatabaseCreation.create_test_db() so test database schema setup sends queued CREATE TABLE, index, and foreign key statements in batched UpdateDatabaseDdl requests while keeping autocommit enabled during test execution. --- .../django_spanner/creation.py | 6 +- .../django_spanner/schema.py | 24 +++ .../tests/mockserver_tests/test_basics.py | 147 +++++++++++++++++- .../unit/django_spanner/test_creation.py | 66 ++++++++ .../tests/unit/django_spanner/test_schema.py | 90 +++++++++++ 5 files changed, 331 insertions(+), 2 deletions(-) create mode 100644 packages/django-google-spanner/tests/unit/django_spanner/test_creation.py diff --git a/packages/django-google-spanner/django_spanner/creation.py b/packages/django-google-spanner/django_spanner/creation.py index 219240195c5f..9cafabb8b777 100644 --- a/packages/django-google-spanner/django_spanner/creation.py +++ b/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): """ diff --git a/packages/django-google-spanner/django_spanner/schema.py b/packages/django-google-spanner/django_spanner/schema.py index 3ed8d8359777..c7f991031b56 100644 --- a/packages/django-google-spanner/django_spanner/schema.py +++ b/packages/django-google-spanner/django_spanner/schema.py @@ -57,6 +57,30 @@ 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): + try: + super().__exit__(exc_type, exc_value, traceback) + finally: + if getattr(self, "_restore_autocommit", False): + try: + if exc_type is not None: + 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 diff --git a/packages/django-google-spanner/tests/mockserver_tests/test_basics.py b/packages/django-google-spanner/tests/mockserver_tests/test_basics.py index e728c1e3570e..dfe59000d472 100644 --- a/packages/django-google-spanner/tests/mockserver_tests/test_basics.py +++ b/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,142 @@ 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 diff --git a/packages/django-google-spanner/tests/unit/django_spanner/test_creation.py b/packages/django-google-spanner/tests/unit/django_spanner/test_creation.py new file mode 100644 index 000000000000..15cd18ee6207 --- /dev/null +++ b/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() diff --git a/packages/django-google-spanner/tests/unit/django_spanner/test_schema.py b/packages/django-google-spanner/tests/unit/django_spanner/test_schema.py index ea4888e6f4a8..befc674274f9 100644 --- a/packages/django-google-spanner/tests/unit/django_spanner/test_schema.py +++ b/packages/django-google-spanner/tests/unit/django_spanner/test_schema.py @@ -517,3 +517,93 @@ def test_skip_default_generated_and_db_default(self): normal_field = mock.MagicMock(generated=False, db_default=None) self.assertTrue(schema_editor.skip_default(normal_field)) + + def test_schema_editor_batches_ddl_when_flag_enabled(self): + """Toggles set_autocommit(False) -> set_autocommit(True) around deferred_sql.""" + events = [] + mock_conn = mock.MagicMock() + mock_conn.in_atomic_block = False + mock_conn.features.can_rollback_ddl = False + mock_conn.get_autocommit.return_value = True + mock_conn._batch_ddl_in_schema_editor = True + mock_conn.set_autocommit.side_effect = lambda val: events.append( + ("set_autocommit", val) + ) + + editor = DatabaseSchemaEditor(mock_conn) + editor.execute = mock.MagicMock( + side_effect=lambda sql, params=None: events.append(("execute", sql)) + ) + + with editor: + editor.execute("CREATE TABLE t1 (id INT64 NOT NULL) PRIMARY KEY(id)", None) + editor.deferred_sql.append("CREATE INDEX idx_t1 ON t1 (id)") + + self.assertEqual( + events, + [ + ("set_autocommit", False), + ("execute", "CREATE TABLE t1 (id INT64 NOT NULL) PRIMARY KEY(id)"), + ("execute", "CREATE INDEX idx_t1 ON t1 (id)"), + ("set_autocommit", True), + ], + ) + + def test_schema_editor_rolls_back_on_exception_when_batching(self): + """Rolls back queued DDLs before restoring autocommit on exception.""" + events = [] + mock_conn = mock.MagicMock() + mock_conn.in_atomic_block = False + mock_conn.features.can_rollback_ddl = False + mock_conn.get_autocommit.return_value = True + mock_conn._batch_ddl_in_schema_editor = True + mock_conn.rollback.side_effect = lambda: events.append("rollback") + mock_conn.set_autocommit.side_effect = lambda val: events.append( + ("set_autocommit", val) + ) + + with self.assertRaises(ValueError): + with DatabaseSchemaEditor(mock_conn): + raise ValueError("ddl failure") + + self.assertEqual( + events, + [ + ("set_autocommit", False), + "rollback", + ("set_autocommit", True), + ], + ) + + def test_schema_editor_skips_autocommit_toggle_by_default(self): + """Does not toggle autocommit when _batch_ddl_in_schema_editor is False/unset.""" + mock_conn = mock.MagicMock() + mock_conn.in_atomic_block = False + mock_conn.features.can_rollback_ddl = False + mock_conn.get_autocommit.return_value = True + mock_conn._batch_ddl_in_schema_editor = False + + with DatabaseSchemaEditor(mock_conn): + pass + + mock_conn.set_autocommit.assert_not_called() + + def test_schema_editor_skips_autocommit_toggle_for_collect_sql_or_atomic_or_manual( + self, + ): + """Skips toggling when collect_sql=True, in_atomic_block=True, or autocommit=False.""" + for collect_sql, in_atomic, is_autocommit in [ + (True, False, True), + (False, True, True), + (False, False, False), + ]: + mock_conn = mock.MagicMock() + mock_conn.in_atomic_block = in_atomic + mock_conn.features.can_rollback_ddl = False + mock_conn.get_autocommit.return_value = is_autocommit + mock_conn._batch_ddl_in_schema_editor = True + + with DatabaseSchemaEditor(mock_conn, collect_sql=collect_sql): + pass + + mock_conn.set_autocommit.assert_not_called() From 8862e994e03f0763019e11d4fa357e2c78ae049e Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Fri, 13 Nov 2026 08:42:36 +0000 Subject: [PATCH 2/2] Addressed comments --- .../django_spanner/schema.py | 9 ++- .../tests/mockserver_tests/test_basics.py | 36 ++++++++++ .../tests/unit/django_spanner/test_schema.py | 69 ++++++++++++++++++- 3 files changed, 111 insertions(+), 3 deletions(-) diff --git a/packages/django-google-spanner/django_spanner/schema.py b/packages/django-google-spanner/django_spanner/schema.py index c7f991031b56..a28e7a117836 100644 --- a/packages/django-google-spanner/django_spanner/schema.py +++ b/packages/django-google-spanner/django_spanner/schema.py @@ -71,12 +71,17 @@ def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): + clean_exit = False try: - super().__exit__(exc_type, exc_value, traceback) + 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 exc_type is not None: + if clean_exit: + self.connection.commit() + else: self.connection.rollback() finally: self.connection.set_autocommit(True) diff --git a/packages/django-google-spanner/tests/mockserver_tests/test_basics.py b/packages/django-google-spanner/tests/mockserver_tests/test_basics.py index dfe59000d472..951f7dd18c88 100644 --- a/packages/django-google-spanner/tests/mockserver_tests/test_basics.py +++ b/packages/django-google-spanner/tests/mockserver_tests/test_basics.py @@ -549,3 +549,39 @@ def fake_call_command(name, *args, **kwargs): 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) diff --git a/packages/django-google-spanner/tests/unit/django_spanner/test_schema.py b/packages/django-google-spanner/tests/unit/django_spanner/test_schema.py index befc674274f9..19fe70524206 100644 --- a/packages/django-google-spanner/tests/unit/django_spanner/test_schema.py +++ b/packages/django-google-spanner/tests/unit/django_spanner/test_schema.py @@ -519,13 +519,14 @@ def test_skip_default_generated_and_db_default(self): self.assertTrue(schema_editor.skip_default(normal_field)) def test_schema_editor_batches_ddl_when_flag_enabled(self): - """Toggles set_autocommit(False) -> set_autocommit(True) around deferred_sql.""" + """Toggles set_autocommit(False), commits, and restores set_autocommit(True).""" events = [] mock_conn = mock.MagicMock() mock_conn.in_atomic_block = False mock_conn.features.can_rollback_ddl = False mock_conn.get_autocommit.return_value = True mock_conn._batch_ddl_in_schema_editor = True + mock_conn.commit.side_effect = lambda: events.append("commit") mock_conn.set_autocommit.side_effect = lambda val: events.append( ("set_autocommit", val) ) @@ -545,6 +546,7 @@ def test_schema_editor_batches_ddl_when_flag_enabled(self): ("set_autocommit", False), ("execute", "CREATE TABLE t1 (id INT64 NOT NULL) PRIMARY KEY(id)"), ("execute", "CREATE INDEX idx_t1 ON t1 (id)"), + "commit", ("set_autocommit", True), ], ) @@ -575,6 +577,71 @@ def test_schema_editor_rolls_back_on_exception_when_batching(self): ], ) + def test_schema_editor_rolls_back_on_deferred_sql_exception_when_batching(self): + """Rolls back when super().__exit__() raises Exception or BaseException.""" + for exc_cls in (RuntimeError, KeyboardInterrupt): + events = [] + mock_conn = mock.MagicMock() + mock_conn.in_atomic_block = False + mock_conn.features.can_rollback_ddl = False + mock_conn.get_autocommit.return_value = True + mock_conn._batch_ddl_in_schema_editor = True + mock_conn.commit.side_effect = lambda: events.append("commit") + mock_conn.rollback.side_effect = lambda: events.append("rollback") + mock_conn.set_autocommit.side_effect = lambda val: events.append( + ("set_autocommit", val) + ) + + editor = DatabaseSchemaEditor(mock_conn) + editor.execute = mock.MagicMock(side_effect=exc_cls("deferred failure")) + + with self.assertRaises(exc_cls): + with editor: + editor.deferred_sql.append("CREATE INDEX idx_bad ON t1 (id)") + + self.assertEqual( + events, + [ + ("set_autocommit", False), + "rollback", + ("set_autocommit", True), + ], + ) + + def test_schema_editor_restores_autocommit_when_commit_raises(self): + """Restores set_autocommit(True) even if commit() raises a DDL error.""" + events = [] + mock_conn = mock.MagicMock() + mock_conn.in_atomic_block = False + mock_conn.features.can_rollback_ddl = False + mock_conn.get_autocommit.return_value = True + mock_conn._batch_ddl_in_schema_editor = True + + def failing_commit(): + events.append("commit") + raise RuntimeError("server DDL error") + + mock_conn.commit.side_effect = failing_commit + mock_conn.set_autocommit.side_effect = lambda val: events.append( + ("set_autocommit", val) + ) + + editor = DatabaseSchemaEditor(mock_conn) + editor.execute = mock.MagicMock() + + with self.assertRaises(RuntimeError): + with editor: + pass + + self.assertEqual( + events, + [ + ("set_autocommit", False), + "commit", + ("set_autocommit", True), + ], + ) + def test_schema_editor_skips_autocommit_toggle_by_default(self): """Does not toggle autocommit when _batch_ddl_in_schema_editor is False/unset.""" mock_conn = mock.MagicMock()