databricks-sdk 0.44.1__py3-none-any.whl → 0.45.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of databricks-sdk might be problematic. Click here for more details.

Files changed (63) hide show
  1. databricks/sdk/__init__.py +123 -115
  2. databricks/sdk/_base_client.py +112 -88
  3. databricks/sdk/_property.py +12 -7
  4. databricks/sdk/_widgets/__init__.py +13 -2
  5. databricks/sdk/_widgets/default_widgets_utils.py +21 -15
  6. databricks/sdk/_widgets/ipywidgets_utils.py +47 -24
  7. databricks/sdk/azure.py +8 -6
  8. databricks/sdk/casing.py +5 -5
  9. databricks/sdk/config.py +152 -99
  10. databricks/sdk/core.py +57 -47
  11. databricks/sdk/credentials_provider.py +300 -205
  12. databricks/sdk/data_plane.py +86 -3
  13. databricks/sdk/dbutils.py +123 -87
  14. databricks/sdk/environments.py +52 -35
  15. databricks/sdk/errors/base.py +61 -35
  16. databricks/sdk/errors/customizer.py +3 -3
  17. databricks/sdk/errors/deserializer.py +38 -25
  18. databricks/sdk/errors/details.py +417 -0
  19. databricks/sdk/errors/mapper.py +1 -1
  20. databricks/sdk/errors/overrides.py +27 -24
  21. databricks/sdk/errors/parser.py +26 -14
  22. databricks/sdk/errors/platform.py +10 -10
  23. databricks/sdk/errors/private_link.py +24 -24
  24. databricks/sdk/logger/round_trip_logger.py +28 -20
  25. databricks/sdk/mixins/compute.py +90 -60
  26. databricks/sdk/mixins/files.py +815 -145
  27. databricks/sdk/mixins/jobs.py +191 -16
  28. databricks/sdk/mixins/open_ai_client.py +26 -20
  29. databricks/sdk/mixins/workspace.py +45 -34
  30. databricks/sdk/oauth.py +372 -196
  31. databricks/sdk/retries.py +14 -12
  32. databricks/sdk/runtime/__init__.py +34 -17
  33. databricks/sdk/runtime/dbutils_stub.py +52 -39
  34. databricks/sdk/service/_internal.py +12 -7
  35. databricks/sdk/service/apps.py +618 -418
  36. databricks/sdk/service/billing.py +827 -604
  37. databricks/sdk/service/catalog.py +6552 -4474
  38. databricks/sdk/service/cleanrooms.py +550 -388
  39. databricks/sdk/service/compute.py +5241 -3531
  40. databricks/sdk/service/dashboards.py +1313 -923
  41. databricks/sdk/service/files.py +442 -309
  42. databricks/sdk/service/iam.py +2115 -1483
  43. databricks/sdk/service/jobs.py +4151 -2588
  44. databricks/sdk/service/marketplace.py +2210 -1517
  45. databricks/sdk/service/ml.py +3364 -2255
  46. databricks/sdk/service/oauth2.py +922 -584
  47. databricks/sdk/service/pipelines.py +1865 -1203
  48. databricks/sdk/service/provisioning.py +1435 -1029
  49. databricks/sdk/service/serving.py +2040 -1278
  50. databricks/sdk/service/settings.py +2846 -1929
  51. databricks/sdk/service/sharing.py +2201 -877
  52. databricks/sdk/service/sql.py +4650 -3103
  53. databricks/sdk/service/vectorsearch.py +816 -550
  54. databricks/sdk/service/workspace.py +1330 -906
  55. databricks/sdk/useragent.py +36 -22
  56. databricks/sdk/version.py +1 -1
  57. {databricks_sdk-0.44.1.dist-info → databricks_sdk-0.45.0.dist-info}/METADATA +31 -31
  58. databricks_sdk-0.45.0.dist-info/RECORD +70 -0
  59. {databricks_sdk-0.44.1.dist-info → databricks_sdk-0.45.0.dist-info}/WHEEL +1 -1
  60. databricks_sdk-0.44.1.dist-info/RECORD +0 -69
  61. {databricks_sdk-0.44.1.dist-info → databricks_sdk-0.45.0.dist-info}/LICENSE +0 -0
  62. {databricks_sdk-0.44.1.dist-info → databricks_sdk-0.45.0.dist-info}/NOTICE +0 -0
  63. {databricks_sdk-0.44.1.dist-info → databricks_sdk-0.45.0.dist-info}/top_level.txt +0 -0
databricks/sdk/config.py CHANGED
@@ -1,5 +1,6 @@
1
1
  import configparser
2
2
  import copy
3
+ import datetime
3
4
  import logging
4
5
  import os
5
6
  import pathlib
@@ -19,11 +20,11 @@ from .oauth import (OidcEndpoints, Token, get_account_endpoints,
19
20
  get_azure_entra_id_workspace_endpoints,
20
21
  get_workspace_endpoints)
21
22
 
22
- logger = logging.getLogger('databricks.sdk')
23
+ logger = logging.getLogger("databricks.sdk")
23
24
 
24
25
 
25
26
  class ConfigAttribute:
26
- """ Configuration attribute metadata and descriptor protocols. """
27
+ """Configuration attribute metadata and descriptor protocols."""
27
28
 
28
29
  # name and transform are discovered from Config.__new__
29
30
  name: str = None
@@ -34,12 +35,12 @@ class ConfigAttribute:
34
35
  self.auth = auth
35
36
  self.sensitive = sensitive
36
37
 
37
- def __get__(self, cfg: 'Config', owner):
38
+ def __get__(self, cfg: "Config", owner):
38
39
  if not cfg:
39
40
  return None
40
41
  return cfg._inner.get(self.name, None)
41
42
 
42
- def __set__(self, cfg: 'Config', value: any):
43
+ def __set__(self, cfg: "Config", value: any):
43
44
  cfg._inner[self.name] = self.transform(value)
44
45
 
45
46
  def __repr__(self) -> str:
@@ -57,71 +58,120 @@ def with_user_agent_extra(key: str, value: str):
57
58
 
58
59
 
59
60
  class Config:
60
- host: str = ConfigAttribute(env='DATABRICKS_HOST')
61
- account_id: str = ConfigAttribute(env='DATABRICKS_ACCOUNT_ID')
62
- token: str = ConfigAttribute(env='DATABRICKS_TOKEN', auth='pat', sensitive=True)
63
- username: str = ConfigAttribute(env='DATABRICKS_USERNAME', auth='basic')
64
- password: str = ConfigAttribute(env='DATABRICKS_PASSWORD', auth='basic', sensitive=True)
65
- client_id: str = ConfigAttribute(env='DATABRICKS_CLIENT_ID', auth='oauth')
66
- client_secret: str = ConfigAttribute(env='DATABRICKS_CLIENT_SECRET', auth='oauth', sensitive=True)
67
- profile: str = ConfigAttribute(env='DATABRICKS_CONFIG_PROFILE')
68
- config_file: str = ConfigAttribute(env='DATABRICKS_CONFIG_FILE')
69
- google_service_account: str = ConfigAttribute(env='DATABRICKS_GOOGLE_SERVICE_ACCOUNT', auth='google')
70
- google_credentials: str = ConfigAttribute(env='GOOGLE_CREDENTIALS', auth='google', sensitive=True)
71
- azure_workspace_resource_id: str = ConfigAttribute(env='DATABRICKS_AZURE_RESOURCE_ID', auth='azure')
72
- azure_use_msi: bool = ConfigAttribute(env='ARM_USE_MSI', auth='azure')
73
- azure_client_secret: str = ConfigAttribute(env='ARM_CLIENT_SECRET', auth='azure', sensitive=True)
74
- azure_client_id: str = ConfigAttribute(env='ARM_CLIENT_ID', auth='azure')
75
- azure_tenant_id: str = ConfigAttribute(env='ARM_TENANT_ID', auth='azure')
76
- azure_environment: str = ConfigAttribute(env='ARM_ENVIRONMENT')
77
- databricks_cli_path: str = ConfigAttribute(env='DATABRICKS_CLI_PATH')
78
- auth_type: str = ConfigAttribute(env='DATABRICKS_AUTH_TYPE')
79
- cluster_id: str = ConfigAttribute(env='DATABRICKS_CLUSTER_ID')
80
- warehouse_id: str = ConfigAttribute(env='DATABRICKS_WAREHOUSE_ID')
81
- serverless_compute_id: str = ConfigAttribute(env='DATABRICKS_SERVERLESS_COMPUTE_ID')
61
+ host: str = ConfigAttribute(env="DATABRICKS_HOST")
62
+ account_id: str = ConfigAttribute(env="DATABRICKS_ACCOUNT_ID")
63
+ token: str = ConfigAttribute(env="DATABRICKS_TOKEN", auth="pat", sensitive=True)
64
+ username: str = ConfigAttribute(env="DATABRICKS_USERNAME", auth="basic")
65
+ password: str = ConfigAttribute(env="DATABRICKS_PASSWORD", auth="basic", sensitive=True)
66
+ client_id: str = ConfigAttribute(env="DATABRICKS_CLIENT_ID", auth="oauth")
67
+ client_secret: str = ConfigAttribute(env="DATABRICKS_CLIENT_SECRET", auth="oauth", sensitive=True)
68
+ profile: str = ConfigAttribute(env="DATABRICKS_CONFIG_PROFILE")
69
+ config_file: str = ConfigAttribute(env="DATABRICKS_CONFIG_FILE")
70
+ google_service_account: str = ConfigAttribute(env="DATABRICKS_GOOGLE_SERVICE_ACCOUNT", auth="google")
71
+ google_credentials: str = ConfigAttribute(env="GOOGLE_CREDENTIALS", auth="google", sensitive=True)
72
+ azure_workspace_resource_id: str = ConfigAttribute(env="DATABRICKS_AZURE_RESOURCE_ID", auth="azure")
73
+ azure_use_msi: bool = ConfigAttribute(env="ARM_USE_MSI", auth="azure")
74
+ azure_client_secret: str = ConfigAttribute(env="ARM_CLIENT_SECRET", auth="azure", sensitive=True)
75
+ azure_client_id: str = ConfigAttribute(env="ARM_CLIENT_ID", auth="azure")
76
+ azure_tenant_id: str = ConfigAttribute(env="ARM_TENANT_ID", auth="azure")
77
+ azure_environment: str = ConfigAttribute(env="ARM_ENVIRONMENT")
78
+ databricks_cli_path: str = ConfigAttribute(env="DATABRICKS_CLI_PATH")
79
+ auth_type: str = ConfigAttribute(env="DATABRICKS_AUTH_TYPE")
80
+ cluster_id: str = ConfigAttribute(env="DATABRICKS_CLUSTER_ID")
81
+ warehouse_id: str = ConfigAttribute(env="DATABRICKS_WAREHOUSE_ID")
82
+ serverless_compute_id: str = ConfigAttribute(env="DATABRICKS_SERVERLESS_COMPUTE_ID")
82
83
  skip_verify: bool = ConfigAttribute()
83
84
  http_timeout_seconds: float = ConfigAttribute()
84
- debug_truncate_bytes: int = ConfigAttribute(env='DATABRICKS_DEBUG_TRUNCATE_BYTES')
85
- debug_headers: bool = ConfigAttribute(env='DATABRICKS_DEBUG_HEADERS')
86
- rate_limit: int = ConfigAttribute(env='DATABRICKS_RATE_LIMIT')
85
+ debug_truncate_bytes: int = ConfigAttribute(env="DATABRICKS_DEBUG_TRUNCATE_BYTES")
86
+ debug_headers: bool = ConfigAttribute(env="DATABRICKS_DEBUG_HEADERS")
87
+ rate_limit: int = ConfigAttribute(env="DATABRICKS_RATE_LIMIT")
87
88
  retry_timeout_seconds: int = ConfigAttribute()
88
- metadata_service_url = ConfigAttribute(env='DATABRICKS_METADATA_SERVICE_URL',
89
- auth='metadata-service',
90
- sensitive=True)
89
+ metadata_service_url = ConfigAttribute(
90
+ env="DATABRICKS_METADATA_SERVICE_URL",
91
+ auth="metadata-service",
92
+ sensitive=True,
93
+ )
91
94
  max_connection_pools: int = ConfigAttribute()
92
95
  max_connections_per_pool: int = ConfigAttribute()
93
96
  databricks_environment: Optional[DatabricksEnvironment] = None
94
97
 
95
- enable_experimental_files_api_client: bool = ConfigAttribute(
96
- env='DATABRICKS_ENABLE_EXPERIMENTAL_FILES_API_CLIENT')
98
+ enable_experimental_files_api_client: bool = ConfigAttribute(env="DATABRICKS_ENABLE_EXPERIMENTAL_FILES_API_CLIENT")
97
99
  files_api_client_download_max_total_recovers = None
98
100
  files_api_client_download_max_total_recovers_without_progressing = 1
99
101
 
102
+ # File multipart upload parameters
103
+ # ----------------------
104
+
105
+ # Minimal input stream size (bytes) to use multipart / resumable uploads.
106
+ # For small files it's more efficient to make one single-shot upload request.
107
+ # When uploading a file, SDK will initially buffer this many bytes from input stream.
108
+ # This parameter can be less or bigger than multipart_upload_chunk_size.
109
+ multipart_upload_min_stream_size: int = 5 * 1024 * 1024
110
+
111
+ # Maximum number of presigned URLs that can be requested at a time.
112
+ #
113
+ # The more URLs we request at once, the higher chance is that some of the URLs will expire
114
+ # before we get to use it. We discover the presigned URL is expired *after* sending the
115
+ # input stream partition to the server. So to retry the upload of this partition we must rewind
116
+ # the stream back. In case of a non-seekable stream we cannot rewind, so we'll abort
117
+ # the upload. To reduce the chance of this, we're requesting presigned URLs one by one
118
+ # and using them immediately.
119
+ multipart_upload_batch_url_count: int = 1
120
+
121
+ # Size of the chunk to use for multipart uploads.
122
+ #
123
+ # The smaller chunk is, the less chance for network errors (or URL get expired),
124
+ # but the more requests we'll make.
125
+ # For AWS, minimum is 5Mb: https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html
126
+ # For GCP, minimum is 256 KiB (and also recommended multiple is 256 KiB)
127
+ # boto uses 8Mb: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/customizations/s3.html#boto3.s3.transfer.TransferConfig
128
+ multipart_upload_chunk_size: int = 10 * 1024 * 1024
129
+
130
+ # use maximum duration of 1 hour
131
+ multipart_upload_url_expiration_duration: datetime.timedelta = datetime.timedelta(hours=1)
132
+
133
+ # This is not a "wall time" cutoff for the whole upload request,
134
+ # but a maximum time between consecutive data reception events (even 1 byte) from the server
135
+ multipart_upload_single_chunk_upload_timeout_seconds: float = 60
136
+
137
+ # Cap on the number of custom retries during incremental uploads:
138
+ # 1) multipart: upload part URL is expired, so new upload URLs must be requested to continue upload
139
+ # 2) resumable: chunk upload produced a retryable response (or exception), so upload status must be
140
+ # retrieved to continue the upload.
141
+ # In these two cases standard SDK retries (which are capped by the `retry_timeout_seconds` option) are not used.
142
+ # Note that retry counter is reset when upload is successfully resumed.
143
+ multipart_upload_max_retries = 3
144
+
100
145
  def __init__(
101
- self,
102
- *,
103
- # Deprecated. Use credentials_strategy instead.
104
- credentials_provider: Optional[CredentialsStrategy] = None,
105
- credentials_strategy: Optional[CredentialsStrategy] = None,
106
- product=None,
107
- product_version=None,
108
- clock: Optional[Clock] = None,
109
- **kwargs):
146
+ self,
147
+ *,
148
+ # Deprecated. Use credentials_strategy instead.
149
+ credentials_provider: Optional[CredentialsStrategy] = None,
150
+ credentials_strategy: Optional[CredentialsStrategy] = None,
151
+ product=None,
152
+ product_version=None,
153
+ clock: Optional[Clock] = None,
154
+ **kwargs,
155
+ ):
110
156
  self._header_factory = None
111
157
  self._inner = {}
112
158
  self._user_agent_other_info = []
113
159
  if credentials_strategy and credentials_provider:
114
- raise ValueError(
115
- "When providing `credentials_strategy` field, `credential_provider` cannot be specified.")
160
+ raise ValueError("When providing `credentials_strategy` field, `credential_provider` cannot be specified.")
116
161
  if credentials_provider:
117
- logger.warning(
118
- "parameter 'credentials_provider' is deprecated. Use 'credentials_strategy' instead.")
162
+ logger.warning("parameter 'credentials_provider' is deprecated. Use 'credentials_strategy' instead.")
119
163
  self._credentials_strategy = next(
120
- s for s in [credentials_strategy, credentials_provider,
121
- DefaultCredentials()] if s is not None)
122
- if 'databricks_environment' in kwargs:
123
- self.databricks_environment = kwargs['databricks_environment']
124
- del kwargs['databricks_environment']
164
+ s
165
+ for s in [
166
+ credentials_strategy,
167
+ credentials_provider,
168
+ DefaultCredentials(),
169
+ ]
170
+ if s is not None
171
+ )
172
+ if "databricks_environment" in kwargs:
173
+ self.databricks_environment = kwargs["databricks_environment"]
174
+ del kwargs["databricks_environment"]
125
175
  self._clock = clock if clock is not None else RealClock()
126
176
  try:
127
177
  self._set_inner_config(kwargs)
@@ -145,15 +195,15 @@ class Config:
145
195
  return message
146
196
 
147
197
  @staticmethod
148
- def parse_dsn(dsn: str) -> 'Config':
198
+ def parse_dsn(dsn: str) -> "Config":
149
199
  uri = urllib.parse.urlparse(dsn)
150
- if uri.scheme != 'databricks':
151
- raise ValueError(f'Expected databricks:// scheme, got {uri.scheme}://')
152
- kwargs = {'host': f'https://{uri.hostname}'}
200
+ if uri.scheme != "databricks":
201
+ raise ValueError(f"Expected databricks:// scheme, got {uri.scheme}://")
202
+ kwargs = {"host": f"https://{uri.hostname}"}
153
203
  if uri.username:
154
- kwargs['username'] = uri.username
204
+ kwargs["username"] = uri.username
155
205
  if uri.password:
156
- kwargs['password'] = uri.password
206
+ kwargs["password"] = uri.password
157
207
  query = dict(urllib.parse.parse_qsl(uri.query))
158
208
  for attr in Config.attributes():
159
209
  if attr.name not in query:
@@ -162,7 +212,7 @@ class Config:
162
212
  return Config(**kwargs)
163
213
 
164
214
  def authenticate(self) -> Dict[str, str]:
165
- """ Returns a list of fresh authentication headers """
215
+ """Returns a list of fresh authentication headers"""
166
216
  return self._header_factory()
167
217
 
168
218
  def as_dict(self) -> dict:
@@ -174,9 +224,9 @@ class Config:
174
224
  env = self.azure_environment.upper()
175
225
  # Compatibility with older versions of the SDK that allowed users to specify AzurePublicCloud or AzureChinaCloud
176
226
  if env.startswith("AZURE"):
177
- env = env[len("AZURE"):]
227
+ env = env[len("AZURE") :]
178
228
  if env.endswith("CLOUD"):
179
- env = env[:-len("CLOUD")]
229
+ env = env[: -len("CLOUD")]
180
230
  return env
181
231
 
182
232
  @property
@@ -241,19 +291,21 @@ class Config:
241
291
 
242
292
  @property
243
293
  def user_agent(self):
244
- """ Returns User-Agent header used by this SDK """
294
+ """Returns User-Agent header used by this SDK"""
245
295
 
246
296
  # global user agent includes SDK version, product name & version, platform info,
247
297
  # and global extra info. Config can have specific extra info associated with it,
248
298
  # such as an override product, auth type, and other user-defined information.
249
- return useragent.to_string(self._product_info,
250
- [("auth", self.auth_type)] + self._user_agent_other_info)
299
+ return useragent.to_string(
300
+ self._product_info,
301
+ [("auth", self.auth_type)] + self._user_agent_other_info,
302
+ )
251
303
 
252
304
  @property
253
305
  def _upstream_user_agent(self) -> str:
254
306
  return " ".join(f"{k}/{v}" for k, v in useragent._get_upstream_user_agent_info())
255
307
 
256
- def with_user_agent_extra(self, key: str, value: str) -> 'Config':
308
+ def with_user_agent_extra(self, key: str, value: str) -> "Config":
257
309
  self._user_agent_other_info.append((key, value))
258
310
  return self
259
311
 
@@ -269,7 +321,7 @@ class Config:
269
321
  return get_workspace_endpoints(self.host)
270
322
 
271
323
  def debug_string(self) -> str:
272
- """ Returns log-friendly representation of configured attributes """
324
+ """Returns log-friendly representation of configured attributes"""
273
325
  buf = []
274
326
  attrs_used = []
275
327
  envs_used = []
@@ -279,13 +331,13 @@ class Config:
279
331
  value = getattr(self, attr.name)
280
332
  if not value:
281
333
  continue
282
- safe = '***' if attr.sensitive else f'{value}'
283
- attrs_used.append(f'{attr.name}={safe}')
334
+ safe = "***" if attr.sensitive else f"{value}"
335
+ attrs_used.append(f"{attr.name}={safe}")
284
336
  if attrs_used:
285
337
  buf.append(f'Config: {", ".join(attrs_used)}')
286
338
  if envs_used:
287
339
  buf.append(f'Env: {", ".join(envs_used)}')
288
- return '. '.join(buf)
340
+ return ". ".join(buf)
289
341
 
290
342
  def to_dict(self) -> Dict[str, any]:
291
343
  return self._inner
@@ -302,16 +354,16 @@ class Config:
302
354
  if (not self.cluster_id) and (not self.warehouse_id):
303
355
  return None
304
356
  if self.cluster_id and self.warehouse_id:
305
- raise ValueError('cannot have both cluster_id and warehouse_id')
357
+ raise ValueError("cannot have both cluster_id and warehouse_id")
306
358
  headers = self.authenticate()
307
- headers['User-Agent'] = f'{self.user_agent} sdk-feature/sql-http-path'
359
+ headers["User-Agent"] = f"{self.user_agent} sdk-feature/sql-http-path"
308
360
  if self.cluster_id:
309
361
  response = requests.get(f"{self.host}/api/2.0/preview/scim/v2/Me", headers=headers)
310
362
  # get workspace ID from the response header
311
- workspace_id = response.headers.get('x-databricks-org-id')
312
- return f'sql/protocolv1/o/{workspace_id}/{self.cluster_id}'
363
+ workspace_id = response.headers.get("x-databricks-org-id")
364
+ return f"sql/protocolv1/o/{workspace_id}/{self.cluster_id}"
313
365
  if self.warehouse_id:
314
- return f'/sql/1.0/warehouses/{self.warehouse_id}'
366
+ return f"/sql/1.0/warehouses/{self.warehouse_id}"
315
367
 
316
368
  @property
317
369
  def clock(self) -> Clock:
@@ -319,17 +371,18 @@ class Config:
319
371
 
320
372
  @classmethod
321
373
  def attributes(cls) -> Iterable[ConfigAttribute]:
322
- """ Returns a list of Databricks SDK configuration metadata """
323
- if hasattr(cls, '_attributes'):
374
+ """Returns a list of Databricks SDK configuration metadata"""
375
+ if hasattr(cls, "_attributes"):
324
376
  return cls._attributes
325
377
  if sys.version_info[1] >= 10:
326
378
  import inspect
379
+
327
380
  anno = inspect.get_annotations(cls)
328
381
  else:
329
382
  # Python 3.7 compatibility: getting type hints require extra hop, as described in
330
383
  # "Accessing The Annotations Dict Of An Object In Python 3.9 And Older" section of
331
384
  # https://docs.python.org/3/howto/annotations.html
332
- anno = cls.__dict__['__annotations__']
385
+ anno = cls.__dict__["__annotations__"]
333
386
  attrs = []
334
387
  for name, v in cls.__dict__.items():
335
388
  if type(v) != ConfigAttribute:
@@ -351,26 +404,25 @@ class Config:
351
404
  If the tenant ID is already set, this method does nothing."""
352
405
  if not self.is_azure or self.azure_tenant_id is not None or self.host is None:
353
406
  return
354
- login_url = f'{self.host}/aad/auth'
355
- logger.debug(f'Loading tenant ID from {login_url}')
407
+ login_url = f"{self.host}/aad/auth"
408
+ logger.debug(f"Loading tenant ID from {login_url}")
356
409
  resp = requests.get(login_url, allow_redirects=False)
357
410
  if resp.status_code // 100 != 3:
358
- logger.debug(
359
- f'Failed to get tenant ID from {login_url}: expected status code 3xx, got {resp.status_code}')
411
+ logger.debug(f"Failed to get tenant ID from {login_url}: expected status code 3xx, got {resp.status_code}")
360
412
  return
361
- entra_id_endpoint = resp.headers.get('Location')
413
+ entra_id_endpoint = resp.headers.get("Location")
362
414
  if entra_id_endpoint is None:
363
- logger.debug(f'No Location header in response from {login_url}')
415
+ logger.debug(f"No Location header in response from {login_url}")
364
416
  return
365
417
  # The Location header has the following form: https://login.microsoftonline.com/<tenant-id>/oauth2/authorize?...
366
418
  # The domain may change depending on the Azure cloud (e.g. login.microsoftonline.us for US Government cloud).
367
419
  url = urllib.parse.urlparse(entra_id_endpoint)
368
- path_segments = url.path.split('/')
420
+ path_segments = url.path.split("/")
369
421
  if len(path_segments) < 2:
370
- logger.debug(f'Invalid path in Location header: {url.path}')
422
+ logger.debug(f"Invalid path in Location header: {url.path}")
371
423
  return
372
424
  self.azure_tenant_id = path_segments[1]
373
- logger.debug(f'Loaded tenant ID: {self.azure_tenant_id}')
425
+ logger.debug(f"Loaded tenant ID: {self.azure_tenant_id}")
374
426
 
375
427
  def _set_inner_config(self, keyword_args: Dict[str, any]):
376
428
  for attr in self.attributes():
@@ -393,11 +445,10 @@ class Config:
393
445
  self.__setattr__(attr.name, value)
394
446
  found = True
395
447
  if found:
396
- logger.debug('Loaded from environment')
448
+ logger.debug("Loaded from environment")
397
449
 
398
450
  def _known_file_config_loader(self):
399
- if not self.profile and (self.is_any_auth_configured or self.host
400
- or self.azure_workspace_resource_id):
451
+ if not self.profile and (self.is_any_auth_configured or self.host or self.azure_workspace_resource_id):
401
452
  # skip loading configuration file if there's any auth configured
402
453
  # directly as part of the Config() constructor.
403
454
  return
@@ -417,15 +468,15 @@ class Config:
417
468
  # from Unified Auth test suite at the moment. Hence, the private variable access.
418
469
  # See: https://docs.python.org/3/library/configparser.html#mapping-protocol-access
419
470
  if not has_explicit_profile and not ini_file.defaults():
420
- logger.debug(f'{config_path} has no DEFAULT profile configured')
471
+ logger.debug(f"{config_path} has no DEFAULT profile configured")
421
472
  return
422
473
  if not has_explicit_profile:
423
474
  profile = "DEFAULT"
424
475
  profiles = ini_file._sections
425
476
  if ini_file.defaults():
426
- profiles['DEFAULT'] = ini_file.defaults()
477
+ profiles["DEFAULT"] = ini_file.defaults()
427
478
  if profile not in profiles:
428
- raise ValueError(f'resolve: {config_path} has no {profile} profile configured')
479
+ raise ValueError(f"resolve: {config_path} has no {profile} profile configured")
429
480
  raw_config = profiles[profile]
430
481
  logger.info(f'loading {profile} profile from {config_file}: {", ".join(raw_config.keys())}')
431
482
  for k, v in raw_config.items():
@@ -448,26 +499,29 @@ class Config:
448
499
  # client has auth preference set
449
500
  return
450
501
  names = " and ".join(sorted(auths_used))
451
- raise ValueError(f'validate: more than one authorization method configured: {names}')
502
+ raise ValueError(f"validate: more than one authorization method configured: {names}")
452
503
 
453
504
  def init_auth(self):
454
505
  try:
455
506
  self._header_factory = self._credentials_strategy(self)
456
507
  self.auth_type = self._credentials_strategy.auth_type()
457
508
  if not self._header_factory:
458
- raise ValueError('not configured')
509
+ raise ValueError("not configured")
459
510
  except ValueError as e:
460
- raise ValueError(f'{self._credentials_strategy.auth_type()} auth: {e}') from e
511
+ raise ValueError(f"{self._credentials_strategy.auth_type()} auth: {e}") from e
461
512
 
462
513
  def _init_product(self, product, product_version):
463
514
  if product is not None or product_version is not None:
464
515
  default_product, default_version = useragent.product()
465
- self._product_info = (product or default_product, product_version or default_version)
516
+ self._product_info = (
517
+ product or default_product,
518
+ product_version or default_version,
519
+ )
466
520
  else:
467
521
  self._product_info = None
468
522
 
469
523
  def __repr__(self):
470
- return f'<{self.debug_string()}>'
524
+ return f"<{self.debug_string()}>"
471
525
 
472
526
  def copy(self):
473
527
  """Creates a copy of the config object.
@@ -480,6 +534,5 @@ class Config:
480
534
  return cpy
481
535
 
482
536
  def deep_copy(self):
483
- """Creates a deep copy of the config object.
484
- """
537
+ """Creates a deep copy of the config object."""
485
538
  return copy.deepcopy(self)
databricks/sdk/core.py CHANGED
@@ -9,9 +9,9 @@ from .credentials_provider import *
9
9
  from .errors import DatabricksError, _ErrorCustomizer
10
10
  from .oauth import retrieve_token
11
11
 
12
- __all__ = ['Config', 'DatabricksError']
12
+ __all__ = ["Config", "DatabricksError"]
13
13
 
14
- logger = logging.getLogger('databricks.sdk')
14
+ logger = logging.getLogger("databricks.sdk")
15
15
 
16
16
  URL_ENCODED_CONTENT_TYPE = "application/x-www-form-urlencoded"
17
17
  JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"
@@ -22,16 +22,18 @@ class ApiClient:
22
22
 
23
23
  def __init__(self, cfg: Config):
24
24
  self._cfg = cfg
25
- self._api_client = _BaseClient(debug_truncate_bytes=cfg.debug_truncate_bytes,
26
- retry_timeout_seconds=cfg.retry_timeout_seconds,
27
- user_agent_base=cfg.user_agent,
28
- header_factory=cfg.authenticate,
29
- max_connection_pools=cfg.max_connection_pools,
30
- max_connections_per_pool=cfg.max_connections_per_pool,
31
- pool_block=True,
32
- http_timeout_seconds=cfg.http_timeout_seconds,
33
- extra_error_customizers=[_AddDebugErrorCustomizer(cfg)],
34
- clock=cfg.clock)
25
+ self._api_client = _BaseClient(
26
+ debug_truncate_bytes=cfg.debug_truncate_bytes,
27
+ retry_timeout_seconds=cfg.retry_timeout_seconds,
28
+ user_agent_base=cfg.user_agent,
29
+ header_factory=cfg.authenticate,
30
+ max_connection_pools=cfg.max_connection_pools,
31
+ max_connections_per_pool=cfg.max_connections_per_pool,
32
+ pool_block=True,
33
+ http_timeout_seconds=cfg.http_timeout_seconds,
34
+ extra_error_customizers=[_AddDebugErrorCustomizer(cfg)],
35
+ clock=cfg.clock,
36
+ )
35
37
 
36
38
  @property
37
39
  def account_id(self) -> str:
@@ -46,44 +48,52 @@ class ApiClient:
46
48
  self._cfg.authenticate()
47
49
  original_token = self._cfg.oauth_token()
48
50
  headers = {"Content-Type": URL_ENCODED_CONTENT_TYPE}
49
- params = urlencode({
50
- "grant_type": JWT_BEARER_GRANT_TYPE,
51
- "authorization_details": auth_details,
52
- "assertion": original_token.access_token
53
- })
54
- return retrieve_token(client_id=self._cfg.client_id,
55
- client_secret=self._cfg.client_secret,
56
- token_url=self._cfg.host + OIDC_TOKEN_PATH,
57
- params=params,
58
- headers=headers)
51
+ params = urlencode(
52
+ {
53
+ "grant_type": JWT_BEARER_GRANT_TYPE,
54
+ "authorization_details": auth_details,
55
+ "assertion": original_token.access_token,
56
+ }
57
+ )
58
+ return retrieve_token(
59
+ client_id=self._cfg.client_id,
60
+ client_secret=self._cfg.client_secret,
61
+ token_url=self._cfg.host + OIDC_TOKEN_PATH,
62
+ params=params,
63
+ headers=headers,
64
+ )
59
65
 
60
- def do(self,
61
- method: str,
62
- path: str = None,
63
- url: str = None,
64
- query: dict = None,
65
- headers: dict = None,
66
- body: dict = None,
67
- raw: bool = False,
68
- files=None,
69
- data=None,
70
- auth: Callable[[requests.PreparedRequest], requests.PreparedRequest] = None,
71
- response_headers: List[str] = None) -> Union[dict, list, BinaryIO]:
66
+ def do(
67
+ self,
68
+ method: str,
69
+ path: Optional[str] = None,
70
+ url: Optional[str] = None,
71
+ query: Optional[dict] = None,
72
+ headers: Optional[dict] = None,
73
+ body: Optional[dict] = None,
74
+ raw: bool = False,
75
+ files=None,
76
+ data=None,
77
+ auth: Optional[Callable[[requests.PreparedRequest], requests.PreparedRequest]] = None,
78
+ response_headers: Optional[List[str]] = None,
79
+ ) -> Union[dict, list, BinaryIO]:
72
80
  if url is None:
73
81
  # Remove extra `/` from path for Files API
74
82
  # Once we've fixed the OpenAPI spec, we can remove this
75
- path = re.sub('^/api/2.0/fs/files//', '/api/2.0/fs/files/', path)
83
+ path = re.sub("^/api/2.0/fs/files//", "/api/2.0/fs/files/", path)
76
84
  url = f"{self._cfg.host}{path}"
77
- return self._api_client.do(method=method,
78
- url=url,
79
- query=query,
80
- headers=headers,
81
- body=body,
82
- raw=raw,
83
- files=files,
84
- data=data,
85
- auth=auth,
86
- response_headers=response_headers)
85
+ return self._api_client.do(
86
+ method=method,
87
+ url=url,
88
+ query=query,
89
+ headers=headers,
90
+ body=body,
91
+ raw=raw,
92
+ files=files,
93
+ data=data,
94
+ auth=auth,
95
+ response_headers=response_headers,
96
+ )
87
97
 
88
98
 
89
99
  class _AddDebugErrorCustomizer(_ErrorCustomizer):
@@ -95,5 +105,5 @@ class _AddDebugErrorCustomizer(_ErrorCustomizer):
95
105
 
96
106
  def customize_error(self, response: requests.Response, kwargs: dict):
97
107
  if response.status_code in (401, 403):
98
- message = kwargs.get('message', 'request failed')
99
- kwargs['message'] = self._cfg.wrap_debug_info(message)
108
+ message = kwargs.get("message", "request failed")
109
+ kwargs["message"] = self._cfg.wrap_debug_info(message)