async def wrapped(*args, **kwargs):
try:
return await coro(*args, **kwargs)
except psycopg2.errors.UniqueViolation as exc:
raise exceptions.UniqueViolation(constraint_name=exc.diag.constraint_name)
except psycopg2.Error as exc:
raise exceptions.ConnectorException from exc
def defer_job_one(
self, task_name, lock, queueing_lock, args, scheduled_at, queue
) -> JobRow:
if any(
job for job in self.jobs.values() if job["queueing_lock"] == queueing_lock
):
raise exceptions.UniqueViolation(
constraint_name=connector.QUEUEING_LOCK_CONSTRAINT
)
id = next(self.job_counter)
self.jobs[id] = job_row = {
"id": id,
"queue_name": queue,
"task_name": task_name,
"lock": lock,
"queueing_lock": queueing_lock,
"args": args,
"status": "todo",
"scheduled_at": scheduled_at,
"attempts": 0,
}
self.events[id] = []
if scheduled_at:
self.events[id].append({"type": "scheduled", "at": scheduled_at})
self.events[id].append({"type": "deferred", "at": utils.utcnow()})
if self.notify_event:
if "procrastinate_any_queue" in self.notify_channels or (
f"procrastinate_queue#{queue}" in self.notify_channels
):
self.notify_event.set()
return job_row
async def test_store_defer_job_unique_violation_exception_other_constraint(
mocker, job_store, job_factory, connector):
connector.execute_query_one = mocker.Mock(
side_effect=exceptions.UniqueViolation(
constraint_name="some_other_constraint"))
with pytest.raises(exceptions.ConnectorException):
await job_store.defer_job(job=job_factory(task_kwargs={"a": "b"}))
async def test_store_defer_job_unique_violation_exception(
mocker, job_store, job_factory, connector):
connector.execute_query_one = mocker.Mock(
side_effect=exceptions.UniqueViolation(
constraint_name="procrastinate_jobs_queueing_lock_idx"))
with pytest.raises(exceptions.AlreadyEnqueued):
await job_store.defer_job(job=job_factory(task_kwargs={"a": "b"}))
def wrapped(*args, **kwargs):
try:
return func(*args, **kwargs)
except sqlalchemy.exc.IntegrityError as exc:
if isinstance(exc.orig, psycopg2.errors.UniqueViolation):
raise exceptions.UniqueViolation(
constraint_name=exc.orig.diag.constraint_name)
except sqlalchemy.exc.SQLAlchemyError as exc:
raise exceptions.ConnectorException from exc
def defer_job_one(
self,
task_name: str,
lock: Optional[str],
queueing_lock: Optional[str],
args: types.JSONDict,
scheduled_at: Optional[datetime.datetime],
queue: str,
) -> JobRow:
if queueing_lock is not None and any(
job["queueing_lock"] == queueing_lock
and job["status"] == "todo" for job in self.jobs.values()):
from . import manager
raise exceptions.UniqueViolation(
constraint_name=manager.QUEUEING_LOCK_CONSTRAINT)
id = next(self.job_counter)
self.jobs[id] = job_row = {
"id": id,
"queue_name": queue,
"task_name": task_name,
"lock": lock,
"queueing_lock": queueing_lock,
"args": args,
"status": "todo",
"scheduled_at": scheduled_at,
"attempts": 0,
}
self.events[id] = []
if scheduled_at:
self.events[id].append({"type": "scheduled", "at": scheduled_at})
self.events[id].append({"type": "deferred", "at": utils.utcnow()})
if self.notify_event:
if "procrastinate_any_queue" in self.notify_channels or (
f"procrastinate_queue#{queue}" in self.notify_channels):
self.notify_event.set()
return job_row
