dataquery-sdk 1.2.0__tar.gz → 1.2.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (49) hide show
  1. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/CHANGELOG.md +6 -0
  2. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/PKG-INFO +8 -5
  3. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/__init__.py +4 -2
  4. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/cli.py +70 -4
  5. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/config/env.py +1 -1
  6. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/constants/download.py +9 -2
  7. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/core/_mixins.py +1 -1
  8. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/core/client.py +22 -27
  9. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/dataquery.py +18 -9
  10. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/download/parallel.py +23 -1
  11. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/sse/client.py +14 -1
  12. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/sse/subscriber.py +78 -2
  13. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/transport/auth.py +3 -3
  14. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/types/models.py +47 -15
  15. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/utils.py +292 -28
  16. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery_sdk.egg-info/PKG-INFO +8 -5
  17. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery_sdk.egg-info/requires.txt +4 -0
  18. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/pyproject.toml +10 -4
  19. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/CONTRIBUTING.md +0 -0
  20. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/LICENSE +0 -0
  21. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/MANIFEST.in +0 -0
  22. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/README.md +0 -0
  23. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/config/__init__.py +0 -0
  24. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/config/logging.py +0 -0
  25. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/constants/__init__.py +0 -0
  26. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/constants/api.py +0 -0
  27. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/constants/rate_limit.py +0 -0
  28. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/constants/sse.py +0 -0
  29. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/core/__init__.py +0 -0
  30. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/core/_sync.py +0 -0
  31. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/data/function.json +0 -0
  32. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/download/__init__.py +0 -0
  33. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/download/utils.py +0 -0
  34. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/export.py +0 -0
  35. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/function_registry.py +0 -0
  36. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/py.typed +0 -0
  37. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/sse/__init__.py +0 -0
  38. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/sse/event_store.py +0 -0
  39. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/transport/__init__.py +0 -0
  40. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/transport/connection_pool.py +0 -0
  41. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/transport/rate_limiter.py +0 -0
  42. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/transport/retry.py +0 -0
  43. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/types/__init__.py +0 -0
  44. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/types/exceptions.py +0 -0
  45. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery_sdk.egg-info/SOURCES.txt +0 -0
  46. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery_sdk.egg-info/dependency_links.txt +0 -0
  47. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery_sdk.egg-info/entry_points.txt +0 -0
  48. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery_sdk.egg-info/top_level.txt +0 -0
  49. {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/setup.cfg +0 -0
@@ -40,5 +40,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
40
40
  - SSE: jittered reconnect backoff to avoid synchronized reconnect storms; stop reconnecting on fatal 403/404 and bound retries on 401; idle `sock_read` timeout distinguished from the heartbeat watchdog; honor the server `retry:` hint; strip a leading UTF-8 BOM; larger read buffer guards against `LineTooLong`; `stop()` is await-safe under concurrent callers
41
41
  ## [1.2.0] - 2026-06-29
42
42
  - DataQuery functions: new `function-help` command for local lookup of all 158 DQ function syntaxes, parameters, and categories (no API call); backed by a static, frozen `dataquery/data/function.json` dataset
43
+ ## [1.2.1] - 2026-07-14
44
+ - Written research: new `download_zip_async` helper that downloads a group over a date range (split into calendar-month windows to fit the available-files endpoint limit) and safely extracts ZIP archives as each download completes, overlapping unzip with in-flight downloads
45
+ - Group downloads: `run_group_download_async` accepts an `on_file_complete` async callback awaited per file on `completed`/`already_exists`
46
+ - Extraction is Zip Slip-guarded, skips current-day archives, and surfaces failures via `extraction_errors` (downgrading overall status to `partial`); date windows with no available files no longer mark a multi-window run as `partial`
47
+ ## [1.2.2] - 2026-07-18
48
+ - MCP: new `mcp-connect` CLI command
43
49
 
44
50
 
@@ -1,12 +1,12 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dataquery-sdk
3
- Version: 1.2.0
3
+ Version: 1.2.2
4
4
  Summary: Python SDK for DATAQUERY Data API - Query, download, and check availability of economic data files
5
5
  Author-email: DATAQUERY SDK Team <dataquery_support@jpmorgan.com>
6
- Project-URL: Homepage, https://github.com/dataquery/dataquery-sdk
7
- Project-URL: Bug Tracker, https://github.com/dataquery/dataquery-sdk/issues
8
- Project-URL: Documentation, https://github.com/dataquery/dataquery-sdk/wiki
9
- Project-URL: Source Code, https://github.com/dataquery/dataquery-sdk
6
+ Project-URL: Homepage, https://github.com/jpmorganchase/dataquery-sdk
7
+ Project-URL: Bug Tracker, https://github.com/jpmorganchase/dataquery-sdk/issues
8
+ Project-URL: Documentation, https://github.com/jpmorganchase/dataquery-sdk/wiki
9
+ Project-URL: Source Code, https://github.com/jpmorganchase/dataquery-sdk
10
10
  Keywords: dataquery,data,api,economic,financial,download,async,oauth,rate-limiting,sdk
11
11
  Classifier: Development Status :: 5 - Production/Stable
12
12
  Classifier: Intended Audience :: Developers
@@ -49,6 +49,8 @@ Requires-Dist: mike>=1.1.0; extra == "docs"
49
49
  Requires-Dist: pymdown-extensions>=10.21.3; extra == "docs"
50
50
  Provides-Extra: pandas
51
51
  Requires-Dist: pandas>=2.0.0; extra == "pandas"
52
+ Provides-Extra: mcp
53
+ Requires-Dist: fastmcp>=2.14; extra == "mcp"
52
54
  Provides-Extra: all
53
55
  Requires-Dist: pytest>=7.0.0; extra == "all"
54
56
  Requires-Dist: pytest-asyncio>=0.21.0; extra == "all"
@@ -61,6 +63,7 @@ Requires-Dist: mkdocs-material>=9.0.0; extra == "all"
61
63
  Requires-Dist: mkdocstrings[python]>=0.24.0; extra == "all"
62
64
  Requires-Dist: mike>=1.1.0; extra == "all"
63
65
  Requires-Dist: pandas>=2.0.0; extra == "all"
66
+ Requires-Dist: fastmcp>=2.14; extra == "all"
64
67
  Requires-Dist: pip-audit>=2.7.0; extra == "all"
65
68
  Requires-Dist: build>=1.3.0; extra == "all"
66
69
  Requires-Dist: urllib3>=2.7.0; extra == "all"
@@ -11,10 +11,10 @@ Quick Start:
11
11
  ... groups = await dq.list_groups_async()
12
12
  ... print(f"Found {len(groups)} groups")
13
13
 
14
- For more information, visit: https://github.com/dataquery/dataquery-sdk
14
+ For more information, visit: https://github.com/jpmorganchase/dataquery-sdk
15
15
  """
16
16
 
17
- __version__ = "1.2.0"
17
+ __version__ = "1.2.2"
18
18
  __author__ = "DATAQUERY SDK Team"
19
19
  __email__ = "dataquery_support@jpmorgan.com"
20
20
  __license__ = "MIT"
@@ -63,6 +63,7 @@ from .types.models import (
63
63
  Group,
64
64
  GroupList,
65
65
  )
66
+ from .utils import download_zip_async
66
67
 
67
68
  __all__ = [
68
69
  "DataQuery",
@@ -103,6 +104,7 @@ __all__ = [
103
104
  "get_function_registry",
104
105
  "list_functions_by_category",
105
106
  "lookup_function",
107
+ "download_zip_async",
106
108
  ]
107
109
 
108
110
  __version_info__ = tuple(int(x) for x in __version__.split("."))
@@ -4,6 +4,7 @@ import argparse
4
4
  import asyncio
5
5
  import json
6
6
  import sys
7
+ from collections.abc import AsyncGenerator
7
8
  from pathlib import Path
8
9
  from typing import Any, Dict, List, Optional
9
10
 
@@ -183,7 +184,7 @@ def create_parser() -> argparse.ArgumentParser:
183
184
  _ = auth_sub.add_parser("test", help="Test authentication by listing groups")
184
185
 
185
186
  p_search = subparsers.add_parser(
186
- "text-search",
187
+ "search",
187
188
  help="Search the DataQuery catalog using a natural-language query (POST /search)",
188
189
  )
189
190
  p_search.add_argument("--query", required=True, help="Free-text search query")
@@ -290,6 +291,20 @@ def create_parser() -> argparse.ArgumentParser:
290
291
 
291
292
  subparsers.add_parser("heartbeat", help="Check if DataQuery is running")
292
293
 
294
+ p_connect = subparsers.add_parser(
295
+ "mcp-connect",
296
+ help="Connect a local stdio MCP client to a remote MCP server (AuthE OAuth)",
297
+ description=(
298
+ "Bridge a desktop MCP client (stdio) to a remote streamable-HTTP MCP\n"
299
+ "server, authenticating with an OAuth client-credentials (AuthE) token\n"
300
+ "minted from the DATAQUERY_* environment. Point your MCP client's\n"
301
+ "`command` at: dataquery mcp-connect --url <MCP_URL>"
302
+ ),
303
+ formatter_class=argparse.RawDescriptionHelpFormatter,
304
+ )
305
+ p_connect.add_argument("--url", required=True, help="Remote MCP endpoint URL")
306
+ p_connect.add_argument("--name", default="dataquery-mcp", help="Proxy server name (default: dataquery-mcp)")
307
+
293
308
  return parser
294
309
 
295
310
 
@@ -483,9 +498,9 @@ async def cmd_auth_test(args: argparse.Namespace) -> int:
483
498
  return 0
484
499
 
485
500
 
486
- async def cmd_text_search(args: argparse.Namespace) -> int:
501
+ async def cmd_search(args: argparse.Namespace) -> int:
487
502
  async with DataQuery(args.env_file) as dq:
488
- result = await dq.text_search_async(args.query)
503
+ result = await dq.search_async(args.query)
489
504
 
490
505
  if args.json:
491
506
  print(json.dumps(result, indent=2))
@@ -716,6 +731,56 @@ def cmd_function_help(args: argparse.Namespace) -> int:
716
731
  return 0
717
732
 
718
733
 
734
+ async def cmd_mcp_connect(args: argparse.Namespace) -> int:
735
+ """Bridge a desktop MCP client (stdio) to a remote streamable-HTTP MCP server.
736
+
737
+ Mints an OAuth client-credentials (AuthE) token with the SDK's own
738
+ TokenManager from the ``DATAQUERY_*`` environment and attaches a fresh bearer
739
+ token to every upstream request. Launched by an MCP client as its stdio
740
+ ``command``; stdout carries only the MCP JSON-RPC stream, so all diagnostics
741
+ go to stderr.
742
+ """
743
+ try:
744
+ import httpx
745
+ from fastmcp import FastMCP
746
+ from fastmcp.client.transports import StreamableHttpTransport
747
+ except ImportError:
748
+ print(
749
+ "The MCP bridge requires the 'mcp' extra. Install it with:\n"
750
+ " pip install 'dataquery-sdk[mcp]'\n"
751
+ "or run it directly with:\n"
752
+ " uvx --from 'dataquery-sdk[mcp]' dataquery mcp-connect --url <MCP_URL>",
753
+ file=sys.stderr,
754
+ )
755
+ return 1
756
+
757
+ from dataquery.config import EnvConfig
758
+ from dataquery.transport.auth import TokenManager
759
+
760
+ config = EnvConfig.create_client_config(env_file=Path(args.env_file) if getattr(args, "env_file", None) else None)
761
+ token_manager = TokenManager(config)
762
+
763
+ class _AutheAuth(httpx.Auth):
764
+ """Stamp a fresh AuthE bearer token (from the SDK TokenManager) per request."""
765
+
766
+ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
767
+ header = await token_manager.get_valid_token() # "Bearer <jwt>"
768
+ if not header:
769
+ raise DataQueryError(
770
+ "Could not obtain an OAuth token \u2014 check DATAQUERY_CLIENT_ID, "
771
+ "DATAQUERY_CLIENT_SECRET, DATAQUERY_OAUTH_TOKEN_URL and "
772
+ "DATAQUERY_OAUTH_AUD."
773
+ )
774
+ request.headers["Authorization"] = header
775
+ yield request
776
+
777
+ transport = StreamableHttpTransport(args.url, auth=_AutheAuth())
778
+ proxy = FastMCP.as_proxy(transport, name=args.name)
779
+ # stdout is reserved for the MCP JSON-RPC stream; keep the banner off it.
780
+ await proxy.run_async(transport="stdio", show_banner=False)
781
+ return 0
782
+
783
+
719
784
  def main_sync(ns: argparse.Namespace) -> int:
720
785
  if ns.command == "config":
721
786
  if ns.config_command == "show":
@@ -734,7 +799,7 @@ _ASYNC_COMMANDS = {
734
799
  "availability": cmd_availability,
735
800
  "download": cmd_download,
736
801
  "download-group": cmd_download_group,
737
- "text-search": cmd_text_search,
802
+ "search": cmd_search,
738
803
  "groups-search": cmd_groups_search,
739
804
  "instruments": cmd_instruments,
740
805
  "instruments-search": cmd_instruments_search,
@@ -745,6 +810,7 @@ _ASYNC_COMMANDS = {
745
810
  "expression-timeseries": cmd_expression_timeseries,
746
811
  "grid-data": cmd_grid_data,
747
812
  "heartbeat": cmd_heartbeat,
813
+ "mcp-connect": cmd_mcp_connect,
748
814
  }
749
815
 
750
816
 
@@ -247,7 +247,7 @@ class EnvConfig:
247
247
  if config.oauth_enabled:
248
248
  if not config.client_id:
249
249
  errors.append("CLIENT_ID is required when OAuth is enabled")
250
- if not config.client_secret:
250
+ if not config.get_client_secret():
251
251
  errors.append("CLIENT_SECRET is required when OAuth is enabled")
252
252
  if not config.oauth_token_url:
253
253
  errors.append("OAUTH_TOKEN_URL is required when OAuth is enabled")
@@ -28,6 +28,13 @@ MBPS_TO_BYTES_PER_SECOND = 125_000
28
28
 
29
29
  # Default size, in days, of each chunk when splitting a date range so that
30
30
  # the available-files endpoint is queried over a smaller window per call.
31
- # Used by ``run_group_download_chunked_async`` to honour group-level
32
- # date-range limits (e.g. RESEARCH_EQUITY_ALL is restricted to 14 days).
31
+ # Used by ``split_date_range_into_chunks``. The endpoint caps each call at
32
+ # one calendar month; smaller windows just keep responses lighter. (The old
33
+ # group-level limits, e.g. 14 days for RESEARCH_EQUITY_ALL, were lifted
34
+ # server-side.)
33
35
  DEFAULT_WRITTEN_RESEARCH_CHUNK_DAYS: int = 7
36
+
37
+ # Error string set on the OperationReport when the available-files endpoint
38
+ # returns nothing for a date window. Chunked workflows match on it to tell
39
+ # a quiet window apart from a real failure.
40
+ NO_FILES_FOUND_ERROR: str = "No available files found for date range"
@@ -609,7 +609,7 @@ class GridMixin(_RequestProto):
609
609
  class SearchMixin(_RequestProto):
610
610
  """Natural-language catalog search via POST /search."""
611
611
 
612
- async def text_search_async(self, query: str) -> Dict[str, Any]:
612
+ async def search_async(self, query: str) -> Dict[str, Any]:
613
613
  """Search the DataQuery catalog using a natural-language query.
614
614
 
615
615
  POSTs ``{"query": query}`` to ``/search`` and returns the parsed JSON body.
@@ -181,7 +181,7 @@ class DataQueryClient(
181
181
  raise ConfigurationError("Invalid base_url format")
182
182
 
183
183
  if strict_oauth_check and self.config.oauth_enabled:
184
- if not self.config.client_id or not self.config.client_secret:
184
+ if not self.config.client_id or not self.config.get_client_secret():
185
185
  raise ConfigurationError("client_id and client_secret are required")
186
186
 
187
187
  if not self.auth_manager.is_authenticated():
@@ -494,33 +494,28 @@ class DataQueryClient(
494
494
  if self.session is None:
495
495
  raise NetworkError("Failed to establish connection")
496
496
 
497
- try:
498
- response = await self.session.request(method, url, **kwargs)
499
-
500
- # Raise 429 inside the retry scope so the retry manager backs off and retries.
501
- if response.status == 429:
502
- self.rate_limiter.handle_rate_limit_response(dict(response.headers))
503
- retry_after = self._parse_retry_after(response.headers)
504
- # Drain body so the connection can be reused; ignore decode errors.
505
- try:
506
- await response.text()
507
- except (UnicodeDecodeError, aiohttp.ClientPayloadError):
508
- pass
509
- raise RateLimitError(f"Rate limit exceeded: {response.status}", retry_after=retry_after)
510
-
511
- if response.status >= 500:
512
- # Drain body so the connection can be reused; ignore decode errors.
513
- try:
514
- await response.text()
515
- except (UnicodeDecodeError, aiohttp.ClientPayloadError):
516
- pass
517
- raise NetworkError(f"Server error: {response.status}", status_code=response.status)
497
+ response = await self.session.request(method, url, **kwargs)
518
498
 
519
- return response
520
- except (NetworkError, RateLimitError):
521
- raise
522
- except Exception:
523
- raise
499
+ # Raise 429 inside the retry scope so the retry manager backs off and retries.
500
+ if response.status == 429:
501
+ self.rate_limiter.handle_rate_limit_response(dict(response.headers))
502
+ retry_after = self._parse_retry_after(response.headers)
503
+ # Drain body so the connection can be reused; ignore decode errors.
504
+ try:
505
+ await response.text()
506
+ except (UnicodeDecodeError, aiohttp.ClientPayloadError):
507
+ pass
508
+ raise RateLimitError(f"Rate limit exceeded: {response.status}", retry_after=retry_after)
509
+
510
+ if response.status >= 500:
511
+ # Drain body so the connection can be reused; ignore decode errors.
512
+ try:
513
+ await response.text()
514
+ except (UnicodeDecodeError, aiohttp.ClientPayloadError):
515
+ pass
516
+ raise NetworkError(f"Server error: {response.status}", status_code=response.status)
517
+
518
+ return response
524
519
 
525
520
  async def list_groups_async(self, limit: Optional[int] = None) -> List[Group]:
526
521
  """
@@ -10,12 +10,14 @@ import time
10
10
  from calendar import monthrange
11
11
  from datetime import date, datetime
12
12
  from pathlib import Path
13
- from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
13
+ from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Tuple, Union
14
14
 
15
15
  import structlog
16
16
  from dotenv import load_dotenv
17
+ from pydantic import SecretStr
17
18
 
18
19
  from .config import EnvConfig
20
+ from .constants.download import NO_FILES_FOUND_ERROR
19
21
  from .core._sync import SyncRunner
20
22
  from .core.client import DataQueryClient
21
23
  from .types.exceptions import ConfigurationError
@@ -186,7 +188,7 @@ class DataQuery:
186
188
  if client_id:
187
189
  self.client_config.client_id = client_id
188
190
  if client_secret:
189
- self.client_config.client_secret = client_secret
191
+ self.client_config.client_secret = SecretStr(client_secret)
190
192
  if not self.client_config.oauth_token_url and self.client_config.base_url:
191
193
  self.client_config.oauth_token_url = f"{self.client_config.base_url.rstrip('/')}/oauth/token"
192
194
 
@@ -717,7 +719,7 @@ class DataQuery:
717
719
  client = self._ensure_client()
718
720
  return await client.get_grid_data_async(expr, grid_id, date)
719
721
 
720
- async def text_search_async(self, query: str) -> Dict[str, Any]:
722
+ async def search_async(self, query: str) -> Dict[str, Any]:
721
723
  """Search the DataQuery catalog using a natural-language query.
722
724
 
723
725
  Args:
@@ -728,7 +730,7 @@ class DataQuery:
728
730
  """
729
731
  await self.connect_async()
730
732
  client = self._ensure_client()
731
- return await client.text_search_async(query)
733
+ return await client.search_async(query)
732
734
 
733
735
  async def run_groups_async(self, max_concurrent: int = 5) -> OperationReport:
734
736
  """Run complete operation for listing all groups."""
@@ -907,6 +909,7 @@ class DataQuery:
907
909
  delay_between_downloads: float = 0.2,
908
910
  max_retries: int = 3,
909
911
  file_group_id: Optional[Union[str, List[str]]] = None,
912
+ on_file_complete: Optional[Callable[["DownloadResult"], Awaitable[None]]] = None,
910
913
  ) -> OperationReport:
911
914
  """
912
915
  Download all files in a group for a date range using parallel HTTP range requests.
@@ -929,6 +932,11 @@ class DataQuery:
929
932
  file_group_id: Optional restriction to specific file-group-id(s). Accepts a
930
933
  single id or a list of ids. When a list is supplied, availability is
931
934
  queried in parallel per id and the union of dates is downloaded.
935
+ on_file_complete: Optional async callback awaited for each file as soon as
936
+ it finishes downloading (status ``completed``/``already_exists``). Runs
937
+ concurrently with the downloads still in flight, so post-processing such
938
+ as unzipping overlaps with subsequent downloads rather than waiting for
939
+ the whole batch.
932
940
 
933
941
  Returns:
934
942
  Dictionary with download results and statistics
@@ -1006,7 +1014,7 @@ class DataQuery:
1006
1014
  total_time_minutes = total_time_seconds / 60.0
1007
1015
 
1008
1016
  logger.warning(
1009
- "No available files found for date range",
1017
+ NO_FILES_FOUND_ERROR,
1010
1018
  group_id=group_id,
1011
1019
  start_date=start_date,
1012
1020
  end_date=end_date,
@@ -1014,7 +1022,7 @@ class DataQuery:
1014
1022
  return OperationReport(
1015
1023
  operation="group_download",
1016
1024
  status="error",
1017
- error="No available files found for date range",
1025
+ error=NO_FILES_FOUND_ERROR,
1018
1026
  subject={"group_id": group_id, "start_date": start_date, "end_date": end_date},
1019
1027
  counts={"total_files": 0, "successful_downloads": 0, "failed_downloads": 0},
1020
1028
  timing={
@@ -1071,6 +1079,7 @@ class DataQuery:
1071
1079
  base_retry_delay=delay_between_downloads,
1072
1080
  max_retries=max_retries,
1073
1081
  progress_callback=progress_callback,
1082
+ on_file_complete=on_file_complete,
1074
1083
  )
1075
1084
 
1076
1085
  operation_end_time = time.time()
@@ -2047,9 +2056,9 @@ class DataQuery:
2047
2056
  """
2048
2057
  return self._run_sync(self.get_grid_data_async(expr, grid_id, date))
2049
2058
 
2050
- def text_search(self, query: str) -> Dict[str, Any]:
2051
- """Synchronous wrapper for :meth:`text_search_async`."""
2052
- return self._run_sync(self.text_search_async(query))
2059
+ def search(self, query: str) -> Dict[str, Any]:
2060
+ """Synchronous wrapper for :meth:`search_async`."""
2061
+ return self._run_sync(self.search_async(query))
2053
2062
 
2054
2063
  def run_groups(self, max_concurrent: int = 5) -> OperationReport:
2055
2064
  """Synchronous wrapper for run_groups_async."""
@@ -16,7 +16,7 @@ import time
16
16
  from dataclasses import dataclass
17
17
  from datetime import datetime
18
18
  from pathlib import Path
19
- from typing import IO, TYPE_CHECKING, Callable, Optional
19
+ from typing import IO, TYPE_CHECKING, Awaitable, Callable, Optional
20
20
 
21
21
  import aiohttp
22
22
  import structlog
@@ -530,6 +530,7 @@ async def _download_one_with_stagger(
530
530
  global_semaphore: asyncio.Semaphore,
531
531
  delay_seconds: float,
532
532
  progress_callback: Optional[Callable],
533
+ on_file_complete: Optional[Callable[[DownloadResult], Awaitable[None]]] = None,
533
534
  ) -> Optional[DownloadResult]:
534
535
  file_group_id = _file_id(file_info)
535
536
  file_datetime = _file_dt(file_info)
@@ -554,6 +555,21 @@ async def _download_one_with_stagger(
554
555
  file_datetime=file_datetime,
555
556
  status=result.status.value if result else "failed",
556
557
  )
558
+ if (
559
+ on_file_complete is not None
560
+ and result is not None
561
+ and result.status.value in ("completed", "already_exists")
562
+ ):
563
+ # Post-processing (e.g. unzip) runs here so it overlaps with the
564
+ # downloads of other files still in flight under the gather().
565
+ try:
566
+ await on_file_complete(result)
567
+ except Exception as cb_err: # pragma: no cover - defensive
568
+ logger.warning(
569
+ "on_file_complete callback failed",
570
+ file_group_id=file_group_id,
571
+ error=str(cb_err),
572
+ )
557
573
  return result
558
574
  except Exception as e:
559
575
  logger.error(
@@ -596,6 +612,7 @@ async def download_files_with_retry(
596
612
  base_retry_delay: float,
597
613
  max_retries: int,
598
614
  progress_callback: Optional[Callable] = None,
615
+ on_file_complete: Optional[Callable[[DownloadResult], Awaitable[None]]] = None,
599
616
  ) -> tuple[list[DownloadResult], list[dict], int]:
600
617
  """Run a staggered, retrying batch of parallel-range downloads.
601
618
 
@@ -605,6 +622,10 @@ async def download_files_with_retry(
605
622
  spread the burst against the API rate limiter. Failures are retried up
606
623
  to ``max_retries`` times with exponential backoff (``base_retry_delay``).
607
624
 
625
+ When ``on_file_complete`` is provided it is awaited for each successfully
626
+ downloaded file as soon as that file finishes, so post-processing (e.g.
627
+ unzipping) overlaps with the downloads of files still in flight.
628
+
608
629
  Returns ``(successful, failed, retry_count)``.
609
630
  """
610
631
 
@@ -618,6 +639,7 @@ async def download_files_with_retry(
618
639
  global_semaphore=global_semaphore,
619
640
  delay_seconds=i * intelligent_delay,
620
641
  progress_callback=progress_callback,
642
+ on_file_complete=on_file_complete,
621
643
  )
622
644
  for i, fi in enumerate(batch)
623
645
  ]
@@ -111,6 +111,7 @@ class SSEClient:
111
111
  params: Optional[dict] = None,
112
112
  event_id_store: Optional[SSEEventIdStore] = None,
113
113
  heartbeat_timeout: float = 0.0,
114
+ defer_event_id_persistence: bool = False,
114
115
  ):
115
116
  """
116
117
  Initialise the SSE client.
@@ -147,6 +148,16 @@ class SSEClient:
147
148
  larger than the server's keep-alive interval. ``0`` (the
148
149
  default) disables the watchdog and relies on the server
149
150
  closing the stream cleanly.
151
+ defer_event_id_persistence: When ``True``, do not persist parsed
152
+ event ids to ``event_id_store`` as they arrive. The owner
153
+ (e.g. :class:`NotificationDownloadManager`) instead commits
154
+ a low-water-mark only after the matching download has
155
+ settled, so a crash mid-download replays the event instead
156
+ of skipping it (at-least-once). ``event_id_store`` is still
157
+ read once on construction to seed the first
158
+ ``Last-Event-ID``. Defaults to ``False`` — persist at parse
159
+ time, which is correct for standalone use where there is no
160
+ download step to gate persistence on.
150
161
  """
151
162
  self.config = config
152
163
  self.auth_manager = auth_manager
@@ -158,6 +169,7 @@ class SSEClient:
158
169
  self.params = dict(params) if params else None
159
170
  self.event_id_store = event_id_store
160
171
  self.heartbeat_timeout = heartbeat_timeout
172
+ self._defer_event_id_persistence = defer_event_id_persistence
161
173
 
162
174
  self._running = False
163
175
  self._task: Optional[asyncio.Task] = None
@@ -481,7 +493,8 @@ class SSEClient:
481
493
  event_id = field_value
482
494
  if field_value and field_value.isdigit():
483
495
  self._last_event_id = field_value
484
- self._persist_event_id(field_value)
496
+ if not self._defer_event_id_persistence:
497
+ self._persist_event_id(field_value)
485
498
  elif field_name == "retry":
486
499
  try:
487
500
  retry_ms = int(field_value)
@@ -6,8 +6,10 @@ downloads files as notifications arrive. Each SSE event uses the standard
6
6
  SSE format with the event type in the ``event:`` field (e.g., ``file-updated``)
7
7
  and the file details in the JSON ``data:`` payload.
8
8
 
9
- The SSE ``id:`` field is persisted to disk so that cross-process replay via
10
- the ``last-event-id`` query parameter works correctly on reconnection.
9
+ The SSE ``id:`` field is persisted to disk as a low-water-mark — advanced past
10
+ an event only once its download has settled — so cross-process replay via the
11
+ ``last-event-id`` query parameter re-delivers any file whose download did not
12
+ complete before a restart (at-least-once delivery).
11
13
  """
12
14
 
13
15
  import asyncio
@@ -245,6 +247,13 @@ class NotificationDownloadManager:
245
247
  self._download_semaphore: Optional[asyncio.Semaphore] = None
246
248
  self._inflight: Set[asyncio.Task] = set()
247
249
 
250
+ # Replay low-water-mark. The persisted last-event-id is only advanced
251
+ # past an event once its download has settled, so a crash mid-download
252
+ # replays that event on restart instead of silently skipping it.
253
+ self._uncommitted_ids: Set[int] = set()
254
+ self._highest_seen_event_id: int = 0
255
+ self._committed_event_id: int = 0
256
+
248
257
  self.stats: Dict[str, Any] = {
249
258
  "start_time": None,
250
259
  "notifications_received": 0,
@@ -285,6 +294,12 @@ class NotificationDownloadManager:
285
294
  if self._event_id_store is not None:
286
295
  stored_event_id = self._event_id_store.load()
287
296
 
297
+ if stored_event_id is not None and stored_event_id.isdigit():
298
+ # Seed the low-water-mark so we neither re-persist the stored id nor
299
+ # ever regress below it as newly received events settle.
300
+ self._committed_event_id = int(stored_event_id)
301
+ self._highest_seen_event_id = int(stored_event_id)
302
+
288
303
  if stored_event_id is not None:
289
304
  logger.info(
290
305
  "Event replay enabled — resuming from event-id %s; skipping initial bulk check.",
@@ -307,6 +322,9 @@ class NotificationDownloadManager:
307
322
  params=self.subscription.query_params(),
308
323
  event_id_store=self._event_id_store,
309
324
  heartbeat_timeout=self._heartbeat_timeout,
325
+ # This manager commits a low-water-mark after downloads settle,
326
+ # so the client must not persist ids itself at parse time.
327
+ defer_event_id_persistence=True,
310
328
  )
311
329
  await self._sse_client.start()
312
330
 
@@ -393,6 +411,7 @@ class NotificationDownloadManager:
393
411
  if not self._running:
394
412
  return
395
413
 
414
+ self._register_event_id(event)
396
415
  task = asyncio.create_task(self._run_handler(event), name="sse-download")
397
416
  self._inflight.add(task)
398
417
  task.add_done_callback(self._inflight.discard)
@@ -407,6 +426,63 @@ class NotificationDownloadManager:
407
426
  except Exception as exc:
408
427
  logger.error("Error handling notification: %s", exc)
409
428
  await self._dispatch_error(exc)
429
+ # The download has settled (success, skip, or handled failure); advance
430
+ # the replay cursor. Deliberately not in a ``finally``: if the task is
431
+ # cancelled mid-download the cursor stays behind this event so it is
432
+ # replayed on restart rather than skipped.
433
+ await self._commit_settled_event(event)
434
+
435
+ def _register_event_id(self, event: SSEEvent) -> None:
436
+ """Record a numeric event id as in-flight before its download starts.
437
+
438
+ The replay cursor is never advanced past an in-flight id, so a crash
439
+ before the download settles leaves the event to be replayed.
440
+ """
441
+ if self._event_id_store is None:
442
+ return
443
+ eid = event.id
444
+ if not eid or not eid.isdigit():
445
+ return
446
+ value = int(eid)
447
+ self._uncommitted_ids.add(value)
448
+ if value > self._highest_seen_event_id:
449
+ self._highest_seen_event_id = value
450
+
451
+ def _next_commit_id(self) -> Optional[str]:
452
+ """Return the highest replay id now safe to persist, or ``None``.
453
+
454
+ The safe id is the low-water-mark: everything below the oldest
455
+ still-in-flight event has settled, so commit ``min(in-flight) - 1``
456
+ (or the highest id seen when nothing is in flight). Fully synchronous
457
+ so two concurrent settlers can't persist a decreasing watermark across
458
+ the ``await`` in :meth:`_commit_settled_event`.
459
+ """
460
+ if self._uncommitted_ids:
461
+ watermark = min(self._uncommitted_ids) - 1
462
+ else:
463
+ watermark = self._highest_seen_event_id
464
+ if watermark > self._committed_event_id:
465
+ self._committed_event_id = watermark
466
+ return str(watermark)
467
+ return None
468
+
469
+ async def _commit_settled_event(self, event: SSEEvent) -> None:
470
+ """Advance the persisted replay cursor after an event's handler settles.
471
+
472
+ Runs whether the download succeeded, was skipped, or failed — the id is
473
+ no longer in-flight either way. A terminal download failure is already
474
+ surfaced via ``error_callback``/``stats``; the cursor still advances so a
475
+ genuinely unfetchable file can't stall replay forever.
476
+ """
477
+ store = self._event_id_store
478
+ if store is None:
479
+ return
480
+ eid = event.id
481
+ if eid and eid.isdigit():
482
+ self._uncommitted_ids.discard(int(eid))
483
+ to_save = self._next_commit_id()
484
+ if to_save is not None:
485
+ await store.save(to_save)
410
486
 
411
487
  async def _on_sse_error(self, exc: Exception) -> None:
412
488
  """Called by SSEClient on connection errors."""
@@ -72,7 +72,7 @@ class TokenManager:
72
72
  Bearer token string or None if no valid token available
73
73
  """
74
74
  if self.config.has_bearer_token:
75
- return f"Bearer {self.config.bearer_token}"
75
+ return f"Bearer {self.config.get_bearer_token()}"
76
76
 
77
77
  if not self.config.has_oauth_credentials:
78
78
  logger.warning("No OAuth credentials or bearer token configured")
@@ -111,7 +111,7 @@ class TokenManager:
111
111
  if not self.config.oauth_token_url:
112
112
  raise ConfigurationError("OAuth token URL not configured")
113
113
 
114
- if not self.config.client_id or not self.config.client_secret:
114
+ if not self.config.client_id or not self.config.get_client_secret():
115
115
  raise ConfigurationError("client_id and client_secret are required for OAuth")
116
116
 
117
117
  token_request = TokenRequest(
@@ -172,7 +172,7 @@ class TokenManager:
172
172
  "grant_type": "refresh_token",
173
173
  "refresh_token": self.current_token.refresh_token,
174
174
  "client_id": self.config.client_id,
175
- "client_secret": self.config.client_secret,
175
+ "client_secret": self.config.get_client_secret(),
176
176
  }
177
177
 
178
178
  if not self.config.oauth_token_url: