Feat: Add support for Snowflake key-pair authentication#702
Conversation
📝 WalkthroughWalkthroughThe SnowflakeConfig class was enhanced to support multiple authentication methods, including password and private key (via string or file path, with optional passphrase). Validation logic and key handling methods were added. Corresponding documentation, tests, and the base SqlConfig interface were updated to accommodate these authentication options. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant SnowflakeConfig
participant FileSystem
participant SnowflakeConnector
User->>SnowflakeConfig: Initialize with auth parameters
SnowflakeConfig->>SnowflakeConfig: _validate_authentication_config()
alt Using password
SnowflakeConfig->>SnowflakeConnector: Connect with password
else Using private_key
SnowflakeConfig->>SnowflakeConfig: _get_private_key_content()
alt private_key_path
SnowflakeConfig->>FileSystem: Read private key file
else private_key string
SnowflakeConfig->>SnowflakeConfig: Use private_key string
end
SnowflakeConfig->>SnowflakeConfig: _get_private_key_bytes()
SnowflakeConfig->>SnowflakeConnector: Connect with private_key_bytes and optional passphrase
end
SnowflakeConnector-->>User: Connection established
Suggested reviewersWould you like to see additional tests for edge cases around malformed private key files, or is the current coverage sufficient, wdyt? 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🧰 Additional context used🧠 Learnings (1)airbyte/_processors/sql/snowflake.py (3)🧬 Code Graph Analysis (1)airbyte/_processors/sql/snowflake.py (2)
⏰ Context from checks skipped due to timeout of 90000ms (2)
🔇 Additional comments (8)
✨ 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: 1
🧹 Nitpick comments (3)
airbyte/_processors/sql/snowflake.py (2)
84-92: Simplify by removing unnecessaryelif.The method logic is correct, but the code can be cleaner by removing the
elifafterreturn.def _get_private_key_content(self) -> bytes: """Get the private key content from either private_key or private_key_path.""" if self.private_key: return str(self.private_key).encode("utf-8") - elif self.private_key_path: + if self.private_key_path: return Path(self.private_key_path).read_bytes() - else: - raise ValueError("No private key provided") + raise ValueError("No private key provided")
93-113: Remove unnecessary assignment before return.The private key handling is secure and correct, but we can simplify the code slightly.
- private_key_bytes = private_key.private_bytes( + return private_key.private_bytes( encoding=serialization.Encoding.DER, format=serialization.PrivateFormat.PKCS8, encryption_algorithm=serialization.NoEncryption(), ) - - return private_key_bytestests/unit_tests/test_snowflake_config.py (1)
1-510: Excellent test coverage! Just needs formatting.The test suite provides comprehensive coverage of all authentication methods, validation logic, and error scenarios. Great job on testing edge cases like wrong passphrases and multiple auth methods!
However, the pipeline indicates formatting issues that need to be fixed. Could you run
ruff formaton this file to resolve the formatting issues?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
airbyte/_processors/sql/snowflake.py(3 hunks)airbyte/caches/snowflake.py(2 hunks)airbyte/shared/sql_processor.py(3 hunks)tests/unit_tests/test_snowflake_config.py(1 hunks)
🧰 Additional context used
🧠 Learnings (3)
airbyte/caches/snowflake.py (1)
Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.
tests/unit_tests/test_snowflake_config.py (1)
Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.
airbyte/_processors/sql/snowflake.py (1)
Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.
🧬 Code Graph Analysis (2)
airbyte/shared/sql_processor.py (1)
airbyte/_processors/sql/snowflake.py (1)
get_sql_alchemy_connect_args(115-119)
airbyte/_processors/sql/snowflake.py (2)
airbyte/secrets/base.py (1)
SecretString(34-139)airbyte/shared/sql_processor.py (2)
get_sql_alchemy_connect_args(131-133)get_vendor_client(147-156)
🪛 Flake8 (7.2.0)
tests/unit_tests/test_snowflake_config.py
[error] 466-466: expected 2 blank lines, found 1
(E302)
airbyte/_processors/sql/snowflake.py
[error] 84-84: too many blank lines (2)
(E303)
🪛 GitHub Actions: Run Linters
tests/unit_tests/test_snowflake_config.py
[error] 1-1: Ruff formatting check failed. File would be reformatted. Run 'ruff format' to fix code style issues.
airbyte/_processors/sql/snowflake.py
[error] 66-66: Ruff E501: Line too long (118 > 100).
[error] 84-84: Ruff E303: Too many blank lines (2). Remove extraneous blank line(s).
[error] 88-88: Ruff RET505: Unnecessary elif after return statement. Remove unnecessary elif.
[error] 112-112: Ruff RET504: Unnecessary assignment to private_key_bytes before return statement. Remove unnecessary assignment.
🪛 Ruff (0.11.9)
airbyte/_processors/sql/snowflake.py
66-66: Line too long (118 > 100)
(E501)
84-84: Too many blank lines (2)
Remove extraneous blank line(s)
(E303)
88-88: Unnecessary elif after return statement
Remove unnecessary elif
(RET505)
112-112: Unnecessary assignment to private_key_bytes before return statement
Remove unnecessary assignment
(RET504)
🪛 Pylint (3.3.7)
airbyte/_processors/sql/snowflake.py
[refactor] 86-91: Unnecessary "elif" after "return", remove the leading "el" from "elif"
(R1705)
⏰ Context from checks skipped due to timeout of 90000ms (2)
- GitHub Check: Pytest (Fast)
- GitHub Check: Pytest (No Creds)
🔇 Additional comments (6)
airbyte/shared/sql_processor.py (2)
131-134: LGTM! Clean extension point for connection arguments.The base implementation returning an empty dict is appropriate and provides a nice interface for subclasses to override when needed.
144-144: Correct integration of connection arguments.The connect_args are properly passed to SQLAlchemy's create_engine, enabling database-specific authentication parameters.
airbyte/caches/snowflake.py (1)
6-57: Excellent documentation examples!The examples clearly demonstrate all three authentication methods and follow security best practices by using
ab.get_secret()for sensitive values. The inline comments for optional parameters are helpful.airbyte/_processors/sql/snowflake.py (3)
114-120: Well-implemented connection arguments override.The method correctly provides private key bytes only when key-based authentication is used, maintaining compatibility with password authentication.
137-152: Correct URL construction for different authentication methods.The method properly validates the configuration and constructs the URL appropriately for both password and key-based authentication.
154-177: Comprehensive vendor client configuration for all auth methods.The implementation correctly handles all three authentication scenarios and properly sets the JWT authenticator for key-based authentication. Nice work on the dynamic parameter building!
There was a problem hiding this comment.
Actionable comments posted: 2
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
poetry.lockis excluded by!**/*.lock
📒 Files selected for processing (4)
airbyte/_processors/sql/snowflake.py(3 hunks)airbyte/caches/snowflake.py(1 hunks)pyproject.toml(1 hunks)tests/unit_tests/test_snowflake_config.py(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- pyproject.toml
🚧 Files skipped from review as they are similar to previous changes (2)
- airbyte/caches/snowflake.py
- tests/unit_tests/test_snowflake_config.py
🧰 Additional context used
🧠 Learnings (1)
airbyte/_processors/sql/snowflake.py (1)
Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.
🧬 Code Graph Analysis (1)
airbyte/_processors/sql/snowflake.py (2)
airbyte/secrets/base.py (1)
SecretString(34-139)airbyte/shared/sql_processor.py (2)
get_sql_alchemy_connect_args(131-133)get_vendor_client(147-156)
🔇 Additional comments (6)
airbyte/_processors/sql/snowflake.py (6)
7-7: LGTM on the new imports!The new imports for
Path, cryptography modules, and typing updates are appropriate for the key-pair authentication functionality.Also applies to: 9-9, 12-13
41-45: Well-structured authentication fields!The new optional fields provide comprehensive support for different key-pair authentication methods while maintaining security with
SecretStringfor sensitive data.
53-82: Excellent validation logic!The authentication validation is comprehensive and handles all edge cases properly. The error messages are clear and the logic prevents conflicting authentication configurations.
84-109: Solid private key handling implementation!The methods properly handle both direct private key strings and file paths, with correct cryptography library usage for PEM loading and DER conversion. The separation of concerns is well done.
111-116: Clean implementation of connect args!The method correctly provides private key bytes only when using key-pair authentication, following SQLAlchemy patterns.
136-149: Good URL method updates!The method correctly validates authentication configuration and conditionally includes password only when using password authentication.
|
/test-pr
|
Aaron ("AJ") Steers (aaronsteers)
left a comment
Aaron ("AJ") Steers (aaronsteers)
left a comment
There was a problem hiding this comment.
Yves Alioune Amoussou (@Alioune05) - Looks great and all tests are passing. Thanks very much for contributing!
Aaron ("AJ") Steers (aaronsteers)
merged commit 6794987
into
airbytehq:main
2 participants
Summary
Related to #654 and #681 and this Slack conversation.
Update SnowflakeConfig and SqlConfig to allow connection with :
It also make the password attribute optional.
Testing
Summary by CodeRabbit
New Features
Documentation
Tests
Chores
cryptographypackage to support key handling and encryption features.