Coverage for src/commands/create_volume.py: 29%

28 statements  

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

1""" 

2Command for creating a new volume in Unity Catalog. 

3""" 

4 

5from typing import Optional, Any 

6from src.clients.databricks import DatabricksAPIClient 

7from src.commands.base import CommandResult 

8from src.config import get_active_catalog, get_active_schema 

9from src.command_registry import CommandDefinition 

10import logging 

11 

12 

13def handle_command( 

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

15) -> CommandResult: 

16 """ 

17 Create a new volume in Unity Catalog. 

18 

19 Args: 

20 client: DatabricksAPIClient instance for API calls 

21 **kwargs: Command parameters 

22 - catalog_name: Name of the catalog where the volume will be created 

23 - schema_name: Name of the schema where the volume will be created 

24 - name: Name for the new volume 

25 - volume_type: Type of volume to create (default: "MANAGED") 

26 

27 Returns: 

28 CommandResult with created volume details if successful 

29 """ 

30 if not client: 

31 return CommandResult( 

32 False, 

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

34 ) 

35 

36 # Extract parameters 

37 catalog_name = kwargs.get("catalog_name") 

38 schema_name = kwargs.get("schema_name") 

39 name = kwargs.get("name") 

40 volume_type = kwargs.get("volume_type", "MANAGED") 

41 

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

43 if not catalog_name: 

44 catalog_name = get_active_catalog() 

45 if not catalog_name: 

46 return CommandResult( 

47 False, 

48 message="No catalog specified and no active catalog selected. Please provide a catalog_name or select a catalog first using /select-catalog.", 

49 ) 

50 

51 # If schema_name not provided, try to use active schema 

52 if not schema_name: 

53 schema_name = get_active_schema() 

54 if not schema_name: 

55 return CommandResult( 

56 False, 

57 message="No schema specified and no active schema selected. Please provide a schema_name or select a schema first using /select-schema.", 

58 ) 

59 

60 try: 

61 # Create the volume 

62 result = client.create_volume( 

63 catalog_name=catalog_name, 

64 schema_name=schema_name, 

65 name=name, 

66 volume_type=volume_type, 

67 ) 

68 

69 return CommandResult( 

70 True, 

71 data=result, 

72 message=f"Successfully created volume '{catalog_name}.{schema_name}.{name}' of type '{volume_type}'.", 

73 ) 

74 except Exception as e: 

75 logging.error(f"Error creating volume: {str(e)}") 

76 return CommandResult( 

77 False, message=f"Failed to create volume: {str(e)}", error=e 

78 ) 

79 

80 

81DEFINITION = CommandDefinition( 

82 name="create-volume", 

83 description="Create a new volume in Unity Catalog.", 

84 handler=handle_command, 

85 parameters={ 

86 "catalog_name": { 

87 "type": "string", 

88 "description": "Name of the catalog where the volume will be created.", 

89 }, 

90 "schema_name": { 

91 "type": "string", 

92 "description": "Name of the schema where the volume will be created.", 

93 }, 

94 "name": {"type": "string", "description": "Name for the new volume."}, 

95 "volume_type": { 

96 "type": "string", 

97 "description": "Type of volume to create (MANAGED or EXTERNAL).", 

98 "default": "MANAGED", 

99 }, 

100 }, 

101 required_params=[ 

102 "name" 

103 ], # Only name is required; catalog and schema can come from active config 

104 tui_aliases=["/create-volume"], 

105 needs_api_client=True, 

106 visible_to_user=True, 

107 visible_to_agent=True, 

108 usage_hint="Usage: /create-volume --name <volume_name> [--catalog_name <catalog>] [--schema_name <schema>] [--volume_type MANAGED|EXTERNAL]\n(Uses active catalog/schema if not specified)", 

109)