Skip to content

Instantly share code, notes, and snippets.

@tsensei
Created July 27, 2024 18:22
Show Gist options
  • Select an option

  • Save tsensei/979109dc542d463758daecd7b3c39234 to your computer and use it in GitHub Desktop.

Select an option

Save tsensei/979109dc542d463758daecd7b3c39234 to your computer and use it in GitHub Desktop.
Transaction with UnitOfWork pattern coupled with Context Manager in Python
from retrying import retry
class Step:
@retry(stop_max_attempt_number=3, wait_exponential_multiplier=1000, wait_exponential_max=10000)
def execute(self):
raise NotImplementedError
def rollback(self):
raise NotImplementedError
class UnitOfWork:
def __init__(self):
self.steps = []
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
if exc_type:
print(f"Exception occurred: {exc_value}. Rolling back.")
self.rollback()
return False # Do not suppress exceptions
def add_step(self, step):
self.steps.append(step)
def commit(self):
for step in self.steps:
step.execute()
def rollback(self):
for step in reversed(self.steps):
try:
step.rollback()
except Exception as e:
print(f"Rollback failed: {e}")
# Example steps
class LogStep(Step):
def __init__(self, message):
self.message = message
self.executed = False
@retry(stop_max_attempt_number=3, wait_exponential_multiplier=1000, wait_exponential_max=10000)
def execute(self):
print(f"Executing step: {self.message}")
if "error" in self.message:
raise Exception("Deliberate error!")
self.executed = True
def rollback(self):
if self.executed:
print(f"Rolling back step: {self.message}")
# Using the UnitOfWork with a context manager
try:
with UnitOfWork() as uow:
uow.add_step(LogStep("Step 1: Initialize"))
uow.add_step(LogStep("Step 2: Process data"))
uow.add_step(LogStep("Step 3: Process data with error")) # This step will cause an error
uow.add_step(LogStep("Step 4: Finalize"))
uow.commit()
except Exception as e:
print(f"Transaction failed: {e}")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment