Cognito support

[pymonger]

PR Summary: Add Cognito Token Support and Bearer Authentication to Unity Initiator

Overview

This PR adds comprehensive token-based authentication support to the Unity Initiator project, replacing basic authentication with modern Bearer token authentication and Cognito integration.

Key Changes

�� New Authentication Utilities (src/unity_initiator/utils/auth_utils.py)

• TokenManager class: Handles Cognito token lifecycle with automatic refresh
  • Caches tokens and refreshes 5 minutes before expiration
  • Supports both username/password and refresh token flows
  • Graceful fallback to credential-based fetching
• fetch_cognito_token(): High-level function for token retrieval
• get_auth_headers(): Generates appropriate authentication headers
• TokenInfo dataclass: Structured token storage with expiration tracking

�� Enhanced SubmitDagByID Action (src/unity_initiator/actions/submit_dag_by_id.py)

• Multiple authentication methods:
  • Bearer token authentication (primary)
  • Cognito token authentication with automatic refresh
  • Basic authentication (fallback)
• Smart token management: Uses TokenManager for automatic token refresh
• Backward compatibility: Maintains existing basic auth support

📋 Router Schema Updates (src/unity_initiator/resources/routers_schema.yaml)

• Added new authentication parameters:
  • airflow_token for Bearer token authentication
  • unity_username, unity_password, unity_client_id, unity_region for Cognito authentication
• Maintained backward compatibility with existing airflow_username and airflow_password
• Added comprehensive documentation explaining authentication methods

📦 Dependencies

• Added httpx: Modern HTTP client for token fetching
• Added requests: Maintained for compatibility
• Updated pyproject.toml with new dependencies

�� Comprehensive Test Coverage

• 16 new test cases for authentication utilities
• 7 new test cases for SubmitDagByID with token support
• Mock-based testing with proper isolation
• Edge case coverage: Token expiration, refresh failures, network errors
• Schema validation tests for new authentication parameters

📚 Documentation and Examples

• Updated README with comprehensive authentication section
• Example router configurations demonstrating all authentication methods
• Inline documentation for all new functions and classes
• Test configurations for validation

Authentication Flow

graph TD
    A[SubmitDagByID.execute] --> B{Token provided?}
    B -->|Yes| C[Use Bearer token]
    B -->|No| D{Cognito credentials?}
    D -->|Yes| E[Fetch Cognito token]
    D -->|No| F[Use Basic auth]
    E --> G[TokenManager.get_valid_token]
    G --> H{Cached token valid?}
    H -->|Yes| I[Use cached token]
    H -->|No| J[Fetch new token]
    J --> K[Cache and return token]
    I --> L[Make API request]
    K --> L
    C --> L
    F --> L

Loading

Configuration Examples

Bearer Token Authentication

actions:
  - name: submit_dag_by_id
    params:
      dag_id: example_dag
      airflow_base_api_endpoint: https://airflow.example.com/api/v1
      airflow_token: ${AIRFLOW_BEARER_TOKEN}

Cognito Authentication

actions:
  - name: submit_dag_by_id
    params:
      dag_id: example_dag
      airflow_base_api_endpoint: https://airflow.example.com/api/v1
      unity_username: ${UNITY_USERNAME}
      unity_password: ${UNITY_PASSWORD}
      unity_client_id: ${UNITY_CLIENT_ID}
      unity_region: us-west-2

Basic Authentication (Fallback)

actions:
  - name: submit_dag_by_id
    params:
      dag_id: example_dag
      airflow_base_api_endpoint: https://airflow.example.com/api/v1
      airflow_username: ${AIRFLOW_USERNAME}
      airflow_password: ${AIRFLOW_PASSWORD}

Benefits

1. 🔒 Enhanced Security: Bearer tokens are more secure than basic auth
2. ⚡ Automatic Token Management: No manual token refresh needed
3. 🔄 Seamless Integration: Works with existing Unity SPS infrastructure
4. 📈 Scalability: Token caching reduces API calls to Cognito
5. 🛡️ Fault Tolerance: Multiple fallback mechanisms ensure reliability
6. 🧪 Testability: Comprehensive test coverage ensures reliability
7. 📋 Schema Validation: Router configurations are properly validated

Breaking Changes

• None: All changes are backward compatible
• Existing basic auth configurations continue to work unchanged
• Router schema maintains compatibility with existing configurations

Testing

• ✅ All authentication tests pass (16/16)
• ✅ All SubmitDagByID tests pass (7/7)
• ✅ All schema validation tests pass (8/8)
• ✅ All existing functionality tests pass
• ✅ Full test suite: 40 passed, 7 errors (unrelated Lambda build issues)

Version Bump

• 0.0.1 → 0.0.2
• Updated src/unity_initiator/__about__.py
• Added comprehensive changelog entry in CHANGELOG.md

Files Changed

• src/unity_initiator/utils/auth_utils.py (new)
• src/unity_initiator/actions/submit_dag_by_id.py (modified)
• src/unity_initiator/resources/routers_schema.yaml (updated)
• pyproject.toml (dependencies added)
• tests/test_auth_utils.py (new)
• tests/test_submit_dag_by_id.py (new)
• tests/resources/test_router_with_auth.yaml (new)
• examples/router_with_cognito_auth.yaml (new)
• README.md (updated with authentication documentation)
• src/unity_initiator/__about__.py (version bump)
• CHANGELOG.md (updated)

This PR provides a robust, production-ready authentication system that enhances security while maintaining full backward compatibility and proper schema validation.