diff --git a/packages/django-google-spanner/README.rst b/packages/django-google-spanner/README.rst index 9e3e8cbd50ab..8bb77001949d 100644 --- a/packages/django-google-spanner/README.rst +++ b/packages/django-google-spanner/README.rst @@ -170,6 +170,30 @@ A transaction is automatically started if you define an To disable transactions in autocommit mode, set the flag `ALLOW_TRANSACTIONS_IN_AUTO_COMMIT` to False in your ``settings.py`` file. +DDL batching +~~~~~~~~~~~~ + +By default, ``DatabaseSchemaEditor`` executes each DDL statement individually. +To batch DDL statements generated within a ``schema_editor()`` context (such as +during ``manage.py migrate`` or test database creation) into batched +``UpdateDatabaseDdl`` requests, set ``'batch_ddl': True`` in ``OPTIONS`` +(defaults to ``False``): + +.. code:: python + + DATABASES = { + 'default': { + 'ENGINE': 'django_spanner', + 'PROJECT': '$PROJECT', + 'INSTANCE': '$INSTANCE', + 'NAME': '$DATABASE', + 'OPTIONS': { + 'batch_ddl': True, + }, + } + } + + Set credentials and project environment variables ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ You'll need to download a service account JSON key file and point to it using an environment variable: diff --git a/packages/django-google-spanner/django_spanner/base.py b/packages/django-google-spanner/django_spanner/base.py index 8e488e41f4f4..9a1bffd38233 100644 --- a/packages/django-google-spanner/django_spanner/base.py +++ b/packages/django-google-spanner/django_spanner/base.py @@ -182,6 +182,7 @@ def get_new_connection(self, conn_params): conn_params.pop("instance", None) conn_params.pop("instance_id", None) conn_params.pop("client", None) + conn_params.pop("batch_ddl", None) # Ensure client is initialized instance = self.instance return self.Database.connect( @@ -209,6 +210,17 @@ def create_cursor(self, name=None): """ return self.connection.cursor() + def set_autocommit( + self, autocommit, force_begin_transaction_with_broken_autocommit=False + ): + editor = getattr(self, "_active_ddl_batch_editor", None) + if editor is not None and editor._ddl_batch_active: + editor._exit_ddl_batch(clean=True) + super().set_autocommit( + autocommit, + force_begin_transaction_with_broken_autocommit=force_begin_transaction_with_broken_autocommit, + ) + def _set_autocommit(self, autocommit): """Set the Spanner transaction autocommit flag. diff --git a/packages/django-google-spanner/django_spanner/creation.py b/packages/django-google-spanner/django_spanner/creation.py index 219240195c5f..2998ab38a733 100644 --- a/packages/django-google-spanner/django_spanner/creation.py +++ b/packages/django-google-spanner/django_spanner/creation.py @@ -48,7 +48,7 @@ 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) + return super().create_test_db(*args, **kwargs) 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..3d2ff8748331 100644 --- a/packages/django-google-spanner/django_spanner/schema.py +++ b/packages/django-google-spanner/django_spanner/schema.py @@ -7,9 +7,24 @@ import uuid import django -from django.db import NotSupportedError +from django.db import ( + IntegrityError, + NotSupportedError, + OperationalError, + ProgrammingError, +) from django.db.backends.base.schema import BaseDatabaseSchemaEditor from django.db.models.fields import NOT_PROVIDED +from google.api_core.exceptions import ( + AlreadyExists, + FailedPrecondition, + GoogleAPICallError, + InternalServerError, + InvalidArgument, + OutOfRange, +) +from google.cloud.spanner_dbapi import parse_utils +from google.cloud.spanner_dbapi.parsed_statement import StatementType from django_spanner import USE_EMULATOR from django_spanner._opentelemetry_tracing import trace_call @@ -57,6 +72,94 @@ 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 __init__(self, connection, collect_sql=False, atomic=True): + self._atomic = atomic + super().__init__(connection, collect_sql=collect_sql, atomic=atomic) + + def _commit_ddls(self): + try: + self.connection.connection.commit() + except (AlreadyExists, FailedPrecondition, OutOfRange) as exc: + raise IntegrityError(getattr(exc, "details", exc)) from exc + except InvalidArgument as exc: + raise ProgrammingError(getattr(exc, "details", exc)) from exc + except (InternalServerError, GoogleAPICallError) as exc: + raise OperationalError(getattr(exc, "details", exc)) from exc + + def _enter_ddl_batch(self): + if not self._ddl_batch_active: + self.connection.ensure_connection() + self.connection._set_autocommit(False) + self._ddl_batch_active = True + + def _exit_ddl_batch(self, clean): + if not self._ddl_batch_active: + return + try: + if clean: + self._commit_ddls() + else: + self.connection.connection.rollback() + finally: + self._ddl_batch_active = False + self.connection._set_autocommit(True) + + def _ddl_execution_wrapper(self, execute, sql, params, many, context): + is_ddl = False + if ( + not many + and sql is not None + and not self.connection.in_atomic_block + and self.connection.get_autocommit() + ): + parsed = parse_utils.classify_statement(str(sql)) + is_ddl = parsed is not None and parsed.statement_type == StatementType.DDL + + if is_ddl: + self._enter_ddl_batch() + elif self._ddl_batch_active: + self._exit_ddl_batch(clean=True) + + return execute(sql, params, many, context) + + def __enter__(self): + super().__enter__() + self._ddl_batch_active = False + self._execute_wrapper_cm = None + batch_ddl = self.connection.settings_dict.get("OPTIONS", {}).get( + "batch_ddl", False + ) + if ( + batch_ddl + and getattr(self, "_atomic", True) + and not self.collect_sql + and not self.connection.in_atomic_block + and self.connection.get_autocommit() + and getattr(self.connection, "_active_ddl_batch_editor", None) is None + ): + self.connection._active_ddl_batch_editor = self + self._execute_wrapper_cm = self.connection.execute_wrapper( + self._ddl_execution_wrapper + ) + self._execute_wrapper_cm.__enter__() + 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: + try: + self._exit_ddl_batch(clean=clean_exit) + finally: + if getattr(self, "_execute_wrapper_cm", None) is not None: + self._execute_wrapper_cm.__exit__(exc_type, exc_value, traceback) + self._execute_wrapper_cm = None + if getattr(self.connection, "_active_ddl_batch_editor", None) is self: + self.connection._active_ddl_batch_editor = None + def create_model(self, model): """ Create a table and any accompanying indexes or unique constraints for diff --git a/packages/django-google-spanner/django_test_suite_5.2.sh b/packages/django-google-spanner/django_test_suite_5.2.sh index 52e489561f31..9e2141c9143d 100755 --- a/packages/django-google-spanner/django_test_suite_5.2.sh +++ b/packages/django-google-spanner/django_test_suite_5.2.sh @@ -47,12 +47,18 @@ DATABASES = { 'PROJECT': "$PROJECT", 'INSTANCE': "$INSTANCE", 'NAME': "$TEST_DBNAME", + 'OPTIONS': { + 'batch_ddl': True, + }, }, 'other': { 'ENGINE': 'django_spanner', 'PROJECT': "$PROJECT", 'INSTANCE': "$INSTANCE", 'NAME': "$TEST_DBNAME_OTHER", + 'OPTIONS': { + 'batch_ddl': True, + }, }, } USE_TZ = False diff --git a/packages/django-google-spanner/django_test_suite_6.0.sh b/packages/django-google-spanner/django_test_suite_6.0.sh index eeab3d35e62d..03414cf7893c 100755 --- a/packages/django-google-spanner/django_test_suite_6.0.sh +++ b/packages/django-google-spanner/django_test_suite_6.0.sh @@ -49,12 +49,18 @@ DATABASES = { 'PROJECT': "$PROJECT", 'INSTANCE': "$INSTANCE", 'NAME': "$TEST_DBNAME", + 'OPTIONS': { + 'batch_ddl': True, + }, }, 'other': { 'ENGINE': 'django_spanner', 'PROJECT': "$PROJECT", 'INSTANCE': "$INSTANCE", 'NAME': "$TEST_DBNAME_OTHER", + 'OPTIONS': { + 'batch_ddl': True, + }, }, } USE_TZ = False diff --git a/packages/django-google-spanner/docs/README.rst b/packages/django-google-spanner/docs/README.rst index 9e3e8cbd50ab..8bb77001949d 100644 --- a/packages/django-google-spanner/docs/README.rst +++ b/packages/django-google-spanner/docs/README.rst @@ -170,6 +170,30 @@ A transaction is automatically started if you define an To disable transactions in autocommit mode, set the flag `ALLOW_TRANSACTIONS_IN_AUTO_COMMIT` to False in your ``settings.py`` file. +DDL batching +~~~~~~~~~~~~ + +By default, ``DatabaseSchemaEditor`` executes each DDL statement individually. +To batch DDL statements generated within a ``schema_editor()`` context (such as +during ``manage.py migrate`` or test database creation) into batched +``UpdateDatabaseDdl`` requests, set ``'batch_ddl': True`` in ``OPTIONS`` +(defaults to ``False``): + +.. code:: python + + DATABASES = { + 'default': { + 'ENGINE': 'django_spanner', + 'PROJECT': '$PROJECT', + 'INSTANCE': '$INSTANCE', + 'NAME': '$DATABASE', + 'OPTIONS': { + 'batch_ddl': True, + }, + } + } + + Set credentials and project environment variables ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ You'll need to download a service account JSON key file and point to it using an environment variable: diff --git a/packages/django-google-spanner/docs/limitations-spanner.rst b/packages/django-google-spanner/docs/limitations-spanner.rst index 75cbbebaf6b3..94fc26ed22e2 100644 --- a/packages/django-google-spanner/docs/limitations-spanner.rst +++ b/packages/django-google-spanner/docs/limitations-spanner.rst @@ -20,9 +20,10 @@ it inherently. The workaround is to use python-spanner objects and run Mutations Batch DDL ~~~~~~~~~ -Migrations (creation and updation) of tables in Django are sequential operations. -Django doesn't inherently support doing batch DDL operations. However, customers -can use python-spanner objects to run batch DDL for unmanaged tables in Django. +By default, migrations (creation and update of tables) in Django execute DDL +statements sequentially. You can opt in to batching DDL statements within each +``schema_editor()`` context by setting ``'OPTIONS': {'batch_ddl': True}`` in +``DATABASES``, or use python-spanner objects directly for unmanaged tables. Stale Reads ~~~~~~~~~~~ @@ -91,7 +92,7 @@ Workaround * - Mutations - Python-spanner database objects can be used to run Mutation queries. `Link `__. * - Batch DDL - - Using unmanaged tables in Django. `Link `__ and Directly using Python spanner objects to execute batch DDL statements. `Link `__. + - Set ``'OPTIONS': {'batch_ddl': True}`` in ``DATABASES`` to batch DDL statements within ``schema_editor()``, or use unmanaged tables (`Link `__) with python-spanner objects directly (`Link `__). * - Stale Reads - Python-spanner database objects can be used to perform stale reads. `Link `__. * - Interleaving diff --git a/packages/django-google-spanner/foreign_key_test.sh b/packages/django-google-spanner/foreign_key_test.sh index 23feb5fdab2d..3469114cd38c 100644 --- a/packages/django-google-spanner/foreign_key_test.sh +++ b/packages/django-google-spanner/foreign_key_test.sh @@ -44,6 +44,9 @@ DATABASES = { 'PROJECT': '$GOOGLE_CLOUD_PROJECT', 'INSTANCE': '$SPANNER_TEST_INSTANCE', 'NAME': 'test_foreign_keys', + 'OPTIONS': { + 'batch_ddl': True, + }, } } " >> foreign_keys/settings.py diff --git a/packages/django-google-spanner/noxfile.py b/packages/django-google-spanner/noxfile.py index 1889aa16ef87..f6ea5270883a 100644 --- a/packages/django-google-spanner/noxfile.py +++ b/packages/django-google-spanner/noxfile.py @@ -63,7 +63,7 @@ ] UNIT_TEST_MOCKSERVER_DEPENDENCIES = [ - "google-cloud-spanner>=3.70.0", + "google-cloud-spanner>=3.72.0", "sqlparse>=0.4.4", ] diff --git a/packages/django-google-spanner/setup.py b/packages/django-google-spanner/setup.py index e5c88ef94c02..811cb0093997 100644 --- a/packages/django-google-spanner/setup.py +++ b/packages/django-google-spanner/setup.py @@ -21,7 +21,7 @@ # (django_spanner/__init__.py currently enforces SUPPORTED_DJANGO_VERSIONS = [(6, 0), (5, 2)]) dependencies = [ "sqlparse >= 0.3.0", - "google-cloud-spanner >= 3.70.0", + "google-cloud-spanner >= 3.72.0", "django >= 5.2, < 6.1", ] extras = { diff --git a/packages/django-google-spanner/tests/mockserver_tests/mock_spanner.py b/packages/django-google-spanner/tests/mockserver_tests/mock_spanner.py index 3f43fff0423b..70e70694195a 100644 --- a/packages/django-google-spanner/tests/mockserver_tests/mock_spanner.py +++ b/packages/django-google-spanner/tests/mockserver_tests/mock_spanner.py @@ -35,10 +35,14 @@ class MockSpanner: def __init__(self): self.results = {} + self.errors = {} def add_result(self, sql: str, result: result_set.ResultSet): self.results[sql.lower().strip()] = result + def add_error(self, sql: str, code: grpc.StatusCode, details: str): + self.errors[sql.lower().strip()] = (code, details) + def get_result(self, sql: str) -> result_set.ResultSet: result = self.results.get(sql.lower().strip()) if result is None: @@ -135,6 +139,9 @@ def ExecuteStreamingSql(self, request: ExecuteSqlRequest, context): started_transaction = self.__create_transaction( request.session, request.transaction.begin ) + error = self.mock_spanner.errors.get(request.sql.lower().strip()) + if error is not None: + context.abort(error[0], error[1]) partials = self.mock_spanner.get_result_as_partial_result_sets( request.sql, started_transaction ) 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..081377cec685 100644 --- a/packages/django-google-spanner/tests/mockserver_tests/test_basics.py +++ b/packages/django-google-spanner/tests/mockserver_tests/test_basics.py @@ -16,14 +16,21 @@ 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, CreateSessionRequest, ExecuteSqlRequest, + RollbackRequest, TypeCode, ) @@ -404,3 +411,332 @@ 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_batch_ddl_option_enabled(self): + cmd = MigrateCommand() + cmd.verbosity = 0 + connection.close() + connection.settings_dict["OPTIONS"]["batch_ddl"] = True + try: + with mock.patch.object( + connection.introspection, "table_names", return_value=[] + ): + cmd.sync_apps(connection, ["auth", "contenttypes", "sessions", "sites"]) + self.assertNotIn("batch_ddl", connection.connection._connection_variables) + finally: + connection.settings_dict["OPTIONS"].pop("batch_ddl", None) + + 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 OPTIONS["batch_ddl"], 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_with_batch_ddl_option(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.settings_dict["OPTIONS"]["batch_ddl"] = True + try: + with connection.schema_editor(atomic=mig.atomic) as editor: + mig.apply(state, editor) + finally: + connection.settings_dict["OPTIONS"].pop("batch_ddl", None) + + 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.settings_dict["OPTIONS"]["batch_ddl"] = True + try: + with self.assertRaises(RuntimeError): + with connection.schema_editor() as editor: + editor.create_model(BrokenModel) + raise RuntimeError("abort schema block") + finally: + connection.settings_dict["OPTIONS"].pop("batch_ddl", None) + + 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_when_batch_ddl_enabled(self): + orig_name = connection.settings_dict["NAME"] + connection.settings_dict["OPTIONS"]["batch_ddl"] = True + + 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(hasattr(connection, "_batch_ddl_in_schema_editor")) + 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["OPTIONS"].pop("batch_ddl", None) + connection.settings_dict["NAME"] = orig_name + + def test_schema_editor_deferred_sql_exception_discards_queued_ddl(self): + connection.settings_dict["OPTIONS"]["batch_ddl"] = 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.settings_dict["OPTIONS"].pop("batch_ddl", None) + + 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 django.db import IntegrityError, ProgrammingError + from google.api_core.exceptions import FailedPrecondition, InvalidArgument + + connection.ensure_connection() + connection.settings_dict["OPTIONS"]["batch_ddl"] = True + try: + with mock.patch.object( + connection.connection.database, + "update_ddl", + side_effect=InvalidArgument("400 Invalid DDL syntax"), + ): + with self.assertRaises(ProgrammingError): + with connection.schema_editor() as editor: + editor.create_model(Singer) + + self.assertTrue(transaction.get_autocommit()) + self.assertTrue(connection.connection.autocommit) + + with mock.patch.object( + connection.connection.database, + "update_ddl", + side_effect=FailedPrecondition("400 Duplicate name in schema"), + ): + with self.assertRaises(IntegrityError): + with connection.schema_editor() as editor: + editor.create_model(Singer) + finally: + connection.settings_dict["OPTIONS"].pop("batch_ddl", None) + + self.assertTrue(transaction.get_autocommit()) + self.assertTrue(connection.connection.autocommit) + + def test_schema_editor_add_field_with_default_commits_rw_tx_before_alter_column( + self, + ): + add_update_count("UPDATE tests_singer SET stage_name=@a0 WHERE 1=1", 2) + field = models.CharField(max_length=100, default="anon") + field.set_attributes_from_name("stage_name") + + event_log = [] + orig_update_ddl = self.database_admin_service.UpdateDatabaseDdl + + def logged_update_ddl(request, context): + event_log.append(("ddl", list(request.statements))) + return orig_update_ddl(request, context) + + connection.settings_dict["OPTIONS"]["batch_ddl"] = True + try: + with connection.schema_editor() as editor: + editor.add_field(Singer, field) + self.assertFalse(connection.connection._spanner_transaction_started) + finally: + connection.settings_dict["OPTIONS"].pop("batch_ddl", None) + + ddl_requests = self.database_admin_service.requests + self.assertEqual(len(ddl_requests), 2) + self.assertTrue( + any(isinstance(r, CommitRequest) for r in self.spanner_service.requests) + ) + self.assertTrue(transaction.get_autocommit()) + + def test_schema_editor_atomic_false_skips_batch_ddl(self): + class NonAtomicModel(models.Model): + code = models.CharField(max_length=50, unique=True) + + class Meta: + app_label = "tests" + + connection.settings_dict["OPTIONS"]["batch_ddl"] = True + try: + with connection.schema_editor(atomic=False) as editor: + self.assertTrue(transaction.get_autocommit()) + editor.create_model(NonAtomicModel) + finally: + connection.settings_dict["OPTIONS"].pop("batch_ddl", None) + + self.assertEqual(len(self.database_admin_service.requests), 2) + self.assertTrue(transaction.get_autocommit()) + + def test_schema_editor_nested_atomic_integrity_error_rolls_back_and_recovers(self): + import grpc + from django.db import IntegrityError + + add_select1_result() + failing_sql = "UPDATE tests_singer SET first_name=@a0 WHERE id=@a1" + self.spanner_service.mock_spanner.add_error( + failing_sql, + grpc.StatusCode.ALREADY_EXISTS, + "UNIQUE violation on index", + ) + + connection.settings_dict["OPTIONS"]["batch_ddl"] = True + try: + with connection.schema_editor(): + self.assertTrue(transaction.get_autocommit()) + with self.assertRaises(IntegrityError): + with transaction.atomic(): + with connection.cursor() as cursor: + cursor.execute( + "UPDATE tests_singer SET first_name=%s WHERE id=%s", + ["dup", 1], + ) + self.assertFalse(connection.needs_rollback) + with connection.cursor() as cursor: + cursor.execute("select 1") + self.assertEqual(cursor.fetchone()[0], 1) + finally: + self.spanner_service.mock_spanner.errors.pop(failing_sql.lower(), None) + connection.settings_dict["OPTIONS"].pop("batch_ddl", None) + + self.assertTrue( + any(isinstance(r, RollbackRequest) for r in self.spanner_service.requests) + ) + self.assertTrue(transaction.get_autocommit()) + self.assertTrue(connection.connection.autocommit) + + def test_migration_create_model_then_runpython_flushes_ddl_before_atomic(self): + add_update_count("UPDATE tests_singer SET first_name=@a0 WHERE id=@a1", 1) + events = [] + + def run_atomic_dml(apps, schema_editor): + events.append( + ( + "runpython_state", + schema_editor.connection.in_atomic_block, + schema_editor.connection.commit_on_exit, + len(self.database_admin_service.requests), + ) + ) + with schema_editor.connection.cursor() as cursor: + cursor.execute( + "UPDATE tests_singer SET first_name=%s WHERE id=%s", + ["updated", 1], + ) + + mig = Migration("0001_create_then_runpython", "tests") + mig.operations = [ + CreateModel( + name="BatchModelBeforeRunPython", + fields=[ + ("id", models.BigAutoField(primary_key=True)), + ("code", models.CharField(max_length=50, unique=True)), + ], + ), + RunPython(run_atomic_dml), + ] + state = ProjectState() + + connection.settings_dict["OPTIONS"]["batch_ddl"] = True + try: + with connection.schema_editor(atomic=mig.atomic) as editor: + mig.apply(state, editor) + finally: + connection.settings_dict["OPTIONS"].pop("batch_ddl", None) + + # CreateModel's CREATE TABLE was flushed before RunPython's atomic() block began, + # and RunPython ran inside a real outermost atomic block (commit_on_exit=True). + self.assertEqual(events, [("runpython_state", True, True, 1)]) + # Deferred CREATE UNIQUE INDEX flushed on schema_editor.__exit__. + 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), 1) + self.assertTrue( + any(isinstance(r, CommitRequest) for r in self.spanner_service.requests) + ) + self.assertTrue(transaction.get_autocommit()) + self.assertTrue(connection.connection.autocommit) diff --git a/packages/django-google-spanner/tests/unit/django_spanner/test_base.py b/packages/django-google-spanner/tests/unit/django_spanner/test_base.py index f6595688646b..6253d257739b 100644 --- a/packages/django-google-spanner/tests/unit/django_spanner/test_base.py +++ b/packages/django-google-spanner/tests/unit/django_spanner/test_base.py @@ -46,12 +46,12 @@ def test_get_connection_params(self): def test_get_new_connection(self): self.db_wrapper.Database = mock_database = mock.MagicMock() mock_database.connect = mock_connection = mock.MagicMock() - conn_params = {"test_param": "dummy"} + conn_params = {"test_param": "dummy", "batch_ddl": True} self.db_wrapper.get_new_connection(conn_params) mock_connection.assert_called_once_with( self.INSTANCE_ID, client=mock_database.connect.call_args[1]["client"], - **conn_params, + test_param="dummy", ) def test_init_connection_state(self): 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..6388fe0e88ab --- /dev/null +++ b/packages/django-google-spanner/tests/unit/django_spanner/test_creation.py @@ -0,0 +1,45 @@ +# 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_returns_db_name_without_private_flag(self): + creation = DatabaseCreation(self.connection) + + with 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") + self.assertFalse(hasattr(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..1987f3ef2d8d 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,288 @@ 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 _make_batch_ddl_mock_conn(self, events, batch_ddl=True): + import contextlib + + 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._active_ddl_batch_editor = None + mock_conn.settings_dict = ( + {"OPTIONS": {"batch_ddl": batch_ddl}} if batch_ddl is not None else {} + ) + wrappers = [] + + @contextlib.contextmanager + def fake_execute_wrapper(wrapper): + wrappers.append(wrapper) + try: + yield + finally: + wrappers.remove(wrapper) + + mock_conn.execute_wrapper.side_effect = fake_execute_wrapper + mock_conn.connection.commit.side_effect = lambda: events.append("dbapi_commit") + mock_conn.connection.rollback.side_effect = lambda: events.append( + "dbapi_rollback" + ) + mock_conn._set_autocommit.side_effect = lambda val: events.append( + ("_set_autocommit", val) + ) + + def run_cursor_execute(sql, params=None): + def base_exec(s, p, many, ctx): + events.append(("cursor_execute", s)) + + def make_wrapped_executor(w, nxt): + def wrapped(s, p, m, c): + return w(nxt, s, p, m, c) + + return wrapped + + executor = base_exec + for wrapper in reversed(wrappers): + executor = make_wrapped_executor(wrapper, executor) + return executor(sql, params, False, {"connection": mock_conn}) + + cursor_cm = mock.MagicMock() + cursor_cm.__enter__.return_value.execute.side_effect = run_cursor_execute + mock_conn.cursor.return_value = cursor_cm + return mock_conn + + def test_schema_editor_batches_ddl_lazily_when_option_enabled(self): + """Lazily toggles _set_autocommit(False) on DDL, commits, and restores _set_autocommit(True).""" + events = [] + mock_conn = self._make_batch_ddl_mock_conn(events) + + with DatabaseSchemaEditor(mock_conn) as editor: + # Entering schema_editor alone does not disable autocommit before any DDL runs. + self.assertEqual(events, []) + 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), + ( + "cursor_execute", + "CREATE TABLE t1 (id INT64 NOT NULL) PRIMARY KEY(id)", + ), + ("cursor_execute", "CREATE INDEX idx_t1 ON t1 (id)"), + "dbapi_commit", + ("_set_autocommit", True), + ], + ) + mock_conn.set_autocommit.assert_not_called() + + def test_schema_editor_noop_without_ddl_does_not_toggle_autocommit(self): + """Does not toggle DB-API autocommit when no DDL statement executes.""" + events = [] + mock_conn = self._make_batch_ddl_mock_conn(events) + + with DatabaseSchemaEditor(mock_conn) as editor: + editor.execute("SELECT 1", None) + + self.assertEqual(events, [("cursor_execute", "SELECT 1")]) + mock_conn._set_autocommit.assert_not_called() + mock_conn.set_autocommit.assert_not_called() + + def test_schema_editor_flushes_ddl_before_non_ddl_and_resumes_batching(self): + """Flushes queued DDLs and restores DB-API autocommit before DML/DQL, then resumes on next DDL.""" + events = [] + mock_conn = self._make_batch_ddl_mock_conn(events) + + with DatabaseSchemaEditor(mock_conn) as editor: + editor.execute("ALTER TABLE t1 ADD COLUMN c1 INT64", None) + editor.execute("UPDATE t1 SET c1 = %s WHERE 1=1", (0,)) + editor.execute("ALTER TABLE t1 ALTER COLUMN c1 INT64 NOT NULL", None) + editor.deferred_sql.append("CREATE INDEX idx_t1_c1 ON t1 (c1)") + + self.assertEqual( + events, + [ + ("_set_autocommit", False), + ("cursor_execute", "ALTER TABLE t1 ADD COLUMN c1 INT64"), + "dbapi_commit", + ("_set_autocommit", True), + ("cursor_execute", "UPDATE t1 SET c1 = %s WHERE 1=1"), + ("_set_autocommit", False), + ("cursor_execute", "ALTER TABLE t1 ALTER COLUMN c1 INT64 NOT NULL"), + ("cursor_execute", "CREATE INDEX idx_t1_c1 ON t1 (c1)"), + "dbapi_commit", + ("_set_autocommit", True), + ], + ) + + def test_schema_editor_rolls_back_on_exception_when_batching(self): + """Rolls back queued DDLs on DB-API connection before restoring autocommit on exception.""" + events = [] + mock_conn = self._make_batch_ddl_mock_conn(events) + + with self.assertRaises(ValueError): + with DatabaseSchemaEditor(mock_conn) as editor: + editor.execute( + "CREATE TABLE t1 (id INT64 NOT NULL) PRIMARY KEY(id)", None + ) + raise ValueError("ddl failure") + + self.assertEqual( + events, + [ + ("_set_autocommit", False), + ( + "cursor_execute", + "CREATE TABLE t1 (id INT64 NOT NULL) PRIMARY KEY(id)", + ), + "dbapi_rollback", + ("_set_autocommit", True), + ], + ) + + 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 = self._make_batch_ddl_mock_conn(events) + + editor = DatabaseSchemaEditor(mock_conn) + + def fail_on_index(sql, params=None, exc_cls=exc_cls, orig=editor.execute): + if "idx_bad" in sql: + raise exc_cls("deferred failure") + return orig(sql, params) + + with self.assertRaises(exc_cls): + with editor: + editor.execute( + "CREATE TABLE t1 (id INT64 NOT NULL) PRIMARY KEY(id)", None + ) + editor.execute = fail_on_index + editor.deferred_sql.append("CREATE INDEX idx_bad ON t1 (id)") + + self.assertEqual( + events, + [ + ("_set_autocommit", False), + ( + "cursor_execute", + "CREATE TABLE t1 (id INT64 NOT NULL) PRIMARY KEY(id)", + ), + "dbapi_rollback", + ("_set_autocommit", True), + ], + ) + + def test_schema_editor_restores_autocommit_when_commit_raises(self): + """Restores _set_autocommit(True) even if DB-API commit() raises a DDL error.""" + events = [] + mock_conn = self._make_batch_ddl_mock_conn(events) + + def failing_commit(): + events.append("dbapi_commit") + raise RuntimeError("server DDL error") + + mock_conn.connection.commit.side_effect = failing_commit + + with self.assertRaises(RuntimeError): + with DatabaseSchemaEditor(mock_conn) as editor: + editor.execute( + "CREATE TABLE t1 (id INT64 NOT NULL) PRIMARY KEY(id)", None + ) + + self.assertEqual( + events, + [ + ("_set_autocommit", False), + ( + "cursor_execute", + "CREATE TABLE t1 (id INT64 NOT NULL) PRIMARY KEY(id)", + ), + "dbapi_commit", + ("_set_autocommit", True), + ], + ) + + def test_schema_editor_skips_autocommit_toggle_by_default(self): + """Does not install wrapper or toggle autocommit when OPTIONS['batch_ddl'] is False or unset.""" + for settings_dict in ( + {}, + {"OPTIONS": {}}, + {"OPTIONS": {"batch_ddl": False}}, + ): + 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.settings_dict = settings_dict + + with DatabaseSchemaEditor(mock_conn): + pass + + mock_conn.execute_wrapper.assert_not_called() + mock_conn._set_autocommit.assert_not_called() + + def test_schema_editor_skips_autocommit_toggle_for_collect_sql_or_atomic_or_manual( + self, + ): + """Skips batching 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.settings_dict = {"OPTIONS": {"batch_ddl": True}} + + with DatabaseSchemaEditor(mock_conn, collect_sql=collect_sql): + pass + + mock_conn.execute_wrapper.assert_not_called() + mock_conn._set_autocommit.assert_not_called() + + def test_schema_editor_skips_autocommit_toggle_when_atomic_false(self): + """Skips batching when schema_editor(atomic=False) is requested.""" + 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.settings_dict = {"OPTIONS": {"batch_ddl": True}} + + with DatabaseSchemaEditor(mock_conn, atomic=False): + pass + + mock_conn.execute_wrapper.assert_not_called() + mock_conn._set_autocommit.assert_not_called() + + def test_schema_editor_maps_google_api_errors_to_django_database_errors(self): + """Maps GoogleAPICallError subclasses on commit() to django.db DatabaseErrors.""" + from django.db import IntegrityError, OperationalError, ProgrammingError + from google.api_core.exceptions import ( + FailedPrecondition, + InternalServerError, + InvalidArgument, + ) + + cases = [ + (FailedPrecondition("Duplicate name"), IntegrityError), + (InvalidArgument("Invalid syntax"), ProgrammingError), + (InternalServerError("Internal failure"), OperationalError), + ] + for api_exc, expected_django_exc in cases: + events = [] + mock_conn = self._make_batch_ddl_mock_conn(events) + mock_conn.connection.commit.side_effect = api_exc + + with self.assertRaises(expected_django_exc): + with DatabaseSchemaEditor(mock_conn) as editor: + editor.execute( + "CREATE TABLE t1 (id INT64 NOT NULL) PRIMARY KEY(id)", None + ) + + mock_conn._set_autocommit.assert_called_with(True)