Repository: keephq/keep
Pull request: fix: Handle race conditions when updating incidents in database
Add retry logic to manage race conditions caused by StaleDataError during incident updates. Ensures consistent database state by rolling back and retrying up to 3 times if a conflict occurs. Logs informative messages for better traceability of retries.

<!-- 
Thanks for creating this pull request 🤗

Please make sure that the pull request is limited to one type (docs, feature, etc.) and keep it as small as possible. You can open multiple prs instead of opening a huge one.
-->

<!-- If this pull request closes an issue, please mention the issue number below -->
Closes #4587 

## 📑 Description
<!-- Add a brief description of the pr -->

<!-- You can also choose to add a list of changes and if they have been completed or not by using the markdown to-do list syntax
- [ ] Not Completed
- [x] Completed
-->

## ✅ Checks
<!-- Make sure your pr passes the CI checks and do check the following fields as needed - -->
- [ ] My pull request adheres to the code style of this project
- [ ] My code requires changes to the documentation
- [ ] I have updated the documentation as required
- [ ] All the tests have passed

## ℹ Additional Information
<!-- Any additional information like breaking changes, dependencies added, screenshots, comparisons between new and old behavior, etc. -->


Problem statement:
[🐛 Bug]: sqlalchemy.orm.exc.StaleDataError once again
```python
Traceback (most recent call last):
  File "/venv/lib/python3.11/site-packages/keep/api/tasks/process_event_task.py", line 525, in __handle_formatted_events
    incidents: List[IncidentDto] = rules_engine.run_rules(
                                   ^^^^^^^^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/keep/rulesengine/rulesengine.py", line 66, in run_rules
    cel_incidents = self._run_cel_rules(events, session)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/keep/rulesengine/rulesengine.py", line 161, in _run_cel_rules
    ).resolve_incident_if_require(incident)
      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/keep/api/bl/incidents_bl.py", line 438, in resolve_incident_if_require
    if incident.resolve_on == ResolveOn.ALL.value and is_all_alerts_resolved(
                                                      ^^^^^^^^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/keep/api/core/db.py", line 5045, in is_all_alerts_resolved
    return is_all_alerts_in_status(
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/keep/api/core/db.py", line 5117, in is_all_alerts_in_status
    ).scalar()
      ^^^^^^^^
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/query.py", line 2805, in scalar
    ret = self.one()
          ^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/query.py", line 2778, in one
    return self._iter().one()  # type: ignore
           ^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/query.py", line 2827, in _iter
    result: Union[ScalarResult[_T], Result[_T]] = self.session.execute(
                                                  ^^^^^^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/sqlmodel/orm/session.py", line 127, in execute
    return super().execute(
           ^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/session.py", line 2362, in execute
    return self._execute_internal(
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/session.py", line 2226, in _execute_internal
    ) = compile_state_cls.orm_pre_session_exec(
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/context.py", line 561, in orm_pre_session_exec
    session._autoflush()
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/session.py", line 3050, in _autoflush
    self.flush()
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/session.py", line 4352, in flush
    self._flush(objects)
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/session.py", line 4487, in _flush
    with util.safe_reraise():
  File "/venv/lib/python3.11/site-packages/sqlalchemy/util/langhelpers.py", line 146, in __exit__
    raise exc_value.with_traceback(exc_tb)
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/session.py", line 4448, in _flush
    flush_context.execute()
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/unitofwork.py", line 466, in execute
    rec.execute(self)
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/unitofwork.py", line 642, in execute
    util.preloaded.orm_persistence.save_obj(
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/persistence.py", line 85, in save_obj
    _emit_update_statements(
  File "/venv/lib/python3.11/site-packages/sqlalchemy/orm/persistence.py", line 948, in _emit_update_statements
    raise orm_exc.StaleDataError(
sqlalchemy.orm.exc.StaleDataError: UPDATE statement on table 'incident' expected to update 1 row(s); 0 were matched.
```

Hints from issue discussion:
<!-- Answer -->
To address the `sqlalchemy.orm.exc.StaleDataError`, you can implement a retry mechanism. This involves attempting to update the incident status and commit the transaction. If a `StaleDataError` occurs, log the error, roll back the session, and retry the operation up to a specified number of times. Here's a code snippet that demonstrates this approach:

```python
if should_resolve:
    for attempt in range(max_retries):
        try:
            incident.status = IncidentStatus.RESOLVED.value
            self.session.add(incident)
            self.session.commit()
            break
        except StaleDataError as ex:
            if "expected to update" in ex.args[0]:
                self.logger.info(
                    f"Phantom read detected while updating incident `{incident_id}`, retry #{attempt}"
                )
                self.session.rollback()
                continue
            else:
                raise
    self.session.refresh(incident)
```

This method helps handle transient issues causing `StaleDataError` by retrying the operation after rolling back the session. Refreshing the session with `session.refresh(incident)` ensures the incident object is updated with the latest data from the database after a successful commit <sup>[[1]](https://github.com/keephq/keep/blob/main/keep/api/bl/incidents_bl.py)</sup>.


<!-- Dosu Comment Footer -->
*To continue talking to [Dosu](https://dosu.dev), mention @dosu.*

---
<sup>Help Dosu learn! Give it feedback: [Great Response](https://app.dosu.dev/response-feedback/fe25a93d-c760-4920-83f6-d23364bf5ecd?feedback_type=great_response) | [Irrelevant Answer](https://app.dosu.dev/response-feedback/fe25a93d-c760-4920-83f6-d23364bf5ecd?feedback_type=irrelevant_answer) | [Incorrect Sources](https://app.dosu.dev/response-feedback/fe25a93d-c760-4920-83f6-d23364bf5ecd?feedback_type=incorrect_sources) | [Too Verbose](https://app.dosu.dev/response-feedback/fe25a93d-c760-4920-83f6-d23364bf5ecd?feedback_type=too_verbose) | [Hallucination](https://app.dosu.dev/response-feedback/fe25a93d-c760-4920-83f6-d23364bf5ecd?feedback_type=hallucination) | [Bug Report](https://app.dosu.dev/response-feedback/fe25a93d-c760-4920-83f6-d23364bf5ecd?feedback_type=bug_report) | [Other](https://app.dosu.dev/response-feedback/fe25a93d-c760-4920-83f6-d23364bf5ecd?feedback_type=other)</sup>

Commit message: Handle race conditions when updating incidents in database

Add retry logic to manage race conditions caused by StaleDataError during incident updates. Ensures consistent database state by rolling back and retrying up to 3 times if a conflict occurs. Logs informative messages for better traceability of retries.

Patch to review:
diff --git a/keep/rulesengine/rulesengine.py b/keep/rulesengine/rulesengine.py
index 7c90104d5f..42ffa31304 100644
--- a/keep/rulesengine/rulesengine.py
+++ b/keep/rulesengine/rulesengine.py
@@ -9,6 +9,7 @@
 import celpy.celparser
 import celpy.celtypes
 import celpy.evaluation
+from sqlalchemy.orm.exc import StaleDataError
 from sqlmodel import Session
 
 from keep.api.bl.incidents_bl import IncidentBl
@@ -156,6 +157,23 @@ def _run_cel_rules(
 
                                 send_created_event = incident.is_visible
 
+                            incident_id = incident
+                            # Incident might change till this moment
+                            for attempt in range(3):
+                                try:
+                                    session.add(incident)
+                                    session.commit()
+                                    break
+                                except StaleDataError as ex:
+                                    if "expected to update" in ex.args[0]:
+                                        self.logger.info(
+                                            f"Race condition met while updating incident `{incident_id}`, retry #{attempt}"
+                                        )
+                                        session.rollback()
+                                        continue
+                                    else:
+                                        raise
+
                             incident = IncidentBl(
                                 self.tenant_id, session
                             ).resolve_incident_if_require(incident)


Code review comments:
diff_hunk: @@ -156,6 +157,23 @@ def _run_cel_rules(
 
                                 send_created_event = incident.is_visible
 
+                            incident_id = incident
+                            # Incident might change till this moment
+                            for attempt in range(3):
+                                try:
+                                    session.add(incident)
+                                    session.commit()
+                                    break
+                                except StaleDataError as ex:
+                                    if "expected to update" in ex.args[0]:
+                                        self.logger.info(
path: keep/rulesengine/rulesengine.py
text: @user1:
```suggestion
                                        self.logger.warning(
```

diff_hunk: @@ -156,6 +157,23 @@ def _run_cel_rules(
 
                                 send_created_event = incident.is_visible
 
+                            incident_id = incident
path: keep/rulesengine/rulesengine.py
text: @user1:
that's also quite confusing, what does `incident_id = incident` mean?

@author:
it should be incident_id = incident.id

Fixed