fix: Only strip config key prefix when present

pyo3-object_store/src/aws/store.rs

@@ -230,7 +230,13 @@ impl<'py> FromPyObject<'_, 'py> for PyAmazonS3ConfigKey {
 
     fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult<Self> {
         let s = obj.extract::<PyBackedStr>()?.to_lowercase();
-        let key = s.parse().map_err(PyObjectStoreError::ObjectStoreError)?;
+        // Some keys (e.g. `aws_endpoint_url_s3`) are only accepted upstream in their prefixed
+        // form, but we strip the `aws_` prefix when converting keys back to Python. Retry with the
+        // prefix so that any key we emit can be parsed back in.
+        let key = s
+            .parse::<AmazonS3ConfigKey>()
+            .or_else(|err| format!("aws_{s}").parse().map_err(|_| err))
+            .map_err(PyObjectStoreError::ObjectStoreError)?;
         Ok(Self(key))
     }
 }
@@ -247,12 +253,10 @@ impl<'py> IntoPyObject<'py> for &PyAmazonS3ConfigKey {
     type Error = PyErr;
 
     fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
-        let s = self
-            .0
-            .as_ref()
-            .strip_prefix("aws_")
-            .expect("Expected config prefix to start with aws_");
-        Ok(PyString::new(py, s))
+        // Client and encryption config keys are not `aws_`-prefixed upstream, so only strip the
+        // prefix when it is actually present.
+        let s = self.0.as_ref();
+        Ok(PyString::new(py, s.strip_prefix("aws_").unwrap_or(s)))
     }
 }
 

pyo3-object_store/src/azure/store.rs

@@ -274,11 +274,9 @@ impl<'py> IntoPyObject<'py> for &PyAzureConfigKey {
         if let Some(stripped) = s.strip_prefix("azure_storage_") {
             return Ok(PyString::new(py, stripped));
         }
-        Ok(PyString::new(
-            py,
-            s.strip_prefix("azure_")
-                .expect("Expected config prefix to start with azure_"),
-        ))
+        // Client config keys are not `azure_`-prefixed upstream, so only strip the prefix when it
+        // is actually present.
+        Ok(PyString::new(py, s.strip_prefix("azure_").unwrap_or(s)))
     }
 }
 

pyo3-object_store/src/gcp/store.rs

@@ -248,12 +248,10 @@ impl<'py> IntoPyObject<'py> for &PyGoogleConfigKey {
     type Error = PyErr;
 
     fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
-        let s = self
-            .0
-            .as_ref()
-            .strip_prefix("google_")
-            .expect("Expected config prefix to start with google_");
-        Ok(PyString::new(py, s))
+        // Client config keys are not `google_`-prefixed upstream, so only strip the prefix when it
+        // is actually present.
+        let s = self.0.as_ref();
+        Ok(PyString::new(py, s.strip_prefix("google_").unwrap_or(s)))
     }
 }
 

tests/store/test_azure.py

@@ -34,6 +34,14 @@
     assert store != store3
 
 
+def test_client_config_key_does_not_panic():
+    # Client config keys are not `azure_`-prefixed upstream, so reading `.config`
+    # used to panic.
+    store = AzureStore("container", account_name="account_name", allow_http=True)
+    assert store.config["allow_http"] == "true"
+    assert AzureStore(config=store.config).config == store.config
+
+
 def test_from_url():
     # https://github.com/developmentseed/obstore/issues/477
     url = "https://overturemapswestus2.blob.core.windows.net/release"

tests/store/test_gcs.py

@@ -21,6 +21,14 @@
     assert store != store3
 
 
+def test_client_config_key_does_not_panic():
+    # Client config keys are not `google_`-prefixed upstream, so reading `.config`
+    # used to panic.
+    store = GCSStore("bucket", allow_http=True)
+    assert store.config["allow_http"] == "true"
+    assert GCSStore(config=store.config).config == store.config
+
+
 def test_application_credentials():
     # The application_credentials parameter should be correctly passed down
     # Finalizing the GCSBuilder should try to load and parse those credentials, which

tests/store/test_s3.py

@@ -85,6 +85,22 @@
     _objects = next(restored.list())
 
 
+def test_client_config_key_does_not_panic():
+    # Client config keys are not `aws_`-prefixed upstream, so reading `.config` used
+    # to panic.
+    store = S3Store("bucket", allow_http=True)
+    assert store.config["allow_http"] == "true"
+
+
+def test_prefixed_only_config_key_round_trip():
+    # `aws_endpoint_url_s3` has no unprefixed alias upstream, but we emit it unprefixed.
+    store = S3Store("bucket", aws_endpoint_url_s3="https://example.com")  # type: ignore
+    assert store.config["endpoint_url_s3"] == "https://example.com"
+
+    assert S3Store(config=store.config).config == store.config
+    assert pickle.loads(pickle.dumps(store)) == store
+
+
 def test_config_round_trip():
     store = S3Store.from_url(
         "s3://ookla-open-data/parquet/performance/type=fixed/year=2024/quarter=1",