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.
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/CHANGELOG.md +6 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/PKG-INFO +8 -5
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/__init__.py +4 -2
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/cli.py +70 -4
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/config/env.py +1 -1
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/constants/download.py +9 -2
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/core/_mixins.py +1 -1
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/core/client.py +22 -27
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/dataquery.py +18 -9
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/download/parallel.py +23 -1
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/sse/client.py +14 -1
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/sse/subscriber.py +78 -2
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/transport/auth.py +3 -3
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/types/models.py +47 -15
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/utils.py +292 -28
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery_sdk.egg-info/PKG-INFO +8 -5
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery_sdk.egg-info/requires.txt +4 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/pyproject.toml +10 -4
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/CONTRIBUTING.md +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/LICENSE +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/MANIFEST.in +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/README.md +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/config/__init__.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/config/logging.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/constants/__init__.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/constants/api.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/constants/rate_limit.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/constants/sse.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/core/__init__.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/core/_sync.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/data/function.json +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/download/__init__.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/download/utils.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/export.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/function_registry.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/py.typed +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/sse/__init__.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/sse/event_store.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/transport/__init__.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/transport/connection_pool.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/transport/rate_limiter.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/transport/retry.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/types/__init__.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery/types/exceptions.py +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery_sdk.egg-info/SOURCES.txt +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery_sdk.egg-info/dependency_links.txt +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery_sdk.egg-info/entry_points.txt +0 -0
- {dataquery_sdk-1.2.0 → dataquery_sdk-1.2.2}/dataquery_sdk.egg-info/top_level.txt +0 -0
- {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.
|
|
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/
|
|
7
|
-
Project-URL: Bug Tracker, https://github.com/
|
|
8
|
-
Project-URL: Documentation, https://github.com/
|
|
9
|
-
Project-URL: Source Code, https://github.com/
|
|
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/
|
|
14
|
+
For more information, visit: https://github.com/jpmorganchase/dataquery-sdk
|
|
15
15
|
"""
|
|
16
16
|
|
|
17
|
-
__version__ = "1.2.
|
|
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
|
-
"
|
|
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
|
|
501
|
+
async def cmd_search(args: argparse.Namespace) -> int:
|
|
487
502
|
async with DataQuery(args.env_file) as dq:
|
|
488
|
-
result = await dq.
|
|
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
|
-
"
|
|
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.
|
|
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 ``
|
|
32
|
-
#
|
|
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
|
|
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.
|
|
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
|
-
|
|
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
|
-
|
|
520
|
-
|
|
521
|
-
|
|
522
|
-
|
|
523
|
-
|
|
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
|
|
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.
|
|
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
|
-
|
|
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=
|
|
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
|
|
2051
|
-
"""Synchronous wrapper for :meth:`
|
|
2052
|
-
return self._run_sync(self.
|
|
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.
|
|
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
|
|
10
|
-
|
|
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.
|
|
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.
|
|
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.
|
|
175
|
+
"client_secret": self.config.get_client_secret(),
|
|
176
176
|
}
|
|
177
177
|
|
|
178
178
|
if not self.config.oauth_token_url:
|