diff --git a/pyo3-object_store/src/aws/store.rs b/pyo3-object_store/src/aws/store.rs index 9e8ca80a..00e1b925 100644 --- a/pyo3-object_store/src/aws/store.rs +++ b/pyo3-object_store/src/aws/store.rs @@ -26,6 +26,9 @@ struct S3Config { config: PyAmazonS3Config, client_options: Option, retry_config: Option, + /// Whether or not this config can accurately be pickled. + /// This is false if a custom CredentialProvider is used. + pickle_safe: bool, } impl S3Config { @@ -38,6 +41,10 @@ impl S3Config { } fn __getnewargs_ex__(&self, py: Python) -> PyResult { + if !self.pickle_safe { + return Err(GenericError::new_err("Instance not safe to pickle.")); + } + let args = PyTuple::empty(py).into_py_any(py)?; let kwargs = PyDict::new(py); @@ -80,6 +87,7 @@ impl PyS3Store { client_options: Option, retry_config: Option, kwargs: Option, + pickle_safe: bool, ) -> PyObjectStoreResult { let mut config = config.unwrap_or_default(); if let Some(bucket) = bucket { @@ -102,6 +110,7 @@ impl PyS3Store { config: combined_config, client_options, retry_config, + pickle_safe, }, }) } @@ -133,6 +142,7 @@ impl PyS3Store { client_options, retry_config, kwargs, + true, ) } @@ -161,6 +171,7 @@ impl PyS3Store { client_options, retry_config, kwargs, + false, ) } @@ -227,6 +238,8 @@ impl PyS3Store { client_options, retry_config, kwargs, + // We use frozen credentials; boto3 isn't a direct credentials provider + true, ) } @@ -258,6 +271,7 @@ impl PyS3Store { client_options, retry_config, kwargs, + true, ) } diff --git a/tests/store/test_s3.py b/tests/store/test_s3.py index 5e070df7..ef52a772 100644 --- a/tests/store/test_s3.py +++ b/tests/store/test_s3.py @@ -84,3 +84,9 @@ def test_config_round_trip(): assert store.prefix == new_store.prefix assert store.client_options == new_store.client_options assert store.retry_config == new_store.retry_config + + +def test_native_credentials_fails_pickle(): + store = S3Store._from_native("bucket") + with pytest.raises(BaseError, match="not safe to pickle"): + pickle.dumps(store)