Coverage for src/commands/list_schemas.py: 91%

33 statements  

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

1""" 

2Command for listing schemas in a Unity Catalog catalog. 

3""" 

4 

5from typing import Optional, Any 

6from src.clients.databricks import DatabricksAPIClient 

7from src.commands.base import CommandResult 

8from src.catalogs import list_schemas as get_schemas_list 

9from src.config import get_active_catalog 

10from src.command_registry import CommandDefinition 

11import logging 

12 

13 

14def handle_command( 

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

16) -> CommandResult: 

17 """ 

18 List schemas in a Unity Catalog catalog. 

19 

20 Args: 

21 client: DatabricksAPIClient instance for API calls 

22 **kwargs: Command parameters 

23 - catalog_name: Name of the catalog to list schemas from (optional, uses active catalog if not provided) 

24 - include_browse: Whether to include schemas with selective metadata access (optional) 

25 - max_results: Maximum number of schemas to return (optional) 

26 - page_token: Opaque pagination token to go to next page (optional) 

27 

28 Returns: 

29 CommandResult with list of schemas if successful 

30 """ 

31 if not client: 

32 return CommandResult( 

33 False, 

34 message="No Databricks client available. Please set up your workspace first.", 

35 ) 

36 

37 # Extract parameters 

38 catalog_name = kwargs.get("catalog_name") 

39 include_browse = kwargs.get("include_browse", False) 

40 max_results = kwargs.get("max_results") 

41 page_token = kwargs.get("page_token") 

42 

43 # If catalog_name not provided, try to use active catalog 

44 if not catalog_name: 

45 catalog_name = get_active_catalog() 

46 if not catalog_name: 

47 return CommandResult( 

48 False, 

49 message="No catalog specified and no active catalog selected. Please provide a catalog_name or select a catalog first.", 

50 ) 

51 

52 try: 

53 # List schemas in the catalog 

54 result = get_schemas_list( 

55 client=client, 

56 catalog_name=catalog_name, 

57 include_browse=include_browse, 

58 max_results=max_results, 

59 page_token=page_token, 

60 ) 

61 

62 schemas = result.get("schemas", []) 

63 next_page_token = result.get("next_page_token") 

64 

65 if not schemas: 

66 return CommandResult( 

67 True, message=f"No schemas found in catalog '{catalog_name}'." 

68 ) 

69 

70 # Format schema information for display 

71 formatted_schemas = [] 

72 for schema in schemas: 

73 formatted_schema = { 

74 "name": schema.get("name"), 

75 "full_name": schema.get("full_name"), 

76 "catalog_name": schema.get("catalog_name"), 

77 "comment": schema.get("comment", ""), 

78 "created_at": schema.get("created_at"), 

79 "created_by": schema.get("created_by", ""), 

80 "owner": schema.get("owner", ""), 

81 } 

82 formatted_schemas.append(formatted_schema) 

83 

84 return CommandResult( 

85 True, 

86 data={ 

87 "schemas": formatted_schemas, 

88 "total_count": len(formatted_schemas), 

89 "catalog_name": catalog_name, 

90 "next_page_token": next_page_token, 

91 }, 

92 message=f"Found {len(formatted_schemas)} schema(s) in catalog '{catalog_name}'." 

93 + ( 

94 f" More schemas available with page token: {next_page_token}" 

95 if next_page_token 

96 else "" 

97 ), 

98 ) 

99 except Exception as e: 

100 logging.error(f"Error listing schemas: {str(e)}") 

101 return CommandResult( 

102 False, message=f"Failed to list schemas: {str(e)}", error=e 

103 ) 

104 

105 

106DEFINITION = CommandDefinition( 

107 name="list-schemas", 

108 description="List schemas in a Unity Catalog catalog.", 

109 handler=handle_command, 

110 parameters={ 

111 "catalog_name": { 

112 "type": "string", 

113 "description": "Name of the catalog to list schemas from (uses active catalog if not provided).", 

114 }, 

115 "include_browse": { 

116 "type": "boolean", 

117 "description": "Whether to include schemas with selective metadata access.", 

118 "default": False, 

119 }, 

120 "max_results": { 

121 "type": "integer", 

122 "description": "Maximum number of schemas to return.", 

123 }, 

124 "page_token": { 

125 "type": "string", 

126 "description": "Opaque pagination token to go to next page.", 

127 }, 

128 }, 

129 required_params=[], 

130 tui_aliases=["/schemas"], 

131 needs_api_client=True, 

132 visible_to_user=True, 

133 visible_to_agent=True, 

134 agent_display="full", # Show full schema list to agents 

135 usage_hint="Usage: /list-schemas [--catalog_name <catalog>] [--include_browse true|false] [--max_results <number>]", 

136)