File size: 5,484 Bytes
969891d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | """
Source Resolver for Multi-Source Data Ingestion
Handles resolution of various source URIs:
- Local paths: /path/to/file.geojson
- S3 URIs: s3://bucket/key.shp
- HTTP URLs: https://example.com/data.csv
"""
import logging
import os
import tempfile
import shutil
from pathlib import Path
from typing import Optional, Tuple
from urllib.parse import urlparse
import requests
logger = logging.getLogger(__name__)
class SourceResolver:
"""Resolves various source URIs to local file paths."""
def __init__(self, cache_dir: Optional[Path] = None):
"""
Initialize the source resolver.
Args:
cache_dir: Directory for caching downloaded files.
Defaults to a temp directory.
"""
self.cache_dir = cache_dir or Path(tempfile.gettempdir()) / "renewhere_cache"
self.cache_dir.mkdir(parents=True, exist_ok=True)
def resolve(self, uri: str) -> Tuple[Path, bool]:
"""
Resolve a URI to a local file path.
Args:
uri: Local path, S3 URI, or HTTP URL
Returns:
Tuple of (local_path, is_temporary)
- local_path: Path to the local file
- is_temporary: True if file was downloaded and should be cleaned up
"""
parsed = urlparse(uri)
if parsed.scheme == 's3':
return self._resolve_s3(uri, parsed)
elif parsed.scheme in ('http', 'https'):
return self._resolve_http(uri, parsed)
else:
# Assume local path
return self._resolve_local(uri)
def _resolve_local(self, path: str) -> Tuple[Path, bool]:
"""Resolve a local file path."""
local_path = Path(path).expanduser().resolve()
if not local_path.exists():
raise FileNotFoundError(f"Local file not found: {local_path}")
return local_path, False
def _resolve_s3(self, uri: str, parsed) -> Tuple[Path, bool]:
"""
Resolve an S3 URI to a local file.
Uses boto3 with credentials from environment variables:
- AWS_ACCESS_KEY_ID
- AWS_SECRET_ACCESS_KEY
- AWS_DEFAULT_REGION (optional)
"""
try:
import boto3
from botocore.exceptions import ClientError
except ImportError:
raise ImportError(
"boto3 is required for S3 support. Install with: pip install boto3"
)
bucket = parsed.netloc
key = parsed.path.lstrip('/')
if not bucket or not key:
raise ValueError(f"Invalid S3 URI: {uri}")
# Generate cache filename
safe_key = key.replace('/', '_')
cache_path = self.cache_dir / f"s3_{bucket}_{safe_key}"
# Check if already cached
if cache_path.exists():
logger.info(f"Using cached S3 file: {cache_path}")
return cache_path, True
logger.info(f"Downloading from S3: {uri}")
# Download from S3
s3 = boto3.client('s3')
try:
s3.download_file(bucket, key, str(cache_path))
except ClientError as e:
raise RuntimeError(f"Failed to download from S3: {e}")
logger.info(f"Downloaded to: {cache_path}")
return cache_path, True
def _resolve_http(self, uri: str, parsed) -> Tuple[Path, bool]:
"""Resolve an HTTP/HTTPS URL to a local file."""
# Generate cache filename from URL
safe_name = parsed.path.split('/')[-1] or 'download'
# Add hash of full URL to avoid collisions
url_hash = str(hash(uri))[-8:]
cache_path = self.cache_dir / f"http_{url_hash}_{safe_name}"
# Check if already cached
if cache_path.exists():
logger.info(f"Using cached HTTP file: {cache_path}")
return cache_path, True
logger.info(f"Downloading from HTTP: {uri}")
# Download with streaming
try:
response = requests.get(uri, stream=True, timeout=60)
response.raise_for_status()
with open(cache_path, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
except requests.RequestException as e:
raise RuntimeError(f"Failed to download from HTTP: {e}")
logger.info(f"Downloaded to: {cache_path}")
return cache_path, True
def cleanup(self, path: Path):
"""Remove a temporary file from cache."""
if path.exists() and str(path).startswith(str(self.cache_dir)):
try:
path.unlink()
logger.debug(f"Cleaned up: {path}")
except Exception as e:
logger.warning(f"Failed to cleanup {path}: {e}")
def clear_cache(self):
"""Clear all cached files."""
if self.cache_dir.exists():
shutil.rmtree(self.cache_dir)
self.cache_dir.mkdir(parents=True, exist_ok=True)
logger.info("Cache cleared")
# Singleton instance
_resolver: Optional[SourceResolver] = None
def get_source_resolver() -> SourceResolver:
"""Get the singleton source resolver instance."""
global _resolver
if _resolver is None:
_resolver = SourceResolver()
return _resolver
|