|
1 | 1 | import json |
2 | 2 | import logging |
3 | 3 | import requests |
4 | | -from typing import Dict, Any, Optional, Union, List, Tuple |
| 4 | +from typing import Callable, Dict, Any, Optional, Union, List, Tuple |
5 | 5 | from urllib.parse import urljoin |
6 | 6 |
|
7 | 7 | from databricks.sql.auth.authenticators import AuthProvider |
@@ -87,6 +87,18 @@ def _get_auth_headers(self) -> Dict[str, str]: |
87 | 87 | self.auth_provider.add_headers(headers) |
88 | 88 | return headers |
89 | 89 |
|
| 90 | + def _get_call(self, method: str) -> Callable: |
| 91 | + """Get the appropriate HTTP method function.""" |
| 92 | + method = method.upper() |
| 93 | + if method == "GET": |
| 94 | + return self.session.get |
| 95 | + elif method == "POST": |
| 96 | + return self.session.post |
| 97 | + elif method == "DELETE": |
| 98 | + return self.session.delete |
| 99 | + else: |
| 100 | + raise ValueError(f"Unsupported HTTP method: {method}") |
| 101 | + |
90 | 102 | def _make_request( |
91 | 103 | self, |
92 | 104 | method: str, |
@@ -116,29 +128,13 @@ def _make_request( |
116 | 128 | logger.debug(f"making {method} request to {url}") |
117 | 129 |
|
118 | 130 | try: |
119 | | - if method.upper() == "GET": |
120 | | - response = self.session.get( |
121 | | - url=url, |
122 | | - headers=headers, |
123 | | - json=data, |
124 | | - params=params, |
125 | | - ) |
126 | | - elif method.upper() == "POST": |
127 | | - response = self.session.post( |
128 | | - url=url, |
129 | | - headers=headers, |
130 | | - json=data, |
131 | | - params=params, |
132 | | - ) |
133 | | - elif method.upper() == "DELETE": |
134 | | - response = self.session.delete( |
135 | | - url=url, |
136 | | - headers=headers, |
137 | | - json=data, |
138 | | - params=params, |
139 | | - ) |
140 | | - else: |
141 | | - raise ValueError(f"Unsupported HTTP method: {method}") |
| 131 | + call = self._get_call(method) |
| 132 | + response = call( |
| 133 | + url=url, |
| 134 | + headers=headers, |
| 135 | + json=data, |
| 136 | + params=params, |
| 137 | + ) |
142 | 138 |
|
143 | 139 | # Check for HTTP errors |
144 | 140 | response.raise_for_status() |
|
0 commit comments