earthscope-sdk 1.0.0b1__py3-none-any.whl → 1.2.0b0__py3-none-any.whl

earthscope_sdk/client/discovery/models.py

@@ -0,0 +1,144 @@
+from datetime import timedelta
+from enum import Enum
+from typing import Annotated, Any, Generic, Iterable, Optional, TypeVar, Union
+
+from pydantic import BaseModel, BeforeValidator, TypeAdapter
+
+from earthscope_sdk.util._itertools import to_set
+
+
+def _coerce_timedelta_ms(v: Union[int, float, timedelta, str]) -> timedelta:
+    if isinstance(v, (int, float)):
+        return timedelta(milliseconds=v)
+
+    # fallback to Pydantic's timedelta parser
+    return v
+
+
+P = TypeVar("P")
+
+
+class Page(BaseModel, Generic[P]):
+    has_next: bool
+    offset: int
+    limit: int
+    items: list[P]
+    total: Optional[int] = None
+
+
+class DatasourceBaseModel(BaseModel):
+    edid: str
+    names: dict[str, str]
+    description: Optional[str] = None
+
+    def to_arrow_columns(
+        self,
+        *,
+        fields: Union[list[str], str] = ["edid", "names"],
+        namespaces: Union[list[str], str] = [],
+    ) -> dict[str, Any]:
+        """
+        Convert the datasource model to a dictionary suitable for use in an Arrow table.
+        """
+        result = {}
+        namespaces = to_set(namespaces)
+        fields = to_set(fields)
+
+        # Add names to fields if namespaces are requested
+        if namespaces:
+            fields.add("names")
+
+        for field in fields:
+            if field != "names":
+                result[field] = getattr(self, field)
+                continue
+
+            # Explode names to own columns
+            if not namespaces:
+                names = {k.lower(): v for k, v in self.names.items()}
+            else:
+                names = {
+                    k_lower: v
+                    for k, v in self.names.items()
+                    if (k_lower := k.lower()) in namespaces
+                }
+
+            result.update(names)
+
+        return result
+
+
+class NetworkDatasource(DatasourceBaseModel): ...
+
+
+ListNetworkDatasourcesResult = TypeAdapter(Union[Page[str], Page[NetworkDatasource]])
+
+
+class StationDatasource(DatasourceBaseModel):
+    network_edids: Optional[list[str]] = None
+    networks: Optional[list[NetworkDatasource]] = None
+
+
+ListStationDatasourcesResult = TypeAdapter(Union[Page[str], Page[StationDatasource]])
+
+
+class _StationDatasourceMember(DatasourceBaseModel):
+    station_edid: Optional[str] = None
+    station: Optional[StationDatasource] = None
+
+    def to_arrow_columns(
+        self,
+        *,
+        fields: list[str] = ["edid", "names"],
+        namespaces: Optional[list[str]] = None,
+    ) -> dict[str, Any]:
+        result = super().to_arrow_columns(fields=fields, namespaces=namespaces)
+        if self.station:
+            parent_columns = self.station.to_arrow_columns(
+                fields=["names"],
+                namespaces=namespaces,
+            )
+            result.update(parent_columns)
+
+        return result
+
+
+class SessionDatasource(_StationDatasourceMember):
+    sample_interval: Annotated[timedelta, BeforeValidator(_coerce_timedelta_ms)]
+    """
+    Session sample interval.
+    """
+
+    roll: timedelta  # already in seconds
+    """
+    Session file roll cadence.
+    """
+
+
+ListSessionDatasourcesResult = TypeAdapter(Union[Page[str], Page[SessionDatasource]])
+
+
+class StreamType(Enum):
+    GNSS_RAW = "gnss_raw"
+    GNSS_PPP = "gnss_ppp"
+
+
+class StreamDatasource(_StationDatasourceMember):
+    stream_type: StreamType
+    facility: str
+    software: str
+    label: str
+    sample_interval: Annotated[timedelta, BeforeValidator(_coerce_timedelta_ms)]
+    """
+    Stream sample interval.
+    """
+
+    def to_arrow_columns(
+        self,
+        *,
+        fields: Iterable[str] = ["edid", "names", "facility", "software", "label"],
+    ) -> dict[str, Any]:
+        return super().to_arrow_columns(fields=fields)
+
+
+ListStreamDatasourcesResult = TypeAdapter(Union[Page[str], Page[StreamDatasource]])

earthscope_sdk/common/context.py

@@ -1,5 +1,6 @@
+import asyncio
 import sys
-from functools import cached_property
+from functools import cached_property, partial
 from typing import TYPE_CHECKING, Any, Callable, Coroutine, Optional, TypeVar, cast
 
 if sys.version_info >= (3, 10):
@@ -9,6 +10,8 @@ else:
 
 
 if TYPE_CHECKING:
+    from concurrent.futures import ThreadPoolExecutor
+
     from httpx import AsyncClient
 
     from earthscope_sdk.common._sync_runner import SyncRunner
@@ -64,6 +67,19 @@ class SdkContext:
 
         return DeviceCodeFlow(ctx=self)
 
+    @cached_property
+    def executor(self):
+        """
+        Thread pool executor for running sync functions in the background
+        """
+        import concurrent.futures
+
+        self._executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=self.settings.thread_pool_max_workers
+        )
+
+        return self._executor
+
     @cached_property
     def httpx_client(self):
         """
@@ -78,6 +94,8 @@ class SdkContext:
         self._httpx_client = httpx.AsyncClient(
             auth=self.auth_flow,
             headers={
+                **self.settings.http.extra_headers,
+                # override anything specified via extra_headers
                 "user-agent": self.settings.http.user_agent,
             },
             limits=self.settings.http.limits,
@@ -103,6 +121,18 @@ class SdkContext:
         """
         return self._settings
 
+    @cached_property
+    def _rate_limit(self):
+        """
+        Rate limit semaphore for rate limiting HTTP requests
+        """
+        from earthscope_sdk.util._concurrency import RateLimitSemaphore
+
+        return RateLimitSemaphore(
+            max_concurrent=self.settings.http.rate_limit.max_concurrent,
+            max_per_second=self.settings.http.rate_limit.max_per_second,
+        )
+
     def __init__(
         self,
         settings: Optional["SdkSettings"] = None,
@@ -113,6 +143,7 @@ class SdkContext:
         from earthscope_sdk.config.settings import SdkSettings
 
         # Local state
+        self._executor: Optional["ThreadPoolExecutor"] = None
         self._httpx_client: Optional["AsyncClient"] = None
         self._runner: Optional["SyncRunner"] = runner
         self._settings = settings or SdkSettings()
@@ -127,6 +158,9 @@ class SdkContext:
         if self._runner:
             self._runner.stop()
 
+        if self._executor:
+            self._executor.shutdown()
+
     def close(self):
         """
         Close this SdkContext to release underlying resources (e.g. connection pools)
@@ -141,6 +175,44 @@ class SdkContext:
         if self._runner:
             self._runner.stop()
 
+        if self._executor:
+            self._executor.shutdown()
+
+    async def run_in_executor(self, fn: Callable[..., Any], *args, **kwargs):
+        """
+        Run a function in the executor
+        """
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(self.executor, partial(fn, *args, **kwargs))
+
+    def asyncify(
+        self,
+        fn: Callable[T_ParamSpec, T_Retval],
+    ) -> Callable[T_ParamSpec, Coroutine[Any, Any, T_Retval]]:
+        """
+        Decorator that wraps a sync function to run in the executor as a coroutine.
+
+        Usage:
+            @ctx.asyncify
+            def some_sync_function(arg1, arg2):
+                # sync code here
+                return result
+
+            # Now some_sync_function is async and runs in executor
+            result = await some_sync_function(arg1, arg2)
+        """
+
+        async def wrapper(
+            *args: T_ParamSpec.args, **kwargs: T_ParamSpec.kwargs
+        ) -> T_Retval:
+            loop = asyncio.get_running_loop()
+            return await loop.run_in_executor(
+                self.executor,
+                partial(fn, *args, **kwargs),
+            )
+
+        return wrapper
+
     def syncify(
         self,
         async_function: Callable[T_ParamSpec, Coroutine[Any, Any, T_Retval]],

earthscope_sdk/common/service.py

@@ -36,17 +36,19 @@ class SdkService:
 
         Performs common response handling.
         """
-        resp = await self.ctx.httpx_client.send(request=request)
+        # Global rate limiting
+        async with self._ctx._rate_limit:
+            resp = await self.ctx.httpx_client.send(request=request)
 
-        # Throw specific errors for certain status codes
+            # Throw specific errors for certain status codes
 
-        if resp.status_code == 401:
-            await resp.aread()  # must read body before using .text prop
-            raise UnauthenticatedError(resp.text)
+            if resp.status_code == 401:
+                await resp.aread()  # must read body before using .text prop
+                raise UnauthenticatedError(resp.text)
 
-        if resp.status_code == 403:
-            await resp.aread()  # must read body before using .text prop
-            raise UnauthorizedError(resp.text)
+            if resp.status_code == 403:
+                await resp.aread()  # must read body before using .text prop
+                raise UnauthorizedError(resp.text)
 
         # Raise HTTP errors
         resp.raise_for_status()

earthscope_sdk/config/_bootstrap.py

@@ -0,0 +1,42 @@
+"""
+This module facilitates bootstrapping SDK settings from a JSON-encoded environment variable.
+"""
+
+import json
+import logging
+import os
+
+from pydantic_settings import PydanticBaseSettingsSource
+
+logger = logging.getLogger(__name__)
+
+
+class BootstrapEnvironmentSettingsSource(PydanticBaseSettingsSource):
+    """
+    This SettingsSource facilitates bootstrapping the SDK from a special environment variable.
+
+    The environment variable should be a JSON string of the expected SDK settings and structure.
+    """
+
+    def __init__(self, settings_cls, env_var: str):
+        super().__init__(settings_cls)
+        self._env_var = env_var
+
+    def __call__(self):
+        try:
+            bootstrap_settings = os.environ[self._env_var]
+        except KeyError:
+            return {}
+
+        try:
+            return json.loads(bootstrap_settings)
+        except json.JSONDecodeError:
+            logger.warning(
+                f"Found bootstrap environment variable '{self._env_var}', but unable to decode content as JSON"
+            )
+            return {}
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(env_var='{self._env_var}')"
+
+    def get_field_value(self, *args, **kwargs): ...  # unused abstract method

earthscope_sdk/config/models.py

@@ -1,13 +1,14 @@
 import base64
 import binascii
 import datetime as dt
+import fnmatch
 import functools
 from contextlib import suppress
 from enum import Enum
 from functools import cached_property
 from typing import Annotated, Any, Optional, Type, Union
 
-from annotated_types import Ge, Gt
+from annotated_types import Ge, Gt, Interval
 from pydantic import (
     AliasChoices,
     BaseModel,
@@ -15,14 +16,12 @@ from pydantic import (
     ConfigDict,
     Field,
     HttpUrl,
-    SecretStr,
-    SerializationInfo,
     ValidationError,
-    field_serializer,
     model_validator,
 )
 
 from earthscope_sdk import __version__
+from earthscope_sdk.model.secret import SecretStr
 
 
 def _try_float(v: Any):
@@ -94,23 +93,6 @@ class Tokens(BaseModel):
 
         raise ValueError("Unable to decode access token body")
 
-    @field_serializer("access_token", "id_token", "refresh_token", when_used="json")
-    def dump_secret_json(self, secret: Optional[SecretStr], info: SerializationInfo):
-        """
-        A special field serializer to dump the actual secret value when writing to JSON.
-
-        Only writes secret in plaintext when `info.context == "plaintext".
-
-        See [Pydantic docs](https://docs.pydantic.dev/latest/concepts/serialization/#serialization-context)
-        """
-        if secret is None:
-            return None
-
-        if info.context == "plaintext":
-            return secret.get_secret_value()
-
-        return str(secret)
-
     @model_validator(mode="after")
     def ensure_one_of(self):
         # allow all fields to be optional in subclasses
@@ -207,6 +189,12 @@ class AuthFlowSettings(Tokens):
     scope: str = "offline_access"
     client_secret: Optional[SecretStr] = None
 
+    # Only inject bearer token for requests to these hosts
+    allowed_hosts: set[str] = {
+        "earthscope.org",
+        "*.earthscope.org",
+    }
+
     # Auth exchange retries
     retry: HttpRetrySettings = HttpRetrySettings(
         attempts=5,
@@ -222,6 +210,46 @@ class AuthFlowSettings(Tokens):
 
         return AuthFlowType.DeviceCode
 
+    @cached_property
+    def allowed_host_patterns(self) -> set[str]:
+        """
+        The subset of allowed hosts that are glob patterns.
+
+        Use `is_host_allowed` to check if a host is allowed by any of these patterns.
+        """
+        return {h for h in self.allowed_hosts if "*" in h or "?" in h}
+
+    def is_host_allowed(self, host: str) -> bool:
+        """
+        Check if a host matches any pattern in the allowed hosts set.
+
+        Supports glob patterns with '?' and '*' characters (e.g., *.earthscope.org).
+
+        Args:
+            host: The hostname to check
+
+        Returns:
+            True if the host matches any allowed pattern, False otherwise
+        """
+        if host in self.allowed_hosts:
+            return True
+
+        for allowed_pattern in self.allowed_host_patterns:
+            if fnmatch.fnmatch(host, allowed_pattern):
+                self.allowed_hosts.add(host)
+                return True
+
+        return False
+
+
+class RateLimitSettings(BaseModel):
+    """
+    Rate limit settings
+    """
+
+    max_concurrent: Annotated[int, Interval(ge=1, le=200)] = 100
+    max_per_second: Annotated[float, Interval(ge=1, le=200)] = 150.0
+
 
 class HttpSettings(BaseModel):
     """
@@ -240,8 +268,12 @@ class HttpSettings(BaseModel):
     # automatically retry requests
     retry: HttpRetrySettings = HttpRetrySettings()
 
+    # rate limit outgoing requests
+    rate_limit: RateLimitSettings = RateLimitSettings()
+
     # Other
     user_agent: str = f"earthscope-sdk py/{__version__}"
+    extra_headers: dict[str, str] = {}
 
     @cached_property
     def limits(self):
@@ -289,3 +321,4 @@ class SdkBaseSettings(BaseModel):
     http: HttpSettings = HttpSettings()
     oauth2: AuthFlowSettings = AuthFlowSettings()
     resources: ResourceRefs = ResourceRefs()
+    thread_pool_max_workers: Optional[int] = None

earthscope_sdk/config/settings.py

@@ -11,11 +11,15 @@ from pydantic_settings import (
     TomlConfigSettingsSource,
 )
 
+from earthscope_sdk.config._bootstrap import BootstrapEnvironmentSettingsSource
 from earthscope_sdk.config._compat import LegacyEarthScopeCLISettingsSource
 from earthscope_sdk.config._util import deep_merge, get_config_dir, slugify
 from earthscope_sdk.config.error import ProfileDoesNotExistError
 from earthscope_sdk.config.models import SdkBaseSettings, Tokens
 
+_BOOTSTRAP_ENV_VAR = "ES_BOOTSTRAP_SETTINGS"
+"""Environment variable for bootstrapping the SDK"""
+
 _DEFAULT_PROFILE = "default"
 """Default profile name"""
 
@@ -269,6 +273,12 @@ class SdkSettings(SdkBaseSettings, BaseSettings):
         alias = SdkSettings.model_fields["profile_name"].validation_alias
         global_settings = _GlobalSettingsSource(settings_cls, "profile_name", alias)
 
+        # Check for bootstrapping configuration
+        bootstrap_settings = BootstrapEnvironmentSettingsSource(
+            settings_cls,
+            _BOOTSTRAP_ENV_VAR,
+        )
+
         # Compatibility with earthscope-cli v0.x.x state:
         # If we find this file, we only care about the access and refresh tokens
         keep_keys = {"access_token", "refresh_token"}
@@ -281,4 +291,5 @@ class SdkSettings(SdkBaseSettings, BaseSettings):
             dotenv_settings,
             global_settings,
             legacy_settings,
+            bootstrap_settings,
         )

earthscope_sdk/model/secret.py

@@ -0,0 +1,29 @@
+from typing import Annotated
+
+from pydantic import PlainSerializer, SerializationInfo
+from pydantic import SecretStr as _SecretStr
+
+
+def _dump_secret_plaintext(secret: _SecretStr, info: SerializationInfo):
+    """
+    A special field serializer to dump the actual secret value.
+
+    Only writes secret in plaintext when `info.context == "plaintext".
+
+    See [Pydantic docs](https://docs.pydantic.dev/latest/concepts/serialization/#serialization-context)
+    """
+
+    if info.context == "plaintext":
+        return secret.get_secret_value()
+
+    return str(secret)
+
+
+SecretStr = Annotated[
+    _SecretStr,
+    PlainSerializer(
+        _dump_secret_plaintext,
+        return_type=str,
+        when_used="json-unless-none",
+    ),
+]

earthscope_sdk/util/_concurrency.py

@@ -0,0 +1,64 @@
+"""
+Rate-limited semaphore for controlling concurrent operations.
+"""
+
+import asyncio
+import time
+
+
+class RateLimitSemaphore(asyncio.Semaphore):
+    """
+    A semaphore that limits both the maximum number of concurrent acquisitions
+    and the maximum rate at which acquisitions can occur.
+    """
+
+    @property
+    def max_concurrent(self) -> int:
+        """Maximum number of concurrent acquisitions."""
+        return self._max_concurrent
+
+    @property
+    def max_per_second(self) -> float:
+        """Maximum acquisitions per second."""
+        return self._max_per_second
+
+    def __init__(self, max_concurrent: int, max_per_second: float):
+        """
+        Initialize the rate-limited semaphore.
+
+        Args:
+            max_concurrent: Maximum number of concurrent acquisitions.
+            max_per_second: Maximum acquisitions per second.
+        """
+        super().__init__(max_concurrent)
+        self._max_concurrent = max_concurrent
+        self._max_per_second = max_per_second
+
+        self._acquire_interval = 1.0 / max_per_second
+        self._next_acquire_time = 0
+        self._rate_lock = asyncio.Lock()
+
+    async def acquire(self):
+        """
+        Acquire a semaphore with rate limiting.
+
+        In addition to traditional semaphore behavior, this will also rate limit
+        the number of acquisitions per second by blocking until the next
+        acquisition interval.
+        """
+        # Rate limit acquisitions (serialized to prevent race conditions)
+        async with self._rate_lock:
+            now = time.monotonic()
+            wait = self._next_acquire_time - now
+
+            if wait > 0:
+                # Sleep until next acquisition interval
+                await asyncio.sleep(wait)
+                self._next_acquire_time += self._acquire_interval
+
+            else:
+                # no wait, set next acquisition time to next interval
+                self._next_acquire_time = now + self._acquire_interval
+
+        # Acquire semaphore for concurrency limiting
+        return await super().acquire()

earthscope_sdk/util/_itertools.py

@@ -0,0 +1,57 @@
+from itertools import islice
+from typing import Generator, Iterable, TypeVar, Union
+
+T = TypeVar("T")
+
+
+def batched(
+    iterable: Iterable[T],
+    n: int = 10,
+) -> Generator[tuple[T, ...], None, None]:
+    """Process an iterable as batches of size `n`
+
+    Args:
+        iterable (Iterable[T]): anything iterable
+        n (Optional[int]): the size of tuples to yield
+
+    Yields (tuple[T, ...]) tuple of size `n` elements from `iterable`
+
+    Example:
+
+    ```py
+    my_list = [1,2,3,4,5]
+    print(list(batched(my_list, 3))) # prints [(1,2,3), (4,5)]
+    ```
+    """
+    if n < 1:
+        raise ValueError("n must be at least 1")
+
+    iterator = iter(iterable)
+    while batch := tuple(islice(iterator, n)):
+        yield batch
+
+
+def to_list(maybe_list: Union[T, list[T], set[T]]) -> list[T]:
+    """
+    Coerce the argument into a list if it is not already
+    """
+    if isinstance(maybe_list, list):
+        return maybe_list
+
+    if isinstance(maybe_list, set):
+        return list(maybe_list)
+
+    return [maybe_list]
+
+
+def to_set(maybe_set: Union[T, list[T], set[T]]) -> set[T]:
+    """
+    Coerce the argument into a set if it is not already
+    """
+    if isinstance(maybe_set, set):
+        return maybe_set
+
+    if isinstance(maybe_set, list):
+        return set(maybe_set)
+
+    return {maybe_set}