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
« 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"""
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
13def handle_command(
14 client: Optional[DatabricksAPIClient], **kwargs: Any
15) -> CommandResult:
16 """
17 Upload a file to Databricks volumes or DBFS.
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)
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 )
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")
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 )
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 )
56 if local_path and not os.path.isfile(local_path):
57 return CommandResult(False, message=f"Local file not found: {local_path}")
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 )
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 )
85 upload_type = "volumes"
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)
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)