storage: add benchwrapper, support STORAGE_EMULATOR_HOST

core/google/cloud/environment_vars.py

@@ -30,6 +30,9 @@
 BIGTABLE_EMULATOR = "BIGTABLE_EMULATOR_HOST"
 """Environment variable defining host for Bigtable emulator."""
 
+STORAGE_EMULATOR = "STORAGE_EMULATOR_HOST"
+"""Environment variable defining host for Storage emulator."""
+
 DISABLE_GRPC = "GOOGLE_CLOUD_DISABLE_GRPC"
 """Environment variable acting as flag to disable gRPC.
 

storage/google/cloud/storage/blob.py

@@ -60,13 +60,19 @@
 from google.cloud.storage.acl import ObjectACL
 
 
+_STORAGE_HOST = (
+    u"https://www.googleapis.com"
+)
+if os.environ.get("STORAGE_EMULATOR_HOST"):
+    _STORAGE_HOST = os.environ.get("STORAGE_EMULATOR_HOST")
+
 _API_ACCESS_ENDPOINT = "https://storage.googleapis.com"
 _DEFAULT_CONTENT_TYPE = u"application/octet-stream"
 _DOWNLOAD_URL_TEMPLATE = (
-    u"https://www.googleapis.com/download/storage/v1{path}?alt=media"
+    _STORAGE_HOST + u"/download/storage/v1{path}?alt=media"
 )
 _BASE_UPLOAD_TEMPLATE = (
-    u"https://www.googleapis.com/upload/storage/v1{bucket_path}/o?uploadType="
+    _STORAGE_HOST + u"/upload/storage/v1{bucket_path}/o?uploadType="
 )
 _MULTIPART_URL_TEMPLATE = _BASE_UPLOAD_TEMPLATE + u"multipart"
 _RESUMABLE_URL_TEMPLATE = _BASE_UPLOAD_TEMPLATE + u"resumable"

storage/google/cloud/storage/client.py

@@ -14,6 +14,8 @@
 
 """Client for interacting with the Google Cloud Storage API."""
 
+import os
+
 import google.api_core.client_options
 
 from google.auth.credentials import AnonymousCredentials
@@ -94,6 +96,10 @@ def __init__(
         )
 
         kw_args = {"client_info": client_info}
+
+        if os.environ.get("STORAGE_EMULATOR_HOST"):
+            kw_args["api_endpoint"] = os.environ.get("STORAGE_EMULATOR_HOST")
+
         if client_options:
             if type(client_options) == dict:
                 client_options = google.api_core.client_options.from_dict(

storage/test_utils/README.md

@@ -0,0 +1,19 @@
+# storage benchwrapp
+
+main.py is a gRPC wrapper around the storage library for benchmarking purposes.
+
+## Running
+
+```
+cd storage
+pip install -e . # install google.cloud.storage locally
+cd test_utils
+python3 benchwrapper.py --port 8081
+```
+
+## Re-generating protos
+
+```
+pip install grpcio-tools
+python -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. *.proto
+```

storage/test_utils/benchwrapper.py

@@ -0,0 +1,48 @@
+import argparse
+import sys
+import time
+import grpc
+import os
+from concurrent import futures
+import storage_pb2_grpc
+import storage_pb2
+from google.cloud import storage
+
+_ONE_DAY_IN_SECONDS = 60 * 60 * 24
+
+parser = argparse.ArgumentParser()
+
+if os.environ.get('STORAGE_EMULATOR_HOST') == None:
+    sys.exit('This benchmarking server only works when connected to an emulator. Please set STORAGE_EMULATOR_HOST.')
+
+parser.add_argument('--port', help='The port to run on.')
+
+args = parser.parse_args()
+
+if args.port == None:
+    sys.exit('Usage: python3 main.py --port 8081')
+
+client = storage.Client()
+
+class StorageBenchWrapperServicer(storage_pb2_grpc.StorageBenchWrapperServicer):
+    def Write(self, request, context):
+        # TODO(deklerk): implement this
+        return storage_pb2.EmptyResponse()
+
+    def Read(self, request, context):
+        bucket = client.get_bucket(request.bucketName)
+        blob = storage.Blob(request.objectName, bucket)
+        blob.download_as_string()
+        return storage_pb2.EmptyResponse()
+
+server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
+storage_pb2_grpc.add_StorageBenchWrapperServicer_to_server(StorageBenchWrapperServicer(), server)
+
+print('listening on localhost:'+args.port)
+server.add_insecure_port('[::]:'+args.port)
+server.start()
+try:
+    while True:
+        time.sleep(_ONE_DAY_IN_SECONDS)
+except KeyboardInterrupt:
+    server.stop(0)

storage/test_utils/storage.proto

@@ -0,0 +1,43 @@
+// Copyright 2019 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+syntax = "proto3";
+
+package storage_bench;
+
+message ObjectRead{
+  // The bucket string identifier.
+  string bucketName = 1;
+  // The object/blob string identifier.
+  string objectName = 2;
+}
+
+message ObjectWrite{
+   // The bucket string identifier.
+  string bucketName = 1;
+  // The object/blob string identifiers.
+  string objectName = 2;
+  // The string containing the upload file path.
+  string destination = 3;
+}
+
+message EmptyResponse{
+}
+
+service StorageBenchWrapper{
+  // Performs an upload from a specific object.
+  rpc Write(ObjectWrite) returns (EmptyResponse) {}
+  // Read a specific object.
+  rpc Read(ObjectRead) returns (EmptyResponse){}
+}