feat(backend): add sqlalchemy backend #20

dadodimauro · 2025-12-30T12:58:12Z

Summary by CodeRabbit

New Features
- Added support for PostgreSQL, MySQL, SQLite and MSSQL result backends alongside Redis
- Database-backed result storage, worker heartbeats, healthchecks and maintenance operations
Documentation
- New comprehensive Postgres example demonstrating tasks, retries, scheduling, batch processing, DLQ and monitoring
- Examples updated to use context-driven retry tracking for flaky tasks
Chores
- Optional dependency groups for database drivers; docker-compose now includes a PostgreSQL service
Tests
- Test matrix extended to cover database optional extras and exports

_{✏️ Tip: You can customize this high-level summary in your review settings.}

coderabbitai · 2025-12-30T12:58:22Z

📝 Walkthrough

Walkthrough

Adds a new SQLAlchemy-backed DatabaseBackend and ORM models, multi-database config and types, optional DB dependency groups, example Postgres demo, docker-compose Postgres service, and updates app/backend exports and backend selection to support Postgres/MySQL/SQLite/MSSQL.

Changes

Cohort / File(s)	Summary
Database backend implementation `src/chicory/backend/database.py`, `src/chicory/backend/models.py`	New DatabaseBackend: async engine/session lifecycle, connect/disconnect, upsert/query/delete TaskResult and heartbeats, healthcheck and cleanup; ORM models TaskResultModel and WorkerHeartbeatModel with shared metadata and naming convention.
Configuration & types `src/chicory/config.py`, `src/chicory/types.py`, `src/chicory/exceptions.py`	Adds Postgres/MySQL/MSSQL/SQLite BackendConfig classes and BackendConfig fields; extends BackendType enum with POSTGRES/MYSQL/MSSQL/SQLITE; adds DbPoolExhaustedException.
App integration & exports `src/chicory/__init__.py`, `src/chicory/app.py`, `src/chicory/backend/__init__.py`	Re-exports DatabaseBackend when available; app._create_backend now returns DatabaseBackend for new backend types; backend package exports DatabaseBackend/Config guarded by ImportError.
Project deps & compose `pyproject.toml`, `docker-compose.yml`	Adds optional dependency groups (postgres/mysql/sqlite/mssql) with asyncio SQLAlchemy + drivers and includes them in `all`; docker-compose adds Postgres service and postgres_data volume.
Examples & tests `examples/postgres_example.py`, `examples/rabbitmq_example.py`, `examples/redis_example.py`, `tests/unit/test_optional_imports.py`	New comprehensive Postgres example demo; flaky_api_call signature changed to accept TaskContext in redis/rabbitmq examples and call sites updated; tests updated to include DatabaseBackend in optional-extras matrix.

Sequence Diagram(s)

sequenceDiagram
    actor App
    participant Backend as DatabaseBackend
    participant Engine as SQLAlchemy Engine
    participant DB as Database

    Note over App,DB: DatabaseBackend connect / store / retrieve / disconnect flows

    rect rgb(220,235,255)
    Note over App,Backend: Connect phase
    App->>Backend: connect()
    Backend->>Engine: create_async_engine(_get_engine_kwargs)
    Engine->>DB: establish connection
    Engine-->>Backend: engine ready
    Backend->>Engine: create async sessionmaker
    Backend->>DB: create_all(metadata)
    Backend-->>App: connected
    end

    rect rgb(210,255,220)
    Note over App,Backend: Store result
    App->>Backend: store_result(task_id, result)
    Backend->>Engine: acquire session
    Engine->>DB: checkout connection
    Backend->>DB: upsert TaskResultModel (INSERT/UPDATE)
    DB-->>Backend: persisted
    Engine->>DB: release connection
    Backend-->>App: ack
    end

    rect rgb(255,245,220)
    Note over App,Backend: Retrieve result
    App->>Backend: get_result(task_id)
    Backend->>Engine: acquire session
    Engine->>DB: query TaskResultModel
    DB-->>Backend: record
    Backend->>Backend: deserialize JSON -> TaskResult
    Engine->>DB: release connection
    Backend-->>App: TaskResult | None
    end

    rect rgb(255,220,230)
    Note over App,Backend: Disconnect
    App->>Backend: disconnect()
    Backend->>Engine: dispose()
    Engine->>DB: close connections
    Backend-->>App: disconnected
    end

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Possibly related PRs

feat: refactor env variable configuration #6: Overlaps on backend export/config changes (modifies chicory.init, backend exports, and backend configuration wiring).

Suggested reviewers

anvouk

Poem

🐰 I hopped in with engine dreams,
New tables, heartbeats, JSON streams,
Postgres, MySQL, SQLite beams,
MSSQL joins the schemes —
Chicory hums with backend gleams 🎉

Pre-merge checks and finishing touches

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately and concisely describes the main change: adding SQLAlchemy-based database backend support to Chicory, which is the core feature across all modified files.
Docstring Coverage	✅ Passed	Docstring coverage is 86.27% which is sufficient. The required threshold is 80.00%.

✨ Finishing touches

📝 Generate docstrings

📜 Recent review details

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a5eb3d6 and f4f6397.

📒 Files selected for processing (3)

src/chicory/app.py
src/chicory/backend/__init__.py
src/chicory/backend/models.py

🚧 Files skipped from review as they are similar to previous changes (1)

src/chicory/backend/models.py

🧰 Additional context used

🧬 Code graph analysis (2)

src/chicory/backend/__init__.py (1)

src/chicory/backend/database.py (1)

DatabaseBackend (40-336)

src/chicory/app.py (2)

src/chicory/types.py (1)

BackendType (239-244)

src/chicory/backend/database.py (1)

DatabaseBackend (40-336)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (12)

GitHub Check: run (macos-latest, 3.12)
GitHub Check: run (windows-latest, 3.14)
GitHub Check: run (windows-latest, 3.11)
GitHub Check: run (windows-latest, 3.13)
GitHub Check: run (windows-latest, 3.12)
GitHub Check: run (ubuntu-latest, 3.11)
GitHub Check: run (ubuntu-latest, 3.13)
GitHub Check: run (macos-latest, 3.13)
GitHub Check: run (ubuntu-latest, 3.14)
GitHub Check: run (ubuntu-latest, 3.12)
GitHub Check: run (macos-latest, 3.14)
GitHub Check: run (macos-latest, 3.11)

🔇 Additional comments (2)

src/chicory/backend/__init__.py (1)

15-20: LGTM! Consistent optional import pattern.

The implementation correctly mirrors the existing RedisBackend import pattern, using a try-except block to conditionally expose DatabaseBackend and DatabaseBackendConfig when the optional database dependencies are installed. Both symbols are properly defined in database.py: DatabaseBackend as a class (line 40) and DatabaseBackendConfig as a TypeAlias (line 32).

src/chicory/app.py (1)

101-115: Well-structured database backend initialization.

The implementation correctly:

Uses pattern matching to handle multiple database backend types (POSTGRES, MYSQL, SQLITE, MSSQL)

Maps each backend type to its corresponding configuration from self.config.backend (postgres, mysql, sqlite, mssql attributes all exist with correct types)

Follows the existing pattern of local imports within match cases

Not passing consumer_name to DatabaseBackend is appropriate here, as the app-level backend will use the default worker name generated in the constructor.

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Actionable comments posted: 9

🧹 Nitpick comments (6)

src/chicory/app.py (1)
101-116: Consider consolidating the DatabaseBackend cases to reduce duplication.

The four database backend cases share the same import statement. You could consolidate them while maintaining lazy loading:
🔎 Proposed refactor
             case BackendType.POSTGRES:
                 from chicory.backend import DatabaseBackend

                 return DatabaseBackend(config=self.config.backend.postgres)
-            case BackendType.MYSQL:
-                from chicory.backend import DatabaseBackend
-
-                return DatabaseBackend(config=self.config.backend.mysql)
-            case BackendType.SQLITE:
-                from chicory.backend import DatabaseBackend
-
-                return DatabaseBackend(config=self.config.backend.sqlite)
-            case BackendType.MSSQL:
-                from chicory.backend import DatabaseBackend
-
-                return DatabaseBackend(config=self.config.backend.mssql)
+            case BackendType.MYSQL:
+                from chicory.backend import DatabaseBackend
+
+                return DatabaseBackend(config=self.config.backend.mysql)
+            case BackendType.SQLITE:
+                from chicory.backend import DatabaseBackend
+
+                return DatabaseBackend(config=self.config.backend.sqlite)
+            case BackendType.MSSQL:
+                from chicory.backend import DatabaseBackend
+
+                return DatabaseBackend(config=self.config.backend.mssql)
Alternatively, for more significant consolidation:
case BackendType.POSTGRES | BackendType.MYSQL | BackendType.SQLITE | BackendType.MSSQL:
    from chicory.backend import DatabaseBackend
    
    config_map = {
        BackendType.POSTGRES: self.config.backend.postgres,
        BackendType.MYSQL: self.config.backend.mysql,
        BackendType.SQLITE: self.config.backend.sqlite,
        BackendType.MSSQL: self.config.backend.mssql,
    }
    return DatabaseBackend(config=config_map[backend_type])
The current implementation is functionally correct; this is purely an optional cleanup for maintainability.
tests/unit/test_optional_imports.py (1)
138-178: Remove unused # noqa: E501 directives.

Static analysis indicates these noqa directives are unused because E501 (line-too-long) is not enabled in the ruff configuration. Remove them to keep the code clean.
🔎 Lines to update

Remove # noqa: E501 from lines: 138, 147, 156, 165, 171, 172, 175, 176, 177, 178

For example:
-                    "from chicory import backend; print('DatabaseBackend' in backend.__all__)",  # noqa: E501
+                    "from chicory import backend; print('DatabaseBackend' in backend.__all__)",
src/chicory/config.py (1)
457-460: Security consideration: trust_server_certificate defaults to True.

Defaulting to True bypasses SSL certificate validation, which is insecure in production environments. Consider defaulting to False to enforce secure connections by default, requiring explicit opt-in for development scenarios.
🔎 Suggested change
     trust_server_certificate: bool = Field(
-        default=True,
+        default=False,
         description="Whether to trust the server certificate (for SSL connections)",
     )
src/chicory/backend/models.py (1)
55-66: Consider adding an index on updated_at.

The cleanup_old_results method in database.py queries by updated_at < cutoff. Without an index, this could become slow with large tables.
🔎 Proposed fix
     updated_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True),
         nullable=False,
         default=lambda: datetime.now(UTC),
         onupdate=lambda: datetime.now(UTC),
+        index=True,
     )
src/chicory/backend/database.py (1)
160-162: Unused session variable in connection test.

The session is created but not used. Consider adding a simple query or using _ to clarify intent.
🔎 Proposed fix
         # Test connection
-        async with self._get_db_session_from_pool() as session:
-            pass
+        async with self._get_db_session_from_pool() as session:
+            await session.execute(select(1))
examples/postgres_example.py (1)

316-330: Consider documenting that this is a workaround.

Manually constructing TaskMessage and publishing bypasses the standard delay() API. If ETA support is planned for the public API, consider adding a comment or TODO noting this workaround can be simplified once that's available.

📜 Review details

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c365df9 and 2128bd1.

⛔ Files ignored due to path filters (1)

uv.lock is excluded by !**/*.lock

📒 Files selected for processing (12)

docker-compose.yml
examples/postgres_example.py
pyproject.toml
src/chicory/__init__.py
src/chicory/app.py
src/chicory/backend/__init__.py
src/chicory/backend/database.py
src/chicory/backend/models.py
src/chicory/config.py
src/chicory/exceptions.py
src/chicory/types.py
tests/unit/test_optional_imports.py

🧰 Additional context used

🧬 Code graph analysis (6)

src/chicory/__init__.py (3)

src/chicory/app.py (1)

backend (127-128)

tests/unit/test_worker.py (1)

backend (44-49)

src/chicory/backend/database.py (1)

DatabaseBackend (40-342)

src/chicory/backend/__init__.py (1)

src/chicory/backend/database.py (1)

DatabaseBackend (40-342)

src/chicory/backend/models.py (2)

src/chicory/config.py (4)

MSSQLBackendConfig (436-524)

MySQLBackendConfig (365-433)

PostgresBackendConfig (294-362)

SQLiteBackendConfig (527-590)

src/chicory/result.py (1)

state (55-61)

examples/postgres_example.py (5)

src/chicory/result.py (2)

AsyncResult (13-71)

get (28-53)

src/chicory/types.py (7)

BackendType (239-244)

BrokerType (234-236)

DeliveryMode (24-26)

RetryBackoff (36-41)

RetryPolicy (44-120)

ValidationMode (29-33)

TaskMessage (123-147)

src/chicory/app.py (6)

Chicory (26-170)

broker (123-124)

backend (127-128)

task (130-152)

connect (160-164)

disconnect (166-170)

src/chicory/context.py (4)

TaskContext (13-84)

remaining_retries (79-84)

retry (22-64)

fail (66-68)

src/chicory/task.py (2)

delay (108-130)

send (132-139)

src/chicory/app.py (2)

src/chicory/types.py (1)

BackendType (239-244)

src/chicory/backend/database.py (1)

DatabaseBackend (40-342)

src/chicory/backend/database.py (5)

src/chicory/config.py (11)

MSSQLBackendConfig (436-524)

MySQLBackendConfig (365-433)

PostgresBackendConfig (294-362)

SQLiteBackendConfig (527-590)

dsn (100-113)

dsn (223-236)

dsn (278-291)

dsn (349-362)

dsn (420-433)

dsn (507-524)

dsn (581-590)

src/chicory/exceptions.py (1)

DbPoolExhaustedException (58-64)

src/chicory/types.py (4)

BackendStatus (202-202)

TaskResult (150-163)

TaskState (15-21)

WorkerStats (205-231)

src/chicory/backend/base.py (1)

Backend (10-68)

src/chicory/backend/models.py (2)

TaskResultModel (47-66)

WorkerHeartbeatModel (69-86)

🪛 Ruff (0.14.10)

examples/postgres_example.py

110-110: Avoid specifying long messages outside the exception class

(TRY003)

166-166: Avoid specifying long messages outside the exception class

(TRY003)

168-168: Avoid specifying long messages outside the exception class

(TRY003)

179-179: Unused function argument: scheduled_time

(ARG001)

231-231: Unused function argument: metadata

(ARG001)

259-259: Avoid specifying long messages outside the exception class

(TRY003)

src/chicory/config.py

497-497: Avoid specifying long messages outside the exception class

(TRY003)

577-577: Avoid specifying long messages outside the exception class

(TRY003)

tests/unit/test_optional_imports.py

138-138: Unused noqa directive (non-enabled: E501)