Add engine_callback to database backend SessionManager#10135
Add engine_callback to database backend SessionManager#10135avolant wants to merge 8 commits intocelery:mainfrom
engine_callback to database backend SessionManager#10135Conversation
Add a new `engine_callback` option to the database result backend that allows users to customize SQLAlchemy engines after creation. The callback is invoked on every engine creation, including after invalidation. This enables use cases like JWT-based authentication where a `do_connect` event listener needs to be registered on each engine to inject fresh credentials before every physical database connection. Supported via config (`database_engine_callback`) as a dotted path or callable, and via constructor kwarg (`engine_callback`).
|
Thanks for the PR @avolant
Can we check it in |
|
Since it’s user-facing, it’d be great if we updated the docs in this PR as well. |
There was a problem hiding this comment.
Pull request overview
This PR adds an engine_callback feature to the database result backend's SessionManager, enabling users to customize SQLAlchemy engines after creation. This is particularly useful for scenarios like JWT-based authentication where event listeners need to be registered on each engine to inject fresh credentials.
Changes:
- Added
engine_callbackparameter toSessionManager.__init__()and invocation logic inget_engine()for both forked and non-forked modes - Integrated callback configuration in
DatabaseBackend.__init__()with support for dotted-path strings and callables from config or constructor kwargs - Added
database_engine_callbackconfiguration option in app defaults
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| celery/backends/database/session.py | Added engine_callback parameter to SessionManager and invocation logic in both forked/non-forked engine creation paths |
| celery/backends/database/init.py | Integrated engine_callback from config or kwargs, with support for dotted-path strings via symbol_by_name |
| celery/app/defaults.py | Added engine_callback configuration option to the database namespace |
| t/unit/backends/test_database.py | Comprehensive test suite covering callback from config/kwargs, dotted paths, fork behavior, invalidation, caching, ordering, and do_connect integration |
|
Good point, added a |
|
Added docs for |
|
Merge conflicts resolved in the latest push (the main conflict was around
There is no existing issue (But I can create one if you need me to). |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #10135 +/- ##
=======================================
Coverage 87.78% 87.79%
=======================================
Files 153 153
Lines 19432 19444 +12
Branches 2233 2237 +4
=======================================
+ Hits 17058 17070 +12
Misses 2076 2076
Partials 298 298
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. |
|
@ChickenBenny hey, would you mind reviewing this PR and let us know if this is needed? |
|
This feature is valuable for modern cloud deployments. It solves the pain point of dynamic database authentication (like AWS RDS IAM or HashiCorp Vault short-lived tokens) by allowing us to hook into SQLAlchemy's LGTM from my side. Minor notes:
|
|
Added a test for the Also opened #10148 for traceability. |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.
Comments suppressed due to low confidence (1)
t/unit/backends/test_database.py:784
- This docstring says the
do_connecthandler “modifies connection params”, but the test only recordscparamswithout changing them. Consider either updating the handler/assertions to verify an actual modification, or adjusting the docstring to reflect what’s being tested (that the event fires).
"""Integration test: engine_callback registers a do_connect event
that modifies connection params on every new connection."""
connect_calls = []
| engine = self._engines[dburi] = create_engine(dburi, **kwargs) | ||
| if self.engine_callback: | ||
| self.engine_callback(engine) | ||
| return engine |
There was a problem hiding this comment.
In forked mode, the new engine is inserted into self._engines before the callback runs. If engine_callback(engine) raises, a partially configured engine remains cached and subsequent calls will return it without reapplying the callback. Consider only caching the engine after the callback succeeds, or popping it from the cache on exception.
Also, if self.engine_callback: can skip valid callables with falsy truthiness; checking self.engine_callback is not None would be more robust.
| engine = create_engine(dburi, poolclass=NullPool, **kwargs) | ||
| if self.engine_callback: | ||
| self.engine_callback(engine) |
There was a problem hiding this comment.
In the non-forked code path, if self.engine_callback: can skip valid callables that define falsy truthiness (e.g., callable objects implementing __bool__). Using an explicit is not None check would avoid surprising behavior.
| .. setting:: database_engine_callback | ||
|
|
||
| ``database_engine_callback`` | ||
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ |
There was a problem hiding this comment.
versionadd:: 5.7 annotation is mandatory for this
Co-authored-by: Copilot <[email protected]>
Add a new
engine_callbackoption to the database result backend that allows users to customize SQLAlchemy engines after creation. The callback is invoked on every engine creation, including after invalidation.This enables use cases like JWT-based authentication where a
do_connectevent listener needs to be registered on each engine to inject fresh credentials before every physical database connection.Supported via config (
database_engine_callback) as a dotted path or callable, and via constructor kwarg (engine_callback).Note: Before submitting this pull request, please review our contributing
guidelines.
Description