Check for pickle safety before pickling

pyo3-object_store/src/aws/store.rs

@@ -26,6 +26,9 @@ struct S3Config {
     config: PyAmazonS3Config,
     client_options: Option<PyClientOptions>,
     retry_config: Option<PyRetryConfig>,
+    /// 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<PyObject> {
+        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<PyClientOptions>,
         retry_config: Option<PyRetryConfig>,
         kwargs: Option<PyAmazonS3Config>,
+        pickle_safe: bool,
     ) -> PyObjectStoreResult<Self> {
         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,
         )
     }
 

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)