feat: add context manager support in client (#470)

- [ ] Regenerate this pull request now.

chore: fix docstring for first attribute of protos

committer: @busunkim96
PiperOrigin-RevId: 401271153

Source-Link: googleapis/googleapis@787f8c9

Source-Link: googleapis/googleapis-gen@81decff
Copy-Tag: eyJwIjoiLmdpdGh1Yi8uT3dsQm90LnlhbWwiLCJoIjoiODFkZWNmZmU5ZmM3MjM5NmE4MTUzZTc1NmQxZDY3YTZlZWNmZDYyMCJ9

google/cloud/firestore_admin_v1/services/firestore_admin/async_client.py

@@ -998,6 +998,12 @@ async def import_documents(
         # Done; return the response.
         return response
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.transport.close()
+
 
 try:
     DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

google/cloud/firestore_admin_v1/services/firestore_admin/client.py

@@ -399,10 +399,7 @@ def __init__(
                 client_cert_source_for_mtls=client_cert_source_func,
                 quota_project_id=client_options.quota_project_id,
                 client_info=client_info,
-                always_use_jwt_access=(
-                    Transport == type(self).get_transport_class("grpc")
-                    or Transport == type(self).get_transport_class("grpc_asyncio")
-                ),
+                always_use_jwt_access=True,
             )
 
     def create_index(
@@ -1169,6 +1166,19 @@ def import_documents(
         # Done; return the response.
         return response
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """Releases underlying transport's resources.
+
+        .. warning::
+            ONLY use as a context manager if the transport is NOT shared
+            with other clients! Exiting the with block will CLOSE the transport
+            and may cause errors in other clients!
+        """
+        self.transport.close()
+
 
 try:
     DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

google/cloud/firestore_admin_v1/services/firestore_admin/transports/base.py

@@ -256,6 +256,15 @@ def _prep_wrapped_messages(self, client_info):
             ),
         }
 
+    def close(self):
+        """Closes resources associated with the transport.
+
+       .. warning::
+            Only call this method if the transport is NOT shared
+            with other clients - this may cause errors in other clients!
+        """
+        raise NotImplementedError()
+
     @property
     def operations_client(self) -> operations_v1.OperationsClient:
         """Return the client designed to process long-running operations."""

google/cloud/firestore_admin_v1/services/firestore_admin/transports/grpc.py

@@ -521,5 +521,8 @@ def import_documents(
             )
         return self._stubs["import_documents"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("FirestoreAdminGrpcTransport",)

google/cloud/firestore_admin_v1/services/firestore_admin/transports/grpc_asyncio.py

@@ -540,5 +540,8 @@ def import_documents(
             )
         return self._stubs["import_documents"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("FirestoreAdminGrpcAsyncIOTransport",)

google/cloud/firestore_admin_v1/types/field.py

@@ -62,6 +62,7 @@ class Field(proto.Message):
 
     class IndexConfig(proto.Message):
         r"""The index configuration for this field.
+
         Attributes:
             indexes (Sequence[google.cloud.firestore_admin_v1.types.Index]):
                 The indexes supported for this field.

google/cloud/firestore_admin_v1/types/location.py

@@ -24,7 +24,8 @@
 class LocationMetadata(proto.Message):
     r"""The metadata message for
     [google.cloud.location.Location.metadata][google.cloud.location.Location.metadata].
-        """
+
+    """
 
 
 __all__ = tuple(sorted(__protobuf__.manifest))

google/cloud/firestore_admin_v1/types/operation.py

@@ -109,6 +109,7 @@ class FieldOperationMetadata(proto.Message):
 
     class IndexConfigDelta(proto.Message):
         r"""Information about an index configuration change.
+
         Attributes:
             change_type (google.cloud.firestore_admin_v1.types.FieldOperationMetadata.IndexConfigDelta.ChangeType):
                 Specifies how the index is changing.

google/cloud/firestore_bundle/types/bundle.py

@@ -34,6 +34,7 @@
 
 class BundledQuery(proto.Message):
     r"""Encodes a query saved in the bundle.
+
     Attributes:
         parent (str):
             The parent resource name.
@@ -84,6 +85,7 @@ class NamedQuery(proto.Message):
 
 class BundledDocumentMetadata(proto.Message):
     r"""Metadata describing a Firestore document saved in the bundle.
+
     Attributes:
         name (str):
             The document key of a bundled document.
@@ -105,6 +107,7 @@ class BundledDocumentMetadata(proto.Message):
 
 class BundleMetadata(proto.Message):
     r"""Metadata describing the bundle file/stream.
+
     Attributes:
         id (str):
             The ID of the bundle.

google/cloud/firestore_v1/services/firestore/async_client.py

@@ -1281,6 +1281,12 @@ async def create_document(
         # Done; return the response.
         return response
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.transport.close()
+
 
 try:
     DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

google/cloud/firestore_v1/services/firestore/client.py

@@ -341,10 +341,7 @@ def __init__(
                 client_cert_source_for_mtls=client_cert_source_func,
                 quota_project_id=client_options.quota_project_id,
                 client_info=client_info,
-                always_use_jwt_access=(
-                    Transport == type(self).get_transport_class("grpc")
-                    or Transport == type(self).get_transport_class("grpc_asyncio")
-                ),
+                always_use_jwt_access=True,
             )
 
     def get_document(
@@ -1289,6 +1286,19 @@ def create_document(
         # Done; return the response.
         return response
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """Releases underlying transport's resources.
+
+        .. warning::
+            ONLY use as a context manager if the transport is NOT shared
+            with other clients! Exiting the with block will CLOSE the transport
+            and may cause errors in other clients!
+        """
+        self.transport.close()
+
 
 try:
     DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

google/cloud/firestore_v1/services/firestore/transports/base.py

@@ -394,6 +394,15 @@ def _prep_wrapped_messages(self, client_info):
             ),
         }
 
+    def close(self):
+        """Closes resources associated with the transport.
+
+       .. warning::
+            Only call this method if the transport is NOT shared
+            with other clients - this may cause errors in other clients!
+        """
+        raise NotImplementedError()
+
     @property
     def get_document(
         self,

google/cloud/firestore_v1/services/firestore/transports/grpc.py

@@ -641,5 +641,8 @@ def create_document(
             )
         return self._stubs["create_document"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("FirestoreGrpcTransport",)

google/cloud/firestore_v1/services/firestore/transports/grpc_asyncio.py

@@ -661,5 +661,8 @@ def create_document(
             )
         return self._stubs["create_document"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("FirestoreGrpcAsyncIOTransport",)

google/cloud/firestore_v1/types/common.py

@@ -65,6 +65,7 @@ class Precondition(proto.Message):
 
 class TransactionOptions(proto.Message):
     r"""Options for creating a new transaction.
+
     Attributes:
         read_only (google.cloud.firestore_v1.types.TransactionOptions.ReadOnly):
             The transaction can only be used for read

google/cloud/firestore_v1/types/document.py

@@ -85,6 +85,7 @@ class Document(proto.Message):
 
 class Value(proto.Message):
     r"""A message that can hold any of the supported value types.
+
     Attributes:
         null_value (google.protobuf.struct_pb2.NullValue):
             A null value.
@@ -149,6 +150,7 @@ class Value(proto.Message):
 
 class ArrayValue(proto.Message):
     r"""An array value.
+
     Attributes:
         values (Sequence[google.cloud.firestore_v1.types.Value]):
             Values in the array.
@@ -159,6 +161,7 @@ class ArrayValue(proto.Message):
 
 class MapValue(proto.Message):
     r"""A map value.
+
     Attributes:
         fields (Sequence[google.cloud.firestore_v1.types.MapValue.FieldsEntry]):
             The map's fields.

google/cloud/firestore_v1/types/firestore.py

@@ -782,6 +782,7 @@ class ListenResponse(proto.Message):
 
 class Target(proto.Message):
     r"""A specification of a set of documents to listen to.
+
     Attributes:
         query (google.cloud.firestore_v1.types.Target.QueryTarget):
             A target specified by a query.
@@ -811,6 +812,7 @@ class Target(proto.Message):
 
     class DocumentsTarget(proto.Message):
         r"""A target specified by a set of documents names.
+
         Attributes:
             documents (Sequence[str]):
                 The names of the documents to retrieve. In the format:
@@ -824,6 +826,7 @@ class DocumentsTarget(proto.Message):
 
     class QueryTarget(proto.Message):
         r"""A target specified by a query.
+
         Attributes:
             parent (str):
                 The parent resource name. In the format:
@@ -861,6 +864,7 @@ class QueryTarget(proto.Message):
 
 class TargetChange(proto.Message):
     r"""Targets being watched have changed.
+
     Attributes:
         target_change_type (google.cloud.firestore_v1.types.TargetChange.TargetChangeType):
             The type of change that occurred.

google/cloud/firestore_v1/types/query.py

@@ -26,6 +26,7 @@
 
 class StructuredQuery(proto.Message):
     r"""A Firestore query.
+
     Attributes:
         select (google.cloud.firestore_v1.types.StructuredQuery.Projection):
             The projection to return.
@@ -77,6 +78,7 @@ class Direction(proto.Enum):
 
     class CollectionSelector(proto.Message):
         r"""A selection of a collection, such as ``messages as m1``.
+
         Attributes:
             collection_id (str):
                 The collection ID.
@@ -93,6 +95,7 @@ class CollectionSelector(proto.Message):
 
     class Filter(proto.Message):
         r"""A filter.
+
         Attributes:
             composite_filter (google.cloud.firestore_v1.types.StructuredQuery.CompositeFilter):
                 A composite filter.
@@ -147,6 +150,7 @@ class Operator(proto.Enum):
 
     class FieldFilter(proto.Message):
         r"""A filter on a specific field.
+
         Attributes:
             field (google.cloud.firestore_v1.types.StructuredQuery.FieldReference):
                 The field to filter by.
@@ -180,6 +184,7 @@ class Operator(proto.Enum):
 
     class UnaryFilter(proto.Message):
         r"""A filter with a single operand.
+
         Attributes:
             op (google.cloud.firestore_v1.types.StructuredQuery.UnaryFilter.Operator):
                 The unary operator to apply.
@@ -207,6 +212,7 @@ class Operator(proto.Enum):
 
     class Order(proto.Message):
         r"""An order on a field.
+
         Attributes:
             field (google.cloud.firestore_v1.types.StructuredQuery.FieldReference):
                 The field to order by.
@@ -221,6 +227,7 @@ class Order(proto.Message):
 
     class FieldReference(proto.Message):
         r"""A reference to a field, such as ``max(messages.time) as max_time``.
+
         Attributes:
             field_path (str):
 
@@ -230,6 +237,7 @@ class FieldReference(proto.Message):
 
     class Projection(proto.Message):
         r"""The projection of document's fields to return.
+
         Attributes:
             fields (Sequence[google.cloud.firestore_v1.types.StructuredQuery.FieldReference]):
                 The fields to return.
@@ -254,6 +262,7 @@ class Projection(proto.Message):
 
 class Cursor(proto.Message):
     r"""A position in a query result set.
+
     Attributes:
         values (Sequence[google.cloud.firestore_v1.types.Value]):
             The values that represent a position, in the

google/cloud/firestore_v1/types/write.py

@@ -36,6 +36,7 @@
 
 class Write(proto.Message):
     r"""A write on a document.
+
     Attributes:
         update (google.cloud.firestore_v1.types.Document):
             A document to write.
@@ -86,6 +87,7 @@ class Write(proto.Message):
 
 class DocumentTransform(proto.Message):
     r"""A transformation of a document.
+
     Attributes:
         document (str):
             The name of the document to transform.
@@ -97,6 +99,7 @@ class DocumentTransform(proto.Message):
 
     class FieldTransform(proto.Message):
         r"""A transformation of a field of the document.
+
         Attributes:
             field_path (str):
                 The path of the field. See
@@ -221,6 +224,7 @@ class ServerValue(proto.Enum):
 
 class WriteResult(proto.Message):
     r"""The result of applying a write.
+
     Attributes:
         update_time (google.protobuf.timestamp_pb2.Timestamp):
             The last update time of the document after applying the
@@ -334,6 +338,7 @@ class DocumentRemove(proto.Message):
 
 class ExistenceFilter(proto.Message):
     r"""A digest of all the documents that match a given target.
+
     Attributes:
         target_id (int):
             The target ID to which this filter applies.