Coverage for src/commands/auth.py: 80%
46 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 22:56 -0700
« 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"""
5import logging
6from typing import Any, Optional
8from src.clients.amperity import AmperityAPIClient
9from src.clients.databricks import DatabricksAPIClient
10from src.commands.base import CommandResult
11from src.command_registry import CommandDefinition
12from src.config import set_databricks_token
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}")
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}")
29 return CommandResult(True, message=message)
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")
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}")
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")
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}")
64 if service in ["all", "amperity"]:
65 from src.config import set_amperity_token
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}")
74 return CommandResult(True, message=f"Successfully logged out from {service}")
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]