Merge branch 'main' of https://github.com/databricks/databricks-sdk-py into update-main

Author: CaymanWilliams

.codegen.json

@@ -1,5 +1,6 @@
 {
   "formatter": "yapf -pri $FILENAMES && autoflake -i $FILENAMES && isort $FILENAMES",
+  "changelog_config": ".codegen/changelog_config.yml",
   "template_libraries": [
     ".codegen/lib.tmpl"
   ],

.codegen/__init__.py.tmpl

@@ -5,18 +5,20 @@ from databricks.sdk.credentials_provider import CredentialsStrategy
 from databricks.sdk.mixins.files import DbfsExt
 from databricks.sdk.mixins.compute import ClustersExt
 from databricks.sdk.mixins.workspace import WorkspaceExt
-{{- range .Services}} {{if not .IsDataPlane}}
-from databricks.sdk.service.{{.Package.Name}} import {{.PascalName}}API{{end}}{{end}}
+from databricks.sdk.mixins.open_ai_client import ServingEndpointsExt
+{{- range .Services}}
+from databricks.sdk.service.{{.Package.Name}} import {{.PascalName}}API{{end}}
 from databricks.sdk.service.provisioning import Workspace
 from databricks.sdk import azure
+from typing import Optional
 
 {{$args := list "host" "account_id" "username" "password" "client_id" "client_secret"
   "token" "profile" "config_file" "azure_workspace_resource_id" "azure_client_secret"
   "azure_client_id" "azure_tenant_id" "azure_environment" "auth_type" "cluster_id"
   "google_credentials" "google_service_account" }}
 
 {{- define "api" -}}
-  {{- $mixins := dict "ClustersAPI" "ClustersExt" "DbfsAPI" "DbfsExt" "WorkspaceAPI" "WorkspaceExt" -}}
+  {{- $mixins := dict "ClustersAPI" "ClustersExt" "DbfsAPI" "DbfsExt" "WorkspaceAPI" "WorkspaceExt" "ServingEndpointsAPI" "ServingEndpointsExt" -}}
   {{- $genApi := concat .PascalName "API" -}}
   {{- getOrDefault $mixins $genApi $genApi -}}
 {{- end -}}
@@ -41,14 +43,14 @@ class WorkspaceClient:
     """
     The WorkspaceClient is a client for the workspace-level Databricks REST API.
     """
-    def __init__(self, *{{range $args}}, {{.}}: str = None{{end}},
-                 debug_truncate_bytes: int = None,
-                 debug_headers: bool = None,
+    def __init__(self, *{{range $args}}, {{.}}: Optional[str] = None{{end}},
+                 debug_truncate_bytes: Optional[int] = None,
+                 debug_headers: Optional[bool] = None,
                  product="unknown",
                  product_version="0.0.0",
-                 credentials_strategy: CredentialsStrategy = None,
-                 credentials_provider: CredentialsStrategy = None,
-                 config: client.Config = None):
+                 credentials_strategy: Optional[CredentialsStrategy] = None,
+                 credentials_provider: Optional[CredentialsStrategy] = None,
+                 config: Optional[client.Config] = None):
         if not config:
           config = client.Config({{range $args}}{{.}}={{.}}, {{end}}
             credentials_strategy=credentials_strategy,
@@ -61,8 +63,20 @@ class WorkspaceClient:
         self._dbutils = _make_dbutils(self._config)
         self._api_client = client.ApiClient(self._config)
 
-        {{- range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane)}}
-        self._{{.SnakeName}} = {{template "api" .}}(self._api_client){{end -}}{{end}}
+        {{- range .Services}}{{if and (not .IsAccounts) (not .HasParent) .HasDataPlaneAPI (not .IsDataPlane)}}
+        {{.SnakeName}} = {{template "api" .}}(self._api_client){{end -}}{{end}}
+
+        {{- range .Services}}
+        {{- if and (not .IsAccounts) (not .HasParent)}}
+        {{- if .IsDataPlane}}
+        self._{{.SnakeName}} = {{template "api" .}}(self._api_client, {{.ControlPlaneService.SnakeName}})
+        {{- else if .HasDataPlaneAPI}}
+        self._{{.SnakeName}} = {{.SnakeName}}
+        {{- else}}
+        self._{{.SnakeName}} = {{template "api" .}}(self._api_client)
+        {{- end -}}
+        {{- end -}}
+        {{end}}
 
     @property
     def config(self) -> client.Config:
@@ -76,7 +90,7 @@ class WorkspaceClient:
     def dbutils(self) -> dbutils.RemoteDbUtils:
         return self._dbutils
 
-    {{- range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane)}}
+    {{- range .Services}}{{if and (not .IsAccounts) (not .HasParent)}}
     @property
     def {{.SnakeName}}(self) -> {{template "api" .}}:
         {{if .Description}}"""{{.Summary}}"""{{end}}
@@ -98,14 +112,14 @@ class AccountClient:
     The AccountClient is a client for the account-level Databricks REST API.
     """
 
-    def __init__(self, *{{range $args}}, {{.}}: str = None{{end}},
-                 debug_truncate_bytes: int = None,
-                 debug_headers: bool = None,
+    def __init__(self, *{{range $args}}, {{.}}: Optional[str] = None{{end}},
+                 debug_truncate_bytes: Optional[int] = None,
+                 debug_headers: Optional[bool] = None,
                  product="unknown",
                  product_version="0.0.0",
-                 credentials_strategy: CredentialsStrategy = None,
-                 credentials_provider: CredentialsStrategy = None,
-                 config: client.Config = None):
+                 credentials_strategy: Optional[CredentialsStrategy] = None,
+                 credentials_provider: Optional[CredentialsStrategy] = None,
+                 config: Optional[client.Config] = None):
         if not config:
           config = client.Config({{range $args}}{{.}}={{.}}, {{end}}
             credentials_strategy=credentials_strategy,
@@ -117,8 +131,20 @@ class AccountClient:
         self._config = config.copy()
         self._api_client = client.ApiClient(self._config)
 
-        {{- range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane)}}
-        self._{{(.TrimPrefix "account").SnakeName}} = {{template "api" .}}(self._api_client){{end -}}{{end}}
+        {{- range .Services}}{{if and .IsAccounts (not .HasParent) .HasDataPlaneAPI (not .IsDataPlane)}}
+        {{(.TrimPrefix "account").SnakeName}} = {{template "api" .}}(self._api_client){{end -}}{{end}}
+
+        {{- range .Services}}
+        {{- if and .IsAccounts (not .HasParent)}}
+        {{- if .IsDataPlane}}
+        self._{{(.TrimPrefix "account").SnakeName}} = {{template "api" .}}(self._api_client, {{.ControlPlaneService.SnakeName}})
+        {{- else if .HasDataPlaneAPI}}
+        self._{{(.TrimPrefix "account").SnakeName}} = {{(.TrimPrefix "account").SnakeName}}
+        {{- else}}
+        self._{{(.TrimPrefix "account").SnakeName}} = {{template "api" .}}(self._api_client)
+        {{- end -}}
+        {{- end -}}
+        {{end}}
 
     @property
     def config(self) -> client.Config:
@@ -128,7 +154,7 @@ class AccountClient:
     def api_client(self) -> client.ApiClient:
         return self._api_client
 
-    {{- range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane)}}
+    {{- range .Services}}{{if and .IsAccounts (not .HasParent)}}
     @property
     def {{(.TrimPrefix "account").SnakeName}}(self) -> {{template "api" .}}:{{if .Description}}
         """{{.Summary}}"""{{end}}

.codegen/_openapi_sha

@@ -1 +1 @@
-7437dabb9dadee402c1fc060df4c1ce8cc5369f0
+cf9c61453990df0f9453670f2fe68e1b128647a2

.codegen/changelog.md.tmpl

@@ -1,13 +1,17 @@
 # Version changelog
 
 ## {{.Version}}
+{{- range .GroupChanges}}
 
-{{range .Changes -}}
+### {{.Type.Message}}
+{{range .Changes}}
  * {{.}}.
-{{end}}{{- if .ApiChanges}}
-API Changes:
-{{range .ApiChanges}}
- * {{.Action}} {{template "what" .}}{{if .Extra}} {{.Extra}}{{with .Other}} {{template "what" .}}{{end}}{{end}}.
+{{- end}}
+{{end}}
+{{if .ApiChanges}}
+### API Changes:
+{{range .ApiChanges.GroupDiff}}
+ * {{.Action}} {{template "group-what" .}}{{if .Extra}} {{.Extra}}{{with .Other}} {{template "other-what" .}}{{end}}{{end}}.
 {{- end}}
 
 OpenAPI SHA: {{.Sha}}, Date: {{.Changed}}
@@ -20,7 +24,35 @@ Dependency updates:
 
 ## {{.PrevVersion}}
 
-{{- define "what" -}}
+{{- define "group-what" -}}
+{{if gt (len .Changes) 1 -}} {{template "single-what" .Changes.First}}{{end -}}
+{{range .Changes.Middle -}}, {{template "single-what" .}}{{end -}}
+{{if gt (len .Changes) 1}} and {{end}}{{template "single-what" .Changes.Last}}{{template "suffix-what" .}}
+{{- end -}}
+
+{{- define "single-what" -}}
+    {{if eq .X "package" -}}
+        `databricks.sdk.service.{{.Package.Name}}`
+    {{- else if eq .X "service" -}}
+        {{template "service" .Service}}
+    {{- else if eq .X "method" -}}
+        `{{.Method.SnakeName}}()`
+    {{- else if eq .X "entity" -}}
+        {{template "entity" .Entity}}
+    {{- else if eq .X "field" -}}
+        `{{.Field.SnakeName}}`
+    {{- end}}
+{{- end -}}
+
+{{- define "suffix-what" -}}
+    {{if eq .Type "package" }} package{{if gt (len .Changes) 1}}s{{end}}
+    {{- else if eq .Type "method" }} method{{if gt (len .Changes) 1}}s{{end}} for {{template "service" .Parent.Service}}
+    {{- else if eq .Type "entity" }} dataclass{{if gt (len .Changes) 1}}es{{end}}
+    {{- else if eq .Type "field" }} field{{if gt (len .Changes) 1}}s{{end}} for {{template "entity" .Parent.Entity}}
+    {{- end}}
+{{- end -}}
+
+{{- define "other-what" -}}
     {{if eq .X "package" -}}
         `databricks.sdk.service.{{.Package.Name}}` package
     {{- else if eq .X "service" -}}

.codegen/changelog_config.yml

@@ -0,0 +1,14 @@
+change_types:
+  - message: New Features and Improvements
+    tag: "[Feature]"
+  - message: Bug Fixes
+    tag: "[Fix]"
+  - message: Documentation
+    tag: "[Doc]"
+  - message: Internal Changes
+    tag: "[Internal]"
+  # Does not appear in the Changelog. Only for PR validation.
+  - message: Release
+    tag: "[Release]"
+  # Default for messages without a tag
+  - message: Other Changes

.codegen/error_overrides.py.tmpl

@@ -11,9 +11,9 @@ _ALL_OVERRIDES = [
         debug_name="{{.Name}}",
         path_regex=re.compile(r'{{.PathRegex}}'),
         verb="{{.Verb}}",
-        status_code_matcher=re.compile(r'{{.StatusCodeMatcher}}'),
-        error_code_matcher=re.compile(r'{{.ErrorCodeMatcher}}'),
-        message_matcher=re.compile(r'{{.MessageMatcher}}'),
+        status_code_matcher=re.compile(r'{{replaceAll "'" "\\'" .StatusCodeMatcher}}'),
+        error_code_matcher=re.compile(r'{{replaceAll "'" "\\'" .ErrorCodeMatcher}}'),
+        message_matcher=re.compile(r'{{replaceAll "'" "\\'" .MessageMatcher}}'),
         custom_error={{.OverrideErrorCode.PascalName}},
     ),
 {{- end }}

.codegen/service.py.tmpl

@@ -8,8 +8,12 @@ from typing import Dict, List, Any, Iterator, Type, Callable, Optional, BinaryIO
 import time
 import random
 import logging
+import requests
+
+from ..data_plane import DataPlaneService
 from ..errors import OperationTimeout, OperationFailed
 from ._internal import _enum, _from_dict, _repeated_dict, _repeated_enum, Wait, _escape_multi_segment_path_parameter
+from ..oauth import Token
 
 _LOG = logging.getLogger('databricks.sdk')
 
@@ -100,12 +104,16 @@ class {{.PascalName}}{{if eq "List" .PascalName}}Request{{end}}:{{if .Descriptio
 	{{- end -}}
 {{- end -}}
 
-{{range .Services}} {{if not .IsDataPlane}}
+{{range .Services}}
 class {{.PascalName}}API:{{if .Description}}
     """{{.Comment "    " 110}}"""
     {{end}}
-    def __init__(self, api_client):
+    def __init__(self, api_client{{if .IsDataPlane}}, control_plane{{end}}):
         self._api = api_client
+        {{if .IsDataPlane -}}
+        self._control_plane = control_plane
+        self._data_plane_service = DataPlaneService()
+        {{end -}}
         {{range .Subservices}}
         self._{{.SnakeName}} = {{.PascalName}}API(self._api){{end}}
 
@@ -183,6 +191,9 @@ class {{.PascalName}}API:{{if .Description}}
         {{if .Request -}}
         {{template "method-serialize" .}}
         {{- end}}
+        {{- if .Service.IsDataPlane}}
+        {{template "data-plane" .}}
+        {{- end}}
         {{template "method-headers" . }}
         {{if .Response.HasHeaderField -}}
         {{template "method-response-headers" . }}
@@ -195,7 +206,27 @@ class {{.PascalName}}API:{{if .Description}}
         return self.{{template "safe-snake-name" .}}({{range $i, $x := .Request.Fields}}{{if $i}}, {{end}}{{template "safe-snake-name" .}}={{template "safe-snake-name" .}}{{end}}).result(timeout=timeout)
     {{end}}
     {{end -}}
-{{end}}
+{{- end}}
+
+{{define "data-plane" -}}
+        def info_getter():
+            response = self._control_plane.{{.Service.DataPlaneInfoMethod.SnakeName}}(
+                {{- range .Service.DataPlaneInfoMethod.Request.Fields }}
+                {{.SnakeName}} = {{.SnakeName}},
+                {{- end}}
+                )
+            if response.{{(index .DataPlaneInfoFields 0).SnakeName}} is None:
+            	raise Exception("Resource does not support direct Data Plane access")
+            return response{{range .DataPlaneInfoFields}}.{{.SnakeName}}{{end}}
+
+        get_params = [{{- range .Service.DataPlaneInfoMethod.Request.Fields }}{{.SnakeName}},{{end}}]
+        data_plane_details = self._data_plane_service.get_data_plane_details('{{.SnakeName}}', get_params, info_getter, self._api.get_oauth_token)
+        token = data_plane_details.token
+
+        def auth(r: requests.PreparedRequest) -> requests.PreparedRequest:
+            authorization = f"{token.token_type} {token.access_token}"
+            r.headers["Authorization"] = authorization
+            return r
 {{- end}}
 
 {{define "method-parameters" -}}
@@ -264,7 +295,7 @@ class {{.PascalName}}API:{{if .Description}}
         {{if .NeedsOffsetDedupe -}}
         # deduplicate items that may have been added during iteration
         seen = set()
-        {{- end}}{{if and .Pagination.Offset (not (eq .Path "/api/2.0/clusters/events")) }}
+        {{- end}}{{if and .Pagination.Offset (not (eq .Path "/api/2.1/clusters/events")) }}
         query['{{.Pagination.Offset.Name}}'] =
         {{- if eq .Pagination.Increment 1 -}}
           1
@@ -290,7 +321,7 @@ class {{.PascalName}}API:{{if .Description}}
           if '{{.Pagination.Token.Bind.Name}}' not in json or not json['{{.Pagination.Token.Bind.Name}}']:
             return
           {{if or (eq "GET" .Verb) (eq "HEAD" .Verb)}}query{{else}}body{{end}}['{{.Pagination.Token.PollField.Name}}'] = json['{{.Pagination.Token.Bind.Name}}']
-          {{- else if eq .Path "/api/2.0/clusters/events" -}}
+          {{- else if eq .Path "/api/2.1/clusters/events" -}}
           if 'next_page' not in json or not json['next_page']:
             return
           body = json['next_page']
@@ -319,25 +350,32 @@ class {{.PascalName}}API:{{if .Description}}
           {{- else if .Response.MapValue -}}
             return res
           {{- else -}}
-            return {{.Response.PascalName}}.from_dict(res)
+            return {{template "type" .Response}}.from_dict(res)
           {{- end}}
         {{- end}}
 {{- end}}
 
 {{define "method-do" -}}
-self._api.do('{{.Verb}}',
-  {{ template "path" . }}
-  {{if .Request}}
-      {{- if .Request.HasQueryField}}, query=query{{end}}
-      {{- if .Request.MapValue}}, body=contents
-      {{- else if .Request.HasJsonField}}, body=body{{end}}
-  {{end}}
-  , headers=headers
-  {{if .Response.HasHeaderField -}}
-  , response_headers=response_headers
-  {{- end}}
-  {{- if and .IsRequestByteStream .RequestBodyField }}, data={{template "safe-snake-name" .RequestBodyField}}{{ end }}
-  {{- if .IsResponseByteStream }}, raw=True{{ end }})
+    self._api.do('{{.Verb}}',
+    {{- if .Service.IsDataPlane -}}
+    url=data_plane_details.endpoint_url
+    {{- else -}}
+    {{ template "path" . }}
+    {{- end -}}
+    {{if .Request}}
+        {{- if .Request.HasQueryField}}, query=query{{end}}
+        {{- if .Request.MapValue}}, body=contents
+        {{- else if .Request.HasJsonField}}, body=body{{end}}
+    {{end}}
+    , headers=headers
+    {{if .Response.HasHeaderField -}}
+    , response_headers=response_headers
+    {{- end}}
+    {{- if and .IsRequestByteStream .RequestBodyField }}, data={{template "safe-snake-name" .RequestBodyField}}{{ end }}
+    {{- if .Service.IsDataPlane -}}
+    ,auth=auth
+    {{- end -}}
+    {{- if .IsResponseByteStream }}, raw=True{{ end }})
 {{- end}}
 
 {{- define "path" -}}