feat: get_obstore_store function for creating Obstore store with Open PC credentials

Author: kylebarron

planetary_computer/__init__.py

@@ -13,12 +13,14 @@
 )
 from planetary_computer.settings import set_subscription_key
 from planetary_computer._adlfs import get_adlfs_filesystem, get_container_client
+from planetary_computer._obstore import get_obstore_store
 
 from planetary_computer.version import __version__
 
 __all__ = [
     "get_adlfs_filesystem",
     "get_container_client",
+    "get_obstore_store",
     "set_subscription_key",
     "sign_asset",
     "sign_assets",

planetary_computer/_obstore.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from planetary_computer.sas import get_token
+
+if TYPE_CHECKING:
+    import sys
+
+    from obstore.store import (
+        AzureStore,
+        AzureConfig,
+        ClientConfig,
+        RetryConfig,
+        AzureSASToken,
+    )
+
+    if sys.version_info >= (3, 11):
+        from typing import Unpack
+    else:
+        from typing_extensions import Unpack
+
+
+def get_obstore_store(  # type: ignore[misc] # overlap with kwargs
+    account_name: str,
+    container_name: str,
+    *,
+    prefix: str | None = None,
+    config: AzureConfig | None = None,
+    client_options: ClientConfig | None = None,
+    retry_config: RetryConfig | None = None,
+    **kwargs: Unpack[AzureConfig],  # type: ignore # noqa: PGH003 (container_name key overlaps with positional arg)
+) -> AzureStore:
+    try:
+        import obstore
+    except ImportError as e:
+        raise ImportError(
+            "'planetary_computer.get_obstore_store' requires "
+            "the optional dependency 'obstore'."
+        ) from e
+
+    def credential_provider() -> AzureSASToken:
+        token = get_token(account_name, container_name)
+        return {
+            "sas_token": token.token,
+            "expires_at": token.expiry,
+        }
+
+    return obstore.store.AzureStore(
+        account_name=account_name,
+        container_name=container_name,
+        prefix=prefix,
+        config=config,
+        client_options=client_options,
+        retry_config=retry_config,
+        credential_provider=credential_provider,
+        **kwargs,  # type: ignore # (container_name key overlaps with positional arg)
+    )

pyproject.toml

@@ -34,6 +34,7 @@ dev = [
     "pytest",
     "responses",
 ]
+obstore = ["obstore"]
 
 [project.scripts]
 planetarycomputer = "planetary_computer.scripts.cli:app"
@@ -42,4 +43,4 @@ planetarycomputer = "planetary_computer.scripts.cli:app"
 include-package-data = false
 
 [tool.setuptools.dynamic]
-version = {attr = "planetary_computer.version.__version__"}
+version = {attr = "planetary_computer.version.__version__"}