Coverage for src/commands/schema_selection.py: 88%

26 statements  

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

1""" 

2Command handler for schema selection. 

3 

4This module contains the handler for setting the active schema 

5for database operations. 

6""" 

7 

8import logging 

9from typing import Optional 

10 

11from src.clients.databricks import DatabricksAPIClient 

12from src.command_registry import CommandDefinition 

13from src.config import get_active_catalog, set_active_schema 

14from .base import CommandResult 

15 

16 

17def handle_command(client: Optional[DatabricksAPIClient], **kwargs) -> CommandResult: 

18 """ 

19 Set the active schema. 

20 

21 Args: 

22 client: API client instance 

23 **kwargs: schema_name (str) 

24 """ 

25 schema_name: str = kwargs.get("schema_name") 

26 if not schema_name: 

27 return CommandResult(False, message="schema_name parameter is required.") 

28 

29 try: 

30 catalog_name = get_active_catalog() 

31 if not catalog_name: 

32 return CommandResult(False, message="No active catalog selected.") 

33 

34 try: 

35 from src.catalogs import get_schema 

36 

37 get_schema(client, f"{catalog_name}.{schema_name}") 

38 except Exception: 

39 set_active_schema(schema_name) # Set anyway 

40 return CommandResult( 

41 True, 

42 message=f"Warning: Could not verify schema '{schema_name}' in catalog '{catalog_name}'. Setting anyway.", 

43 data={"schema_name": schema_name, "catalog_name": catalog_name}, 

44 ) 

45 

46 set_active_schema(schema_name) 

47 return CommandResult( 

48 True, 

49 message=f"Active schema is now set to '{schema_name}' in catalog '{catalog_name}'.", 

50 data={"schema_name": schema_name, "catalog_name": catalog_name}, 

51 ) 

52 except Exception as e: 

53 logging.error(f"Failed to set schema '{schema_name}': {e}", exc_info=True) 

54 return CommandResult(False, error=e, message=str(e)) 

55 

56 

57DEFINITION = CommandDefinition( 

58 name="set-schema", 

59 description="Set the active schema for database operations", 

60 handler=handle_command, 

61 parameters={ 

62 "schema_name": { 

63 "type": "string", 

64 "description": "Name of the schema to set as active", 

65 } 

66 }, 

67 required_params=["schema_name"], 

68 tui_aliases=["/select-schema"], 

69 visible_to_user=True, 

70 visible_to_agent=True, 

71 condensed_action="Setting schema", 

72)