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

1""" 

2Reusable Databricks API client for authentication and requests. 

3""" 

4 

5import logging 

6import requests 

7 

8 

9class APIClient: 

10 """Reusable Databricks API client for authentication and requests.""" 

11 

12 def __init__(self, workspace_url, token): 

13 """ 

14 Initialize the API client. 

15 

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 } 

27 

28 def get(self, endpoint): 

29 """ 

30 Send a GET request to the Databricks API. 

31 

32 Args: 

33 endpoint: API endpoint (starting with /) 

34 

35 Returns: 

36 JSON response from the API 

37 

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 

44 

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}") 

55 

56 def post(self, endpoint, data): 

57 """ 

58 Send a POST request to the Databricks API. 

59 

60 Args: 

61 endpoint: API endpoint (starting with /) 

62 data: JSON data to send in the request body 

63 

64 Returns: 

65 JSON response from the API 

66 

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}") 

73 

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}") 

84 

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. 

88 

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 

94 

95 Returns: 

96 True if successful (API returns no content on success) 

97 

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") 

105 

106 # URL encode the path and make sure it starts with a slash 

107 import urllib.parse 

108 

109 if not path.startswith("/"): 

110 path = f"/{path}" 

111 

112 # Remove duplicate slashes if any 

113 while "//" in path: 

114 path = path.replace("//", "/") 

115 

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}" 

121 

122 url = f"{self.workspace_url}/api/2.0/fs/files{encoded_path}" 

123 

124 if overwrite: 

125 url += "?overwrite=true" 

126 

127 logging.debug(f"File upload request to: {url}") 

128 

129 headers = self.headers.copy() 

130 headers.update({"Content-Type": "application/octet-stream"}) 

131 

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") 

139 

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}")