diff --git a/packages/django-google-spanner/README.rst b/packages/django-google-spanner/README.rst index 9e3e8cbd50ab..789db7c8b1de 100644 --- a/packages/django-google-spanner/README.rst +++ b/packages/django-google-spanner/README.rst @@ -170,6 +170,51 @@ 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 (``AUTO_BATCH_DDL = False``), ``django_spanner`` executes each DDL +statement emitted by ``DatabaseSchemaEditor`` individually. To automatically +batch DDL statements generated inside each ``connection.schema_editor()`` block +(including ``manage.py migrate`` and ``manage.py migrate --run-syncdb``) into a +single Cloud Spanner ``UpdateDatabaseDdl`` operation—flushing automatically +before any interleaved queries, DML backfills, or ``RunPython`` +transactions—set ``AUTO_BATCH_DDL`` to ``True`` in your ``DATABASES`` setting +(or ``'auto_batch_ddl': True`` inside ``OPTIONS``): + +.. code:: python + + DATABASES = { + 'default': { + 'ENGINE': 'django_spanner', + 'PROJECT': '$PROJECT', + 'INSTANCE': '$INSTANCE', + 'NAME': '$DATABASE', + 'AUTO_BATCH_DDL': True, + } + } + +Regardless of ``AUTO_BATCH_DDL``, you can always group DDL statements across +multiple ``schema_editor()`` blocks or raw cursors programmatically using +``with connection.ddl_batch():``, or via client-side ``START BATCH DDL`` / +``RUN BATCH`` / ``ABORT BATCH`` SQL statements on ``connection.cursor()``. + +Important operational behavior: + +- **Partial execution on failure**: Cloud Spanner applies statements inside an + ``UpdateDatabaseDdl`` call in order without transactional DDL rollback. If a + later statement in a batch fails, earlier statements in that batch may remain + applied while the Django migration remains unrecorded. +- **Per-connection visibility**: Buffered DDL is tracked on a single Django + database connection until flushed. Queries run on a separate connection, + thread, or external Spanner client mid-migration will not observe pending + schema changes until a non-DDL flush boundary, ``schema_editor()`` exit, or + ``RUN BATCH``. +- **Strict SQL batch mode**: ``START BATCH DDL`` rejects non-DDL queries, DML, + and transactions with ``ProgrammingError`` until ``RUN BATCH`` or + ``ABORT BATCH`` completes the batch. + + 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..4da17069b76e 100644 --- a/packages/django-google-spanner/django_spanner/base.py +++ b/packages/django-google-spanner/django_spanner/base.py @@ -4,24 +4,346 @@ # license that can be found in the LICENSE file or at # https://developers.google.com/open-source/licenses/bsd +import logging import os +import re +from contextlib import contextmanager +import sqlparse from asgiref.sync import sync_to_async +from django.core.exceptions import ImproperlyConfigured from django.db.backends.base.base import BaseDatabaseWrapper +from google.api_core.exceptions import ( + AlreadyExists, + FailedPrecondition, + GoogleAPICallError, + InvalidArgument, + OutOfRange, +) from google.cloud import spanner, spanner_dbapi +from google.cloud.spanner_dbapi import parse_utils +from google.cloud.spanner_dbapi.connection import BatchMode +from google.cloud.spanner_dbapi.parsed_statement import StatementType +from ._opentelemetry_tracing import trace_call from .client import DatabaseClient from .creation import DatabaseCreation from .features import DatabaseFeatures from .introspection import DatabaseIntrospection from .operations import DatabaseOperations from .schema import DatabaseSchemaEditor +from .utils import quote_ddl_literal from .version import __version__ +logger = logging.getLogger(__name__) + # Global cache for Spanner client to prevent multiple initializations # which can cause OpenTelemetry 'MeterProvider override' crashes. _SPANNER_CLIENT_CACHE = None +_BOOL_STRINGS = {"true": True, "false": False} +_DBAPI_CURSOR_RESULT_ATTRS = ( + "_itr", + "_result_set", + "_row_count", + "_batch_dml_rows_count", +) +_RE_BATCH_PREFIX = re.compile(r"^\s*(?:--|/\*|(?:START|RUN|ABORT)\b)", re.IGNORECASE) +_RE_START_BATCH_DDL = re.compile(r"^START\s+BATCH\s+DDL$", re.IGNORECASE) +_RE_RUN_BATCH = re.compile(r"^RUN\s+BATCH$", re.IGNORECASE) +_RE_ABORT_BATCH = re.compile(r"^ABORT\s+BATCH$", re.IGNORECASE) + + +def _parse_bool_setting(value, setting_name): + """Validate and coerce a Django Spanner boolean setting value. + + Accepts ``bool`` or case-insensitive ``"true"``/``"false"`` strings matching + ``RANDOM_ID_GENERATION_ENABLED``, rejecting ambiguous values instead of + silently treating non-empty strings such as ``"false"`` as truthy. + + :param value: Raw setting value from ``DATABASES[alias]``. + :type setting_name: str + :param setting_name: Setting key used in error messages. + :rtype: bool + :returns: Parsed boolean setting value. + :raises django.core.exceptions.ImproperlyConfigured: If ``value`` is not a + ``bool`` or ``"true"``/``"false"`` string. + """ + if type(value) is bool: + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in _BOOL_STRINGS: + return _BOOL_STRINGS[normalized] + raise ImproperlyConfigured( + f"Invalid value {value!r} for {setting_name}; expected a bool or " + "'true'/'false'." + ) + + +def _format_ddl_sql(sql, args, error_cls): + """Inline query parameters into a Cloud Spanner DDL SQL string. + + Cloud Spanner ``UpdateDatabaseDdl`` accepts only literal DDL strings without + bind parameters; parameters passed via ``SchemaEditor`` or ``RunSQL`` are + quoted as GoogleSQL literals before delegating to ``spanner_dbapi``. + + :type sql: str + :param sql: Parameterized DDL SQL template using ``%s`` or ``%(name)s``. + :param args: Non-empty positional sequence or keyword mapping of values. + :param error_cls: PEP 249 ``ProgrammingError`` class to raise on mismatch. + :rtype: str + :returns: DDL statement string with all placeholders replaced by literals. + :raises ProgrammingError: If ``args`` is neither a sequence nor a mapping, + or placeholder substitution fails. + """ + try: + if isinstance(args, dict): + quoted = {key: quote_ddl_literal(val) for key, val in args.items()} + elif isinstance(args, (list, tuple)): + quoted = tuple(quote_ddl_literal(val) for val in args) + else: + raise error_cls("DDL query parameters must be a sequence or mapping.") + return sql % quoted + except error_cls: + raise + except (KeyError, TypeError, ValueError) as exc: + raise error_cls(str(exc)) from exc + + +def _split_and_classify_sql(sql, error_cls): + """Split a SQL string into ``(statement, parsed)`` tuples. + + External SQL comments and trailing semicolons are stripped so DDL templates + can be classified cleanly before delegating to ``Cursor._batch_DDLs``, while + inputs with zero executable tokens raise ``error_cls``. + + :type sql: str + :param sql: Raw SQL text passed to ``cursor.execute``. + :param error_cls: PEP 249 ``ProgrammingError`` class to raise on invalid SQL. + :rtype: list[tuple[str, object]] + :returns: Non-empty ordered list of ``(statement, parsed)`` tuples. + :raises ProgrammingError: If no executable statement exists. + """ + parts = [] + for piece in sqlparse.split(sql): + stripped = sqlparse.format(piece, strip_comments=True).strip() + stmt = stripped.rstrip(";").strip() + if not stmt: + continue + parts.append((stmt, parse_utils.classify_statement(stmt))) + if not parts: + raise error_cls("Invalid Statement.") + return parts + + +def _match_batch_control_sql(sql_str): + """Classify client-side DDL batch SQL without parsing normal DQL/DML.""" + if _RE_BATCH_PREFIX.match(sql_str) is None: + return None + stmts = [] + for piece in sqlparse.split(sql_str): + stripped = sqlparse.format(piece, strip_comments=True).strip() + stmt = stripped.rstrip(";").strip() + if stmt: + stmts.append(stmt) + if len(stmts) > 1: + return None + if len(stmts) != 1: + return None + candidate = stmts[0] + if _RE_START_BATCH_DDL.match(candidate): + return "START_DDL" + if _RE_RUN_BATCH.match(candidate): + return "RUN" + if _RE_ABORT_BATCH.match(candidate): + return "ABORT" + return None + + +class _DdlBatchCursor: + """DB-API cursor proxy coordinating scoped ``spanner_dbapi`` DDL batching. + + Buffers DDL into ``spanner_dbapi.Connection._ddl_statements`` via a scoped + client transaction (``Connection.begin()``), flushes pending DDL via + ``Connection.commit()`` immediately before interleaved non-DDL operations in + ``auto_flush`` mode, intercepts client-side ``START BATCH DDL`` / + ``RUN BATCH`` / ``ABORT BATCH`` SQL, and maps ``spanner_dbapi`` and gRPC + errors to ``django.db`` exception classes. + """ + + def __init__(self, db_wrapper, cursor): + self._db = db_wrapper + self._cursor = cursor + + def __getattr__(self, attr): + return getattr(self._cursor, attr) + + def __setattr__(self, attr, value): + if attr in ("_db", "_cursor"): + object.__setattr__(self, attr, value) + else: + setattr(self._cursor, attr, value) + + def __iter__(self): + return self + + def __next__(self): + with self._db._wrap_spanner_errors(): + return next(self._cursor) + + def __enter__(self): + self._cursor.__enter__() + return self + + def __exit__(self, exc_type, exc_value, traceback): + return self._cursor.__exit__(exc_type, exc_value, traceback) + + def _ensure_open_and_reset(self): + """Validate open cursor state and clear prior query result attributes.""" + if not getattr(self._cursor, "connection", True): + raise self._db.Database.ProgrammingError( + "Cursor is not connected to the database" + ) + if getattr(self._cursor, "is_closed", False) is True: + raise self._db.Database.InterfaceError( + "Cursor and/or connection is already closed." + ) + if ( + self._db.connection is not None + and getattr(self._db.connection, "database", False) is None + ): + raise ValueError("Database needs to be passed for this operation") + for attr in _DBAPI_CURSOR_RESULT_ATTRS: + if hasattr(self._cursor, attr): + setattr(self._cursor, attr, None) + + def execute(self, sql, args=None): + """Execute ``sql``, coordinating scoped ``spanner_dbapi`` DDL batching. + + :type sql: str + :param sql: SQL query, DDL statement(s), or client-side batch command. + :param args: Optional positional or mapping query parameters. + :returns: Result of ``spanner_dbapi.Cursor.execute``, or ``None`` for + buffered DDL and client-side DDL batch statements. + :raises django.db.DatabaseError: On closed cursors, invalid parameters, + or rejected batch transitions. + """ + error_cls = self._db.Database.ProgrammingError + with self._db._wrap_spanner_errors(): + sql_str = str(sql) + batch_cmd = _match_batch_control_sql(sql_str) + if batch_cmd == "START_DDL": + self._ensure_open_and_reset() + if args: + raise error_cls( + "Client-side batch statements do not accept parameters." + ) + self._db.start_batch_ddl(auto_flush=False) + return None + if batch_cmd in ("RUN", "ABORT") and not self._db._in_dml_batch(): + self._ensure_open_and_reset() + if args: + raise error_cls( + "Client-side batch statements do not accept parameters." + ) + if batch_cmd == "RUN": + return self._db.run_batch() + return self._db.abort_batch() + + if self._db.in_ddl_batch and not self._db._in_external_tx_or_dml_batch(): + self._ensure_open_and_reset() + parts = _split_and_classify_sql(sql_str, error_cls) + if all(p.statement_type == StatementType.DDL for _, p in parts): + if args: + rendered_sql = _format_ddl_sql(sql_str, args, error_cls) + rendered_parts = _split_and_classify_sql( + rendered_sql, error_cls + ) + # Re-validate statement count and DDL classification after + # literal substitution so parameter values cannot inject new + # statements or non-DDL operations. + if len(rendered_parts) != len(parts) or any( + p.statement_type != StatementType.DDL + for _, p in rendered_parts + ): + raise error_cls( + "Only DDL statements are allowed in a DDL batch." + ) + parts = rendered_parts + self._db._stage_ddl_statements( + self._cursor, [stmt for stmt, _ in parts] + ) + return None + + if not self._db.ddl_batch_auto_flush: + raise error_cls("Only DDL statements are allowed in a DDL batch.") + + if len(parts) > 1: + if args: + raise error_cls( + "Parameters are not supported for multi-statement SQL " + "in a DDL batch." + ) + for stmt, parsed in parts: + if ( + parsed.statement_type == StatementType.DDL + and not self._db._in_external_tx_or_dml_batch() + ): + self._ensure_open_and_reset() + self._db._stage_ddl_statements(self._cursor, [stmt]) + else: + self._db._before_non_ddl_in_batch() + self._cursor.execute(stmt) + return None + + self._db.flush_ddl_batch() + + if args is None: + return self._cursor.execute(sql) + return self._cursor.execute(sql, args) + + def executemany(self, operation, seq_of_params): + """Delegate ``executemany`` after flushing pending DDL for non-DDL SQL.""" + if self._db.in_ddl_batch and not self._db._in_external_tx_or_dml_batch(): + stmt_type = parse_utils.classify_statement(str(operation)).statement_type + if stmt_type not in (StatementType.DDL, StatementType.CLIENT_SIDE): + self._db._before_non_ddl_in_batch() + with self._db._wrap_spanner_errors(): + return self._cursor.executemany(operation, seq_of_params) + + def callproc(self, procname, args=None): + """Delegate ``callproc`` after flushing or rejecting pending DDL.""" + self._db._before_non_ddl_in_batch() + with self._db._wrap_spanner_errors(): + if args is None: + return self._cursor.callproc(procname) + return self._cursor.callproc(procname, args) + + def run_sql_in_snapshot(self, sql, params=None, param_types=None): + """Delegate ``run_sql_in_snapshot`` after flushing pending DDL.""" + self._db._before_non_ddl_in_batch() + with self._db._wrap_spanner_errors(): + return self._cursor.run_sql_in_snapshot( + sql, params=params, param_types=param_types + ) + + def list_tables(self, schema_name="", include_views=True): + """Delegate ``list_tables`` after flushing pending DDL.""" + self._db._before_non_ddl_in_batch() + with self._db._wrap_spanner_errors(): + return self._cursor.list_tables( + schema_name=schema_name, include_views=include_views + ) + + def get_table_column_schema(self, table_name, schema_name=""): + """Delegate ``get_table_column_schema`` after flushing pending DDL.""" + self._db._before_non_ddl_in_batch() + with self._db._wrap_spanner_errors(): + return self._cursor.get_table_column_schema( + table_name, schema_name=schema_name + ) + class DatabaseWrapper(BaseDatabaseWrapper): vendor = "spanner" @@ -118,6 +440,113 @@ class DatabaseWrapper(BaseDatabaseWrapper): ops_class = DatabaseOperations client_class = DatabaseClient + def __init__(self, *args, **kwargs): + """Initialize per-connection Cloud Spanner wrapper and DDL batch state.""" + super().__init__(*args, **kwargs) + self._ddl_batch_active = False + self._ddl_batch_auto_flush = False + self._ddl_batch_flush_count = 0 + self._ddl_tx_begun = False + + @contextmanager + def _wrap_spanner_errors(self): + """Translate both gRPC ``GoogleAPICallError`` and PEP 249 errors to Django. + + ``Connection.commit()``, ``Connection.run_prior_DDL_statements()``, and + ``Cursor.run_sql_in_snapshot()`` raise unmapped ``GoogleAPICallError`` + subclasses directly, whereas Django's ``wrap_database_errors`` only + catches ``self.Database.Error`` subclasses. + """ + with self.wrap_database_errors: + try: + yield + except (AlreadyExists, FailedPrecondition, OutOfRange) as exc: + raise self.Database.IntegrityError(str(exc)) from exc + except InvalidArgument as exc: + raise self.Database.ProgrammingError(str(exc)) from exc + except GoogleAPICallError as exc: + raise self.Database.OperationalError(str(exc)) from exc + + def _has_open_dbapi_conn(self): + """Return ``True`` when ``self.connection`` holds an open DB-API connection.""" + return ( + self.connection is not None + and getattr(self.connection, "is_closed", False) is not True + ) + + def _in_dml_batch(self): + """Return ``True`` when ``self.connection`` is inside a DB-API DML batch.""" + return ( + self._has_open_dbapi_conn() + and getattr(self.connection, "_batch_mode", BatchMode.NONE) is BatchMode.DML + ) + + def _in_external_tx_or_dml_batch(self): + """Return ``True`` during caller transactions or DML batches outside DDL.""" + return self._in_dml_batch() or ( + self._has_open_dbapi_conn() + and getattr(self.connection, "_client_transaction_started", False) is True + and not self._ddl_tx_begun + ) + + def _stage_ddl_statements(self, raw_cursor, statements): + """Buffer validated DDL via ``spanner_dbapi`` client-transaction mode. + + Starts ``self.connection.begin()`` lazily on the first DDL statement of + a burst so ``Cursor.execute`` appends directly to + ``self.connection._ddl_statements`` via ``Cursor._batch_DDLs`` without + opening a Spanner read-write transaction. + """ + if getattr(self.connection, "read_only", False) is True: + raise self.Database.ProgrammingError( + "Cannot execute DDL statements on a read-only connection." + ) + began_here = not self._ddl_tx_begun + if began_here: + self.connection.begin() + self._ddl_tx_begun = True + try: + raw_cursor.execute("; ".join(statements)) + except BaseException: + if began_here and not getattr(self.connection, "_ddl_statements", None): + self._rollback_ddl_burst() + raise + + def _rollback_ddl_burst(self): + """Clear un-flushed DB-API DDL statements and roll back ``begin()``.""" + was_begun = self._ddl_tx_begun + self._ddl_tx_begun = False + if not self._has_open_dbapi_conn(): + return + # Connection.rollback() resets _transaction_begin_marked without clearing + # _ddl_statements, so clear the list explicitly before rolling back. + ddl_list = getattr(self.connection, "_ddl_statements", None) + if ddl_list: + ddl_list.clear() + if ( + was_begun + and getattr(self.connection, "_client_transaction_started", False) is True + ): + self.connection.rollback() + + def _clear_ddl_batch(self): + """Discard any active DDL burst and reset wrapper DDL batch state.""" + self._rollback_ddl_burst() + self._ddl_batch_active = False + self._ddl_batch_auto_flush = False + self._ddl_batch_flush_count = 0 + + def _before_non_ddl_in_batch(self): + """Flush buffered DDL or raise ``ProgrammingError`` before non-DDL work.""" + if not self.in_ddl_batch or self._in_external_tx_or_dml_batch(): + return + if not self.ddl_batch_auto_flush: + with self.wrap_database_errors: + raise self.Database.ProgrammingError( + "Only DDL statements are allowed in a DDL batch." + ) + self.flush_ddl_batch() + def _get_project_id(self): return os.environ.get("GOOGLE_CLOUD_PROJECT", "test-project") @@ -146,6 +575,267 @@ def allow_transactions_in_auto_commit(self): return self.settings_dict["ALLOW_TRANSACTIONS_IN_AUTO_COMMIT"] return True + @property + def auto_batch_ddl(self): + """Whether ``DatabaseSchemaEditor`` automatically batches DDL statements. + + Reads ``DATABASES[alias]["AUTO_BATCH_DDL"]`` first, falling back to + ``DATABASES[alias]["OPTIONS"]["auto_batch_ddl"]`` and defaulting to + ``False`` (opt-in automatic schema-editor batching). + + :rtype: bool + :returns: ``True`` if schema editor contexts auto-batch DDL. + :raises django.core.exceptions.ImproperlyConfigured: If configured with + a non-boolean value other than ``"true"``/``"false"``. + """ + if "AUTO_BATCH_DDL" in self.settings_dict: + return _parse_bool_setting( + self.settings_dict["AUTO_BATCH_DDL"], "AUTO_BATCH_DDL" + ) + options = self.settings_dict.get("OPTIONS") or {} + if "auto_batch_ddl" in options: + return _parse_bool_setting( + options["auto_batch_ddl"], "OPTIONS['auto_batch_ddl']" + ) + return False + + @property + def in_ddl_batch(self): + """Return ``True`` while a DDL batch is open on this connection.""" + return self._ddl_batch_active and ( + self.connection is None or self._has_open_dbapi_conn() + ) + + @property + def ddl_batch_auto_flush(self): + """Return ``True`` if the active DDL batch flushes before non-DDL SQL.""" + return self.in_ddl_batch and self._ddl_batch_auto_flush + + @property + def buffered_ddl_statements(self): + """Return an immutable snapshot of DDL statements currently buffered.""" + if not self._ddl_tx_begun or not self._has_open_dbapi_conn(): + return () + return tuple(getattr(self.connection, "_ddl_statements", ())) + + def _can_auto_batch_ddl(self): + """Return whether connection state permits starting a SchemaEditor DDL batch. + + Inspects ``settings_dict["AUTOCOMMIT"]`` before ``connect()`` so lazy + unopened wrappers do not open network sessions on empty schema editors, + and inspects live transaction/read-only/batch state once connected. + """ + if not self._has_open_dbapi_conn(): + return bool(self.settings_dict.get("AUTOCOMMIT", True)) + if not self.autocommit: + return False + if getattr(self.connection, "_client_transaction_started", False) is True: + return False + if getattr(self.connection, "read_only", False) is True: + return False + return not self._in_dml_batch() + + def checkpoint_ddl_batch(self): + """Capture a ``(flush_count, statement_count)`` checkpoint for rollback. + + :rtype: tuple[int, int] + :returns: Current flush generation and buffered DDL statement count. + """ + return (self._ddl_batch_flush_count, len(self.buffered_ddl_statements)) + + def rollback_ddl_checkpoint(self, checkpoint): + """Discard un-flushed DDL statements queued after ``checkpoint``. + + :type checkpoint: tuple[int, int] + :param checkpoint: Snapshot returned by :meth:`checkpoint_ddl_batch`. + """ + saved_flushes, saved_len = checkpoint + if ( + self._ddl_tx_begun + and self._ddl_batch_flush_count == saved_flushes + and saved_len > 0 + ): + del self.connection._ddl_statements[saved_len:] + else: + self._rollback_ddl_burst() + + def start_batch_ddl(self, *, auto_flush=True): + """Open a DDL batch backed by scoped ``spanner_dbapi`` transaction mode. + + Remains lazy until the first DDL statement executes so empty schema + editors and unit tests open zero Spanner sessions or ``UpdateDatabaseDdl`` + calls. + + :type auto_flush: bool + :param auto_flush: When ``True`` (default for Python callers and + ``SchemaEditor``), interleaved DQL/DML/transaction boundaries commit + queued DDL automatically before executing. When ``False`` (used by + ``START BATCH DDL``), non-DDL operations raise ``ProgrammingError``. + :raises django.db.utils.ProgrammingError: If a DDL or DML batch is + already active, if a transaction is active or ``AUTOCOMMIT`` is + disabled, or if the connection is ``read_only``. + """ + if self.in_ddl_batch: + with self.wrap_database_errors: + raise self.Database.ProgrammingError( + "Cannot start a DDL batch when a DDL batch is already active." + ) + if self.in_atomic_block: + with self.wrap_database_errors: + raise self.Database.ProgrammingError( + "Cannot start a DDL batch while in a transaction." + ) + if not self._has_open_dbapi_conn(): + if not self.settings_dict.get("AUTOCOMMIT", True): + with self.wrap_database_errors: + raise self.Database.ProgrammingError( + "Cannot start a DDL batch when AUTOCOMMIT is disabled." + ) + else: + if ( + not self.autocommit + or getattr(self.connection, "_client_transaction_started", False) + is True + ): + with self.wrap_database_errors: + raise self.Database.ProgrammingError( + "Cannot start a DDL batch while in a transaction." + ) + if getattr(self.connection, "read_only", False) is True: + with self.wrap_database_errors: + raise self.Database.ProgrammingError( + "Cannot start a DDL batch when the connection is in " + "read-only mode." + ) + if self._in_dml_batch(): + with self.wrap_database_errors: + raise self.Database.ProgrammingError( + "Cannot start a DDL batch when a DML batch is active." + ) + self._ddl_batch_active = True + self._ddl_batch_auto_flush = bool(auto_flush) + self._ddl_batch_flush_count = 0 + self._ddl_tx_begun = False + + def flush_ddl_batch(self): + """Commit buffered DDL via ``spanner_dbapi.Connection.commit()``. + + ``Connection.commit()`` executes ``run_prior_DDL_statements()`` + (submitting one ``database.update_ddl(ddl_statements)`` RPC) and resets + ``_client_transaction_started`` without starting a Spanner read-write + transaction, while leaving :attr:`in_ddl_batch` open for subsequent DDL. + + :returns: ``None``. + :raises ValueError: If the underlying connection has no ``database``. + :raises django.db.DatabaseError: On read-only connections or if Spanner + rejects the DDL batch. + """ + if not self._has_open_dbapi_conn(): + return None + if getattr(self.connection, "database", False) is None: + raise ValueError("Database needs to be passed for this operation") + if getattr(self.connection, "read_only", False) is True: + with self.wrap_database_errors: + raise self.Database.ProgrammingError( + "Cannot execute DDL statements on a read-only connection." + ) + count = len(self.buffered_ddl_statements) + if not count: + self._rollback_ddl_burst() + return None + try: + self._ddl_batch_flush_count += 1 + with self._wrap_spanner_errors(): + with trace_call( + "CloudSpannerDjango.flush_ddl_batch", + self, + {"statement_count": count}, + ): + self.connection.commit() + self._ddl_tx_begun = False + return None + finally: + # Connection.commit() calls run_prior_DDL_statements() before its + # own try/finally reset; roll back on error so _transaction_begin_marked + # is never stranded on the connection. + if self._ddl_tx_begun: + self._rollback_ddl_burst() + + def run_batch(self): + """Submit the active DML or DDL batch and close batch mode. + + :returns: ``None`` for DDL batches, or the DB-API DML batch result. + :raises django.db.utils.ProgrammingError: If neither a DDL nor a DML + batch is active. + :raises django.db.DatabaseError: If executing the batch fails. + """ + if self._in_dml_batch(): + with self._wrap_spanner_errors(): + return self.connection.run_batch() + if not self.in_ddl_batch: + with self.wrap_database_errors: + raise self.Database.ProgrammingError( + "Cannot run a batch when no DDL or DML batch is active." + ) + try: + return self.flush_ddl_batch() + finally: + self._clear_ddl_batch() + + def abort_batch(self): + """Discard the active DML or DDL batch without sending an RPC. + + :returns: ``None`` when aborting a DDL or DML batch. + :raises django.db.utils.ProgrammingError: If neither a DDL nor a DML + batch is active. + """ + if self._in_dml_batch(): + with self._wrap_spanner_errors(): + return self.connection.abort_batch() + if not self.in_ddl_batch: + with self.wrap_database_errors: + raise self.Database.ProgrammingError( + "Cannot abort a batch when no DDL or DML batch is active." + ) + self._clear_ddl_batch() + return None + + @contextmanager + def ddl_batch(self, *, auto_flush=True): + """Group DDL statements across nested schema editors or cursors into one RPC. + + When outermost, opens a DDL batch on entry, calls :meth:`run_batch` on + clean exit, and aborts on uncaught exceptions. When nested inside an + already-open DDL batch, records a checkpoint for rollback on error and + scopes ``auto_flush`` strictness without flushing until outermost exit. + + :type auto_flush: bool + :param auto_flush: Whether interleaved non-DDL operations flush prior + DDL automatically (cannot weaken an already-strict outer batch). + :raises django.db.DatabaseError: On invalid entry state or DDL failure. + """ + if not self.in_ddl_batch: + self.start_batch_ddl(auto_flush=auto_flush) + try: + yield self + except BaseException: + self._clear_ddl_batch() + raise + self.run_batch() + return + + checkpoint = self.checkpoint_ddl_batch() + prev_auto_flush = self._ddl_batch_auto_flush + self._ddl_batch_auto_flush = prev_auto_flush and bool(auto_flush) + try: + yield self + except BaseException: + self.rollback_ddl_checkpoint(checkpoint) + raise + finally: + if self._ddl_batch_active: + self._ddl_batch_auto_flush = prev_auto_flush + @property def _nodb_connection(self): raise NotImplementedError('Spanner does not have a "no db" connection.') @@ -168,9 +858,9 @@ def get_connection_params(self): def get_new_connection(self, conn_params): """Create a new connection with corresponding connection parameters. - :type conn_params: list - :param conn_params: A List of the connection parameters for - :class:`~google.cloud.spanner_dbapi.connection.Connection` + :type conn_params: dict + :param conn_params: Connection parameters for + :class:`~google.cloud.spanner_dbapi.connection.Connection`. :rtype: :class:`google.cloud.spanner_dbapi.connection.Connection` :returns: A new Spanner DB API Connection object associated with the @@ -182,6 +872,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("auto_batch_ddl", None) # Ensure client is initialized instance = self.instance return self.Database.connect( @@ -190,32 +881,74 @@ def get_new_connection(self, conn_params): **conn_params, ) + def _discard_ddl_batch_on_teardown(self, reason): + """Warn when dropping un-flushed DDL statements and clear DDL batch state.""" + pending = len(self.buffered_ddl_statements) + if pending: + logger.warning( + "Discarding %d un-flushed Cloud Spanner DDL statement(s) on " + "connection %s.", + pending, + reason, + ) + self._clear_ddl_batch() + def init_connection_state(self): - """Initialize the state of the existing connection.""" + """Initialize the state of the existing connection. + + Discards buffered statements when resetting an already-used connection, + while preserving a pre-connect lazy DDL batch unless the newly opened + DB-API connection is read-only. + """ + if self.buffered_ddl_statements: + self._discard_ddl_batch_on_teardown("re-initialization") + else: + self._rollback_ddl_burst() autocommit = self.connection.autocommit self.connection.close() database = self.connection.database - self.connection.__init__(self.instance, database) + read_only = getattr(self.connection, "read_only", False) is True + self.connection.__init__(self.instance, database, read_only=read_only) self.connection.autocommit = autocommit + if self._ddl_batch_active and read_only: + self._clear_ddl_batch() + with self.wrap_database_errors: + raise self.Database.ProgrammingError( + "Cannot start a DDL batch when the connection is in read-only mode." + ) + + def close(self): + """Close the connection and unconditionally discard any open DDL batch.""" + try: + return super().close() + finally: + self._discard_ddl_batch_on_teardown("close") + + def _close(self): + """Abort un-flushed DDL before closing the underlying DB-API connection.""" + self._discard_ddl_batch_on_teardown("close") + return super()._close() def create_cursor(self, name=None): - """Create a new Database cursor. + """Create a new Database cursor wrapped for DDL batch coordination. :type name: str :param name: Currently not used. - :rtype: :class:`~google.cloud.spanner_dbapi.cursor.Cursor` - :returns: The Cursor for this connection. + :rtype: :class:`_DdlBatchCursor` + :returns: DDL-batch-aware cursor proxy around ``self.connection.cursor()``. """ - return self.connection.cursor() + return _DdlBatchCursor(self, self.connection.cursor()) def _set_autocommit(self, autocommit): - """Set the Spanner transaction autocommit flag. + """Set the Spanner transaction autocommit flag on ``self.connection``. :type autocommit: bool :param autocommit: The new value of the autocommit flag. """ - with self.wrap_database_errors: + if not autocommit: + self._before_non_ddl_in_batch() + with self._wrap_spanner_errors(): self.connection.autocommit = autocommit async def _a_set_autocommit(self, autocommit): @@ -241,11 +974,11 @@ def is_usable(self): return True def _start_transaction_under_autocommit(self): - """ - Start a transaction explicitly in autocommit mode. - """ - if self.allow_transactions_in_auto_commit: - self.connection.cursor().execute("BEGIN") - else: - # transactions in autocommit mode for Django 3.2. - self.connection.cursor().execute("SELECT 1") + """Start a transaction explicitly in autocommit mode via ``spanner_dbapi``.""" + self._before_non_ddl_in_batch() + with self._wrap_spanner_errors(): + if self.allow_transactions_in_auto_commit: + self.connection.cursor().execute("BEGIN") + else: + # transactions in autocommit mode for Django 3.2. + self.connection.cursor().execute("SELECT 1") diff --git a/packages/django-google-spanner/django_spanner/operations.py b/packages/django-google-spanner/django_spanner/operations.py index 5a5b8e3e4612..0b3cc3c23149 100644 --- a/packages/django-google-spanner/django_spanner/operations.py +++ b/packages/django-google-spanner/django_spanner/operations.py @@ -10,6 +10,7 @@ from datetime import datetime, time from uuid import UUID +import sqlparse from django.conf import settings from django.db.backends.base.operations import BaseDatabaseOperations from django.db.utils import DatabaseError @@ -99,6 +100,31 @@ def quote_name(self, name): name = name.replace(" ", "_").replace("-", "_") return escape_name(name) + def prepare_sql_script(self, sql): + """Split a raw SQL script into executable non-empty statements. + + Strips SQL comments before checking emptiness so trailing comments and + standalone semicolons in ``RunSQL`` migrations do not emit empty strings + that Cloud Spanner rejects as invalid statements. + + :type sql: str + :param sql: Raw SQL script containing zero or more ``;``-delimited + statements and optional SQL comments. + + :rtype: list[str] + :returns: Ordered list of comment-stripped SQL statements; returns + ``[]`` when ``sql`` contains only whitespace, comments, or + semicolons. + """ + statements = [] + for raw_statement in sqlparse.split(sql): + formatted = sqlparse.format(raw_statement, strip_comments=True).strip() + # Filter comment-only chunks and bare semicolons produced after + # comment stripping so callers receive only executable SQL. + if formatted and formatted.rstrip(";").strip(): + statements.append(formatted) + return statements + def bulk_batch_size(self, fields, objs): """ Override the base class method. Returns the maximum number of objects diff --git a/packages/django-google-spanner/django_spanner/schema.py b/packages/django-google-spanner/django_spanner/schema.py index 15c5b3ce1816..d10b21982162 100644 --- a/packages/django-google-spanner/django_spanner/schema.py +++ b/packages/django-google-spanner/django_spanner/schema.py @@ -13,6 +13,7 @@ from django_spanner import USE_EMULATOR from django_spanner._opentelemetry_tracing import trace_call +from django_spanner.utils import quote_ddl_literal class DatabaseSchemaEditor(BaseDatabaseSchemaEditor): @@ -57,6 +58,68 @@ 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): + """Enter the schema editing context and coordinate DDL batching. + + Automatically opens a connection-scoped auto-flushing DDL batch when + ``connection.auto_batch_ddl`` is enabled outside ``collect_sql`` and + ``atomic`` transaction blocks, or records a rollback checkpoint when + joining a caller-managed outer ``ddl_batch()``. + + :rtype: :class:`DatabaseSchemaEditor` + :returns: ``self`` initialized with an empty ``deferred_sql`` queue. + """ + super().__enter__() + self._owns_ddl_batch = False + self._ddl_batch_checkpoint = None + # Skip starting a DDL batch during sqlmigrate (collect_sql=True) or + # inside atomic blocks so BaseDatabaseSchemaEditor.execute raises + # TransactionManagementError on disallowed transactional DDL. + if not self.collect_sql and not self.connection.in_atomic_block: + if self.connection.in_ddl_batch: + self._ddl_batch_checkpoint = self.connection.checkpoint_ddl_batch() + elif ( + self.connection.auto_batch_ddl and self.connection._can_auto_batch_ddl() + ): + self.connection.start_batch_ddl(auto_flush=True) + self._owns_ddl_batch = True + return self + + def __exit__(self, exc_type, exc_value, traceback): + """Exit the schema editing context, flushing or discarding batched DDL. + + On clean exit, appends ``deferred_sql`` into the active batch before + submitting a single ``UpdateDatabaseDdl`` RPC when this editor started + the batch. On error, discards un-flushed DDL from this editor without + sending an RPC. + + :param exc_type: Exception class raised inside the ``with`` block, or + ``None`` on normal completion. + :param exc_value: Exception instance raised inside the ``with`` block, + or ``None``. + :param traceback: Traceback object for ``exc_value``, or ``None``. + :returns: ``None`` (never suppresses exceptions). + :raises django.db.DatabaseError: If flushing the batched DDL statements + fails on Cloud Spanner. + """ + try: + super().__exit__(exc_type, exc_value, traceback) + except BaseException: + self._discard_editor_ddl_batch() + raise + if exc_type is not None: + self._discard_editor_ddl_batch() + return + if self._owns_ddl_batch and self.connection.in_ddl_batch: + self.connection.run_batch() + + def _discard_editor_ddl_batch(self): + """Abort an owned DDL batch or roll back un-flushed nested DDL statements.""" + if self._owns_ddl_batch and self.connection.in_ddl_batch: + self.connection.abort_batch() + elif self._ddl_batch_checkpoint is not None: + self.connection.rollback_ddl_checkpoint(self._ddl_batch_checkpoint) + def create_model(self, model): """ Create a table and any accompanying indexes or unique constraints for @@ -390,39 +453,22 @@ def column_sql(self, model, field, include_default=False, exclude_not_null=False # Handle GeneratedField if getattr(field, "generated", False): generated_sql, generated_params = field.generated_sql(self.connection) - quoted_params = [] - for p in generated_params: - if isinstance(p, str): - quoted_params.append( - "'%s'" % p.replace("\\", "\\\\").replace("'", "\\'") - ) - elif isinstance(p, bool): - quoted_params.append("TRUE" if p else "FALSE") - elif p is None: - quoted_params.append("NULL") - else: - quoted_params.append(str(p)) - sql += " AS (%s) STORED" % (generated_sql % tuple(quoted_params)) + quoted_params = tuple(quote_ddl_literal(p) for p in generated_params) + sql += " AS (%s) STORED" % ( + generated_sql % quoted_params if quoted_params else generated_sql + ) # Handle db_default db_default = getattr(field, "db_default", None) if db_default is not None and db_default is not NOT_PROVIDED: default_sql, default_params = self.db_default_sql(field) if default_sql: - # Spanner DDL doesn't support parameters, so we must inline them. - quoted_params = [] - for p in default_params: - if isinstance(p, str): - quoted_params.append( - "'%s'" % p.replace("\\", "\\\\").replace("'", "\\'") - ) - elif isinstance(p, bool): - quoted_params.append("TRUE" if p else "FALSE") - elif p is None: - quoted_params.append("NULL") - else: - quoted_params.append(str(p)) - sql += " DEFAULT (%s)" % (default_sql % tuple(quoted_params)) + # Spanner DDL doesn't support parameters, so inline literals directly. + quoted_params = tuple(quote_ddl_literal(p) for p in default_params) + rendered_default = ( + default_sql % quoted_params if quoted_params else default_sql + ) + sql += " DEFAULT (%s)" % rendered_default # Return the sql return sql, params @@ -455,12 +501,8 @@ def add_index(self, model, index): super().add_index(model, index) def quote_value(self, value): - # A more complete implementation isn't currently required. - if isinstance(value, str): - return "'%s'" % value.replace("'", "''") - if isinstance(value, bool): - return "TRUE" if value else "FALSE" - return str(value) + """Quote a Python scalar as a Cloud Spanner GoogleSQL literal for DDL/SQL.""" + return quote_ddl_literal(value) def prepare_default(self, value): return self.quote_value(value) diff --git a/packages/django-google-spanner/django_spanner/utils.py b/packages/django-google-spanner/django_spanner/utils.py index 877de815b057..3800a3023f22 100644 --- a/packages/django-google-spanner/django_spanner/utils.py +++ b/packages/django-google-spanner/django_spanner/utils.py @@ -4,9 +4,22 @@ # license that can be found in the LICENSE file or at # https://developers.google.com/open-source/licenses/bsd +import datetime +import math +import re +import uuid +from decimal import Decimal + import django import sqlparse +from django.conf import settings from django.core.exceptions import ImproperlyConfigured +from django.utils import timezone +from google.cloud.spanner_dbapi.parse_utils import DateStr, TimestampStr +from google.cloud.spanner_v1 import JsonObject + +_RE_CONTROL_CHARS = re.compile(r"[\x00-\x1f\x7f]") +_NAMED_CONTROL_ESCAPES = {"\n": r"\n", "\r": r"\r", "\t": r"\t"} def check_django_compatability(supported_django_versions): @@ -40,3 +53,94 @@ def add_dummy_where(sql): if any(isinstance(token, sqlparse.sql.Where) for token in sqlparse.parse(sql)[0]): return sql return sql + " WHERE 1=1" + + +def _escape_control_char(match): + ch = match.group(0) + return _NAMED_CONTROL_ESCAPES.get(ch, f"\\x{ord(ch):02x}") + + +def _escape_googlesql_string(text): + """Escape backslashes, single quotes, and ASCII control chars for GoogleSQL.""" + escaped = text.replace("\\", "\\\\").replace("'", "\\'") + return _RE_CONTROL_CHARS.sub(_escape_control_char, escaped) + + +def quote_ddl_literal(value): + """Format a Python scalar as an inline Cloud Spanner GoogleSQL DDL literal. + + Cloud Spanner ``UpdateDatabaseDdl`` statements do not accept bind + parameters, and GoogleSQL string literals escape quotes with backslashes + rather than SQL-standard doubled single quotes. Naive :class:`datetime.datetime` + values convert from ``settings.TIME_ZONE`` to UTC when ``settings.USE_TZ`` + is active, matching :meth:`DatabaseOperations.adapt_datetimefield_value`. + + :type value: object + :param value: Scalar value to inline into a DDL expression (``None``, + ``bool``, ``int``, ``float``, ``Decimal``, ``UUID``, + ``JsonObject``, ``TimestampStr``, ``DateStr``, ``str``, + ``bytes``/``bytearray``/``memoryview``, + ``datetime.datetime``, ``datetime.date``, or + ``datetime.time``). + + :rtype: str + :returns: A self-contained GoogleSQL literal token safe to embed inside a + DDL statement; ``None`` renders as ``"NULL"`` and ``UUID`` + renders as its 32-character hex ``STRING(32)`` literal. + + :raises ValueError: If ``value`` is a non-finite :class:`~decimal.Decimal` + (``NaN`` or ``Infinity``), which Spanner ``NUMERIC`` + rejects. + :raises TypeError: If ``value`` is not a supported scalar DDL literal type. + """ + # Check bool before int, TimestampStr/DateStr before str, and datetime + # before date to preserve subclass specificity in Python's type hierarchy. + if isinstance(value, bool): + return "TRUE" if value else "FALSE" + if value is None: + return "NULL" + if isinstance(value, int): + return f"{int(value):d}" + if isinstance(value, TimestampStr): + return f"TIMESTAMP '{_escape_googlesql_string(str(value))}'" + if isinstance(value, DateStr): + return f"DATE '{_escape_googlesql_string(str(value))}'" + if isinstance(value, str): + return f"'{_escape_googlesql_string(value)}'" + if isinstance(value, uuid.UUID): + return f"'{value.hex}'" + if isinstance(value, JsonObject): + serialized = value.serialize() + if serialized is None: + return "NULL" + return f"JSON '{_escape_googlesql_string(serialized)}'" + if isinstance(value, (bytes, bytearray, memoryview)): + hex_escapes = "".join(f"\\x{byte:02x}" for byte in bytes(value)) + return f"b'{hex_escapes}'" + if isinstance(value, datetime.datetime): + if timezone.is_naive(value) and settings.configured and settings.USE_TZ: + value = timezone.make_aware(value, timezone.get_default_timezone()) + if timezone.is_aware(value): + value = value.astimezone(datetime.timezone.utc).replace(tzinfo=None) + return f"TIMESTAMP '{value.isoformat()}Z'" + if isinstance(value, datetime.date): + return f"DATE '{value.isoformat()}'" + if isinstance(value, datetime.time): + # TimeField maps to Spanner TIMESTAMP prefixed with dummy date 0001-01-01. + return f"TIMESTAMP '0001-01-01T{value.isoformat(timespec='microseconds')}Z'" + if isinstance(value, Decimal): + if not value.is_finite(): + raise ValueError( + f"Spanner NUMERIC DDL literals must be finite, got {value!r}." + ) + return f"NUMERIC '{format(value, 'f')}'" + if isinstance(value, float): + if math.isnan(value): + return "CAST('NaN' AS FLOAT64)" + if math.isinf(value): + sign = "-" if value < 0 else "" + return f"CAST('{sign}Infinity' AS FLOAT64)" + return repr(float(value)) + raise TypeError( + f"Unsupported Cloud Spanner DDL literal type: {type(value).__name__}." + ) diff --git a/packages/django-google-spanner/docs/limitations-spanner.rst b/packages/django-google-spanner/docs/limitations-spanner.rst index 75cbbebaf6b3..36a6a660d75f 100644 --- a/packages/django-google-spanner/docs/limitations-spanner.rst +++ b/packages/django-google-spanner/docs/limitations-spanner.rst @@ -18,12 +18,6 @@ Mutations Mutations are not a concept that is understood by Django so it does not support it inherently. The workaround is to use python-spanner objects and run Mutations using that. -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. - Stale Reads ~~~~~~~~~~~ Stale Reads are not a concept that is understood by Django so it does not support @@ -90,8 +84,6 @@ Workaround - Raw Queries can be used for Query hints `Link `__. * - 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 `__. * - Stale Reads - Python-spanner database objects can be used to perform stale reads. `Link `__. * - Interleaving diff --git a/packages/django-google-spanner/tests/mockserver_tests/mock_database_admin.py b/packages/django-google-spanner/tests/mockserver_tests/mock_database_admin.py index e18d4c3a334d..160411e7aba7 100644 --- a/packages/django-google-spanner/tests/mockserver_tests/mock_database_admin.py +++ b/packages/django-google-spanner/tests/mockserver_tests/mock_database_admin.py @@ -12,6 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import threading +from collections import deque + from google.longrunning import operations_pb2 as operations_pb2 from google.protobuf import empty_pb2 @@ -20,18 +23,38 @@ # An in-memory mock DatabaseAdmin server that can be used for testing. class DatabaseAdminServicer(database_admin_grpc.DatabaseAdminServicer): + """Thread-safe in-memory Spanner DatabaseAdmin servicer for mock tests.""" + def __init__(self): + self._lock = threading.Lock() self._requests = [] + self._errors = deque() @property def requests(self): - return self._requests + """Return a snapshot list of recorded UpdateDatabaseDdl requests.""" + with self._lock: + return list(self._requests) + + def add_error(self, status_code, details="Injected DDL failure"): + """Queue a gRPC error to abort the next UpdateDatabaseDdl call.""" + with self._lock: + self._errors.append((status_code, details)) def clear_requests(self): - self._requests = [] + """Clear recorded UpdateDatabaseDdl requests and queued gRPC errors.""" + with self._lock: + self._requests = [] + self._errors.clear() def UpdateDatabaseDdl(self, request, context): - self._requests.append(request) + """Record an UpdateDatabaseDdl RPC or abort if an error was queued.""" + with self._lock: + self._requests.append(request) + queued_error = self._errors.popleft() if self._errors else None + if queued_error is not None: + status_code, details = queued_error + context.abort(status_code, details) operation = operations_pb2.Operation() operation.done = True operation.name = "projects/test-project/operations/test-operation" diff --git a/packages/django-google-spanner/tests/mockserver_tests/test_ddl_batching.py b/packages/django-google-spanner/tests/mockserver_tests/test_ddl_batching.py new file mode 100644 index 000000000000..78afb49aced0 --- /dev/null +++ b/packages/django-google-spanner/tests/mockserver_tests/test_ddl_batching.py @@ -0,0 +1,774 @@ +# Copyright 2026 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. + +import threading +from contextlib import contextmanager +from unittest import mock + +import google.cloud.spanner_v1.types.result_set as result_set +import google.cloud.spanner_v1.types.type as spanner_type +import grpc +from django.db import ( + IntegrityError, + InterfaceError, + OperationalError, + ProgrammingError, + connection, + migrations, + models, + transaction, +) +from django.db.migrations.executor import MigrationExecutor +from django.db.migrations.state import ProjectState +from django.test.utils import CaptureQueriesContext +from google.cloud.spanner_v1 import TypeCode + +from django_spanner.base import DatabaseWrapper +from tests.mockserver_tests.mock_server_test_base import ( + MockServerTestBase, + add_result, + add_select1_result, + add_update_count, +) + + +def _add_empty_result(sql: str): + """Register an empty ResultSet on the mock Spanner service for ``sql``.""" + res = result_set.ResultSet( + dict( + metadata=result_set.ResultSetMetadata( + dict( + row_type=spanner_type.StructType( + dict( + fields=[ + spanner_type.StructType.Field( + dict( + name="c", + type=spanner_type.Type( + dict(code=TypeCode.STRING) + ), + ) + ) + ] + ) + ) + ) + ) + ) + ) + add_result(sql, res) + + +class TestDdlBatching(MockServerTestBase): + """End-to-end MockServer tests for Django Spanner DDL batching.""" + + def _ddl_batches(self): + return [list(req.statements) for req in self.database_admin_service.requests] + + @contextmanager + def _auto_batch_ddl(self, enabled=True): + had_key = "AUTO_BATCH_DDL" in connection.settings_dict + prev = connection.settings_dict.get("AUTO_BATCH_DDL") + connection.settings_dict["AUTO_BATCH_DDL"] = enabled + try: + yield + finally: + if had_key: + connection.settings_dict["AUTO_BATCH_DDL"] = prev + else: + connection.settings_dict.pop("AUTO_BATCH_DDL", None) + + def test_multi_model_migration_batches_all_ddl_into_single_rpc(self): + mig = migrations.Migration("0001_initial", "batch_app") + mig.operations = [ + migrations.CreateModel( + name="Author", + fields=[ + ("id", models.BigAutoField(primary_key=True)), + ("email", models.CharField(max_length=120, unique=True)), + ("name", models.CharField(max_length=80, db_index=True)), + ], + ), + migrations.CreateModel( + name="Tag", + fields=[ + ("id", models.BigAutoField(primary_key=True)), + ("slug", models.SlugField(unique=True)), + ], + ), + migrations.CreateModel( + name="Book", + fields=[ + ("id", models.BigAutoField(primary_key=True)), + ("title", models.CharField(max_length=200, db_index=True)), + ( + "author", + models.ForeignKey("batch_app.Author", on_delete=models.CASCADE), + ), + ("tags", models.ManyToManyField("batch_app.Tag")), + ], + ), + ] + + state = ProjectState() + with self._auto_batch_ddl(True): + with CaptureQueriesContext(connection) as captured: + with connection.schema_editor(atomic=mig.atomic) as editor: + mig.apply(state, editor) + + batches = self._ddl_batches() + self.assertEqual(len(batches), 1) + stmts = batches[0] + self.assertGreaterEqual(len(stmts), 10) + # All CREATE TABLE statements execute before deferred indexes and FKs. + first_index_or_fk = next( + i for i, s in enumerate(stmts) if not s.startswith("CREATE TABLE ") + ) + for s in stmts[:first_index_or_fk]: + self.assertTrue(s.startswith("CREATE TABLE ")) + self.assertTrue(all(not s.endswith(";") for s in stmts)) + # CursorDebugWrapper continues to observe every individual DDL statement. + self.assertEqual(len(captured), len(stmts)) + self.assertFalse(connection.in_ddl_batch) + self.assertEqual(connection.buffered_ddl_statements, ()) + + def test_auto_batch_ddl_false_executes_each_ddl_individually(self): + mig = migrations.Migration("0001_no_batch", "nobatch_app") + mig.operations = [ + migrations.CreateModel( + name="Item", + fields=[ + ("id", models.BigAutoField(primary_key=True)), + ("code", models.CharField(max_length=40, unique=True)), + ("name", models.CharField(max_length=40, db_index=True)), + ], + ), + ] + # 1. Default (unset) AUTO_BATCH_DDL is False: each DDL statement runs individually. + self.assertNotIn("AUTO_BATCH_DDL", connection.settings_dict) + self.assertFalse(connection.auto_batch_ddl) + with connection.schema_editor(atomic=mig.atomic) as editor: + mig.apply(ProjectState(), editor) + batches = self._ddl_batches() + self.assertEqual(len(batches), 3) + self.assertTrue(all(len(b) == 1 for b in batches)) + + # 2. OPTIONS["auto_batch_ddl"] = True opts in; top-level AUTO_BATCH_DDL = False overrides it. + self.database_admin_service.clear_requests() + connection.settings_dict["OPTIONS"]["auto_batch_ddl"] = True + try: + self.assertTrue(connection.auto_batch_ddl) + with connection.schema_editor(atomic=mig.atomic) as editor: + mig.apply(ProjectState(), editor) + self.assertEqual(len(self._ddl_batches()), 1) + self.assertEqual(len(self._ddl_batches()[0]), 3) + + self.database_admin_service.clear_requests() + with self._auto_batch_ddl(False): + self.assertFalse(connection.auto_batch_ddl) + with connection.schema_editor(atomic=mig.atomic) as editor: + mig.apply(ProjectState(), editor) + batches = self._ddl_batches() + self.assertEqual(len(batches), 3) + self.assertTrue(all(len(b) == 1 for b in batches)) + finally: + connection.settings_dict["OPTIONS"].pop("auto_batch_ddl", None) + + def test_apply_migration_without_deferred_sql_flushes_before_recording(self): + mig = migrations.Migration("0001_simple", "rec_app") + mig.operations = [ + migrations.CreateModel( + name="Simple", + fields=[ + ("id", models.BigAutoField(primary_key=True)), + ("label", models.CharField(max_length=50)), + ], + ), + migrations.CreateModel( + name="SimpleTwo", + fields=[ + ("id", models.BigAutoField(primary_key=True)), + ("label", models.CharField(max_length=50)), + ], + ), + ] + + executor = MigrationExecutor.__new__(MigrationExecutor) + executor.connection = connection + executor.progress_callback = None + ddl_count_at_record = [] + + def fake_record(*args): + # Django 5.2 passes (migration,); Django 6.0 passes (app_label, name). + app_label = args[0].app_label if len(args) == 1 else args[0] + mig_name = args[0].name if len(args) == 1 else args[1] + # Simulate MigrationRecorder.record_applied issuing DML via cursor + # inside the open SchemaEditor block (when deferred_sql is empty). + add_update_count( + "INSERT INTO django_migrations (app, name) VALUES (@a0, @a1)", 1 + ) + with connection.cursor() as cur: + cur.execute( + "INSERT INTO django_migrations (app, name) VALUES (%s, %s)", + [app_label, mig_name], + ) + ddl_count_at_record.append(len(self._ddl_batches())) + + executor.record_migration = fake_record + with self._auto_batch_ddl(True): + executor.apply_migration(ProjectState(), mig) + + # Both CREATE TABLEs flushed in 1 batch BEFORE django_migrations INSERT! + self.assertEqual(ddl_count_at_record, [1]) + self.assertEqual(len(self._ddl_batches()), 1) + self.assertEqual(len(self._ddl_batches()[0]), 2) + + def test_apply_migration_ddl_failure_prevents_recording_migration(self): + self.database_admin_service.add_error( + grpc.StatusCode.FAILED_PRECONDITION, + "Duplicate column or precondition failed", + ) + mig = migrations.Migration("0001_fail", "fail_app") + mig.operations = [ + migrations.CreateModel( + name="Broken", + fields=[("id", models.BigAutoField(primary_key=True))], + ), + ] + executor = MigrationExecutor.__new__(MigrationExecutor) + executor.connection = connection + executor.progress_callback = None + recorded = [] + + def fake_record(*args): + app_label = args[0].app_label if len(args) == 1 else args[0] + mig_name = args[0].name if len(args) == 1 else args[1] + with connection.cursor() as cur: + cur.execute( + "INSERT INTO django_migrations (app, name) VALUES (%s, %s)", + [app_label, mig_name], + ) + recorded.append(mig_name) + + executor.record_migration = fake_record + with self._auto_batch_ddl(True): + with self.assertRaises(IntegrityError): + executor.apply_migration(ProjectState(), mig) + + # DDL batch flushed before INSERT, raised IntegrityError, never recorded! + self.assertEqual(recorded, []) + self.assertFalse(connection.in_ddl_batch) + self.assertEqual(connection.buffered_ddl_statements, ()) + + # Subsequent DDL on the same connection succeeds cleanly with zero stale leak. + with connection.schema_editor() as editor: + editor.execute( + "CREATE TABLE clean_after_err (id INT64) PRIMARY KEY (id)" + ) + batches = self._ddl_batches() + self.assertEqual(len(batches), 2) + self.assertEqual( + batches[1], ["CREATE TABLE clean_after_err (id INT64) PRIMARY KEY (id)"] + ) + + def test_delete_model_remove_field_and_alter_field_batch_index_drops(self): + state = ProjectState() + create_mig = migrations.Migration("0001_init", "edit_app") + create_mig.operations = [ + migrations.CreateModel( + name="Article", + fields=[ + ("id", models.BigAutoField(primary_key=True)), + ("title", models.CharField(max_length=100, db_index=True)), + ("code", models.CharField(max_length=40, db_index=True)), + ], + ) + ] + state = create_mig.mutate_state(state, preserve=True) + + with self._auto_batch_ddl(True): + # 1. AlterField changing nullability of indexed column drops index, + # alters column, and recreates index in ONE UpdateDatabaseDdl RPC. + alter_mig = migrations.Migration("0002_alter", "edit_app") + alter_mig.operations = [ + migrations.AlterField( + model_name="article", + name="title", + field=models.CharField(max_length=100, null=True, db_index=True), + ) + ] + with mock.patch.object( + connection.SchemaEditorClass, + "_constraint_names", + return_value=["article_title_idx"], + ): + with connection.schema_editor() as editor: + state = alter_mig.apply(state, editor) + + self.assertEqual(len(self._ddl_batches()), 1) + self.assertEqual(len(self._ddl_batches()[0]), 3) + self.assertTrue(self._ddl_batches()[0][0].startswith("DROP INDEX")) + self.assertTrue(self._ddl_batches()[0][1].startswith("ALTER TABLE")) + self.assertTrue(self._ddl_batches()[0][2].startswith("CREATE INDEX")) + + # 2. RemoveField with index drops index + column in ONE RPC. + self.database_admin_service.clear_requests() + rem_mig = migrations.Migration("0003_rem", "edit_app") + rem_mig.operations = [ + migrations.RemoveField(model_name="article", name="code"), + ] + with mock.patch.object( + connection.SchemaEditorClass, + "_constraint_names", + side_effect=[["article_code_idx"], []], + ): + with connection.schema_editor() as editor: + state = rem_mig.apply(state, editor) + + self.assertEqual(len(self._ddl_batches()), 1) + self.assertEqual( + self._ddl_batches()[0], + [ + "DROP INDEX article_code_idx", + "ALTER TABLE edit_app_article DROP COLUMN code", + ], + ) + + # 3. DeleteModel with multiple indexes drops all indexes + table in ONE RPC. + self.database_admin_service.clear_requests() + del_mig = migrations.Migration("0004_del", "edit_app") + del_mig.operations = [migrations.DeleteModel(name="Article")] + with mock.patch.object( + connection.SchemaEditorClass, + "_constraint_names", + return_value=["idx_a", "idx_b"], + ): + with connection.schema_editor() as editor: + del_mig.apply(state, editor) + + self.assertEqual(len(self._ddl_batches()), 1) + self.assertEqual( + self._ddl_batches()[0], + [ + "DROP INDEX idx_a", + "DROP INDEX idx_b", + "DROP TABLE edit_app_article", + ], + ) + + def test_add_field_with_default_flushes_before_update_and_batches_rest(self): + state = ProjectState() + init_mig = migrations.Migration("0001_init", "def_app") + init_mig.operations = [ + migrations.CreateModel( + name="Member", + fields=[("id", models.BigAutoField(primary_key=True))], + ) + ] + state = init_mig.mutate_state(state, preserve=True) + + add_update_count("UPDATE def_app_member SET score=@a0 WHERE 1=1", 2) + # Note: SchemaEditor.add_field emits "UPDATE SET =%s" + add_update_count("UPDATE def_app_member SET score=@a0", 2) + + mig = migrations.Migration("0002_add_score", "def_app") + mig.operations = [ + migrations.AddField( + model_name="member", + name="score", + field=models.IntegerField(default=10, db_index=True), + ), + migrations.AddField( + model_name="member", + name="nickname", + field=models.CharField(max_length=30, null=True, unique=True), + ), + ] + with self._auto_batch_ddl(True): + with connection.schema_editor() as editor: + mig.apply(state, editor) + + batches = self._ddl_batches() + # Batch 1: ADD COLUMN score (flushed automatically before UPDATE DML). + # Batch 2: ALTER COLUMN score NOT NULL + ADD COLUMN nickname + + # CREATE INDEX score + CREATE UNIQUE INDEX nickname! + self.assertEqual(len(batches), 2) + self.assertEqual(len(batches[0]), 1) + self.assertIn("ADD COLUMN score INT64", batches[0][0]) + self.assertEqual(len(batches[1]), 4) + self.assertIn("ALTER COLUMN score INT64 NOT NULL", batches[1][0]) + self.assertIn("ADD COLUMN nickname STRING(30)", batches[1][1]) + + def test_runpython_atomic_and_introspection_flush_prior_ddl(self): + batches_seen_in_python = [] + add_select1_result() + + def forward_fn(apps, schema_editor): + # Entering RunPython (with atomic=True on Migration) has ALREADY + # flushed CREATE TABLE prior to BEGIN! + batches_seen_in_python.append(len(self._ddl_batches())) + with schema_editor.connection.cursor() as cur: + cur.execute("select 1") + cur.fetchall() + + mig = migrations.Migration("0001_rp", "rp_app") + mig.operations = [ + migrations.CreateModel( + name="StageOne", + fields=[("id", models.BigAutoField(primary_key=True))], + ), + migrations.RunPython(forward_fn), + migrations.CreateModel( + name="StageTwo", + fields=[ + ("id", models.BigAutoField(primary_key=True)), + ("code", models.CharField(max_length=20, db_index=True)), + ], + ), + ] + with self._auto_batch_ddl(True): + with connection.schema_editor(atomic=mig.atomic) as editor: + mig.apply(ProjectState(), editor) + + self.assertEqual(batches_seen_in_python, [1]) + batches = self._ddl_batches() + self.assertEqual(len(batches), 2) + self.assertEqual(len(batches[0]), 1) # StageOne CREATE TABLE + self.assertEqual(len(batches[1]), 2) # StageTwo CREATE TABLE + index + + # Also verify cursor.run_sql_in_snapshot / list_tables auto-flushes buffered DDL. + self.database_admin_service.clear_requests() + _add_empty_result("SELECT table_name FROM information_schema.tables") + with connection.ddl_batch(): + with connection.cursor() as cur: + cur.execute("CREATE TABLE snap_t (id INT64) PRIMARY KEY (id)") + self.assertEqual(len(self._ddl_batches()), 0) + cur.run_sql_in_snapshot( + "SELECT table_name FROM information_schema.tables" + ) + self.assertEqual(len(self._ddl_batches()), 1) + + def test_runsql_comments_multi_statement_and_parameterized_ddl(self): + mig = migrations.Migration("0001_sql", "sql_app") + mig.operations = [ + migrations.RunSQL( + """ + -- Create two tables in a raw SQL script + CREATE TABLE raw_a (id INT64 NOT NULL) PRIMARY KEY (id); + /* second table */ + CREATE TABLE raw_b (id INT64 NOT NULL, a_id INT64) PRIMARY KEY (id); + -- trailing comment only + """ + ), + migrations.RunSQL( + [ + ( + "ALTER TABLE raw_a ADD COLUMN label STRING(50) DEFAULT (%s)", + ["it's \\ ok"], + ) + ] + ), + ] + with self._auto_batch_ddl(True): + with connection.schema_editor() as editor: + mig.apply(ProjectState(), editor) + + self.assertEqual( + self._ddl_batches(), + [ + [ + "CREATE TABLE raw_a (id INT64 NOT NULL) PRIMARY KEY (id)", + "CREATE TABLE raw_b (id INT64 NOT NULL, a_id INT64) PRIMARY KEY (id)", + r"ALTER TABLE raw_a ADD COLUMN label STRING(50) DEFAULT ('it\'s \\ ok')", + ] + ], + ) + + def test_explicit_ddl_batch_groups_multiple_schema_editors_and_rollbacks(self): + # 1. Explicit ddl_batch() remains active under default AUTO_BATCH_DDL=False, + # grouping multiple schema_editor() contexts with nested checkpoint rollback. + self.assertFalse(connection.auto_batch_ddl) + with connection.ddl_batch(): + with connection.schema_editor() as ed1: + ed1.execute("CREATE TABLE outer_1 (id INT64) PRIMARY KEY (id)") + try: + with connection.schema_editor() as ed_fail: + ed_fail.execute( + "CREATE TABLE discarded_inner (id INT64) PRIMARY KEY (id)" + ) + raise RuntimeError("abort inner editor only") + except RuntimeError: + pass + with connection.ddl_batch(): + with connection.schema_editor() as ed2: + ed2.execute("CREATE TABLE outer_2 (id INT64) PRIMARY KEY (id)") + self.assertEqual(len(self._ddl_batches()), 0) + + self.assertEqual( + self._ddl_batches(), + [ + [ + "CREATE TABLE outer_1 (id INT64) PRIMARY KEY (id)", + "CREATE TABLE outer_2 (id INT64) PRIMARY KEY (id)", + ] + ], + ) + + # 2. Uncaught exception inside auto-batched schema_editor() or explicit ddl_batch() sends 0 RPCs. + self.database_admin_service.clear_requests() + with self._auto_batch_ddl(True): + with self.assertRaises(ValueError): + with connection.schema_editor() as editor: + editor.execute( + "CREATE TABLE never_sent (id INT64) PRIMARY KEY (id)" + ) + raise ValueError("boom") + self.assertEqual(self._ddl_batches(), []) + self.assertFalse(connection.in_ddl_batch) + + # 3. Invalid deferred_sql failing inside SchemaEditor.__exit__ aborts owned batch. + with self.assertRaises(ProgrammingError): + with connection.schema_editor() as editor: + editor.execute( + "CREATE TABLE never_sent_3 (id INT64) PRIMARY KEY (id)" + ) + editor.deferred_sql.append(";") + self.assertEqual(self._ddl_batches(), []) + self.assertFalse(connection.in_ddl_batch) + + # Empty auto-batched schema_editor() sends 0 RPCs. + with connection.schema_editor(): + pass + self.assertEqual(self._ddl_batches(), []) + + with self.assertRaises(ValueError): + with connection.ddl_batch(): + with connection.cursor() as cur: + cur.execute("CREATE TABLE never_sent_2 (id INT64) PRIMARY KEY (id)") + raise ValueError("boom outer") + self.assertEqual(self._ddl_batches(), []) + self.assertFalse(connection.in_ddl_batch) + + # 4. Empty ddl_batch() and default schema_editor() send 0 RPCs. + with connection.ddl_batch(): + pass + with connection.schema_editor(): + pass + self.assertEqual(self._ddl_batches(), []) + + # 5. ddl_batch() inside transaction.atomic() raises ProgrammingError. + with transaction.atomic(): + with self.assertRaises(ProgrammingError): + with connection.ddl_batch(): + pass + + def test_client_side_start_run_abort_batch_ddl_and_negative_cases(self): + add_select1_result() + with connection.cursor() as cur: + # First run a DQL query on the same cursor to verify buffered DDL + # resets cursor description/rowcount cleanly. + cur.execute("select 1") + self.assertIsNotNone(cur.description) + + cur.execute("/* header */ START BATCH DDL;") + self.assertTrue(connection.in_ddl_batch) + self.assertFalse(connection.ddl_batch_auto_flush) + self.assertIsNone(cur.description) + + cur.execute( + "CREATE TABLE cs_1 (id INT64) PRIMARY KEY (id); " + "CREATE INDEX cs_1_idx ON cs_1 (id); -- comment" + ) + cur.execute("ALTER TABLE cs_1 ADD COLUMN name STRING(50)") + self.assertEqual(len(self._ddl_batches()), 0) + + # Strict mode rejects DQL, DML, nested START BATCH DDL, START BATCH DML, + # empty/invalid SQL, parameterized batch control, and executemany + # without dropping already-queued valid DDL statements! + for bad_sql in ( + "SELECT 1", + "INSERT INTO cs_1 (id) VALUES (1)", + "START BATCH DDL", + "START BATCH DML", + "", + " ; ; ", + "-- comment only", + "CREATE TABLE a (id INT64) PRIMARY KEY (id); SELECT 1", + ): + with self.assertRaises(ProgrammingError): + cur.execute(bad_sql) + with self.assertRaises(ProgrammingError): + cur.execute("RUN BATCH", [1]) + with self.assertRaises(ProgrammingError): + cur.execute("ABORT BATCH", [1]) + with self.assertRaises(ProgrammingError): + cur.executemany("INSERT INTO cs_1 (id) VALUES (%s)", [(1,), (2,)]) + with self.assertRaises(ProgrammingError): + cur.run_sql_in_snapshot("SELECT 1") + with self.assertRaises(ProgrammingError): + cur.list_tables() + + # Valid queued DDL executes atomically on RUN BATCH. + cur.execute("run batch") + self.assertFalse(connection.in_ddl_batch) + + self.assertEqual( + self._ddl_batches(), + [ + [ + "CREATE TABLE cs_1 (id INT64) PRIMARY KEY (id)", + "CREATE INDEX cs_1_idx ON cs_1 (id)", + "ALTER TABLE cs_1 ADD COLUMN name STRING(50)", + ] + ], + ) + + # ABORT BATCH discards buffered statements with 0 RPCs. + self.database_admin_service.clear_requests() + with connection.cursor() as cur: + with self.assertRaises(ProgrammingError): + cur.execute("START BATCH DDL", [1]) + cur.execute("START BATCH DDL") + cur.execute("CREATE TABLE discarded_cs (id INT64) PRIMARY KEY (id)") + cur.execute("ABORT BATCH;") + self.assertFalse(connection.in_ddl_batch) + # RUN / ABORT when no batch is active raises ProgrammingError. + with self.assertRaises(ProgrammingError): + cur.execute("RUN BATCH") + with self.assertRaises(ProgrammingError): + cur.execute("ABORT BATCH") + self.assertEqual(self._ddl_batches(), []) + + # Closed cursor raises InterfaceError even for START BATCH DDL / DDL batch. + closed_cur = connection.create_cursor() + closed_cur.close() + with self.assertRaises(InterfaceError): + closed_cur.execute("START BATCH DDL") + + # Server gRPC INVALID_ARGUMENT and INTERNAL map to ProgrammingError and OperationalError. + self.database_admin_service.add_error( + grpc.StatusCode.INVALID_ARGUMENT, "Syntax error in DDL" + ) + with self.assertRaises(ProgrammingError): + with connection.ddl_batch(): + with connection.cursor() as cur: + cur.execute("CREATE TABLE bad_syn (id INT64) PRIMARY KEY (id)") + + self.database_admin_service.add_error( + grpc.StatusCode.INTERNAL, "Internal DDL error" + ) + with self.assertRaises(OperationalError): + with connection.ddl_batch(): + with connection.cursor() as cur: + cur.execute("CREATE TABLE int_err (id INT64) PRIMARY KEY (id)") + + def test_close_discards_unflushed_ddl_and_mixed_multistatement_auto_flush(self): + add_update_count("INSERT INTO mix_t (id) VALUES (1)", 1) + connection.start_batch_ddl() + with connection.cursor() as cur: + cur.execute("CREATE TABLE drop_on_close (id INT64) PRIMARY KEY (id)") + with self.assertLogs("django_spanner.base", level="WARNING") as cm: + connection.close() + self.assertEqual(len(cm.output), 1) + self.assertIn("1 un-flushed Cloud Spanner DDL statement(s)", cm.output[0]) + self.assertFalse(connection.in_ddl_batch) + self.assertEqual(connection.buffered_ddl_statements, ()) + self.assertEqual(self._ddl_batches(), []) + + # Reconnects cleanly and handles mixed multi-statement DDL + DML in auto_flush mode. + with connection.ddl_batch(auto_flush=True): + with connection.cursor() as cur: + cur.execute( + "CREATE TABLE mix_t (id INT64) PRIMARY KEY (id); " + "INSERT INTO mix_t (id) VALUES (1); " + "CREATE INDEX mix_idx ON mix_t (id)" + ) + self.assertEqual( + self._ddl_batches(), + [ + ["CREATE TABLE mix_t (id INT64) PRIMARY KEY (id)"], + ["CREATE INDEX mix_idx ON mix_t (id)"], + ], + ) + + def test_concurrent_connections_batch_ddl_independently(self): + barrier = threading.Barrier(2, timeout=10) + errors = [] + + def make_wrapper(alias): + from google.cloud.spanner_v1 import PingingPool + + cfg = dict(connection.settings_dict) + cfg["OPTIONS"] = dict(connection.settings_dict["OPTIONS"]) + # Each DatabaseWrapper receives its own session pool because + # spanner_v1 PingingPool binds strictly to a single Database object. + cfg["OPTIONS"]["pool"] = PingingPool(size=5) + return DatabaseWrapper(cfg, alias=alias) + + def worker_commit(): + conn_a = make_wrapper("worker_a") + try: + with conn_a.ddl_batch(): + with conn_a.cursor() as cur: + cur.execute( + "CREATE TABLE thread_a1 (id INT64) PRIMARY KEY (id)" + ) + barrier.wait() + cur.execute( + "CREATE TABLE thread_a2 (id INT64) PRIMARY KEY (id)" + ) + except Exception as exc: + errors.append(exc) + finally: + conn_a.close() + + def worker_abort_then_commit(): + conn_b = make_wrapper("worker_b") + try: + conn_b.start_batch_ddl() + with conn_b.cursor() as cur: + cur.execute( + "CREATE TABLE thread_b_aborted (id INT64) PRIMARY KEY (id)" + ) + barrier.wait() + conn_b.abort_batch() + with conn_b.ddl_batch(): + with conn_b.cursor() as cur: + cur.execute( + "CREATE TABLE thread_b_kept (id INT64) PRIMARY KEY (id)" + ) + except Exception as exc: + errors.append(exc) + finally: + conn_b.close() + + t1 = threading.Thread(target=worker_commit) + t2 = threading.Thread(target=worker_abort_then_commit) + t1.start() + t2.start() + t1.join(timeout=15) + t2.join(timeout=15) + + self.assertEqual(errors, []) + batches = sorted(self._ddl_batches()) + self.assertEqual( + batches, + [ + [ + "CREATE TABLE thread_a1 (id INT64) PRIMARY KEY (id)", + "CREATE TABLE thread_a2 (id INT64) PRIMARY KEY (id)", + ], + ["CREATE TABLE thread_b_kept (id INT64) PRIMARY KEY (id)"], + ], + ) 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..dcf08a3d0407 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 @@ -99,3 +99,582 @@ def test_start_transaction_under_autocommit(self): mock_connection.cursor = mock_cursor = mock.MagicMock() self.db_wrapper._start_transaction_under_autocommit() mock_cursor.assert_called_once_with() + + def test_auto_batch_ddl_default_and_settings_precedence(self): + from django.core.exceptions import ImproperlyConfigured + + self.assertFalse(self.db_wrapper.auto_batch_ddl) + + self.db_wrapper.settings_dict["OPTIONS"]["auto_batch_ddl"] = True + self.assertTrue(self.db_wrapper.auto_batch_ddl) + + # Top-level AUTO_BATCH_DDL takes precedence over OPTIONS["auto_batch_ddl"]. + self.db_wrapper.settings_dict["AUTO_BATCH_DDL"] = "FALSE" + self.assertFalse(self.db_wrapper.auto_batch_ddl) + + self.db_wrapper.settings_dict["OPTIONS"]["auto_batch_ddl"] = False + self.db_wrapper.settings_dict["AUTO_BATCH_DDL"] = "true" + self.assertTrue(self.db_wrapper.auto_batch_ddl) + + for invalid in ("maybe", None, 1): + self.db_wrapper.settings_dict["AUTO_BATCH_DDL"] = invalid + with self.assertRaises(ImproperlyConfigured): + _ = self.db_wrapper.auto_batch_ddl + + self.db_wrapper.settings_dict.pop("AUTO_BATCH_DDL", None) + self.db_wrapper.settings_dict["OPTIONS"].pop("auto_batch_ddl", None) + + def test_get_new_connection_strips_auto_batch_ddl_option(self): + self.db_wrapper.Database = mock_database = mock.MagicMock() + mock_connection = mock_database.connect + conn_params = {"auto_batch_ddl": True, "test_param": "dummy"} + self.db_wrapper.get_new_connection(conn_params) + self.assertNotIn("auto_batch_ddl", mock_connection.call_args[1]) + self.assertEqual(mock_connection.call_args[1]["test_param"], "dummy") + + def test_ddl_batch_reuses_scoped_dbapi_transaction_lifecycle(self): + import warnings + + from django.db.utils import InterfaceError, ProgrammingError + from google.cloud.spanner_dbapi.connection import Connection + + mock_db = mock.MagicMock() + snapshot_cm = mock.MagicMock() + snapshot_obj = mock.MagicMock() + snapshot_obj.execute_sql.return_value = iter([]) + snapshot_cm.__enter__.return_value = snapshot_obj + mock_db.snapshot.return_value = snapshot_cm + + real_conn = Connection(self.db_wrapper.instance, mock_db) + real_conn.autocommit = True + self.db_wrapper.connection = real_conn + self.db_wrapper.autocommit = True + + with ( + mock.patch.object(real_conn, "begin", wraps=real_conn.begin) as spy_begin, + mock.patch.object( + real_conn, "commit", wraps=real_conn.commit + ) as spy_commit, + mock.patch.object( + real_conn, "rollback", wraps=real_conn.rollback + ) as spy_rollback, + mock.patch.object( + real_conn, + "run_prior_DDL_statements", + wraps=real_conn.run_prior_DDL_statements, + ) as spy_run_prior, + warnings.catch_warnings(record=True) as caught_warnings, + ): + warnings.simplefilter("always") + self.db_wrapper.start_batch_ddl(auto_flush=True) + # Lazy entry: Connection.begin() is not called until the first DDL. + self.assertEqual(spy_begin.call_count, 0) + + with self.db_wrapper.create_cursor() as cur: + cur.execute("CREATE TABLE t1 (id INT64) PRIMARY KEY (id)") + self.assertEqual(spy_begin.call_count, 1) + cur.execute("CREATE INDEX t1_idx ON t1 (id)") + self.assertEqual(spy_begin.call_count, 1) + mock_db.update_ddl.assert_not_called() + + # Snapshot query commits the scoped DDL burst, then executes cleanly. + cur.run_sql_in_snapshot("SELECT 1") + self.assertEqual(spy_commit.call_count, 1) + self.assertEqual(spy_run_prior.call_count, 2) + mock_db.update_ddl.assert_called_once_with( + [ + "CREATE TABLE t1 (id INT64) PRIMARY KEY (id)", + "CREATE INDEX t1_idx ON t1 (id)", + ] + ) + + # Explicit client transaction runs unimpeded while outer auto_flush + # DDL batch stays open, then subsequent DDL opens a new burst. + cur.execute("BEGIN") + cur.execute("ROLLBACK") + self.assertTrue(self.db_wrapper.in_ddl_batch) + + cur.execute("CREATE TABLE aborted_t (id INT64) PRIMARY KEY (id)") + self.db_wrapper.abort_batch() + self.assertGreaterEqual(spy_rollback.call_count, 2) + # Verify abort_batch cleared _ddl_statements so later snapshot SQL + # never leaks the aborted statement into update_ddl. + mock_db.update_ddl.reset_mock() + cur.run_sql_in_snapshot("SELECT 1") + mock_db.update_ddl.assert_not_called() + + self.assertEqual(caught_warnings, []) + + # First-statement failure in raw cursor rolls back Connection.begin(); + # second-statement failure preserves already-queued DDL until abort/exit. + self.db_wrapper.start_batch_ddl() + fail_first = mock.MagicMock( + spec=["connection", "is_closed", "execute", "close"] + ) + fail_first.connection = real_conn + fail_first.is_closed = False + fail_first.execute.side_effect = RuntimeError("raw execute boom") + with mock.patch.object(real_conn, "cursor", return_value=fail_first): + with self.assertRaises(RuntimeError): + self.db_wrapper.create_cursor().execute( + "CREATE TABLE bad1 (id INT64) PRIMARY KEY (id)" + ) + self.assertEqual(self.db_wrapper.buffered_ddl_statements, ()) + + with self.db_wrapper.create_cursor() as cur: + cur.execute("CREATE TABLE keep1 (id INT64) PRIMARY KEY (id)") + with mock.patch.object(real_conn, "cursor", return_value=fail_first): + with self.assertRaises(RuntimeError): + self.db_wrapper.create_cursor().execute( + "CREATE TABLE bad2 (id INT64) PRIMARY KEY (id)" + ) + self.assertEqual( + self.db_wrapper.buffered_ddl_statements, + ("CREATE TABLE keep1 (id INT64) PRIMARY KEY (id)",), + ) + self.db_wrapper.abort_batch() + + # Closed and disconnected cursors raise PEP 249 InterfaceError / ProgrammingError. + closed_cur = self.db_wrapper.create_cursor() + closed_cur.close() + with self.assertRaises(InterfaceError): + closed_cur.execute("START BATCH DDL") + + disconnected_raw = mock.MagicMock( + spec=["connection", "is_closed", "execute", "close"] + ) + disconnected_raw.connection = None + disconnected_raw.is_closed = False + with mock.patch.object(real_conn, "cursor", return_value=disconnected_raw): + with self.assertRaises(ProgrammingError): + self.db_wrapper.create_cursor().execute("START BATCH DDL") + + def test_ddl_batch_state_lifecycle_and_errors(self): + from django.db.utils import ( + IntegrityError, + OperationalError, + ProgrammingError, + ) + from google.api_core.exceptions import ( + AlreadyExists, + FailedPrecondition, + InternalServerError, + InvalidArgument, + ServiceUnavailable, + ) + from google.cloud.spanner_dbapi.connection import Connection + + # Empty flush is a no-op without opening a connection. + self.db_wrapper.connection = None + self.assertIsNone(self.db_wrapper.flush_ddl_batch()) + + # Running or aborting when inactive raises Django ProgrammingError. + with self.assertRaises(ProgrammingError): + self.db_wrapper.run_batch() + with self.assertRaises(ProgrammingError): + self.db_wrapper.abort_batch() + + # Starting while in atomic block or non-autocommit fails. + self.db_wrapper.in_atomic_block = True + try: + with self.assertRaises(ProgrammingError): + self.db_wrapper.start_batch_ddl() + finally: + self.db_wrapper.in_atomic_block = False + + self.db_wrapper.settings_dict["AUTOCOMMIT"] = False + try: + with self.assertRaises(ProgrammingError): + self.db_wrapper.start_batch_ddl() + finally: + self.db_wrapper.settings_dict["AUTOCOMMIT"] = True + + # Start lazy DDL batch before connection exists: nested lazy ddl_batch, + # lazy checkpoint/rollback, lazy run_batch, lazy abort_batch, and close(). + self.db_wrapper.start_batch_ddl() + self.assertTrue(self.db_wrapper.in_ddl_batch) + self.assertTrue(self.db_wrapper.ddl_batch_auto_flush) + with self.assertRaises(ProgrammingError): + self.db_wrapper.start_batch_ddl() + cp = self.db_wrapper.checkpoint_ddl_batch() + self.assertEqual(cp, (0, 0)) + with self.assertRaises(RuntimeError): + with self.db_wrapper.ddl_batch(auto_flush=False): + self.assertFalse(self.db_wrapper.ddl_batch_auto_flush) + raise RuntimeError("lazy nested rollback") + self.assertTrue(self.db_wrapper.ddl_batch_auto_flush) + self.db_wrapper.rollback_ddl_checkpoint(cp) + self.assertIsNone(self.db_wrapper.run_batch()) + self.assertFalse(self.db_wrapper.in_ddl_batch) + + self.db_wrapper.start_batch_ddl() + self.assertIsNone(self.db_wrapper.abort_batch()) + self.assertFalse(self.db_wrapper.in_ddl_batch) + + self.db_wrapper.start_batch_ddl() + self.db_wrapper.close() + self.assertFalse(self.db_wrapper.in_ddl_batch) + self.assertEqual(self.db_wrapper.buffered_ddl_statements, ()) + + # Bind a real spanner_dbapi.Connection backed by a mock Database. + mock_db = mock.MagicMock() + real_conn = Connection(self.db_wrapper.instance, mock_db) + real_conn.autocommit = True + self.db_wrapper.connection = real_conn + self.db_wrapper.autocommit = True + + # Reject starting, staging, flushing, or lazy-initializing DDL on read-only connection. + ro_conn = Connection(self.db_wrapper.instance, mock_db, read_only=True) + ro_conn.autocommit = True + ro_conn.close = mock.MagicMock() + self.db_wrapper.connection = ro_conn + with self.assertRaises(ProgrammingError): + self.db_wrapper.start_batch_ddl() + self.db_wrapper.connection = None + self.db_wrapper.start_batch_ddl() + self.db_wrapper.connection = ro_conn + with self.assertRaises(ProgrammingError): + self.db_wrapper.create_cursor().execute( + "CREATE TABLE ro_t (id INT64) PRIMARY KEY (id)" + ) + with self.assertRaises(ProgrammingError): + self.db_wrapper.flush_ddl_batch() + with self.assertRaises(ProgrammingError): + self.db_wrapper.init_connection_state() + self.assertFalse(self.db_wrapper.in_ddl_batch) + self.db_wrapper.connection = real_conn + + mock_dml_exec = mock.MagicMock() + mock_dml_exec.run_batch_dml.return_value = "dml_res" + with mock.patch( + "google.cloud.spanner_dbapi.connection.BatchDmlExecutor", + return_value=mock_dml_exec, + ): + real_conn.start_batch_dml(real_conn.cursor()) + with self.assertRaises(ProgrammingError): + self.db_wrapper.start_batch_ddl() + self.assertEqual(self.db_wrapper.run_batch(), "dml_res") + mock_dml_exec.run_batch_dml.assert_called_once_with() + + real_conn.start_batch_dml(real_conn.cursor()) + self.db_wrapper.abort_batch() + + with self.assertRaises(ProgrammingError): + self.db_wrapper.run_batch() + with self.assertRaises(ProgrammingError): + self.db_wrapper.abort_batch() + + # Empty flush while connected and active returns None without update_ddl RPC. + self.db_wrapper.start_batch_ddl() + self.assertIsNone(self.db_wrapper.flush_ddl_batch()) + self.assertIsNone(self.db_wrapper.run_batch()) + mock_db.update_ddl.assert_not_called() + + # Map GoogleAPICallError subclasses to Django DatabaseError types and emit OTel spans. + error_cases = [ + (AlreadyExists("dup"), IntegrityError), + (FailedPrecondition("precond"), IntegrityError), + (InvalidArgument("bad syntax"), ProgrammingError), + (InternalServerError("boom"), OperationalError), + (ServiceUnavailable("transient"), OperationalError), + ] + for api_exc, expected_django_exc in error_cases: + self.ot_exporter.clear() + mock_db.update_ddl.side_effect = api_exc + with self.assertRaises(expected_django_exc) as exc_cm: + with self.db_wrapper.ddl_batch(): + with self.db_wrapper.create_cursor() as cur: + cur.execute("CREATE TABLE t (id INT64) PRIMARY KEY (id)") + self.assertIn(api_exc.message, str(exc_cm.exception)) + self.assertFalse(self.db_wrapper.in_ddl_batch) + self.assertEqual(self.db_wrapper.buffered_ddl_statements, ()) + spans = self.ot_exporter.get_finished_spans() + self.assertEqual(spans[-1].name, "CloudSpannerDjango.flush_ddl_batch") + self.assertEqual(spans[-1].attributes["statement_count"], 1) + + mock_db.update_ddl.side_effect = None + mock_db.update_ddl.reset_mock() + + # Flush DDL batch traces non-empty flush while keeping batch active. + self.ot_exporter.clear() + self.db_wrapper.start_batch_ddl() + with self.db_wrapper.create_cursor() as cur: + cur.execute("CREATE TABLE t_flush (id INT64) PRIMARY KEY (id)") + self.db_wrapper.flush_ddl_batch() + spans = self.ot_exporter.get_finished_spans() + self.assertEqual(spans[-1].name, "CloudSpannerDjango.flush_ddl_batch") + self.assertEqual(spans[-1].attributes["statement_count"], 1) + self.assertTrue(self.db_wrapper.in_ddl_batch) + self.db_wrapper.run_batch() + + # Missing database raises ValueError on cursor execute and flush_ddl_batch. + no_db_conn = Connection(self.db_wrapper.instance, None) + no_db_conn.autocommit = True + self.db_wrapper.connection = no_db_conn + self.db_wrapper.start_batch_ddl() + with self.assertRaises(ValueError): + self.db_wrapper.create_cursor().execute( + "CREATE TABLE no_db_t (id INT64) PRIMARY KEY (id)" + ) + with self.assertRaises(ValueError): + self.db_wrapper.flush_ddl_batch() + self.db_wrapper.abort_batch() + self.db_wrapper.connection = real_conn + + # _set_autocommit(False) and _start_transaction_under_autocommit flush + # in auto_flush mode and raise in strict mode. + mock_db.update_ddl.reset_mock() + self.db_wrapper.start_batch_ddl(auto_flush=True) + with self.db_wrapper.create_cursor() as cur: + cur.execute("CREATE TABLE t_auto (id INT64) PRIMARY KEY (id)") + self.db_wrapper._set_autocommit(False) + mock_db.update_ddl.assert_called_with( + ["CREATE TABLE t_auto (id INT64) PRIMARY KEY (id)"] + ) + self.db_wrapper._set_autocommit(True) + self.db_wrapper.abort_batch() + + self.db_wrapper.start_batch_ddl(auto_flush=True) + with self.db_wrapper.create_cursor() as cur: + cur.execute("CREATE TABLE t_tx (id INT64) PRIMARY KEY (id)") + self.db_wrapper._start_transaction_under_autocommit() + mock_db.update_ddl.assert_called_with( + ["CREATE TABLE t_tx (id INT64) PRIMARY KEY (id)"] + ) + real_conn.rollback() + self.db_wrapper.abort_batch() + + self.db_wrapper.start_batch_ddl(auto_flush=False) + with self.assertRaises(ProgrammingError): + self.db_wrapper._set_autocommit(False) + with self.assertRaises(ProgrammingError): + self.db_wrapper._start_transaction_under_autocommit() + self.db_wrapper.abort_batch() + + # Lazy pre-connect DDL batch stays active across init_connection_state, + # and manual init_connection_state warns when dropping un-flushed DDL. + real_conn.close = mock.MagicMock() + self.db_wrapper.connection = None + self.db_wrapper.start_batch_ddl(auto_flush=False) + self.db_wrapper.connection = real_conn + self.db_wrapper.init_connection_state() + self.assertTrue(self.db_wrapper.in_ddl_batch) + self.assertFalse(self.db_wrapper.ddl_batch_auto_flush) + with self.db_wrapper.create_cursor() as cur: + cur.execute("CREATE TABLE stale_t (id INT64) PRIMARY KEY (id)") + self.assertEqual(len(self.db_wrapper.buffered_ddl_statements), 1) + with self.assertLogs("django_spanner.base", level="WARNING") as cm: + self.db_wrapper.init_connection_state() + self.assertTrue(any("re-initialization" in msg for msg in cm.output)) + self.assertFalse(self.db_wrapper.in_ddl_batch) + self.assertEqual(self.db_wrapper.buffered_ddl_statements, ()) + + def test_ddl_cursor_proxy_formatting_and_guard_branches(self): + import asyncio + + from django.db.utils import ProgrammingError + from google.cloud.spanner_dbapi import batch_dml_executor + from google.cloud.spanner_dbapi.connection import Connection + + mock_db = mock.MagicMock() + snapshot_cm = mock.MagicMock() + snapshot_obj = mock.MagicMock() + snapshot_obj.execute_sql.return_value = iter([]) + snapshot_cm.__enter__.return_value = snapshot_obj + mock_db.snapshot.return_value = snapshot_cm + + real_conn = Connection(self.db_wrapper.instance, mock_db) + real_conn.autocommit = True + self.db_wrapper.connection = real_conn + self.db_wrapper.autocommit = True + + # Async autocommit helper delegates cleanly to _set_autocommit. + asyncio.run(self.db_wrapper._a_set_autocommit(True)) + + # Default AUTO_BATCH_DDL=False keeps schema_editor() unbatched; + # lazy unopened wrapper with AUTO_BATCH_DDL=True inspects AUTOCOMMIT setting. + with self.db_wrapper.schema_editor(): + self.assertFalse(self.db_wrapper.in_ddl_batch) + + self.db_wrapper.settings_dict["AUTO_BATCH_DDL"] = True + try: + self.db_wrapper.connection = None + with self.db_wrapper.schema_editor(): + self.assertTrue(self.db_wrapper.in_ddl_batch) + self.assertFalse(self.db_wrapper.in_ddl_batch) + + # Lazy nested ddl_batch connecting mid-block rolls back checkpoint on error. + self.db_wrapper.start_batch_ddl(auto_flush=True) + with self.assertRaises(RuntimeError): + with self.db_wrapper.ddl_batch(auto_flush=False): + self.db_wrapper.connection = real_conn + with self.db_wrapper.create_cursor() as lazy_cur: + lazy_cur.execute( + "CREATE TABLE lazy_drop (id INT64) PRIMARY KEY (id)" + ) + raise RuntimeError("abort after lazy connect") + self.assertEqual(self.db_wrapper.buffered_ddl_statements, ()) + self.db_wrapper.abort_batch() + + self.db_wrapper.connection = real_conn + self.db_wrapper.autocommit = False + with self.db_wrapper.schema_editor(): + self.assertFalse(self.db_wrapper.in_ddl_batch) + with self.assertRaises(ProgrammingError): + self.db_wrapper.start_batch_ddl() + self.db_wrapper.autocommit = True + + real_conn.begin() + with self.db_wrapper.schema_editor(): + self.assertFalse(self.db_wrapper.in_ddl_batch) + with self.assertRaises(ProgrammingError): + self.db_wrapper.start_batch_ddl() + real_conn.rollback() + + ro_conn = Connection(self.db_wrapper.instance, mock_db, read_only=True) + ro_conn.autocommit = True + self.db_wrapper.connection = ro_conn + with self.db_wrapper.schema_editor(): + self.assertFalse(self.db_wrapper.in_ddl_batch) + self.db_wrapper.connection = real_conn + + real_conn.start_batch_dml(real_conn.cursor()) + with self.db_wrapper.schema_editor(): + self.assertFalse(self.db_wrapper.in_ddl_batch) + real_conn.abort_batch() + finally: + self.db_wrapper.settings_dict.pop("AUTO_BATCH_DDL", None) + + # Proxy cursor attributes, iteration, client-side comments, callproc, and executemany. + cur = self.db_wrapper.create_cursor() + cur.arraysize = 50 + self.assertEqual(cur.arraysize, 50) + with self.assertRaises(ProgrammingError): + next(iter(cur)) + with self.assertRaises(ProgrammingError): + next(cur) + cur.callproc("sp_noop") + cur.callproc("sp_args", [1]) + cur.execute("/* leading */ START BATCH DDL; -- trailing") + with self.assertRaises(ProgrammingError): + cur.execute("START BATCH DDL; SELECT 1") + cur.execute("ABORT BATCH") + + # Nested ddl_batch(auto_flush=False) tightens strictness, formats %s / %(k)s + # GoogleSQL literals safely, blocks parameter injection, and rolls back inner DDL. + with self.db_wrapper.ddl_batch(auto_flush=True): + cur.execute("CREATE TABLE outer_ok (id INT64) PRIMARY KEY (id)") + cur.execute( + "ALTER TABLE outer_ok ADD COLUMN c STRING(20) DEFAULT (%(d)s)", + {"d": "a'b"}, + ) + with self.assertRaises(ProgrammingError): + cur.execute("ALTER TABLE outer_ok ADD COLUMN bad INT64", 123) + with self.assertRaises(ProgrammingError): + cur.execute( + "ALTER TABLE outer_ok ADD COLUMN bad INT64 DEFAULT (%s)", + [1, 2], + ) + with self.assertRaises(ProgrammingError): + cur.execute(" ; -- empty with params", [1]) + with self.assertRaises(RuntimeError): + with self.db_wrapper.ddl_batch(auto_flush=False): + self.assertFalse(self.db_wrapper.ddl_batch_auto_flush) + cur.execute("CREATE TABLE inner_drop (id INT64) PRIMARY KEY (id)") + with self.assertRaises(ProgrammingError): + cur.get_table_column_schema("outer_ok") + raise RuntimeError("rollback inner") + self.assertTrue(self.db_wrapper.ddl_batch_auto_flush) + self.assertEqual( + self.db_wrapper.buffered_ddl_statements, + ( + "CREATE TABLE outer_ok (id INT64) PRIMARY KEY (id)", + r"ALTER TABLE outer_ok ADD COLUMN c STRING(20) DEFAULT ('a\'b')", + ), + ) + + # Snapshot helpers, executemany, and START BATCH DML flush cleanly. + cur.list_tables() + self.assertEqual(self.db_wrapper.buffered_ddl_statements, ()) + cur.execute("CREATE TABLE after_list (id INT64) PRIMARY KEY (id)") + cur.get_table_column_schema("after_list") + self.assertEqual(self.db_wrapper.buffered_ddl_statements, ()) + + with mock.patch.object( + batch_dml_executor, "run_batch_dml", return_value=mock.MagicMock() + ): + cur.execute("CREATE TABLE before_many (id INT64) PRIMARY KEY (id)") + cur.executemany("INSERT INTO outer_ok (id) VALUES (%s)", [(1,)]) + self.assertEqual(self.db_wrapper.buffered_ddl_statements, ()) + + with self.assertRaises(ProgrammingError): + cur.executemany("DROP TABLE outer_ok", [()]) + + # Parameterized DDL double-quote mistake / injection guard and mixed SQL guard. + with self.assertRaises(ProgrammingError): + cur.execute( + "ALTER TABLE outer_ok ADD COLUMN bad STRING(40) DEFAULT ('%s')", + ["x); SELECT 1; --"], + ) + with self.assertRaises(ProgrammingError): + cur.execute( + "CREATE TABLE m1 (id INT64) PRIMARY KEY (id); " + "INSERT INTO m1 (id) VALUES (%s)", + [1], + ) + + with self.assertRaises(_SimulatedInterrupt): + with self.db_wrapper.ddl_batch(auto_flush=True): + cur.execute( + "CREATE TABLE flushed_in_inner (id INT64) PRIMARY KEY (id)" + ) + cur.list_tables() + cur.execute( + "CREATE TABLE dropped_post_flush (id INT64) PRIMARY KEY (id)" + ) + raise _SimulatedInterrupt("interrupt inner after flush") + self.assertEqual(self.db_wrapper.buffered_ddl_statements, ()) + + cur.execute("CREATE TABLE pre_dml_b (id INT64) PRIMARY KEY (id)") + cur.execute("START BATCH DML") + self.assertEqual(self.db_wrapper.buffered_ddl_statements, ()) + cur.executemany("SELECT 1", []) + with self.assertRaises(ProgrammingError): + cur.execute("START BATCH DDL") + cur.execute("ABORT BATCH") + + # BaseException unwinds outer ddl_batch cleanly; outer error preventing + # inner finally from restoring auto_flush handles inactive flag safely. + with self.assertRaises(_SimulatedInterrupt): + with self.db_wrapper.ddl_batch(): + with self.db_wrapper.ddl_batch(auto_flush=False): + with self.db_wrapper.create_cursor() as inner_cur: + inner_cur.execute( + "CREATE TABLE base_exc_t (id INT64) PRIMARY KEY (id)" + ) + self.db_wrapper.abort_batch() + raise _SimulatedInterrupt("ctrl-c") + self.assertFalse(self.db_wrapper.in_ddl_batch) + self.assertEqual(self.db_wrapper.buffered_ddl_statements, ()) + + # _close() helper discards active DDL batch before super()._close(). + self.db_wrapper.start_batch_ddl() + real_conn.close = mock.MagicMock() + self.db_wrapper._close() + self.assertFalse(self.db_wrapper.in_ddl_batch) + + # _start_transaction_under_autocommit fallback when ALLOW_TRANSACTIONS_IN_AUTO_COMMIT=False. + raw_cursor = mock.MagicMock() + mock_conn = mock.MagicMock() + mock_conn.cursor.return_value = raw_cursor + self.db_wrapper.connection = mock_conn + self.db_wrapper.settings_dict["ALLOW_TRANSACTIONS_IN_AUTO_COMMIT"] = False + try: + self.db_wrapper._start_transaction_under_autocommit() + raw_cursor.execute.assert_called_with("SELECT 1") + finally: + self.db_wrapper.settings_dict.pop("ALLOW_TRANSACTIONS_IN_AUTO_COMMIT", None) + + +class _SimulatedInterrupt(BaseException): + pass 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 754baa60c735..20111f3ef8ae 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 @@ -503,3 +503,185 @@ 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_quote_value_googlesql_literals(self): + import datetime + import uuid + from decimal import Decimal + + from django.apps.registry import Apps + from django.db import models + from django.db.models.functions import Concat, Lower, Now + from django.test import override_settings + from google.cloud.spanner_dbapi.parse_utils import DateStr, TimestampStr + from google.cloud.spanner_v1 import JsonObject + + schema_editor = DatabaseSchemaEditor(self.connection) + self.assertEqual(schema_editor.quote_value(True), "TRUE") + self.assertEqual(schema_editor.quote_value(False), "FALSE") + self.assertEqual(schema_editor.quote_value(None), "NULL") + self.assertEqual(schema_editor.quote_value(42), "42") + self.assertEqual(schema_editor.quote_value(1.5), "1.5") + self.assertEqual( + schema_editor.quote_value("it's a \\ test\n\r\t\x00\x7f"), + r"'it\'s a \\ test\n\r\t\x00\x7f'", + ) + self.assertEqual( + schema_editor.quote_value(JsonObject({"k": "it's"})), + r"""JSON '{"k":"it\'s"}'""", + ) + self.assertEqual(schema_editor.quote_value(JsonObject(None)), "NULL") + self.assertEqual( + schema_editor.quote_value( + uuid.UUID("12345678-1234-5678-1234-567812345678") + ), + "'12345678123456781234567812345678'", + ) + self.assertEqual(schema_editor.quote_value(b"\x00\xffa"), r"b'\x00\xff\x61'") + self.assertEqual( + schema_editor.quote_value(DateStr("2026-11-10")), + "DATE '2026-11-10'", + ) + self.assertEqual( + schema_editor.quote_value(TimestampStr("2026-11-10T01:02:03.000000Z")), + "TIMESTAMP '2026-11-10T01:02:03.000000Z'", + ) + self.assertEqual( + schema_editor.quote_value(datetime.date(2026, 11, 10)), + "DATE '2026-11-10'", + ) + with override_settings(USE_TZ=False): + self.assertEqual( + schema_editor.quote_value(datetime.datetime(2026, 11, 10, 1, 2, 3)), + "TIMESTAMP '2026-11-10T01:02:03Z'", + ) + with override_settings(USE_TZ=True, TIME_ZONE="America/New_York"): + self.assertEqual( + schema_editor.quote_value(datetime.datetime(2026, 1, 15, 12, 0, 0)), + "TIMESTAMP '2026-01-15T17:00:00Z'", + ) + tz_plus2 = datetime.timezone(datetime.timedelta(hours=2)) + self.assertEqual( + schema_editor.quote_value( + datetime.datetime(2026, 11, 10, 3, 2, 3, tzinfo=tz_plus2) + ), + "TIMESTAMP '2026-11-10T01:02:03Z'", + ) + self.assertEqual( + schema_editor.quote_value(datetime.time(12, 30, 0)), + "TIMESTAMP '0001-01-01T12:30:00.000000Z'", + ) + self.assertEqual( + schema_editor.quote_value(Decimal("12.34")), + "NUMERIC '12.34'", + ) + self.assertEqual( + schema_editor.quote_value(float("nan")), + "CAST('NaN' AS FLOAT64)", + ) + self.assertEqual( + schema_editor.quote_value(float("inf")), + "CAST('Infinity' AS FLOAT64)", + ) + self.assertEqual( + schema_editor.quote_value(float("-inf")), + "CAST('-Infinity' AS FLOAT64)", + ) + with self.assertRaises(ValueError): + schema_editor.quote_value(Decimal("NaN")) + with self.assertRaises(TypeError): + schema_editor.quote_value(object()) + + class _LiteralModel(models.Model): + name = models.CharField(max_length=40) + tag = models.CharField(max_length=20, db_default="it's") + created_at = models.DateTimeField(db_default=Now()) + gen_tag = models.GeneratedField( + expression=Concat(models.Value("a'b"), "name"), + output_field=models.CharField(max_length=80), + db_persist=True, + ) + gen_lower = models.GeneratedField( + expression=Lower("name"), + output_field=models.CharField(max_length=80), + db_persist=True, + ) + + class Meta: + app_label = "tests" + apps = Apps() + + def_sql, _ = schema_editor.column_sql( + _LiteralModel, _LiteralModel._meta.get_field("tag") + ) + self.assertIn(r"DEFAULT ('it\'s')", def_sql) + now_sql, _ = schema_editor.column_sql( + _LiteralModel, _LiteralModel._meta.get_field("created_at") + ) + self.assertIn("DEFAULT (", now_sql) + gen_sql, _ = schema_editor.column_sql( + _LiteralModel, _LiteralModel._meta.get_field("gen_tag") + ) + self.assertIn(r"'a\'b'", gen_sql) + self.assertTrue(gen_sql.endswith("STORED")) + gen_np_sql, _ = schema_editor.column_sql( + _LiteralModel, _LiteralModel._meta.get_field("gen_lower") + ) + self.assertTrue(gen_np_sql.endswith("STORED")) + + def test_prepare_sql_script_strips_comment_and_empty_fragments(self): + script = ( + "-- leading comment\n" + "CREATE TABLE a (id INT64 NOT NULL) PRIMARY KEY (id);\n" + " ; \n" + "CREATE INDEX a_idx ON a (id);\n" + "-- trailing comment\n" + ) + self.assertEqual( + self.connection.ops.prepare_sql_script(script), + [ + "CREATE TABLE a (id INT64 NOT NULL) PRIMARY KEY (id);", + "CREATE INDEX a_idx ON a (id);", + ], + ) + + def test_schema_editor_collect_sql_and_atomic_guard(self): + from django.db.transaction import TransactionManagementError + + self.connection.settings_dict["AUTO_BATCH_DDL"] = True + try: + # Lazy unopened connection opens an auto-flushing batch only outside + # collect_sql=True, atomic blocks, and AUTOCOMMIT=False. + with DatabaseSchemaEditor(self.connection): + self.assertTrue(self.connection.in_ddl_batch) + self.assertFalse(self.connection.in_ddl_batch) + + self.connection.settings_dict["AUTOCOMMIT"] = False + try: + with DatabaseSchemaEditor(self.connection): + self.assertFalse(self.connection.in_ddl_batch) + finally: + self.connection.settings_dict["AUTOCOMMIT"] = True + + with DatabaseSchemaEditor(self.connection, collect_sql=True) as editor: + self.assertFalse(self.connection.in_ddl_batch) + editor.create_model(Author) + with self.assertRaises(TypeError): + editor.execute( + "INSERT INTO tests_author (name) VALUES (%s)", [object()] + ) + self.assertFalse(self.connection.in_ddl_batch) + self.assertTrue( + any("CREATE TABLE tests_author" in s for s in editor.collected_sql) + ) + + self.connection.in_atomic_block = True + try: + with self.assertRaises(TransactionManagementError): + with DatabaseSchemaEditor(self.connection) as editor: + self.assertFalse(self.connection.in_ddl_batch) + editor.create_model(Author) + finally: + self.connection.in_atomic_block = False + finally: + self.connection.settings_dict.pop("AUTO_BATCH_DDL", None)