Cache object stores per scheme + bucket

In this PR, we cache object stores per scheme (e.g. s3) + bucket (container) combination
in each Postgres session. This will reduce authentication costs by only doing it at the first
time.

object_store does not perform sts assume_role to get temp token, so pg_parquet
make use of aws sdk to perform it. And then configure object_store with the temp
token that it fetched. This is why, pg_parquet checks expiration of the tokens for s3 store.
It will fetch the temp token if it expires. (obviously if you configured temp token auth via config)

For azure blob store, we do not need to recreate tokens because object_store handles it.

Author: aykut-bozkurt

.devcontainer/entrypoint.sh

@@ -4,5 +4,6 @@ trap "echo 'Caught termination signal. Exiting...'; exit 0" SIGINT SIGTERM
 
 # create azurite container
 az storage container create -n $AZURE_TEST_CONTAINER_NAME --connection-string $AZURE_STORAGE_CONNECTION_STRING
+az storage container create -n ${AZURE_TEST_CONTAINER_NAME}2 --connection-string $AZURE_STORAGE_CONNECTION_STRING
 
 sleep infinity

.devcontainer/minio-entrypoint.sh

@@ -14,7 +14,8 @@ done
 # set access key and secret key
 mc alias set local $AWS_ENDPOINT_URL $MINIO_ROOT_USER $MINIO_ROOT_PASSWORD
 
-# create bucket
+# create buckets
 mc mb local/$AWS_S3_TEST_BUCKET
+mc mb local/${AWS_S3_TEST_BUCKET}2
 
 wait $minio_pid

.github/workflows/ci.yml

@@ -138,6 +138,7 @@ jobs:
 
           # create container
           az storage container create -n $AZURE_TEST_CONTAINER_NAME --connection-string $AZURE_STORAGE_CONNECTION_STRING
+          az storage container create -n ${AZURE_TEST_CONTAINER_NAME}2 --connection-string $AZURE_STORAGE_CONNECTION_STRING
 
       - name: Run tests
         run: |

src/arrow_parquet/uri_utils.rs

@@ -18,8 +18,8 @@ use pgrx::{
 use url::Url;
 
 use crate::{
-    arrow_parquet::parquet_writer::DEFAULT_ROW_GROUP_SIZE, object_store::create_object_store,
-    PG_BACKEND_TOKIO_RUNTIME,
+    arrow_parquet::parquet_writer::DEFAULT_ROW_GROUP_SIZE,
+    object_store::object_store_cache::get_or_create_object_store, PG_BACKEND_TOKIO_RUNTIME,
 };
 
 const PARQUET_OBJECT_STORE_READ_ROLE: &str = "parquet_object_store_read";
@@ -58,7 +58,7 @@ pub(crate) fn parquet_schema_from_uri(uri: &Url) -> SchemaDescriptor {
 
 pub(crate) fn parquet_metadata_from_uri(uri: &Url) -> Arc<ParquetMetaData> {
     let copy_from = true;
-    let (parquet_object_store, location) = create_object_store(uri, copy_from);
+    let (parquet_object_store, location) = get_or_create_object_store(uri, copy_from);
 
     PG_BACKEND_TOKIO_RUNTIME.block_on(async {
         let object_store_meta = parquet_object_store
@@ -81,7 +81,7 @@ pub(crate) fn parquet_metadata_from_uri(uri: &Url) -> Arc<ParquetMetaData> {
 
 pub(crate) fn parquet_reader_from_uri(uri: &Url) -> ParquetRecordBatchStream<ParquetObjectReader> {
     let copy_from = true;
-    let (parquet_object_store, location) = create_object_store(uri, copy_from);
+    let (parquet_object_store, location) = get_or_create_object_store(uri, copy_from);
 
     PG_BACKEND_TOKIO_RUNTIME.block_on(async {
         let object_store_meta = parquet_object_store
@@ -113,7 +113,7 @@ pub(crate) fn parquet_writer_from_uri(
     writer_props: WriterProperties,
 ) -> AsyncArrowWriter<ParquetObjectWriter> {
     let copy_from = false;
-    let (parquet_object_store, location) = create_object_store(uri, copy_from);
+    let (parquet_object_store, location) = get_or_create_object_store(uri, copy_from);
 
     let parquet_object_writer = ParquetObjectWriter::new(parquet_object_store, location);
 

src/object_store.rs

@@ -1,8 +1,3 @@
-use std::sync::Arc;
-
-use object_store::{path::Path, ObjectStore, ObjectStoreScheme};
-use url::Url;
-
 use crate::{
     arrow_parquet::uri_utils::uri_as_string,
     object_store::{
@@ -15,42 +10,4 @@ use crate::{
 pub(crate) mod aws;
 pub(crate) mod azure;
 pub(crate) mod local_file;
-
-pub(crate) fn create_object_store(uri: &Url, copy_from: bool) -> (Arc<dyn ObjectStore>, Path) {
-    let (scheme, path) = ObjectStoreScheme::parse(uri).unwrap_or_else(|_| {
-        panic!(
-            "unrecognized uri {}. pg_parquet supports local paths, s3:// or azure:// schemes.",
-            uri
-        )
-    });
-
-    // object_store crate can recognize a bunch of different schemes and paths, but we only support
-    // local, azure, and s3 schemes with a subset of all supported paths.
-    match scheme {
-        ObjectStoreScheme::AmazonS3 => {
-            let storage_container = Arc::new(create_s3_object_store(uri));
-
-            (storage_container, path)
-        }
-        ObjectStoreScheme::MicrosoftAzure => {
-            let storage_container = Arc::new(create_azure_object_store(uri));
-
-            (storage_container, path)
-        }
-        ObjectStoreScheme::Local => {
-            let storage_container = Arc::new(create_local_file_object_store(uri, copy_from));
-
-            let path =
-                Path::from_filesystem_path(uri_as_string(uri)).unwrap_or_else(|e| panic!("{}", e));
-
-            (storage_container, path)
-        }
-        _ => {
-            panic!(
-                "unsupported scheme {} in uri {}. pg_parquet supports local paths, s3:// or azure:// schemes.",
-                uri.scheme(),
-                uri
-            );
-        }
-    }
-}
+pub(crate) mod object_store_cache;

src/object_store/aws.rs

@@ -1,9 +1,11 @@
+use std::{sync::Arc, time::SystemTime};
+
 use aws_config::BehaviorVersion;
 use aws_credential_types::provider::ProvideCredentials;
-use object_store::aws::{AmazonS3, AmazonS3Builder};
+use object_store::aws::AmazonS3Builder;
 use url::Url;
 
-use super::PG_BACKEND_TOKIO_RUNTIME;
+use super::{object_store_cache::ObjectStoreWithExpiration, PG_BACKEND_TOKIO_RUNTIME};
 
 // create_s3_object_store creates an AmazonS3 object store with the given bucket name.
 // It is configured by environment variables and aws config files as fallback method.
@@ -19,7 +21,7 @@ use super::PG_BACKEND_TOKIO_RUNTIME;
 // - AWS_CONFIG_FILE (env var only)
 // - AWS_PROFILE (env var only)
 // - AWS_ALLOW_HTTP (env var only, object_store specific)
-pub(crate) fn create_s3_object_store(uri: &Url) -> AmazonS3 {
+pub(crate) fn create_s3_object_store(uri: &Url) -> ObjectStoreWithExpiration {
     let bucket_name = parse_s3_bucket(uri).unwrap_or_else(|| {
         panic!("unsupported s3 uri: {}", uri);
     });
@@ -58,10 +60,17 @@ pub(crate) fn create_s3_object_store(uri: &Url) -> AmazonS3 {
         aws_s3_builder = aws_s3_builder.with_region(region);
     }
 
-    aws_s3_builder.build().unwrap_or_else(|e| panic!("{}", e))
+    let object_store = aws_s3_builder.build().unwrap_or_else(|e| panic!("{}", e));
+
+    let expire_at = aws_s3_config.expire_at;
+
+    ObjectStoreWithExpiration {
+        object_store: Arc::new(object_store),
+        expire_at,
+    }
 }
 
-fn parse_s3_bucket(uri: &Url) -> Option<String> {
+pub(crate) fn parse_s3_bucket(uri: &Url) -> Option<String> {
     let host = uri.host_str()?;
 
     // s3(a)://{bucket}/key
@@ -98,6 +107,7 @@ struct AwsS3Config {
     access_key_id: Option<String>,
     secret_access_key: Option<String>,
     session_token: Option<String>,
+    expire_at: Option<SystemTime>,
     endpoint_url: Option<String>,
     allow_http: bool,
 }
@@ -121,6 +131,7 @@ impl AwsS3Config {
         let mut access_key_id = None;
         let mut secret_access_key = None;
         let mut session_token = None;
+        let mut expire_at = None;
 
         if let Some(credential_provider) = sdk_config.credentials_provider() {
             if let Ok(credentials) = PG_BACKEND_TOKIO_RUNTIME
@@ -129,6 +140,7 @@ impl AwsS3Config {
                 access_key_id = Some(credentials.access_key_id().to_string());
                 secret_access_key = Some(credentials.secret_access_key().to_string());
                 session_token = credentials.session_token().map(|t| t.to_string());
+                expire_at = credentials.expiry();
             }
         }
 
@@ -141,6 +153,7 @@ impl AwsS3Config {
             access_key_id,
             secret_access_key,
             session_token,
+            expire_at,
             endpoint_url,
             allow_http,
         }

src/object_store/azure.rs

@@ -1,9 +1,13 @@
+use std::sync::Arc;
+
 use azure_storage::{ConnectionString, EndpointProtocol};
 use home::home_dir;
 use ini::Ini;
-use object_store::azure::{AzureConfigKey, MicrosoftAzure, MicrosoftAzureBuilder};
+use object_store::azure::{AzureConfigKey, MicrosoftAzureBuilder};
 use url::Url;
 
+use super::object_store_cache::ObjectStoreWithExpiration;
+
 // create_azure_object_store creates a MicrosoftAzure object store with the given container name.
 // It is configured by environment variables and azure config files as fallback method.
 // We need to read the config files to make the fallback method work since object_store
@@ -16,7 +20,7 @@ use url::Url;
 // - AZURE_CONFIG_FILE (env var only, object_store specific)
 // - AZURE_STORAGE_ENDPOINT (env var only, object_store specific)
 // - AZURE_ALLOW_HTTP (env var only, object_store specific)
-pub(crate) fn create_azure_object_store(uri: &Url) -> MicrosoftAzure {
+pub(crate) fn create_azure_object_store(uri: &Url) -> ObjectStoreWithExpiration {
     let container_name = parse_azure_blob_container(uri).unwrap_or_else(|| {
         panic!("unsupported azure blob storage uri: {}", uri);
     });
@@ -63,10 +67,18 @@ pub(crate) fn create_azure_object_store(uri: &Url) -> MicrosoftAzure {
         azure_builder = azure_builder.with_client_secret(client_secret);
     }
 
-    azure_builder.build().unwrap_or_else(|e| panic!("{}", e))
+    let object_store = azure_builder.build().unwrap_or_else(|e| panic!("{}", e));
+
+    // object store handles refreshing bearer token, so we do not need to handle expiry here
+    let expire_at = None;
+
+    ObjectStoreWithExpiration {
+        object_store: Arc::new(object_store),
+        expire_at,
+    }
 }
 
-fn parse_azure_blob_container(uri: &Url) -> Option<String> {
+pub(crate) fn parse_azure_blob_container(uri: &Url) -> Option<String> {
     let host = uri.host_str()?;
 
     // az(ure)://{container}/key

src/object_store/local_file.rs

@@ -1,10 +1,15 @@
+use std::sync::Arc;
+
 use object_store::local::LocalFileSystem;
 use url::Url;
 
-use super::uri_as_string;
+use super::{object_store_cache::ObjectStoreWithExpiration, uri_as_string};
 
 // create_local_file_object_store creates a LocalFileSystem object store with the given path.
-pub(crate) fn create_local_file_object_store(uri: &Url, copy_from: bool) -> LocalFileSystem {
+pub(crate) fn create_local_file_object_store(
+    uri: &Url,
+    copy_from: bool,
+) -> ObjectStoreWithExpiration {
     let path = uri_as_string(uri);
 
     if !copy_from {
@@ -17,5 +22,11 @@ pub(crate) fn create_local_file_object_store(uri: &Url, copy_from: bool) -> Loca
             .unwrap_or_else(|e| panic!("{}", e));
     }
 
-    LocalFileSystem::new()
+    let object_store = LocalFileSystem::new();
+    let expire_at = None;
+
+    ObjectStoreWithExpiration {
+        object_store: Arc::new(object_store),
+        expire_at,
+    }
 }