From dd6bd527f29742a54468e2183fba1146dea02a64 Mon Sep 17 00:00:00 2001
From: Hector Castejon Diaz <hector.castejon@databricks.com>
Date: Fri, 16 Feb 2024 14:56:13 +0100
Subject: [PATCH] Only use from_dict when using headers

---
 .codegen/service.py.tmpl          | 6 +++++-
 databricks/sdk/service/billing.py | 4 +---
 databricks/sdk/service/files.py   | 4 +---
 databricks/sdk/service/sql.py     | 1 -
 tests/integration/test_files.py   | 2 +-
 5 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/.codegen/service.py.tmpl b/.codegen/service.py.tmpl
index 4584f2e47..fe43be702 100644
--- a/.codegen/service.py.tmpl
+++ b/.codegen/service.py.tmpl
@@ -328,9 +328,13 @@ class {{.Name}}API:{{if .Description}}
 {{- end}}
 
 {{define "deserialize-stream" -}}
-deserialized = {{.Response.PascalName}}.from_dict(res)
+    {{if or .Response.HasJsonField .Response.HasHeaderField -}}
+        deserialized = {{.Response.PascalName}}.from_dict(res)
         {{.Response.PascalName}}.{{.ResponseBodyField.CamelName}}=content
         return deserialized
+    {{- else -}}
+        return {{.Response.PascalName}}({{.ResponseBodyField.CamelName}}=res)
+    {{end}}
 {{- end}}
 
 {{define "method-do" -}}
diff --git a/databricks/sdk/service/billing.py b/databricks/sdk/service/billing.py
index 61f4f48e4..662b89ccd 100755
--- a/databricks/sdk/service/billing.py
+++ b/databricks/sdk/service/billing.py
@@ -711,9 +711,7 @@ def download(self,
                                     headers=headers,
                                     response_headers=response_headers,
                                     raw=True)
-        deserialized = DownloadResponse.from_dict(res)
-        DownloadResponse.contents = content
-        return deserialized
+        return DownloadResponse(contents=res)
 
 
 class BudgetsAPI:
diff --git a/databricks/sdk/service/files.py b/databricks/sdk/service/files.py
index 26880079a..02fff643d 100755
--- a/databricks/sdk/service/files.py
+++ b/databricks/sdk/service/files.py
@@ -729,9 +729,7 @@ def download(self, file_path: str) -> DownloadResponse:
                                     headers=headers,
                                     response_headers=response_headers,
                                     raw=True)
-        deserialized = DownloadResponse.from_dict(res)
-        DownloadResponse.contents = content
-        return deserialized
+        return DownloadResponse(contents=res)
 
     def get_directory_metadata(self, directory_path: str):
         """Get directory metadata.
diff --git a/databricks/sdk/service/sql.py b/databricks/sdk/service/sql.py
index 90b5196e9..5cd289905 100755
--- a/databricks/sdk/service/sql.py
+++ b/databricks/sdk/service/sql.py
@@ -346,7 +346,6 @@ def from_dict(cls, d: Dict[str, any]) -> ChannelInfo:
 
 
 class ChannelName(Enum):
-    """Name of the channel"""
 
     CHANNEL_NAME_CURRENT = 'CHANNEL_NAME_CURRENT'
     CHANNEL_NAME_CUSTOM = 'CHANNEL_NAME_CUSTOM'
diff --git a/tests/integration/test_files.py b/tests/integration/test_files.py
index 31fdc7d2e..85a0203d1 100644
--- a/tests/integration/test_files.py
+++ b/tests/integration/test_files.py
@@ -222,7 +222,7 @@ def test_files_api_upload_download(ucws, random):
             target_file = f'/Volumes/main/{schema}/{volume}/filesit-{random()}.txt'
             w.files.upload(target_file, f)
             m = w.files.get_metadata(target_file)
-            assert m.content_type == 'text/plain'
+            assert m.content_type == 'application/octet-stream'
             with w.files.download(target_file).contents as f:
                 assert f.read() == b"some text data"