Conversation
piccolo/engine/cockroach.py
Outdated
| 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 |
There was a problem hiding this comment.
Is there some way of setting autocommit_before_ddl = off for all sessions, so we don't have to do it for each connection?
There was a problem hiding this comment.
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:
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.- 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
initargument.
I hope that makes sense.
There was a problem hiding this comment.
Could you run it immediately after creating the transaction which wraps the migration code?
There was a problem hiding this comment.
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 hereWrong:
BEGIN;
SET autocommit_before_ddl = off;
-- DDL hereWe 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.
There was a problem hiding this comment.
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)
There was a problem hiding this comment.
Have you found another way to do this?
There was a problem hiding this comment.
@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.
There was a problem hiding this comment.
@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.
There was a problem hiding this comment.
Yes, I think something like that, but I'll double check tomorrow.
There was a problem hiding this comment.
@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.
|
@dantownsend Can you please take a look at this PR code? I have |
Yes, I'll prioritise this over the other issues - sorry about that. |
2 participants
Piccolo CI uses an unmaintained version. This PR upgrades Cockroach to the latest regular release v25.4.2. Also related to #1254