Coverage for src/chuck_data/commands/auth.py: 0%

46 statements  

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

1""" 

2Commands for user authentication with Amperity and Databricks. 

3""" 

4 

5import logging 

6from typing import Any, Optional 

7 

8from ..clients.amperity import AmperityAPIClient 

9from ..clients.databricks import DatabricksAPIClient 

10from ..commands.base import CommandResult 

11from ..command_registry import CommandDefinition 

12from ..config import set_databricks_token 

13 

14 

15def handle_amperity_login( 

16 client: Optional[DatabricksAPIClient], **kwargs: Any 

17) -> CommandResult: 

18 """Handle the login command for Amperity.""" 

19 auth_client = AmperityAPIClient() 

20 success, message = auth_client.start_auth() 

21 if not success: 

22 return CommandResult(False, message=f"Login failed: {message}") 

23 

24 # Wait for auth completion 

25 success, message = auth_client.wait_for_auth_completion() 

26 if not success: 

27 return CommandResult(False, message=f"Login failed: {message}") 

28 

29 return CommandResult(True, message=message) 

30 

31 

32def handle_databricks_login( 

33 client: Optional[DatabricksAPIClient], **kwargs: Any 

34) -> CommandResult: 

35 """Handle the login command for Databricks.""" 

36 token = kwargs.get("token") 

37 if not token: 

38 return CommandResult(False, message="Token parameter is required") 

39 

40 # Save token to config 

41 try: 

42 set_databricks_token(token) 

43 logging.info("Databricks token set successfully") 

44 return CommandResult(True, message="Databricks token set successfully") 

45 except Exception as e: 

46 logging.error("Failed to set Databricks token: %s", e) 

47 return CommandResult(False, message=f"Failed to set token: {e}") 

48 

49 

50def handle_logout( 

51 client: Optional[DatabricksAPIClient], **kwargs: Any 

52) -> CommandResult: 

53 """Handle the logout command for Amperity by default.""" 

54 service = kwargs.get("service", "amperity") 

55 

56 if service in ["all", "databricks"]: 

57 try: 

58 set_databricks_token("") 

59 logging.info("Databricks token cleared") 

60 except Exception as e: 

61 logging.error("Error clearing Databricks token: %s", e) 

62 return CommandResult(False, message=f"Error clearing Databricks token: {e}") 

63 

64 if service in ["all", "amperity"]: 

65 from ..config import set_amperity_token 

66 

67 try: 

68 set_amperity_token("") 

69 logging.info("Amperity token cleared") 

70 except Exception as e: 

71 logging.error("Error clearing Amperity token: %s", e) 

72 return CommandResult(False, message=f"Error clearing Amperity token: {e}") 

73 

74 return CommandResult(True, message=f"Successfully logged out from {service}") 

75 

76 

77DEFINITION = [ 

78 CommandDefinition( 

79 name="amperity-login", 

80 description="Log in to Amperity", 

81 handler=handle_amperity_login, 

82 parameters={}, 

83 required_params=[], 

84 tui_aliases=["/login", "/amperity-login"], 

85 needs_api_client=False, 

86 visible_to_user=True, 

87 visible_to_agent=False, 

88 ), 

89 CommandDefinition( 

90 name="databricks-login", 

91 description="Set Databricks API token", 

92 handler=handle_databricks_login, 

93 parameters={ 

94 "token": {"type": "string", "description": "Your Databricks API token"} 

95 }, 

96 required_params=["token"], 

97 tui_aliases=["/databricks-login", "/set-token"], 

98 needs_api_client=False, 

99 visible_to_user=True, 

100 visible_to_agent=False, 

101 ), 

102 CommandDefinition( 

103 name="logout", 

104 description="Log out from Amperity (default) or other authentication services", 

105 handler=handle_logout, 

106 parameters={ 

107 "service": { 

108 "type": "string", 

109 "description": "Service to log out from (amperity, databricks, or all)", 

110 "default": "amperity", 

111 } 

112 }, 

113 required_params=[], 

114 tui_aliases=["/logout"], 

115 needs_api_client=False, 

116 visible_to_user=True, 

117 visible_to_agent=False, 

118 ), 

119]