feat: batch update schedule state

Author: jorwoods

tableauserverclient/models/schedule_item.py

@@ -1,6 +1,6 @@
 import xml.etree.ElementTree as ET
 from datetime import datetime
-from typing import Optional, Union
+from typing import Optional, Union, TYPE_CHECKING
 
 from defusedxml.ElementTree import fromstring
 
@@ -16,6 +16,10 @@
     property_is_enum,
 )
 
+if TYPE_CHECKING:
+    from requests import Response
+
+
 Interval = Union[HourlyInterval, DailyInterval, WeeklyInterval, MonthlyInterval]
 
 
@@ -407,3 +411,8 @@ def _read_warnings(parsed_response, ns):
         for warning_xml in all_warning_xml:
             warnings.append(warning_xml.get("message", None))
         return warnings
+
+
+def parse_batch_schedule_state(response: "Response", ns) -> list[str]:
+    xml = fromstring(response.content)
+    return [text for tag in xml.findall(".//t:scheduleLuid", namespaces=ns) if (text := tag.text)]

tableauserverclient/server/endpoint/schedules_endpoint.py

@@ -1,13 +1,15 @@
+from collections.abc import Iterable
 import copy
 import logging
 import warnings
 from collections import namedtuple
-from typing import TYPE_CHECKING, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union, overload
 
 from .endpoint import Endpoint, api, parameter_added_in
 from .exceptions import MissingRequiredFieldError
 from tableauserverclient.server import RequestFactory
 from tableauserverclient.models import PaginationItem, ScheduleItem, TaskItem, ExtractItem
+from tableauserverclient.models.schedule_item import parse_batch_schedule_state
 
 from tableauserverclient.helpers.logging import logger
 
@@ -279,3 +281,48 @@ def get_extract_refresh_tasks(
         extract_items = ExtractItem.from_response(server_response.content, self.parent_srv.namespace)
 
         return extract_items, pagination_item
+
+    @overload
+    def batch_update_state(
+        self,
+        schedules: Iterable[ScheduleItem | str],
+        state: Literal["active", "suspended"],
+        update_all: Literal[False] = False,
+    ) -> list[str]: ...
+
+    @overload
+    def batch_update_state(
+        self, schedules: Any, state: Literal["active", "suspended"], update_all: Literal[True]
+    ) -> list[str]: ...
+
+    @api(version="3.27")
+    def batch_update_state(self, schedules, state, update_all=False) -> list[str]:
+        """
+        Batch update the status of one or more scheudles. If update_all is set,
+        all schedules on the Tableau Server are affected.
+
+        Parameters
+        ----------
+        schedules: Iterable[ScheudleItem | str] | Any
+            The schedules to be updated. If update_all=True, this is ignored.
+
+        state: Literal["active", "suspended"]
+            The state of the schedules, whether active or suspended.
+
+        update_all: bool
+            Whether or not to apply the status to all schedules.
+
+        Returns
+        -------
+        List[str]
+            The IDs of the affected schedules.
+        """
+        params = {"state": state}
+        if update_all:
+            params["updateAll"] = "true"
+            payload = RequestFactory.Empty.empty_req()
+        else:
+            payload = RequestFactory.Schedule.batch_update_state(schedules)
+
+        response = self.put_request(self.baseurl, payload, parameters={"params": params})
+        return parse_batch_schedule_state(response, self.parent_srv.namespace)

tableauserverclient/server/request_factory.py

@@ -643,6 +643,16 @@ def add_datasource_req(self, id_: Optional[str], task_type: str = TaskItem.Type.
     def add_flow_req(self, id_: Optional[str], task_type: str = TaskItem.Type.RunFlow) -> bytes:
         return self._add_to_req(id_, "flow", task_type)
 
+    @_tsrequest_wrapped
+    def batch_update_state(self, xml: ET.Element, schedules: Iterable[ScheduleItem | str]) -> None:
+        luids = ET.SubElement(xml, "scheduleLuids")
+        for schedule in schedules:
+            luid = getattr(schedule, "id", schedule)
+            if not isinstance(luid, str):
+                continue
+            luid_tag = ET.SubElement(luids, "scheduleLuid")
+            luid_tag.text = luid
+
 
 class SiteRequest:
     def update_req(self, site_item: "SiteItem", parent_srv: Optional["Server"] = None):

test/assets/schedule_batch_update_state.xml

@@ -0,0 +1,7 @@
+<tsResponse xmlns="http://tableau.com/api" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://tableau.com/api http://tableau.com/api/ts-api-2.3.xsd">
+  <scheduleLuids>
+    <scheduleLuid>593d2ebf-0d18-4deb-9d21-b113a4902583</scheduleLuid>
+    <scheduleLuid>cecbb71e-def0-4030-8068-5ae50f51db1c</scheduleLuid>
+    <scheduleLuid>f39a6e7d-405e-4c07-8c18-95845f9da80e</scheduleLuid>
+  </scheduleLuids>
+</tsResponse>

test/test_schedule.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 from datetime import time
 
+from defusedxml.ElementTree import fromstring
 import pytest
 import requests_mock
 
@@ -26,7 +27,7 @@
 ADD_DATASOURCE_TO_SCHEDULE = TEST_ASSET_DIR / "schedule_add_datasource.xml"
 ADD_FLOW_TO_SCHEDULE = TEST_ASSET_DIR / "schedule_add_flow.xml"
 GET_EXTRACT_TASKS_XML = TEST_ASSET_DIR / "schedule_get_extract_refresh_tasks.xml"
-BATCH_UPDATE_STATE = TEST_ASSET_DIR / "schedules_batch_update_state.xml"
+BATCH_UPDATE_STATE = TEST_ASSET_DIR / "schedule_batch_update_state.xml"
 
 WORKBOOK_GET_BY_ID_XML = TEST_ASSET_DIR / "workbook_get_by_id.xml"
 DATASOURCE_GET_BY_ID_XML = TEST_ASSET_DIR / "datasource_get_by_id.xml"
@@ -421,3 +422,58 @@ def test_get_extract_refresh_tasks(server: TSC.Server) -> None:
         assert isinstance(extracts[0], list)
         assert 2 == len(extracts[0])
         assert "task1" == extracts[0][0].id
+
+
+def test_batch_update_state_items(server: TSC.Server) -> None:
+    server.version = "3.27"
+    hourly_interval = TSC.HourlyInterval(start_time=time(2, 30), end_time=time(23, 0), interval_value=2)
+    args = ("hourly", 50, TSC.ScheduleItem.Type.Extract, TSC.ScheduleItem.ExecutionOrder.Parallel, hourly_interval)
+    new_schedules = [TSC.ScheduleItem(*args), TSC.ScheduleItem(*args), TSC.ScheduleItem(*args)]
+    new_schedules[0]._id = "593d2ebf-0d18-4deb-9d21-b113a4902583"
+    new_schedules[1]._id = "cecbb71e-def0-4030-8068-5ae50f51db1c"
+    new_schedules[2]._id = "f39a6e7d-405e-4c07-8c18-95845f9da80e"
+
+    state = "active"
+    with requests_mock.mock() as m:
+        m.put(f"{server.schedules.baseurl}?state={state}", text=BATCH_UPDATE_STATE.read_text())
+        resp = server.schedules.batch_update_state(new_schedules, state)
+
+    assert len(resp) == 3
+    for sch, r in zip(new_schedules, resp):
+        assert sch.id == r
+
+
+def test_batch_update_state_str(server: TSC.Server) -> None:
+    server.version = "3.27"
+    new_schedules = [
+        "593d2ebf-0d18-4deb-9d21-b113a4902583",
+        "cecbb71e-def0-4030-8068-5ae50f51db1c",
+        "f39a6e7d-405e-4c07-8c18-95845f9da80e",
+    ]
+
+    state = "suspended"
+    with requests_mock.mock() as m:
+        m.put(f"{server.schedules.baseurl}?state={state}", text=BATCH_UPDATE_STATE.read_text())
+        resp = server.schedules.batch_update_state(new_schedules, state)
+
+    assert len(resp) == 3
+    for sch, r in zip(new_schedules, resp):
+        assert sch == r
+
+
+def test_batch_update_state_all(server: TSC.Server) -> None:
+    server.version = "3.27"
+    new_schedules = [
+        "593d2ebf-0d18-4deb-9d21-b113a4902583",
+        "cecbb71e-def0-4030-8068-5ae50f51db1c",
+        "f39a6e7d-405e-4c07-8c18-95845f9da80e",
+    ]
+
+    state = "suspended"
+    with requests_mock.mock() as m:
+        m.put(f"{server.schedules.baseurl}?state={state}&updateAll=true", text=BATCH_UPDATE_STATE.read_text())
+        _ = server.schedules.batch_update_state(new_schedules, state, True)
+
+        history = m.request_history[0]
+
+    assert history.text == "<tsRequest />"