Support for nested transactions (savepoints) in cockroachdb.

Allow cockroachdb to expose nested transaction APIs. Requires
cockroachdb version 20.1 or newer.

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

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

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:

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)
 

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:

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