Adding multi-ds support to sample code

Mariano Teixeira · Mariano Teixeira · commit ad7962276aff · 2017-10-09T14:25:11.000-07:00
diff --git a/tableauserverclient/server/endpoint/datasources_endpoint.py b/tableauserverclient/server/endpoint/datasources_endpoint.py
@@ -116,7 +116,7 @@ def update(self, datasource_item):
 
     # Publish datasource
     @api(version="2.0")
-    def publish(self, datasource_item, file_path, mode, connection_credentials=None):
+    def publish(self, datasource_item, file_path, mode, connections=None):
         if not os.path.isfile(file_path):
             error = "File path does not lead to an existing file."
             raise IOError(error)
@@ -145,15 +145,15 @@ def publish(self, datasource_item, file_path, mode, connection_credentials=None)
             upload_session_id = Fileuploads.upload_chunks(self.parent_srv, file_path)
             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()
             xml_request, content_type = RequestFactory.Datasource.publish_req(datasource_item,
                                                                               filename,
                                                                               file_contents,
-                                                                              connection_credentials)
+                                                                              connections)
         server_response = self.post_request(url, xml_request, content_type)
         new_datasource = DatasourceItem.from_response(server_response.content)[0]
         logger.info('Published {0} (ID: {1})'.format(filename, new_datasource.id))
diff --git a/tableauserverclient/server/request_factory.py b/tableauserverclient/server/request_factory.py
@@ -278,22 +278,30 @@ def add_req(self, user_item):
 
 
 class WorkbookRequest(object):
-    def _generate_xml(self, workbook_item, connection_credentials=None):
+    def _generate_xml(self, workbook_item, connections=None):
         xml_request = ET.Element('tsRequest')
         workbook_element = ET.SubElement(xml_request, 'workbook')
         workbook_element.attrib['name'] = workbook_item.name
         if workbook_item.show_tabs:
             workbook_element.attrib['showTabs'] = str(workbook_item.show_tabs).lower()
         project_element = ET.SubElement(workbook_element, 'project')
         project_element.attrib['id'] = workbook_item.project_id
-        if connection_credentials:
-            credentials_element = ET.SubElement(workbook_element, 'connectionCredentials')
-            credentials_element.attrib['name'] = connection_credentials.name
-            credentials_element.attrib['password'] = connection_credentials.password
-            credentials_element.attrib['embed'] = 'true' if connection_credentials.embed else 'false'
-
-            if connection_credentials.oauth:
-                credentials_element.attrib['oAuth'] = 'true'
+        if connections is not None:
+            connections_element = ET.SubElement(workbook_element, 'connections')
+            for connection in connections:
+                if connection:
+                    connection_element = ET.SubElement(connections_element, 'connection')
+                    connection_element.attrib['serverAddress'] = connection.server_address
+                    if connection.server_port:
+                        connection_element.attrib['serverPort'] = connection.server_port
+                    if connection.connection_credentials:
+                        connection_credentials = connection.connection_credentials
+                        credentials_element = ET.SubElement(connection_element, 'connectionCredentials')
+                        credentials_element.attrib['name'] = connection_credentials.name
+                        credentials_element.attrib['password'] = connection_credentials.password
+                        credentials_element.attrib['embed'] = 'true' if connection_credentials.embed else 'false'
+                        if connection_credentials.oauth:
+                            credentials_element.attrib['oAuth'] = 'true'
         return ET.tostring(xml_request)
 
     def update_req(self, workbook_item):
@@ -309,15 +317,15 @@ def update_req(self, workbook_item):
             owner_element.attrib['id'] = workbook_item.owner_id
         return ET.tostring(xml_request)
 
-    def publish_req(self, workbook_item, filename, file_contents, connection_credentials=None):
-        xml_request = self._generate_xml(workbook_item, connection_credentials)
+    def publish_req(self, workbook_item, filename, file_contents, connections=None):
+        xml_request = self._generate_xml(workbook_item, connections)
 
         parts = {'request_payload': ('', xml_request, 'text/xml'),
                  'tableau_workbook': (filename, file_contents, 'application/octet-stream')}
         return _add_multipart(parts)
 
-    def publish_req_chunked(self, workbook_item, connection_credentials=None):
-        xml_request = self._generate_xml(workbook_item, connection_credentials)
+    def publish_req_chunked(self, workbook_item, connections=None):
+        xml_request = self._generate_xml(workbook_item, connections)
 
         parts = {'request_payload': ('', xml_request, 'text/xml')}
         return _add_multipart(parts)
