Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support HEAD operation and response Headers #547

Merged
merged 18 commits into from
Feb 19, 2024
Merged
2 changes: 1 addition & 1 deletion .codegen/__init__.py.tmpl
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,7 @@ class AccountClient:
:param workspace: The workspace to construct a client for.
:return: A ``WorkspaceClient`` for the given workspace.
"""
config = self._config.copy()
config = self._config.deep_copy()
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug fix. The generated code was already updated here

9af2630#diff-3847ef3d682f5902da2252a8d7b3a0745d15a4b81ef42cd73f73be158622c5d6R812

config.host = config.environment.deployment_url(workspace.deployment_name)
config.azure_workspace_resource_id = azure.get_azure_resource_id(workspace)
config.account_id = None
Expand Down
30 changes: 20 additions & 10 deletions .codegen/service.py.tmpl
Original file line number Diff line number Diff line change
Expand Up @@ -21,26 +21,26 @@ from databricks.sdk.service import {{.Package.Name}}{{end}}
{{if .Fields -}}{{if not .IsRequest}}@dataclass
class {{.PascalName}}{{if eq "List" .PascalName}}Request{{end}}:{{if .Description}}
"""{{.Comment " " 100}}"""
{{end}}{{- range .RequiredFields | alphanumOnly}}
{{end}}{{- range .RequiredFields}}
hectorcast-db marked this conversation as resolved.
Show resolved Hide resolved
{{template "safe-snake-name" .}}: {{template "type" .Entity}}{{if .Description}}
"""{{.Comment " " 100 | trimSuffix "\""}}"""{{end}}
{{end}}
{{- range .NonRequiredFields | alphanumOnly}}
{{- range .NonRequiredFields}}
{{template "safe-snake-name" .}}: Optional[{{template "type" .Entity}}] = None{{if .Description}}
"""{{.Comment " " 100 | trimSuffix "\""}}"""{{end}}
{{end}}
{{if .HasJsonField -}}
{{if or .HasJsonField .HasHeaderField .HasByteStreamField -}}
def as_dict(self) -> dict:
"""Serializes the {{.PascalName}}{{if eq "List" .PascalName}}Request{{end}} into a dictionary suitable for use as a JSON request body."""
body = {}
{{range .Fields | alphanumOnly}}if self.{{template "safe-snake-name" .}}{{with .Entity.IsPrimitive}} is not None{{end}}: body['{{.Name}}'] = {{template "as_request_type" .}}
{{range .Fields}}if self.{{template "safe-snake-name" .}}{{with .Entity.IsPrimitive}} is not None{{end}}: body['{{.Name}}'] = {{template "as_request_type" .}}
{{end -}}
return body

@classmethod
def from_dict(cls, d: Dict[str, any]) -> {{.PascalName}}{{if eq "List" .PascalName}}Request{{end}}:
"""Deserializes the {{.PascalName}}{{if eq "List" .PascalName}}Request{{end}} from a dictionary."""
return cls({{range $i, $f := .Fields | alphanumOnly}}{{if $i}}, {{end}}{{template "safe-snake-name" $f}}={{template "from_dict_type" $f}}{{end}})
return cls({{range $i, $f := .Fields}}{{if $i}}, {{end}}{{template "safe-snake-name" $f}}={{template "from_dict_type" $f}}{{end}})
{{end}}
{{end}}
{{else if .ArrayValue}}type {{.PascalName}} []{{template "type" .ArrayValue}}
Expand Down Expand Up @@ -187,6 +187,9 @@ class {{.Name}}API:{{if .Description}}
{{template "method-serialize" .}}
{{- end}}
{{template "method-headers" . }}
{{if and .Response .Response.HasHeaderField -}}
{{template "method-response-headers" . }}
{{- end}}
{{template "method-call" .}}

{{if and .Wait (and (not .IsCrudRead) (not (eq .SnakeName "get_run"))) }}
Expand All @@ -211,7 +214,7 @@ class {{.Name}}API:{{if .Description}}
{{if or .Request.HasJsonField .Request.HasQueryField -}}
{{if .Request.HasJsonField}}body = {}{{end}}{{if .Request.HasQueryField}}
query = {}{{end}}
{{- range .Request.Fields}}{{ if not .IsPath }}
{{- range .Request.Fields}}{{ if and (not .IsPath) (not .IsHeader) }}
{{- if .IsQuery }}
if {{template "safe-snake-name" .}} is not None: query['{{.Name}}'] = {{template "method-param-bind" .}}{{end}}
{{- if .IsJson }}
Expand All @@ -227,6 +230,12 @@ class {{.Name}}API:{{if .Description}}
}
{{- end }}

{{ define "method-response-headers" -}}
hectorcast-db marked this conversation as resolved.
Show resolved Hide resolved
response_headers = [
{{- range $h := .ResponseHeaders}}'{{ $h.Name }}',{{ end -}}
]
{{- end }}

{{- define "method-param-bind" -}}
{{- if not .Entity }}None # ERROR: No Type
{{- else if .Entity.ArrayValue }}[
Expand Down Expand Up @@ -282,7 +291,7 @@ class {{.Name}}API:{{if .Description}}
{{ if .Pagination.Token -}}
if '{{.Pagination.Token.Bind.Name}}' not in json or not json['{{.Pagination.Token.Bind.Name}}']:
return
{{if eq "GET" .Verb}}query{{else}}body{{end}}['{{.Pagination.Token.PollField.Name}}'] = json['{{.Pagination.Token.Bind.Name}}']
{{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" -}}
if 'next_page' not in json or not json['next_page']:
return
Expand All @@ -307,9 +316,7 @@ class {{.Name}}API:{{if .Description}}
{{if .Response -}}
res = {{end}}{{template "method-do" .}}
{{if .Response -}}
{{- if .IsResponseByteStream -}}
return {{.Response.PascalName}}({{.ResponseBodyField.CamelName}}=res)
{{- else if .Response.ArrayValue -}}
{{- if .Response.ArrayValue -}}
return [{{.Response.ArrayValue.PascalName}}.from_dict(v) for v in res]
{{- else if .Response.MapValue -}}
return res
Expand All @@ -330,6 +337,9 @@ self._api.do('{{.Verb}}',
{{- else if .Request.HasJsonField}}, body=body{{end}}
{{end}}
, headers=headers
{{if and .Response .Response.HasHeaderField -}}
, response_headers=response_headers
{{- end}}
{{- if and .IsRequestByteStream .RequestBodyField }}, data={{template "safe-snake-name" .RequestBodyField}}{{ end }}
{{- if .IsResponseByteStream }}, raw=True{{ end }})
{{- end}}
Expand Down
1 change: 1 addition & 0 deletions .gitattributes
Original file line number Diff line number Diff line change
Expand Up @@ -285,6 +285,7 @@ examples/workspace/list_workspace_integration.py linguist-generated=true
examples/workspace_assignment/list_workspace_assignment_on_aws.py linguist-generated=true
examples/workspace_assignment/update_workspace_assignment_on_aws.py linguist-generated=true
examples/workspace_bindings/get_catalog_workspace_bindings.py linguist-generated=true
examples/workspace_bindings/update_catalog_workspace_bindings.py linguist-generated=true
examples/workspace_conf/get_status_repos.py linguist-generated=true
examples/workspaces/create_workspaces.py linguist-generated=true
examples/workspaces/get_workspaces.py linguist-generated=true
Expand Down
6 changes: 4 additions & 2 deletions databricks/sdk/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,11 +135,13 @@ def do(self,
raw=raw,
files=files,
data=data)
if raw:
return StreamingResponse(response)

resp = dict()
for header in response_headers if response_headers else []:
resp[header] = response.headers.get(Casing.to_header_case(header))
if raw:
resp["contents"] = StreamingResponse(response)
return resp
if not len(response.content):
return resp

Expand Down
3 changes: 2 additions & 1 deletion databricks/sdk/mixins/workspace.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,4 +102,5 @@ def download(self, path: str, *, format: Optional[ExportFormat] = None) -> Binar
"""
query = {'path': path, 'direct_download': 'true'}
if format: query['format'] = format.value
return self._api.do('GET', '/api/2.0/workspace/export', query=query, raw=True)
response = self._api.do('GET', '/api/2.0/workspace/export', query=query, raw=True)
return response["contents"]
23 changes: 22 additions & 1 deletion databricks/sdk/service/billing.py

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading
Loading