Improvements to CRDB documentation.

Prefer the CockroachDatabase.run_transaction() method to the standalone
function in documentation / examples (though both are available).

docs/peewee/crdb.rst

@@ -3,16 +3,24 @@
 Cockroach Database
 ------------------
 
-`CockroachDB <https://www.cockroachlabs.com>` (CRDB) is well supported by
-peewee. The ``playhouse.cockroachdb`` extension module provides the following
-classes and helpers:
+`CockroachDB <https://www.cockroachlabs.com>`_ (CRDB) is well supported by
+peewee.
+
+.. code-block:: python
+
+    from playhouse.cockroachdb import CockroachDatabase
+
+    db = CockroachDatabase('my_app', user='root', host='10.1.0.8')
+
+The ``playhouse.cockroachdb`` extension module provides the following classes
+and helpers:
 
 * :py:class:`CockroachDatabase` - a subclass of :py:class:`PostgresqlDatabase`,
   designed specifically for working with CRDB.
 * :py:class:`PooledCockroachDatabase` - like the above, but implements
   connection-pooling.
-* :py:func:`run_transaction` - runs a function inside a transaction and
-  provides automatic client-side retry logic.
+* :py:meth:`~CockroachDatabase.run_transaction` - runs a function inside a
+  transaction and provides automatic client-side retry logic.
 
 Special field-types that may be useful when using CRDB:
 
@@ -26,11 +34,10 @@ Special field-types that may be useful when using CRDB:
   multi-dimensional arrays).
 
 CRDB is compatible with Postgres' wire protocol and exposes a very similar
-SQL interface, so it is possible (though not recommended) to use
-:py:class:`PostgresqlDatabase` with CRDB. There are a number of reasons for
-this:
+SQL interface, so it is possible (though **not recommended**) to use
+:py:class:`PostgresqlDatabase` with CRDB:
 
-1. CRDB does not support nested transactions, so the
+1. CRDB does not support nested transactions (savepoints), so the
    :py:meth:`~Database.atomic` method has been implemented to enforce this when
    using :py:class:`CockroachDatabase`. For more info :ref:`crdb-transactions`.
 2. CRDB may have subtle differences in field-types, date functions and
@@ -44,12 +51,13 @@ this:
 CRDB Transactions
 ^^^^^^^^^^^^^^^^^
 
-CRDB does not support nested transactions, so the :py:meth:`~Database.atomic`
-method on the :py:class:`CockroachDatabase` has been modified to raise an
-exception if an invalid nesting is encountered. If you would like to be able to
-nest transactional code, you can use the :py:meth:`~Database.transaction`
-method, which will ensure that the outer-most block will manage the
-transaction (e.g., exiting a nested-block will not cause an early commit).
+CRDB does not support nested transactions (savepoints), so the
+:py:meth:`~Database.atomic` method on the :py:class:`CockroachDatabase` has
+been modified to raise an exception if an invalid nesting is encountered. If
+you would like to be able to nest transactional code, you can use the
+:py:meth:`~Database.transaction` method, which will ensure that the outer-most
+block will manage the transaction (e.g., exiting a nested-block will not cause
+an early commit).
 
 Example:
 
@@ -63,9 +71,10 @@ Example:
         with db.transaction() as txn:
             # do some stuff...
 
-            # This function is wrapped in a transaction, but it will be
-            # ignored, as we are already in a wrapped-block (via the context
-            # manager).
+            # This function is wrapped in a transaction, but the nested
+            # transaction will be ignored and folded into the outer
+            # transaction, as we are already in a wrapped-block (via the
+            # context manager).
             create_user('[email protected]')
 
             # do other stuff.
@@ -76,16 +85,40 @@ Example:
 
 
 CRDB provides client-side transaction retries, which are available using a
-special :py:func:`run_transaction` helper. This helper function accepts a
-database instance as well as a callable, which is responsible for executing any
-transactional statements that may need to be retried.
+special :py:meth:`~CockroachDatabase.run_transaction` helper. This helper
+method accepts a callable, which is responsible for executing any transactional
+statements that may need to be retried.
 
-Example of using :py:func:`run_transaction` to implement client-side retries
-for a transaction that transfers an amount from one account to another:
+Simplest possible example of :py:meth:`~CockroachDatabase.run_transaction`:
 
 .. code-block:: python
 
-    from playhouse.cockroachdb import run_transaction
+    def create_user(email):
+        # Callable that accepts a single argument (the database instance) and
+        # which is responsible for executing the transactional SQL.
+        def callback(db_ref):
+            return User.create(email=email)
+
+        return db.run_transaction(callback, max_attempts=10)
+
+    huey = create_user('[email protected]')
+
+.. note::
+    The ``cockroachdb.ExceededMaxAttempts`` exception will be raised if the
+    transaction cannot be committed after the given number of attempts. If the
+    SQL is mal-formed, violates a constraint, etc., then the function will
+    raise the exception to the caller.
+
+Example of using :py:meth:`~CockroachDatabase.run_transaction` to implement
+client-side retries for a transaction that transfers an amount from one account
+to another:
+
+.. code-block:: python
+
+    from playhouse.cockroachdb import CockroachDatabase
+
+    db = CockroachDatabase('my_app')
+
 
     def transfer_funds(from_id, to_id, amt):
         """
@@ -119,7 +152,7 @@ for a transaction that transfers an amount from one account to another:
         # Perform the queries that comprise a logical transaction. In the
         # event the transaction fails due to contention, it will be auto-
         # matically retried (up to 10 times).
-        return run_transaction(db, thunk, max_attempts=10)
+        return db.run_transaction(thunk, max_attempts=10)
 
 CRDB APIs
 ^^^^^^^^^
@@ -133,6 +166,45 @@ CRDB APIs
     constructor, and may be used to specify the database ``user``, ``port``,
     etc.
 
+    .. py:method:: run_transaction(callback[, max_attempts=None[, system_time=None[, priority=None]]])
+
+        :param CockroachDatabase db: database instance.
+        :param callback: callable that accepts a single ``db`` parameter (which
+            will be the same as the value passed above).
+        :param int max_attempts: max number of times to try before giving up.
+        :param datetime system_time: execute the transaction ``AS OF SYSTEM TIME``
+            with respect to the given value.
+        :param str priority: either "low", "normal" or "high".
+        :return: returns the value returned by the callback.
+        :raises: ``ExceededMaxAttempts`` if ``max_attempts`` is exceeded.
+
+        Run SQL in a transaction with automatic client-side retries.
+
+        User-provided ``callback``:
+
+        * **Must** accept one parameter, the ``db`` instance representing the
+          connection the transaction is running under.
+        * **Must** not attempt to commit, rollback or otherwise manage the
+          transaction.
+        * **May** be called more than one time.
+        * **Should** ideally only contain SQL operations.
+
+        Additionally, the database must not have any open transactions at the
+        time this function is called, as CRDB does not support nested
+        transactions. Attempting to do so will raise a ``NotImplementedError``.
+
+        Simplest possible example:
+
+        .. code-block:: python
+
+            def create_user(email):
+                def callback(db_ref):
+                    return User.create(email=email)
+
+                return db.run_transaction(callback, max_attempts=10)
+
+            user = create_user('[email protected]')
+
 .. py:class:: PooledCockroachDatabase(database[, **kwargs])
 
     CockroachDB connection-pooling implementation, based on
@@ -141,30 +213,12 @@ CRDB APIs
 
 .. py:function:: run_transaction(db, callback[, max_attempts=None[, system_time=None[, priority=None]]])
 
-    :param CockroachDatabase db: database instance.
-    :param callback: callable that accepts a single ``db`` parameter (which
-        will be the same as the value passed above).
-    :param int max_attempts: max number of times to try before giving up.
-    :param datetime system_time: execute the transaction ``AS OF SYSTEM TIME``
-        with respect to the given value.
-    :param str priority: either "low", "normal" or "high".
-    :return: returns the value returned by the callback.
-    :raises: ``ExceededMaxAttempts`` if ``max_attempts`` is exceeded.
-
-    Run transactional SQL in a transaction with automatic retries.
-
-    User-provided ``callback``:
-
-    * **Must** accept one parameter, the ``db`` instance representing the
-      connection the transaction is running under.
-    * **Must** not attempt to commit, rollback or otherwise manage the
-      transaction.
-    * **May** be called more than one time.
-    * **Should** ideally only contain SQL operations.
-
-    Additionally, the database must not have any open transactions at the time
-    this function is called, as CRDB does not support nested transactions.
-    Attempting to do so will raise a ``NotImplementedError``.
+    Run SQL in a transaction with automatic client-side retries. See
+    :py:meth:`CockroachDatabase.run_transaction` for details.
+
+    .. note::
+        This function is equivalent to the identically-named method on
+        the :py:class:`CockroachDatabase` class.
 
 .. py:class:: UUIDKeyField()
 

docs/peewee/database.rst

@@ -193,47 +193,30 @@ class, defined in ``playhouse.cockroachdb``:
 
     db = CockroachDatabase('my_app', user='root', port=26257, host='localhost')
 
+CRDB provides client-side transaction retries, which are available using a
+special :py:meth:`CockroachDatabase.run_transaction` helper-method. This method
+accepts a callable, which is responsible for executing any transactional
+statements that may need to be retried.
 
-The ``playhouse.cockroachdb`` module also contains a :py:func:`run_transaction`
-helper for running a transaction with client-side retry logic:
+Simplest possible example of :py:meth:`~CockroachDatabase.run_transaction`:
 
 .. code-block:: python
 
-    from playhouse.cockroachdb import run_transaction
-
-    def transfer_funds(from_id, to_id, amt):
-        """
-        Returns a 3-tuple of (success?, from balance, to balance). If there are
-        not sufficient funds, then the original balances are returned.
-        """
-        def thunk(db_ref):
-            src, dest = (Account
-                         .select()
-                         .where(Account.id.in_([from_id, to_id])))
-            if src.id != from_id:
-                src, dest = dest, src  # Swap order.
-
-            # Cannot perform transfer, insufficient funds!
-            if src.balance < amt:
-                return False, src.balance, dest.balance
-
-            # Update each account, returning the new balance.
-            src, = (Account
-                    .update(balance=Account.balance - amt)
-                    .where(Account.id == from_id)
-                    .returning(Account.balance)
-                    .execute())
-            dest, = (Account
-                     .update(balance=Account.balance + amt)
-                     .where(Account.id == to_id)
-                     .returning(Account.balance)
-                     .execute())
-            return True, src.balance, dest.balance
-
-        # Perform the queries that comprise a logical transaction. In the
-        # event the transaction fails due to contention, it will be auto-
-        # matically retried (up to 10 times).
-        return run_transaction(db, thunk, max_attempts=10)
+    def create_user(email):
+        # Callable that accepts a single argument (the database instance) and
+        # which is responsible for executing the transactional SQL.
+        def callback(db_ref):
+            return User.create(email=email)
+
+        return db.run_transaction(callback, max_attempts=10)
+
+    huey = create_user('[email protected]')
+
+.. note::
+    The ``cockroachdb.ExceededMaxAttempts`` exception will be raised if the
+    transaction cannot be committed after the given number of attempts. If the
+    SQL is mal-formed, violates a constraint, etc., then the function will
+    raise the exception to the caller.
 
 For more information, see:
 

playhouse/cockroachdb.py

@@ -155,6 +155,10 @@ def new_fn():
             return new_fn
         return deco
 
+    def run_transaction(self, cb, max_attempts=None, system_time=None,
+                        priority=None):
+        return run_transaction(self, cb, max_attempts, system_time, priority)
+
 
 class _crdb_atomic(_atomic):
     def __enter__(self):

tests/cockroachdb.py

@@ -10,6 +10,7 @@
 from .base import db
 from .base import requires_models
 from .base import skip_unless
+from .base_models import User
 from .postgres_helpers import BaseBinaryJsonFieldTestCase
 
 
@@ -155,6 +156,19 @@ def insert_row(db):
                               self.database, insert_row)
         self.assertEqual(KV.select().count(), 0)
 
+    @requires_models(User)
+    def test_retry_transaction_docs_example(self):
+        def create_user(username):
+            def thunk(db_ref):
+                return User.create(username=username)
+            return self.database.run_transaction(thunk, max_attempts=5)
+
+        users = [create_user(u) for u in 'abc']
+        self.assertEqual([u.username for u in users], ['a', 'b', 'c'])
+
+        query = User.select().order_by(User.username)
+        self.assertEqual([u.username for u in query], ['a', 'b', 'c'])
+
     @requires_models(KV)
     def test_retry_transaction_decorator(self):
         @self.database.retry_transaction()