@@ -5,18 +5,20 @@ from databricks.sdk.credentials_provider import CredentialsStrategy
55from databricks.sdk.mixins.files import DbfsExt
66from databricks.sdk.mixins.compute import ClustersExt
77from databricks.sdk.mixins.workspace import WorkspaceExt
8- {{- range .Services}} {{if not .IsDataPlane}}
9- from databricks.sdk.service.{{.Package.Name}} import {{.PascalName}}API{{end}}{{end}}
8+ from databricks.sdk.mixins.open_ai_client import ServingEndpointsExt
9+ {{- range .Services}}
10+ from databricks.sdk.service.{{.Package.Name}} import {{.PascalName}}API{{end}}
1011from databricks.sdk.service.provisioning import Workspace
1112from databricks.sdk import azure
13+ from typing import Optional
1214
1315{{$args := list "host" "account_id" "username" "password" "client_id" "client_secret"
1416 "token" "profile" "config_file" "azure_workspace_resource_id" "azure_client_secret"
1517 "azure_client_id" "azure_tenant_id" "azure_environment" "auth_type" "cluster_id"
1618 "google_credentials" "google_service_account" }}
1719
1820{{- define "api" -}}
19- {{- $mixins := dict "ClustersAPI" "ClustersExt" "DbfsAPI" "DbfsExt" "WorkspaceAPI" "WorkspaceExt" -}}
21+ {{- $mixins := dict "ClustersAPI" "ClustersExt" "DbfsAPI" "DbfsExt" "WorkspaceAPI" "WorkspaceExt" "ServingEndpointsAPI" "ServingEndpointsExt" -}}
2022 {{- $genApi := concat .PascalName "API" -}}
2123 {{- getOrDefault $mixins $genApi $genApi -}}
2224{{- end -}}
@@ -41,14 +43,14 @@ class WorkspaceClient:
4143 """
4244 The WorkspaceClient is a client for the workspace-level Databricks REST API.
4345 """
44- def __init__(self, *{{range $args}}, {{.}}: str = None{{end}},
45- debug_truncate_bytes: int = None,
46- debug_headers: bool = None,
46+ def __init__(self, *{{range $args}}, {{.}}: Optional[ str] = None{{end}},
47+ debug_truncate_bytes: Optional[ int] = None,
48+ debug_headers: Optional[ bool] = None,
4749 product="unknown",
4850 product_version="0.0.0",
49- credentials_strategy: CredentialsStrategy = None,
50- credentials_provider: CredentialsStrategy = None,
51- config: client.Config = None):
51+ credentials_strategy: Optional[ CredentialsStrategy] = None,
52+ credentials_provider: Optional[ CredentialsStrategy] = None,
53+ config: Optional[ client.Config] = None):
5254 if not config:
5355 config = client.Config({{range $args}}{{.}}={{.}}, {{end}}
5456 credentials_strategy=credentials_strategy,
@@ -61,8 +63,20 @@ class WorkspaceClient:
6163 self._dbutils = _make_dbutils(self._config)
6264 self._api_client = client.ApiClient(self._config)
6365
64- {{- range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane)}}
65- self._{{.SnakeName}} = {{template "api" .}}(self._api_client){{end -}}{{end}}
66+ {{- range .Services}}{{if and (not .IsAccounts) (not .HasParent) .HasDataPlaneAPI (not .IsDataPlane)}}
67+ {{.SnakeName}} = {{template "api" .}}(self._api_client){{end -}}{{end}}
68+
69+ {{- range .Services}}
70+ {{- if and (not .IsAccounts) (not .HasParent)}}
71+ {{- if .IsDataPlane}}
72+ self._{{.SnakeName}} = {{template "api" .}}(self._api_client, {{.ControlPlaneService.SnakeName}})
73+ {{- else if .HasDataPlaneAPI}}
74+ self._{{.SnakeName}} = {{.SnakeName}}
75+ {{- else}}
76+ self._{{.SnakeName}} = {{template "api" .}}(self._api_client)
77+ {{- end -}}
78+ {{- end -}}
79+ {{end}}
6680
6781 @property
6882 def config(self) -> client.Config:
@@ -76,7 +90,7 @@ class WorkspaceClient:
7690 def dbutils(self) -> dbutils.RemoteDbUtils:
7791 return self._dbutils
7892
79- {{- range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane) }}
93+ {{- range .Services}}{{if and (not .IsAccounts) (not .HasParent)}}
8094 @property
8195 def {{.SnakeName}}(self) -> {{template "api" .}}:
8296 {{if .Description}}"""{{.Summary}}"""{{end}}
@@ -98,14 +112,14 @@ class AccountClient:
98112 The AccountClient is a client for the account-level Databricks REST API.
99113 """
100114
101- def __init__(self, *{{range $args}}, {{.}}: str = None{{end}},
102- debug_truncate_bytes: int = None,
103- debug_headers: bool = None,
115+ def __init__(self, *{{range $args}}, {{.}}: Optional[ str] = None{{end}},
116+ debug_truncate_bytes: Optional[ int] = None,
117+ debug_headers: Optional[ bool] = None,
104118 product="unknown",
105119 product_version="0.0.0",
106- credentials_strategy: CredentialsStrategy = None,
107- credentials_provider: CredentialsStrategy = None,
108- config: client.Config = None):
120+ credentials_strategy: Optional[ CredentialsStrategy] = None,
121+ credentials_provider: Optional[ CredentialsStrategy] = None,
122+ config: Optional[ client.Config] = None):
109123 if not config:
110124 config = client.Config({{range $args}}{{.}}={{.}}, {{end}}
111125 credentials_strategy=credentials_strategy,
@@ -117,8 +131,20 @@ class AccountClient:
117131 self._config = config.copy()
118132 self._api_client = client.ApiClient(self._config)
119133
120- {{- range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane)}}
121- self._{{(.TrimPrefix "account").SnakeName}} = {{template "api" .}}(self._api_client){{end -}}{{end}}
134+ {{- range .Services}}{{if and .IsAccounts (not .HasParent) .HasDataPlaneAPI (not .IsDataPlane)}}
135+ {{(.TrimPrefix "account").SnakeName}} = {{template "api" .}}(self._api_client){{end -}}{{end}}
136+
137+ {{- range .Services}}
138+ {{- if and .IsAccounts (not .HasParent)}}
139+ {{- if .IsDataPlane}}
140+ self._{{(.TrimPrefix "account").SnakeName}} = {{template "api" .}}(self._api_client, {{.ControlPlaneService.SnakeName}})
141+ {{- else if .HasDataPlaneAPI}}
142+ self._{{(.TrimPrefix "account").SnakeName}} = {{(.TrimPrefix "account").SnakeName}}
143+ {{- else}}
144+ self._{{(.TrimPrefix "account").SnakeName}} = {{template "api" .}}(self._api_client)
145+ {{- end -}}
146+ {{- end -}}
147+ {{end}}
122148
123149 @property
124150 def config(self) -> client.Config:
@@ -128,7 +154,7 @@ class AccountClient:
128154 def api_client(self) -> client.ApiClient:
129155 return self._api_client
130156
131- {{- range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane) }}
157+ {{- range .Services}}{{if and .IsAccounts (not .HasParent)}}
132158 @property
133159 def {{(.TrimPrefix "account").SnakeName}}(self) -> {{template "api" .}}:{{if .Description}}
134160 """{{.Summary}}"""{{end}}
0 commit comments