From 82c613833e7255d80e29c8faf36d144088b54ce1 Mon Sep 17 00:00:00 2001 From: Charles Leifer Date: Wed, 17 Jun 2020 15:55:19 +0000 Subject: [PATCH] Support for nested transactions (savepoints) in cockroachdb. Allow cockroachdb to expose nested transaction APIs. Requires cockroachdb version 20.1 or newer. --- .travis.yml | 6 +++--- playhouse/cockroachdb.py | 15 +++++++++------ tests/base.py | 8 ++++++++ tests/cockroachdb.py | 2 +- tests/fields.py | 2 +- tests/transactions.py | 4 +++- 6 files changed, 25 insertions(+), 12 deletions(-) diff --git a/.travis.yml b/.travis.yml index e67b5f6c..7f1897b3 100644 --- a/.travis.yml +++ b/.travis.yml @@ -37,9 +37,9 @@ matrix: env: - PEEWEE_TEST_BACKEND=cockroachdb before_install: - - wget -qO- https://binaries.cockroachdb.com/cockroach-v19.2.0.linux-amd64.tgz | tar xvz - - ./cockroach-v19.2.0.linux-amd64/cockroach start --insecure --background - - ./cockroach-v19.2.0.linux-amd64/cockroach sql --insecure -e 'create database peewee_test;' + - wget -qO- https://binaries.cockroachdb.com/cockroach-v20.1.1.linux-amd64.tgz | tar xvz + - ./cockroach-v20.1.1.linux-amd64/cockroach start --insecure --background + - ./cockroach-v20.1.1.linux-amd64/cockroach sql --insecure -e 'create database peewee_test;' allow_failures: addons: postgresql: "9.6" diff --git a/playhouse/cockroachdb.py b/playhouse/cockroachdb.py index afb11076..01988d13 100644 --- a/playhouse/cockroachdb.py +++ b/playhouse/cockroachdb.py @@ -4,7 +4,6 @@ import re from peewee import * from peewee import _atomic from peewee import _manual -from peewee import _transaction from peewee import ColumnMetadata # (name, data_type, null, primary_key, table, default) from peewee import ForeignKeyMetadata # (column, dest_table, dest_column, table). from peewee import IndexMetadata @@ -18,6 +17,8 @@ except ImportError: # psycopg2 not installed, ignore. ArrayField = BinaryJSONField = IntervalField = JSONField = None +NESTED_TX_MIN_VERSION = 200100 + TXN_ERR_MSG = ('CockroachDB does not support nested transactions. You may ' 'alternatively use the @transaction context-manager/decorator, ' 'which only wraps the outer-most block in transactional logic. ' @@ -56,6 +57,7 @@ class CockroachDatabase(PostgresqlDatabase): for_update = False nulls_ordering = False + release_after_rollback = True def __init__(self, *args, **kwargs): kwargs.setdefault('user', 'root') @@ -137,13 +139,14 @@ class CockroachDatabase(PostgresqlDatabase): commit=False) def atomic(self, system_time=None, priority=None): - return _crdb_atomic(self, system_time, priority) - - def transaction(self, system_time=None, priority=None): - return _transaction(self, system_time, priority) + if self.server_version < NESTED_TX_MIN_VERSION: + return _crdb_atomic(self, system_time, priority) + return super(CockroachDatabase, self).atomic(system_time, priority) def savepoint(self): - raise NotImplementedError(TXN_ERR_MSG) + if self.server_version < NESTED_TX_MIN_VERSION: + raise NotImplementedError(TXN_ERR_MSG) + return super(CockroachDatabase, self).savepoint() def retry_transaction(self, max_attempts=None, system_time=None, priority=None): diff --git a/tests/base.py b/tests/base.py index c7b782e9..91ee8700 100644 --- a/tests/base.py +++ b/tests/base.py @@ -14,6 +14,7 @@ from peewee import * from peewee import sqlite3 from playhouse.mysql_ext import MySQLConnectorDatabase from playhouse.cockroachdb import CockroachDatabase +from playhouse.cockroachdb import NESTED_TX_MIN_VERSION logger = logging.getLogger('peewee') @@ -109,6 +110,13 @@ if IS_MYSQL: if not IS_MYSQL_ADVANCED_FEATURES: logger.warning('MySQL too old to test certain advanced features.') +if IS_CRDB: + db.connect() + IS_CRDB_NESTED_TX = db.server_version >= NESTED_TX_MIN_VERSION + db.close() +else: + IS_CRDB_NESTED_TX = False + class TestModel(Model): class Meta: diff --git a/tests/cockroachdb.py b/tests/cockroachdb.py index a3bc759b..021125bb 100644 --- a/tests/cockroachdb.py +++ b/tests/cockroachdb.py @@ -152,7 +152,7 @@ class TestCockroachDatabase(ModelTestCase): KV.create(k='k1', v=1) with self.database.atomic(): - self.assertRaises(NotImplementedError, run_transaction, + self.assertRaises(Exception, run_transaction, self.database, insert_row) self.assertEqual(KV.select().count(), 0) diff --git a/tests/fields.py b/tests/fields.py index 6f7e899e..80881fe2 100644 --- a/tests/fields.py +++ b/tests/fields.py @@ -243,7 +243,7 @@ class TestDateFields(ModelTestCase): .tuples()) row, = query - if IS_SQLITE or IS_MYSQL or IS_CRDB: + if IS_SQLITE or IS_MYSQL: self.assertEqual(row, (2011, 1, 2, 11, 12, 13, 2012, 2, 3, 3, 13, 37)) else: diff --git a/tests/transactions.py b/tests/transactions.py index da65dd42..8ff01290 100644 --- a/tests/transactions.py +++ b/tests/transactions.py @@ -2,6 +2,7 @@ from peewee import * from .base import DatabaseTestCase from .base import IS_CRDB +from .base import IS_CRDB_NESTED_TX from .base import IS_SQLITE from .base import ModelTestCase from .base import db @@ -23,7 +24,8 @@ class BaseTransactionTestCase(ModelTestCase): def requires_nested(fn): - return skip_if(IS_CRDB, 'nested transaction support is required')(fn) + return skip_if(IS_CRDB and not IS_CRDB_NESTED_TX, + 'nested transaction support is required')(fn) class TestTransaction(BaseTransactionTestCase):