Skip to content

Commit

Permalink
Update OpenAPI Spec & Merge Files API (#304)
Browse files Browse the repository at this point in the history
## Changes
Support for Files API.

## Tests
- [x] Integration tests for FilesAPI updated to create their own
schema/volume.

- [ ] `make test` run locally
- [ ] `make fmt` applied
- [ ] relevant integration tests applied
  • Loading branch information
mgyucht authored Aug 28, 2023
1 parent 0f11a96 commit e05c8c0
Show file tree
Hide file tree
Showing 25 changed files with 1,430 additions and 578 deletions.
3 changes: 1 addition & 2 deletions .codegen/__init__.py.tmpl
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import databricks.sdk.core as client
import databricks.sdk.dbutils as dbutils

from databricks.sdk.mixins.files import DbfsExt, FilesMixin
from databricks.sdk.mixins.files import DbfsExt
from databricks.sdk.mixins.compute import ClustersExt
from databricks.sdk.mixins.workspace import WorkspaceExt
{{- range .Services}}
Expand Down Expand Up @@ -51,7 +51,6 @@ class WorkspaceClient:
self.config = config.copy()
self.dbutils = _make_dbutils(self.config)
self.api_client = client.ApiClient(self.config)
self.files = FilesMixin(self.api_client)
{{- range .Services}}{{if not .IsAccounts}}
self.{{.SnakeName}} = {{template "api" .}}(self.api_client){{end -}}{{end}}

Expand Down
2 changes: 1 addition & 1 deletion .codegen/_openapi_sha
Original file line number Diff line number Diff line change
@@ -1 +1 @@
beff621d7b3e1d59244e2e34fc53a496f310e130
be3d4a799362f0d5ddbfeb0a0acdfb91f8736a3b
9 changes: 7 additions & 2 deletions .codegen/service.py.tmpl
Original file line number Diff line number Diff line change
Expand Up @@ -200,8 +200,13 @@ 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}}
if {{.SnakeName}} is not None: {{if .IsQuery}}query{{else}}body{{end}}['{{.Name}}'] = {{template "method-param-bind" .}}{{end}}{{end}}
{{- range .Request.Fields}}{{ if not .IsPath }}
{{- if .IsQuery }}
if {{.SnakeName}} is not None: query['{{.Name}}'] = {{template "method-param-bind" .}}{{end}}
{{- if .IsJson }}
if {{.SnakeName}} is not None: body['{{.Name}}'] = {{template "method-param-bind" .}}{{end}}
{{- end}}
{{- end}}
{{- end}}
{{- end}}

Expand Down
15 changes: 8 additions & 7 deletions databricks/sdk/__init__.py

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

5 changes: 3 additions & 2 deletions databricks/sdk/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -988,11 +988,12 @@ def do(self,
raw: bool = False,
files=None,
data=None) -> Union[dict, BinaryIO]:
# Remove extra `/` from path for Files API
# Once we've fixed the OpenAPI spec, we can remove this
path = re.sub('^/api/2.0/fs/files//', '/api/2.0/fs/files/', path)
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),
Expand Down
17 changes: 1 addition & 16 deletions databricks/sdk/mixins/files.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from types import TracebackType
from typing import TYPE_CHECKING, AnyStr, BinaryIO, Iterable, Iterator, Type

from databricks.sdk.core import ApiClient, DatabricksError
from databricks.sdk.core import DatabricksError

from ..service import files

Expand Down Expand Up @@ -397,18 +397,3 @@ def move_(self, src: str, dst: str, *, recursive=False, overwrite=False):
# do cross-fs moving
self.copy(src, dst, recursive=recursive, overwrite=overwrite)
source.delete(recursive=recursive)


class FilesMixin:

def __init__(self, api_client: ApiClient):
self._api = api_client

def upload(self, path: str, src: BinaryIO):
self._api.do('PUT', f'/api/2.0/fs/files{path}', data=src)

def download(self, path: str) -> BinaryIO:
return self._api.do('GET', f'/api/2.0/fs/files{path}', raw=True)

def delete(self, path: str):
self._api.do('DELETE', f'/api/2.0/fs/files{path}')
Loading

0 comments on commit e05c8c0

Please sign in to comment.