Refactor S3 handling in GeoZarr conversion: update create_s3_store to return S3 path and introduce get_s3_storage_options for configuration. Enhance tests for path handling and storage options retrieval.

Author: emmanuelmathot

eopf_geozarr/conversion/geozarr.py

@@ -282,14 +282,15 @@ def recursive_copy(
         # Handle S3 vs local paths for zarr operations
         group_path = f"{output_path}/{group_prefix}"
         if s3_utils.is_s3_path(output_path):
-            # For S3, use the S3 store
-            store = s3_utils.create_s3_store(group_path)
+            # For S3, use storage_options
+            storage_options = s3_utils.get_s3_storage_options(group_path)
             ds.to_zarr(
-                store,
+                group_path,
                 mode="w" if no_children else "a",  # Write if no children, append otherwise
                 consolidated=is_dataset,  # Consolidate metadata if it's a dataset
                 zarr_format=3,
                 encoding=encoding,
+                storage_options=storage_options,
             )
         else:
             ds.to_zarr(
@@ -1165,15 +1166,16 @@ def write_dataset_band_by_band_with_validation(
 
                 # Handle S3 vs local paths for zarr operations
                 if s3_utils.is_s3_path(output_path):
-                    # For S3, use the S3 store
-                    store = s3_utils.create_s3_store(output_path)
+                    # For S3, use storage_options
+                    storage_options = s3_utils.get_s3_storage_options(output_path)
                     single_var_ds.to_zarr(
-                        store,
+                        output_path,
                         mode=mode,
                         consolidated=False,
                         zarr_format=3,
                         encoding=var_encoding,
                         align_chunks=True,
+                        storage_options=storage_options,
                     )
                 else:
                     single_var_ds.to_zarr(
@@ -1225,14 +1227,15 @@ def write_dataset_band_by_band_with_validation(
             try:
                 # Handle S3 vs local paths for zarr operations
                 if s3_utils.is_s3_path(output_path):
-                    # For S3, use the S3 store
-                    store = s3_utils.create_s3_store(output_path)
+                    # For S3, use storage_options
+                    storage_options = s3_utils.get_s3_storage_options(output_path)
                     single_var_ds.to_zarr(
-                        store,
+                        output_path,
                         mode="a",  # Always append for grid_mapping variables
                         consolidated=False,
                         zarr_format=3,
                         encoding=var_encoding,
+                        storage_options=storage_options,
                     )
                 else:
                     single_var_ds.to_zarr(

eopf_geozarr/conversion/s3_utils.py

@@ -51,9 +51,9 @@ def parse_s3_path(s3_path: str) -> tuple[str, str]:
     return bucket, key
 
 
-def create_s3_store(s3_path: str, **s3_kwargs) -> FsspecStore:
+def get_s3_storage_options(s3_path: str, **s3_kwargs) -> Dict[str, Any]:
     """
-    Create an S3 store for Zarr operations.
+    Get storage options for S3 access with xarray.
 
     Parameters
     ----------
@@ -64,8 +64,8 @@ def create_s3_store(s3_path: str, **s3_kwargs) -> FsspecStore:
 
     Returns
     -------
-    FsspecStore
-        Zarr store backed by S3
+    Dict[str, Any]
+        Storage options dictionary for xarray
     """
     # Set up default S3 configuration
     default_s3_kwargs = {
@@ -84,17 +84,30 @@ def create_s3_store(s3_path: str, **s3_kwargs) -> FsspecStore:
     # Merge with user-provided kwargs
     s3_config = {**default_s3_kwargs, **s3_kwargs}
 
-    # Create S3 filesystem
-    fs = s3fs.S3FileSystem(**s3_config)
-    
-    # Remove the s3:// scheme from the path for FsspecStore
-    bucket, key = parse_s3_path(s3_path)
-    path_without_scheme = f"{bucket}/{key}" if key else bucket
-    
-    # Create FsspecStore
-    store = FsspecStore(fs=fs, path=path_without_scheme)
+    return s3_config
+
+
+def create_s3_store(s3_path: str, **s3_kwargs) -> str:
+    """
+    Create an S3 path with storage options for Zarr operations.
     
-    return store
+    This function now returns the S3 path directly, to be used with
+    xarray's storage_options parameter instead of creating a store.
+
+    Parameters
+    ----------
+    s3_path : str
+        S3 path in format s3://bucket/key
+    **s3_kwargs
+        Additional keyword arguments for s3fs.S3FileSystem
+
+    Returns
+    -------
+    str
+        S3 path to be used with storage_options
+    """
+    # Just return the S3 path - storage options will be handled separately
+    return s3_path
 
 
 def write_s3_json_metadata(s3_path: str, metadata: Dict[str, Any], **s3_kwargs) -> None:
@@ -117,6 +130,7 @@ def write_s3_json_metadata(s3_path: str, metadata: Dict[str, Any], **s3_kwargs)
     default_s3_kwargs = {
         "anon": False,
         "use_ssl": True,
+        "asynchronous": False,  # Force synchronous mode
         "client_kwargs": {
             "region_name": os.environ.get("AWS_DEFAULT_REGION", "us-east-1")
         }
@@ -155,6 +169,7 @@ def s3_path_exists(s3_path: str, **s3_kwargs) -> bool:
     default_s3_kwargs = {
         "anon": False,
         "use_ssl": True,
+        "asynchronous": False,  # Force synchronous mode
         "client_kwargs": {
             "region_name": os.environ.get("AWS_DEFAULT_REGION", "us-east-1")
         }
@@ -234,6 +249,7 @@ def validate_s3_access(s3_path: str, **s3_kwargs) -> tuple[bool, Optional[str]]:
         default_s3_kwargs = {
             "anon": False,
             "use_ssl": True,
+            "asynchronous": False,  # Force synchronous mode
             "client_kwargs": {
                 "region_name": os.environ.get("AWS_DEFAULT_REGION", "us-east-1")
             }

eopf_geozarr/tests/test_s3_utils.py

@@ -9,6 +9,7 @@
     get_s3_credentials_info,
     validate_s3_access,
     create_s3_store,
+    get_s3_storage_options,
 )
 
 
@@ -77,22 +78,27 @@ def test_validate_s3_access_failure(mock_s3fs):
     assert "Access denied" in error
 
 
-@patch('eopf_geozarr.conversion.s3_utils.s3fs.S3FileSystem')
-@patch('eopf_geozarr.conversion.s3_utils.FsspecStore')
-def test_create_s3_store_path_handling(mock_fsspec_store, mock_s3fs):
-    """Test that create_s3_store correctly handles S3 path schemes."""
-    mock_fs = Mock()
-    mock_s3fs.return_value = mock_fs
-    mock_store = Mock()
-    mock_fsspec_store.return_value = mock_store
-    
+def test_create_s3_store_path_handling():
+    """Test that create_s3_store returns the S3 path correctly."""
     # Test with S3 path
-    store = create_s3_store("s3://test-bucket/path/to/data")
-    
-    # Verify that FsspecStore was called with path without scheme
-    mock_fsspec_store.assert_called_once_with(fs=mock_fs, path="test-bucket/path/to/data")
+    result = create_s3_store("s3://test-bucket/path/to/data")
+    assert result == "s3://test-bucket/path/to/data"
 
     # Test with bucket only
-    mock_fsspec_store.reset_mock()
-    store = create_s3_store("s3://test-bucket")
-    mock_fsspec_store.assert_called_with(fs=mock_fs, path="test-bucket")
+    result = create_s3_store("s3://test-bucket")
+    assert result == "s3://test-bucket"
+
+
+def test_get_s3_storage_options():
+    """Test that get_s3_storage_options returns correct configuration."""
+    with patch.dict('os.environ', {
+        'AWS_DEFAULT_REGION': 'us-west-2',
+        'AWS_S3_ENDPOINT': 'https://s3.example.com'
+    }):
+        options = get_s3_storage_options("s3://test-bucket/path")
+        
+        assert options['anon'] is False
+        assert options['use_ssl'] is True
+        assert options['client_kwargs']['region_name'] == 'us-west-2'
+        assert options['endpoint_url'] == 'https://s3.example.com'
+        assert options['client_kwargs']['endpoint_url'] == 'https://s3.example.com'