UN-3157 [FIX] Fix resource leaks in platform-service #1748
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
- Add try/finally blocks to ensure database cursors are always closed - Add shared Redis connection pool to prevent per-request connection creation - Fix execute_query(), validate_bearer_token() in platform.py - Fix get_adapter_instance_from_db() in adapter_instance.py - Fix get_prompt_instance_from_db(), get_llm_profile_instance_from_db() in prompt_studio.py - Add get_redis_client() and safe_cursor() utilities in extensions.py Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- Add test_memory_leak_simulation.py demonstrating OLD vs NEW code behavior - Tests verify cursor/connection cleanup in exception scenarios - Uses tracemalloc for actual memory growth measurement - Includes load test simulating 1000 requests with 30% failure rate - Add tests/README.md with instructions for running tests Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Contributor
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings. WalkthroughIntroduce a shared Redis pool and client, add a safe_cursor context manager for DB cursor lifecycle, update controller and helper modules to use these utilities (replacing per-call Redis and manual cursor handling), and add tests plus README for leak/memory simulations. Changes
Sequence Diagram(s)mermaid Estimated code review effort🎯 4 (Complex) | ⏱️ ~40 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
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 |
johnyrahul
requested review from
jaseemjaskp,
muhammad-ali-e and
vishnuszipstack
and removed request for
muhammad-ali-e
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@platform-service/src/unstract/platform_service/controller/platform.py`:
- Around line 87-90: The query construction for platform_key_table interpolates
token directly into the SQL (query = f"""... WHERE key = '{token}'"""), creating
an SQL injection risk; change this to use a parameterized query consistent with
execute_query: build the SQL string using a %s placeholder for the key and pass
token as a parameter to execute_query (keep Env.DB_SCHEMA and platform_key_table
usage intact but do not interpolate token), ensuring execute_query receives the
parameter list/tuple so the DB driver performs safe binding.
🧹 Nitpick comments (4)
platform-service/src/unstract/platform_service/extensions.py (2)
14-33: Thread safety consideration for lazy initialization.The lazy initialization of
_redis_poolusing a global variable without synchronization could lead to a race condition in multi-threaded environments (e.g., multiple Flask workers). While unlikely to cause issues in practice (worst case: pool created twice, one discarded), consider using a lock for thread-safe initialization if this service runs with threading.♻️ Thread-safe lazy initialization (optional)
+import threading + # Redis connection pool (initialized lazily) _redis_pool: redis.ConnectionPool | None = None +_redis_pool_lock = threading.Lock() def get_redis_pool() -> redis.ConnectionPool: """Get or create the Redis connection pool. Returns: redis.ConnectionPool: Shared connection pool for Redis operations. """ global _redis_pool if _redis_pool is None: - # Import here to avoid circular imports - from unstract.platform_service.env import Env - - _redis_pool = redis.ConnectionPool( - host=Env.REDIS_HOST, - port=Env.REDIS_PORT, - username=Env.REDIS_USERNAME, - password=Env.REDIS_PASSWORD, - max_connections=10, - decode_responses=False, - ) + with _redis_pool_lock: + if _redis_pool is None: # Double-check locking + from unstract.platform_service.env import Env + + _redis_pool = redis.ConnectionPool( + host=Env.REDIS_HOST, + port=Env.REDIS_PORT, + username=Env.REDIS_USERNAME, + password=Env.REDIS_PASSWORD, + max_connections=10, + decode_responses=False, + ) return _redis_pool
45-60:safe_cursorcontext manager is defined but not used.The
safe_cursorcontext manager is a useful utility, but the changed files (adapter_instance.py,prompt_studio.py,platform.py) all use inline try/finally blocks instead of this context manager. Consider either usingsafe_cursorin those files for consistency, or documenting that it's available for future use.platform-service/tests/test_memory_leak_simulation.py (1)
16-44: Consider annotating mutable class attribute withClassVar.The
_open_cursorsclass attribute is used for test tracking and works correctly, but for type correctness it should be annotated withClassVaras flagged by static analysis.♻️ Add ClassVar annotation
+from typing import ClassVar + class FakeCursor: """Simulates a database cursor that tracks if it was closed.""" - _open_cursors: list["FakeCursor"] = [] + _open_cursors: ClassVar[list["FakeCursor"]] = []platform-service/tests/README.md (1)
56-67: Add language specifier to fenced code block.The sample output code block is missing a language specifier. Use
textorconsolefor output examples.♻️ Add language specifier
**Sample Output:** -``` +```text [OLD CODE] Open cursors after 100 failed requests: 100 ❌ LEAKED [NEW CODE] Open cursors after 100 failed requests: 0 ✅ NO LEAK
jaseemjaskp
approved these changes
Jan 19, 2026
muhammad-ali-e
approved these changes
Jan 20, 2026
platform-service/src/unstract/platform_service/controller/platform.py
Outdated
Show resolved
Hide resolved
Contributor
|
@johnyrahul please check sonarqubecloud |
- Fix SQL injection vulnerability by using parameterized queries - Refactor cursor handling to use safe_cursor context manager - Add ClassVar type annotation in test file - Add language specifier to README code block Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@platform-service/tests/test_memory_leak_simulation.py`:
- Around line 196-260: Tests use absolute tracemalloc snapshot diffs which are
flaky; update both test_memory_growth_with_leaked_objects and
test_no_memory_growth_with_proper_cleanup to call gc.collect() before sampling,
use tracemalloc.get_traced_memory() (start tracing, call gc.collect(), read
current/peak), compute a relative delta (e.g., percent change) instead of
summing snapshot compare_to results, and assert using relative thresholds (e.g.,
growth >X% for leaked case and <Y% for cleanup case) to make the assertions
stable across CI allocators.
🧹 Nitpick comments (2)
platform-service/src/unstract/platform_service/controller/platform.py (2)
84-96: Consider using explicit column names instead ofSELECT *.The query uses
SELECT *with positional indexing (result_row[1],result_row[2]), which is fragile if the table schema changes. Consider selecting explicit columns to make the code more maintainable:♻️ Suggested improvement
platform_key_table = DBTable.PLATFORM_KEY query = f""" - SELECT * FROM \"{Env.DB_SCHEMA}\".{platform_key_table} + SELECT key, is_active FROM \"{Env.DB_SCHEMA}\".{platform_key_table} WHERE key = %s """ try: with safe_cursor(query, (token,)) as cursor: result_row = cursor.fetchone() if not result_row or len(result_row) == 0: app.logger.error(f"Authentication failed. bearer token not found {token}") return False - platform_key = str(result_row[1]) - is_active = bool(result_row[2]) + platform_key = str(result_row[0]) + is_active = bool(result_row[1])
340-346: Consider moving the success return to anelseblock for clarity.The static analysis hint (TRY300) suggests moving
return value, 200to anelseblock. This is a minor style preference that makes it clearer that the return only happens when no exception occurs.♻️ Optional style improvement
try: redis_key = f"{account_id}:{key}" app.logger.info(f"Getting cached data for key: {redis_key}") value = r.get(redis_key) if value is None: return "Not Found", 404 - return value, 200 except Exception as e: raise APIError(message=f"Error while getting cached data: {e}") from e + else: + return value, 200
Contributor
Test ResultsSummary
Runner Tests - Full Report
SDK1 Tests - Full Report
|
|
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
try/finallyblocks to ensure cursors are always closed, even when exceptions occurChanges
Resource Leak Fixes
execute_query()- Wrap cursor operations in try/finallyvalidate_bearer_token()- Ensure cursor is closed on all code pathsget_adapter_instance_from_db()- Fix cursor leak when adapter not found (APIError)get_prompt_instance_from_db()- Fix cursor leak when prompt not foundget_llm_profile_instance_from_db()- Fix cursor leak when LLM profile not foundRedis Connection Pool
get_redis_pool()andget_redis_client()in extensions.pyTests Added
test_memory_leak_simulation.py- 9 tests demonstrating OLD vs NEW code behaviortracemallocfor actual memory growth measurementTest plan
uv run pytest tests/test_memory_leak_simulation.py -v -sTest Results
🤖 Generated with Claude Code