arraylake 0.30.0__tar.gz → 0.32.0__tar.gz

{arraylake-0.30.0 → arraylake-0.32.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arraylake
-Version: 0.30.0
+Version: 0.32.0
 Summary: Python client for ArrayLake
 Project-URL: Documentation, https://docs.earthmover.io/
 Project-URL: Changelog, https://docs.earthmover.io/changelog

{arraylake-0.30.0 → arraylake-0.32.0}/arraylake/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.30.0'
-__version_tuple__ = version_tuple = (0, 30, 0)
+__version__ = version = '0.32.0'
+__version_tuple__ = version_tuple = (0, 32, 0)
 
 __commit_id__ = commit_id = None

{arraylake-0.30.0 → arraylake-0.32.0}/arraylake/client.py

@@ -38,7 +38,7 @@ from arraylake.credentials import (
     _use_hmac_credentials,
 )
 from arraylake.display.repolist import RepoList
-from arraylake.exceptions import ArraylakeClientError, BucketNotFoundError
+from arraylake.exceptions import BucketNotFoundError
 from arraylake.log_util import get_logger
 from arraylake.metastore import HttpMetastore, HttpMetastoreConfig
 from arraylake.repos.icechunk.storage import (
@@ -57,12 +57,12 @@ from arraylake.types import (
     ICECHUNK_ANY_CREDENTIAL,
     URI,
     ApiClientResponse,
-    ApiTokenInfo,
     Author,
     AzureCredentials,
     BucketNickname,
     BucketPrefix,
     BucketResponse,
+    ExplicitVirtualChunkAccessPolicyResponse,
     GSCredentials,
     NewBucket,
     OptimizationConfig,
@@ -78,7 +78,6 @@ from arraylake.types import (
     S3Credentials,
     StorageOptions,
     TempCredentials,
-    UserInfo,
     validate_name,
     validate_org_and_repo_name,
 )
@@ -491,6 +490,8 @@ class AsyncClient:
         Returns:
             A icechunk.Repository object.
         """
+        from arraylake import __version__ as arraylake_version
+
         if isinstance(config, RepositoryConfig):
             pass
         elif config is None:
@@ -500,70 +501,84 @@ class AsyncClient:
             raise ValueError(f"config must be an icechunk.RepositoryConfig object or None: {config}.")
 
         org, repo_name = _parse_org_and_repo(name)
-
         mstore = self._metastore_for_org(org)
 
-        async def _get_vccs_or_empty_on_missing_repo() -> dict[BucketPrefix, BucketNickname | None]:
-            """Get VCCs, returning empty dict if repo doesn't exist (let get_database raise that error)."""
-            try:
-                return await mstore.get_virtual_chunk_containers(repo_name)
-            except ArraylakeClientError as e:
-                if e.status_code == 404:
-                    return {}  # Repo doesn't exist - get_database will raise the proper error
-                raise
-
-        # optimization: fetch everything we might need from the metastore concurrently
-        if authorize_virtual_chunk_access is None or authorize_virtual_chunk_access:
-            # TODO these 3 concurrent API requests could instead be handled by a single new endpoint
-            repo_model_raw, user_raw, preset_vccs = await async_gather_tasks(
-                mstore.get_database(repo_name),
-                mstore.get_user(),
-                # This call is slower because it requires the server to also fetch the repo config from object storage
-                _get_vccs_or_empty_on_missing_repo(),
-            )
-        else:
-            # authorize_virtual_chunk_access is an empty dict, so no VCCs to authorize, and no need to fetch them from repo config
-            repo_model_raw, user_raw = await async_gather_tasks(
-                mstore.get_database(repo_name),
-                mstore.get_user(),
-            )
-            preset_vccs = {}
-
-        # TODO this "_raw" business is just for mypy? Surely it's not necessary...
-        if not isinstance(repo_model_raw, RepoModel):
-            raise ValueError("Expected Repo object from metastore")
-        if not isinstance(user_raw, (ApiTokenInfo, UserInfo)):
-            raise ValueError("Expected UserInfo or ApiTokenInfo from metastore")
-        if not isinstance(preset_vccs, dict):
-            raise ValueError("Expected dict from metastore")
-
-        repo_model = repo_model_raw
-        user = user_raw
-        author: Author = user.as_author()
+        # Single API call to get everything needed to open the repo
+        response = await mstore.open_repo(repo_name)
 
-        # Get the icechunk storage for the repo
-        user_id = UUID(str(user.id)) if not isinstance(user.id, UUID) else user.id
-        icechunk_storage = await self._get_icechunk_storage_from_repo_model(repo_model, user_id=user_id, storage_options=storage_options)
+        if response.status.mode == RepoOperationMode.OFFLINE:
+            raise ValueError(f"Repository is offline, check back after {response.status.estimated_end_time}")
+        elif response.status.mode == RepoOperationMode.MAINTENANCE:
+            warnings.warn(
+                f"Repository is in maintenance mode until {response.status.estimated_end_time}, do not attempt to write to the store"
+            )
 
-        # fetch credentials for any buckets containing virtual chunks referred to by this repo
-        prefixes_credential_mapping = await self._maybe_fetch_credentials_for_virtual_chunks(
-            org=org,
-            config=config,
-            pre_authorized_vccs=preset_vccs,
-            authorize_virtual_chunk_access=authorize_virtual_chunk_access,
+        # Build icechunk storage from the response
+        credential_refresh_func = self._maybe_get_credential_refresh_func_for_icechunk(
+            bucket=response.repo_bucket, org=org, repo_name=repo_name
+        )
+        credentials = response.repo_credentials if credential_refresh_func is None else None
+        icechunk_storage = _get_icechunk_storage_obj(
+            bucket_config=response.repo_bucket,
+            prefix=response.repo_prefix,
+            credential_type=_get_credential_type(credentials, credential_refresh_func),
+            credentials=credentials,
+            credential_refresh_func=credential_refresh_func,
+            scatter_initial_credentials=arraylake_config.get("icechunk.scatter_initial_credentials", True),
+            arraylake_version=arraylake_version,
+            user_id=UUID(response.principal_id),
+            storage_options=storage_options,
         )
 
-        if not isinstance(config, RepositoryConfig) and config is not None:
-            raise ValueError(f"config must be an icechunk.RepositoryConfig object or None: {config}.")
+        # Check if the user passed a custom config with VCCs that weren't authorized by the repo writer
+        server_authorized_vcc_prefixes = set(response.virtual_chunk_credentials.keys())
+        config_vccs = config.virtual_chunk_containers if config.virtual_chunk_containers is not None else {}
+        if config_vccs:
+            unauthorized_vccs = set(config_vccs.keys()) - server_authorized_vcc_prefixes
+            if unauthorized_vccs:
+                # TODO this is likely uneccessarily restrictive
+                raise ValueError(
+                    f"All virtual chunk access must be pre-authorized by a user with repo writer permissions, not just by the user reading the repo. "
+                    f"The following virtual chunk container prefixes in your custom config were not authorized by a repo writer: {unauthorized_vccs}"
+                )
+
+        # Resolve VCC credentials
+        vcc_credentials: dict[BucketPrefix, icechunk.AnyCredential | None]
+        if authorize_virtual_chunk_access is not None:
+            # User explicitly provided VCC overrides — fetch credentials for those buckets
+            # This incurs an extra roundtrip to the DB but it's not the happy-path so that's okay.
+            vcc_credentials = await self._containers_credentials_for_buckets(
+                org,
+                containers_to_buckets_map=dict(authorize_virtual_chunk_access),
+                forbid_configs_unsafe_for_virtual_chunks=True,
+            )
+        else:
+            # Use server-provided VCC credentials
+            vcc_credentials = {}
+            if response.virtual_chunk_credentials:
+                credentials_map: dict[str, icechunk.AnyS3Credential | icechunk.AnyGcsCredential | icechunk.AnyAzureCredential | None] = {}
+                for vcc_prefix, vcc_creds in response.virtual_chunk_credentials.items():
+                    # TODO: build refresh closures here using vcc_creds.org, vcc_creds.bucket_nickname,
+                    # and vcc_creds.platform so that credentials can be refreshed when they expire. (See https://github.com/earth-mover/arraylake/issues/5122)
+                    # TODO: support Azure virtual chunk credentials
+                    if isinstance(vcc_creds.credentials, AzureCredentials):
+                        raise NotImplementedError(f"Azure virtual chunk credentials are not yet supported. (prefix: {vcc_prefix}).")
+                    # TODO: fix return type of get_icechunk_container_credentials (returns GcsBearerCredential
+                    # which is not in icechunk's AnyGcsCredential union — icechunk typing gap)
+                    credentials_map[vcc_prefix] = get_icechunk_container_credentials(  # type: ignore[assignment]
+                        bucket_platform=vcc_creds.platform,
+                        credentials=vcc_creds.credentials,
+                        credential_refresh_func=None,
+                    )
+                vcc_credentials = icechunk.containers_credentials(credentials_map)
 
-        # TODO skip fetching the config again in the case that we have already fetched it previously? See https://github.com/earth-mover/arraylake/pull/3590#discussion_r2397969758
         ic_repo: IcechunkRepository = await IcechunkRepository.open_async(
             icechunk_storage,
-            config=config,  # The config passed here takes precedence over the stored config
-            authorize_virtual_chunk_access=prefixes_credential_mapping,
+            config=config,
+            authorize_virtual_chunk_access=vcc_credentials,
         )
 
-        ic_repo.set_default_commit_metadata({"author_name": author.name, "author_email": author.email})
+        ic_repo.set_default_commit_metadata({"author_name": response.author.name, "author_email": response.author.email})
 
         return ic_repo
 
@@ -1366,6 +1381,65 @@ class AsyncClient:
         bucket_id = await self._bucket_id_for_nickname(mstore, nickname)
         await mstore.delete_bucket_config(bucket_id)
 
+    async def set_virtual_chunk_access_policy(
+        self,
+        org: OrgName,
+        *,
+        bucket_nickname: BucketNickname,
+        subprefix: str = "",
+        public: bool = False,
+    ) -> ExplicitVirtualChunkAccessPolicyResponse:
+        """Set an explicit virtual chunk access policy on this org.
+
+        Args:
+            org: Name of the organization.
+            bucket_nickname: Nickname of the bucket config to associate with this VCAP.
+            subprefix: Subprefix within the bucket for this VCAP.
+            public: Whether this VCAP allows public access.
+        """
+        mstore = self._metastore_for_org(org)
+        bucket_id = await self._bucket_id_for_nickname(mstore, bucket_nickname)
+        vcap = await mstore.set_virtual_chunk_access_policy(
+            bucket_id=str(bucket_id),
+            subprefix=subprefix.strip("/"),
+            public=public,
+        )
+        return ExplicitVirtualChunkAccessPolicyResponse.model_validate(vcap.model_dump())
+
+    async def list_virtual_chunk_access_policies(self, org: OrgName) -> list[ExplicitVirtualChunkAccessPolicyResponse]:
+        """List explicit virtual chunk access policies on this org.
+
+        Args:
+            org: Name of the organization.
+        """
+        mstore = self._metastore_for_org(org)
+        vcaps = await mstore.list_virtual_chunk_access_policies()
+        return [ExplicitVirtualChunkAccessPolicyResponse.model_validate(v.model_dump()) for v in vcaps if v.use_case == "explicit"]
+
+    async def delete_virtual_chunk_access_policy(
+        self,
+        org: OrgName,
+        *,
+        bucket_nickname: BucketNickname,
+        subprefix: str,
+        public: bool,
+    ) -> None:
+        """Delete a virtual chunk access policy from this org (idempotent).
+
+        Args:
+            org: Name of the organization.
+            bucket_nickname: Nickname of the bucket config for the VCAP.
+            subprefix: Subprefix of the VCAP to delete.
+            public: Whether this VCAP allows public access.
+        """
+        mstore = self._metastore_for_org(org)
+        bucket_id = await self._bucket_id_for_nickname(mstore, bucket_nickname)
+        await mstore.delete_virtual_chunk_access_policy(
+            bucket_id=str(bucket_id),
+            subprefix=subprefix.strip("/"),
+            public=public,
+        )
+
     async def login(self, *, browser: bool = False) -> None:
         """Login to ArrayLake.
 
@@ -1878,6 +1952,52 @@ class Client:
         """
         return sync(self.aclient.delete_bucket_config, org=org, nickname=nickname, imsure=imsure, imreallysure=imreallysure)
 
+    def set_virtual_chunk_access_policy(
+        self,
+        org: OrgName,
+        *,
+        bucket_nickname: BucketNickname,
+        subprefix: str,
+        public: bool,
+    ) -> ExplicitVirtualChunkAccessPolicyResponse:
+        """Set an explicit virtual chunk access policy on this org.
+
+        Args:
+            org: Name of the organization.
+            bucket_nickname: Nickname of the bucket config to associate with this VCAP.
+            subprefix: Subprefix within the bucket for this VCAP.
+            public: Whether this VCAP allows public access.
+        """
+        return sync(self.aclient.set_virtual_chunk_access_policy, org, bucket_nickname=bucket_nickname, subprefix=subprefix, public=public)
+
+    def list_virtual_chunk_access_policies(self, org: OrgName) -> list[ExplicitVirtualChunkAccessPolicyResponse]:
+        """List explicit virtual chunk access policies on this org.
+
+        Args:
+            org: Name of the organization.
+        """
+        return sync(self.aclient.list_virtual_chunk_access_policies, org)
+
+    def delete_virtual_chunk_access_policy(
+        self,
+        org: OrgName,
+        *,
+        bucket_nickname: BucketNickname,
+        subprefix: str,
+        public: bool,
+    ) -> None:
+        """Delete a virtual chunk access policy from this org.
+
+        Args:
+            org: Name of the organization.
+            bucket_nickname: Nickname of the bucket config for the VCAP.
+            subprefix: Subprefix of the VCAP to delete.
+            public: Whether this VCAP allows public access.
+        """
+        return sync(
+            self.aclient.delete_virtual_chunk_access_policy, org, bucket_nickname=bucket_nickname, subprefix=subprefix, public=public
+        )
+
     def login(self, *, browser: bool = False) -> None:
         """Login to ArrayLake.
 

{arraylake-0.30.0 → arraylake-0.32.0}/arraylake/credentials.py

@@ -8,7 +8,7 @@ def _is_r2_bucket(bucket: BucketResponse | None) -> bool:
         return False
 
     endpoint_url = bucket.extra_config.get("endpoint_url")
-    if endpoint_url is not None and isinstance(endpoint_url, str) and "r2.cloudflarestorage.com" in endpoint_url:
+    if endpoint_url is not None and isinstance(endpoint_url, str) and ".r2." in endpoint_url:
         return True
     return False
 

{arraylake-0.30.0 → arraylake-0.32.0}/arraylake/metastore/http_metastore.py

@@ -17,6 +17,7 @@ from arraylake.types import (
     GSCredentials,
     NewBucket,
     NewRepoOperationStatus,
+    OpenRepoResponse,
     OptimizationConfig,
     OrgActions,
     PaginatedResponse,
@@ -32,6 +33,7 @@ from arraylake.types import (
     RepoOperationStatusResponse,
     S3Credentials,
     TokenAuthenticateBody,
+    VirtualChunkAccessPolicyResponse,
 )
 
 # type adapters
@@ -39,6 +41,7 @@ LIST_DATABASES_ADAPTER = TypeAdapter(list[Repo])
 LIST_DATABASES_PAGINATED_ADAPTER = TypeAdapter(PaginatedResponse[Repo])
 LIST_BUCKETS_ADAPTER = TypeAdapter(list[BucketResponse])
 LIST_REPOS_FOR_BUCKET_ADAPTER = TypeAdapter(list[Repo])
+LIST_VCAPS_PAGINATED_ADAPTER = TypeAdapter(PaginatedResponse[VirtualChunkAccessPolicyResponse])
 
 
 @dataclass
@@ -82,6 +85,19 @@ class HttpMetastore(ArraylakeHttpClient, Metastore):
         handle_response(response)
         return Repo.model_validate_json(response.content)
 
+    async def open_repo(self, name: str) -> OpenRepoResponse:
+        """Fetch everything needed to open an icechunk repo in a single request.
+
+        Args:
+            name: Name of the repository (without org prefix).
+
+        Returns:
+            OpenRepoResponse with storage config, credentials, VCC credentials, and author.
+        """
+        response = await self._request("GET", f"/repos/{self._config.org}/{name}/open")
+        handle_response(response)
+        return OpenRepoResponse.model_validate_json(response.content)
+
     async def list_databases(
         self,
         filter_metadata: RepoMetadataT | None = None,
@@ -311,3 +327,68 @@ class HttpMetastore(ArraylakeHttpClient, Metastore):
             content=json.dumps(payload),
         )
         handle_response(response)
+
+    async def set_virtual_chunk_access_policy(
+        self,
+        *,
+        bucket_id: str,
+        subprefix: str,
+        public: bool,
+    ) -> VirtualChunkAccessPolicyResponse:
+        """Set an explicit virtual chunk access policy on this org (idempotent).
+
+        Args:
+            bucket_id: Bucket ID for this VCAP.
+            subprefix: Subprefix within the bucket for this VCAP.
+            public: Whether this VCAP allows public access.
+        """
+        payload: dict[str, str | bool] = {
+            "use_case": "explicit",
+            "bucket_id": bucket_id,
+            "subprefix": subprefix,
+            "public": public,
+        }
+        response = await self._request(
+            "PUT",
+            f"/orgs/{self._config.org}/settings/virtual-chunk-access-policies",
+            content=json.dumps(payload),
+        )
+        handle_response(response)
+        return VirtualChunkAccessPolicyResponse.model_validate_json(response.content)
+
+    async def list_virtual_chunk_access_policies(self) -> list[VirtualChunkAccessPolicyResponse]:
+        """List all virtual chunk access policies on this org."""
+        response = await self._request(
+            "GET",
+            f"/orgs/{self._config.org}/settings/virtual-chunk-access-policies",
+            params={"size": 100},
+        )
+        handle_response(response)
+        return LIST_VCAPS_PAGINATED_ADAPTER.validate_json(response.content).items
+
+    async def delete_virtual_chunk_access_policy(
+        self,
+        *,
+        bucket_id: str,
+        subprefix: str,
+        public: bool,
+    ) -> None:
+        """Delete an explicit virtual chunk access policy from this org (idempotent).
+
+        Args:
+            bucket_id: Bucket ID of the VCAP to delete.
+            subprefix: Subprefix of the VCAP to delete.
+            public: Whether this VCAP allows public access.
+        """
+        payload: dict[str, str | bool] = {
+            "use_case": "explicit",
+            "bucket_id": bucket_id,
+            "subprefix": subprefix,
+            "public": public,
+        }
+        response = await self._request(
+            "DELETE",
+            f"/orgs/{self._config.org}/settings/virtual-chunk-access-policies",
+            content=json.dumps(payload),
+        )
+        handle_response(response)

{arraylake-0.30.0 → arraylake-0.32.0}/arraylake/repos/icechunk/storage.py

@@ -143,7 +143,7 @@ def _get_icechunk_storage_obj(
                 tigris_kwargs["network_stream_timeout_seconds"] = network_stream_timeout_seconds
             return icechunk.tigris_storage(**tigris_kwargs)
         # Use r2_storage to create the storage object for R2 buckets
-        elif endpoint_url and "r2.cloudflarestorage.com" in endpoint_url:
+        elif endpoint_url and ".r2." in endpoint_url:
             r2_kwargs: dict = dict(
                 bucket=bucket_config.name,
                 prefix=prefix,
@@ -297,7 +297,7 @@ def create_icechunk_store_config(
             )
             return icechunk.ObjectStoreConfig.S3Compatible(options)
 
-        elif endpoint_url and "r2.cloudflarestorage.com" in endpoint_url:
+        elif endpoint_url and ".r2." in endpoint_url:
             options = icechunk.S3Options(
                 region=region,
                 endpoint_url=endpoint_url,

{arraylake-0.30.0 → arraylake-0.32.0}/arraylake/repos/icechunk/virtual.py

@@ -8,7 +8,6 @@ from arraylake.types import (
     ICECHUNK_ANY_CREDENTIAL,
     URI,
     UUID,
-    AnonymousAuth,
     BucketNickname,
     BucketPrefix,
     BucketResponse,
@@ -62,7 +61,7 @@ def get_icechunk_container_credentials(
         # https://github.com/earth-mover/icechunk/pull/776
         if credentials:
             assert isinstance(credentials, GSCredentials)
-            return icechunk.GcsBearerCredential(bearer=credentials.access_token, expires_after=credentials.expiration)
+            return icechunk.GcsCredentials.Static(icechunk.GcsStaticCredentials.BearerToken(credentials.access_token))
         else:
             return icechunk.GcsCredentials.Anonymous()
     else:
@@ -164,13 +163,6 @@ def forbid_unsafe_virtual_bucket_configs(bucket: BucketResponse, bucket_nickname
             f"But bucket {bucket_nickname} uses HMAC credentials."
         )
 
-    # TODO: in future we could relax this to allow buckets with delegated credentials
-    elif not isinstance(auth_method, AnonymousAuth):
-        raise ValueError(
-            "Cannot use virtual chunk references that refer to a bucket which does not support anonymous public access, "
-            f"But bucket {bucket_nickname} uses {auth_method.method if auth_method else 'unknown'} credentials."
-        )
-
 
 def forbid_unsafe_virtual_chunk_containers(
     container_prefixes: Iterable[BucketPrefix],
@@ -178,9 +170,6 @@ def forbid_unsafe_virtual_chunk_containers(
     """
     Raises if any virtual chunk container prefix is somehow unsafe.
     """
-    # TODO Once we allow auto-authorization of existing VCCs (needed for Flux), this fn will need to be changed.
-    # TODO It will need to actually fetch the existing VCCs, check they are all safe, and only then open the repo with authorization.
-
     for url_prefix in container_prefixes:
         if any(url_prefix.startswith(forbidden_prefix) for forbidden_prefix in FORBIDDEN_VIRTUAL_CHUNK_PREFIXES):
             raise ValueError(

{arraylake-0.30.0 → arraylake-0.32.0}/arraylake/token.py

@@ -117,8 +117,10 @@ class TokenHandler:
                 response = client.get(f"{self.api_endpoint}/auth/config?target=client")
                 response.raise_for_status()
                 return AuthProviderConfig.model_validate_json(response.text)
-        except Exception:
-            raise AuthException("⚠️ There was an error getting the auth configuration!")
+        except httpx.HTTPStatusError as e:
+            raise AuthException(f"⚠️ Error getting auth configuration: {e.response.status_code} from {e.request.url!r}")
+        except httpx.ConnectError:
+            raise AuthException(f"⚠️ Could not connect to {self.api_endpoint}. Check your network connection and service URI.")
 
     async def get_authorize_info(self) -> Auth0UrlCode:
         """
@@ -263,7 +265,7 @@ class TokenHandler:
         async with self._create_async_client() as client:
             response = await client.get(f"{self.api_endpoint}/user", headers=headers)
         if response.status_code != httpx.codes.OK:
-            raise AuthException("There was an error getting your user!\nPlease contact support at support@earthmover.io.")
+            raise AuthException(f"Error fetching user: {response.status_code} {response.text}")
         return UserInfo.model_validate_json(response.content)
 
     async def _update_user_profile(self) -> UserInfo:
@@ -273,7 +275,7 @@ class TokenHandler:
         async with self._create_async_client() as client:
             response = await client.post(f"{self.api_endpoint}/auth/update-user", headers=headers)
         if response.status_code != httpx.codes.OK:
-            raise AuthException("There was an error getting your user!\nPlease contact support at support@earthmover.io.")
+            raise AuthException(f"Error updating user profile: {response.status_code} {response.text}")
         return UserInfo.model_validate_json(response.content)
 
     async def login(self, *, browser: bool = False) -> None:
@@ -321,9 +323,9 @@ class TokenHandler:
             self.console.print(f"[green][bold]Successfully logged in![/bold][/green] [dim]Token stored at {self.token_path}[/dim]")
             print_user_details(user)
 
-        except Exception:
-            self.error_console.print("\n[bold red]❌ There was an error fetching your user![/bold red]")
-            self.error_console.print("[dim]Please contact support: support@earthmover.io[/dim]")
+        except AuthException as e:
+            self.console.print(f"[green][bold]Successfully logged in![/bold][/green] [dim]Token stored at {self.token_path}[/dim]")
+            self.error_console.print(f"\n[bold red]⚠️ Could not update user profile: {e}[/bold red]")
 
     async def _logout(self) -> None:
         try: