Adds Support for COPY TO/FROM Google Cloud Storage (#61)

aykut-bozkurt · web-flow · commit 0de2de171cff · 2025-04-09T13:35:15.000Z
Supports following Google Cloud Storage uri forms: - gs:// \<bucket\> / \<path\> **Configuration** The simplest way to configure object storage is by creating a json config file like `~/.config/gcloud/application_default_credentials.json` (can be generated by `gcloud auth application-default login`): ```bash $ ~/.config/gcloud/application_default_credentials.json { "gcs_base_url": "gs://testbucket/test.parquet", "disable_oauth": false, "client_email": "...", "private_key_id": "...", "private_key": "..." } ``` Alternatively, you can use the following environment variables when starting postgres to configure the Google Cloud Storage client: - `GOOGLE_SERVICE_ACCOUNT_KEY`: json serialized service account key - `GOOGLE_SERVICE_ACCOUNT_PATH`: an alternative location for the config file Closes #62
diff --git a/.devcontainer/.env b/.devcontainer/.env
@@ -22,5 +22,10 @@ AZURE_TEST_READ_WRITE_SAS="se=2100-05-05&sp=rcw&sv=2022-11-02&sr=c&sig=TPz2jEz0t
 ALLOW_HTTP=true
 HTTP_ENDPOINT=http://localhost:8080
 
+# GCS tests
+GOOGLE_TEST_BUCKET=testbucket
+GOOGLE_SERVICE_ACCOUNT_KEY='{"gcs_base_url": "http://localhost:4443","disable_oauth": true,"client_email": "","private_key_id": "","private_key": ""}'
+GOOGLE_SERVICE_ENDPOINT=http://localhost:4443
+
 # Others
 RUST_TEST_THREADS=1
diff --git a/.devcontainer/docker-compose.yml b/.devcontainer/docker-compose.yml
@@ -12,6 +12,7 @@ services:
       - ${USERPROFILE}${HOME}/.gitconfig:/home/rust/.gitconfig:ro
       - ${USERPROFILE}${HOME}/.aws:/home/rust/.aws:rw
       - ${USERPROFILE}${HOME}/.azure:/home/rust/.azure:rw
+      - ${USERPROFILE}${HOME}/.config/gcloud:/home/rust/.config/gcloud:rw
       - ./entrypoint.sh:/entrypoint.sh
     env_file:
       - .env
@@ -21,6 +22,7 @@ services:
       - minio
       - azurite
       - webdav
+      - fake-gcs-server
 
   minio:
     image: minio/minio
@@ -61,3 +63,16 @@ services:
       interval: 6s
       timeout: 2s
       retries: 3
+
+  fake-gcs-server:
+     image: tustvold/fake-gcs-server
+     env_file:
+       - .env
+     network_mode: host
+     command: -scheme http -public-host localhost:4443
+     restart: unless-stopped
+     healthcheck:
+       test: ["CMD", "nc", "-z", "localhost", "4443"]
+       interval: 6s
+       timeout: 2s
+       retries: 3
diff --git a/.devcontainer/entrypoint.sh b/.devcontainer/entrypoint.sh
@@ -6,4 +6,8 @@ trap "echo 'Caught termination signal. Exiting...'; exit 0" SIGINT SIGTERM
 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
 
+# create fake-gcs bucket
+curl -v -X POST --data-binary "{\"name\":\"$GOOGLE_TEST_BUCKET\"}" -H "Content-Type: application/json" "$GOOGLE_SERVICE_ENDPOINT/storage/v1/b"
+curl -v -X POST --data-binary "{\"name\":\"${GOOGLE_TEST_BUCKET}2\"}" -H "Content-Type: application/json" "$GOOGLE_SERVICE_ENDPOINT/storage/v1/b"
+
 sleep infinity
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -152,6 +152,22 @@ jobs:
             sleep 1
           done
 
+      - name: Start fake-gcs-server for Google Cloud Storage emulator tests
+        run: |
+          docker run -d \
+            --env-file .devcontainer/.env \
+            -p 4443:4443 \
+            tustvold/fake-gcs-server -scheme http -public-host localhost:4443
+
+          while ! curl $GOOGLE_SERVICE_ENDPOINT; do   
+            echo "Waiting for $GOOGLE_SERVICE_ENDPOINT..."
+            sleep 1
+          done
+
+          # create bucket
+          curl -v -X POST --data-binary "{\"name\":\"$GOOGLE_TEST_BUCKET\"}" -H "Content-Type: application/json" "$GOOGLE_SERVICE_ENDPOINT/storage/v1/b"
+          curl -v -X POST --data-binary "{\"name\":\"${GOOGLE_TEST_BUCKET}2\"}" -H "Content-Type: application/json" "$GOOGLE_SERVICE_ENDPOINT/storage/v1/b"
+
       - name: Run tests
         run: |
           # Run tests with coverage tool
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -28,7 +28,7 @@ aws-credential-types = {version = "1", default-features = false}
 azure_storage = {version = "0.21", default-features = false}
 futures = "0.3"
 home = "0.5"
-object_store = {version = "0.11", default-features = false, features = ["aws", "azure", "http"]}
+object_store = {version = "0.11", default-features = false, features = ["aws", "azure", "gcp", "http"]}
 once_cell = "1"
 parquet = {version = "54", default-features = false, features = [
     "arrow",
diff --git a/README.md b/README.md
@@ -183,7 +183,7 @@ SELECT * FROM parquet.column_stats('/tmp/product_example.parquet')
 ```
 
 ## Object Store Support
-`pg_parquet` supports reading and writing Parquet files from/to `S3`, `Azure Blob Storage` and `http(s)` object stores.
+`pg_parquet` supports reading and writing Parquet files from/to `S3`, `Azure Blob Storage`, `http(s)` and `Google Cloud Storage` object stores.
 
 > [!NOTE]
 > To be able to write into a object store location, you need to grant `parquet_object_store_write` role to your current postgres user.
@@ -270,6 +270,28 @@ Supported authorization methods' priority order is shown below:
 
 `Https` uris are supported by default. You can set `ALLOW_HTTP` environment variable to allow `http` uris.
 
+#### Google Cloud Storage
+
+The simplest way to configure object storage is by creating a json config file like [`~/.config/gcloud/application_default_credentials.json`] (can be generated by `gcloud auth application-default login`):
+
+```bash
+$ cat ~/.config/gcloud/application_default_credentials.json
+{
+  "gcs_base_url": "http://localhost:4443",     
+  "disable_oauth": true,
+  "client_email": "",
+  "private_key_id": "",
+  "private_key": ""
+}
+```
+
+Alternatively, you can use the following environment variables when starting postgres to configure the Google Cloud Storage client:
+- `GOOGLE_SERVICE_ACCOUNT_KEY`: json serialized service account key **(only via environment variables)**
+- `GOOGLE_SERVICE_ACCOUNT_PATH`: an alternative location for the config file **(only via environment variables)**
+
+Supported Google Cloud Storage uri formats are shown below:
+- gs:// \<bucket\> / \<path\>
+
 ## Copy Options
 `pg_parquet` supports the following options in the `COPY TO` command:
 - `format parquet`: you need to specify this option to read or write Parquet files which does not end with `.parquet[.<compression>]` extension,
diff --git a/src/arrow_parquet/uri_utils.rs b/src/arrow_parquet/uri_utils.rs
@@ -19,8 +19,8 @@ use url::Url;
 
 use crate::{
     object_store::{
-        aws::parse_s3_bucket, azure::parse_azure_blob_container, http::parse_http_base_uri,
-        object_store_cache::get_or_create_object_store,
+        aws::parse_s3_bucket, azure::parse_azure_blob_container, gcs::parse_gcs_bucket,
+        http::parse_http_base_uri, object_store_cache::get_or_create_object_store,
     },
     PG_BACKEND_TOKIO_RUNTIME,
 };
@@ -50,7 +50,7 @@ impl ParsedUriInfo {
     fn try_parse_scheme(uri: &Url) -> Result<(ObjectStoreScheme, Path), String> {
         ObjectStoreScheme::parse(uri).map_err(|_| {
             format!(
-                "unrecognized uri {}. pg_parquet supports local paths, https://, s3:// or az:// schemes.",
+                "unrecognized uri {}. pg_parquet supports local paths, https://, s3://, az:// or gs:// schemes.",
                 uri
             )
         })
@@ -67,8 +67,11 @@ impl ParsedUriInfo {
             ObjectStoreScheme::Http => parse_http_base_uri(uri).
                 ok_or(format!("unsupported http storage uri: {uri}"))
                 .map(Some),
+            ObjectStoreScheme::GoogleCloudStorage => parse_gcs_bucket(uri)
+                .ok_or(format!("unsupported gcs uri {uri}"))
+                .map(Some),
             ObjectStoreScheme::Local => Ok(None),
-            _ => Err(format!("unsupported scheme {} in uri {}. pg_parquet supports local paths, https://, s3:// or az:// schemes.",
+            _ => Err(format!("unsupported scheme {} in uri {}. pg_parquet supports local paths, https://, s3://, az:// or gs:// schemes.",
                             uri.scheme(), uri))
         }
     }
diff --git a/src/object_store.rs b/src/object_store.rs
@@ -1,14 +1,6 @@
-use crate::{
-    arrow_parquet::uri_utils::uri_as_string,
-    object_store::{
-        aws::create_s3_object_store, azure::create_azure_object_store,
-        local_file::create_local_file_object_store,
-    },
-    PG_BACKEND_TOKIO_RUNTIME,
-};
-
 pub(crate) mod aws;
 pub(crate) mod azure;
+pub(crate) mod gcs;
 pub(crate) mod http;
 pub(crate) mod local_file;
 pub(crate) mod object_store_cache;
diff --git a/src/object_store/aws.rs b/src/object_store/aws.rs
@@ -5,7 +5,9 @@ use aws_credential_types::provider::ProvideCredentials;
 use object_store::aws::AmazonS3Builder;
 use url::Url;
 
-use super::{object_store_cache::ObjectStoreWithExpiration, PG_BACKEND_TOKIO_RUNTIME};
+use crate::PG_BACKEND_TOKIO_RUNTIME;
+
+use super::object_store_cache::ObjectStoreWithExpiration;
 
 // 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.
diff --git a/src/object_store/gcs.rs b/src/object_store/gcs.rs
@@ -0,0 +1,69 @@
+use std::sync::Arc;
+
+use object_store::gcp::GoogleCloudStorageBuilder;
+use url::Url;
+
+use super::object_store_cache::ObjectStoreWithExpiration;
+
+// create_gcs_object_store a GoogleCloudStorage object store from given uri.
+// It is configured by environment variables. Currently, we only support
+// following environment variables:
+// - GOOGLE_SERVICE_ACCOUNT_KEY
+// - GOOGLE_SERVICE_ACCOUNT_PATH
+pub(crate) fn create_gcs_object_store(uri: &Url) -> ObjectStoreWithExpiration {
+    let bucket_name = parse_gcs_bucket(uri).unwrap_or_else(|| {
+        panic!("unsupported gcs uri: {}", uri);
+    });
+
+    let mut gcs_builder = GoogleCloudStorageBuilder::new().with_bucket_name(bucket_name);
+
+    let gcs_config = GoogleStorageConfig::load();
+
+    // service account key
+    if let Some(service_account_key) = gcs_config.service_account_key {
+        gcs_builder = gcs_builder.with_service_account_key(&service_account_key);
+    }
+
+    // service account path
+    if let Some(service_account_path) = gcs_config.service_account_path {
+        gcs_builder = gcs_builder.with_service_account_path(&service_account_path);
+    }
+
+    let object_store = gcs_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,
+    }
+}
+
+pub(crate) fn parse_gcs_bucket(uri: &Url) -> Option<String> {
+    let host = uri.host_str()?;
+
+    // gs://{bucket}/key
+    if uri.scheme() == "gs" {
+        return Some(host.to_string());
+    }
+
+    None
+}
+
+// GoogleStorageConfig is a struct that holds the configuration that is
+// used to configure the Google Storage object store.
+struct GoogleStorageConfig {
+    service_account_key: Option<String>,
+    service_account_path: Option<String>,
+}
+
+impl GoogleStorageConfig {
+    // load loads the Google Storage configuration from the environment.
+    fn load() -> Self {
+        Self {
+            service_account_key: std::env::var("GOOGLE_SERVICE_ACCOUNT_KEY").ok(),
+            service_account_path: std::env::var("GOOGLE_SERVICE_ACCOUNT_PATH").ok(),
+        }
+    }
+}
diff --git a/src/object_store/local_file.rs b/src/object_store/local_file.rs
@@ -3,7 +3,9 @@ use std::sync::Arc;
 use object_store::local::LocalFileSystem;
 use url::Url;
 
-use super::{object_store_cache::ObjectStoreWithExpiration, uri_as_string};
+use crate::arrow_parquet::uri_utils::uri_as_string;
+
+use super::object_store_cache::ObjectStoreWithExpiration;
 
 // create_local_file_object_store creates a LocalFileSystem object store with the given path.
 pub(crate) fn create_local_file_object_store(
diff --git a/src/object_store/object_store_cache.rs b/src/object_store/object_store_cache.rs
@@ -13,8 +13,8 @@ use url::Url;
 use crate::arrow_parquet::uri_utils::ParsedUriInfo;
 
 use super::{
-    create_azure_object_store, create_local_file_object_store, create_s3_object_store,
-    http::create_http_object_store,
+    aws::create_s3_object_store, azure::create_azure_object_store, gcs::create_gcs_object_store,
+    http::create_http_object_store, local_file::create_local_file_object_store,
 };
 
 // OBJECT_STORE_CACHE is a global cache for object stores per Postgres session.
@@ -85,14 +85,15 @@ impl ObjectStoreCache {
 
     fn create(scheme: ObjectStoreScheme, uri: &Url, copy_from: bool) -> ObjectStoreWithExpiration {
         // 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.
+        // local, s3, azure, https and gs schemes with a subset of all supported paths.
         match scheme {
             ObjectStoreScheme::AmazonS3 => create_s3_object_store(uri),
             ObjectStoreScheme::MicrosoftAzure => create_azure_object_store(uri),
             ObjectStoreScheme::Http => create_http_object_store(uri),
+            ObjectStoreScheme::GoogleCloudStorage => create_gcs_object_store(uri),
             ObjectStoreScheme::Local => create_local_file_object_store(uri, copy_from),
             _ => panic!(
-                    "unsupported scheme {} in uri {}. pg_parquet supports local paths, https, s3:// or az:// schemes.",
+                    "unsupported scheme {} in uri {}. pg_parquet supports local paths, https://, s3://, az:// or gs:// schemes.",
                     uri.scheme(),
                     uri
                 ),
diff --git a/src/pgrx_tests/object_store.rs b/src/pgrx_tests/object_store.rs
