feat: Add RSA key-pair authentication support for Snowflake

[kajbaf]

This PR introduces support for key-pair authentication with Snowflake.
Snowflake now enforces key-pair authentication for service accounts and MFA for human accounts.
Defining RSA key-pair authentication provides a secure and straightforward solution for both cases. Particularly helpful since approving MFA prompts for every query can quickly become cumbersome.

[llamapreview]

AI Code Review by LlamaPReview

🎯 TL;DR & Recommendation

Recommendation: Request Changes

This PR adds valuable RSA key-pair authentication but introduces a critical silent failure in password handling that could break existing deployments relying on environment variables.

📄 Documentation Diagram

This diagram documents the new RSA key-pair authentication flow for Snowflake connections.

sequenceDiagram
    participant U as User
    participant V as Vanna
    participant S as Snowflake
    U->>V: connect_to_snowflake()
    V->>V: Check authentication method
    alt Private key provided
        V->>V: Set authenticator = SNOWFLAKE_JWT
        note over V: PR #35;959 added key-pair authentication
        V->>S: Authenticate with private key
    else Password provided
        V->>S: Authenticate with password
    end
    S-->>V: Connection established

Loading

🌟 Strengths

• Improved code organization with explicit connection parameter management
• Clear separation between password and key-pair authentication methods

Priority	File	Category	Impact Summary	Anchors
P1	src/vanna/base/base.py	Bug	Silent failure breaks deployments if env var not set	path:tests/test_vanna.py
P2	src/vanna/base/base.py	Testing	New RSA auth path lacks test coverage	path:tests/test_vanna.py
P2	src/vanna/base/base.py	Architecture	Clear authentication separation with user-friendly errors
P2	src/vanna/base/base.py	Maintainability	Improved parameter organization for future extensions

🔍 Notable Themes

• Testing gaps: New authentication methods lack dedicated test coverage, increasing production risk
• Authentication robustness: While architecture is clean, error handling inconsistencies could impact user experience

📈 Risk Diagram

This diagram illustrates the silent failure risk in password environment variable handling.

sequenceDiagram
    participant U as User
    participant V as Vanna
    participant S as Snowflake
    U->>V: connect_to_snowflake(password="mypassword")
    V->>V: password = os.getenv("SNOWFLAKE_PASSWORD")
    note over V: R1(P1): If SNOWFLAKE_PASSWORD not set,<br/>password becomes None
    V->>S: Attempt connection with password=None
    S-->>V: Authentication failed

Loading

---

💡 Have feedback? We'd love to hear it in our GitHub Discussions.
✨ This review was generated by LlamaPReview Advanced, which is free for all open-source projects. Learn more.

[llamapreview]

Follow-up Review: PR #959 - RSA Key-Pair Authentication for Snowflake

1. Status Update

Since my previous review, I can see the PR has evolved significantly with three new commits addressing both my concerns and adding new functionality:

Key Improvements:

• ✅ Added comprehensive test coverage for key-pair authentication (addressing my P2 concern)
• ✅ Improved password handling logic (cleaner implementation)
• ✅ Enhanced connection parameter management (more structured approach)
• ✅ Added proper validation for private key file existence

New Additions:

• Added private_key_file_pwd parameter support
• Implemented proper error handling for missing authentication credentials
• Created a robust test suite with mocking for key-pair authentication

2. Detailed Analysis

2.1 Test Coverage Improvements

The new test test_snowflake_keypair_auth_mock is excellent - it:

• Properly mocks the Snowflake connector
• Verifies JWT authenticator is set correctly
• Tests both key file and password parameters
• Ensures password isn't included when using key-pair auth

Code Quality Observation:

# Good pattern for parameter capture in tests
captured_params = {}
def fake_connect(**kwargs):
    nonlocal captured_params
    captured_params = kwargs
    return FakeConn()

2.2 Connection Parameter Management

The refactored connection logic is much cleaner:

Before:

conn = snowflake.connector.connect(
    user=username,
    password=password,
    account=account,
    database=database,
    client_session_keep_alive=True,
    **kwargs
)

After:

conn_params = {
    "user": username,
    "account": account,
    "database": database,
    "warehouse": warehouse,
    "role": role,
    "client_session_keep_alive": True,
}

# Conditional auth logic
if private_key_file:
    conn_params["authenticator"] = "SNOWFLAKE_JWT"
    conn_params["private_key_file"] = private_key_file
    if private_key_file_pwd:
        conn_params["private_key_file_pwd"] = private_key_file_pwd
else:
    if password:
        conn_params["password"] = password

This is a significant improvement because:

1. Makes the authentication flow more explicit
2. Separates concerns between connection parameters and auth method
3. Easier to maintain and extend

2.3 Error Handling

The validation logic is now more robust:

if not os.path.isfile(private_key_file):
    raise ImproperlyConfigured("private_key_file path is invalid.")

if "password" not in conn_params and "private_key_file" not in conn_params:
    raise ImproperlyConfigured(
        "Snowflake authentication not configured. Provide a password or a valid private_key_file."
    )

2.4 Security Considerations

The implementation properly:

• Validates file existence before attempting to use it
• Only includes password in connection when not using key-pair auth
• Uses environment variables as fallback (good security practice)

One minor suggestion: Consider adding validation that the private key file has proper permissions (not world-readable).

2.5 Backward Compatibility

The changes maintain excellent backward compatibility:

• All existing password-based auth continues to work
• New parameters are optional with None defaults
• Existing environment variable support preserved

3. Outstanding Concerns

3.1 Minor Documentation Gap

While the code is well-structured, there's no:

• Docstring explaining the new parameters
• Example usage in the class/module docstring
• Mention of the environment variable pattern for private key auth

Suggestion: Add a class-level docstring example:

"""
Example with key-pair authentication:
    vn.connect_to_snowflake(
        account='my-account',
        username='service-user',
        database='my-db',
        private_key_file='/path/to/rsa_key.p8',
        private_key_file_pwd='key-password-if-any'
    )
"""

3.2 Edge Case Consideration

The current implementation doesn't handle:

• Empty string vs None for private_key_file_pwd
• Very long paths for private key files (potential path length issues on Windows)

Suggestion: Add explicit handling:

if private_key_file_pwd == '':
    private_key_file_pwd = None  # Treat empty string as no password

3.3 Test Coverage Expansion

While the new tests are good, consider adding:

1. Test for invalid private key file path
2. Test for missing both password and key file
3. Test for empty password with key file (should work)
4. Test for environment variable fallback scenarios

4. Recommendations

4.1 High Priority (Should Address Before Merge)

1. Add docstring documentation for the new parameters (5 min fix)
2. Explicit empty string handling for private_key_file_pwd (5 min fix)

4.2 Medium Priority (Nice to Have)

1. Expand test coverage for additional edge cases (30 min)
2. Add file permission validation for private key (15 min)
3. Consider path length validation for Windows compatibility (10 min)

4.3 Future Considerations

1. Key rotation support: Consider adding parameters for key rotation scenarios
2. Temporary credential support: Could extend to support temporary credentials
3. Connection pooling: The current implementation could be enhanced with connection pooling for key-pair auth

5. Final Assessment

Overall Rating: 9/10 (Excellent implementation with minor polish needed)

Strengths:
✅ Comprehensive test coverage added
✅ Clean, maintainable code structure
✅ Proper security considerations
✅ Excellent backward compatibility
✅ Robust error handling

Areas for Polish:
📝 Minor documentation additions
📝 Edge case handling for empty strings
📝 Expanded test scenarios

Ready to Merge: Almost - just needs the high-priority documentation additions. The implementation is production-ready from a functional perspective.

Kudos to the author for addressing all major concerns from the previous review and implementing a robust solution that significantly enhances the Snowflake authentication capabilities while maintaining backward compatibility. The test coverage is particularly well-done and demonstrates good software engineering practices.

---

💡 LlamaPReview Community
Have feedback on this AI Code review tool? Join our GitHub Discussions to share your thoughts and help shape the future of LlamaPReview.