From c92eeade3afc15f66db740ba3675d7aca0525a52 Mon Sep 17 00:00:00 2001 From: sinisaos Date: Tue, 30 Dec 2025 10:01:06 +0100 Subject: [PATCH 01/12] upgrade crdb to the latest regular release --- .github/workflows/tests.yaml | 4 ++-- piccolo/engine/cockroach.py | 29 +++++++++++++++-------------- 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 2f6655c2d..7fcee0aa5 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -2,7 +2,7 @@ name: Test Suite on: push: - branches: ["master", "v1"] + branches: ["master", "v1", "cockroach_upgrade"] paths-ignore: - "docs/**" pull_request: @@ -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.2"] steps: - uses: actions/checkout@v3 - name: Set up Python ${{ matrix.python-version }} diff --git a/piccolo/engine/cockroach.py b/piccolo/engine/cockroach.py index e823ced5a..4f55f9d4a 100644 --- a/piccolo/engine/cockroach.py +++ b/piccolo/engine/cockroach.py @@ -1,15 +1,17 @@ from __future__ import annotations from collections.abc import Sequence -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from piccolo.utils.lazy_loader import LazyLoader -from piccolo.utils.warnings import Level, colored_warning from .postgres import PostgresEngine asyncpg = LazyLoader("asyncpg", globals(), "asyncpg") +if TYPE_CHECKING: # pragma: no cover + from asyncpg.connection import Connection + class CockroachEngine(PostgresEngine): """ @@ -35,15 +37,14 @@ 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, - ) + async def get_new_connection(self) -> Connection: + """ + Set `autocommit_before_ddl` to off (enabled by default since v25.2) + to prevent automatic DDL commits in transactions and enable rollback + """ + connection = await super().get_new_connection() + await connection.execute( + "SET autocommit_before_ddl = off;" + "ALTER ROLE ALL SET autocommit_before_ddl = false;" + ) + return connection From 3b8e6bfc33a37ec4463a67b3eeb1924f7954a12e Mon Sep 17 00:00:00 2001 From: sinisaos Date: Tue, 30 Dec 2025 10:21:57 +0100 Subject: [PATCH 02/12] remove cockroach_upgrade branch from CI --- .github/workflows/tests.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 7fcee0aa5..9ad4a88e4 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -2,7 +2,7 @@ name: Test Suite on: push: - branches: ["master", "v1", "cockroach_upgrade"] + branches: ["master", "v1"] paths-ignore: - "docs/**" pull_request: From 6f33b8a7e4f2204b198efa519e203b36fd7de8a1 Mon Sep 17 00:00:00 2001 From: sinisaos Date: Wed, 7 Jan 2026 08:02:51 +0100 Subject: [PATCH 03/12] settings applied to the current transaction only --- piccolo/engine/cockroach.py | 30 ++++++++++++------------------ 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/piccolo/engine/cockroach.py b/piccolo/engine/cockroach.py index 4f55f9d4a..b7253c01e 100644 --- a/piccolo/engine/cockroach.py +++ b/piccolo/engine/cockroach.py @@ -1,16 +1,19 @@ from __future__ import annotations from collections.abc import Sequence -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional -from piccolo.utils.lazy_loader import LazyLoader +from .postgres import PostgresEngine, PostgresTransaction -from .postgres import PostgresEngine -asyncpg = LazyLoader("asyncpg", globals(), "asyncpg") - -if TYPE_CHECKING: # pragma: no cover - from asyncpg.connection import Connection +class CockroachTransaction(PostgresTransaction): + async def begin(self): + await self.transaction.start() + # Set `autocommit_before_ddl` to off (enabled by default since v25.2) + # to prevent automatic DDL commits in transactions and enable rollback. + # Applies only to the current transaction and automatically reverted + # when the transaction commits or rollback. + await self.connection.execute("SET LOCAL autocommit_before_ddl = off") class CockroachEngine(PostgresEngine): @@ -37,14 +40,5 @@ def __init__( self.engine_type = "cockroach" self.min_version_number = 0 - async def get_new_connection(self) -> Connection: - """ - Set `autocommit_before_ddl` to off (enabled by default since v25.2) - to prevent automatic DDL commits in transactions and enable rollback - """ - connection = await super().get_new_connection() - await connection.execute( - "SET autocommit_before_ddl = off;" - "ALTER ROLE ALL SET autocommit_before_ddl = false;" - ) - return connection + def transaction(self, allow_nested: bool = True) -> CockroachTransaction: + return CockroachTransaction(engine=self, allow_nested=allow_nested) From 3e9300d4b9589bc9e3e53bfe4800561b383780a1 Mon Sep 17 00:00:00 2001 From: sinisaos Date: Fri, 16 Jan 2026 08:22:14 +0100 Subject: [PATCH 04/12] make changes to transaction and atomic --- piccolo/engine/cockroach.py | 119 +++++++++++++++++++++++++++++++++--- tests/postgres_conf.py | 2 +- 2 files changed, 112 insertions(+), 9 deletions(-) diff --git a/piccolo/engine/cockroach.py b/piccolo/engine/cockroach.py index b7253c01e..a19aa5c20 100644 --- a/piccolo/engine/cockroach.py +++ b/piccolo/engine/cockroach.py @@ -3,17 +3,103 @@ from collections.abc import Sequence from typing import Any, Optional -from .postgres import PostgresEngine, PostgresTransaction +from piccolo.query.base import DDL, Query + +from .postgres import Atomic, PostgresEngine, PostgresTransaction + + +class CockroachAtomic(Atomic): + """ + :param autocommit_before_ddl: + Default to ``False`` to prevent automatic DDL commits + in transactions and enable rollback. + Applies only to the current transaction and automatically + reverted when the transaction commits or rollback. + + Usage:: + + # Default to ``False`` (``autocommit_before_ddl = off``) + transaction = engine.atomic() + transaction.add(Foo.create_table()) + + # If we want set ``autocommit_before_ddl = on`` + # which is default Cockroach session setting. + transaction = engine.atomic(autocommit_before_ddl=True) + transaction.add(Foo.create_table()) + + """ + + __slots__ = ("autocommit_before_ddl",) + + def __init__( + self, + engine: CockroachEngine, + autocommit_before_ddl: Optional[bool] = False, + ): + super().__init__(engine) + self.autocommit_before_ddl = autocommit_before_ddl + + async def run(self): + from piccolo.query.methods.objects import Create, GetOrCreate + + try: + async with self.engine.transaction( + autocommit_before_ddl=self.autocommit_before_ddl + ): + for query in self.queries: + if isinstance(query, (Query, DDL, Create, GetOrCreate)): + await query.run() + else: + raise ValueError("Unrecognised query") + + self.queries = [] + + except Exception as exception: + self.queries = [] + raise exception from exception class CockroachTransaction(PostgresTransaction): + """ + :param autocommit_before_ddl: + Default to ``False`` to prevent automatic DDL commits + in transactions and enable rollback. Applies only + to the current transaction and automatically + reverted when the transaction commits or rollback. + + Usage:: + + # Default to ``False`` (``autocommit_before_ddl = off``) + async with engine.transaction(): + # Run some queries: + await Band.select().run() + + # If we want set ``autocommit_before_ddl = on`` + # which is default Cockroach session setting. + async with engine.transaction(autocommit_before_ddl=True): + # Run some queries: + await Band.select().run() + + """ + + __slots__ = ("autocommit_before_ddl",) + + def __init__( + self, + autocommit_before_ddl: Optional[bool] = False, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.autocommit_before_ddl = autocommit_before_ddl + async def begin(self): await self.transaction.start() - # Set `autocommit_before_ddl` to off (enabled by default since v25.2) - # to prevent automatic DDL commits in transactions and enable rollback. - # Applies only to the current transaction and automatically reverted - # when the transaction commits or rollback. - await self.connection.execute("SET LOCAL autocommit_before_ddl = off") + + value = "on" if self.autocommit_before_ddl else "off" + await self.connection.execute( + f"SET LOCAL autocommit_before_ddl = {value}" + ) class CockroachEngine(PostgresEngine): @@ -40,5 +126,22 @@ def __init__( self.engine_type = "cockroach" self.min_version_number = 0 - def transaction(self, allow_nested: bool = True) -> CockroachTransaction: - return CockroachTransaction(engine=self, allow_nested=allow_nested) + 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, + autocommit_before_ddl: Optional[bool] = False, + ) -> CockroachTransaction: + return CockroachTransaction( + engine=self, + allow_nested=allow_nested, + autocommit_before_ddl=autocommit_before_ddl, + ) diff --git a/tests/postgres_conf.py b/tests/postgres_conf.py index af21dcbc5..36763b7eb 100644 --- a/tests/postgres_conf.py +++ b/tests/postgres_conf.py @@ -8,7 +8,7 @@ "host": os.environ.get("PG_HOST", "localhost"), "port": os.environ.get("PG_PORT", "5432"), "user": os.environ.get("PG_USER", "postgres"), - "password": os.environ.get("PG_PASSWORD", ""), + "password": os.environ.get("PG_PASSWORD", "postgres"), "database": os.environ.get("PG_DATABASE", "piccolo"), } ) From a420277043d3a591e8567b1fafcf4ff200fbff95 Mon Sep 17 00:00:00 2001 From: sinisaos Date: Fri, 16 Jan 2026 08:23:11 +0100 Subject: [PATCH 05/12] remove postgres password --- tests/postgres_conf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/postgres_conf.py b/tests/postgres_conf.py index 36763b7eb..af21dcbc5 100644 --- a/tests/postgres_conf.py +++ b/tests/postgres_conf.py @@ -8,7 +8,7 @@ "host": os.environ.get("PG_HOST", "localhost"), "port": os.environ.get("PG_PORT", "5432"), "user": os.environ.get("PG_USER", "postgres"), - "password": os.environ.get("PG_PASSWORD", "postgres"), + "password": os.environ.get("PG_PASSWORD", ""), "database": os.environ.get("PG_DATABASE", "piccolo"), } ) From 69c8c822aeaf926308219c82c0a09ab8868448cb Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Tue, 10 Feb 2026 23:02:19 +0000 Subject: [PATCH 06/12] `autocommit_before_ddl` experiments --- .../apps/migrations/auto/migration_manager.py | 7 +- piccolo/engine/cockroach.py | 102 +++++------------- piccolo/engine/postgres.py | 6 +- 3 files changed, 36 insertions(+), 79 deletions(-) diff --git a/piccolo/apps/migrations/auto/migration_manager.py b/piccolo/apps/migrations/auto/migration_manager.py index ce085783f..6f2e89099 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(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 a19aa5c20..1b05d9bd9 100644 --- a/piccolo/engine/cockroach.py +++ b/piccolo/engine/cockroach.py @@ -3,31 +3,10 @@ from collections.abc import Sequence from typing import Any, Optional -from piccolo.query.base import DDL, Query - from .postgres import Atomic, PostgresEngine, PostgresTransaction class CockroachAtomic(Atomic): - """ - :param autocommit_before_ddl: - Default to ``False`` to prevent automatic DDL commits - in transactions and enable rollback. - Applies only to the current transaction and automatically - reverted when the transaction commits or rollback. - - Usage:: - - # Default to ``False`` (``autocommit_before_ddl = off``) - transaction = engine.atomic() - transaction.add(Foo.create_table()) - - # If we want set ``autocommit_before_ddl = on`` - # which is default Cockroach session setting. - transaction = engine.atomic(autocommit_before_ddl=True) - transaction.add(Foo.create_table()) - - """ __slots__ = ("autocommit_before_ddl",) @@ -36,67 +15,38 @@ def __init__( 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:: + # Default 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 run(self): - from piccolo.query.methods.objects import Create, GetOrCreate - - try: - async with self.engine.transaction( - autocommit_before_ddl=self.autocommit_before_ddl - ): - for query in self.queries: - if isinstance(query, (Query, DDL, Create, GetOrCreate)): - await query.run() - else: - raise ValueError("Unrecognised query") - - self.queries = [] - - except Exception as exception: - self.queries = [] - raise exception from exception + async def setup_transaction(self, transaction: CockroachTransaction): + if self.autocommit_before_ddl is not None: + await transaction.autocommit_before_ddl( + enabled=self.autocommit_before_ddl + ) class CockroachTransaction(PostgresTransaction): - """ - :param autocommit_before_ddl: - Default to ``False`` to prevent automatic DDL commits - in transactions and enable rollback. Applies only - to the current transaction and automatically - reverted when the transaction commits or rollback. - - Usage:: - - # Default to ``False`` (``autocommit_before_ddl = off``) - async with engine.transaction(): - # Run some queries: - await Band.select().run() - - # If we want set ``autocommit_before_ddl = on`` - # which is default Cockroach session setting. - async with engine.transaction(autocommit_before_ddl=True): - # Run some queries: - await Band.select().run() - """ - - __slots__ = ("autocommit_before_ddl",) - - def __init__( - self, - autocommit_before_ddl: Optional[bool] = False, - *args, - **kwargs, - ): - super().__init__(*args, **kwargs) - self.autocommit_before_ddl = autocommit_before_ddl - - async def begin(self): - await self.transaction.start() - - value = "on" if self.autocommit_before_ddl else "off" + 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}" ) @@ -138,10 +88,8 @@ def atomic( def transaction( self, allow_nested: bool = True, - autocommit_before_ddl: Optional[bool] = False, ) -> CockroachTransaction: return CockroachTransaction( engine=self, allow_nested=allow_nested, - autocommit_before_ddl=autocommit_before_ddl, ) 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() From f37edcdcb5cb019ae4a7d718b1536a55d2991ff0 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Tue, 10 Feb 2026 23:07:05 +0000 Subject: [PATCH 07/12] arg -> kwarg --- piccolo/apps/migrations/auto/migration_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/piccolo/apps/migrations/auto/migration_manager.py b/piccolo/apps/migrations/auto/migration_manager.py index 6f2e89099..57b97fa2d 100644 --- a/piccolo/apps/migrations/auto/migration_manager.py +++ b/piccolo/apps/migrations/auto/migration_manager.py @@ -988,7 +988,7 @@ async def run(self, backwards: bool = False): ) as transaction: if isinstance(transaction, CockroachTransaction): # To enable DDL rollbacks in CockroachDB. - await transaction.autocommit_before_ddl(False) + await transaction.autocommit_before_ddl(enabled=False) if not self.preview: if direction == "backwards": From abe43e8a886088161085b971741ce18241dc4aeb Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Tue, 10 Feb 2026 23:13:49 +0000 Subject: [PATCH 08/12] tweak docs --- piccolo/engine/cockroach.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/piccolo/engine/cockroach.py b/piccolo/engine/cockroach.py index 1b05d9bd9..25fcbd2be 100644 --- a/piccolo/engine/cockroach.py +++ b/piccolo/engine/cockroach.py @@ -22,15 +22,16 @@ def __init__( transaction and is automatically reverted when the transaction commits or is rolled back. - Usage:: - # Default 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()) + 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) From cba02da6c24a8b2a3ab24217f0a0af4c4928b529 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Tue, 10 Feb 2026 23:20:46 +0000 Subject: [PATCH 09/12] fix linter error --- piccolo/engine/cockroach.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/piccolo/engine/cockroach.py b/piccolo/engine/cockroach.py index 25fcbd2be..d47336382 100644 --- a/piccolo/engine/cockroach.py +++ b/piccolo/engine/cockroach.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Sequence -from typing import Any, Optional +from typing import Any, Optional, cast from .postgres import Atomic, PostgresEngine, PostgresTransaction @@ -37,8 +37,9 @@ def __init__( super().__init__(engine) self.autocommit_before_ddl = autocommit_before_ddl - async def setup_transaction(self, transaction: CockroachTransaction): + 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 ) From 9d1167461fe498519791a4b2c1204e3add0aee4f Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Tue, 10 Feb 2026 23:32:51 +0000 Subject: [PATCH 10/12] fix test --- tests/engine/test_transaction.py | 4 ++++ 1 file changed, 4 insertions(+) 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() From 5f819d90a153cc68f4d38c24a2cf9b7810677e7f Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Tue, 10 Feb 2026 23:49:27 +0000 Subject: [PATCH 11/12] Update tests.yaml --- .github/workflows/tests.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 9ad4a88e4..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: ["v25.4.2"] + cockroachdb-version: ["v25.4.3"] steps: - uses: actions/checkout@v3 - name: Set up Python ${{ matrix.python-version }} From 5c3bbdc0a073967b0e4851b97740f0ed1b581491 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Tue, 10 Feb 2026 23:56:29 +0000 Subject: [PATCH 12/12] more test fixes --- piccolo/testing/test_case.py | 5 +++++ 1 file changed, 5 insertions(+) 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()