Skip to content

Cockroach upgrade #1318

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
sinisaos wants to merge 8 commits into from
Closed

Cockroach upgrade #1318

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
c92eead
upgrade crdb to the latest regular release
sinisaos Dec 30, 2025
3b8e6bf
remove cockroach_upgrade branch from CI
sinisaos Dec 30, 2025
6f33b8a
settings applied to the current transaction only
sinisaos Jan 7, 2026
7878492
Merge branch 'piccolo-orm:master' into cockroach_upgrade
sinisaos Jan 10, 2026
5be8d40
Merge branch 'piccolo-orm:master' into cockroach_upgrade
sinisaos Jan 12, 2026
3e9300d
make changes to transaction and atomic
sinisaos Jan 16, 2026
a420277
remove postgres password
sinisaos Jan 16, 2026
f19e3be
Merge branch 'piccolo-orm:master' into cockroach_upgrade
sinisaos Feb 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/tests.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}
Expand Down
29 changes: 15 additions & 14 deletions piccolo/engine/cockroach.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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):
"""
Expand All @@ -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
Copy link
Copy Markdown
Member

@dantownsend dantownsend Dec 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there some way of setting autocommit_before_ddl = off for all sessions, so we don't have to do it for each connection?

Copy link
Copy Markdown
Member Author

@sinisaos sinisaos Dec 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's a good point, but CockroachDB doesn't have a built-in setting to globally configure autocommit_before_ddl for all connections by default. We have two options:

  1. get_new_connection() like the code in this PR which works because it runs right when Piccolo acquires a usable session, just before transactions and migrations.
  2. Use connection pooling with an asyncpg init argument, but that doesn't work because we can't guarantee that every connection used in a transaction goes through that init argument.

I hope that makes sense.

Copy link
Copy Markdown
Member

@dantownsend dantownsend Dec 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you run it immediately after creating the transaction which wraps the migration code?

Copy link
Copy Markdown
Member Author

@sinisaos sinisaos Dec 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

CockroachDB requires that autocommit_before_ddl be set before a transaction exists, so we cannot set it after the transaction is created.

Correct:

SET autocommit_before_ddl = off;
BEGIN;
-- DDL here

Wrong:

BEGIN;
SET autocommit_before_ddl = off;
-- DDL here

We can’t make SET autocommit before ddl = off apply to all sessions globally, because it’s a session variable, not a cluster setting. I haven't found any other way other than overriding get_new_connection because that sets the session variable before the transaction and DDL. That's the only way I can pass this transaction tests. If you have a better way, feel free to close this PR.

Copy link
Copy Markdown
Member Author

@sinisaos sinisaos Jan 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If this PR is merged, this PR can also be merged because that PR works without issues with this PR changes (works with the latest Cockroach regular release v25.4.2)

Copy link
Copy Markdown
Member Author

@sinisaos sinisaos Jan 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Have you found another way to do this?

Copy link
Copy Markdown
Member

@dantownsend dantownsend Jan 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@sinisaos I was thinking of having an argument for transactions and Atomic which can be used to turn this behaviour on or off (off by default), and we turn it on for migrations and those Atomic tests.

Copy link
Copy Markdown
Member Author

@sinisaos sinisaos Jan 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@dantownsend Do you have something like this in mind? If autocommit_before_ddl is None (or True), the Cockroach defaults apply, otherwise we disable it.

cockroach_engine 
from __future__ import annotations

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):
    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):
    def __init__(
        self,
        autocommit_before_ddl: Optional[bool] = None,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.autocommit_before_ddl = autocommit_before_ddl

    async def begin(self):
        await self.transaction.start()

        # we use None to leave the default Cockroach behavior
        if self.autocommit_before_ddl is not None:
            value = "on" if self.autocommit_before_ddl else "off"
            await self.connection.execute(
                f"SET LOCAL autocommit_before_ddl = {value}"
            )


class CockroachEngine(PostgresEngine):
    """
    An extension of
    :class:`PostgresEngine <piccolo.engine.postgres.PostgresEngine>`.
    """

    def __init__(
        self,
        config: dict[str, Any],
        extensions: Sequence[str] = (),
        log_queries: bool = False,
        log_responses: bool = False,
        extra_nodes: Optional[dict[str, CockroachEngine]] = None,
    ) -> None:
        super().__init__(
            config=config,
            extensions=extensions,
            log_queries=log_queries,
            log_responses=log_responses,
            extra_nodes=extra_nodes,
        )
        self.engine_type = "cockroach"
        self.min_version_number = 0

    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] = None,
    ) -> CockroachTransaction:
        return CockroachTransaction(
            engine=self,
            allow_nested=allow_nested,
            autocommit_before_ddl=autocommit_before_ddl,
        )

Of course, we also need to modify the tests where we want to disable the default behavior.

Copy link
Copy Markdown
Member

@dantownsend dantownsend Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I think something like that, but I'll double check tomorrow.

Copy link
Copy Markdown
Member Author

@sinisaos sinisaos Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@dantownsend I made another commit where I made small changes (the default is always False for both transaction and atomic so we don't have to change the tests) that you can easily try out. I hope that's better.