Add support for 'STORAGE_EMULATOR_HOST' env var; add 'benchwrapper' s…

…cript. (#9219)

storage/google/cloud/storage/_helpers.py

@@ -19,6 +19,16 @@
 
 import base64
 from hashlib import md5
+import os
+
+STORAGE_EMULATOR_ENV_VAR = "STORAGE_EMULATOR_HOST"
+"""Environment variable defining host for Storage emulator."""
+
+_DEFAULT_STORAGE_HOST = u"https://www.googleapis.com"
+
+
+def _get_storage_host():
+    return os.environ.get(STORAGE_EMULATOR_ENV_VAR, _DEFAULT_STORAGE_HOST)
 
 
 def _validate_name(name):

storage/google/cloud/storage/blob.py

@@ -47,27 +47,25 @@
 from google.resumable_media.requests import ResumableUpload
 
 from google.cloud import exceptions
+from google.cloud._helpers import _bytes_to_unicode
 from google.cloud._helpers import _rfc3339_to_datetime
 from google.cloud._helpers import _to_bytes
-from google.cloud._helpers import _bytes_to_unicode
 from google.cloud.exceptions import NotFound
 from google.api_core.iam import Policy
+from google.cloud.storage._helpers import _get_storage_host
 from google.cloud.storage._helpers import _PropertyMixin
 from google.cloud.storage._helpers import _scalar_property
 from google.cloud.storage._signing import generate_signed_url_v2
 from google.cloud.storage._signing import generate_signed_url_v4
 from google.cloud.storage.acl import ACL
 from google.cloud.storage.acl import ObjectACL
 
+_STORAGE_HOST = _get_storage_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"
-)
-_BASE_UPLOAD_TEMPLATE = (
-    u"https://www.googleapis.com/upload/storage/v1{bucket_path}/o?uploadType="
-)
+_DOWNLOAD_URL_TEMPLATE = _STORAGE_HOST + u"/download/storage/v1{path}?alt=media"
+_BASE_UPLOAD_TEMPLATE = _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"
 # NOTE: "acl" is also writeable but we defer ACL management to

storage/google/cloud/storage/client.py

@@ -22,6 +22,7 @@
 from google.cloud._helpers import _LocalStack
 from google.cloud.client import ClientWithProject
 from google.cloud.exceptions import NotFound
+from google.cloud.storage._helpers import _get_storage_host
 from google.cloud.storage._http import Connection
 from google.cloud.storage.batch import Batch
 from google.cloud.storage.bucket import Bucket
@@ -94,6 +95,9 @@ def __init__(
         )
 
         kw_args = {"client_info": client_info}
+
+        kw_args["api_endpoint"] = _get_storage_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,21 @@
+# storage benchwrapp
+
+main.py is a gRPC wrapper around the storage library for benchmarking purposes.
+
+## Running
+
+```
+export STORAGE_EMULATOR_HOST=localhost:8080
+pip install grpcio
+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.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){}
+}