File objects for publish datasource

tableauserverclient/filesys_helpers.py

@@ -30,13 +30,33 @@ def get_file_object_size(file):
     return file_size
 
 
-def file_is_compressed(file):
-    # Determine if file is a zip file or not
-    # This reference lists magic file signatures: https://www.garykessler.net/library/file_sigs.html
-
-    zip_file_signature = b'PK\x03\x04'
+def get_file_type(file):
+    # Tableau workbooks (twb) and data sources (tds) are both stored as xml files.
+    # Packaged workbooks (twbx) and data sources (tdsx) are zip files
+    # containing original files accompanied with supporting local files.
 
-    is_zip_file = file.read(len(zip_file_signature)) == zip_file_signature
+    # This reference lists magic file signatures: https://www.garykessler.net/library/file_sigs.html
+    MAGIC_BYTES = {
+        'zip': bytes.fromhex("504b0304"),
+        'tde': bytes.fromhex("20020162"),
+        'xml': bytes.fromhex("3c3f786d6c20"),
+        'hyper': bytes.fromhex("487970657208000001000000")
+    }
+
+    # Peek first bytes of a file
+    first_bytes = file.read(32)
+
+    file_type = None
+    for ft, signature in MAGIC_BYTES.items():
+        if first_bytes.startswith(signature):
+            file_type = ft
+            break
+
+    # Return pointer back to start
     file.seek(0)
 
-    return is_zip_file
+    if file_type is None:
+        error = "Unknown file type!"
+        raise ValueError(error)
+
+    return file_type

tableauserverclient/server/endpoint/datasources_endpoint.py

@@ -5,7 +5,7 @@
 from .resource_tagger import _ResourceTagger
 from .. import RequestFactory, DatasourceItem, PaginationItem, ConnectionItem
 from ..query import QuerySet
-from ...filesys_helpers import to_filename, make_download_path
+from ...filesys_helpers import to_filename, make_download_path, get_file_type, get_file_object_size
 from ...models.job_item import JobItem
 
 import os
@@ -173,22 +173,45 @@ def delete_extract(self, datasource_item):
     @api(version="2.0")
     @parameter_added_in(connections="2.8")
     @parameter_added_in(as_job='3.0')
-    def publish(self, datasource_item, file_path, mode, connection_credentials=None, connections=None, as_job=False):
-        if not os.path.isfile(file_path):
-            error = "File path does not lead to an existing file."
-            raise IOError(error)
-        if not mode or not hasattr(self.parent_srv.PublishMode, mode):
-            error = 'Invalid mode defined.'
-            raise ValueError(error)
+    def publish(self, datasource_item, file, mode, connection_credentials=None, connections=None, as_job=False):
+
+        try:
+
+            if not os.path.isfile(file):
+                error = "File path does not lead to an existing file."
+                raise IOError(error)
+
+            filename = os.path.basename(file)
+            file_extension = os.path.splitext(filename)[1][1:]
+            file_size = os.path.getsize(file)
+
+            # If name is not defined, grab the name from the file to publish
+            if not datasource_item.name:
+                datasource_item.name = os.path.splitext(filename)[0]
+            if file_extension not in ALLOWED_FILE_EXTENSIONS:
+                error = "Only {} files can be published as datasources.".format(', '.join(ALLOWED_FILE_EXTENSIONS))
+                raise ValueError(error)
 
-        filename = os.path.basename(file_path)
-        file_extension = os.path.splitext(filename)[1][1:]
+        except TypeError:
 
-        # If name is not defined, grab the name from the file to publish
-        if not datasource_item.name:
-            datasource_item.name = os.path.splitext(filename)[0]
-        if file_extension not in ALLOWED_FILE_EXTENSIONS:
-            error = "Only {} files can be published as datasources.".format(', '.join(ALLOWED_FILE_EXTENSIONS))
+            if not datasource_item.name:
+                error = "Datasource item must have a name when passing a file object"
+                raise ValueError(error)
+
+            file_type = get_file_type(file)
+            if file_type == 'zip':
+                file_extension = 'tdsx'
+            elif file_type == 'xml':
+                file_extension = 'tds'
+            else:
+                error = "Unsupported file type {}".format(file_type)
+                raise ValueError(error)
+
+            filename = "{}.{}".format(datasource_item.name, file_extension)
+            file_size = get_file_object_size(file)
+
+        if not mode or not hasattr(self.parent_srv.PublishMode, mode):
+            error = 'Invalid mode defined.'
             raise ValueError(error)
 
         # Construct the url with the defined mode
@@ -200,17 +223,22 @@ def publish(self, datasource_item, file_path, mode, connection_credentials=None,
             url += '&{0}=true'.format('asJob')
 
         # Determine if chunking is required (64MB is the limit for single upload method)
-        if os.path.getsize(file_path) >= FILESIZE_LIMIT:
+        if file_size >= FILESIZE_LIMIT:
             logger.info('Publishing {0} to server with chunking method (datasource over 64MB)'.format(filename))
-            upload_session_id = Fileuploads.upload_chunks(self.parent_srv, file_path)
+            upload_session_id = Fileuploads.upload_chunks(self.parent_srv, file)
             url = "{0}&uploadSessionId={1}".format(url, upload_session_id)
             xml_request, content_type = RequestFactory.Datasource.publish_req_chunked(datasource_item,
                                                                                       connection_credentials,
                                                                                       connections)
         else:
             logger.info('Publishing {0} to server'.format(filename))
-            with open(file_path, 'rb') as f:
-                file_contents = f.read()
+
+            try:
+                with open(file, 'rb') as f:
+                    file_contents = f.read()
+            except TypeError:
+                file_contents = file.read()
+
             xml_request, content_type = RequestFactory.Datasource.publish_req(datasource_item,
                                                                               filename,
                                                                               file_contents,

tableauserverclient/server/endpoint/workbooks_endpoint.py

@@ -5,7 +5,7 @@
 from .resource_tagger import _ResourceTagger
 from .. import RequestFactory, WorkbookItem, ConnectionItem, ViewItem, PaginationItem
 from ...models.job_item import JobItem
-from ...filesys_helpers import to_filename, make_download_path, file_is_compressed, get_file_object_size
+from ...filesys_helpers import to_filename, make_download_path, get_file_type, get_file_object_size
 
 import os
 import logging
@@ -284,7 +284,16 @@ def publish(
         except TypeError:
             # Expect file to be a file object
             file_size = get_file_object_size(file)
-            file_extension = 'twbx' if file_is_compressed(file) else 'twb'
+
+            file_type = get_file_type(file)
+
+            if file_type == 'zip':
+                file_extension = 'twbx'
+            elif file_type == 'xml':
+                file_extension = 'twb'
+            else:
+                error = 'Unsupported file type {}!'.format(file_type)
+                raise ValueError(error)
 
             if not workbook_item.name:
                 error = "Workbook item must have a name when passing a file object"