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

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/CHANGELOG.md

@@ -40,5 +40,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - 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
 ## [1.2.0] - 2026-06-29
 - 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
+## [1.2.1] - 2026-07-14
+- 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
+- Group downloads: `run_group_download_async` accepts an `on_file_complete` async callback awaited per file on `completed`/`already_exists`
+- 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`
 
 

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataquery-sdk
-Version: 1.2.0
+Version: 1.2.1
 Summary: Python SDK for DATAQUERY Data API - Query, download, and check availability of economic data files
 Author-email: DATAQUERY SDK Team <dataquery_support@jpmorgan.com>
 Project-URL: Homepage, https://github.com/dataquery/dataquery-sdk

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery/__init__.py

@@ -14,7 +14,7 @@ Quick Start:
 For more information, visit: https://github.com/dataquery/dataquery-sdk
 """
 
-__version__ = "1.2.0"
+__version__ = "1.2.1"
 __author__ = "DATAQUERY SDK Team"
 __email__ = "dataquery_support@jpmorgan.com"
 __license__ = "MIT"
@@ -63,6 +63,7 @@ from .types.models import (
     Group,
     GroupList,
 )
+from .utils import download_zip_async
 
 __all__ = [
     "DataQuery",
@@ -103,6 +104,7 @@ __all__ = [
     "get_function_registry",
     "list_functions_by_category",
     "lookup_function",
+    "download_zip_async",
 ]
 
 __version_info__ = tuple(int(x) for x in __version__.split("."))

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery/cli.py

@@ -183,7 +183,7 @@ def create_parser() -> argparse.ArgumentParser:
     _ = auth_sub.add_parser("test", help="Test authentication by listing groups")
 
     p_search = subparsers.add_parser(
-        "text-search",
+        "search",
         help="Search the DataQuery catalog using a natural-language query (POST /search)",
     )
     p_search.add_argument("--query", required=True, help="Free-text search query")
@@ -483,9 +483,9 @@ async def cmd_auth_test(args: argparse.Namespace) -> int:
     return 0
 
 
-async def cmd_text_search(args: argparse.Namespace) -> int:
+async def cmd_search(args: argparse.Namespace) -> int:
     async with DataQuery(args.env_file) as dq:
-        result = await dq.text_search_async(args.query)
+        result = await dq.search_async(args.query)
 
     if args.json:
         print(json.dumps(result, indent=2))
@@ -734,7 +734,7 @@ _ASYNC_COMMANDS = {
     "availability": cmd_availability,
     "download": cmd_download,
     "download-group": cmd_download_group,
-    "text-search": cmd_text_search,
+    "search": cmd_search,
     "groups-search": cmd_groups_search,
     "instruments": cmd_instruments,
     "instruments-search": cmd_instruments_search,

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery/config/env.py

@@ -247,7 +247,7 @@ class EnvConfig:
         if config.oauth_enabled:
             if not config.client_id:
                 errors.append("CLIENT_ID is required when OAuth is enabled")
-            if not config.client_secret:
+            if not config.get_client_secret():
                 errors.append("CLIENT_SECRET is required when OAuth is enabled")
             if not config.oauth_token_url:
                 errors.append("OAUTH_TOKEN_URL is required when OAuth is enabled")

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery/constants/download.py

@@ -28,6 +28,13 @@ MBPS_TO_BYTES_PER_SECOND = 125_000
 
 # Default size, in days, of each chunk when splitting a date range so that
 # the available-files endpoint is queried over a smaller window per call.
-# Used by ``run_group_download_chunked_async`` to honour group-level
-# date-range limits (e.g. RESEARCH_EQUITY_ALL is restricted to 14 days).
+# Used by ``split_date_range_into_chunks``. The endpoint caps each call at
+# one calendar month; smaller windows just keep responses lighter. (The old
+# group-level limits, e.g. 14 days for RESEARCH_EQUITY_ALL, were lifted
+# server-side.)
 DEFAULT_WRITTEN_RESEARCH_CHUNK_DAYS: int = 7
+
+# Error string set on the OperationReport when the available-files endpoint
+# returns nothing for a date window. Chunked workflows match on it to tell
+# a quiet window apart from a real failure.
+NO_FILES_FOUND_ERROR: str = "No available files found for date range"

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery/core/_mixins.py

@@ -609,7 +609,7 @@ class GridMixin(_RequestProto):
 class SearchMixin(_RequestProto):
     """Natural-language catalog search via POST /search."""
 
-    async def text_search_async(self, query: str) -> Dict[str, Any]:
+    async def search_async(self, query: str) -> Dict[str, Any]:
         """Search the DataQuery catalog using a natural-language query.
 
         POSTs ``{"query": query}`` to ``/search`` and returns the parsed JSON body.

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery/core/client.py

@@ -181,7 +181,7 @@ class DataQueryClient(
             raise ConfigurationError("Invalid base_url format")
 
         if strict_oauth_check and self.config.oauth_enabled:
-            if not self.config.client_id or not self.config.client_secret:
+            if not self.config.client_id or not self.config.get_client_secret():
                 raise ConfigurationError("client_id and client_secret are required")
 
         if not self.auth_manager.is_authenticated():
@@ -494,33 +494,28 @@ class DataQueryClient(
         if self.session is None:
             raise NetworkError("Failed to establish connection")
 
-        try:
-            response = await self.session.request(method, url, **kwargs)
-
-            # Raise 429 inside the retry scope so the retry manager backs off and retries.
-            if response.status == 429:
-                self.rate_limiter.handle_rate_limit_response(dict(response.headers))
-                retry_after = self._parse_retry_after(response.headers)
-                # Drain body so the connection can be reused; ignore decode errors.
-                try:
-                    await response.text()
-                except (UnicodeDecodeError, aiohttp.ClientPayloadError):
-                    pass
-                raise RateLimitError(f"Rate limit exceeded: {response.status}", retry_after=retry_after)
-
-            if response.status >= 500:
-                # Drain body so the connection can be reused; ignore decode errors.
-                try:
-                    await response.text()
-                except (UnicodeDecodeError, aiohttp.ClientPayloadError):
-                    pass
-                raise NetworkError(f"Server error: {response.status}", status_code=response.status)
+        response = await self.session.request(method, url, **kwargs)
 
-            return response
-        except (NetworkError, RateLimitError):
-            raise
-        except Exception:
-            raise
+        # Raise 429 inside the retry scope so the retry manager backs off and retries.
+        if response.status == 429:
+            self.rate_limiter.handle_rate_limit_response(dict(response.headers))
+            retry_after = self._parse_retry_after(response.headers)
+            # Drain body so the connection can be reused; ignore decode errors.
+            try:
+                await response.text()
+            except (UnicodeDecodeError, aiohttp.ClientPayloadError):
+                pass
+            raise RateLimitError(f"Rate limit exceeded: {response.status}", retry_after=retry_after)
+
+        if response.status >= 500:
+            # Drain body so the connection can be reused; ignore decode errors.
+            try:
+                await response.text()
+            except (UnicodeDecodeError, aiohttp.ClientPayloadError):
+                pass
+            raise NetworkError(f"Server error: {response.status}", status_code=response.status)
+
+        return response
 
     async def list_groups_async(self, limit: Optional[int] = None) -> List[Group]:
         """

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery/dataquery.py

@@ -10,12 +10,14 @@ import time
 from calendar import monthrange
 from datetime import date, datetime
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
+from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Tuple, Union
 
 import structlog
 from dotenv import load_dotenv
+from pydantic import SecretStr
 
 from .config import EnvConfig
+from .constants.download import NO_FILES_FOUND_ERROR
 from .core._sync import SyncRunner
 from .core.client import DataQueryClient
 from .types.exceptions import ConfigurationError
@@ -186,7 +188,7 @@ class DataQuery:
             if client_id:
                 self.client_config.client_id = client_id
             if client_secret:
-                self.client_config.client_secret = client_secret
+                self.client_config.client_secret = SecretStr(client_secret)
             if not self.client_config.oauth_token_url and self.client_config.base_url:
                 self.client_config.oauth_token_url = f"{self.client_config.base_url.rstrip('/')}/oauth/token"
 
@@ -717,7 +719,7 @@ class DataQuery:
         client = self._ensure_client()
         return await client.get_grid_data_async(expr, grid_id, date)
 
-    async def text_search_async(self, query: str) -> Dict[str, Any]:
+    async def search_async(self, query: str) -> Dict[str, Any]:
         """Search the DataQuery catalog using a natural-language query.
 
         Args:
@@ -728,7 +730,7 @@ class DataQuery:
         """
         await self.connect_async()
         client = self._ensure_client()
-        return await client.text_search_async(query)
+        return await client.search_async(query)
 
     async def run_groups_async(self, max_concurrent: int = 5) -> OperationReport:
         """Run complete operation for listing all groups."""
@@ -907,6 +909,7 @@ class DataQuery:
         delay_between_downloads: float = 0.2,
         max_retries: int = 3,
         file_group_id: Optional[Union[str, List[str]]] = None,
+        on_file_complete: Optional[Callable[["DownloadResult"], Awaitable[None]]] = None,
     ) -> OperationReport:
         """
         Download all files in a group for a date range using parallel HTTP range requests.
@@ -929,6 +932,11 @@ class DataQuery:
             file_group_id: Optional restriction to specific file-group-id(s). Accepts a
                 single id or a list of ids. When a list is supplied, availability is
                 queried in parallel per id and the union of dates is downloaded.
+            on_file_complete: Optional async callback awaited for each file as soon as
+                it finishes downloading (status ``completed``/``already_exists``). Runs
+                concurrently with the downloads still in flight, so post-processing such
+                as unzipping overlaps with subsequent downloads rather than waiting for
+                the whole batch.
 
         Returns:
             Dictionary with download results and statistics
@@ -1006,7 +1014,7 @@ class DataQuery:
                 total_time_minutes = total_time_seconds / 60.0
 
                 logger.warning(
-                    "No available files found for date range",
+                    NO_FILES_FOUND_ERROR,
                     group_id=group_id,
                     start_date=start_date,
                     end_date=end_date,
@@ -1014,7 +1022,7 @@ class DataQuery:
                 return OperationReport(
                     operation="group_download",
                     status="error",
-                    error="No available files found for date range",
+                    error=NO_FILES_FOUND_ERROR,
                     subject={"group_id": group_id, "start_date": start_date, "end_date": end_date},
                     counts={"total_files": 0, "successful_downloads": 0, "failed_downloads": 0},
                     timing={
@@ -1071,6 +1079,7 @@ class DataQuery:
                 base_retry_delay=delay_between_downloads,
                 max_retries=max_retries,
                 progress_callback=progress_callback,
+                on_file_complete=on_file_complete,
             )
 
             operation_end_time = time.time()
@@ -2047,9 +2056,9 @@ class DataQuery:
         """
         return self._run_sync(self.get_grid_data_async(expr, grid_id, date))
 
-    def text_search(self, query: str) -> Dict[str, Any]:
-        """Synchronous wrapper for :meth:`text_search_async`."""
-        return self._run_sync(self.text_search_async(query))
+    def search(self, query: str) -> Dict[str, Any]:
+        """Synchronous wrapper for :meth:`search_async`."""
+        return self._run_sync(self.search_async(query))
 
     def run_groups(self, max_concurrent: int = 5) -> OperationReport:
         """Synchronous wrapper for run_groups_async."""

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery/download/parallel.py

@@ -16,7 +16,7 @@ import time
 from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
-from typing import IO, TYPE_CHECKING, Callable, Optional
+from typing import IO, TYPE_CHECKING, Awaitable, Callable, Optional
 
 import aiohttp
 import structlog
@@ -530,6 +530,7 @@ async def _download_one_with_stagger(
     global_semaphore: asyncio.Semaphore,
     delay_seconds: float,
     progress_callback: Optional[Callable],
+    on_file_complete: Optional[Callable[[DownloadResult], Awaitable[None]]] = None,
 ) -> Optional[DownloadResult]:
     file_group_id = _file_id(file_info)
     file_datetime = _file_dt(file_info)
@@ -554,6 +555,21 @@ async def _download_one_with_stagger(
             file_datetime=file_datetime,
             status=result.status.value if result else "failed",
         )
+        if (
+            on_file_complete is not None
+            and result is not None
+            and result.status.value in ("completed", "already_exists")
+        ):
+            # Post-processing (e.g. unzip) runs here so it overlaps with the
+            # downloads of other files still in flight under the gather().
+            try:
+                await on_file_complete(result)
+            except Exception as cb_err:  # pragma: no cover - defensive
+                logger.warning(
+                    "on_file_complete callback failed",
+                    file_group_id=file_group_id,
+                    error=str(cb_err),
+                )
         return result
     except Exception as e:
         logger.error(
@@ -596,6 +612,7 @@ async def download_files_with_retry(
     base_retry_delay: float,
     max_retries: int,
     progress_callback: Optional[Callable] = None,
+    on_file_complete: Optional[Callable[[DownloadResult], Awaitable[None]]] = None,
 ) -> tuple[list[DownloadResult], list[dict], int]:
     """Run a staggered, retrying batch of parallel-range downloads.
 
@@ -605,6 +622,10 @@ async def download_files_with_retry(
     spread the burst against the API rate limiter. Failures are retried up
     to ``max_retries`` times with exponential backoff (``base_retry_delay``).
 
+    When ``on_file_complete`` is provided it is awaited for each successfully
+    downloaded file as soon as that file finishes, so post-processing (e.g.
+    unzipping) overlaps with the downloads of files still in flight.
+
     Returns ``(successful, failed, retry_count)``.
     """
 
@@ -618,6 +639,7 @@ async def download_files_with_retry(
                 global_semaphore=global_semaphore,
                 delay_seconds=i * intelligent_delay,
                 progress_callback=progress_callback,
+                on_file_complete=on_file_complete,
             )
             for i, fi in enumerate(batch)
         ]

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery/sse/client.py

@@ -111,6 +111,7 @@ class SSEClient:
         params: Optional[dict] = None,
         event_id_store: Optional[SSEEventIdStore] = None,
         heartbeat_timeout: float = 0.0,
+        defer_event_id_persistence: bool = False,
     ):
         """
         Initialise the SSE client.
@@ -147,6 +148,16 @@ class SSEClient:
                     larger than the server's keep-alive interval. ``0`` (the
                     default) disables the watchdog and relies on the server
                     closing the stream cleanly.
+            defer_event_id_persistence: When ``True``, do not persist parsed
+                    event ids to ``event_id_store`` as they arrive. The owner
+                    (e.g. :class:`NotificationDownloadManager`) instead commits
+                    a low-water-mark only after the matching download has
+                    settled, so a crash mid-download replays the event instead
+                    of skipping it (at-least-once). ``event_id_store`` is still
+                    read once on construction to seed the first
+                    ``Last-Event-ID``. Defaults to ``False`` — persist at parse
+                    time, which is correct for standalone use where there is no
+                    download step to gate persistence on.
         """
         self.config = config
         self.auth_manager = auth_manager
@@ -158,6 +169,7 @@ class SSEClient:
         self.params = dict(params) if params else None
         self.event_id_store = event_id_store
         self.heartbeat_timeout = heartbeat_timeout
+        self._defer_event_id_persistence = defer_event_id_persistence
 
         self._running = False
         self._task: Optional[asyncio.Task] = None
@@ -481,7 +493,8 @@ class SSEClient:
                 event_id = field_value
                 if field_value and field_value.isdigit():
                     self._last_event_id = field_value
-                    self._persist_event_id(field_value)
+                    if not self._defer_event_id_persistence:
+                        self._persist_event_id(field_value)
             elif field_name == "retry":
                 try:
                     retry_ms = int(field_value)

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery/sse/subscriber.py

@@ -6,8 +6,10 @@ downloads files as notifications arrive. Each SSE event uses the standard
 SSE format with the event type in the ``event:`` field (e.g., ``file-updated``)
 and the file details in the JSON ``data:`` payload.
 
-The SSE ``id:`` field is persisted to disk so that cross-process replay via
-the ``last-event-id`` query parameter works correctly on reconnection.
+The SSE ``id:`` field is persisted to disk as a low-water-mark — advanced past
+an event only once its download has settled — so cross-process replay via the
+``last-event-id`` query parameter re-delivers any file whose download did not
+complete before a restart (at-least-once delivery).
 """
 
 import asyncio
@@ -245,6 +247,13 @@ class NotificationDownloadManager:
         self._download_semaphore: Optional[asyncio.Semaphore] = None
         self._inflight: Set[asyncio.Task] = set()
 
+        # Replay low-water-mark. The persisted last-event-id is only advanced
+        # past an event once its download has settled, so a crash mid-download
+        # replays that event on restart instead of silently skipping it.
+        self._uncommitted_ids: Set[int] = set()
+        self._highest_seen_event_id: int = 0
+        self._committed_event_id: int = 0
+
         self.stats: Dict[str, Any] = {
             "start_time": None,
             "notifications_received": 0,
@@ -285,6 +294,12 @@ class NotificationDownloadManager:
             if self._event_id_store is not None:
                 stored_event_id = self._event_id_store.load()
 
+        if stored_event_id is not None and stored_event_id.isdigit():
+            # Seed the low-water-mark so we neither re-persist the stored id nor
+            # ever regress below it as newly received events settle.
+            self._committed_event_id = int(stored_event_id)
+            self._highest_seen_event_id = int(stored_event_id)
+
         if stored_event_id is not None:
             logger.info(
                 "Event replay enabled — resuming from event-id %s; skipping initial bulk check.",
@@ -307,6 +322,9 @@ class NotificationDownloadManager:
             params=self.subscription.query_params(),
             event_id_store=self._event_id_store,
             heartbeat_timeout=self._heartbeat_timeout,
+            # This manager commits a low-water-mark after downloads settle,
+            # so the client must not persist ids itself at parse time.
+            defer_event_id_persistence=True,
         )
         await self._sse_client.start()
 
@@ -393,6 +411,7 @@ class NotificationDownloadManager:
         if not self._running:
             return
 
+        self._register_event_id(event)
         task = asyncio.create_task(self._run_handler(event), name="sse-download")
         self._inflight.add(task)
         task.add_done_callback(self._inflight.discard)
@@ -407,6 +426,63 @@ class NotificationDownloadManager:
             except Exception as exc:
                 logger.error("Error handling notification: %s", exc)
                 await self._dispatch_error(exc)
+        # The download has settled (success, skip, or handled failure); advance
+        # the replay cursor. Deliberately not in a ``finally``: if the task is
+        # cancelled mid-download the cursor stays behind this event so it is
+        # replayed on restart rather than skipped.
+        await self._commit_settled_event(event)
+
+    def _register_event_id(self, event: SSEEvent) -> None:
+        """Record a numeric event id as in-flight before its download starts.
+
+        The replay cursor is never advanced past an in-flight id, so a crash
+        before the download settles leaves the event to be replayed.
+        """
+        if self._event_id_store is None:
+            return
+        eid = event.id
+        if not eid or not eid.isdigit():
+            return
+        value = int(eid)
+        self._uncommitted_ids.add(value)
+        if value > self._highest_seen_event_id:
+            self._highest_seen_event_id = value
+
+    def _next_commit_id(self) -> Optional[str]:
+        """Return the highest replay id now safe to persist, or ``None``.
+
+        The safe id is the low-water-mark: everything below the oldest
+        still-in-flight event has settled, so commit ``min(in-flight) - 1``
+        (or the highest id seen when nothing is in flight). Fully synchronous
+        so two concurrent settlers can't persist a decreasing watermark across
+        the ``await`` in :meth:`_commit_settled_event`.
+        """
+        if self._uncommitted_ids:
+            watermark = min(self._uncommitted_ids) - 1
+        else:
+            watermark = self._highest_seen_event_id
+        if watermark > self._committed_event_id:
+            self._committed_event_id = watermark
+            return str(watermark)
+        return None
+
+    async def _commit_settled_event(self, event: SSEEvent) -> None:
+        """Advance the persisted replay cursor after an event's handler settles.
+
+        Runs whether the download succeeded, was skipped, or failed — the id is
+        no longer in-flight either way. A terminal download failure is already
+        surfaced via ``error_callback``/``stats``; the cursor still advances so a
+        genuinely unfetchable file can't stall replay forever.
+        """
+        store = self._event_id_store
+        if store is None:
+            return
+        eid = event.id
+        if eid and eid.isdigit():
+            self._uncommitted_ids.discard(int(eid))
+        to_save = self._next_commit_id()
+        if to_save is not None:
+            await store.save(to_save)
 
     async def _on_sse_error(self, exc: Exception) -> None:
         """Called by SSEClient on connection errors."""

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery/transport/auth.py

@@ -72,7 +72,7 @@ class TokenManager:
             Bearer token string or None if no valid token available
         """
         if self.config.has_bearer_token:
-            return f"Bearer {self.config.bearer_token}"
+            return f"Bearer {self.config.get_bearer_token()}"
 
         if not self.config.has_oauth_credentials:
             logger.warning("No OAuth credentials or bearer token configured")
@@ -111,7 +111,7 @@ class TokenManager:
         if not self.config.oauth_token_url:
             raise ConfigurationError("OAuth token URL not configured")
 
-        if not self.config.client_id or not self.config.client_secret:
+        if not self.config.client_id or not self.config.get_client_secret():
             raise ConfigurationError("client_id and client_secret are required for OAuth")
 
         token_request = TokenRequest(
@@ -172,7 +172,7 @@ class TokenManager:
                 "grant_type": "refresh_token",
                 "refresh_token": self.current_token.refresh_token,
                 "client_id": self.config.client_id,
-                "client_secret": self.config.client_secret,
+                "client_secret": self.config.get_client_secret(),
             }
 
             if not self.config.oauth_token_url:

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery/types/models.py

@@ -5,16 +5,30 @@ Data models for the DATAQUERY SDK based on the OpenAPI specification.
 import asyncio
 import base64
 import time
-from datetime import date, datetime, timedelta
+from datetime import date, datetime, timedelta, timezone
 from enum import Enum
 from pathlib import Path
 from typing import Any, Dict, List, Literal, Optional, Union, cast
 
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator
 
 from .. import constants as C
 
 
+def _reveal_secret(value: Union[str, SecretStr, None]) -> Optional[str]:
+    """Return the plain-text value of a secret field, or ``None``.
+
+    Accepts a ``SecretStr`` (the normal case), a bare ``str`` (e.g. assigned
+    directly to a config attribute, which bypasses pydantic's coercion), or
+    ``None`` — so callers never have to care which form is stored.
+    """
+    if value is None:
+        return None
+    if isinstance(value, SecretStr):
+        return value.get_secret_value()
+    return value
+
+
 class DownloadStatus(str, Enum):
     """Status of a download operation."""
 
@@ -59,14 +73,14 @@ class ClientConfig(BaseModel):
         description="OAuth token endpoint URL",
     )
     client_id: Optional[str] = Field(default=None, description="OAuth client ID")
-    client_secret: Optional[str] = Field(default=None, description="OAuth client secret")
+    client_secret: Optional[SecretStr] = Field(default=None, description="OAuth client secret")
     aud: Optional[str] = Field(
         default="JPMC:URI:RS-06785-DataQueryExternalApi-PROD",
         description="OAuth audience (aud)",
     )
     grant_type: str = Field(default="client_credentials", description="OAuth grant type")
 
-    bearer_token: Optional[str] = Field(default=None, description="Bearer token for API access")
+    bearer_token: Optional[SecretStr] = Field(default=None, description="Bearer token for API access")
     token_refresh_threshold: int = Field(default=300, description="Seconds before expiry to refresh token")
 
     timeout: float = Field(default=600.0, description="Default request timeout in seconds")
@@ -86,7 +100,7 @@ class ClientConfig(BaseModel):
         description="Proxy URL (e.g., http://proxy:8080, socks5://proxy:1080)",
     )
     proxy_username: Optional[str] = Field(default="", description="Proxy username for authentication")
-    proxy_password: Optional[str] = Field(default="", description="Proxy password for authentication")
+    proxy_password: Optional[SecretStr] = Field(default=None, description="Proxy password for authentication")
     proxy_verify_ssl: bool = Field(default=True, description="Verify SSL certificates for proxy connections")
 
     log_level: str = Field(default="INFO", description="Logging level")
@@ -205,12 +219,25 @@ class ClientConfig(BaseModel):
     @property
     def has_bearer_token(self) -> bool:
         """Check if bearer token is configured."""
-        return self.bearer_token is not None and self.bearer_token.strip() != ""
+        token = self.get_bearer_token()
+        return bool(token and token.strip())
 
     @property
     def has_proxy_credentials(self) -> bool:
         """Check if proxy credentials are configured."""
-        return bool(self.proxy_username) and bool(self.proxy_password)
+        return bool(self.proxy_username) and bool(self.get_proxy_password())
+
+    def get_client_secret(self) -> Optional[str]:
+        """Return the plain-text OAuth client secret, or ``None``."""
+        return _reveal_secret(self.client_secret)
+
+    def get_bearer_token(self) -> Optional[str]:
+        """Return the plain-text bearer token, or ``None``."""
+        return _reveal_secret(self.bearer_token)
+
+    def get_proxy_password(self) -> Optional[str]:
+        """Return the plain-text proxy password, or ``None``."""
+        return _reveal_secret(self.proxy_password)
 
     def get_proxy_kwargs(self) -> Dict[str, Any]:
         """Return aiohttp request/session kwargs for proxy routing.
@@ -227,7 +254,7 @@ class ClientConfig(BaseModel):
         kwargs: Dict[str, Any] = {"proxy": self.proxy_url}
         if self.has_proxy_credentials:
             # Encode Basic proxy auth ourselves: aiohttp.BasicAuth is removed in 4.0 (latin1).
-            raw = f"{self.proxy_username or ''}:{self.proxy_password or ''}".encode("latin1")
+            raw = f"{self.proxy_username or ''}:{self.get_proxy_password() or ''}".encode("latin1")
             token = base64.b64encode(raw).decode("ascii")
             kwargs["proxy_headers"] = {"Proxy-Authorization": f"Basic {token}"}
         return kwargs
@@ -263,9 +290,14 @@ class OAuthToken(BaseModel):
 
     @property
     def expires_at(self) -> Optional[datetime]:
-        """Get token expiry time."""
+        """Get token expiry time (timezone-aware, UTC)."""
         if self.expires_in and self.issued_at:
-            return self.issued_at + timedelta(seconds=self.expires_in)
+            issued = self.issued_at
+            if issued.tzinfo is None:
+                # Older SDK versions persisted naive local timestamps; interpret
+                # them in the system timezone before normalizing to UTC.
+                issued = issued.astimezone(timezone.utc)
+            return issued + timedelta(seconds=self.expires_in)
         return None
 
     @property
@@ -273,13 +305,13 @@ class OAuthToken(BaseModel):
         """Check if token is expired."""
         if not self.expires_at:
             return False
-        return datetime.now() >= self.expires_at
+        return datetime.now(timezone.utc) >= self.expires_at
 
     def is_expiring_soon(self, threshold: int = 300) -> bool:
         """Check if token is expiring soon."""
         if not self.expires_at:
             return False
-        remaining_time = (self.expires_at - datetime.now()).total_seconds()
+        remaining_time = (self.expires_at - datetime.now(timezone.utc)).total_seconds()
         if self.expires_in and threshold > self.expires_in:
             return False
         return remaining_time < threshold
@@ -301,7 +333,7 @@ class TokenRequest(BaseModel):
 
     grant_type: str = Field("client_credentials", description="OAuth grant type")
     client_id: Optional[str] = Field(None, description="OAuth client ID")
-    client_secret: Optional[str] = Field(None, description="OAuth client secret")
+    client_secret: Optional[SecretStr] = Field(None, description="OAuth client secret")
     aud: Optional[str] = Field(None, description="OAuth audience (aud)")
 
     model_config = ConfigDict(extra="allow", populate_by_name=True)
@@ -314,7 +346,7 @@ class TokenRequest(BaseModel):
         if self.client_id:
             data["client_id"] = cast(str, self.client_id)
         if self.client_secret:
-            data["client_secret"] = cast(str, self.client_secret)
+            data["client_secret"] = self.client_secret.get_secret_value()
         if getattr(self, "aud", None):
             data["aud"] = cast(str, self.aud)
         return data
@@ -337,7 +369,7 @@ class TokenResponse(BaseModel):
             token_type=(self.token_type or "Bearer"),
             expires_in=self.expires_in,
             refresh_token=self.refresh_token,
-            issued_at=datetime.now(),
+            issued_at=datetime.now(timezone.utc),
         )
 
 

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery/utils.py

@@ -2,22 +2,26 @@
 Utility functions for the DATAQUERY SDK.
 """
 
+import asyncio
 import os
 import re
 import urllib.parse
-from datetime import datetime, timedelta
+import zipfile
+from calendar import monthrange
+from datetime import date, datetime, timedelta
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
 import aiohttp
 import structlog
 
-from .constants.download import DEFAULT_WRITTEN_RESEARCH_CHUNK_DAYS
+from .constants.download import DEFAULT_WRITTEN_RESEARCH_CHUNK_DAYS, NO_FILES_FOUND_ERROR
 from .types.exceptions import ValidationError
 from .types.models import ClientConfig
 
 if TYPE_CHECKING:
     from .dataquery import DataQuery
+    from .types.models import DownloadResult
 
 logger = structlog.get_logger(__name__)
 
@@ -174,8 +178,8 @@ DATAQUERY_DEFAULT_DIR=files
 # =============================================================================
 
 # User agent string for HTTP requests
-# Default: DATAQUERY-SDK/1.2.0
-DATAQUERY_USER_AGENT=DATAQUERY-SDK/1.2.0
+# Default: DATAQUERY-SDK/1.2.1
+DATAQUERY_USER_AGENT=DATAQUERY-SDK/1.2.1
 
 # Enable HTTP/2 support (true/false)
 # Default: true
@@ -252,12 +256,12 @@ DATAQUERY_BASE_URL={config.base_url}
 DATAQUERY_OAUTH_ENABLED={str(config.oauth_enabled).lower()}
 DATAQUERY_OAUTH_TOKEN_URL={config.oauth_token_url or ""}
 DATAQUERY_CLIENT_ID={config.client_id or ""}
-DATAQUERY_CLIENT_SECRET={config.client_secret or ""}
+DATAQUERY_CLIENT_SECRET={config.get_client_secret() or ""}
 DATAQUERY_OAUTH_AUD={getattr(config, "aud", "") or ""}
 DATAQUERY_GRANT_TYPE={config.grant_type}
 
 # Bearer Token Configuration
-DATAQUERY_BEARER_TOKEN={config.bearer_token or ""}
+DATAQUERY_BEARER_TOKEN={config.get_bearer_token() or ""}
 DATAQUERY_TOKEN_REFRESH_THRESHOLD={config.token_refresh_threshold}
 
 # HTTP Configuration
@@ -712,6 +716,287 @@ def split_date_range_into_chunks(
     return chunks
 
 
+def _split_into_monthly_ranges(from_date: str, to_date: str) -> List[Tuple[str, str]]:
+    """Split an inclusive YYYYMMDD date range into calendar-month windows.
+
+    Each returned window spans at most one calendar month, matching the
+    ``available-files`` endpoint's one-month range limit.
+
+    Args:
+        from_date: Inclusive start date in YYYYMMDD format.
+        to_date: Inclusive end date in YYYYMMDD format.
+
+    Returns:
+        Ordered list of ``(chunk_start, chunk_end)`` YYYYMMDD tuples covering
+        ``[from_date, to_date]`` with no gaps or overlaps.
+    """
+    try:
+        start = datetime.strptime(from_date, "%Y%m%d").date()
+        end = datetime.strptime(to_date, "%Y%m%d").date()
+    except ValueError as exc:
+        raise ValidationError(f"Invalid date format (expected YYYYMMDD): {exc}")
+
+    if end < start:
+        raise ValidationError("end_date must be on or after start_date")
+
+    ranges: List[Tuple[str, str]] = []
+    current = start
+    while current <= end:
+        last_day = monthrange(current.year, current.month)[1]
+        month_end = date(current.year, current.month, last_day)
+        chunk_end = min(month_end, end)
+        ranges.append((current.strftime("%Y%m%d"), chunk_end.strftime("%Y%m%d")))
+        if current.month == 12:
+            current = date(current.year + 1, 1, 1)
+        else:
+            current = date(current.year, current.month + 1, 1)
+    return ranges
+
+
+def _extract_zip_safely(zip_path: Path, target_dir: Path) -> List[str]:
+    """Extract a ZIP archive into ``target_dir``, guarding against Zip Slip.
+
+    Args:
+        zip_path: Path to the ``.zip`` archive to extract.
+        target_dir: Directory the archive contents are written to.
+
+    Returns:
+        List of member names that were extracted.
+
+    Raises:
+        ValidationError: If any archive member would be written outside
+            ``target_dir`` (path traversal attempt).
+    """
+    target_dir.mkdir(parents=True, exist_ok=True)
+    resolved_target = target_dir.resolve()
+
+    with zipfile.ZipFile(zip_path) as zf:
+        members = zf.namelist()
+        for member in members:
+            member_path = (target_dir / member).resolve()
+            if resolved_target not in member_path.parents and member_path != resolved_target:
+                raise ValidationError(f"Unsafe path in archive '{zip_path.name}': {member}")
+        zf.extractall(target_dir)
+    return members
+
+
+def _extract_single_zip(
+    zip_path: Path,
+    today: str,
+    remove_zip_after_extract: bool,
+) -> Optional[Dict[str, Any]]:
+    """Extract one top-level archive, returning a record or ``None`` if skipped.
+
+    Leaves current-day archives untouched (their content may still be updating)
+    and silently skips archives that are not valid ZIP files.
+    """
+    # The embedded file-datetime (YYYYMMDDThhmm[ss]) is the last such token in
+    # the name; earlier digit runs (e.g. the GPS id) must not be matched.
+    matches = re.findall(r"(\d{8})T\d{4,6}", zip_path.stem)
+    if matches and matches[-1] == today:
+        logger.info("Skipping current-day zip archive", zip=str(zip_path))
+        return None
+
+    target_dir = zip_path.parent
+    try:
+        members = _extract_zip_safely(zip_path, target_dir)
+    except zipfile.BadZipFile:
+        logger.warning("Skipping invalid zip archive", zip=str(zip_path))
+        return None
+
+    if remove_zip_after_extract:
+        zip_path.unlink()
+
+    logger.info("Extracted zip archive", zip=str(zip_path), files=len(members))
+    return {
+        "zip": str(zip_path),
+        "target_dir": str(target_dir),
+        "files": members,
+    }
+
+
+async def _run_download_over_ranges(
+    dq: "DataQuery",
+    group_id: str,
+    ranges: List[Tuple[str, str]],
+    **kwargs: Any,
+) -> Tuple[Dict[str, int], List[Dict[str, Any]]]:
+    """Run :meth:`DataQuery.run_group_download_async` once per date window.
+
+    Shared by :func:`download_zip_async` and
+    :func:`run_group_download_chunked_async`.
+
+    Returns:
+        ``(totals, per_chunk)`` where ``totals`` aggregates the download counts
+        across all windows and ``per_chunk`` holds one summary dict (dates,
+        status, error, counts) per window, in order.
+    """
+    totals = {"total_files": 0, "successful_downloads": 0, "failed_downloads": 0}
+    per_chunk: List[Dict[str, Any]] = []
+
+    for chunk_start, chunk_end in ranges:
+        report = await dq.run_group_download_async(
+            group_id=group_id,
+            start_date=chunk_start,
+            end_date=chunk_end,
+            **kwargs,
+        )
+        counts = report.counts or {}
+        for key in totals:
+            totals[key] += int(counts.get(key, 0) or 0)
+        per_chunk.append(
+            {
+                "start_date": chunk_start,
+                "end_date": chunk_end,
+                "status": report.status,
+                "error": report.error,
+                "counts": counts,
+            }
+        )
+
+    return totals, per_chunk
+
+
+def _aggregate_chunk_status(per_chunk: List[Dict[str, Any]]) -> str:
+    """Combine per-window statuses into one ``success``/``partial``/``error``.
+
+    Windows that merely had no available files (a quiet window, reported by
+    the API as an error) are ignored so they don't mark an otherwise clean
+    multi-window run as ``partial``.
+    """
+    relevant = [
+        c["status"] for c in per_chunk if not (c["status"] == "error" and c.get("error") == NO_FILES_FOUND_ERROR)
+    ]
+    if not relevant or all(s == "success" for s in relevant):
+        return "success"
+    if all(s == "error" for s in relevant):
+        return "error"
+    return "partial"
+
+
+async def download_zip_async(
+    dq: "DataQuery",
+    group_id: str,
+    start_date: str,
+    end_date: str,
+    destination_dir: Path = Path("./downloads"),
+    remove_zip_after_extract: bool = True,
+    **kwargs: Any,
+) -> Dict[str, Any]:
+    """Download all files in a group and unzip any ZIP archives as they arrive.
+
+    Runs :meth:`DataQuery.run_group_download_async` over the
+    ``[start_date, end_date]`` range, splitting it into calendar-month windows
+    because the ``available-files`` endpoint rejects ranges longer than one
+    month. Each ``.zip`` archive is extracted in a worker thread the moment its
+    download completes, so unzipping overlaps with the downloads of the files
+    still in flight instead of waiting for the whole batch to finish. A final
+    sweep extracts any archives not already handled by the per-file hook.
+
+    Note that with ``remove_zip_after_extract=True`` the archives are gone
+    after the run, so re-running the same range downloads them again (the
+    downloader can't see they were already fetched).
+
+    Args:
+        dq: An already-constructed ``DataQuery`` instance (used inside its own
+            ``async with`` context by the caller).
+        group_id: Group ID to download files from.
+        start_date: Inclusive start date in YYYYMMDD format.
+        end_date: Inclusive end date in YYYYMMDD format.
+        destination_dir: Base destination directory for downloads. Files are
+            written to ``destination_dir / group_id``.
+        remove_zip_after_extract: When True (default), delete each ``.zip``
+            archive after it has been successfully extracted.
+        **kwargs: Forwarded as-is to ``DataQuery.run_group_download_async``.
+
+    Returns:
+        Dict with the aggregated ``status``/``counts``, a ``chunks`` list of
+        per-window summaries, an ``extracted`` list describing each archive
+        that was unzipped, and an ``extraction_errors`` list for archives that
+        downloaded fine but could not be extracted (any entry here downgrades
+        an overall ``success`` to ``partial``).
+    """
+    ranges = _split_into_monthly_ranges(start_date, end_date)
+
+    today = datetime.now().strftime("%Y%m%d")
+    extracted: List[Dict[str, Any]] = []
+    extraction_errors: List[Dict[str, str]] = []
+    processed: set[Path] = set()
+
+    async def _on_file_complete(result: "DownloadResult") -> None:
+        local_path = result.local_path
+        if not local_path or local_path.suffix.lower() != ".zip":
+            return
+        key = local_path.resolve()
+        if key in processed:
+            return
+        processed.add(key)
+        # Run the blocking unzip off the event loop so concurrent downloads
+        # keep progressing while this archive is extracted.
+        try:
+            record = await asyncio.to_thread(_extract_single_zip, local_path, today, remove_zip_after_extract)
+        except Exception as exc:
+            # Leave the archive for the final sweep to retry; only the sweep
+            # records a definitive extraction error.
+            processed.discard(key)
+            logger.warning(
+                "Extraction failed, will retry in final sweep",
+                zip=str(local_path),
+                error=str(exc),
+            )
+            return
+        if record is not None:
+            extracted.append(record)
+
+    totals, per_chunk = await _run_download_over_ranges(
+        dq,
+        group_id,
+        ranges,
+        destination_dir=destination_dir,
+        on_file_complete=_on_file_complete,
+        **kwargs,
+    )
+    status = _aggregate_chunk_status(per_chunk)
+
+    group_dir = destination_dir / group_id
+
+    # Fallback sweep: extract any archives the per-file hook did not handle
+    # (e.g. files whose local path was unavailable, or first-attempt extraction
+    # failures). Already-extracted zips were removed/deduplicated, so this is
+    # idempotent.
+    if group_dir.is_dir():
+        # Snapshot the top-level archives before extracting so nested zips
+        # unpacked from these archives are left untouched (no recursion).
+        top_level_zips = sorted(group_dir.glob("*.zip"))
+        for zip_path in top_level_zips:
+            key = zip_path.resolve()
+            if key in processed:
+                continue
+            processed.add(key)
+            try:
+                record = await asyncio.to_thread(_extract_single_zip, zip_path, today, remove_zip_after_extract)
+            except Exception as exc:
+                logger.error("Extraction failed", zip=str(zip_path), error=str(exc))
+                extraction_errors.append({"zip": str(zip_path), "error": str(exc)})
+                continue
+            if record is not None:
+                extracted.append(record)
+
+    if extraction_errors and status == "success":
+        status = "partial"
+
+    return {
+        "group_id": group_id,
+        "start_date": start_date,
+        "end_date": end_date,
+        "status": status,
+        "counts": totals,
+        "chunks": per_chunk,
+        "extracted": extracted,
+        "extraction_errors": extraction_errors,
+    }
+
+
 async def run_group_download_chunked_async(
     dq: "DataQuery",
     group_id: str,
@@ -742,28 +1027,7 @@ async def run_group_download_chunked_async(
         ``OperationReport`` dumps.
     """
     chunks = split_date_range_into_chunks(start_date, end_date, chunk_days)
-
-    totals = {"total_files": 0, "successful_downloads": 0, "failed_downloads": 0}
-    per_chunk: List[Dict[str, Any]] = []
-
-    for chunk_start, chunk_end in chunks:
-        report = await dq.run_group_download_async(
-            group_id=group_id,
-            start_date=chunk_start,
-            end_date=chunk_end,
-            **kwargs,
-        )
-        counts = report.counts or {}
-        for key in totals:
-            totals[key] += int(counts.get(key, 0) or 0)
-        per_chunk.append(
-            {
-                "start_date": chunk_start,
-                "end_date": chunk_end,
-                "status": report.status,
-                "counts": counts,
-            }
-        )
+    totals, per_chunk = await _run_download_over_ranges(dq, group_id, chunks, **kwargs)
 
     return {
         "group_id": group_id,

{dataquery_sdk-1.2.0 → dataquery_sdk-1.2.1}/dataquery_sdk.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataquery-sdk
-Version: 1.2.0
+Version: 1.2.1
 Summary: Python SDK for DATAQUERY Data API - Query, download, and check availability of economic data files
 Author-email: DATAQUERY SDK Team <dataquery_support@jpmorgan.com>
 Project-URL: Homepage, https://github.com/dataquery/dataquery-sdk