Skip to content

Commit

Permalink
Merge pull request #703 from pangeo-forge/gcorradini/credential_fix
Browse files Browse the repository at this point in the history
Fix WriteCombinedReference FSSpec Credentials Yet Again
  • Loading branch information
ranchodeluxe authored Feb 27, 2024
2 parents e39a205 + 675a357 commit 0507455
Showing 1 changed file with 10 additions and 7 deletions.
17 changes: 10 additions & 7 deletions pangeo_forge_recipes/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -581,7 +581,11 @@ class WriteCombinedReference(beam.PTransform, ZarrWriterMixin):
:param concat_dims: Dimensions along which to concatenate inputs.
:param identical_dims: Dimensions shared among all inputs.
:param mzz_kwargs: Additional kwargs to pass to ``kerchunk.combine.MultiZarrToZarr``.
:param target_root: Root path the Zarr store will be created inside; ``store_name``
:param remote_options: options to pass to ``kerchunk.combine.MultiZarrToZarr``
to read reference inputs (can include credentials).
:param remote_protocol: If files are accessed over the network, provide the remote protocol
over which they are accessed. e.g.: "s3", "https", etc.
:param target_root: Output root path the store will be created inside; ``store_name``
will be appended to this prefix to create a full path.
:param output_file_name: Name to give the output references file
(``.json`` or ``.parquet`` suffix).
Expand All @@ -591,23 +595,22 @@ class WriteCombinedReference(beam.PTransform, ZarrWriterMixin):
concat_dims: List[str]
identical_dims: List[str]
mzz_kwargs: dict = field(default_factory=dict)
remote_options: Optional[Dict] = field(default_factory=dict)
remote_protocol: Optional[str] = None
target_root: Union[str, FSSpecTarget, RequiredAtRuntimeDefault] = field(
default_factory=RequiredAtRuntimeDefault
)
output_file_name: str = "reference.json"

def expand(self, references: beam.PCollection) -> beam.PCollection[zarr.storage.FSStore]:
# unpack fsspec options that will be used below for transforms without dep injection
storage_options = self.target_root.fsspec_kwargs # type: ignore[union-attr]
remote_protocol = self.target_root.get_fsspec_remote_protocol() # type: ignore[union-attr]
return (
references
| CombineReferences(
concat_dims=self.concat_dims,
identical_dims=self.identical_dims,
target_options=storage_options,
remote_options=storage_options,
remote_protocol=remote_protocol,
target_options=self.remote_options,
remote_options=self.remote_options,
remote_protocol=self.remote_protocol,
mzz_kwargs=self.mzz_kwargs,
)
| WriteReference(
Expand Down

0 comments on commit 0507455

Please sign in to comment.