Coverage for src/chuck_data/chuck_data/api_client.py: 0%
62 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"""
2Reusable Databricks API client for authentication and requests.
3"""
5import logging
6import requests
9class APIClient:
10 """Reusable Databricks API client for authentication and requests."""
12 def __init__(self, workspace_url, token):
13 """
14 Initialize the API client.
16 Args:
17 workspace_url: Databricks workspace URL
18 token: Databricks API token
19 """
20 # Initialize with workspace URL and token
21 self.workspace_url = workspace_url
22 self.token = token
23 self.headers = {
24 "Authorization": f"Bearer {self.token}",
25 "User-Agent": "amperity",
26 }
28 def get(self, endpoint):
29 """
30 Send a GET request to the Databricks API.
32 Args:
33 endpoint: API endpoint (starting with /)
35 Returns:
36 JSON response from the API
38 Raises:
39 ValueError: If an HTTP error occurs
40 ConnectionError: If a connection error occurs
41 """
42 url = f"{self.workspace_url}{endpoint}"
43 # Construct the full URL for the API request
45 try:
46 response = requests.get(url, headers=self.headers)
47 response.raise_for_status()
48 return response.json()
49 except requests.exceptions.HTTPError as e:
50 logging.error(f"HTTP error: {e}, Response: {response.text}")
51 raise ValueError(f"HTTP error occurred: {e}, Response: {response.text}")
52 except requests.RequestException as e:
53 logging.error(f"Connection error: {e}")
54 raise ConnectionError(f"Connection error occurred: {e}")
56 def post(self, endpoint, data):
57 """
58 Send a POST request to the Databricks API.
60 Args:
61 endpoint: API endpoint (starting with /)
62 data: JSON data to send in the request body
64 Returns:
65 JSON response from the API
67 Raises:
68 ValueError: If an HTTP error occurs
69 ConnectionError: If a connection error occurs
70 """
71 url = f"{self.workspace_url}{endpoint}"
72 logging.debug(f"POST request to: {url}")
74 try:
75 response = requests.post(url, headers=self.headers, json=data)
76 response.raise_for_status()
77 return response.json()
78 except requests.exceptions.HTTPError as e:
79 logging.error(f"HTTP error: {e}, Response: {response.text}")
80 raise ValueError(f"HTTP error occurred: {e}, Response: {response.text}")
81 except requests.RequestException as e:
82 logging.error(f"Connection error: {e}")
83 raise ConnectionError(f"Connection error occurred: {e}")
85 def upload_file(self, path, file_path=None, content=None, overwrite=False):
86 """
87 Upload a file using the /api/2.0/fs/files endpoint.
89 Args:
90 path: The destination path (e.g., "/Volumes/my-catalog/my-schema/my-volume/file.txt")
91 file_path: Local file path to upload (mutually exclusive with content)
92 content: String content to upload (mutually exclusive with file_path)
93 overwrite: Whether to overwrite an existing file
95 Returns:
96 True if successful (API returns no content on success)
98 Raises:
99 ValueError: If both file_path and content are provided or neither is provided
100 ValueError: If an HTTP error occurs
101 ConnectionError: If a connection error occurs
102 """
103 if (file_path and content) or (not file_path and not content):
104 raise ValueError("Exactly one of file_path or content must be provided")
106 # URL encode the path and make sure it starts with a slash
107 import urllib.parse
109 if not path.startswith("/"):
110 path = f"/{path}"
112 # Remove duplicate slashes if any
113 while "//" in path:
114 path = path.replace("//", "/")
116 # URL encode path components but preserve the slashes
117 encoded_path = "/".join(
118 urllib.parse.quote(component) for component in path.split("/") if component
119 )
120 encoded_path = f"/{encoded_path}"
122 url = f"{self.workspace_url}/api/2.0/fs/files{encoded_path}"
124 if overwrite:
125 url += "?overwrite=true"
127 logging.debug(f"File upload request to: {url}")
129 headers = self.headers.copy()
130 headers.update({"Content-Type": "application/octet-stream"})
132 # Get binary data to upload
133 if file_path:
134 with open(file_path, "rb") as f:
135 binary_data = f.read()
136 else:
137 # Convert string content to bytes
138 binary_data = content.encode("utf-8")
140 try:
141 # Use PUT request with raw binary data in the body
142 response = requests.put(url, headers=headers, data=binary_data)
143 response.raise_for_status()
144 # API returns 204 No Content on success
145 return True
146 except requests.exceptions.HTTPError as e:
147 logging.error(f"HTTP error: {e}, Response: {response.text}")
148 raise ValueError(f"HTTP error occurred: {e}, Response: {response.text}")
149 except requests.RequestException as e:
150 logging.error(f"Connection error: {e}")
151 raise ConnectionError(f"Connection error occurred: {e}")