diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/base_client.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/base_client.py index 4d6892623a1c..95166266bef2 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/base_client.py +++ b/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) diff --git a/packages/google-cloud-firestore/tests/system/test_system.py b/packages/google-cloud-firestore/tests/system/test_system.py index cd16279fa988..68a57829b243 100644 --- a/packages/google-cloud-firestore/tests/system/test_system.py +++ b/packages/google-cloud-firestore/tests/system/test_system.py @@ -3836,3 +3836,37 @@ def in_transaction(transaction, rollback): assert len(result) == 1 assert len(result[0]) == 1 assert result[0][0].value == expected + + +@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 large 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" * (900 * 1024) + doc_ref.set({"payload": large_payload}) + + snapshot = doc_ref.get() + assert snapshot.exists + assert snapshot.to_dict() == {"payload": large_payload} + + +@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 large 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" * (900 * 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}] diff --git a/packages/google-cloud-firestore/tests/system/test_system_async.py b/packages/google-cloud-firestore/tests/system/test_system_async.py index 1003f2a5a015..4723b06f3bcc 100644 --- a/packages/google-cloud-firestore/tests/system/test_system_async.py +++ b/packages/google-cloud-firestore/tests/system/test_system_async.py @@ -3708,3 +3708,39 @@ 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.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 large 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" * (900 * 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.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 large 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" * (900 * 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}] diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_base_client.py b/packages/google-cloud-firestore/tests/unit/v1/test_base_client.py index 856c771a195b..1d3938ac809e 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_base_client.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_base_client.py @@ -110,9 +110,13 @@ 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, @@ -236,7 +240,12 @@ 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), + ], )