From b8a684cec75760fcb47c4005a65c52d21908a057 Mon Sep 17 00:00:00 2001
From: Ethan Paul <24588726+enpaul@users.noreply.github.com>
Date: Thu, 20 Jan 2022 00:59:02 -0500
Subject: [PATCH] Add flat transaction decorator function

---
 peewee_plus.py | 45 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/peewee_plus.py b/peewee_plus.py
index 0107011..d36abc7 100644
--- a/peewee_plus.py
+++ b/peewee_plus.py
@@ -12,7 +12,9 @@
 
 .. _`Peewee documentation`: http://docs.peewee-orm.com/en/latest/peewee/database.html#recommended-settings
 """
+import contextlib
 import enum
+import functools
 import json
 from pathlib import Path
 from typing import Any
@@ -42,6 +44,7 @@ __all__ = [
     "__authors__",
     "calc_batch_size",
     "EnumField",
+    "flat_transaction",
     "JSONField",
     "PathField",
     "PrecisionFloatField",
@@ -131,6 +134,48 @@ def calc_batch_size(
     return len(models)
 
 
+def flat_transaction(interface: peewee.Database):
+    """Database transaction wrapper that avoids nested transactions
+
+    A decorator that can be used to decorate functions or methods so that the entire callable
+    is executed in a single transaction context. If a transaction is already open then it will
+    be reused rather than opening a nested transaction.
+
+    Example usage:
+
+    .. code-block:: python
+
+      db = peewee.SqliteDatabase("test.db")
+
+
+      @flat_transaction(db)
+      def subquery():
+          ...
+
+
+      @flat_transaction(db)
+      def my_query():
+          ...
+          subquery()
+
+
+      # This call opens only a single transaction
+      my_query()
+
+    :param interface: Peewee database interface that should be used to open the transaction
+    """
+
+    def outer(func):
+        @functools.wraps(func)
+        def inner(*args, **kwargs):
+            with interface.atomic() if not interface.in_transaction() else contextlib.nullcontext():
+                return func(*args, **kwargs)
+
+        return inner
+
+    return outer
+
+
 class PathField(peewee.CharField):
     """Field class for storing file paths