Merge pull request #2244 from dhermes/fix-ds-emulator

dhermes · web-flow · commit 4673a436ffea · 2016-09-07T18:46:42.000-07:00
Upgrading datastore emulator to work with gRPC
diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst
@@ -208,13 +208,14 @@ Running System Tests
   ``datastore`` emulator::
 
    $ tox -e datastore-emulator
+   $ GOOGLE_CLOUD_DISABLE_GRPC=true tox -e datastore-emulator
 
   This also requires that the ``gcloud`` command line tool is
   installed. If you'd like to run them directly (outside of a
   ``tox`` environment), first start the emulator and
   take note of the process ID::
 
-   $ gcloud beta emulators datastore start 2>&1 > log.txt &
+   $ gcloud beta emulators datastore start --no-legacy 2>&1 > log.txt &
    [1] 33333
 
   then determine the environment variables needed to interact with
diff --git a/docs/logging-usage.rst b/docs/logging-usage.rst
@@ -16,8 +16,8 @@ Authentication and Configuration
 - The library now enables the ``gRPC`` transport for the logging API by
   default, assuming that the required dependencies are installed and
   importable.  To *disable* this transport, set the
-  :envvar:`GOOGLE_CLOUD_DISABLE_GAX` environment variable to a non-empty string,
-  e.g.:  ``$ export GOOGLE_CLOUD_DISABLE_GAX=1``.
+  :envvar:`GOOGLE_CLOUD_DISABLE_GRPC` environment variable to a
+  non-empty string, e.g.:  ``$ export GOOGLE_CLOUD_DISABLE_GRPC=true``.
 
 - After configuring your environment, create a
   :class:`Client <google.cloud.logging.client.Client>`
diff --git a/docs/pubsub-usage.rst b/docs/pubsub-usage.rst
@@ -15,8 +15,8 @@ Authentication / Configuration
 - The library now enables the ``gRPC`` transport for the pubsub API by
   default, assuming that the required dependencies are installed and
   importable.  To *disable* this transport, set the
-  :envvar:`GOOGLE_CLOUD_DISABLE_GAX` environment variable to a non-empty string,
-  e.g.:  ``$ export GOOGLE_CLOUD_DISABLE_GAX=1``.
+  :envvar:`GOOGLE_CLOUD_DISABLE_GRPC` environment variable to a
+  non-empty string, e.g.:  ``$ export GOOGLE_CLOUD_DISABLE_GRPC=true``.
 
 - :class:`Client <google.cloud.pubsub.client.Client>` objects hold both a ``project``
   and an authenticated connection to the PubSub service.
diff --git a/google/cloud/_helpers.py b/google/cloud/_helpers.py
@@ -39,7 +39,7 @@
     grpc = None
     _Rendezvous = Exception
 import six
-from six.moves.http_client import HTTPConnection
+from six.moves import http_client
 from six.moves import configparser
 
 # pylint: disable=ungrouped-imports
@@ -269,7 +269,7 @@ def _compute_engine_id():
     host = '169.254.169.254'
     uri_path = '/computeMetadata/v1/project/project-id'
     headers = {'Metadata-Flavor': 'Google'}
-    connection = HTTPConnection(host, timeout=0.1)
+    connection = http_client.HTTPConnection(host, timeout=0.1)
 
     try:
         connection.request('GET', uri_path, headers=headers)
@@ -612,8 +612,8 @@ def __call__(self, unused_context, callback):
         callback(headers, None)
 
 
-def make_stub(credentials, user_agent, stub_class, host, port):
-    """Makes a stub for an RPC service.
+def make_secure_stub(credentials, user_agent, stub_class, host):
+    """Makes a secure stub for an RPC service.
 
     Uses / depends on gRPC.
 
@@ -630,25 +630,49 @@ def make_stub(credentials, user_agent, stub_class, host, port):
     :type host: str
     :param host: The host for the service.
 
-    :type port: int
-    :param port: The port for the service.
-
     :rtype: object, instance of ``stub_class``
     :returns: The stub object used to make gRPC requests to a given API.
     """
-    # Leaving the first argument to ssl_channel_credentials() as None
-    # loads root certificates from `grpc/_adapter/credentials/roots.pem`.
+    # ssl_channel_credentials() loads root certificates from
+    # `grpc/_adapter/credentials/roots.pem`.
     transport_creds = grpc.ssl_channel_credentials()
     custom_metadata_plugin = MetadataPlugin(credentials, user_agent)
     auth_creds = grpc.metadata_call_credentials(
         custom_metadata_plugin, name='google_creds')
     channel_creds = grpc.composite_channel_credentials(
         transport_creds, auth_creds)
-    target = '%s:%d' % (host, port)
+    target = '%s:%d' % (host, http_client.HTTPS_PORT)
     channel = grpc.secure_channel(target, channel_creds)
     return stub_class(channel)
 
 
+def make_insecure_stub(stub_class, host, port=None):
+    """Makes an insecure stub for an RPC service.
+
+    Uses / depends on gRPC.
+
+    :type stub_class: type
+    :param stub_class: A gRPC stub type for a given service.
+
+    :type host: str
+    :param host: The host for the service. May also include the port
+                 if ``port`` is unspecified.
+
+    :type port: int
+    :param port: (Optional) The port for the service.
+
+    :rtype: object, instance of ``stub_class``
+    :returns: The stub object used to make gRPC requests to a given API.
+    """
+    if port is None:
+        target = host
+    else:
+        # NOTE: This assumes port != http_client.HTTPS_PORT:
+        target = '%s:%d' % (host, port)
+    channel = grpc.insecure_channel(target)
+    return stub_class(channel)
+
+
 def exc_to_code(exc):
     """Retrieves the status code from a gRPC exception.
 
diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py
@@ -29,7 +29,7 @@
 
 from pkg_resources import get_distribution
 
-from google.cloud._helpers import make_stub
+from google.cloud._helpers import make_secure_stub
 from google.cloud.bigtable._generated import bigtable_instance_admin_pb2
 from google.cloud.bigtable._generated import bigtable_pb2
 from google.cloud.bigtable._generated import bigtable_table_admin_pb2
@@ -44,21 +44,14 @@
 
 TABLE_ADMIN_HOST = 'bigtableadmin.googleapis.com'
 """Table Admin API request host."""
-TABLE_ADMIN_PORT = 443
-"""Table Admin API request port."""
 
 INSTANCE_ADMIN_HOST = 'bigtableadmin.googleapis.com'
 """Cluster Admin API request host."""
-INSTANCE_ADMIN_PORT = 443
-"""Cluster Admin API request port."""
 
 DATA_API_HOST = 'bigtable.googleapis.com'
 """Data API request host."""
-DATA_API_PORT = 443
-"""Data API request port."""
 
 OPERATIONS_API_HOST = INSTANCE_ADMIN_HOST
-OPERATIONS_API_PORT = INSTANCE_ADMIN_PORT
 
 ADMIN_SCOPE = 'https://www.googleapis.com/auth/bigtable.admin'
 """Scope for interacting with the Cluster Admin and Table Admin APIs."""
@@ -81,9 +74,8 @@ def _make_data_stub(client):
     :rtype: :class:`._generated.bigtable_pb2.BigtableStub`
     :returns: A gRPC stub object.
     """
-    return make_stub(client.credentials, client.user_agent,
-                     bigtable_pb2.BigtableStub,
-                     DATA_API_HOST, DATA_API_PORT)
+    return make_secure_stub(client.credentials, client.user_agent,
+                            bigtable_pb2.BigtableStub, DATA_API_HOST)
 
 
 def _make_instance_stub(client):
@@ -95,9 +87,10 @@ def _make_instance_stub(client):
     :rtype: :class:`.bigtable_instance_admin_pb2.BigtableInstanceAdminStub`
     :returns: A gRPC stub object.
     """
-    return make_stub(client.credentials, client.user_agent,
-                     bigtable_instance_admin_pb2.BigtableInstanceAdminStub,
-                     INSTANCE_ADMIN_HOST, INSTANCE_ADMIN_PORT)
+    return make_secure_stub(
+        client.credentials, client.user_agent,
+        bigtable_instance_admin_pb2.BigtableInstanceAdminStub,
+        INSTANCE_ADMIN_HOST)
 
 
 def _make_operations_stub(client):
@@ -112,9 +105,9 @@ def _make_operations_stub(client):
     :rtype: :class:`._generated.operations_grpc_pb2.OperationsStub`
     :returns: A gRPC stub object.
     """
-    return make_stub(client.credentials, client.user_agent,
-                     operations_grpc_pb2.OperationsStub,
-                     OPERATIONS_API_HOST, OPERATIONS_API_PORT)
+    return make_secure_stub(client.credentials, client.user_agent,
+                            operations_grpc_pb2.OperationsStub,
+                            OPERATIONS_API_HOST)
 
 
 def _make_table_stub(client):
@@ -126,9 +119,9 @@ def _make_table_stub(client):
     :rtype: :class:`.bigtable_instance_admin_pb2.BigtableTableAdminStub`
     :returns: A gRPC stub object.
     """
-    return make_stub(client.credentials, client.user_agent,
-                     bigtable_table_admin_pb2.BigtableTableAdminStub,
-                     TABLE_ADMIN_HOST, TABLE_ADMIN_PORT)
+    return make_secure_stub(client.credentials, client.user_agent,
+                            bigtable_table_admin_pb2.BigtableTableAdminStub,
+                            TABLE_ADMIN_HOST)
 
 
 class Client(_ClientFactoryMixin, _ClientProjectMixin):
diff --git a/google/cloud/datastore/connection.py b/google/cloud/datastore/connection.py
@@ -18,8 +18,10 @@
 
 from google.rpc import status_pb2
 
-from google.cloud._helpers import make_stub
+from google.cloud._helpers import make_insecure_stub
+from google.cloud._helpers import make_secure_stub
 from google.cloud import connection as connection_module
+from google.cloud.environment_vars import DISABLE_GRPC
 from google.cloud.environment_vars import GCD_HOST
 from google.cloud.exceptions import Conflict
 from google.cloud.exceptions import make_exception
@@ -41,8 +43,9 @@
 
 DATASTORE_API_HOST = 'datastore.googleapis.com'
 """Datastore API request host."""
-DATASTORE_API_PORT = 443
-"""Datastore API request port."""
+
+_DISABLE_GRPC = os.getenv(DISABLE_GRPC, False)
+_USE_GRPC = _HAVE_GRPC and not _DISABLE_GRPC
 
 
 class _DatastoreAPIOverHttp(object):
@@ -230,12 +233,20 @@ class _DatastoreAPIOverGRPC(object):
     :type connection: :class:`google.cloud.datastore.connection.Connection`
     :param connection: A connection object that contains helpful
                        information for making requests.
+
+    :type secure: bool
+    :param secure: Flag indicating if a secure stub connection is needed.
     """
 
-    def __init__(self, connection):
-        self._stub = make_stub(connection.credentials, connection.USER_AGENT,
-                               datastore_grpc_pb2.DatastoreStub,
-                               DATASTORE_API_HOST, DATASTORE_API_PORT)
+    def __init__(self, connection, secure):
+        if secure:
+            self._stub = make_secure_stub(connection.credentials,
+                                          connection.USER_AGENT,
+                                          datastore_grpc_pb2.DatastoreStub,
+                                          connection.host)
+        else:
+            self._stub = make_insecure_stub(datastore_grpc_pb2.DatastoreStub,
+                                            connection.host)
 
     def lookup(self, project, request_pb):
         """Perform a ``lookup`` request.
@@ -352,10 +363,6 @@ class Connection(connection_module.Connection):
 
     :type http: :class:`httplib2.Http` or class that defines ``request()``.
     :param http: An optional HTTP object to make requests.
-
-    :type api_base_url: string
-    :param api_base_url: The base of the API call URL. Defaults to
-                         :attr:`API_BASE_URL`.
     """
 
     API_BASE_URL = 'https://' + DATASTORE_API_HOST
@@ -371,18 +378,18 @@ class Connection(connection_module.Connection):
     SCOPE = ('https://www.googleapis.com/auth/datastore',)
     """The scopes required for authenticating as a Cloud Datastore consumer."""
 
-    def __init__(self, credentials=None, http=None, api_base_url=None):
+    def __init__(self, credentials=None, http=None):
         super(Connection, self).__init__(credentials=credentials, http=http)
-        if api_base_url is None:
-            try:
-                # gcd.sh has /datastore/ in the path still since it supports
-                # v1beta2 and v1beta3 simultaneously.
-                api_base_url = '%s/datastore' % (os.environ[GCD_HOST],)
-            except KeyError:
-                api_base_url = self.__class__.API_BASE_URL
-        self.api_base_url = api_base_url
-        if _HAVE_GRPC:
-            self._datastore_api = _DatastoreAPIOverGRPC(self)
+        try:
+            self.host = os.environ[GCD_HOST]
+            self.api_base_url = 'http://' + self.host
+            secure = False
+        except KeyError:
+            self.host = DATASTORE_API_HOST
+            self.api_base_url = self.__class__.API_BASE_URL
+            secure = True
+        if _USE_GRPC:
+            self._datastore_api = _DatastoreAPIOverGRPC(self, secure=secure)
         else:
             self._datastore_api = _DatastoreAPIOverHttp(self)
 
diff --git a/google/cloud/environment_vars.py b/google/cloud/environment_vars.py
@@ -27,11 +27,18 @@
 GCD_DATASET = 'DATASTORE_DATASET'
 """Environment variable defining default dataset ID under GCD."""
 
-GCD_HOST = 'DATASTORE_HOST'
+GCD_HOST = 'DATASTORE_EMULATOR_HOST'
 """Environment variable defining host for GCD dataset server."""
 
 PUBSUB_EMULATOR = 'PUBSUB_EMULATOR_HOST'
 """Environment variable defining host for Pub/Sub emulator."""
 
 CREDENTIALS = 'GOOGLE_APPLICATION_CREDENTIALS'
 """Environment variable defining location of Google credentials."""
+
+DISABLE_GRPC = 'GOOGLE_CLOUD_DISABLE_GRPC'
+"""Environment variable acting as flag to disable gRPC.
+
+To be used for APIs where both an HTTP and gRPC implementation
+exist.
+"""
diff --git a/google/cloud/logging/client.py b/google/cloud/logging/client.py
@@ -35,6 +35,7 @@
     _HAVE_GAX = True
 
 from google.cloud.client import JSONClient
+from google.cloud.environment_vars import DISABLE_GRPC
 from google.cloud.logging.connection import Connection
 from google.cloud.logging.connection import _LoggingAPI as JSONLoggingAPI
 from google.cloud.logging.connection import _MetricsAPI as JSONMetricsAPI
@@ -47,7 +48,7 @@
 from google.cloud.logging.sink import Sink
 
 
-_DISABLE_GAX = os.getenv('GOOGLE_CLOUD_DISABLE_GAX', False)
+_DISABLE_GAX = os.getenv(DISABLE_GRPC, False)
 _USE_GAX = _HAVE_GAX and not _DISABLE_GAX
 ASCENDING = 'timestamp asc'
 """Query string to order by ascending timestamps."""
diff --git a/google/cloud/pubsub/client.py b/google/cloud/pubsub/client.py
@@ -17,6 +17,7 @@
 import os
 
 from google.cloud.client import JSONClient
+from google.cloud.environment_vars import DISABLE_GRPC
 from google.cloud.pubsub.connection import Connection
 from google.cloud.pubsub.connection import _PublisherAPI as JSONPublisherAPI
 from google.cloud.pubsub.connection import _SubscriberAPI as JSONSubscriberAPI
@@ -41,7 +42,7 @@
 # pylint: enable=ungrouped-imports
 
 
-_DISABLE_GAX = os.getenv('GOOGLE_CLOUD_DISABLE_GAX', False)
+_DISABLE_GAX = os.getenv(DISABLE_GRPC, False)
 _USE_GAX = _HAVE_GAX and not _DISABLE_GAX
 
 
diff --git a/system_tests/run_emulator.py b/system_tests/run_emulator.py
@@ -35,7 +35,10 @@
     'datastore': (GCD_DATASET, GCD_HOST),
     'pubsub': (PUBSUB_EMULATOR,)
 }
-_DS_READY_LINE = '[datastore] INFO: Dev App Server is now running\n'
+EXTRA = {
+    'datastore': ('--no-legacy',),
+}
+_DS_READY_LINE = '[datastore] Dev App Server is now running.\n'
 _PS_READY_LINE_PREFIX = '[pubsub] INFO: Server started, listening on '
 
 
@@ -62,7 +65,9 @@ def get_start_command(package):
     :rtype: tuple
     :returns: The arguments to be used, in a tuple.
     """
-    return 'gcloud', 'beta', 'emulators', package, 'start'
+    result = ('gcloud', 'beta', 'emulators', package, 'start')
+    extra = EXTRA.get(package, ())
+    return result + extra
 
 
 def get_env_init_command(package):
@@ -74,7 +79,9 @@ def get_env_init_command(package):
     :rtype: tuple
     :returns: The arguments to be used, in a tuple.
     """
-    return 'gcloud', 'beta', 'emulators', package, 'env-init'
+    result = ('gcloud', 'beta', 'emulators', package, 'env-init')
+    extra = EXTRA.get(package, ())
+    return result + extra
 
 
 def datastore_wait_ready(popen):
diff --git a/tox.ini b/tox.ini
@@ -124,6 +124,8 @@ commands =
     python {toxinidir}/system_tests/run_emulator.py --package=datastore
 setenv =
     GOOGLE_CLOUD_NO_PRINT=true
+passenv =
+    GOOGLE_CLOUD_DISABLE_GRPC
 deps =
     {[testenv]deps}
     psutil
diff --git a/unit_tests/bigtable/test_client.py b/unit_tests/bigtable/test_client.py
diff --git a/unit_tests/datastore/test_connection.py b/unit_tests/datastore/test_connection.py
diff --git a/unit_tests/test__helpers.py b/unit_tests/test__helpers.py
