Cockroach upgrade - autocommit_before_ddl (#1348)

* upgrade crdb to the latest regular release

* remove cockroach_upgrade branch from CI

* settings applied to the current transaction only

* make changes to transaction and atomic

* remove postgres password

* `autocommit_before_ddl` experiments

* arg -> kwarg

* tweak docs

* fix linter error

* fix test

* Update tests.yaml

* more test fixes

---------

Co-authored-by: sinisaos <[email protected]>

Author: dantownsend

.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 }}

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

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,
+        )

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()

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()

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()