diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 2f6655c2d..07491e85e 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -142,7 +142,7 @@ jobs: strategy: matrix: python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] - cockroachdb-version: ["v24.1.0"] + cockroachdb-version: ["v25.4.3"] steps: - uses: actions/checkout@v3 - name: Set up Python ${{ matrix.python-version }} diff --git a/piccolo/apps/migrations/auto/migration_manager.py b/piccolo/apps/migrations/auto/migration_manager.py index ce085783f..57b97fa2d 100644 --- a/piccolo/apps/migrations/auto/migration_manager.py +++ b/piccolo/apps/migrations/auto/migration_manager.py @@ -18,6 +18,7 @@ from piccolo.columns import Column, column_types from piccolo.columns.column_types import ForeignKey, Serial from piccolo.engine import engine_finder +from piccolo.engine.cockroach import CockroachTransaction from piccolo.query import Query from piccolo.query.base import DDL from piccolo.query.constraints import get_fk_constraint_name @@ -984,7 +985,11 @@ async def run(self, backwards: bool = False): engine.transaction() if self.wrap_in_transaction else SkippedTransaction() - ): + ) as transaction: + if isinstance(transaction, CockroachTransaction): + # To enable DDL rollbacks in CockroachDB. + await transaction.autocommit_before_ddl(enabled=False) + if not self.preview: if direction == "backwards": raw_list = self.raw_backwards diff --git a/piccolo/engine/cockroach.py b/piccolo/engine/cockroach.py index e823ced5a..d47336382 100644 --- a/piccolo/engine/cockroach.py +++ b/piccolo/engine/cockroach.py @@ -1,14 +1,57 @@ from __future__ import annotations from collections.abc import Sequence -from typing import Any, Optional +from typing import Any, Optional, cast -from piccolo.utils.lazy_loader import LazyLoader -from piccolo.utils.warnings import Level, colored_warning +from .postgres import Atomic, PostgresEngine, PostgresTransaction -from .postgres import PostgresEngine -asyncpg = LazyLoader("asyncpg", globals(), "asyncpg") +class CockroachAtomic(Atomic): + + __slots__ = ("autocommit_before_ddl",) + + def __init__( + self, + engine: CockroachEngine, + autocommit_before_ddl: Optional[bool] = False, + ): + """ + :param autocommit_before_ddl: + Defaults to ``False`` to prevent automatic DDL commits + in transactions (preventing rollbacks). Applies only to the current + transaction and is automatically reverted when the transaction + commits or is rolled back. + + Usage:: + + # Defaults to ``False`` (``autocommit_before_ddl = off``) + transaction = engine.atomic() + transaction.add(Foo.create_table()) + + # If we want to set ``autocommit_before_ddl = on``, + # which is the default Cockroach session setting. + transaction = engine.atomic(autocommit_before_ddl=True) + transaction.add(Foo.create_table()) + + """ + super().__init__(engine) + self.autocommit_before_ddl = autocommit_before_ddl + + async def setup_transaction(self, transaction: PostgresTransaction): + if self.autocommit_before_ddl is not None: + transaction = cast(CockroachTransaction, transaction) + await transaction.autocommit_before_ddl( + enabled=self.autocommit_before_ddl + ) + + +class CockroachTransaction(PostgresTransaction): + + async def autocommit_before_ddl(self, enabled: bool = True): + value = "on" if enabled else "off" + await self.connection.execute( + f"SET LOCAL autocommit_before_ddl = {value}" + ) class CockroachEngine(PostgresEngine): @@ -35,15 +78,20 @@ def __init__( self.engine_type = "cockroach" self.min_version_number = 0 - async def prep_database(self): - try: - await self._run_in_new_connection( - "SET CLUSTER SETTING sql.defaults.experimental_alter_column_type.enabled = true;" # noqa: E501 - ) - except asyncpg.exceptions.InsufficientPrivilegeError: - colored_warning( - "=> Unable to set up Cockroach DB " - "functionality may not behave as expected. Make sure " - "your database user has permission to set cluster options.", - level=Level.medium, - ) + def atomic( + self, + autocommit_before_ddl: Optional[bool] = False, + ) -> CockroachAtomic: + return CockroachAtomic( + engine=self, + autocommit_before_ddl=autocommit_before_ddl, + ) + + def transaction( + self, + allow_nested: bool = True, + ) -> CockroachTransaction: + return CockroachTransaction( + engine=self, + allow_nested=allow_nested, + ) diff --git a/piccolo/engine/postgres.py b/piccolo/engine/postgres.py index f4ac2d17e..ec8c11b91 100644 --- a/piccolo/engine/postgres.py +++ b/piccolo/engine/postgres.py @@ -113,11 +113,15 @@ def __init__(self, engine: PostgresEngine): def add(self, *query: Union[Query, DDL]): self.queries += list(query) + async def setup_transaction(self, transaction: PostgresTransaction): + pass + async def run(self): from piccolo.query.methods.objects import Create, GetOrCreate try: - async with self.engine.transaction(): + async with self.engine.transaction() as transaction: + await self.setup_transaction(transaction=transaction) for query in self.queries: if isinstance(query, (Query, DDL, Create, GetOrCreate)): await query.run() diff --git a/piccolo/testing/test_case.py b/piccolo/testing/test_case.py index 08bd61a5c..b123e72da 100644 --- a/piccolo/testing/test_case.py +++ b/piccolo/testing/test_case.py @@ -4,6 +4,7 @@ from unittest import IsolatedAsyncioTestCase, TestCase from piccolo.engine import Engine, engine_finder +from piccolo.engine.cockroach import CockroachTransaction from piccolo.table import ( Table, create_db_tables, @@ -112,9 +113,13 @@ async def asyncSetUp(self) -> None: db = self.db or engine_finder() assert db is not None self.transaction = db.transaction() + # This is only available in Python 3.11 and above: await self.enterAsyncContext(cm=self.transaction) # type: ignore + if isinstance(self.transaction, CockroachTransaction): + await self.transaction.autocommit_before_ddl(False) + async def asyncTearDown(self): await super().asyncTearDown() await self.transaction.rollback() diff --git a/tests/engine/test_transaction.py b/tests/engine/test_transaction.py index d381f5d14..e9f930837 100644 --- a/tests/engine/test_transaction.py +++ b/tests/engine/test_transaction.py @@ -4,6 +4,7 @@ import pytest +from piccolo.engine.cockroach import CockroachTransaction from piccolo.engine.sqlite import SQLiteEngine, TransactionType from piccolo.table import drop_db_tables_sync from piccolo.utils.sync import run_sync @@ -132,6 +133,9 @@ def test_manual_rollback(self): async def run_transaction(): async with Band._meta.db.transaction() as transaction: + if isinstance(transaction, CockroachTransaction): + await transaction.autocommit_before_ddl(enabled=False) + await Manager.create_table() await transaction.rollback()