Firestore: support emulator in client. (#8721)

firestore/google/cloud/firestore_v1/client.py

@@ -23,6 +23,8 @@
 * a :class:`~google.cloud.firestore_v1.client.Client` owns a
   :class:`~google.cloud.firestore_v1.document.DocumentReference`
 """
+import os
+
 from google.api_core.gapic_v1 import client_info
 from google.cloud.client import ClientWithProject
 
@@ -51,6 +53,7 @@
 _ACTIVE_TXN = "There is already an active transaction."
 _INACTIVE_TXN = "There is no active transaction."
 _CLIENT_INFO = client_info.ClientInfo(client_library_version=__version__)
+_FIRESTORE_EMULATOR_HOST = "FIRESTORE_EMULATOR_HOST"
 
 
 class Client(ClientWithProject):
@@ -103,6 +106,7 @@ def __init__(
         )
         self._client_info = client_info
         self._database = database
+        self._emulator_host = os.getenv(_FIRESTORE_EMULATOR_HOST)
 
     @property
     def _firestore_api(self):
@@ -115,11 +119,17 @@ def _firestore_api(self):
         if self._firestore_api_internal is None:
             # Use a custom channel.
             # We need this in order to set appropriate keepalive options.
-            channel = firestore_grpc_transport.FirestoreGrpcTransport.create_channel(
-                self._target,
-                credentials=self._credentials,
-                options={"grpc.keepalive_time_ms": 30000}.items(),
-            )
+
+            if self._emulator_host is not None:
+                channel = firestore_grpc_transport.firestore_pb2_grpc.grpc.insecure_channel(
+                    self._emulator_host
+                )
+            else:
+                channel = firestore_grpc_transport.FirestoreGrpcTransport.create_channel(
+                    self._target,
+                    credentials=self._credentials,
+                    options={"grpc.keepalive_time_ms": 30000}.items(),
+                )
 
             self._transport = firestore_grpc_transport.FirestoreGrpcTransport(
                 address=self._target, channel=channel
@@ -138,6 +148,9 @@ def _target(self):
         Returns:
             str: The location of the API.
         """
+        if self._emulator_host is not None:
+            return self._emulator_host
+
         return firestore_client.FirestoreClient.SERVICE_ADDRESS
 
     @property
@@ -179,6 +192,10 @@ def _rpc_metadata(self):
                 self._database_string
             )
 
+            if self._emulator_host is not None:
+                # The emulator requires additional metadata to be set.
+                self._rpc_metadata_internal.append(("authorization", "Bearer owner"))
+
         return self._rpc_metadata_internal
 
     def collection(self, *collection_path):

firestore/tests/unit/v1/test_client.py

@@ -47,6 +47,18 @@ def test_constructor(self):
         self.assertEqual(client._credentials, credentials)
         self.assertEqual(client._database, DEFAULT_DATABASE)
         self.assertIs(client._client_info, _CLIENT_INFO)
+        self.assertIsNone(client._emulator_host)
+
+    def test_constructor_with_emulator_host(self):
+        from google.cloud.firestore_v1.client import _FIRESTORE_EMULATOR_HOST
+
+        credentials = _make_credentials()
+        emulator_host = "localhost:8081"
+        with mock.patch("os.getenv") as getenv:
+            getenv.return_value = emulator_host
+            client = self._make_one(project=self.PROJECT, credentials=credentials)
+            self.assertEqual(client._emulator_host, emulator_host)
+            getenv.assert_called_once_with(_FIRESTORE_EMULATOR_HOST)
 
     def test_constructor_explicit(self):
         credentials = _make_credentials()
@@ -64,7 +76,7 @@ def test_constructor_explicit(self):
         self.assertIs(client._client_info, client_info)
 
     @mock.patch(
-        "google.cloud.firestore_v1.gapic.firestore_client." "FirestoreClient",
+        "google.cloud.firestore_v1.gapic.firestore_client.FirestoreClient",
         autospec=True,
         return_value=mock.sentinel.firestore_api,
     )
@@ -84,6 +96,34 @@ def test__firestore_api_property(self, mock_client):
         self.assertIs(client._firestore_api, mock_client.return_value)
         self.assertEqual(mock_client.call_count, 1)
 
+    @mock.patch(
+        "google.cloud.firestore_v1.gapic.firestore_client.FirestoreClient",
+        autospec=True,
+        return_value=mock.sentinel.firestore_api,
+    )
+    @mock.patch(
+        "google.cloud.firestore_v1.gapic.transports.firestore_grpc_transport.firestore_pb2_grpc.grpc.insecure_channel",
+        autospec=True,
+    )
+    def test__firestore_api_property_with_emulator(
+        self, mock_insecure_channel, mock_client
+    ):
+        emulator_host = "localhost:8081"
+        with mock.patch("os.getenv") as getenv:
+            getenv.return_value = emulator_host
+            client = self._make_default_one()
+
+        self.assertIsNone(client._firestore_api_internal)
+        firestore_api = client._firestore_api
+        self.assertIs(firestore_api, mock_client.return_value)
+        self.assertIs(firestore_api, client._firestore_api_internal)
+
+        mock_insecure_channel.assert_called_once_with(emulator_host)
+
+        # Call again to show that it is cached, but call count is still 1.
+        self.assertIs(client._firestore_api, mock_client.return_value)
+        self.assertEqual(mock_client.call_count, 1)
+
     def test___database_string_property(self):
         credentials = _make_credentials()
         database = "cheeeeez"
@@ -112,6 +152,25 @@ def test___rpc_metadata_property(self):
             [("google-cloud-resource-prefix", client._database_string)],
         )
 
+    def test__rpc_metadata_property_with_emulator(self):
+        emulator_host = "localhost:8081"
+        with mock.patch("os.getenv") as getenv:
+            getenv.return_value = emulator_host
+
+            credentials = _make_credentials()
+            database = "quanta"
+            client = self._make_one(
+                project=self.PROJECT, credentials=credentials, database=database
+            )
+
+        self.assertEqual(
+            client._rpc_metadata,
+            [
+                ("google-cloud-resource-prefix", client._database_string),
+                ("authorization", "Bearer owner"),
+            ],
+        )
+
     def test_collection_factory(self):
         from google.cloud.firestore_v1.collection import CollectionReference