dataquery-sdk 0.0.7__py3-none-any.whl

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/dataquery.py ADDED
@@ -0,0 +1,2467 @@
1
+ """
2
+ Main DataQuery class for the DATAQUERY SDK.
3
+
4
+ This module provides the main DataQuery class that serves as the primary interface
5
+ for all API interactions, encapsulating the client and providing high-level operations.
6
+ """
7
+
8
+ import asyncio
9
+ import time
10
+ from pathlib import Path
11
+ from typing import Any, Callable, Dict, List, Optional, Union
12
+
13
+ import structlog
14
+ from dotenv import load_dotenv
15
+
16
+ from .client import DataQueryClient, format_file_size
17
+ from .config import EnvConfig
18
+ from .exceptions import ConfigurationError
19
+ from .models import (
20
+ AttributesResponse,
21
+ AvailabilityInfo,
22
+ ClientConfig,
23
+ DownloadOptions,
24
+ DownloadResult,
25
+ DownloadStatus,
26
+ FileInfo,
27
+ FiltersResponse,
28
+ GridDataResponse,
29
+ Group,
30
+ InstrumentsResponse,
31
+ TimeSeriesResponse,
32
+ )
33
+
34
+ # Load environment variables from .env file
35
+ load_dotenv()
36
+
37
+ logger = structlog.get_logger(__name__)
38
+
39
+
40
+ def setup_logging(log_level: str = "INFO") -> structlog.BoundLogger:
41
+ """Deprecated shim: prefer `LoggingManager`.
42
+
43
+ This function intentionally does not configure structlog. It returns a
44
+ namespaced logger tagged to indicate deprecated usage.
45
+ """
46
+ return structlog.get_logger(__name__).bind(
47
+ deprecated_setup_logging=True,
48
+ level=log_level,
49
+ )
50
+
51
+
52
+ """
53
+ Note: format_file_size, format_duration, and ensure_directory are imported from
54
+ utils to maintain a single source of truth for these helpers across the SDK.
55
+ """
56
+
57
+
58
+ def get_download_paths() -> Dict[str, Path]:
59
+ """Get download paths from environment variables with defaults."""
60
+ base_download_dir = EnvConfig.get_path("DOWNLOAD_DIR", "./downloads")
61
+
62
+ return {
63
+ "base": base_download_dir,
64
+ "workflow": base_download_dir
65
+ / (EnvConfig.get_env_var("WORKFLOW_DIR", "workflow") or "workflow"),
66
+ "groups": base_download_dir
67
+ / (EnvConfig.get_env_var("GROUPS_DIR", "groups") or "groups"),
68
+ "availability": base_download_dir
69
+ / (EnvConfig.get_env_var("AVAILABILITY_DIR", "availability") or "availability"),
70
+ "default": base_download_dir
71
+ / (EnvConfig.get_env_var("DEFAULT_DIR", "files") or "files"),
72
+ }
73
+
74
+
75
+ class ConfigManager:
76
+ """Configuration manager for DATAQUERY SDK."""
77
+
78
+ def __init__(self, env_file: Optional[Path] = None):
79
+ """
80
+ Initialize ConfigManager.
81
+
82
+ Args:
83
+ env_file: Optional path to .env file. If None, will look for .env in current directory.
84
+ """
85
+ self.env_file = env_file
86
+
87
+ def get_client_config(self) -> ClientConfig:
88
+ """Get client configuration from environment variables."""
89
+ try:
90
+ # Pass env_file as keyword argument to match expected signature
91
+ config = EnvConfig.create_client_config(env_file=self.env_file)
92
+ EnvConfig.validate_config(config)
93
+ return config
94
+ except Exception as e:
95
+ logger.warning(
96
+ "Failed to load configuration from environment", error=str(e)
97
+ )
98
+ return self._get_default_config()
99
+
100
+ def _get_default_config(self) -> ClientConfig:
101
+ """Get default configuration for examples."""
102
+ return ClientConfig(
103
+ base_url="https://api.dataquery.com",
104
+ # Set oauth_enabled False by default for fallback to satisfy tests
105
+ oauth_enabled=False,
106
+ # All other fields will use their default values from the model
107
+ )
108
+
109
+
110
+ class ProgressTracker:
111
+ """Progress tracking for batch operations."""
112
+
113
+ def __init__(self, log_interval: int = 10):
114
+ self.log_interval = log_interval
115
+ self.last_log_time = 0.0
116
+
117
+ def create_progress_callback(self) -> Callable:
118
+ """Create a progress callback function."""
119
+
120
+ def progress_callback(progress: Any):
121
+ current_time = time.time()
122
+ if current_time - self.last_log_time >= self.log_interval:
123
+ logger.info(
124
+ "Batch progress",
125
+ completed=getattr(progress, "completed_files", 0),
126
+ total=getattr(progress, "total_files", 0),
127
+ percentage=f"{getattr(progress, 'percentage', 0):.1f}%",
128
+ current_file=getattr(progress, "current_file", "unknown"),
129
+ )
130
+ self.last_log_time = current_time
131
+
132
+ return progress_callback
133
+
134
+
135
+ class DataQuery:
136
+ """
137
+ Main DataQuery class for all API interactions.
138
+
139
+ This class serves as the primary interface for the DATAQUERY SDK,
140
+ encapsulating the client and providing high-level operations for
141
+ listing, searching, downloading, and managing data files.
142
+
143
+ Supports both async and sync operations with proper event loop management.
144
+ """
145
+
146
+ def __init__(
147
+ self,
148
+ config_or_env_file: Optional[Union[ClientConfig, str, Path]] = None,
149
+ client_id: Optional[str] = None,
150
+ client_secret: Optional[str] = None,
151
+ **overrides: Any,
152
+ ):
153
+ """
154
+ Initialize DataQuery with configuration.
155
+
156
+ Args:
157
+ config_or_env_file: Either a ClientConfig object, a Path, or a str to .env file.
158
+ If None, will look for .env in current directory.
159
+ """
160
+ # Handle different input types
161
+ if isinstance(config_or_env_file, ClientConfig):
162
+ # Direct ClientConfig provided
163
+ self.client_config = config_or_env_file
164
+ else:
165
+ # env_file provided (Path, str, or None)
166
+ env_file = None
167
+ if isinstance(config_or_env_file, (str, Path)):
168
+ env_file = Path(config_or_env_file)
169
+ config_manager = ConfigManager(env_file)
170
+ self.client_config = config_manager.get_client_config()
171
+
172
+ # Apply default-first initialization pattern with optional overrides.
173
+ # Credentials are never defaulted; if provided, enable OAuth and set them.
174
+ if client_id or client_secret:
175
+ if client_id and not isinstance(client_id, str):
176
+ raise ConfigurationError("client_id must be a string")
177
+ if client_secret and not isinstance(client_secret, str):
178
+ raise ConfigurationError("client_secret must be a string")
179
+
180
+ self.client_config.oauth_enabled = True
181
+ if client_id:
182
+ self.client_config.client_id = client_id
183
+ if client_secret:
184
+ self.client_config.client_secret = client_secret
185
+ # Auto-derive token URL if missing
186
+ if not self.client_config.oauth_token_url and self.client_config.base_url:
187
+ try:
188
+ self.client_config.oauth_token_url = (
189
+ f"{self.client_config.base_url.rstrip('/')}/oauth/token"
190
+ )
191
+ except Exception:
192
+ pass
193
+
194
+ # Apply any non-credential overrides (e.g., base_url, context_path, files_base_url, etc.)
195
+ for key, value in (overrides or {}).items():
196
+ if key in {"client_id", "client_secret"}:
197
+ continue
198
+ if hasattr(self.client_config, key) and value is not None:
199
+ try:
200
+ setattr(self.client_config, key, value)
201
+ except Exception:
202
+ # Best-effort: ignore invalid overrides silently to preserve backward compatibility
203
+ pass
204
+
205
+ # Validate configuration
206
+ try:
207
+ EnvConfig.validate_config(self.client_config)
208
+ except Exception as e:
209
+ logger.error("Configuration validation failed", error=str(e))
210
+ raise ConfigurationError(f"Configuration validation failed: {e}")
211
+
212
+ self._client: Optional[DataQueryClient] = None
213
+ self._loop: Optional[asyncio.AbstractEventLoop] = None
214
+ self._own_loop: bool = False
215
+
216
+ async def __aenter__(self):
217
+ """Async context manager entry."""
218
+ await self.connect_async()
219
+ return self
220
+
221
+ def __enter__(self):
222
+ """Sync context manager entry."""
223
+ self.connect()
224
+ return self
225
+
226
+ async def __aexit__(self, exc_type, exc_val, exc_tb):
227
+ """Async context manager exit."""
228
+ await self.close_async()
229
+
230
+ async def connect_async(self):
231
+ """Connect to the API."""
232
+ if self._client is None:
233
+ self._client = DataQueryClient(self.client_config)
234
+ await self._client.connect()
235
+
236
+ async def close_async(self):
237
+ """Close the connection and cleanup resources."""
238
+ if self._client:
239
+ await self._client.close()
240
+ self._client = None
241
+
242
+ async def cleanup_async(self):
243
+ """Cleanup resources and ensure proper shutdown."""
244
+ await self.close_async()
245
+
246
+ # Force garbage collection to clean up any remaining references
247
+ import gc
248
+
249
+ gc.collect()
250
+
251
+ def _run_sync(self, coro):
252
+ """
253
+ Run an async coroutine in a new event loop.
254
+
255
+ Args:
256
+ coro: Coroutine to run
257
+
258
+ Returns:
259
+ Result of the coroutine
260
+ """
261
+ try:
262
+ return asyncio.run(coro)
263
+ except RuntimeError as e:
264
+ if "cannot run loop while another loop is running" in str(e):
265
+ raise RuntimeError(
266
+ "Cannot run sync method when an asyncio event loop is already running. "
267
+ "Use the async version of the method instead."
268
+ ) from e
269
+ raise
270
+
271
+ # Core API Methods
272
+
273
+ async def list_groups_async(self, limit: Optional[int] = 100) -> List[Group]:
274
+ """
275
+ List all available data groups with pagination support.
276
+
277
+ Args:
278
+ limit: Maximum number of groups to return (default: 100). If None, returns all groups.
279
+
280
+ Returns:
281
+ List of group information
282
+ """
283
+ await self.connect_async()
284
+
285
+ if limit is None:
286
+ # Fetch all groups using pagination
287
+ assert self._client is not None
288
+ return await self._client.list_all_groups_async()
289
+ else:
290
+ # Fetch limited number of groups
291
+ assert self._client is not None
292
+ return await self._client.list_groups_async(limit=limit)
293
+
294
+ async def search_groups_async(
295
+ self, keywords: str, limit: Optional[int] = 100, offset: Optional[int] = None
296
+ ) -> List[Group]:
297
+ """
298
+ Search groups by keywords.
299
+
300
+ Args:
301
+ keywords: Search keywords
302
+ limit: Maximum number of results to return (default: 100)
303
+ offset: Number of results to skip
304
+
305
+ Returns:
306
+ List of matching groups
307
+ """
308
+ await self.connect_async()
309
+ assert self._client is not None
310
+ return await self._client.search_groups_async(keywords, limit, offset)
311
+
312
+ async def list_files_async(
313
+ self, group_id: str, file_group_id: Optional[str] = None
314
+ ) -> List[FileInfo]:
315
+ """
316
+ List all files in a group.
317
+
318
+ Args:
319
+ group_id: Group ID to list files for
320
+ file_group_id: Optional specific file ID to filter by
321
+
322
+ Returns:
323
+ List of file information
324
+ """
325
+ await self.connect_async()
326
+ assert self._client is not None
327
+ file_list = await self._client.list_files_async(group_id, file_group_id)
328
+ return file_list.file_group_ids
329
+
330
+ async def check_availability_async(
331
+ self, file_group_id: str, file_datetime: str
332
+ ) -> AvailabilityInfo:
333
+ """
334
+ Check file availability for a specific datetime.
335
+
336
+ Args:
337
+ file_group_id: File ID to check availability for
338
+ file_datetime: File datetime in YYYYMMDD, YYYYMMDDTHHMM, or YYYYMMDDTHHMMSS format
339
+
340
+ Returns:
341
+ Availability response with status for the datetime
342
+ """
343
+ await self.connect_async()
344
+ assert self._client is not None
345
+ return await self._client.check_availability_async(file_group_id, file_datetime)
346
+
347
+ async def download_file_async(
348
+ self,
349
+ file_group_id: str,
350
+ file_datetime: Optional[str] = None,
351
+ destination_path: Optional[Path] = None,
352
+ options: Optional[DownloadOptions] = None,
353
+ num_parts: int = 5,
354
+ progress_callback: Optional[Callable] = None,
355
+ ) -> DownloadResult:
356
+ """
357
+ Download a specific file using parallel HTTP range requests.
358
+
359
+ Args:
360
+ file_group_id: File ID to download
361
+ file_datetime: Optional datetime of the file (YYYYMMDD, YYYYMMDDTHHMM, or YYYYMMDDTHHMMSS)
362
+ destination_path: Optional download destination directory. The filename will be extracted
363
+ from the Content-Disposition header in the response. If not provided,
364
+ uses the default download directory from configuration.
365
+ options: Download options
366
+ num_parts: Number of parallel parts to split the file into (default 5)
367
+ progress_callback: Optional progress callback function
368
+
369
+ Returns:
370
+ DownloadResult with download information
371
+ """
372
+ await self.connect_async()
373
+
374
+ if destination_path and options is None:
375
+ options = DownloadOptions(
376
+ destination_path=destination_path,
377
+ create_directories=True,
378
+ overwrite_existing=False,
379
+ chunk_size=8192,
380
+ max_retries=3,
381
+ retry_delay=1.0,
382
+ timeout=600.0,
383
+ enable_range_requests=True,
384
+ range_start=None,
385
+ range_end=None,
386
+ range_header=None,
387
+ show_progress=True,
388
+ progress_callback=None,
389
+ )
390
+
391
+ assert self._client is not None
392
+ # Pass parameters in correct order: file_group_id, file_datetime, options, num_parts, progress_callback
393
+ # Use default num_parts=5
394
+ return await self._client.download_file_async(
395
+ file_group_id, file_datetime, options, num_parts, progress_callback
396
+ )
397
+
398
+ async def list_available_files_async(
399
+ self,
400
+ group_id: str,
401
+ file_group_id: Optional[str] = None,
402
+ start_date: Optional[str] = None,
403
+ end_date: Optional[str] = None,
404
+ ) -> List[Dict[str, Any]]:
405
+ """
406
+ List available files by date range.
407
+
408
+ Args:
409
+ group_id: Group ID to list files for
410
+ file_group_id: Optional specific file ID to filter by
411
+ start_date: Optional start date in YYYYMMDD format
412
+ end_date: Optional end date in YYYYMMDD format
413
+
414
+ Returns:
415
+ List of available file information
416
+ """
417
+ await self.connect_async()
418
+ assert self._client is not None
419
+ return await self._client.list_available_files_async(
420
+ group_id, file_group_id, start_date, end_date
421
+ )
422
+
423
+ async def health_check_async(self) -> bool:
424
+ """
425
+ Check if the API is healthy.
426
+
427
+ Returns:
428
+ True if API is healthy
429
+ """
430
+ await self.connect_async()
431
+ assert self._client is not None
432
+ return await self._client.health_check_async()
433
+
434
+ # Instrument Collection Endpoints
435
+ async def list_instruments_async(
436
+ self,
437
+ group_id: str,
438
+ instrument_id: Optional[str] = None,
439
+ page: Optional[str] = None,
440
+ ) -> "InstrumentsResponse":
441
+ """
442
+ Request the complete list of instruments and identifiers for a given dataset.
443
+
444
+ Args:
445
+ group_id: Catalog data group identifier
446
+ instrument_id: Optional instrument identifier to filter results
447
+ page: Optional page token for pagination
448
+
449
+ Returns:
450
+ InstrumentsResponse containing the list of instruments
451
+ """
452
+ await self.connect_async()
453
+ assert self._client is not None
454
+ return await self._client.list_instruments_async(group_id, instrument_id, page)
455
+
456
+ async def search_instruments_async(
457
+ self, group_id: str, keywords: str, page: Optional[str] = None
458
+ ) -> "InstrumentsResponse":
459
+ """
460
+ Search within a dataset using keywords to create subsets of matching instruments.
461
+
462
+ Args:
463
+ group_id: Catalog data group identifier
464
+ keywords: Keywords to narrow scope of results
465
+ page: Optional page token for pagination
466
+
467
+ Returns:
468
+ InstrumentsResponse containing the matching instruments
469
+ """
470
+ await self.connect_async()
471
+ assert self._client is not None
472
+ return await self._client.search_instruments_async(group_id, keywords, page)
473
+
474
+ async def get_instrument_time_series_async(
475
+ self,
476
+ instruments: List[str],
477
+ attributes: List[str],
478
+ data: str = "REFERENCE_DATA",
479
+ format: str = "JSON",
480
+ start_date: Optional[str] = None,
481
+ end_date: Optional[str] = None,
482
+ calendar: str = "CAL_USBANK",
483
+ frequency: str = "FREQ_DAY",
484
+ conversion: str = "CONV_LASTBUS_ABS",
485
+ nan_treatment: str = "NA_NOTHING",
486
+ page: Optional[str] = None,
487
+ ) -> "TimeSeriesResponse":
488
+ """
489
+ Retrieve time-series data for explicit list of instruments and attributes using identifiers.
490
+
491
+ Args:
492
+ instruments: List of instrument identifiers
493
+ attributes: List of attribute identifiers
494
+ data: Data type (REFERENCE_DATA, NO_REFERENCE_DATA, ALL)
495
+ format: Response format (JSON)
496
+ start_date: Start date in YYYYMMDD or TODAY-Nx format
497
+ end_date: End date in YYYYMMDD or TODAY-Nx format
498
+ calendar: Calendar convention
499
+ frequency: Frequency convention
500
+ conversion: Conversion convention
501
+ nan_treatment: Missing data treatment
502
+ page: Optional page token for pagination
503
+
504
+ Returns:
505
+ TimeSeriesResponse containing the time series data
506
+ """
507
+ await self.connect_async()
508
+ assert self._client is not None
509
+ return await self._client.get_instrument_time_series_async(
510
+ instruments,
511
+ attributes,
512
+ data,
513
+ format,
514
+ start_date,
515
+ end_date,
516
+ calendar,
517
+ frequency,
518
+ conversion,
519
+ nan_treatment,
520
+ page,
521
+ )
522
+
523
+ async def get_expressions_time_series_async(
524
+ self,
525
+ expressions: List[str],
526
+ format: str = "JSON",
527
+ start_date: Optional[str] = None,
528
+ end_date: Optional[str] = None,
529
+ calendar: str = "CAL_USBANK",
530
+ frequency: str = "FREQ_DAY",
531
+ conversion: str = "CONV_LASTBUS_ABS",
532
+ nan_treatment: str = "NA_NOTHING",
533
+ data: str = "REFERENCE_DATA",
534
+ page: Optional[str] = None,
535
+ ) -> "TimeSeriesResponse":
536
+ """
537
+ Retrieve time-series data using an explicit list of traditional DataQuery expressions.
538
+
539
+ Args:
540
+ expressions: List of traditional DataQuery expressions
541
+ format: Response format (JSON)
542
+ start_date: Start date in YYYYMMDD or TODAY-Nx format
543
+ end_date: End date in YYYYMMDD or TODAY-Nx format
544
+ calendar: Calendar convention
545
+ frequency: Frequency convention
546
+ conversion: Conversion convention
547
+ nan_treatment: Missing data treatment
548
+ data: Data type (REFERENCE_DATA, NO_REFERENCE_DATA, ALL)
549
+ page: Optional page token for pagination
550
+
551
+ Returns:
552
+ TimeSeriesResponse containing the time series data
553
+ """
554
+ await self.connect_async()
555
+ assert self._client is not None
556
+ return await self._client.get_expressions_time_series_async(
557
+ expressions,
558
+ format,
559
+ start_date,
560
+ end_date,
561
+ calendar,
562
+ frequency,
563
+ conversion,
564
+ nan_treatment,
565
+ data,
566
+ page,
567
+ )
568
+
569
+ # Group Collection Additional Endpoints
570
+ async def get_group_filters_async(
571
+ self, group_id: str, page: Optional[str] = None
572
+ ) -> "FiltersResponse":
573
+ """
574
+ Request the unique list of filter dimensions that are available for a given dataset.
575
+
576
+ Args:
577
+ group_id: Catalog data group identifier
578
+ page: Optional page token for pagination
579
+
580
+ Returns:
581
+ FiltersResponse containing the available filters
582
+ """
583
+ await self.connect_async()
584
+ assert self._client is not None
585
+ return await self._client.get_group_filters_async(group_id, page)
586
+
587
+ async def get_group_attributes_async(
588
+ self,
589
+ group_id: str,
590
+ instrument_id: Optional[str] = None,
591
+ page: Optional[str] = None,
592
+ ) -> "AttributesResponse":
593
+ """
594
+ Request the unique list of analytic attributes for each instrument of a given dataset.
595
+
596
+ Args:
597
+ group_id: Catalog data group identifier
598
+ instrument_id: Optional instrument identifier to filter results
599
+ page: Optional page token for pagination
600
+
601
+ Returns:
602
+ AttributesResponse containing the attributes for each instrument
603
+ """
604
+ await self.connect_async()
605
+ assert self._client is not None
606
+ return await self._client.get_group_attributes_async(
607
+ group_id, instrument_id, page
608
+ )
609
+
610
+ async def get_group_time_series_async(
611
+ self,
612
+ group_id: str,
613
+ attributes: List[str],
614
+ filter: Optional[str] = None,
615
+ data: str = "REFERENCE_DATA",
616
+ format: str = "JSON",
617
+ start_date: Optional[str] = None,
618
+ end_date: Optional[str] = None,
619
+ calendar: str = "CAL_USBANK",
620
+ frequency: str = "FREQ_DAY",
621
+ conversion: str = "CONV_LASTBUS_ABS",
622
+ nan_treatment: str = "NA_NOTHING",
623
+ page: Optional[str] = None,
624
+ ) -> "TimeSeriesResponse":
625
+ """
626
+ Request time-series data across a subset of instruments and analytics of a given dataset.
627
+
628
+ Args:
629
+ group_id: Catalog data group identifier
630
+ attributes: List of attribute identifiers
631
+ filter: Optional filter string (e.g., "currency(USD)")
632
+ data: Data type (REFERENCE_DATA, NO_REFERENCE_DATA, ALL)
633
+ format: Response format (JSON)
634
+ start_date: Start date in YYYYMMDD or TODAY-Nx format
635
+ end_date: End date in YYYYMMDD or TODAY-Nx format
636
+ calendar: Calendar convention
637
+ frequency: Frequency convention
638
+ conversion: Conversion convention
639
+ nan_treatment: Missing data treatment
640
+ page: Optional page token for pagination
641
+
642
+ Returns:
643
+ TimeSeriesResponse containing the time series data
644
+ """
645
+ await self.connect_async()
646
+ assert self._client is not None
647
+ return await self._client.get_group_time_series_async(
648
+ group_id,
649
+ attributes,
650
+ filter,
651
+ data,
652
+ format,
653
+ start_date,
654
+ end_date,
655
+ calendar,
656
+ frequency,
657
+ conversion,
658
+ nan_treatment,
659
+ page,
660
+ )
661
+
662
+ # Grid Collection Endpoints
663
+ async def get_grid_data_async(
664
+ self,
665
+ expr: Optional[str] = None,
666
+ grid_id: Optional[str] = None,
667
+ date: Optional[str] = None,
668
+ ) -> "GridDataResponse":
669
+ """
670
+ Retrieve grid data using an expression or a grid ID.
671
+
672
+ Args:
673
+ expr: The grid expression (mutually exclusive with grid_id)
674
+ grid_id: The grid ID (mutually exclusive with expr)
675
+ date: Optional specific snapshot date in YYYYMMDD format
676
+
677
+ Returns:
678
+ GridDataResponse containing the grid data
679
+
680
+ Raises:
681
+ ValueError: If both expr and grid_id are provided or neither is provided
682
+ """
683
+ await self.connect_async()
684
+ assert self._client is not None
685
+ return await self._client.get_grid_data_async(expr, grid_id, date)
686
+
687
+ # Workflow Methods
688
+
689
+ async def run_groups_async(self, max_concurrent: int = 5) -> Dict[str, Any]:
690
+ """Run complete operation for listing all groups."""
691
+ logger.info("=== Starting Groups Operation ===")
692
+
693
+ try:
694
+ # Step 1: List all groups
695
+ logger.info("Step 1: Listing All Groups")
696
+ groups = await self.list_groups_async()
697
+
698
+ if not groups:
699
+ logger.warning("No groups found")
700
+ return {"error": "No groups found"}
701
+
702
+ # Step 2: Generate summary report
703
+ logger.info("Step 2: Summary Report")
704
+ report = {
705
+ "total_groups": len(groups),
706
+ "total_files": sum((g.file_groups or 0) for g in groups),
707
+ "groups": [g.model_dump() for g in groups],
708
+ "file_types": [], # Groups don't have file_types attribute
709
+ "providers": list(set(g.provider for g in groups if g.provider)),
710
+ }
711
+
712
+ logger.info("Groups operation completed successfully!")
713
+ logger.info("Summary report", **report)
714
+
715
+ return report
716
+
717
+ except Exception as e:
718
+ logger.error("Groups operation failed", error=str(e))
719
+ raise
720
+
721
+ async def run_group_files_async(
722
+ self, group_id: str, max_concurrent: int = 5
723
+ ) -> Dict[str, Any]:
724
+ """Run complete operation for a specific group."""
725
+ logger.info("=== Starting Group Files Operation ===", group_id=group_id)
726
+
727
+ try:
728
+ # Step 1: List files in the group
729
+ logger.info("Step 1: Listing Files")
730
+ files = await self.list_files_async(group_id)
731
+
732
+ if not files:
733
+ logger.warning("No files found for group", group_id=group_id)
734
+ return {"error": "No files found"}
735
+
736
+ # Step 2: Generate summary report
737
+ logger.info("Step 2: Summary Report")
738
+ # Collect file types robustly (handles str or List[str])
739
+ _collected_types = []
740
+ try:
741
+ for _f in files:
742
+ _ft = getattr(_f, "file_type", None)
743
+ if isinstance(_ft, list):
744
+ _collected_types.extend([t for t in _ft if isinstance(t, str)])
745
+ elif isinstance(_ft, str):
746
+ _collected_types.append(_ft)
747
+ except Exception:
748
+ pass
749
+ report = {
750
+ "group_id": group_id,
751
+ "total_files": len(files),
752
+ "file_types": list(set(_collected_types)),
753
+ "date_range": None, # FileInfo doesn't have date_range attribute
754
+ "files": [f.model_dump() for f in files],
755
+ }
756
+
757
+ logger.info("Group files operation completed successfully!")
758
+ logger.info("Summary report", **report)
759
+
760
+ return report
761
+
762
+ except Exception as e:
763
+ logger.error(
764
+ "Group files operation failed", group_id=group_id, error=str(e)
765
+ )
766
+ raise
767
+
768
+ async def run_availability_async(
769
+ self, file_group_id: str, file_datetime: str
770
+ ) -> Dict[str, Any]:
771
+ """Run operation for checking file availability."""
772
+ logger.info(
773
+ "=== Starting Availability Operation ===",
774
+ file_group_id=file_group_id,
775
+ file_datetime=file_datetime,
776
+ )
777
+
778
+ try:
779
+ # Step 1: Check availability
780
+ logger.info("Step 1: Checking Availability")
781
+ availability = await self.check_availability_async(
782
+ file_group_id, file_datetime
783
+ )
784
+
785
+ # Step 2: Generate summary report
786
+ logger.info("Step 2: Summary Report")
787
+ report = {
788
+ "file_group_id": file_group_id,
789
+ "file_datetime": file_datetime,
790
+ "is_available": bool(getattr(availability, "is_available", False)),
791
+ "file_name": getattr(availability, "file_name", None),
792
+ "first_created_on": getattr(availability, "first_created_on", None),
793
+ "last_modified": getattr(availability, "last_modified", None),
794
+ }
795
+
796
+ logger.info("Availability operation completed successfully!")
797
+ logger.info("Summary report", **report)
798
+
799
+ return report
800
+
801
+ except Exception as e:
802
+ logger.error(
803
+ "Availability operation failed",
804
+ file_group_id=file_group_id,
805
+ error=str(e),
806
+ )
807
+ raise
808
+
809
+ async def run_download_async(
810
+ self,
811
+ file_group_id: str,
812
+ file_datetime: Optional[str] = None,
813
+ destination_path: Optional[Path] = None,
814
+ max_concurrent: int = 1,
815
+ ) -> Dict[str, Any]:
816
+ """Run operation for downloading a single file."""
817
+ logger.info(
818
+ "=== Starting Download Operation ===",
819
+ file_group_id=file_group_id,
820
+ file_datetime=file_datetime,
821
+ )
822
+
823
+ try:
824
+ # Step 1: Download file
825
+ logger.info("Step 1: Downloading File")
826
+ download_options = (
827
+ DownloadOptions(
828
+ destination_path=destination_path,
829
+ create_directories=True,
830
+ overwrite_existing=False,
831
+ chunk_size=8192,
832
+ max_retries=3,
833
+ retry_delay=1.0,
834
+ timeout=600.0,
835
+ enable_range_requests=True,
836
+ range_start=None,
837
+ range_end=None,
838
+ range_header=None,
839
+ show_progress=True,
840
+ progress_callback=None,
841
+ )
842
+ if destination_path
843
+ else None
844
+ )
845
+ result = await self.download_file_async(
846
+ file_group_id, file_datetime, destination_path, download_options
847
+ )
848
+
849
+ # Step 2: Generate summary report
850
+ logger.info("Step 2: Summary Report")
851
+ report = {
852
+ "file_group_id": file_group_id,
853
+ "file_datetime": file_datetime,
854
+ "download_successful": result.status == DownloadStatus.COMPLETED,
855
+ "local_path": str(result.local_path),
856
+ "file_size": result.file_size,
857
+ "download_time": result.download_time,
858
+ "speed_mbps": result.speed_mbps,
859
+ "error_message": result.error_message,
860
+ }
861
+
862
+ logger.info("Download operation completed successfully!")
863
+ logger.info("Summary report", **report)
864
+
865
+ return report
866
+
867
+ except Exception as e:
868
+ logger.error(
869
+ "Download operation failed", file_group_id=file_group_id, error=str(e)
870
+ )
871
+ raise
872
+
873
+ async def run_group_download_async(
874
+ self,
875
+ group_id: str,
876
+ start_date: str,
877
+ end_date: str,
878
+ destination_dir: Path = Path("./downloads"),
879
+ max_concurrent: int = 5,
880
+ num_parts: int = 5,
881
+ progress_callback: Optional[Callable] = None,
882
+ delay_between_downloads: float = 1.0,
883
+ ) -> dict:
884
+ """
885
+ Download all files in a group for a date range using parallel HTTP range requests.
886
+
887
+ This method implements a flattened concurrency model where the total concurrent
888
+ HTTP requests = max_concurrent × num_parts, providing true parallelism across
889
+ all file parts rather than hierarchical file-then-parts concurrency.
890
+
891
+ Args:
892
+ group_id: Group ID to download files from
893
+ start_date: Start date in YYYYMMDD format
894
+ end_date: End date in YYYYMMDD format
895
+ destination_dir: Destination directory for downloads
896
+ max_concurrent: Maximum concurrent files (multiplied by num_parts for total concurrency)
897
+ num_parts: Number of HTTP range parts per file (default 5)
898
+ progress_callback: Optional progress callback for individual parts aggregation
899
+ delay_between_downloads: Delay in seconds between starting each file download (default 1.0)
900
+
901
+ Returns:
902
+ Dictionary with download results and statistics
903
+ """
904
+ # Record start time for total operation timing
905
+ operation_start_time = time.time()
906
+
907
+ logger.info(
908
+ "=== Starting Group Parallel Download for Date Range Operation ===",
909
+ group_id=group_id,
910
+ start_date=start_date,
911
+ end_date=end_date,
912
+ max_concurrent=max_concurrent,
913
+ num_parts=num_parts,
914
+ total_concurrency=max_concurrent * num_parts,
915
+ )
916
+
917
+ await self.connect_async()
918
+ assert self._client is not None
919
+
920
+ try:
921
+ # Step 1: Get available files for the date range
922
+ logger.info("Step 1: Getting Available Files for Date Range")
923
+ available_files = await self.list_available_files_async(
924
+ group_id=group_id, start_date=start_date, end_date=end_date
925
+ )
926
+
927
+ # Filter to only files explicitly marked available
928
+ try:
929
+ filtered_files = [
930
+ f
931
+ for f in (available_files or [])
932
+ if (f.get("is-available") is True)
933
+ or (f.get("is_available") is True)
934
+ ]
935
+ except Exception:
936
+ filtered_files = []
937
+
938
+ if not filtered_files:
939
+ # Calculate timing even when no files are found
940
+ operation_end_time = time.time()
941
+ total_time_seconds = operation_end_time - operation_start_time
942
+ total_time_minutes = total_time_seconds / 60.0
943
+
944
+ logger.warning(
945
+ "No available files found for date range",
946
+ group_id=group_id,
947
+ start_date=start_date,
948
+ end_date=end_date,
949
+ )
950
+ return {
951
+ "error": "No available files found for date range",
952
+ "group_id": group_id,
953
+ "start_date": start_date,
954
+ "end_date": end_date,
955
+ "total_files": 0,
956
+ "successful_downloads": 0,
957
+ "failed_downloads": 0,
958
+ "success_rate": 0.0,
959
+ "total_time_seconds": round(total_time_seconds, 2),
960
+ "total_time_minutes": round(total_time_minutes, 2),
961
+ "total_time_formatted": (
962
+ f"{int(total_time_minutes)}m {int(total_time_seconds % 60)}s"
963
+ if total_time_minutes >= 1
964
+ else f"{total_time_seconds:.1f}s"
965
+ ),
966
+ }
967
+
968
+ logger.info("Found available files", count=len(filtered_files))
969
+
970
+ # Step 2: Download all available files using flattened concurrency model
971
+ logger.info(
972
+ "Step 2: Downloading Available Files with flattened concurrency model"
973
+ )
974
+
975
+ # Create destination directory
976
+ dest_dir = destination_dir / group_id
977
+ dest_dir.mkdir(parents=True, exist_ok=True)
978
+
979
+ # Rate-limit-aware flattened concurrency with delay-based protection
980
+ total_concurrent_requests = max_concurrent * num_parts
981
+
982
+ # Use full concurrency but with intelligent delay calculation
983
+ rate_limit_capacity = self._calculate_rate_limit_capacity()
984
+ intelligent_delay = self._calculate_intelligent_delay(
985
+ total_concurrent_requests, rate_limit_capacity, delay_between_downloads
986
+ )
987
+
988
+ # Use full concurrency - delays will manage rate limiting
989
+ global_semaphore = asyncio.Semaphore(total_concurrent_requests)
990
+
991
+ logger.info(
992
+ "Using delay-based rate limit protection with full concurrency",
993
+ requested_concurrency=total_concurrent_requests,
994
+ rate_limit_capacity=rate_limit_capacity,
995
+ base_delay=delay_between_downloads,
996
+ intelligent_delay=intelligent_delay,
997
+ files=len(filtered_files),
998
+ )
999
+
1000
+ # Create all download tasks with flattened concurrency and staggered delays
1001
+ async def download_file_with_flattened_concurrency(
1002
+ file_info, delay_seconds: float = 0.0
1003
+ ):
1004
+ file_group_id = file_info.get(
1005
+ "file-group-id", file_info.get("file_group_id")
1006
+ )
1007
+ file_datetime = file_info.get(
1008
+ "file-datetime", file_info.get("file_datetime")
1009
+ )
1010
+
1011
+ if not file_group_id:
1012
+ logger.error("File info missing file-group-id", file_info=file_info)
1013
+ return None
1014
+
1015
+ # Apply delay before starting download
1016
+ if delay_seconds > 0:
1017
+ logger.debug(
1018
+ "Applying delay before download",
1019
+ file_group_id=file_group_id,
1020
+ delay_seconds=delay_seconds,
1021
+ )
1022
+ await asyncio.sleep(delay_seconds)
1023
+
1024
+ dest_path = dest_dir
1025
+
1026
+ try:
1027
+ # Use the client's parallel download method but with our global semaphore
1028
+ # We need to modify the client call to use our flattened concurrency
1029
+ result = await self._download_file_parallel_flattened(
1030
+ file_group_id=file_group_id,
1031
+ file_datetime=file_datetime,
1032
+ dest_path=dest_path,
1033
+ num_parts=num_parts,
1034
+ global_semaphore=global_semaphore,
1035
+ progress_callback=progress_callback,
1036
+ )
1037
+ logger.info(
1038
+ "Downloaded file (flattened concurrency)",
1039
+ file_group_id=file_group_id,
1040
+ file_datetime=file_datetime,
1041
+ status=result.status.value if result else "failed",
1042
+ )
1043
+ return result
1044
+ except Exception as e:
1045
+ logger.error(
1046
+ "Flattened parallel download failed",
1047
+ file_group_id=file_group_id,
1048
+ file_datetime=file_datetime,
1049
+ error=str(e),
1050
+ )
1051
+ return None
1052
+
1053
+ # Execute all downloads concurrently with intelligent delay-based rate limiting
1054
+ download_tasks = []
1055
+ for i, file_info in enumerate(filtered_files):
1056
+ # Calculate intelligent delay: combines base delay with rate limit protection
1057
+ delay_seconds = i * intelligent_delay
1058
+ task = download_file_with_flattened_concurrency(
1059
+ file_info, delay_seconds
1060
+ )
1061
+ download_tasks.append(task)
1062
+
1063
+ logger.info(
1064
+ "Starting downloads with intelligent delay-based rate limiting",
1065
+ total_files=len(filtered_files),
1066
+ base_delay=delay_between_downloads,
1067
+ intelligent_delay=intelligent_delay,
1068
+ total_delay_range=f"0-{(len(filtered_files) - 1) * intelligent_delay:.1f}s",
1069
+ rate_limit_protection="enabled",
1070
+ )
1071
+
1072
+ download_results = await asyncio.gather(
1073
+ *download_tasks, return_exceptions=True
1074
+ )
1075
+
1076
+ # Process results
1077
+ successful = []
1078
+ failed = []
1079
+
1080
+ for file_info, result in zip(filtered_files, download_results):
1081
+ if isinstance(result, BaseException):
1082
+ failed.append(file_info)
1083
+ elif (
1084
+ result
1085
+ and hasattr(result, "status")
1086
+ and hasattr(result, "file_group_id")
1087
+ and result.status.value == "completed"
1088
+ ):
1089
+ successful.append(result)
1090
+ else:
1091
+ failed.append(file_info)
1092
+
1093
+ # Calculate total operation time
1094
+ operation_end_time = time.time()
1095
+ total_time_seconds = operation_end_time - operation_start_time
1096
+ total_time_minutes = total_time_seconds / 60.0
1097
+
1098
+ # Generate rate limit recommendations
1099
+ recommendations = self._get_rate_limit_recommendations(
1100
+ total_concurrent_requests
1101
+ )
1102
+
1103
+ # Calculate per-file timing statistics
1104
+ file_times = []
1105
+ total_download_time = 0.0
1106
+ for result in successful:
1107
+ if result.download_time:
1108
+ file_times.append(
1109
+ {
1110
+ "file_group_id": result.file_group_id,
1111
+ "download_time_seconds": round(result.download_time, 2),
1112
+ "file_size_bytes": result.file_size or 0,
1113
+ "speed_mbps": (
1114
+ round(result.speed_mbps, 2)
1115
+ if result.speed_mbps
1116
+ else 0.0
1117
+ ),
1118
+ }
1119
+ )
1120
+ total_download_time += result.download_time
1121
+
1122
+ # Calculate timing statistics
1123
+ avg_file_time = total_download_time / len(successful) if successful else 0.0
1124
+ min_file_time = (
1125
+ min([ft["download_time_seconds"] for ft in file_times])
1126
+ if file_times
1127
+ else 0.0
1128
+ )
1129
+ max_file_time = (
1130
+ max([ft["download_time_seconds"] for ft in file_times])
1131
+ if file_times
1132
+ else 0.0
1133
+ )
1134
+
1135
+ report = {
1136
+ "group_id": group_id,
1137
+ "start_date": start_date,
1138
+ "end_date": end_date,
1139
+ "total_files": len(filtered_files),
1140
+ "successful_downloads": len(successful),
1141
+ "failed_downloads": len(failed),
1142
+ "success_rate": (
1143
+ (len(successful) / len(filtered_files)) * 100
1144
+ if filtered_files
1145
+ else 0
1146
+ ),
1147
+ "downloaded_files": [r.file_group_id for r in successful],
1148
+ "failed_files": [
1149
+ f.get("file-group-id", f.get("file_group_id", "unknown"))
1150
+ for f in failed
1151
+ ],
1152
+ "num_parts": num_parts,
1153
+ "max_concurrent": max_concurrent,
1154
+ "total_concurrent_requests": total_concurrent_requests,
1155
+ "concurrency_model": "delay_based_rate_limit_protection",
1156
+ "rate_limit_protection": "enabled",
1157
+ "base_delay": delay_between_downloads,
1158
+ "intelligent_delay": intelligent_delay,
1159
+ "delay_range": f"0-{(len(filtered_files) - 1) * intelligent_delay:.1f}s",
1160
+ "rate_limit_capacity": rate_limit_capacity,
1161
+ "rate_limit_recommendations": recommendations,
1162
+ "total_time_seconds": round(total_time_seconds, 2),
1163
+ "total_time_minutes": round(total_time_minutes, 2),
1164
+ "total_time_formatted": (
1165
+ f"{int(total_time_minutes)}m {int(total_time_seconds % 60)}s"
1166
+ if total_time_minutes >= 1
1167
+ else f"{total_time_seconds:.1f}s"
1168
+ ),
1169
+ "per_file_timing": {
1170
+ "file_times": file_times,
1171
+ "total_download_time_seconds": round(total_download_time, 2),
1172
+ "average_file_time_seconds": round(avg_file_time, 2),
1173
+ "min_file_time_seconds": round(min_file_time, 2),
1174
+ "max_file_time_seconds": round(max_file_time, 2),
1175
+ "total_download_time_formatted": (
1176
+ f"{int(total_download_time // 60)}m {int(total_download_time % 60)}s"
1177
+ if total_download_time >= 60
1178
+ else f"{total_download_time:.1f}s"
1179
+ ),
1180
+ },
1181
+ }
1182
+ logger.info(
1183
+ "Group parallel download for date range operation completed!", **report
1184
+ )
1185
+ return report
1186
+ except Exception as e:
1187
+ logger.error(
1188
+ "Group parallel download for date range operation failed",
1189
+ group_id=group_id,
1190
+ start_date=start_date,
1191
+ end_date=end_date,
1192
+ error=str(e),
1193
+ )
1194
+ raise
1195
+
1196
+ async def _download_file_parallel_flattened(
1197
+ self,
1198
+ file_group_id: str,
1199
+ file_datetime: Optional[str],
1200
+ dest_path: Path,
1201
+ num_parts: int,
1202
+ global_semaphore: asyncio.Semaphore,
1203
+ progress_callback: Optional[Callable] = None,
1204
+ ) -> Optional[DownloadResult]:
1205
+ """
1206
+ Download a file using parallel parts with flattened concurrency control.
1207
+
1208
+ This method implements the core flattened concurrency logic where each
1209
+ HTTP range request competes for the global semaphore rather than being
1210
+ grouped by file.
1211
+
1212
+ Args:
1213
+ file_group_id: File group ID to download
1214
+ file_datetime: Optional file datetime
1215
+ dest_path: Destination path (directory)
1216
+ num_parts: Number of parts to download in parallel
1217
+ global_semaphore: Global semaphore controlling total HTTP concurrency
1218
+ progress_callback: Optional progress callback
1219
+
1220
+ Returns:
1221
+ DownloadResult if successful, None if failed
1222
+ """
1223
+ import time
1224
+ from datetime import datetime
1225
+
1226
+ from .client import get_filename_from_response, validate_file_datetime
1227
+ from .models import (
1228
+ DownloadOptions,
1229
+ DownloadProgress,
1230
+ DownloadResult,
1231
+ DownloadStatus,
1232
+ )
1233
+
1234
+ try:
1235
+ assert self._client is not None
1236
+
1237
+ if file_datetime:
1238
+ validate_file_datetime(file_datetime)
1239
+
1240
+ if num_parts is None or num_parts <= 0:
1241
+ num_parts = 5
1242
+
1243
+ # Build params for API call
1244
+ params = {"file-group-id": file_group_id}
1245
+ if file_datetime:
1246
+ params["file-datetime"] = file_datetime
1247
+
1248
+ # Determine destination directory
1249
+ download_options = DownloadOptions(destination_path=dest_path)
1250
+ if download_options.destination_path:
1251
+ dest_path = Path(download_options.destination_path)
1252
+ if dest_path.suffix:
1253
+ destination_dir = dest_path.parent
1254
+ else:
1255
+ destination_dir = dest_path
1256
+ else:
1257
+ destination_dir = Path(self._client.config.download_dir)
1258
+
1259
+ if download_options.create_directories:
1260
+ destination_dir.mkdir(parents=True, exist_ok=True)
1261
+
1262
+ start_time = time.time()
1263
+ bytes_downloaded = 0
1264
+ destination: Optional[Path] = None
1265
+ temp_destination: Optional[Path] = None
1266
+ total_bytes: int = 0
1267
+ shared_fh = None
1268
+
1269
+ # Step 1: Probe file size with a 1-byte range request
1270
+ url = self._client._build_files_api_url("group/file/download")
1271
+ probe_headers = {"Range": "bytes=0-0"}
1272
+
1273
+ async with global_semaphore: # Use global semaphore for probe request
1274
+ async with await self._client._enter_request_cm(
1275
+ "GET", url, params=params, headers=probe_headers
1276
+ ) as probe_resp:
1277
+ await self._client._handle_response(probe_resp)
1278
+ content_range = probe_resp.headers.get(
1279
+ "content-range"
1280
+ ) or probe_resp.headers.get("Content-Range")
1281
+ if content_range and "/" in content_range:
1282
+ try:
1283
+ total_bytes = int(content_range.split("/")[-1])
1284
+ except Exception:
1285
+ total_bytes = int(
1286
+ probe_resp.headers.get("content-length", "0")
1287
+ )
1288
+ else:
1289
+ # Fallback to single-stream download if range not supported
1290
+ return await self._client.download_file_async(
1291
+ file_group_id=file_group_id,
1292
+ file_datetime=file_datetime,
1293
+ options=download_options,
1294
+ progress_callback=progress_callback,
1295
+ )
1296
+
1297
+ # If file is small (<10MB), prefer a single-stream download
1298
+ ten_mb = 10 * 1024 * 1024
1299
+ if total_bytes and total_bytes < ten_mb:
1300
+ return await self._client.download_file_async(
1301
+ file_group_id=file_group_id,
1302
+ file_datetime=file_datetime,
1303
+ options=download_options,
1304
+ progress_callback=progress_callback,
1305
+ )
1306
+
1307
+ # Determine filename from headers
1308
+ filename = get_filename_from_response(
1309
+ probe_resp, file_group_id, file_datetime
1310
+ )
1311
+ destination = destination_dir / filename
1312
+
1313
+ if destination.exists() and not download_options.overwrite_existing:
1314
+ raise FileExistsError(f"File already exists: {destination}")
1315
+
1316
+ # Step 2: Prepare temp file with full size for random access writes
1317
+ temp_destination = destination.with_suffix(destination.suffix + ".part")
1318
+ with open(temp_destination, "wb") as f:
1319
+ f.truncate(total_bytes)
1320
+
1321
+ # Step 3: Compute ranges for parallel download
1322
+ part_size = total_bytes // num_parts
1323
+ ranges = []
1324
+ start = 0
1325
+ for i in range(num_parts):
1326
+ end = (
1327
+ (start + part_size - 1) if i < num_parts - 1 else (total_bytes - 1)
1328
+ )
1329
+ if start > end:
1330
+ break
1331
+ ranges.append((start, end))
1332
+ start = end + 1
1333
+
1334
+ progress = DownloadProgress(
1335
+ file_group_id=file_group_id,
1336
+ total_bytes=total_bytes,
1337
+ start_time=datetime.now(),
1338
+ )
1339
+
1340
+ bytes_lock = asyncio.Lock()
1341
+ file_lock = asyncio.Lock()
1342
+ shared_fh = open(temp_destination, "r+b")
1343
+
1344
+ # Progress callback optimization: track last callback state
1345
+ last_callback_bytes = 0
1346
+ last_callback_time = time.time()
1347
+ callback_threshold_bytes = 1024 * 1024 # 1MB
1348
+ callback_threshold_time = 0.5 # 0.5 seconds
1349
+
1350
+ # Step 4: Download each range with global semaphore control
1351
+ async def download_range_with_global_semaphore(
1352
+ start_byte: int, end_byte: int
1353
+ ):
1354
+ nonlocal bytes_downloaded, last_callback_bytes, last_callback_time
1355
+ headers = {"Range": f"bytes={start_byte}-{end_byte}"}
1356
+
1357
+ assert self._client is not None # Already checked above
1358
+ # Each range request uses the global semaphore
1359
+ async with global_semaphore:
1360
+ async with await self._client._enter_request_cm(
1361
+ "GET", url, params=params, headers=headers
1362
+ ) as resp:
1363
+ await self._client._handle_response(resp)
1364
+ # Stream and write to correct offset
1365
+ current_pos = start_byte
1366
+ chunk_size = download_options.chunk_size or 8192
1367
+ async for chunk in resp.content.iter_chunked(chunk_size):
1368
+ async with file_lock:
1369
+ shared_fh.seek(current_pos)
1370
+ shared_fh.write(chunk)
1371
+ current_pos += len(chunk)
1372
+ async with bytes_lock:
1373
+ bytes_downloaded += len(chunk)
1374
+ progress.update_progress(bytes_downloaded)
1375
+
1376
+ # Optimized progress callback: only call every 1MB or 0.5s
1377
+ current_time = time.time()
1378
+ bytes_diff = bytes_downloaded - last_callback_bytes
1379
+ time_diff = current_time - last_callback_time
1380
+
1381
+ should_callback = (
1382
+ bytes_diff >= callback_threshold_bytes
1383
+ or time_diff >= callback_threshold_time
1384
+ or bytes_downloaded
1385
+ == total_bytes # Always callback on completion
1386
+ )
1387
+
1388
+ if should_callback:
1389
+ if progress_callback:
1390
+ progress_callback(progress)
1391
+ elif download_options.show_progress:
1392
+ logger.info(
1393
+ "Download progress (flattened)",
1394
+ file=file_group_id,
1395
+ percentage=f"{progress.percentage:.1f}%",
1396
+ downloaded=format_file_size(
1397
+ bytes_downloaded
1398
+ ),
1399
+ )
1400
+ last_callback_bytes = bytes_downloaded
1401
+ last_callback_time = current_time
1402
+
1403
+ # Step 5: Execute all range downloads concurrently (each will acquire global semaphore)
1404
+ await asyncio.gather(
1405
+ *(download_range_with_global_semaphore(s, e) for s, e in ranges)
1406
+ )
1407
+
1408
+ # Step 6: Finalize file
1409
+ try:
1410
+ async with file_lock:
1411
+ shared_fh.flush()
1412
+ finally:
1413
+ try:
1414
+ shared_fh.close()
1415
+ except Exception:
1416
+ pass
1417
+
1418
+ temp_destination.replace(destination)
1419
+
1420
+ download_time = time.time() - start_time
1421
+ return DownloadResult(
1422
+ file_group_id=file_group_id,
1423
+ group_id="",
1424
+ local_path=destination,
1425
+ file_size=total_bytes,
1426
+ download_time=download_time,
1427
+ bytes_downloaded=bytes_downloaded,
1428
+ status=DownloadStatus.COMPLETED,
1429
+ error_message=None,
1430
+ )
1431
+
1432
+ except Exception as e:
1433
+ # Cleanup on error
1434
+ try:
1435
+ if shared_fh:
1436
+ shared_fh.close()
1437
+ if temp_destination and temp_destination.exists():
1438
+ # Attempt salvage if file appears complete
1439
+ if total_bytes and temp_destination.stat().st_size >= total_bytes:
1440
+ if destination is None:
1441
+ destination = temp_destination.with_suffix("")
1442
+ temp_destination.replace(destination)
1443
+ return DownloadResult(
1444
+ file_group_id=file_group_id,
1445
+ group_id="",
1446
+ local_path=destination,
1447
+ file_size=total_bytes,
1448
+ download_time=time.time() - start_time,
1449
+ bytes_downloaded=bytes_downloaded,
1450
+ status=DownloadStatus.COMPLETED,
1451
+ error_message=None,
1452
+ )
1453
+ except Exception:
1454
+ pass
1455
+
1456
+ logger.error(
1457
+ "Flattened parallel download failed for file",
1458
+ file_group_id=file_group_id,
1459
+ error=str(e),
1460
+ )
1461
+ return None
1462
+
1463
+ def _calculate_rate_limit_capacity(self) -> Dict[str, Any]:
1464
+ """
1465
+ Calculate the rate limit capacity without reducing concurrency.
1466
+
1467
+ Returns:
1468
+ Dictionary with rate limit capacity information
1469
+ """
1470
+ if not self._client:
1471
+ return {
1472
+ "requests_per_minute": 1000,
1473
+ "burst_capacity": 100,
1474
+ "safe_interval": 0.1,
1475
+ }
1476
+
1477
+ try:
1478
+ rate_config = self._client.rate_limiter.config
1479
+
1480
+ if not rate_config.enable_rate_limiting:
1481
+ return {
1482
+ "requests_per_minute": 1000,
1483
+ "burst_capacity": 100,
1484
+ "safe_interval": 0.1,
1485
+ }
1486
+
1487
+ requests_per_minute = rate_config.requests_per_minute
1488
+ burst_capacity = rate_config.burst_capacity
1489
+
1490
+ # Calculate safe interval between requests to stay within rate limits
1491
+ # This ensures we don't exceed the per-minute limit
1492
+ safe_interval = 60.0 / requests_per_minute # seconds between requests
1493
+
1494
+ return {
1495
+ "requests_per_minute": requests_per_minute,
1496
+ "burst_capacity": burst_capacity,
1497
+ "safe_interval": safe_interval,
1498
+ "queuing_enabled": rate_config.enable_queuing,
1499
+ }
1500
+
1501
+ except Exception as e:
1502
+ logger.warning(
1503
+ "Failed to calculate rate limit capacity, using conservative defaults",
1504
+ error=str(e),
1505
+ )
1506
+ return {
1507
+ "requests_per_minute": 100,
1508
+ "burst_capacity": 10,
1509
+ "safe_interval": 0.6,
1510
+ }
1511
+
1512
+ def _calculate_intelligent_delay(
1513
+ self,
1514
+ total_concurrent_requests: int,
1515
+ rate_limit_capacity: Dict[str, Any],
1516
+ base_delay: float,
1517
+ ) -> float:
1518
+ """
1519
+ Calculate intelligent delay that ensures rate limit compliance while maintaining concurrency.
1520
+
1521
+ Burst capacity is the absolute number of requests that can be made immediately,
1522
+ not per minute or per second. It's a one-time allowance before throttling begins.
1523
+
1524
+ Args:
1525
+ total_concurrent_requests: Total number of concurrent requests
1526
+ rate_limit_capacity: Rate limit capacity information
1527
+ base_delay: Base delay requested by user
1528
+
1529
+ Returns:
1530
+ Intelligent delay that ensures rate limit compliance
1531
+ """
1532
+ try:
1533
+ requests_per_minute = rate_limit_capacity["requests_per_minute"]
1534
+ burst_capacity = rate_limit_capacity[
1535
+ "burst_capacity"
1536
+ ] # Absolute number of immediate requests
1537
+ safe_interval = rate_limit_capacity["safe_interval"]
1538
+
1539
+ # Calculate minimum delay needed to stay within rate limits
1540
+ # We need to ensure that even with full concurrency, we don't exceed RPM
1541
+ min_delay_for_rpm = safe_interval
1542
+
1543
+ # Calculate delay needed for burst capacity
1544
+ # Burst capacity is the absolute number of requests that can be made immediately
1545
+ # If we have more concurrent requests than burst capacity, we need to spread them over time
1546
+ if total_concurrent_requests > burst_capacity:
1547
+ # Calculate delay to spread the excess requests over time
1548
+ # The excess requests need to be spread out to avoid exceeding burst capacity
1549
+ excess_requests = total_concurrent_requests - burst_capacity
1550
+
1551
+ # Calculate how much time we need to spread the excess requests
1552
+ # We want to spread them over a reasonable time window based on the excess
1553
+ # More excess requests = longer time window to spread them out
1554
+ time_window = max(
1555
+ 5.0, excess_requests * 0.1
1556
+ ) # At least 5s, or 0.1s per excess request
1557
+ min_delay_for_burst = (
1558
+ time_window / excess_requests if excess_requests > 0 else 0.0
1559
+ )
1560
+
1561
+ logger.debug(
1562
+ "Burst capacity exceeded, calculating spread delay",
1563
+ total_concurrent_requests=total_concurrent_requests,
1564
+ burst_capacity=burst_capacity,
1565
+ excess_requests=excess_requests,
1566
+ time_window=time_window,
1567
+ min_delay_for_burst=min_delay_for_burst,
1568
+ )
1569
+ else:
1570
+ min_delay_for_burst = 0.0
1571
+
1572
+ # Use the maximum of user's base delay and calculated minimum delays
1573
+ intelligent_delay = max(base_delay, min_delay_for_rpm, min_delay_for_burst)
1574
+
1575
+ # Log the calculation details for debugging
1576
+ logger.debug(
1577
+ "Intelligent delay calculation details",
1578
+ base_delay=base_delay,
1579
+ min_delay_for_rpm=min_delay_for_rpm,
1580
+ min_delay_for_burst=min_delay_for_burst,
1581
+ final_intelligent_delay=intelligent_delay,
1582
+ )
1583
+
1584
+ # Add some safety margin (10%)
1585
+ intelligent_delay *= 1.1
1586
+
1587
+ logger.info(
1588
+ "Calculated intelligent delay for rate limit protection",
1589
+ total_concurrent_requests=total_concurrent_requests,
1590
+ requests_per_minute=requests_per_minute,
1591
+ burst_capacity=burst_capacity,
1592
+ safe_interval=safe_interval,
1593
+ base_delay=base_delay,
1594
+ intelligent_delay=intelligent_delay,
1595
+ burst_capacity_type="absolute_requests",
1596
+ rate_limit_protection="enabled",
1597
+ )
1598
+
1599
+ return intelligent_delay
1600
+
1601
+ except Exception as e:
1602
+ logger.warning(
1603
+ "Failed to calculate intelligent delay, using base delay", error=str(e)
1604
+ )
1605
+ return base_delay
1606
+
1607
+ def _calculate_safe_concurrency_limit(self, requested_concurrency: int) -> int:
1608
+ """
1609
+ Calculate a safe concurrency limit that respects rate limiting constraints.
1610
+
1611
+ This method ensures that the total concurrent HTTP requests don't overwhelm
1612
+ the rate limiter, which could cause request queuing or failures.
1613
+
1614
+ Args:
1615
+ requested_concurrency: The desired number of concurrent requests
1616
+
1617
+ Returns:
1618
+ Safe concurrency limit that respects rate limiting
1619
+ """
1620
+ if not self._client:
1621
+ return requested_concurrency
1622
+
1623
+ try:
1624
+ # Get rate limiter configuration
1625
+ rate_config = self._client.rate_limiter.config
1626
+
1627
+ # If rate limiting is disabled, use requested concurrency
1628
+ if not rate_config.enable_rate_limiting:
1629
+ logger.info(
1630
+ "Rate limiting disabled, using requested concurrency",
1631
+ requested=requested_concurrency,
1632
+ )
1633
+ return requested_concurrency
1634
+
1635
+ # Calculate safe limits based on rate limiter configuration
1636
+ requests_per_minute = rate_config.requests_per_minute
1637
+ burst_capacity = rate_config.burst_capacity
1638
+ queue_size = rate_config.max_queue_size if rate_config.enable_queuing else 0
1639
+
1640
+ # Conservative calculation:
1641
+ # - Don't exceed burst capacity for immediate requests
1642
+ # - Consider queue capacity for sustained load
1643
+ # - Leave some headroom for other operations
1644
+
1645
+ immediate_capacity = burst_capacity
1646
+ sustained_capacity = min(
1647
+ requests_per_minute
1648
+ // 4, # Quarter of per-minute limit for sustained load
1649
+ (
1650
+ queue_size // 2 if queue_size > 0 else immediate_capacity
1651
+ ), # Half queue size
1652
+ )
1653
+
1654
+ # Use the more conservative of immediate or sustained capacity
1655
+ safe_limit = max(
1656
+ 1, min(immediate_capacity, sustained_capacity, requested_concurrency)
1657
+ )
1658
+
1659
+ # Add warning if we're significantly reducing concurrency
1660
+ if safe_limit < requested_concurrency * 0.5:
1661
+ logger.warning(
1662
+ "Significantly reducing concurrency due to rate limits",
1663
+ requested=requested_concurrency,
1664
+ safe_limit=safe_limit,
1665
+ rate_limit_rpm=requests_per_minute,
1666
+ burst_capacity=burst_capacity,
1667
+ recommendation="Consider reducing max_concurrent or num_parts, or increasing rate limits",
1668
+ )
1669
+
1670
+ logger.info(
1671
+ "Calculated safe concurrency limit",
1672
+ requested=requested_concurrency,
1673
+ safe_limit=safe_limit,
1674
+ rate_limit_rpm=requests_per_minute,
1675
+ burst_capacity=burst_capacity,
1676
+ queuing_enabled=rate_config.enable_queuing,
1677
+ )
1678
+
1679
+ return safe_limit
1680
+
1681
+ except Exception as e:
1682
+ logger.warning(
1683
+ "Failed to calculate safe concurrency limit, using conservative default",
1684
+ error=str(e),
1685
+ requested=requested_concurrency,
1686
+ )
1687
+ # Conservative fallback: use a small fraction of requested
1688
+ return max(1, min(10, requested_concurrency // 4))
1689
+
1690
+ def _get_rate_limit_recommendations(
1691
+ self, requested_concurrency: int
1692
+ ) -> Dict[str, Any]:
1693
+ """
1694
+ Get recommendations for optimizing rate limit settings for the requested concurrency.
1695
+
1696
+ Args:
1697
+ requested_concurrency: Desired concurrent requests
1698
+
1699
+ Returns:
1700
+ Dictionary with recommendations
1701
+ """
1702
+ if not self._client:
1703
+ return {}
1704
+
1705
+ try:
1706
+ rate_config = self._client.rate_limiter.config
1707
+
1708
+ recommendations: Dict[str, Any] = {
1709
+ "current_settings": {
1710
+ "requests_per_minute": rate_config.requests_per_minute,
1711
+ "burst_capacity": rate_config.burst_capacity,
1712
+ "queuing_enabled": rate_config.enable_queuing,
1713
+ "max_queue_size": rate_config.max_queue_size,
1714
+ },
1715
+ "requested_concurrency": requested_concurrency,
1716
+ "recommendations": [],
1717
+ }
1718
+
1719
+ # Recommend burst capacity increase if needed
1720
+ if requested_concurrency > rate_config.burst_capacity:
1721
+ recommendations["recommendations"].append(
1722
+ {
1723
+ "type": "increase_burst_capacity",
1724
+ "current": rate_config.burst_capacity,
1725
+ "recommended": max(
1726
+ requested_concurrency, rate_config.burst_capacity * 2
1727
+ ),
1728
+ "reason": "Burst capacity should accommodate concurrent requests",
1729
+ }
1730
+ )
1731
+
1732
+ # Recommend requests per minute increase for sustained load
1733
+ sustained_rpm_needed = (
1734
+ requested_concurrency * 15
1735
+ ) # Assume 15 requests per minute per concurrent slot
1736
+ if sustained_rpm_needed > rate_config.requests_per_minute:
1737
+ recommendations["recommendations"].append(
1738
+ {
1739
+ "type": "increase_requests_per_minute",
1740
+ "current": rate_config.requests_per_minute,
1741
+ "recommended": sustained_rpm_needed,
1742
+ "reason": "Higher RPM needed for sustained concurrent load",
1743
+ }
1744
+ )
1745
+
1746
+ # Recommend enabling queuing if not enabled
1747
+ if (
1748
+ not rate_config.enable_queuing
1749
+ and requested_concurrency > rate_config.burst_capacity
1750
+ ):
1751
+ recommendations["recommendations"].append(
1752
+ {
1753
+ "type": "enable_queuing",
1754
+ "current": False,
1755
+ "recommended": True,
1756
+ "reason": "Queuing helps manage burst requests beyond capacity",
1757
+ }
1758
+ )
1759
+
1760
+ return recommendations
1761
+
1762
+ except Exception as e:
1763
+ logger.warning(
1764
+ "Failed to generate rate limit recommendations", error=str(e)
1765
+ )
1766
+ return {}
1767
+
1768
+ # Utility Methods
1769
+
1770
+ def get_pool_stats(self) -> Dict[str, Any]:
1771
+ """Get connection pool statistics including active, idle, and total connections."""
1772
+ if self._client:
1773
+ return self._client.get_pool_stats()
1774
+ return {"error": "Client not connected"}
1775
+
1776
+ def get_stats(self) -> Dict[str, Any]:
1777
+ """Get comprehensive client statistics."""
1778
+ if self._client:
1779
+ stats = self._client.get_stats()
1780
+ # Add context path information
1781
+ stats["api_config"] = {
1782
+ "base_url": self.client_config.base_url,
1783
+ "context_path": self.client_config.context_path,
1784
+ "api_base_url": self.client_config.api_base_url,
1785
+ }
1786
+ return stats
1787
+ return {"status": "not_connected"}
1788
+
1789
+ def create_progress_callback(self, log_interval: int = 10) -> Callable:
1790
+ """Create a progress callback function."""
1791
+ tracker = ProgressTracker(log_interval)
1792
+ return tracker.create_progress_callback()
1793
+
1794
+ def get_rate_limit_info(self) -> Dict[str, Any]:
1795
+ """
1796
+ Get current rate limit configuration and status.
1797
+
1798
+ Returns:
1799
+ Dictionary with rate limit information
1800
+ """
1801
+ if not self._client:
1802
+ return {"error": "Client not connected"}
1803
+
1804
+ try:
1805
+ rate_config = self._client.rate_limiter.config
1806
+ rate_state = self._client.rate_limiter.state
1807
+
1808
+ return {
1809
+ "configuration": {
1810
+ "requests_per_minute": rate_config.requests_per_minute,
1811
+ "burst_capacity": rate_config.burst_capacity,
1812
+ "enable_rate_limiting": rate_config.enable_rate_limiting,
1813
+ "enable_queuing": rate_config.enable_queuing,
1814
+ "max_queue_size": rate_config.max_queue_size,
1815
+ "adaptive_rate_limiting": rate_config.adaptive_rate_limiting,
1816
+ },
1817
+ "current_state": {
1818
+ "available_tokens": rate_state.tokens,
1819
+ "queue_size": len(rate_state.queue),
1820
+ "total_requests": rate_state.total_requests,
1821
+ "successful_requests": rate_state.successful_requests,
1822
+ "failed_requests": rate_state.failed_requests,
1823
+ "rate_limited_requests": rate_state.rate_limited_requests,
1824
+ },
1825
+ }
1826
+ except Exception as e:
1827
+ return {"error": f"Failed to get rate limit info: {e}"}
1828
+
1829
+ def optimize_concurrency_for_rate_limits(
1830
+ self, max_concurrent: int, num_parts: int
1831
+ ) -> Dict[str, Any]:
1832
+ """
1833
+ Get optimized concurrency settings that respect rate limits.
1834
+
1835
+ Args:
1836
+ max_concurrent: Desired maximum concurrent files
1837
+ num_parts: Desired number of parts per file
1838
+
1839
+ Returns:
1840
+ Dictionary with optimized settings and recommendations
1841
+ """
1842
+ requested_concurrency = max_concurrent * num_parts
1843
+ safe_concurrency = self._calculate_safe_concurrency_limit(requested_concurrency)
1844
+ recommendations = self._get_rate_limit_recommendations(requested_concurrency)
1845
+
1846
+ # Calculate optimal max_concurrent and num_parts that fit within safe limits
1847
+ if safe_concurrency < requested_concurrency:
1848
+ # Try to maintain the ratio but reduce total
1849
+ ratio = safe_concurrency / requested_concurrency
1850
+ optimal_max_concurrent = max(1, int(max_concurrent * ratio))
1851
+ optimal_num_parts = max(1, int(num_parts * ratio))
1852
+
1853
+ # Ensure we don't exceed safe limits
1854
+ while optimal_max_concurrent * optimal_num_parts > safe_concurrency:
1855
+ if optimal_num_parts > optimal_max_concurrent:
1856
+ optimal_num_parts -= 1
1857
+ else:
1858
+ optimal_max_concurrent -= 1
1859
+ else:
1860
+ optimal_max_concurrent = max_concurrent
1861
+ optimal_num_parts = num_parts
1862
+
1863
+ return {
1864
+ "requested": {
1865
+ "max_concurrent": max_concurrent,
1866
+ "num_parts": num_parts,
1867
+ "total_concurrency": requested_concurrency,
1868
+ },
1869
+ "optimized": {
1870
+ "max_concurrent": optimal_max_concurrent,
1871
+ "num_parts": optimal_num_parts,
1872
+ "total_concurrency": optimal_max_concurrent * optimal_num_parts,
1873
+ },
1874
+ "safe_limit": safe_concurrency,
1875
+ "rate_limit_applied": safe_concurrency < requested_concurrency,
1876
+ "recommendations": recommendations,
1877
+ }
1878
+
1879
+ # Synchronous Wrapper Methods
1880
+
1881
+ def connect(self):
1882
+ """Connect to the API."""
1883
+ return self._run_sync(self.connect_async())
1884
+
1885
+ def close(self):
1886
+ """Close the connection and cleanup resources."""
1887
+ if self._client:
1888
+ return self._run_sync(self.close_async())
1889
+
1890
+ def list_groups(self, limit: Optional[int] = 100) -> List[Group]:
1891
+ """Synchronous wrapper for list_groups."""
1892
+ return self._run_sync(self.list_groups_async(limit))
1893
+
1894
+ def search_groups(
1895
+ self, keywords: str, limit: Optional[int] = 100, offset: Optional[int] = None
1896
+ ) -> List[Group]:
1897
+ """Synchronous wrapper for search_groups."""
1898
+ return self._run_sync(self.search_groups_async(keywords, limit, offset))
1899
+
1900
+ def list_files(
1901
+ self, group_id: str, file_group_id: Optional[str] = None
1902
+ ) -> List[FileInfo]:
1903
+ """Synchronous wrapper for list_files."""
1904
+ return self._run_sync(self.list_files_async(group_id, file_group_id))
1905
+
1906
+ def check_availability(
1907
+ self, file_group_id: str, file_datetime: str
1908
+ ) -> AvailabilityInfo:
1909
+ """Synchronous wrapper for check_availability."""
1910
+ return self._run_sync(
1911
+ self.check_availability_async(file_group_id, file_datetime)
1912
+ )
1913
+
1914
+ def download_file(
1915
+ self,
1916
+ file_group_id: str,
1917
+ file_datetime: Optional[str] = None,
1918
+ destination_path: Optional[Path] = None,
1919
+ options: Optional[DownloadOptions] = None,
1920
+ num_parts: int = 5,
1921
+ progress_callback: Optional[Callable] = None,
1922
+ ) -> DownloadResult:
1923
+ """
1924
+ Synchronous wrapper for download_file.
1925
+ Note: Will raise an error if called from within an existing event loop.
1926
+
1927
+ Args:
1928
+ file_group_id: File ID to download
1929
+ file_datetime: Optional datetime of the file (YYYYMMDD, YYYYMMDDTHHMM, or YYYYMMDDTHHMMSS)
1930
+ destination_path: Optional download destination directory. The filename will be extracted
1931
+ from the Content-Disposition header in the response. If not provided,
1932
+ uses the default download directory from configuration.
1933
+ options: Download options
1934
+ num_parts: Number of parallel parts for download (default: 5)
1935
+ progress_callback: Optional progress callback function
1936
+
1937
+ Returns:
1938
+ DownloadResult with download information
1939
+ """
1940
+ return self._run_sync(
1941
+ self.download_file_async(
1942
+ file_group_id,
1943
+ file_datetime,
1944
+ destination_path,
1945
+ options,
1946
+ num_parts,
1947
+ progress_callback,
1948
+ )
1949
+ )
1950
+
1951
+ def list_available_files(
1952
+ self,
1953
+ group_id: str,
1954
+ file_group_id: Optional[str] = None,
1955
+ start_date: Optional[str] = None,
1956
+ end_date: Optional[str] = None,
1957
+ ) -> List[Dict[str, Any]]:
1958
+ """Synchronous wrapper for list_available_files."""
1959
+ return self._run_sync(
1960
+ self.list_available_files_async(
1961
+ group_id, file_group_id, start_date, end_date
1962
+ )
1963
+ )
1964
+
1965
+ def health_check(self) -> bool:
1966
+ """Synchronous wrapper for health_check."""
1967
+ return self._run_sync(self.health_check_async())
1968
+
1969
+ # Instrument Collection Endpoints - Synchronous wrappers
1970
+ def list_instruments(
1971
+ self,
1972
+ group_id: str,
1973
+ instrument_id: Optional[str] = None,
1974
+ page: Optional[str] = None,
1975
+ ) -> "InstrumentsResponse":
1976
+ """
1977
+ Request the complete list of instruments and identifiers for a given dataset.
1978
+
1979
+ Args:
1980
+ group_id: Catalog data group identifier
1981
+ instrument_id: Optional instrument identifier to filter results
1982
+ page: Optional page token for pagination
1983
+
1984
+ Returns:
1985
+ InstrumentsResponse containing the list of instruments
1986
+ """
1987
+ return self._run_sync(
1988
+ self.list_instruments_async(group_id, instrument_id, page)
1989
+ )
1990
+
1991
+ def search_instruments(
1992
+ self, group_id: str, keywords: str, page: Optional[str] = None
1993
+ ) -> "InstrumentsResponse":
1994
+ """
1995
+ Search within a dataset using keywords to create subsets of matching instruments.
1996
+
1997
+ Args:
1998
+ group_id: Catalog data group identifier
1999
+ keywords: Keywords to narrow scope of results
2000
+ page: Optional page token for pagination
2001
+
2002
+ Returns:
2003
+ InstrumentsResponse containing the matching instruments
2004
+ """
2005
+ return self._run_sync(self.search_instruments_async(group_id, keywords, page))
2006
+
2007
+ def get_instrument_time_series(
2008
+ self,
2009
+ instruments: List[str],
2010
+ attributes: List[str],
2011
+ data: str = "REFERENCE_DATA",
2012
+ format: str = "JSON",
2013
+ start_date: Optional[str] = None,
2014
+ end_date: Optional[str] = None,
2015
+ calendar: str = "CAL_USBANK",
2016
+ frequency: str = "FREQ_DAY",
2017
+ conversion: str = "CONV_LASTBUS_ABS",
2018
+ nan_treatment: str = "NA_NOTHING",
2019
+ page: Optional[str] = None,
2020
+ ) -> "TimeSeriesResponse":
2021
+ """
2022
+ Retrieve time-series data for explicit list of instruments and attributes using identifiers.
2023
+
2024
+ Args:
2025
+ instruments: List of instrument identifiers
2026
+ attributes: List of attribute identifiers
2027
+ data: Data type (REFERENCE_DATA, NO_REFERENCE_DATA, ALL)
2028
+ format: Response format (JSON)
2029
+ start_date: Start date in YYYYMMDD or TODAY-Nx format
2030
+ end_date: End date in YYYYMMDD or TODAY-Nx format
2031
+ calendar: Calendar convention
2032
+ frequency: Frequency convention
2033
+ conversion: Conversion convention
2034
+ nan_treatment: Missing data treatment
2035
+ page: Optional page token for pagination
2036
+
2037
+ Returns:
2038
+ TimeSeriesResponse containing the time series data
2039
+ """
2040
+ return self._run_sync(
2041
+ self.get_instrument_time_series_async(
2042
+ instruments,
2043
+ attributes,
2044
+ data,
2045
+ format,
2046
+ start_date,
2047
+ end_date,
2048
+ calendar,
2049
+ frequency,
2050
+ conversion,
2051
+ nan_treatment,
2052
+ page,
2053
+ )
2054
+ )
2055
+
2056
+ def get_expressions_time_series(
2057
+ self,
2058
+ expressions: List[str],
2059
+ format: str = "JSON",
2060
+ start_date: Optional[str] = None,
2061
+ end_date: Optional[str] = None,
2062
+ calendar: str = "CAL_USBANK",
2063
+ frequency: str = "FREQ_DAY",
2064
+ conversion: str = "CONV_LASTBUS_ABS",
2065
+ nan_treatment: str = "NA_NOTHING",
2066
+ data: str = "REFERENCE_DATA",
2067
+ page: Optional[str] = None,
2068
+ ) -> "TimeSeriesResponse":
2069
+ """
2070
+ Retrieve time-series data using an explicit list of traditional DataQuery expressions.
2071
+
2072
+ Args:
2073
+ expressions: List of traditional DataQuery expressions
2074
+ format: Response format (JSON)
2075
+ start_date: Start date in YYYYMMDD or TODAY-Nx format
2076
+ end_date: End date in YYYYMMDD or TODAY-Nx format
2077
+ calendar: Calendar convention
2078
+ frequency: Frequency convention
2079
+ conversion: Conversion convention
2080
+ nan_treatment: Missing data treatment
2081
+ data: Data type (REFERENCE_DATA, NO_REFERENCE_DATA, ALL)
2082
+ page: Optional page token for pagination
2083
+
2084
+ Returns:
2085
+ TimeSeriesResponse containing the time series data
2086
+ """
2087
+ return self._run_sync(
2088
+ self.get_expressions_time_series_async(
2089
+ expressions,
2090
+ format,
2091
+ start_date,
2092
+ end_date,
2093
+ calendar,
2094
+ frequency,
2095
+ conversion,
2096
+ nan_treatment,
2097
+ data,
2098
+ page,
2099
+ )
2100
+ )
2101
+
2102
+ # Group Collection Additional Endpoints - Synchronous wrappers
2103
+ def get_group_filters(
2104
+ self, group_id: str, page: Optional[str] = None
2105
+ ) -> "FiltersResponse":
2106
+ """
2107
+ Request the unique list of filter dimensions that are available for a given dataset.
2108
+
2109
+ Args:
2110
+ group_id: Catalog data group identifier
2111
+ page: Optional page token for pagination
2112
+
2113
+ Returns:
2114
+ FiltersResponse containing the available filters
2115
+ """
2116
+ return self._run_sync(self.get_group_filters_async(group_id, page))
2117
+
2118
+ def get_group_attributes(
2119
+ self,
2120
+ group_id: str,
2121
+ instrument_id: Optional[str] = None,
2122
+ page: Optional[str] = None,
2123
+ ) -> "AttributesResponse":
2124
+ """
2125
+ Request the unique list of analytic attributes for each instrument of a given dataset.
2126
+
2127
+ Args:
2128
+ group_id: Catalog data group identifier
2129
+ instrument_id: Optional instrument identifier to filter results
2130
+ page: Optional page token for pagination
2131
+
2132
+ Returns:
2133
+ AttributesResponse containing the attributes for each instrument
2134
+ """
2135
+ return self._run_sync(
2136
+ self.get_group_attributes_async(group_id, instrument_id, page)
2137
+ )
2138
+
2139
+ def get_group_time_series(
2140
+ self,
2141
+ group_id: str,
2142
+ attributes: List[str],
2143
+ filter: Optional[str] = None,
2144
+ data: str = "REFERENCE_DATA",
2145
+ format: str = "JSON",
2146
+ start_date: Optional[str] = None,
2147
+ end_date: Optional[str] = None,
2148
+ calendar: str = "CAL_USBANK",
2149
+ frequency: str = "FREQ_DAY",
2150
+ conversion: str = "CONV_LASTBUS_ABS",
2151
+ nan_treatment: str = "NA_NOTHING",
2152
+ page: Optional[str] = None,
2153
+ ) -> "TimeSeriesResponse":
2154
+ """
2155
+ Request time-series data across a subset of instruments and analytics of a given dataset.
2156
+
2157
+ Args:
2158
+ group_id: Catalog data group identifier
2159
+ attributes: List of attribute identifiers
2160
+ filter: Optional filter string (e.g., "currency(USD)")
2161
+ data: Data type (REFERENCE_DATA, NO_REFERENCE_DATA, ALL)
2162
+ format: Response format (JSON)
2163
+ start_date: Start date in YYYYMMDD or TODAY-Nx format
2164
+ end_date: End date in YYYYMMDD or TODAY-Nx format
2165
+ calendar: Calendar convention
2166
+ frequency: Frequency convention
2167
+ conversion: Conversion convention
2168
+ nan_treatment: Missing data treatment
2169
+ page: Optional page token for pagination
2170
+
2171
+ Returns:
2172
+ TimeSeriesResponse containing the time series data
2173
+ """
2174
+ return self._run_sync(
2175
+ self.get_group_time_series_async(
2176
+ group_id,
2177
+ attributes,
2178
+ filter,
2179
+ data,
2180
+ format,
2181
+ start_date,
2182
+ end_date,
2183
+ calendar,
2184
+ frequency,
2185
+ conversion,
2186
+ nan_treatment,
2187
+ page,
2188
+ )
2189
+ )
2190
+
2191
+ # Grid Collection Endpoints - Synchronous wrappers
2192
+ def get_grid_data(
2193
+ self,
2194
+ expr: Optional[str] = None,
2195
+ grid_id: Optional[str] = None,
2196
+ date: Optional[str] = None,
2197
+ ) -> "GridDataResponse":
2198
+ """
2199
+ Retrieve grid data using an expression or a grid ID.
2200
+
2201
+ Args:
2202
+ expr: The grid expression (mutually exclusive with grid_id)
2203
+ grid_id: The grid ID (mutually exclusive with expr)
2204
+ date: Optional specific snapshot date in YYYYMMDD format
2205
+
2206
+ Returns:
2207
+ GridDataResponse containing the grid data
2208
+
2209
+ Raises:
2210
+ ValueError: If both expr and grid_id are provided or neither is provided
2211
+ """
2212
+ return self._run_sync(self.get_grid_data_async(expr, grid_id, date))
2213
+
2214
+ def run_groups(self, max_concurrent: int = 5) -> Dict[str, Any]:
2215
+ """Synchronous wrapper for run_groups_async."""
2216
+ return self._run_sync(self.run_groups_async(max_concurrent))
2217
+
2218
+ def run_group_files(self, group_id: str, max_concurrent: int = 5) -> Dict[str, Any]:
2219
+ """Synchronous wrapper for run_group_files_async."""
2220
+ return self._run_sync(self.run_group_files_async(group_id, max_concurrent))
2221
+
2222
+ def run_availability(
2223
+ self, file_group_id: str, file_datetime: str
2224
+ ) -> Dict[str, Any]:
2225
+ """Synchronous wrapper for run_availability_async."""
2226
+ return self._run_sync(self.run_availability_async(file_group_id, file_datetime))
2227
+
2228
+ def run_download(
2229
+ self,
2230
+ file_group_id: str,
2231
+ file_datetime: Optional[str] = None,
2232
+ destination_path: Optional[Path] = None,
2233
+ max_concurrent: int = 1,
2234
+ ) -> Dict[str, Any]:
2235
+ """Synchronous wrapper for run_download_async."""
2236
+ return self._run_sync(
2237
+ self.run_download_async(
2238
+ file_group_id, file_datetime, destination_path, max_concurrent
2239
+ )
2240
+ )
2241
+
2242
+ def run_group_download(
2243
+ self,
2244
+ group_id: str,
2245
+ start_date: str,
2246
+ end_date: str,
2247
+ destination_dir: Path = Path("./downloads"),
2248
+ max_concurrent: int = 5,
2249
+ num_parts: int = 5,
2250
+ progress_callback: Optional[Callable] = None,
2251
+ delay_between_downloads: float = 1.0,
2252
+ ) -> dict:
2253
+ """Synchronous wrapper for run_group_download_async."""
2254
+ return self._run_sync(
2255
+ self.run_group_download_async(
2256
+ group_id,
2257
+ start_date,
2258
+ end_date,
2259
+ destination_dir,
2260
+ max_concurrent,
2261
+ num_parts,
2262
+ progress_callback,
2263
+ delay_between_downloads,
2264
+ )
2265
+ )
2266
+
2267
+ def cleanup(self):
2268
+ """Synchronous cleanup resources and ensure proper shutdown."""
2269
+ if self._client:
2270
+ self._run_sync(self.close_async())
2271
+ self._client = None
2272
+
2273
+ # Force garbage collection to clean up any remaining references
2274
+ import gc
2275
+
2276
+ gc.collect()
2277
+
2278
+ # Sync wrapper methods with _sync suffix for testing compatibility
2279
+
2280
+ def connect_sync(self):
2281
+ """Synchronous wrapper for connect with _sync suffix."""
2282
+ return asyncio.run(self.connect_async())
2283
+
2284
+ def close_sync(self):
2285
+ """Synchronous wrapper for close with _sync suffix."""
2286
+ if self._client:
2287
+ return asyncio.run(self.close_async())
2288
+
2289
+ def list_groups_sync(self, limit: Optional[int] = 100) -> List[Group]:
2290
+ """Synchronous wrapper for list_groups with _sync suffix."""
2291
+ return asyncio.run(self.list_groups_async(limit))
2292
+
2293
+ def search_groups_sync(
2294
+ self, keywords: str, limit: Optional[int] = 100, offset: Optional[int] = None
2295
+ ) -> List[Group]:
2296
+ """Synchronous wrapper for search_groups with _sync suffix."""
2297
+ return asyncio.run(self.search_groups_async(keywords, limit, offset))
2298
+
2299
+ def list_files_sync(
2300
+ self, group_id: str, file_group_id: Optional[str] = None
2301
+ ) -> List[FileInfo]:
2302
+ """Synchronous wrapper for list_files with _sync suffix."""
2303
+ return asyncio.run(self.list_files_async(group_id, file_group_id))
2304
+
2305
+ def check_availability_sync(
2306
+ self, file_group_id: str, file_datetime: str
2307
+ ) -> AvailabilityInfo:
2308
+ """Synchronous wrapper for check_availability with _sync suffix."""
2309
+ return asyncio.run(self.check_availability_async(file_group_id, file_datetime))
2310
+
2311
+ def download_file_sync(
2312
+ self,
2313
+ file_group_id: str,
2314
+ file_datetime: Optional[str] = None,
2315
+ destination_path: Optional[Path] = None,
2316
+ options: Optional[DownloadOptions] = None,
2317
+ num_parts: int = 5,
2318
+ progress_callback: Optional[Callable] = None,
2319
+ ) -> DownloadResult:
2320
+ """Synchronous wrapper for download_file with _sync suffix."""
2321
+ return asyncio.run(
2322
+ self.download_file_async(
2323
+ file_group_id,
2324
+ file_datetime,
2325
+ destination_path,
2326
+ options,
2327
+ num_parts,
2328
+ progress_callback,
2329
+ )
2330
+ )
2331
+
2332
+ def list_available_files_sync(
2333
+ self,
2334
+ group_id: str,
2335
+ file_group_id: Optional[str] = None,
2336
+ start_date: Optional[str] = None,
2337
+ end_date: Optional[str] = None,
2338
+ ) -> List[Dict[str, Any]]:
2339
+ """Synchronous wrapper for list_available_files with _sync suffix."""
2340
+ return asyncio.run(
2341
+ self.list_available_files_async(
2342
+ group_id, file_group_id, start_date, end_date
2343
+ )
2344
+ )
2345
+
2346
+ def health_check_sync(self) -> bool:
2347
+ """Synchronous wrapper for health_check with _sync suffix."""
2348
+ return asyncio.run(self.health_check_async())
2349
+
2350
+ def run_group_download_sync(
2351
+ self,
2352
+ group_id: str,
2353
+ start_date: str,
2354
+ end_date: str,
2355
+ destination_dir: Path = Path("./downloads"),
2356
+ max_concurrent: int = 5,
2357
+ ) -> dict:
2358
+ """Synchronous wrapper for run_group_download with _sync suffix."""
2359
+ return asyncio.run(
2360
+ self.run_group_download_async(
2361
+ group_id, start_date, end_date, destination_dir, max_concurrent
2362
+ )
2363
+ )
2364
+
2365
+ # Auto-Download wrappers
2366
+ async def start_auto_download_async(
2367
+ self,
2368
+ group_id: str,
2369
+ destination_dir: str = "./downloads",
2370
+ interval_minutes: int = 30,
2371
+ file_filter: Optional[Callable] = None,
2372
+ progress_callback: Optional[Callable] = None,
2373
+ error_callback: Optional[Callable] = None,
2374
+ max_retries: int = 3,
2375
+ check_current_date_only: bool = True,
2376
+ max_concurrent_downloads: Optional[int] = None,
2377
+ ):
2378
+ """Proxy to client's start_auto_download_async."""
2379
+ await self.connect_async()
2380
+ assert self._client is not None
2381
+ return await self._client.start_auto_download_async(
2382
+ group_id=group_id,
2383
+ destination_dir=destination_dir,
2384
+ interval_minutes=interval_minutes,
2385
+ file_filter=file_filter,
2386
+ progress_callback=progress_callback,
2387
+ error_callback=error_callback,
2388
+ max_retries=max_retries,
2389
+ check_current_date_only=check_current_date_only,
2390
+ max_concurrent_downloads=max_concurrent_downloads,
2391
+ )
2392
+
2393
+ def start_auto_download(
2394
+ self,
2395
+ group_id: str,
2396
+ destination_dir: str = "./downloads",
2397
+ interval_minutes: int = 30,
2398
+ file_filter: Optional[Callable] = None,
2399
+ progress_callback: Optional[Callable] = None,
2400
+ error_callback: Optional[Callable] = None,
2401
+ max_retries: int = 3,
2402
+ check_current_date_only: bool = True,
2403
+ max_concurrent_downloads: Optional[int] = None,
2404
+ ):
2405
+ """Synchronous proxy to client's start_auto_download_async."""
2406
+ return self._run_sync(
2407
+ self.start_auto_download_async(
2408
+ group_id,
2409
+ destination_dir,
2410
+ interval_minutes,
2411
+ file_filter,
2412
+ progress_callback,
2413
+ error_callback,
2414
+ max_retries,
2415
+ check_current_date_only,
2416
+ max_concurrent_downloads,
2417
+ )
2418
+ )
2419
+
2420
+ # DataFrame conversion proxies
2421
+ def to_dataframe(
2422
+ self,
2423
+ response_data,
2424
+ flatten_nested: bool = True,
2425
+ include_metadata: bool = False,
2426
+ date_columns: Optional[List[str]] = None,
2427
+ numeric_columns: Optional[List[str]] = None,
2428
+ custom_transformations: Optional[Dict[str, Callable]] = None,
2429
+ ):
2430
+ """Proxy to client's to_dataframe utility."""
2431
+ if self._client is None:
2432
+ # Use a temporary client for utilities if not connected yet
2433
+ self._client = DataQueryClient(self.client_config)
2434
+ return self._client.to_dataframe(
2435
+ response_data,
2436
+ flatten_nested=flatten_nested,
2437
+ include_metadata=include_metadata,
2438
+ date_columns=date_columns,
2439
+ numeric_columns=numeric_columns,
2440
+ custom_transformations=custom_transformations,
2441
+ )
2442
+
2443
+ def groups_to_dataframe(self, groups, include_metadata: bool = False):
2444
+ if self._client is None:
2445
+ self._client = DataQueryClient(self.client_config)
2446
+ return self._client.groups_to_dataframe(
2447
+ groups, include_metadata=include_metadata
2448
+ )
2449
+
2450
+ def files_to_dataframe(self, files, include_metadata: bool = False):
2451
+ if self._client is None:
2452
+ self._client = DataQueryClient(self.client_config)
2453
+ return self._client.files_to_dataframe(files, include_metadata=include_metadata)
2454
+
2455
+ def instruments_to_dataframe(self, instruments, include_metadata: bool = False):
2456
+ if self._client is None:
2457
+ self._client = DataQueryClient(self.client_config)
2458
+ return self._client.instruments_to_dataframe(
2459
+ instruments, include_metadata=include_metadata
2460
+ )
2461
+
2462
+ def time_series_to_dataframe(self, time_series, include_metadata: bool = False):
2463
+ if self._client is None:
2464
+ self._client = DataQueryClient(self.client_config)
2465
+ return self._client.time_series_to_dataframe(
2466
+ time_series, include_metadata=include_metadata
2467
+ )