diff --git a/README.rst b/README.rst index a10f8be..ea1929b 100644 --- a/README.rst +++ b/README.rst @@ -116,6 +116,16 @@ them: In [5]: response[:215] Out[5]: b'\n\n\n' +You can override the report's date range per request with ``period`` (the +``p`` parameter - a lookback in days, up to 365) or ``fd``/``td`` (explicit +``yyyymmdd`` from/to dates, given together); the two forms are mutually +exclusive. This is useful because the v3 Flex Web Service may not honor a +query's saved "Last N Calendar Days" period: + +.. code:: python + + In [6]: response = client.download(token, query_id, period=90) + ``pip install ibflex[web]`` also installs a ``flexget`` console script: diff --git a/ibflex/client.py b/ibflex/client.py index 56b681f..a34781e 100755 --- a/ibflex/client.py +++ b/ibflex/client.py @@ -99,15 +99,32 @@ class StatementError: ############################################################################### # FUNCTIONS ############################################################################### -def download(token: str, query_id: str, max_tries: int | None = 5) -> bytes: +def download( + token: str, + query_id: str, + max_tries: int | None = 5, + period: int | None = None, + fd: str | None = None, + td: str | None = None, +) -> bytes: """2-step FlexQueryReport download process. Args: token: Current access token from Reports > Settings > FlexWeb Service. query_id: Flex Query ID from Reports > Flex Queries > Custom Flex Queries > Configure. + max_tries: Number of times to poll for statement generation before + raising StatementGenerationTimeout. + period: Optional lookback override, in days (max 365), sent as the + ``p`` SendRequest parameter. + fd, td: Optional explicit from/to date override (``yyyymmdd``, at most + 365 days apart), sent as the ``fd``/``td`` SendRequest + parameters; must be given together. ``period`` and + ``fd``/``td`` are mutually exclusive. """ - stmt_access = request_statement(token, query_id) + stmt_access = request_statement( + token, query_id, period=period, fd=fd, td=td + ) status = 0 tries = 0 while status is not True: @@ -127,12 +144,22 @@ def download(token: str, query_id: str, max_tries: int | None = 5) -> bytes: def request_statement( - token: str, query_id: str, url: str | None = None + token: str, + query_id: str, + url: str | None = None, + period: int | None = None, + fd: str | None = None, + td: str | None = None, ) -> StatementAccess: """First part of the 2-step download process. + + ``period``/``fd``/``td`` are optional date-range overrides applied only + to this SendRequest, not to the GetStatement polls. """ url = url or REQUEST_URL - response = submit_request(url, token, query=query_id) + response = submit_request( + url, token, query=query_id, period=period, fd=fd, td=td + ) stmt_access = parse_stmt_response(response) if isinstance(stmt_access, StatementError): raise ResponseCodeError( @@ -142,21 +169,46 @@ def request_statement( return stmt_access -def submit_request(url: str, token: str, query: str) -> requests.Response: +def submit_request( + url: str, + token: str, + query: str, + period: int | None = None, + fd: str | None = None, + td: str | None = None, +) -> requests.Response: """Post a query to an API access point, along with an authentication token. + ``period`` (``p``), ``fd`` and ``td`` are optional date-range overrides + added to the query string after the required ``v``/``t``/``q`` params. + ``period`` and ``fd``/``td`` are mutually exclusive, and ``fd``/``td`` + must be given together. + Retry with a progressive timeout window. """ + if period is not None and (fd is not None or td is not None): + raise ValueError("Pass either `period` or `fd`/`td`, not both") + if (fd is None) != (td is None): + raise ValueError("Pass both `fd` and `td`, or neither") + MAX_REQUESTS = 3 TIMEOUT_INCREMENT = 5 + params = {"v": "3", "t": token, "q": query} + if period is not None: + params["p"] = str(period) + if fd is not None: + params["fd"] = fd + if td is not None: + params["td"] = td + response = None req_count = 1 while (not response): try: response = requests.get( url, - params={"v": "3", "t": token, "q": query}, + params=params, headers={"user-agent": "Java"}, timeout=req_count * TIMEOUT_INCREMENT, ) @@ -243,9 +295,20 @@ def main(): help='Current Flex Web Service token') argparser.add_argument('--query', '-q', required=True, help='Flex Query ID#') + argparser.add_argument('--period', '-p', type=int, default=None, + help='Lookback override in days (max 365)') + argparser.add_argument('--fromdate', '--fd', dest='fd', default=None, + help='From-date override (yyyymmdd)') + argparser.add_argument('--todate', '--td', dest='td', default=None, + help='To-date override (yyyymmdd)') args = argparser.parse_args() - statement = download(args.token, args.query) + try: + statement = download( + args.token, args.query, period=args.period, fd=args.fd, td=args.td + ) + except ValueError as exc: + argparser.error(str(exc)) print(statement.decode()) diff --git a/tests/test_client.py b/tests/test_client.py index ceb75e1..90e4310 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -150,5 +150,85 @@ def test_request_statement(self, mock_requests_get: Mock): self.assertIsInstance(response, Types.FlexQueryResponse) +@patch("requests.get", side_effect=mock_response) +class DateRangeOverrideTestCase(unittest.TestCase): + def test_period_override_sent_on_send_request_only(self, mock_requests_get): + # `period` is forwarded to SendRequest as `p`, but the GetStatement + # poll must NOT carry it. + output = client.download( + token="DEADBEEF", + query_id="0987654321", + period=90, + ) + self.assertIsInstance(output, bytes) + + self.assertEqual( + mock_requests_get.call_args_list, + [ + call( + client.REQUEST_URL, + params={ + "v": "3", "t": "DEADBEEF", "q": "0987654321", "p": "90", + }, + headers={"user-agent": "Java"}, + timeout=5, + ), + call( + client.STMT_URL, + params={"v": "3", "t": "DEADBEEF", "q": "1234567890"}, + headers={"user-agent": "Java"}, + timeout=5, + ), + ], + ) + + def test_from_to_date_override(self, mock_requests_get): + client.request_statement( + token="DEADBEEF", + query_id="0987654321", + fd="20260101", + td="20260331", + ) + + mock_requests_get.assert_called_once_with( + client.REQUEST_URL, + params={ + "v": "3", "t": "DEADBEEF", "q": "0987654321", + "fd": "20260101", "td": "20260331", + }, + headers={"user-agent": "Java"}, + timeout=5, + ) + + def test_period_and_dates_are_mutually_exclusive(self, mock_requests_get): + with self.assertRaises(ValueError): + client.download( + token="DEADBEEF", + query_id="0987654321", + period=90, + fd="20260101", + td="20260131", + ) + mock_requests_get.assert_not_called() + + def test_fd_and_td_must_be_given_together(self, mock_requests_get): + for kwargs in ({"fd": "20260101"}, {"td": "20260131"}): + with self.assertRaises(ValueError): + client.download( + token="DEADBEEF", query_id="0987654321", **kwargs + ) + mock_requests_get.assert_not_called() + + def test_no_override_leaves_params_unchanged(self, mock_requests_get): + client.request_statement(token="DEADBEEF", query_id="0987654321") + + mock_requests_get.assert_called_once_with( + client.REQUEST_URL, + params={"v": "3", "t": "DEADBEEF", "q": "0987654321"}, + headers={"user-agent": "Java"}, + timeout=5, + ) + + if __name__ == '__main__': unittest.main(verbosity=3)