Coverage for src/chuck_data/databricks_auth.py: 0%

24 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-06-05 22:56 -0700

1"""Helper functions for Databricks authentication.""" 

2 

3import logging 

4import os 

5 

6from .clients.databricks import DatabricksAPIClient 

7from .config import get_workspace_url, get_databricks_token as get_token_from_config 

8 

9 

10def get_databricks_token() -> str: 

11 """Return the Databricks token from config or environment.""" 

12 token = get_token_from_config() 

13 if not token: 

14 token = os.getenv("DATABRICKS_TOKEN") 

15 if token: 

16 logging.info( 

17 "Using Databricks token from environment variable. Consider using /set-token to save it to config." 

18 ) 

19 if not token: 

20 raise EnvironmentError("Databricks token not found in config or environment!") 

21 return token 

22 

23 

24def validate_databricks_token(token: str) -> bool: 

25 """Validate the provided Databricks token.""" 

26 try: 

27 workspace_url = get_workspace_url() 

28 client = DatabricksAPIClient(workspace_url, token) 

29 return client.validate_token() 

30 except ValueError as e: 

31 logging.error("Token validation failed: %s", e) 

32 return False 

33 except ConnectionError as e: 

34 logging.error("Connection error during token validation: %s", e) 

35 raise ConnectionError(f"Connection Error: {e}")