feat: Remove long standing connections on sql state and catalog backends#675
Conversation
📝 WalkthroughWalkthroughThe changes refactor backend classes to accept a Changes
Sequence Diagram(s)sequenceDiagram
participant CacheBase
participant SqlCatalogBackend
participant SqlStateBackend
participant SqlConfig
participant SqlEngine
CacheBase->>SqlCatalogBackend: __init__(sql_config=self, table_prefix)
CacheBase->>SqlStateBackend: __init__(sql_config=self, table_prefix)
SqlCatalogBackend->>SqlConfig: get_sql_engine()
SqlConfig->>SqlEngine: returns engine
SqlStateBackend->>SqlConfig: get_sql_engine()
SqlConfig->>SqlEngine: returns engine
Would you like me to help draft a more detailed sequence diagram showing specific method calls inside the backends, or does this high-level overview work for your review, wdyt? Note ⚡️ AI Code Reviews for VS Code, Cursor, WindsurfCodeRabbit now has a plugin for VS Code, Cursor and Windsurf. This brings AI code reviews directly in the code editor. Each commit is reviewed immediately, finding bugs before the PR is raised. Seamless context handoff to your AI code agent ensures that you can easily incorporate review feedback. Note ⚡️ Faster reviews with cachingCodeRabbit now supports caching for code and dependencies, helping speed up reviews. This means quicker feedback, reduced wait times, and a smoother review experience overall. Cached data is encrypted and stored securely. This feature will be automatically enabled for all accounts on May 30th. To opt out, configure 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
⏰ Context from checks skipped due to timeout of 90000ms (2)
✨ Finishing Touches
🧪 Generate Unit Tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
airbyte/caches/_catalog_backend.py (1)
179-226: Consider connection pooling implications.While creating a new engine for each operation solves the token expiration issue, have you considered the performance implications? SQLAlchemy engines typically use connection pooling, so this might not be as expensive as it seems, but it's worth considering whether setting appropriate pool options (like
pool_recycle) might be a more efficient alternative in some cases, wdyt?airbyte/caches/_state_backend.py (1)
1-276: Consider adding tests for this change.The PR objectives mention that testing is yet to be completed. Would you consider adding specific tests that simulate long-running operations to verify this fix actually resolves the token expiration issue? This could be particularly valuable to prevent regression in the future.
🧰 Tools
🪛 Ruff (0.11.9)
23-23: Move application import
airbyte.shared.sql_processor.SqlConfiginto a type-checking blockMove into type-checking block
(TC001)
🪛 GitHub Actions: Run Linters
[error] 23-23: Ruff TC001: Move application import
airbyte.shared.sql_processor.SqlConfiginto a type-checking block.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to data retention organization setting
Knowledge Base: Disabled due to data retention organization setting
📒 Files selected for processing (3)
airbyte/caches/_catalog_backend.py(5 hunks)airbyte/caches/_state_backend.py(4 hunks)airbyte/caches/base.py(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
airbyte/caches/_catalog_backend.py (4)
airbyte/shared/sql_processor.py (3)
SqlConfig(86-151)get_sql_engine(131-140)get_sql_engine(384-386)airbyte/cloud/sync_results.py (1)
get_sql_engine(279-281)airbyte/_processors/sql/duckdb.py (1)
get_sql_engine(78-87)airbyte/results.py (1)
get_sql_engine(80-82)
🪛 Ruff (0.11.9)
airbyte/caches/_state_backend.py
23-23: Move application import airbyte.shared.sql_processor.SqlConfig into a type-checking block
Move into type-checking block
(TC001)
airbyte/caches/_catalog_backend.py
26-26: Move application import airbyte.shared.sql_processor.SqlConfig into a type-checking block
Move into type-checking block
(TC001)
🪛 GitHub Actions: Run Linters
airbyte/caches/_state_backend.py
[error] 23-23: Ruff TC001: Move application import airbyte.shared.sql_processor.SqlConfig into a type-checking block.
airbyte/caches/_catalog_backend.py
[error] 26-26: Ruff TC001: Move application import airbyte.shared.sql_processor.SqlConfig into a type-checking block.
⏰ Context from checks skipped due to timeout of 90000ms (2)
- GitHub Check: Pytest (No Creds)
- GitHub Check: Pytest (Fast)
🔇 Additional comments (10)
airbyte/caches/base.py (2)
91-94: LGTM! Dynamic SQL engine retrieval instead of reusing a single connection.This change looks good - passing
sql_config=selfto theSqlCatalogBackendinstead of directly injecting an engine instance. This allows the backend to create a fresh connection each time it needs one, which should help address the Snowflake token expiration issue mentioned in the PR description.
95-98: LGTM! Consistent approach for both backend components.Perfect consistency - applying the same pattern to the state backend as well. This ensures all SQL operations will now use fresh connections, preventing authentication token expirations during long-running syncs.
airbyte/caches/_catalog_backend.py (4)
107-113: LGTM! Improved connection handling with SqlConfig.Great refactoring to use
sql_configinstead of a direct engine instance. This changes the backend from eagerly connecting and reusing a single connection to lazily connecting for each database operation.This addresses the Snowflake token expiration issue mentioned in the PR objectives.
123-125: LGTM! Dynamic retrieval of SQL engine.Perfect - now retrieving a fresh engine instance when needed rather than reusing a stored engine.
179-181: LGTM! Consistent engine retrieval pattern.Correctly using the dynamic engine retrieval pattern here as well.
215-217: LGTM! Consistent implementation throughout the class.Final instance of the pattern applied correctly.
airbyte/caches/_state_backend.py (4)
137-139: LGTM! Dynamic SQL engine retrieval in the state writer.Perfect implementation - this change ensures that the state writer creates a fresh connection each time it writes state, avoiding the token expiration issue.
187-194: LGTM! Consistent SqlConfig pattern in state backend.Excellent implementation of the same pattern in the state backend constructor. This change keeps the pattern consistent across both catalog and state backends.
197-201: LGTM! Dynamic engine retrieval when ensuring tables exist.Good implementation - retrieving a fresh engine when creating internal tables.
224-226: LGTM! Final instance of the dynamic engine retrieval pattern.Consistently applied the pattern throughout the codebase.
Ben Wilen (benjaminwilen)
force-pushed
the
ben.wilen/feat_Remove_long_standing_connections_on_sql_state_and_catalog_backends
branch
2 times, most recently
from
66ed0d1 to
2f3b247
Compare
Ben Wilen (benjaminwilen)
force-pushed
the
ben.wilen/feat_Remove_long_standing_connections_on_sql_state_and_catalog_backends
branch
2 times, most recently
from
c0d9e0d to
8bfa5a2
Compare
Ben Wilen (benjaminwilen)
force-pushed
the
ben.wilen/feat_Remove_long_standing_connections_on_sql_state_and_catalog_backends
branch
from
8bfa5a2 to
872f893
Compare
1 participant
What does this PR do?
This PR modifies how
SqlCatalogBackendandSqlStateBackendconnect to an sql engine to do so lazily and per query execution instead of eagerly and in a single connection. This allows long running syncs to complete without receiving the following error:Testing
Summary by CodeRabbit