dbt token authentication commands

[pgoell]

new dbt token command and underlying module structure

• auth: handles azure context for AzureAD authentication & databricks rest endpoints
• dbt: Profiles parser
• token: token manager & storage

exceptions for errors in utils

[claude]

Pull Request Review: dbt token authentication commands

I've reviewed PR #2 and have the following feedback organized by category:

---

🔒 Security Concerns

1. Token Storage in Environment Variables (manager.py:163)

   • Tokens are stored in os.environ which persists for the process lifetime
   • If the process is compromised or dumps are taken, tokens could be exposed
   • Recommendation: Consider if tokens should be cleared after use or document this security consideration

2. Sensitive Data in Token Storage Files (storage.py:50-52)

   • Token metadata files are stored as JSON with default permissions
   • No explicit file permission setting (should be 0600/owner-only)
   • Recommendation: Set restrictive permissions after creating token files:

   file_path.write_text(json.dumps(data, indent=2, default=str))
   file_path.chmod(0o600)  # Owner read/write only

3. Broad Exception Handling (manager.py:182-188, azure.py:73-77)

   • Catching generic Exception could mask security-relevant errors
   • Recommendation: Be more specific with exception types or at least log the full exception for debugging

4. Hardcoded Databricks Resource ID (azure.py:16)

   • Has a TODO comment but should be addressed before merging
   • Recommendation: Make this configurable via environment variable or config file

---

🐛 Potential Bugs

1. Incorrect Expiration Calculation (databricks.py:75-83)

   • Lines 75-80 perform redundant operations that don't change the datetime
   • The import statement on line 81 is in the middle of the function
   • Critical Issue: This calculation appears incorrect
   • Fix needed:

   from datetime import timedelta  # Move to top
   
   expires_at = datetime.now(timezone.utc).replace(microsecond=0)
   expires_at = expires_at + timedelta(hours=lifetime_hours)

2. Timezone Handling Inconsistency (manager.py:74-76, storage.py:90-91)

   • Multiple places check if tzinfo is None and add UTC
   • This is defensive but suggests timezone handling might not be consistent
   • Recommendation: Ensure TokenInfo.expires_at always stores timezone-aware datetimes via Pydantic validators

3. Silent Failures in Token Loading (storage.py:72-76)

   • Returns None for both missing files and corrupted JSON
   • Users can't distinguish between "never created" vs "corrupted storage"
   • Recommendation: Consider logging warnings for corrupted files or exposing different error states

---

⚡ Performance Considerations

1. Sequential Azure Auth in Concurrent Refresh (manager.py:213-215)

   • refresh_all_tokens creates concurrent tasks, but each calls get_credential() and get_azure_token() separately
   • This could result in multiple redundant Azure auth attempts for the same credential type
   • Recommendation: Consider acquiring Azure token once and reusing it for all environments (with the same workspace URL)

2. Synchronous File I/O (storage.py:52, profiles.py:107)

   • Token storage and profile loading use synchronous file operations in async context
   • Not critical for small files, but could use aiofiles for consistency

---

🧪 Test Coverage

CRITICAL: There are no tests for this PR.

Given the sensitive nature of authentication and token management, tests are essential:

• Unit tests for:
  • Token expiration logic
  • Profile parsing (especially env_var extraction regex)
  • Error handling paths
  • Timezone conversions
• Integration tests for:
  • Token refresh workflow
  • File storage operations
  • CLI command outputs
• Mock tests for:
  • Azure authentication
  • Databricks API calls

Recommendation: Add comprehensive test coverage before merging.

---

📋 Code Quality & Best Practices

Positives:

• ✅ Excellent use of type hints throughout
• ✅ Google-style docstrings are comprehensive
• ✅ Good separation of concerns (auth/dbt/token/storage modules)
• ✅ Proper use of Pydantic models for validation
• ✅ Clean CLI interface with Typer
• ✅ Good error handling with custom exceptions

Improvements:

1. Import Organization (databricks.py:81)

   • from datetime import timedelta should be at the top with other imports

2. Regex Validation (profiles.py:57)

   • The env_var extraction regex is untested and critical for functionality
   • Consider adding validation or more specific patterns

3. Magic Values (databricks.py:30)

   • Lifetime hours constraint (1-24) is validated in CLI but docstring in create_databricks_token doesn't mention this limitation
   • Recommendation: Add validation or document the Databricks API constraint

4. File Path Construction (storage.py:50)

   • Using token_variable as filename might cause issues if variable names contain invalid characters
   • Recommendation: Sanitize the variable name before using as filename

5. Documentation

   • Missing information about where tokens are stored (though the code has good docstrings)
   • No README update explaining the new commands

---

🎯 Specific File Comments

src/databricks_dbt_cli/modules/auth/azure.py

• Line 16: Remove TODO or make resource ID configurable before merge

src/databricks_dbt_cli/modules/auth/databricks.py

• Lines 75-83: Fix the expiration calculation logic
• Line 81: Move import to top of file

src/databricks_dbt_cli/modules/token/storage.py

• Line 52: Add file permission setting after write
• Lines 72-76: Consider logging corrupted file warnings

src/databricks_dbt_cli/modules/token/manager.py

• Line 163: Document security implications of env var storage
• Lines 149-161: Consider reusing Azure token for multiple environments

src/databricks_dbt_cli/commands/dbt/token.py

• Clean implementation, no issues

pyproject.toml

• Good dependency additions
• CLI entry point properly configured

---

🎬 Summary

This PR implements a solid foundation for Databricks token management with good architecture and separation of concerns. However, it has:

Must Fix Before Merge:

1. ❌ Token expiration calculation bug (databricks.py:75-83)
2. ❌ Import placement (databricks.py:81)
3. ❌ File permission security issue (storage.py)
4. ❌ Complete lack of tests

Should Fix:

1. ⚠️ Hardcoded Azure resource ID
2. ⚠️ Broad exception handling
3. ⚠️ Performance optimization for concurrent token refresh

Nice to Have:

• Documentation/README updates
• More specific error messages for debugging
• Filename sanitization in token storage

Overall Assessment: The code quality is good, but the lack of tests and the critical expiration calculation bug need to be addressed before merging. Once these issues are resolved, this will be a valuable addition to the project.

[claude]

Pull Request Review: dbt token authentication commands

Overview

This PR introduces a comprehensive token management system for Databricks with Azure authentication. The implementation is well-structured with clear separation of concerns across auth, dbt profile parsing, and token management modules.

Code Quality & Best Practices

Strengths

✅ Excellent code organization: Clean separation into modules (auth, dbt, token, commands)
✅ Type safety: Strict type hints throughout, following project's ANN rules
✅ Good docstrings: Google-style docstrings on all public functions
✅ Error handling: Custom exception hierarchy with descriptive error messages
✅ Async support: Proper use of asyncio for concurrent token refresh operations
✅ Pydantic models: Leverages Pydantic for validation and data modeling

Areas for Improvement

1. CRITICAL: Empty Azure Resource ID (src/databricks_dbt_cli/modules/auth/azure.py:16)

DATABRICKS_RESOURCE_ID = ""  #!TODO: make configurable
DATABRICKS_SCOPE = f"{DATABRICKS_RESOURCE_ID}/.default"

This will create a malformed scope "/.default" and authentication will fail. This needs to be:

• Either hardcoded to the correct Azure Databricks resource ID (typically 2ff814a6-3304-4ab8-85cb-cd0e6f879c1d)
• Or made configurable via environment variable/config file before merge

2. Redundant datetime manipulation (src/databricks_dbt_cli/modules/auth/databricks.py:75-83)

expires_at = datetime.now(timezone.utc).replace(microsecond=0)
expires_at = expires_at.replace(
    hour=expires_at.hour,
    minute=expires_at.minute,
    second=expires_at.second,
)
from datetime import timedelta  # Import should be at top
expires_at = expires_at + timedelta(hours=lifetime_hours)

Issues:

• The second replace() call does nothing (replaces values with themselves)
• timedelta import should be at module level
• Could be simplified to: expires_at = datetime.now(timezone.utc).replace(microsecond=0) + timedelta(hours=lifetime_hours)

3. Inconsistent timezone handling

Multiple places check and add timezone info defensively (e.g., manager.py:74-76, storage.py:90-91). Consider:

• Always storing datetimes as timezone-aware from creation
• Using a helper function to normalize timezone-aware datetimes

4. Bare Exception catching (src/databricks_dbt_cli/modules/token/manager.py:182)

except Exception as e:
    return TokenRefreshResult(...)

This catches all exceptions including KeyboardInterrupt and system errors. Consider catching specific exceptions or at minimum re-raising system exceptions.

5. Missing newline at EOF (CLAUDE.md, pyproject.toml)

Project should enforce newlines at end of files via ruff configuration.

Security Concerns

HIGH PRIORITY

🔴 Token stored in environment variables (manager.py:163)

os.environ[target.token_env_var] = db_token.token_value

• Setting env vars in os.environ only affects the current process - won't persist for dbt runs
• Consider documenting how users should actually export these for dbt to use
• Environment variables can be exposed via process listings, logging, or error messages
• Consider alternative secure storage (system keychain, encrypted file)

MEDIUM PRIORITY

⚠️ No token cleanup/revocation: When refreshing tokens, old tokens are not revoked via Databricks API, leading to accumulation of valid tokens

⚠️ Token storage permissions: Token metadata files don't explicitly set restrictive permissions (should be 0600)

⚠️ No rate limiting: Concurrent refresh operations could trigger API rate limits

Performance Considerations

✅ Good: Concurrent token refresh using asyncio.gather() (manager.py:215)
✅ Good: Appropriate HTTP timeout of 30 seconds (databricks.py:57)

⚠️ Could improve: Azure credential creation happens for each token refresh. For AUTO mode with ChainedTokenCredential, consider caching the successful credential to avoid retrying all methods.

Test Coverage

❌ CRITICAL: No tests included
This PR adds 1,575 lines of code with 0 test coverage. Essential test scenarios needed:

Unit Tests Needed

1. Auth module:

   • Each auth method credential creation
   • Azure token acquisition (mock azure.identity)
   • Databricks token creation API (mock httpx)
   • Error handling for auth failures

2. Profiles module:

   • YAML parsing with various profile configurations
   • env_var() extraction regex
   • Missing/invalid profile handling
   • Target validation

3. Token storage:

   • Save/load token info
   • Expiration checking
   • Missing file handling
   • Malformed JSON handling

4. Token manager:

   • Token check logic
   • Refresh logic with force flag
   • Concurrent refresh
   • Error propagation

5. CLI commands:

   • Each command with various flags
   • Error message formatting
   • Exit codes

Integration Tests Needed

• End-to-end token refresh flow with mocked Azure/Databricks APIs
• Profile loading with real test YAML files

Additional Issues

1. pyproject.toml: Added D100 to ignore list (module docstrings). Main __init__.py files now have docstrings, but is this needed elsewhere?

2. Hard-coded profile name: DDBT is the default in commands (token.py:57, 109, 154) but should match actual usage patterns or be documented

3. Unused function: get_hours_remaining() in storage.py is defined but never called

4. Missing validation: lifetime_hours parameter allows 1-24 via typer.Option but databricks.py doesn't validate this constraint

5. No logging: Consider adding logging for debugging auth flows and API calls

Recommendations

Must Fix Before Merge

1. ✅ Fix empty DATABRICKS_RESOURCE_ID
2. ✅ Add comprehensive test suite
3. ✅ Document how environment variables are meant to be used by dbt

Should Fix Before Merge

4. ✅ Fix redundant datetime code and move imports to top
5. ✅ Improve exception handling specificity
6. ✅ Add token revocation on refresh
7. ✅ Set restrictive permissions on token storage files

Nice to Have

8. Consider caching Azure credentials
9. Add logging throughout
10. Add rate limiting for API calls

Conclusion

This is a solid foundation for Databricks token management with good architecture and code organization. However, the empty Azure resource ID and lack of tests are blocking issues that must be addressed before merge. The security concern about token storage in environment variables also needs documentation or a better solution.

The code quality is high overall and follows the project's style guidelines well. Once the critical issues are resolved and tests are added, this will be a valuable addition to the project.

---

Recommendation: Request Changes ⚠️