Repository navigation
Conversation
|
Warning Review limit reached
Next review available in: 53 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. 📝 WalkthroughWalkthrough
ChangesDatabase engine caching
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
trapdata/db/base.py (1)
147-155: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd types to the changed
get_dbsignature.Proposed fix
-def get_db(db_path): +def get_db(db_path: DatabaseURL) -> sa.engine.Engine:As per coding guidelines, “Use type hints in function signatures to document expected types without requiring extensive documentation.” <coding_guidelines>
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@trapdata/db/base.py` around lines 147 - 155, Add type annotations to the get_db signature, specifying the expected type of db_path and the function’s return type. Keep the existing path normalization, dialect selection, and _get_engine behavior unchanged.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@trapdata/db/base.py`:
- Around line 97-155: Update the fork-safety handling around _get_engine and
get_db so a child process disposes inherited connection pools with
engine.dispose(close=False) before reusing them. Add a PID-aware child-startup
or checkout hook that detects the post-fork process, preserves the parent engine
reference only as needed for detection, and ensures the child creates fresh
connections without closing the parent’s sockets.
---
Nitpick comments:
In `@trapdata/db/base.py`:
- Around line 147-155: Add type annotations to the get_db signature, specifying
the expected type of db_path and the function’s return type. Keep the existing
path normalization, dialect selection, and _get_engine behavior unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
| @functools.lru_cache(maxsize=None) | ||
| def _get_engine(db_path_str: str, dialect: str, pid: int) -> sa.engine.Engine: | ||
| """ | ||
| Build (and cache) a single Engine per (db_path, dialect, pid). | ||
|
|
||
| IMPORTANT: `get_db()` used to call `sa.create_engine(...)` on every | ||
| invocation, which meant every call to `get_session()` (there are ~90 | ||
| call sites, including inside per-batch loops like | ||
| `ImageQueue.pull_n_from_queue`) spun up a brand new connection pool | ||
| (pool_size=20, max_overflow=30 for postgres). Old engines/pools were | ||
| never explicitly disposed, so connections accumulated across batches | ||
| until Postgres hit `max_connections` ("sorry, too many clients | ||
| already"). Caching the engine here means we reuse one pool for the | ||
| lifetime of the process instead of leaking a new one per call. | ||
|
|
||
| Keying on `pid` makes this cache fork-safe. A process that inherits | ||
| this cache via fork() (e.g. a persistent DataLoader worker) will | ||
| compute a *different* key the first time it calls get_db(), since | ||
| os.getpid() differs from the parent's -- so it transparently builds | ||
| its own fresh Engine/pool instead of reusing (and potentially | ||
| corrupting, via shared sockets) connections the parent already had | ||
| open before the fork happened. | ||
|
|
||
| Arguments: | ||
| db_path_str: database connection string (filepath or URL) | ||
| dialect: database dialect (sqlite, postgresql, etc.) | ||
| pid: process ID (used to key the cache for fork safety) | ||
| Even though this is not used in the function itself, | ||
| lru_cache will key the cache on it, so that a forked process will | ||
| get a new engine instead of reusing the parent's engine. | ||
| """ | ||
| engine_kwargs = { | ||
| "echo": False, | ||
| "future": True, | ||
| "connect_args": DIALECT_CONNECTION_ARGS.get(dialect, {}), | ||
| } | ||
|
|
||
| # Add PostgreSQL-specific engine optimizations | ||
| if dialect == "postgresql": | ||
| engine_kwargs.update( | ||
| { | ||
| "pool_size": 20, # Connection pooling for better performance | ||
| "max_overflow": 30, | ||
| "pool_size": 5, # Reused across all calls now, so this can be modest | ||
| "max_overflow": 10, | ||
| "pool_pre_ping": True, | ||
| "pool_recycle": 3600, | ||
| } | ||
| ) | ||
|
|
||
| db = sa.create_engine(db_path, **engine_kwargs) | ||
| return sa.create_engine(db_path_str, **engine_kwargs) | ||
|
|
||
|
|
||
| def get_db(db_path): | ||
| """ """ | ||
| db_path = get_safe_db_path(db_path) | ||
| dialect = get_dialect(db_path) | ||
|
|
||
| # Reuse a single cached engine (and its connection pool) per | ||
| # (db_path, pid), rather than creating a new engine/pool on every | ||
| # call -- see _get_engine's docstring for why pid is included. | ||
| db = _get_engine(db_path.render_as_string(hide_password=False), dialect, os.getpid()) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## file outline\n'
ast-grep outline trapdata/db/base.py --view expanded || true
printf '\n## relevant search\n'
rg -n "sessionmaker|scoped_session|create_engine|dispose\\(|os\\.getpid\\(|pool_pre_ping|pool_recycle|get_db\\(" trapdata/db -S || true
printf '\n## focused read: trapdata/db/base.py\n'
wc -l trapdata/db/base.py
sed -n '1,260p' trapdata/db/base.pyRepository: RolnickLab/ami-data-companion
Length of output: 11110
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## SQLAlchemy fork-safety docs snippets\n'
python3 - <<'PY'
import urllib.request, re
url = "https://docs.sqlalchemy.org/en/20/core/pooling.html"
html = urllib.request.urlopen(url, timeout=20).read().decode("utf-8", "replace")
for needle in [
"dispose(close=False)",
"os.getpid()",
"pool checkout",
"fork",
]:
print(f"\n### {needle}")
idx = html.lower().find(needle.lower())
if idx == -1:
print("not found")
else:
start = max(0, idx - 600)
end = min(len(html), idx + 1400)
snippet = re.sub(r"<[^>]+>", " ", html[start:end])
snippet = re.sub(r"\s+", " ", snippet)
print(snippet)
PYRepository: RolnickLab/ami-data-companion
Length of output: 2802
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("trapdata/db/base.py")
print(p.read_text())
PYRepository: RolnickLab/ami-data-companion
Length of output: 10014
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the exact db/session wiring in trapdata/db/base.py and related modules.
rg -n "def get_db|def get_session|sessionmaker|scoped_session|`@functools`\\.lru_cache|dispose\\(|os\\.getpid\\(" trapdata/db trapdata -S || true
# Show the relevant portion of trapdata/db/base.py with line numbers.
nl -ba trapdata/db/base.py | sed -n '1,260p'Repository: RolnickLab/ami-data-companion
Length of output: 1349
🌐 Web query:
SQLAlchemy pooling fork safety dispose(close=False) pid checkout docs
💡 Result:
When using SQLAlchemy with processes created via os.fork (including multiprocessing), it is critical to ensure that database connections from the parent process are not shared with the child process, as they are represented by file descriptors that can be concurrently accessed by both processes, leading to corruption or unexpected behavior [1][2]. The recommended approach to handle this is to invoke Engine.dispose(close=False) within the child process during its initialization phase [3][4]. Behavior of dispose(close=False): The close=False parameter (introduced in SQLAlchemy 1.4.33) instructs the engine to discard the current connection pool without attempting to close the underlying database connections [4][2][5]. This is significant because, in a forked child process, attempting to close connections that were inherited from the parent can inadvertently send commands (like a quit/rollback message) to the database that affect the parent process's active connections [6][5][7]. By passing close=False, the child process effectively orphans the inherited connections and initializes a fresh pool for its own use, avoiding interference with the parent [3][1]. Additional Safety Mechanisms: Beyond using engine.dispose(close=False), you can implement a more robust defense using SQLAlchemy event listeners to detect and invalidate connections that were created in a different process [3][4]. This involves: 1. Tracking the PID: Use the connect event to store the process ID (os.getpid) in the connection_record.info dictionary [3][6]. 2. Validation on Checkout: Use the checkout event to compare the current process ID with the stored PID. If they do not match, invalidate the connection by setting it to None and raising a DisconnectionError [3][4]. This event-based approach forces the pool to discard any inherited connections and establish new ones when they are first checked out in the child process, providing a pessimistic layer of protection [3][6].
Citations:
- 1: https://docs.sqlalchemy.org/en/20/core/pooling.html
- 2: https://dokk.org/documentation/sqlalchemy/rel_1_4_50/core/pooling/
- 3: https://docs.sqlalchemy.org/en/21/core/pooling.html
- 4: https://dokk.org/documentation/sqlalchemy/rel_2_0_23/core/pooling/
- 5: add close=True argument to engine.dispose() sqlalchemy/sqlalchemy#7877
- 6: https://docs.sqlalchemy.org/en/13/core/pooling.html
- 7: Documentation on using connection pools with forking is incorrect sqlalchemy/sqlalchemy#7815
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## usages of get_session_class / get_db\n'
rg -n "\bget_session_class\b|\bget_db\b" trapdata -S
printf '\n## any fork hooks or dispose calls in repo\n'
rg -n "dispose\\(close=False\\)|os\\.fork|multiprocessing|register_after_fork|checkout.*pid|connection_record\\.info" trapdata -SRepository: RolnickLab/ami-data-companion
Length of output: 4276
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("trapdata/db/base.py")
for i, line in enumerate(p.read_text().splitlines(), start=1):
if 90 <= i <= 180:
print(f"{i:4d}: {line}")
PYRepository: RolnickLab/ami-data-companion
Length of output: 4256
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "\bget_session_class\b|\bget_db\b|os\.fork|multiprocessing|DataLoader|fork" . -SRepository: RolnickLab/ami-data-companion
Length of output: 8279
Dispose inherited pools after fork. PID-keying only affects new get_db() calls; if the process forks after an engine is created, the child still inherits the parent pool’s open connections. Add engine.dispose(close=False) in child startup or a PID-check checkout hook before reusing the engine.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@trapdata/db/base.py` around lines 97 - 155, Update the fork-safety handling
around _get_engine and get_db so a child process disposes inherited connection
pools with engine.dispose(close=False) before reusing them. Add a PID-aware
child-startup or checkout hook that detects the post-fork process, preserves the
parent engine reference only as needed for detection, and ensures the child
creates fresh connections without closing the parent’s sockets.
1 participant
Fix for #157
Creates a helper function,
_get_engine()that's decorated with@functools.lru_cache. This caches the engine each process uses, so that a new process isn't spawned every timeget_dbis called. Note that_get_engineis keyed on PID even though the function doesn't use it - this just creates a separate entry for each process so that multiple processes don't try to write with the same engine.Also note that I removed the unused
createandupdatearguments toget_db.Summary by CodeRabbit
Performance
Compatibility