From f476c504739c9e5d4edc3784f8e2e54d32304d37 Mon Sep 17 00:00:00 2001
From: Adam Wallner <adam.wallner@gmail.com>
Date: Sat, 23 Nov 2024 11:54:34 +0100
Subject: [PATCH] You can now raise esorm.error.ConflictError to retry in
 retry_on_conflict decorated functions and methods.

---
 esorm/error.py | 8 ++++++++
 esorm/model.py | 4 ++--
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/esorm/error.py b/esorm/error.py
index 0427468..f95916c 100644
--- a/esorm/error.py
+++ b/esorm/error.py
@@ -31,6 +31,14 @@ class InvalidModelError(Exception):
     """
 
 
+class ConflictError(Exception):
+    """
+    Raised when a conflict occurs.
+
+    You can manually raise this to retry operation with `retry_on_conflict` decorator.
+    """
+
+
 class BulkOperationError(TypedDict):
     """
     A dictionary type to represent an error in a bulk operation response from Elasticsearch.
diff --git a/esorm/model.py b/esorm/model.py
index e8fd541..fb0db5f 100644
--- a/esorm/model.py
+++ b/esorm/model.py
@@ -32,7 +32,7 @@
 from .utils import snake_case, utcnow
 from .aggs import ESAggs, ESAggsResponse
 
-from .error import InvalidResponseError, NotFoundError
+from .error import InvalidResponseError, NotFoundError, ConflictError as ESORMConflictError
 from .esorm import es, get_es_version
 from .query import ESQuery
 from .response import ESResponse
@@ -981,7 +981,7 @@ async def wrapper(*args, **kwargs):
             while True:
                 try:
                     return await func(*args, **kwargs)
-                except ElasticConflictError:
+                except (ElasticConflictError, ESORMConflictError):
                     # Reload the document if it is a method of ESModel
                     if reload_on_conflict and isinstance(args[0], ESModel):
                         await args[0].reload()