dataquery-sdk 0.0.7__py3-none-any.whl

dataquery/auto_download.py

@@ -0,0 +1,448 @@
+"""
+Auto-Download Manager for dataquery-sdk
+
+This module provides automatic file download functionality that continuously
+monitors data groups and downloads new files when they become available.
+"""
+
+import asyncio
+import logging
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Set
+
+from .config import EnvConfig
+from .models import DownloadOptions, DownloadProgress
+
+
+class AutoDownloadManager:
+    """
+    Manages automatic file download monitoring and downloading.
+
+    This class continuously monitors a data group for new files and automatically
+    downloads them if they don't already exist in the destination folder.
+
+    Features:
+    - Continuous monitoring with configurable intervals
+    - File filtering capabilities
+    - Progress tracking and callbacks
+    - Error handling and retry logic
+    - Statistics and monitoring
+    - Graceful shutdown
+    """
+
+    def __init__(
+        self,
+        client,  # DataQueryClient instance
+        group_id: str,
+        destination_dir: str = "./downloads",
+        interval_minutes: int = 30,
+        file_filter: Optional[Callable] = None,
+        progress_callback: Optional[Callable] = None,
+        error_callback: Optional[Callable] = None,
+        max_retries: int = 3,
+        check_current_date_only: bool = True,
+        max_concurrent_downloads: Optional[int] = None,
+    ):
+        """
+        Initialize the AutoDownloadManager.
+
+        Args:
+            client: DataQueryClient instance
+            group_id: ID of the data group to monitor
+            destination_dir: Directory to download files to
+            interval_minutes: Check interval in minutes
+            file_filter: Optional function to filter files
+            progress_callback: Optional callback for download progress
+            error_callback: Optional callback for errors
+            max_retries: Maximum retry attempts for failed downloads
+            check_current_date_only: If True, only check files for current date
+            max_concurrent_downloads: Maximum concurrent downloads (uses SDK default if None)
+        """
+        self.client = client
+        self.group_id = group_id
+        self.destination_dir = Path(destination_dir)
+        self.interval_minutes = interval_minutes
+        self.file_filter = file_filter
+        self.progress_callback = progress_callback
+        self.error_callback = error_callback
+        self.max_retries = max_retries
+        self.check_current_date_only = check_current_date_only
+        # Use SDK default if not specified
+        self.max_concurrent_downloads = (
+            max_concurrent_downloads
+            if max_concurrent_downloads is not None
+            else EnvConfig.get_int("MAX_CONCURRENT_DOWNLOADS", "5")
+        )
+
+        # Create destination directory if it doesn't exist
+        self.destination_dir.mkdir(parents=True, exist_ok=True)
+
+        # Internal state
+        self._running = False
+        self._task: Optional[asyncio.Task] = None
+        self._stop_event = asyncio.Event()
+        self._downloaded_files: Set[str] = set()
+        self._failed_files: Dict[str, int] = {}  # file_group_id -> retry_count
+
+        # Statistics
+        self.stats: Dict[str, Any] = {
+            "start_time": None,
+            "last_check_time": None,
+            "total_checks": 0,
+            "files_discovered": 0,
+            "files_downloaded": 0,
+            "files_skipped": 0,
+            "download_failures": 0,
+            "total_bytes_downloaded": 0,
+            "errors": [],
+        }
+
+        # Setup logging
+        self.logger = logging.getLogger(f"AutoDownloadManager.{group_id}")
+
+        # Download options
+        # Note: destination_path is used by the client to treat this as a directory target
+        self.download_options = DownloadOptions(
+            destination_path=self.destination_dir,
+            overwrite_existing=False,  # Don't overwrite existing files
+            chunk_size=8192,
+        )
+        # Backward-compat extra attribute used by tests; harmless extra field
+        setattr(self.download_options, "output_dir", str(self.destination_dir))
+
+    async def start(self):
+        """Start the auto-download monitoring."""
+        if self._running:
+            raise RuntimeError("AutoDownloadManager is already running")
+
+        self._running = True
+        self.stats["start_time"] = datetime.now()
+        self._stop_event.clear()
+
+        self.logger.info(
+            f"Starting auto-download for group '{self.group_id}' "
+            f"(interval: {self.interval_minutes} minutes)"
+        )
+
+        # Start the monitoring task
+        self._task = asyncio.create_task(self._monitoring_loop())
+
+        return self
+
+    async def stop(self):
+        """Stop the auto-download monitoring."""
+        if not self._running:
+            return
+
+        self.logger.info("Stopping auto-download monitoring...")
+
+        self._running = False
+        self._stop_event.set()
+
+        # Wait for the monitoring task to complete
+        if self._task:
+            try:
+                await asyncio.wait_for(self._task, timeout=5.0)
+            except asyncio.TimeoutError:
+                self.logger.warning(
+                    "Monitoring task did not stop gracefully, cancelling..."
+                )
+                self._task.cancel()
+                try:
+                    await self._task
+                except asyncio.CancelledError:
+                    pass
+
+        self.logger.info("Auto-download monitoring stopped")
+
+    async def _monitoring_loop(self):
+        """Main monitoring loop that runs continuously."""
+        try:
+            while self._running:
+                try:
+                    await self._check_and_download_files()
+                    self.stats["total_checks"] += 1
+                    self.stats["last_check_time"] = datetime.now()
+
+                except Exception as e:
+                    self.logger.error(f"Error in monitoring loop: {e}")
+                    self.stats["errors"].append(
+                        {
+                            "timestamp": datetime.now().isoformat(),
+                            "error": str(e),
+                            "type": "monitoring_loop",
+                        }
+                    )
+
+                    # Call error callback if provided
+                    if self.error_callback:
+                        try:
+                            await self._safe_call_callback(self.error_callback, e)
+                        except Exception as cb_error:
+                            self.logger.error(f"Error in error callback: {cb_error}")
+
+                # Wait for the specified interval or until stopped
+                try:
+                    await asyncio.wait_for(
+                        self._stop_event.wait(), timeout=self.interval_minutes * 60
+                    )
+                    # If we reach here, stop was requested
+                    break
+                except asyncio.TimeoutError:
+                    # Timeout is normal, continue to next iteration
+                    continue
+
+        except asyncio.CancelledError:
+            self.logger.info("Monitoring loop cancelled")
+        except Exception as e:
+            self.logger.error(f"Fatal error in monitoring loop: {e}")
+        finally:
+            self._running = False
+
+    async def _check_and_download_files(self):
+        """Use available-files to find and download files if needed."""
+        try:
+            if not self._running:
+                return
+
+            # Determine date window to check
+            dates_to_check = self._get_dates_to_check()
+            start_date = min(dates_to_check)
+            end_date = max(dates_to_check)
+
+            # Query available files for the date window
+            available_files = await self.client.list_available_files_async(
+                group_id=self.group_id,
+                file_group_id=None,
+                start_date=start_date,
+                end_date=end_date,
+            )
+
+            if not isinstance(available_files, list) or not available_files:
+                self.logger.debug(
+                    f"No available files for group '{self.group_id}' between {start_date} and {end_date}"
+                )
+                return
+
+            self.logger.debug(
+                f"Found {len(available_files)} available file entries for group '{self.group_id}'"
+            )
+
+            # Build eligible downloads
+            eligible: List[tuple[str, str, str]] = []  # (file_id, date_str, file_key)
+            for item in available_files:
+                if not self._running:
+                    break
+                file_id = item.get("file-group-id")
+                date_str = item.get("file-datetime")
+                avail_flag = item.get("is-available")
+                is_available = bool(avail_flag)
+                if not file_id or not date_str or not is_available:
+                    continue
+                if self.file_filter and not self.file_filter(item):
+                    continue
+                file_key = f"{file_id}_{date_str}"
+                if file_key in self._downloaded_files:
+                    continue
+                if self._failed_files.get(file_key, 0) >= self.max_retries:
+                    continue
+                if self._file_exists_locally(file_id, date_str):
+                    self.stats["files_skipped"] += 1
+                    self._downloaded_files.add(file_key)
+                    continue
+                eligible.append((file_id, date_str, file_key))
+
+            if not eligible:
+                return
+
+            # Parallel downloads with concurrency limit
+            semaphore = asyncio.Semaphore(max(1, int(self.max_concurrent_downloads)))
+
+            async def worker(fid: str, dstr: str, fkey: str):
+                if not self._running:
+                    return
+                async with semaphore:
+                    await self._download_file(fid, dstr, fkey)
+
+            tasks = [
+                asyncio.create_task(worker(fid, dstr, fkey))
+                for fid, dstr, fkey in eligible
+            ]
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+        except Exception as e:
+            self.logger.error(f"Error checking available files: {e}")
+            raise
+
+    async def _process_file(self, file_info):
+        """Deprecated path; available-files flow supersedes this."""
+        return
+
+    def _get_dates_to_check(self) -> List[str]:
+        """Get list of dates to check for files."""
+        if self.check_current_date_only:
+            # Only check current date
+            return [datetime.now().strftime("%Y%m%d")]
+        else:
+            # Check current date and previous few days
+            dates: List[str] = []
+            today = datetime.now()
+            for i in range(3):  # Check today and last 2 days
+                date = today - timedelta(days=i)
+                dates.append(date.strftime("%Y%m%d"))
+            return dates
+
+    async def _check_and_download_file_for_date(self, file_id: str, date_str: str):
+        """Deprecated path; available-files flow supersedes this."""
+        return
+
+    def _file_exists_locally(self, file_id: str, date_str: str) -> bool:
+        """Check if file already exists in the destination directory."""
+        # Look for files that might match this file_group_id and date
+        # This is a heuristic approach since we don't know the exact filename
+
+        for file_path in self.destination_dir.glob("*"):
+            if file_path.is_file():
+                filename = file_path.name
+                # Check if filename contains the file_group_id and date
+                if file_id in filename and date_str in filename:
+                    return True
+
+        return False
+
+    async def _download_file(self, file_id: str, date_str: str, file_key: str):
+        """Download a file."""
+        self.logger.info(f"Downloading '{file_id}' for {date_str}...")
+
+        # Create progress callback that updates statistics
+        def progress_wrapper(progress: DownloadProgress):
+            # Update statistics
+            current_total = int(self.stats.get("total_bytes_downloaded", 0) or 0)
+            self.stats["total_bytes_downloaded"] = max(
+                current_total, int(getattr(progress, "bytes_downloaded", 0) or 0)
+            )
+
+            # Call user callback if provided
+            if self.progress_callback:
+                try:
+                    if asyncio.iscoroutinefunction(self.progress_callback):
+                        asyncio.create_task(self.progress_callback(progress))
+                    else:
+                        self.progress_callback(progress)
+                except Exception as e:
+                    self.logger.error(f"Error in progress callback: {e}")
+
+        try:
+            result = await self.client.download_file_async(
+                file_group_id=file_id,
+                file_datetime=date_str,
+                options=self.download_options,
+                progress_callback=progress_wrapper,
+            )
+
+            succeeded = False
+            if result is not None and getattr(result, "status", None) is not None:
+                try:
+                    from .models import (  # import inside to avoid cycles in some tools
+                        DownloadStatus,
+                    )
+
+                    succeeded = result.status == DownloadStatus.COMPLETED
+                except Exception:
+                    succeeded = False
+
+            if succeeded:
+                self.logger.info(f"Successfully downloaded '{file_id}' for {date_str}")
+                self.stats["files_downloaded"] += 1
+                self._downloaded_files.add(file_key)
+                # Accumulate total bytes if provided
+                try:
+                    if getattr(result, "file_size", None):
+                        self.stats["total_bytes_downloaded"] += (
+                            int(result.file_size) or 0
+                        )
+                except Exception:
+                    pass
+
+                # Reset failure count on success
+                if file_key in self._failed_files:
+                    del self._failed_files[file_key]
+
+            else:
+                error_msg = (
+                    getattr(result, "error_message", "Unknown error")
+                    if result
+                    else "No result"
+                )
+                self.logger.error(
+                    f"Download failed for '{file_id}' for {date_str}: {error_msg}"
+                )
+                self._failed_files[file_key] = self._failed_files.get(file_key, 0) + 1
+                self.stats["download_failures"] += 1
+
+        except Exception as e:
+            self.logger.error(f"Exception downloading '{file_id}' for {date_str}: {e}")
+            self._failed_files[file_key] = self._failed_files.get(file_key, 0) + 1
+            self.stats["download_failures"] += 1
+            raise
+
+    async def _safe_call_callback(self, callback: Callable, *args, **kwargs):
+        """Safely call a callback function."""
+        try:
+            if asyncio.iscoroutinefunction(callback):
+                await callback(*args, **kwargs)
+            else:
+                callback(*args, **kwargs)
+        except Exception as e:
+            self.logger.error(f"Error in callback: {e}")
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get auto-download statistics."""
+        runtime = None
+        if self.stats["start_time"]:
+            try:
+                start = self.stats["start_time"]
+                if isinstance(start, datetime):
+                    runtime = (datetime.now() - start).total_seconds()
+            except Exception:
+                runtime = None
+
+        return {
+            **self.stats,
+            "runtime_seconds": runtime,
+            "is_running": self._running,
+            "downloaded_files_count": len(self._downloaded_files),
+            "failed_files_count": len(self._failed_files),
+            "group_id": self.group_id,
+            "destination_dir": str(self.destination_dir),
+            "interval_minutes": self.interval_minutes,
+            "check_current_date_only": self.check_current_date_only,
+        }
+
+    def get_downloaded_files(self) -> Set[str]:
+        """Get set of successfully downloaded file keys."""
+        return self._downloaded_files.copy()
+
+    def get_failed_files(self) -> Dict[str, int]:
+        """Get dictionary of failed file keys and their retry counts."""
+        return self._failed_files.copy()
+
+    @property
+    def is_running(self) -> bool:
+        """Check if the auto-download manager is running."""
+        return self._running
+
+    def __str__(self) -> str:
+        """String representation of the manager."""
+        status = "running" if self._running else "stopped"
+        return f"AutoDownloadManager(group='{self.group_id}', status={status}, interval={self.interval_minutes}min)"
+
+    def __repr__(self) -> str:
+        """Detailed representation of the manager."""
+        return (
+            f"AutoDownloadManager(group_id='{self.group_id}', "
+            f"destination_dir='{self.destination_dir}', "
+            f"interval_minutes={self.interval_minutes}, "
+            f"running={self._running})"
+        )

dataquery/cli.py

@@ -0,0 +1,244 @@
+"""Command Line Interface for DataQuery SDK."""
+
+import argparse
+import asyncio
+import json
+import sys
+from pathlib import Path
+
+from dataquery import DataQuery
+
+
+def create_parser() -> argparse.ArgumentParser:
+    """Create the top-level CLI parser with subcommands."""
+    parser = argparse.ArgumentParser(
+        description="Command Line Interface for the DataQuery SDK"
+    )
+    parser.add_argument("--env-file", type=str, default=None, help="Path to .env file")
+    subparsers = parser.add_subparsers(dest="command")
+
+    # groups
+    p_groups = subparsers.add_parser("groups", help="List or search groups")
+    p_groups.add_argument("--json", action="store_true", help="Output JSON")
+    p_groups.add_argument(
+        "--limit", type=int, default=None, help="Limit number of results"
+    )
+    p_groups.add_argument("--search", type=str, default=None, help="Search keywords")
+
+    # files
+    p_files = subparsers.add_parser("files", help="List files in a group")
+    p_files.add_argument("--group-id", required=True)
+    p_files.add_argument("--file-group-id", default=None)
+    p_files.add_argument("--limit", type=int, default=None)
+    p_files.add_argument("--json", action="store_true")
+
+    # availability
+    p_avail = subparsers.add_parser("availability", help="Check file availability")
+    p_avail.add_argument("--file-group-id", required=True)
+    p_avail.add_argument("--file-datetime", required=True)
+    p_avail.add_argument("--json", action="store_true")
+
+    # download
+    p_dl = subparsers.add_parser("download", help="Download a file or start watch mode")
+    p_dl.add_argument("--file-group-id", default=None)
+    p_dl.add_argument("--file-datetime", default=None)
+    p_dl.add_argument("--destination", type=str, default=None)
+    p_dl.add_argument(
+        "--watch", action="store_true", help="Watch a group for new files"
+    )
+    p_dl.add_argument("--group-id", default=None, help="Required with --watch")
+    p_dl.add_argument("--json", action="store_true")
+
+    # config
+    p_cfg = subparsers.add_parser("config", help="Config utilities")
+    cfg_sub = p_cfg.add_subparsers(dest="config_command")
+    _ = cfg_sub.add_parser("show", help="Show resolved config")
+    _ = cfg_sub.add_parser("validate", help="Validate config")
+    p_tmpl = cfg_sub.add_parser("template", help="Write .env template")
+    p_tmpl.add_argument("--output", type=str, required=True)
+
+    # auth
+    p_auth = subparsers.add_parser("auth", help="Auth utilities")
+    auth_sub = p_auth.add_subparsers(dest="auth_command")
+    _ = auth_sub.add_parser("test", help="Test authentication by listing groups")
+
+    return parser
+
+
+async def cmd_groups(args: argparse.Namespace) -> int:
+    async with DataQuery(args.env_file) as dq:
+        if args.search:
+            items = await dq.search_groups_async(args.search, limit=args.limit)
+        else:
+            items = await dq.list_groups_async(limit=args.limit)
+        if args.json:
+            payload = []
+            for g in items:
+                try:
+                    payload.append(g.model_dump())
+                except Exception:
+                    payload.append(str(g))
+            print(json.dumps(payload, indent=2))
+        else:
+            for g in items:
+                try:
+                    d = g.model_dump()
+                    print(
+                        f"{d.get('group_id') or d.get('group-id')}\t{d.get('group_name') or d.get('group-name')}"
+                    )
+                except Exception:
+                    print(str(g))
+    return 0
+
+
+async def cmd_files(args: argparse.Namespace) -> int:
+    async with DataQuery(args.env_file) as dq:
+        files = await dq.list_files_async(args.group_id, args.file_group_id)
+        if args.json:
+            payload = []
+            for f in files:
+                try:
+                    payload.append(f.model_dump())
+                except Exception:
+                    payload.append(str(f))
+            print(json.dumps(payload, indent=2))
+        else:
+            print(f"Found {len(files)} files")
+            for f in files:
+                try:
+                    d = f.model_dump()
+                    print(
+                        f"{d.get('file_group_id') or d.get('file-group-id')}\t{d.get('file_type')}"
+                    )
+                except Exception:
+                    print(str(f))
+    return 0
+
+
+async def cmd_availability(args: argparse.Namespace) -> int:
+    async with DataQuery(args.env_file) as dq:
+        avail = await dq.check_availability_async(
+            args.file_group_id, args.file_datetime
+        )
+        if args.json:
+            try:
+                print(json.dumps(getattr(avail, "model_dump")(), indent=2))
+            except Exception:
+                print(str(avail))
+        else:
+            print(f"{args.file_group_id} @ {args.file_datetime}")
+    return 0
+
+
+async def cmd_download(args: argparse.Namespace) -> int:
+    # Watch mode requires group-id
+    if args.watch and not args.group_id:
+        print("--group-id is required when using --watch")
+        return 1
+
+    async with DataQuery(args.env_file) as dq:
+        if args.watch:
+            try:
+                mgr = await dq.start_auto_download_async(
+                    group_id=args.group_id,
+                    destination_dir=(args.destination or "./downloads"),
+                )
+                # Simulate quick watch then Ctrl+C
+                await asyncio.sleep(1)
+            except KeyboardInterrupt:
+                try:
+                    await mgr.stop()
+                except Exception:
+                    pass
+                stats: dict = getattr(mgr, "get_stats", lambda: {})()
+                print(json.dumps(stats))
+                return 0
+            return 0
+        # single download
+        dest_path = Path(args.destination) if args.destination else None
+        result = await dq.download_file_async(
+            args.file_group_id, args.file_datetime, dest_path, None
+        )
+        if args.json:
+            print(json.dumps(getattr(result, "model_dump")(), indent=2))
+        else:
+            print(f"Downloaded to {result.local_path}")
+        return 0
+
+
+def cmd_config_show(args: argparse.Namespace) -> int:
+    from dataquery.config import EnvConfig
+
+    EnvConfig.create_client_config(
+        env_file=Path(args.env_file) if getattr(args, "env_file", None) else None
+    )
+    print("Configuration loaded")
+    return 0
+
+
+def cmd_config_validate(args: argparse.Namespace) -> int:
+    from dataquery.config import (
+        EnvConfig,
+    )
+
+    try:
+        EnvConfig.validate_config(EnvConfig.create_client_config())
+        print("Configuration valid")
+        return 0
+    except Exception as e:
+        print(f"Configuration invalid: {e}")
+        return 1
+
+
+def cmd_config_template(args: argparse.Namespace) -> int:
+    # Import inside function to allow monkeypatch of dataquery.utils.create_env_template
+    import dataquery.utils as utils
+
+    out = utils.create_env_template(Path(args.output))
+    print(f"Template written to {out}")
+    return 0
+
+
+async def cmd_auth_test(args: argparse.Namespace) -> int:
+    async with DataQuery(args.env_file) as dq:
+        _ = await dq.list_groups_async(limit=1)
+    return 0
+
+
+def main_sync(ns: argparse.Namespace) -> int:
+    if ns.command == "config":
+        if ns.config_command == "show":
+            return cmd_config_show(ns)
+        if ns.config_command == "validate":
+            return cmd_config_validate(ns)
+        if ns.config_command == "template":
+            return cmd_config_template(ns)
+        return 1
+    return 0
+
+
+def main() -> int:
+    parser = create_parser()
+    args = parser.parse_args()
+    if not args.command:
+        parser.print_help()
+        return 1
+    # Dispatch
+    if args.command == "groups":
+        return asyncio.run(cmd_groups(args))
+    if args.command == "files":
+        return asyncio.run(cmd_files(args))
+    if args.command == "availability":
+        return asyncio.run(cmd_availability(args))
+    if args.command == "download":
+        return asyncio.run(cmd_download(args))
+    if args.command == "config":
+        return main_sync(args)
+    if args.command == "auth" and args.auth_command == "test":
+        return asyncio.run(cmd_auth_test(args))
+    parser.print_help()
+    return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())