feat(firestore): configure gRPC message length limits for large documents (#18220)

- Configure `grpc.max_send_message_length` and
`grpc.max_receive_message_length` options on production and emulator
gRPC channels.
- Update unit tests in `test_base_client.py` to assert gRPC message size
options.
- Add system tests for large document CRUD and pipeline execution across
sync and async clients.

Author: ohmayr

packages/google-cloud-firestore/google/cloud/firestore_v1/base_client.py

@@ -77,6 +77,14 @@
 _INACTIVE_TXN: str = "There is no active transaction."
 _CLIENT_INFO: Any = client_info.ClientInfo(client_library_version=__version__)
 _FIRESTORE_EMULATOR_HOST: str = "FIRESTORE_EMULATOR_HOST"
+_GRPC_MSG_SIZE_OPTIONS: List[Tuple[str, int]] = [
+    ("grpc.max_send_message_length", -1),
+    ("grpc.max_receive_message_length", -1),
+]
+_DEFAULT_CHANNEL_OPTIONS: List[Tuple[str, Any]] = [
+    ("grpc.keepalive_time_ms", 30000),
+    *_GRPC_MSG_SIZE_OPTIONS,
+]
 
 
 class BaseClient(ClientWithProject):
@@ -173,7 +181,7 @@ def _firestore_api_helper(self, transport, client_class, client_module) -> Any:
                 channel = transport.create_channel(
                     self._target,
                     credentials=self._credentials,
-                    options={"grpc.keepalive_time_ms": 30000}.items(),
+                    options=_DEFAULT_CHANNEL_OPTIONS,
                 )
 
             self._transport = transport(host=self._target, channel=channel)
@@ -204,7 +212,10 @@ def _emulator_channel(self, transport):
             and getattr(self._credentials, "id_token", None) is not None
         ):
             token = self._credentials.id_token
-        options = [("Authorization", f"Bearer {token}")]
+        options = [
+            ("Authorization", f"Bearer {token}"),
+            *_GRPC_MSG_SIZE_OPTIONS,
+        ]
 
         if "GrpcAsyncIOTransport" in str(transport.__name__):
             return grpc.aio.insecure_channel(self._emulator_host, options=options)

packages/google-cloud-firestore/tests/system/test_system.py

@@ -3836,3 +3836,39 @@ def in_transaction(transaction, rollback):
     assert len(result) == 1
     assert len(result[0]) == 1
     assert result[0][0].value == expected
+
+
+@pytest.mark.skip(reason="Temporarily skipped. Not yet in production.")
+@pytest.mark.parametrize("database", [FIRESTORE_ENTERPRISE_DB], indirect=True)
+def test_large_document_standard_writes(client, cleanup, database):
+    """Test standard write and read operations for 5MB document on Enterprise DB."""
+    collection_id = "large_docs_" + UNIQUE_RESOURCE_ID
+    doc_ref = client.collection(collection_id).document("large_doc")
+    cleanup(doc_ref.delete)
+
+    large_payload = "a" * (5 * 1024 * 1024)
+    doc_ref.set({"payload": large_payload})
+
+    snapshot = doc_ref.get()
+    assert snapshot.exists
+    assert snapshot.to_dict() == {"payload": large_payload}
+
+
+@pytest.mark.skip(reason="Temporarily skipped. Not yet in production.")
+@pytest.mark.parametrize("method", ["execute", "stream"])
+@pytest.mark.parametrize("database", [FIRESTORE_ENTERPRISE_DB], indirect=True)
+def test_large_document_pipeline(client, cleanup, database, method):
+    """Test pipeline execution over 5MB document on Enterprise DB."""
+    collection_id = "large_pipeline_" + UNIQUE_RESOURCE_ID
+    col_ref = client.collection(collection_id)
+    doc_ref = col_ref.document("large_doc")
+    cleanup(doc_ref.delete)
+
+    large_payload = "b" * (5 * 1024 * 1024)
+    doc_ref.set({"payload": large_payload})
+
+    pipeline = client.pipeline().collection(collection_id)
+    method_under_test = getattr(pipeline, method)
+
+    results = list(method_under_test())
+    assert [doc.data() for doc in results] == [{"payload": large_payload}]

packages/google-cloud-firestore/tests/system/test_system_async.py

@@ -3708,3 +3708,41 @@ async def in_transaction(transaction):
         await in_transaction(transaction)
         # make sure we didn't skip assertions in inner function
         assert inner_fn_ran is True
+
+
+@pytest.mark.skip(reason="Temporarily skipped. Not yet in production.")
+@pytest.mark.parametrize("database", [FIRESTORE_ENTERPRISE_DB], indirect=True)
+async def test_large_document_standard_writes_async(client, cleanup, database):
+    """Test standard write and read operations for 5MB document on Enterprise DB (async)."""
+    collection_id = "large_docs_async_" + UNIQUE_RESOURCE_ID
+    doc_ref = client.collection(collection_id).document("large_doc")
+    cleanup(doc_ref.delete)
+
+    large_payload = "c" * (5 * 1024 * 1024)
+    await doc_ref.set({"payload": large_payload})
+
+    snapshot = await doc_ref.get()
+    assert snapshot.exists
+    assert snapshot.to_dict() == {"payload": large_payload}
+
+
+@pytest.mark.skip(reason="Temporarily skipped. Not yet in production.")
+@pytest.mark.parametrize("method", ["execute", "stream"])
+@pytest.mark.parametrize("database", [FIRESTORE_ENTERPRISE_DB], indirect=True)
+async def test_large_document_pipeline_async(client, cleanup, database, method):
+    """Test async pipeline execution over 5MB document on Enterprise DB."""
+    collection_id = "large_pipeline_async_" + UNIQUE_RESOURCE_ID
+    col_ref = client.collection(collection_id)
+    doc_ref = col_ref.document("large_doc")
+    cleanup(doc_ref.delete)
+
+    large_payload = "d" * (5 * 1024 * 1024)
+    await doc_ref.set({"payload": large_payload})
+
+    pipeline = client.pipeline().collection(collection_id)
+    if method == "execute":
+        results = await pipeline.execute()
+    else:
+        results = [doc async for doc in pipeline.stream()]
+
+    assert [doc.data() for doc in results] == [{"payload": large_payload}]

packages/google-cloud-firestore/tests/unit/v1/test_base_client.py

@@ -110,9 +110,48 @@ def test_baseclient__firestore_api_helper_wo_emulator():
 
     assert api is client_class.return_value
     assert client._firestore_api_internal is api
-    channel_options = {"grpc.keepalive_time_ms": 30000}
+    channel_options = [
+        ("grpc.keepalive_time_ms", 30000),
+        ("grpc.max_send_message_length", -1),
+        ("grpc.max_receive_message_length", -1),
+    ]
     transport_class.create_channel.assert_called_once_with(
-        target, credentials=client._credentials, options=channel_options.items()
+        target, credentials=client._credentials, options=channel_options
+    )
+    transport_class.assert_called_once_with(
+        host=target,
+        channel=transport_class.create_channel.return_value,
+    )
+    client_class.assert_called_once_with(
+        transport=transport_class.return_value, client_options=client_options
+    )
+
+
+def test_baseclient__firestore_api_helper_async_wo_emulator():
+    from google.cloud.firestore_v1.services.firestore.transports.grpc_asyncio import (
+        FirestoreGrpcAsyncIOTransport,
+    )
+
+    client = _make_default_base_client()
+    client_options = client._client_options = mock.Mock()
+    target = client._target
+    assert client._firestore_api_internal is None
+
+    transport_class = mock.Mock(spec=FirestoreGrpcAsyncIOTransport)
+    client_class = mock.Mock()
+    client_module = mock.Mock()
+
+    api = client._firestore_api_helper(transport_class, client_class, client_module)
+
+    assert api is client_class.return_value
+    assert client._firestore_api_internal is api
+    channel_options = [
+        ("grpc.keepalive_time_ms", 30000),
+        ("grpc.max_send_message_length", -1),
+        ("grpc.max_receive_message_length", -1),
+    ]
+    transport_class.create_channel.assert_called_once_with(
+        target, credentials=client._credentials, options=channel_options
     )
     transport_class.assert_called_once_with(
         host=target,
@@ -236,7 +275,23 @@ def test_baseclient__emulator_channel():
     with mock.patch("grpc.insecure_channel") as insecure_channel:
         channel = client._emulator_channel(FirestoreGrpcTransport)
         insecure_channel.assert_called_once_with(
-            emulator_host, options=[("Authorization", "Bearer test")]
+            emulator_host,
+            options=[
+                ("Authorization", "Bearer test"),
+                ("grpc.max_send_message_length", -1),
+                ("grpc.max_receive_message_length", -1),
+            ],
+        )
+
+    with mock.patch("grpc.aio.insecure_channel") as aio_insecure_channel:
+        channel = client._emulator_channel(FirestoreGrpcAsyncIOTransport)
+        aio_insecure_channel.assert_called_once_with(
+            emulator_host,
+            options=[
+                ("Authorization", "Bearer test"),
+                ("grpc.max_send_message_length", -1),
+                ("grpc.max_receive_message_length", -1),
+            ],
         )