Skip to content

Commit

Permalink
Support BinaryIO for streaming request and response bodies (#303)
Browse files Browse the repository at this point in the history
## Changes
This PR adds support for non-application/json requests and responses.
This PR is divided into two parts:

**Support setting headers per request**. The `do` method of the
ApiClient class is expanded to accept one new parameter, `headers:
dict`, and all callers are expected to pass a map of headers to be
included in the request. As the allowed content types for requests and
responses are known from the OpenAPI specification, callers must
construct the request header map in code generation and pass it to the
ApiClient. The default "Content-Type: application/json" header is
removed, as each request should specify its own Content-Type and Accept
headers. This is done in
#302.

**Add support for streaming request and response bodies to support
non-application/json requests/responses**. Today, serialization of
requests is done in the ApiClient. This implies that ApiClient needs to
be able to serialize a request purely based on the parameters provided
to it via headers, the request body, etc. In this proposal, the
serialization is specified by the caller depending on whether the `body`
or `data` field of `do()` is used: `data` should be used for
already-serialized data, such as BinaryIO or file-like objects, and
`body` should be used for JSON objects.

One important caveat about streaming responses is that callers are
required to close the streams when they are done using them. To do this,
we may need to expose the raw Response object so that users can use
`with w.service.method() as response:` to automatically close the
response when they are done reading.

## Tests
<!-- 
How is this tested? Please see the checklist below and also describe any
other relevant tests
-->

- [ ] `make test` run locally
- [ ] `make fmt` applied
- [ ] relevant integration tests applied
  • Loading branch information
mgyucht authored Aug 23, 2023
1 parent 27356dd commit 821b02a
Show file tree
Hide file tree
Showing 17 changed files with 799 additions and 804 deletions.
18 changes: 12 additions & 6 deletions .codegen/service.py.tmpl
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from dataclasses import dataclass
from datetime import timedelta
from enum import Enum
from typing import Dict, List, Any, Iterator, Type, Callable, Optional
from typing import Dict, List, Any, Iterator, Type, Callable, Optional, BinaryIO
import time
import random
import logging
Expand Down Expand Up @@ -78,6 +78,7 @@ class {{.PascalName}}{{if eq "List" .PascalName}}Request{{end}}:{{if .Descriptio
{{- else if .IsInt64}}int
{{- else if .IsFloat64}}float
{{- else if .IsInt}}int
{{- else if .IsByteStream}}BinaryIO
{{- else if .ArrayValue }}List[{{template "type-nq" .ArrayValue}}]
{{- else if .MapValue }}Dict[str,{{template "type-nq" .MapValue}}]
{{- else if .IsExternal }}{{.Package.Name}}.{{.PascalName}}
Expand All @@ -95,6 +96,7 @@ class {{.PascalName}}{{if eq "List" .PascalName}}Request{{end}}:{{if .Descriptio
{{- else if .IsInt64}}int
{{- else if .IsFloat64}}float
{{- else if .IsInt}}int
{{- else if .IsByteStream}}BinaryIO
{{- else if .ArrayValue }}List[{{template "type-doc" .ArrayValue}}]
{{- else if .MapValue }}Dict[str,{{template "type-doc" .MapValue}}]
{{- else if .IsExternal }}:class:`{{.PascalName}}`
Expand Down Expand Up @@ -274,12 +276,14 @@ class {{.Name}}API:{{if .Description}}

{{define "method-call-default" -}}
{{if .Response -}}
json = {{end}}{{template "method-do" .}}
res = {{end}}{{template "method-do" .}}
{{if .Response -}}
{{if .Response.ArrayValue -}}
return [{{.Response.ArrayValue.PascalName}}.from_dict(v) for v in json]
{{- if .IsResponseByteStream -}}
return {{.Response.PascalName}}({{.ResponseBodyField.CamelName}}=res)
{{- else if .Response.ArrayValue -}}
return [{{.Response.ArrayValue.PascalName}}.from_dict(v) for v in res]
{{- else -}}
return {{.Response.PascalName}}.from_dict(json)
return {{.Response.PascalName}}.from_dict(res)
{{- end}}
{{- end}}
{{- end}}
Expand All @@ -293,7 +297,9 @@ self._api.do('{{.Verb}}',
{{- if .Request.HasQueryField}}, query=query{{end}}
{{- if .Request.HasJsonField}}, body=body{{end}}
{{end}}
, headers=headers)
, headers=headers
{{- if and .IsRequestByteStream .RequestBodyField }}, data={{.RequestBodyField.SnakeName}}{{ end }}
{{- if .IsResponseByteStream }}, raw=True{{ end }})
{{- end}}

{{define "method-return-type" -}}
Expand Down
4 changes: 3 additions & 1 deletion databricks/sdk/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -991,14 +991,16 @@ def do(self,
if headers is None:
headers = {}
headers['User-Agent'] = self._user_agent_base
# Replace // with / in path (for Files API where users specify filenames beginning with /)
path = re.sub(r'//', '/', path)
response = self._session.request(method,
f"{self._cfg.host}{path}",
params=self._fix_query_string(query),
json=body,
headers=headers,
files=files,
data=data,
stream=True if raw else False)
stream=raw)
try:
self._record_request_log(response, raw=raw or data is not None or files is not None)
if not response.ok:
Expand Down
32 changes: 16 additions & 16 deletions databricks/sdk/service/billing.py

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

Loading

0 comments on commit 821b02a

Please sign in to comment.