Coverage for src/commands/upload_file.py: 21%

38 statements  

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

1""" 

2Command for uploading files to Databricks volumes or DBFS. 

3""" 

4 

5from typing import Optional, Any 

6from src.clients.databricks import DatabricksAPIClient 

7from src.commands.base import CommandResult 

8from src.command_registry import CommandDefinition 

9import os 

10import logging 

11 

12 

13def handle_command( 

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

15) -> CommandResult: 

16 """ 

17 Upload a file to Databricks volumes or DBFS. 

18 

19 Args: 

20 client: DatabricksAPIClient instance for API calls 

21 **kwargs: Command parameters 

22 - local_path: Path to local file to upload 

23 - destination_path: Path in Databricks where the file should be uploaded 

24 - overwrite: Whether to overwrite existing files (optional) 

25 - use_dbfs: Whether to use DBFS API instead of volumes API (optional) 

26 - contents: String content to upload instead of a file (optional, mutually exclusive with local_path) 

27 

28 Returns: 

29 CommandResult with upload status 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 local_path = kwargs.get("local_path") 

39 destination_path = kwargs.get("destination_path") 

40 overwrite = kwargs.get("overwrite", False) 

41 use_dbfs = kwargs.get("use_dbfs", False) 

42 contents = kwargs.get("contents") 

43 

44 # Validation 

45 if local_path and contents: 

46 return CommandResult( 

47 False, 

48 message="You cannot specify both local_path and contents. Choose one method to provide file content.", 

49 ) 

50 

51 if not local_path and not contents: 

52 return CommandResult( 

53 False, message="You must provide either local_path or contents to upload." 

54 ) 

55 

56 if local_path and not os.path.isfile(local_path): 

57 return CommandResult(False, message=f"Local file not found: {local_path}") 

58 

59 try: 

60 # Choose the appropriate upload method based on parameters 

61 if use_dbfs: 

62 if contents: 

63 client.store_dbfs_file( 

64 path=destination_path, contents=contents, overwrite=overwrite 

65 ) 

66 else: 

67 with open(local_path, "r") as file: 

68 file_contents = file.read() 

69 client.store_dbfs_file( 

70 path=destination_path, contents=file_contents, overwrite=overwrite 

71 ) 

72 

73 upload_type = "DBFS" 

74 else: 

75 # Upload to volumes 

76 if contents: 

77 client.upload_file( 

78 path=destination_path, content=contents, overwrite=overwrite 

79 ) 

80 else: 

81 client.upload_file( 

82 path=destination_path, file_path=local_path, overwrite=overwrite 

83 ) 

84 

85 upload_type = "volumes" 

86 

87 source = local_path if local_path else "provided content" 

88 return CommandResult( 

89 True, 

90 message=f"Successfully uploaded {source} to {upload_type} path: {destination_path}", 

91 ) 

92 except Exception as e: 

93 logging.error(f"Error uploading file: {str(e)}") 

94 return CommandResult(False, message=f"Failed to upload file: {str(e)}", error=e) 

95 

96 

97DEFINITION = CommandDefinition( 

98 name="upload-file", 

99 description="Upload a file to Databricks volumes or DBFS.", 

100 handler=handle_command, 

101 parameters={ 

102 "local_path": { 

103 "type": "string", 

104 "description": "Path to local file to upload.", 

105 }, 

106 "destination_path": { 

107 "type": "string", 

108 "description": "Path in Databricks where the file should be uploaded.", 

109 }, 

110 "contents": { 

111 "type": "string", 

112 "description": "String content to upload instead of a file (mutually exclusive with local_path).", 

113 }, 

114 "overwrite": { 

115 "type": "boolean", 

116 "description": "Whether to overwrite existing files.", 

117 "default": False, 

118 }, 

119 "use_dbfs": { 

120 "type": "boolean", 

121 "description": "Whether to use DBFS API instead of volumes API.", 

122 "default": False, 

123 }, 

124 }, 

125 required_params=["destination_path"], 

126 tui_aliases=["/upload", "/upload-file"], 

127 needs_api_client=True, 

128 visible_to_user=True, 

129 visible_to_agent=True, 

130 usage_hint="Usage: /upload --local_path <file_path> --destination_path <dbx_path> [--overwrite true|false] [--use_dbfs true|false]\n" 

131 + 'Or: /upload --contents "file content" --destination_path <dbx_path> [--overwrite true|false] [--use_dbfs true|false]', 

132)