diff --git a/.github/.OwlBot.lock.yaml b/.github/.OwlBot.lock.yaml index 8b90899d21..597e0c3261 100644 --- a/.github/.OwlBot.lock.yaml +++ b/.github/.OwlBot.lock.yaml @@ -13,5 +13,5 @@ # limitations under the License. docker: image: gcr.io/cloud-devrel-public-resources/owlbot-python:latest - digest: sha256:2dc6f67639bee669c33c6277a624ab9857d363e2fd33ac5b02d417b7d25f1ffc -# created: 2024-08-15T17:41:26.438340772Z + digest: sha256:e8dcfd7cbfd8beac3a3ff8d3f3185287ea0625d859168cc80faccfc9a7a00455 +# created: 2024-09-16T21:04:09.091105552Z diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index f4a337c496..dd8bd76922 100644 --- a/.github/workflows/unittest.yml +++ b/.github/workflows/unittest.yml @@ -30,6 +30,7 @@ jobs: with: name: coverage-artifact-${{ matrix.python }} path: .coverage-${{ matrix.python }} + include-hidden-files: true cover: runs-on: ubuntu-latest diff --git a/.kokoro/release.sh b/.kokoro/release.sh index 0be2271b27..85315bb58e 100755 --- a/.kokoro/release.sh +++ b/.kokoro/release.sh @@ -23,7 +23,7 @@ python3 -m releasetool publish-reporter-script > /tmp/publisher-script; source / export PYTHONUNBUFFERED=1 # Move into the package, build the distribution and upload. -TWINE_PASSWORD=$(cat "${KOKORO_KEYSTORE_DIR}/73713_google-cloud-pypi-token-keystore-1") +TWINE_PASSWORD=$(cat "${KOKORO_KEYSTORE_DIR}/73713_google-cloud-pypi-token-keystore-2") cd github/python-firestore python3 setup.py sdist bdist_wheel twine upload --username __token__ --password "${TWINE_PASSWORD}" dist/* diff --git a/.kokoro/release/common.cfg b/.kokoro/release/common.cfg index 46d49fdc69..8f9b40e16f 100644 --- a/.kokoro/release/common.cfg +++ b/.kokoro/release/common.cfg @@ -28,7 +28,7 @@ before_action { fetch_keystore { keystore_resource { keystore_config_id: 73713 - keyname: "google-cloud-pypi-token-keystore-1" + keyname: "google-cloud-pypi-token-keystore-2" } } } diff --git a/.release-please-manifest.json b/.release-please-manifest.json index a627e662e0..b7f666a684 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -1,3 +1,3 @@ { - ".": "2.18.0" + ".": "2.19.0" } \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 786b1399b7..d8b96a9383 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,15 @@ [1]: https://pypi.org/project/google-cloud-firestore/#history +## [2.19.0](https://github.com/googleapis/python-firestore/compare/v2.18.0...v2.19.0) (2024-09-20) + + +### Features + +* Add Database.SourceInfo and Database.source_info (information about database provenance, specifically for restored databases) ([#963](https://github.com/googleapis/python-firestore/issues/963)) ([4e15714](https://github.com/googleapis/python-firestore/commit/4e15714cd70b0577d1450b081ad26a8678fe1a9e)) +* Query profiling part 1: synchronous ([#938](https://github.com/googleapis/python-firestore/issues/938)) ([1614b3f](https://github.com/googleapis/python-firestore/commit/1614b3f15311f9eee39c8b72b8dc81f259498dcb)) +* Query profiling part 2: asynchronous ([#961](https://github.com/googleapis/python-firestore/issues/961)) ([060a3ef](https://github.com/googleapis/python-firestore/commit/060a3efa7df4eb6b4ef0701a246ff630dde432c7)) + ## [2.18.0](https://github.com/googleapis/python-firestore/compare/v2.17.2...v2.18.0) (2024-08-26) diff --git a/google/cloud/firestore/__init__.py b/google/cloud/firestore/__init__.py index 79095778db..314a138cbc 100644 --- a/google/cloud/firestore/__init__.py +++ b/google/cloud/firestore/__init__.py @@ -38,6 +38,7 @@ from google.cloud.firestore_v1 import DocumentSnapshot from google.cloud.firestore_v1 import DocumentTransform from google.cloud.firestore_v1 import ExistsOption +from google.cloud.firestore_v1 import ExplainOptions from google.cloud.firestore_v1 import FieldFilter from google.cloud.firestore_v1 import GeoPoint from google.cloud.firestore_v1 import Increment @@ -78,6 +79,7 @@ "DocumentSnapshot", "DocumentTransform", "ExistsOption", + "ExplainOptions", "FieldFilter", "GeoPoint", "Increment", diff --git a/google/cloud/firestore/gapic_version.py b/google/cloud/firestore/gapic_version.py index f09943f6bd..0f1a446f38 100644 --- a/google/cloud/firestore/gapic_version.py +++ b/google/cloud/firestore/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "2.18.0" # {x-release-please-version} +__version__ = "2.19.0" # {x-release-please-version} diff --git a/google/cloud/firestore_admin_v1/gapic_version.py b/google/cloud/firestore_admin_v1/gapic_version.py index f09943f6bd..0f1a446f38 100644 --- a/google/cloud/firestore_admin_v1/gapic_version.py +++ b/google/cloud/firestore_admin_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "2.18.0" # {x-release-please-version} +__version__ = "2.19.0" # {x-release-please-version} diff --git a/google/cloud/firestore_admin_v1/services/firestore_admin/async_client.py b/google/cloud/firestore_admin_v1/services/firestore_admin/async_client.py index 34a30d378c..db6037da34 100644 --- a/google/cloud/firestore_admin_v1/services/firestore_admin/async_client.py +++ b/google/cloud/firestore_admin_v1/services/firestore_admin/async_client.py @@ -128,6 +128,8 @@ class FirestoreAdminAsyncClient: parse_index_path = staticmethod(FirestoreAdminClient.parse_index_path) location_path = staticmethod(FirestoreAdminClient.location_path) parse_location_path = staticmethod(FirestoreAdminClient.parse_location_path) + operation_path = staticmethod(FirestoreAdminClient.operation_path) + parse_operation_path = staticmethod(FirestoreAdminClient.parse_operation_path) common_billing_account_path = staticmethod( FirestoreAdminClient.common_billing_account_path ) @@ -834,7 +836,7 @@ async def sample_get_field(): database. Fields are grouped by their "Collection Group", which represent all collections - in the database with the same id. + in the database with the same ID. """ # Create or coerce a protobuf request object. @@ -967,7 +969,7 @@ async def sample_update_field(): Fields are grouped by their "Collection Group", which represent all collections in the database with the - same id. + same ID. """ # Create or coerce a protobuf request object. @@ -1632,7 +1634,7 @@ async def sample_create_database(): last a letter or a number. Must not be UUID-like /[0-9a-f]{8}(-[0-9a-f]{4}){3}-[0-9a-f]{12}/. - "(default)" database id is also valid. + "(default)" database ID is also valid. This corresponds to the ``database_id`` field on the ``request`` instance; if ``request`` is provided, this @@ -2548,7 +2550,7 @@ async def sample_restore_database(): Args: request (Optional[Union[google.cloud.firestore_admin_v1.types.RestoreDatabaseRequest, dict]]): The request object. The request message for - [FirestoreAdmin.RestoreDatabase][google.firestore.admin.v1.RestoreDatabase]. + [FirestoreAdmin.RestoreDatabase][google.firestore.admin.v1.FirestoreAdmin.RestoreDatabase]. retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. diff --git a/google/cloud/firestore_admin_v1/services/firestore_admin/client.py b/google/cloud/firestore_admin_v1/services/firestore_admin/client.py index b7bcfd80a5..2b4fa5890c 100644 --- a/google/cloud/firestore_admin_v1/services/firestore_admin/client.py +++ b/google/cloud/firestore_admin_v1/services/firestore_admin/client.py @@ -378,6 +378,28 @@ def parse_location_path(path: str) -> Dict[str, str]: m = re.match(r"^projects/(?P.+?)/locations/(?P.+?)$", path) return m.groupdict() if m else {} + @staticmethod + def operation_path( + project: str, + database: str, + operation: str, + ) -> str: + """Returns a fully-qualified operation string.""" + return "projects/{project}/databases/{database}/operations/{operation}".format( + project=project, + database=database, + operation=operation, + ) + + @staticmethod + def parse_operation_path(path: str) -> Dict[str, str]: + """Parses a operation path into its component segments.""" + m = re.match( + r"^projects/(?P.+?)/databases/(?P.+?)/operations/(?P.+?)$", + path, + ) + return m.groupdict() if m else {} + @staticmethod def common_billing_account_path( billing_account: str, @@ -1354,7 +1376,7 @@ def sample_get_field(): database. Fields are grouped by their "Collection Group", which represent all collections - in the database with the same id. + in the database with the same ID. """ # Create or coerce a protobuf request object. @@ -1484,7 +1506,7 @@ def sample_update_field(): Fields are grouped by their "Collection Group", which represent all collections in the database with the - same id. + same ID. """ # Create or coerce a protobuf request object. @@ -2134,7 +2156,7 @@ def sample_create_database(): last a letter or a number. Must not be UUID-like /[0-9a-f]{8}(-[0-9a-f]{4}){3}-[0-9a-f]{12}/. - "(default)" database id is also valid. + "(default)" database ID is also valid. This corresponds to the ``database_id`` field on the ``request`` instance; if ``request`` is provided, this @@ -3026,7 +3048,7 @@ def sample_restore_database(): Args: request (Union[google.cloud.firestore_admin_v1.types.RestoreDatabaseRequest, dict]): The request object. The request message for - [FirestoreAdmin.RestoreDatabase][google.firestore.admin.v1.RestoreDatabase]. + [FirestoreAdmin.RestoreDatabase][google.firestore.admin.v1.FirestoreAdmin.RestoreDatabase]. retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. diff --git a/google/cloud/firestore_admin_v1/services/firestore_admin/transports/rest.py b/google/cloud/firestore_admin_v1/services/firestore_admin/transports/rest.py index 0003a5c13a..127f42b2a1 100644 --- a/google/cloud/firestore_admin_v1/services/firestore_admin/transports/rest.py +++ b/google/cloud/firestore_admin_v1/services/firestore_admin/transports/rest.py @@ -2157,7 +2157,7 @@ def __call__( database. Fields are grouped by their "Collection Group", which represent all collections - in the database with the same id. + in the database with the same ID. """ @@ -2861,7 +2861,7 @@ def __call__( Args: request (~.firestore_admin.RestoreDatabaseRequest): The request object. The request message for - [FirestoreAdmin.RestoreDatabase][google.firestore.admin.v1.RestoreDatabase]. + [FirestoreAdmin.RestoreDatabase][google.firestore.admin.v1.FirestoreAdmin.RestoreDatabase]. retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. diff --git a/google/cloud/firestore_admin_v1/types/database.py b/google/cloud/firestore_admin_v1/types/database.py index 58e0e20985..32901f729f 100644 --- a/google/cloud/firestore_admin_v1/types/database.py +++ b/google/cloud/firestore_admin_v1/types/database.py @@ -50,6 +50,10 @@ class Database(proto.Message): database was most recently updated. Note this only includes updates to the database resource and not data contained by the database. + delete_time (google.protobuf.timestamp_pb2.Timestamp): + Output only. The timestamp at which this + database was deleted. Only set if the database + has been deleted. location_id (str): The location of the database. Available locations are listed at @@ -93,8 +97,8 @@ class Database(proto.Message): this database. key_prefix (str): Output only. The key_prefix for this database. This - key_prefix is used, in combination with the project id ("~") - to construct the application id that is returned from the + key_prefix is used, in combination with the project ID ("~") + to construct the application ID that is returned from the Cloud Datastore APIs in Google App Engine first generation runtimes. @@ -103,6 +107,16 @@ class Database(proto.Message): v~foo). delete_protection_state (google.cloud.firestore_admin_v1.types.Database.DeleteProtectionState): State of delete protection for the database. + cmek_config (google.cloud.firestore_admin_v1.types.Database.CmekConfig): + Optional. Presence indicates CMEK is enabled + for this database. + previous_id (str): + Output only. The database resource's prior + database ID. This field is only populated for + deleted databases. + source_info (google.cloud.firestore_admin_v1.types.Database.SourceInfo): + Output only. Information about the provenance + of this database. etag (str): This checksum is computed by the server based on the value of other fields, and may be sent on @@ -120,8 +134,7 @@ class DatabaseType(proto.Enum): Values: DATABASE_TYPE_UNSPECIFIED (0): - The default value. This value is used if the - database type is omitted. + Not used. FIRESTORE_NATIVE (1): Firestore Native Mode DATASTORE_MODE (2): @@ -225,6 +238,173 @@ class DeleteProtectionState(proto.Enum): DELETE_PROTECTION_DISABLED = 1 DELETE_PROTECTION_ENABLED = 2 + class CmekConfig(proto.Message): + r"""The CMEK (Customer Managed Encryption Key) configuration for + a Firestore database. If not present, the database is secured by + the default Google encryption key. + + Attributes: + kms_key_name (str): + Required. Only keys in the same location as this database + are allowed to be used for encryption. + + For Firestore's nam5 multi-region, this corresponds to Cloud + KMS multi-region us. For Firestore's eur3 multi-region, this + corresponds to Cloud KMS multi-region europe. See + https://cloud.google.com/kms/docs/locations. + + The expected format is + ``projects/{project_id}/locations/{kms_location}/keyRings/{key_ring}/cryptoKeys/{crypto_key}``. + active_key_version (MutableSequence[str]): + Output only. Currently in-use `KMS key + versions `__. + During `key + rotation `__, + there can be multiple in-use key versions. + + The expected format is + ``projects/{project_id}/locations/{kms_location}/keyRings/{key_ring}/cryptoKeys/{crypto_key}/cryptoKeyVersions/{key_version}``. + """ + + kms_key_name: str = proto.Field( + proto.STRING, + number=1, + ) + active_key_version: MutableSequence[str] = proto.RepeatedField( + proto.STRING, + number=2, + ) + + class SourceInfo(proto.Message): + r"""Information about the provenance of this database. + + .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields + + Attributes: + backup (google.cloud.firestore_admin_v1.types.Database.SourceInfo.BackupSource): + If set, this database was restored from the + specified backup (or a snapshot thereof). + + This field is a member of `oneof`_ ``source``. + operation (str): + The associated long-running operation. This field may not be + set after the operation has completed. Format: + ``projects/{project}/databases/{database}/operations/{operation}``. + """ + + class BackupSource(proto.Message): + r"""Information about a backup that was used to restore a + database. + + Attributes: + backup (str): + The resource name of the backup that was used to restore + this database. Format: + ``projects/{project}/locations/{location}/backups/{backup}``. + """ + + backup: str = proto.Field( + proto.STRING, + number=1, + ) + + backup: "Database.SourceInfo.BackupSource" = proto.Field( + proto.MESSAGE, + number=1, + oneof="source", + message="Database.SourceInfo.BackupSource", + ) + operation: str = proto.Field( + proto.STRING, + number=3, + ) + + class EncryptionConfig(proto.Message): + r"""Encryption configuration for a new database being created from + another source. + + The source could be a [Backup][google.firestore.admin.v1.Backup] . + + This message has `oneof`_ fields (mutually exclusive fields). + For each oneof, at most one member field can be set at the same time. + Setting any member of the oneof automatically clears all other + members. + + .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields + + Attributes: + google_default_encryption (google.cloud.firestore_admin_v1.types.Database.EncryptionConfig.GoogleDefaultEncryptionOptions): + Use Google default encryption. + + This field is a member of `oneof`_ ``encryption_type``. + use_source_encryption (google.cloud.firestore_admin_v1.types.Database.EncryptionConfig.SourceEncryptionOptions): + The database will use the same encryption + configuration as the source. + + This field is a member of `oneof`_ ``encryption_type``. + customer_managed_encryption (google.cloud.firestore_admin_v1.types.Database.EncryptionConfig.CustomerManagedEncryptionOptions): + Use Customer Managed Encryption Keys (CMEK) + for encryption. + + This field is a member of `oneof`_ ``encryption_type``. + """ + + class GoogleDefaultEncryptionOptions(proto.Message): + r"""The configuration options for using Google default + encryption. + + """ + + class SourceEncryptionOptions(proto.Message): + r"""The configuration options for using the same encryption + method as the source. + + """ + + class CustomerManagedEncryptionOptions(proto.Message): + r"""The configuration options for using CMEK (Customer Managed + Encryption Key) encryption. + + Attributes: + kms_key_name (str): + Required. Only keys in the same location as the database are + allowed to be used for encryption. + + For Firestore's nam5 multi-region, this corresponds to Cloud + KMS multi-region us. For Firestore's eur3 multi-region, this + corresponds to Cloud KMS multi-region europe. See + https://cloud.google.com/kms/docs/locations. + + The expected format is + ``projects/{project_id}/locations/{kms_location}/keyRings/{key_ring}/cryptoKeys/{crypto_key}``. + """ + + kms_key_name: str = proto.Field( + proto.STRING, + number=1, + ) + + google_default_encryption: "Database.EncryptionConfig.GoogleDefaultEncryptionOptions" = proto.Field( + proto.MESSAGE, + number=1, + oneof="encryption_type", + message="Database.EncryptionConfig.GoogleDefaultEncryptionOptions", + ) + use_source_encryption: "Database.EncryptionConfig.SourceEncryptionOptions" = ( + proto.Field( + proto.MESSAGE, + number=2, + oneof="encryption_type", + message="Database.EncryptionConfig.SourceEncryptionOptions", + ) + ) + customer_managed_encryption: "Database.EncryptionConfig.CustomerManagedEncryptionOptions" = proto.Field( + proto.MESSAGE, + number=3, + oneof="encryption_type", + message="Database.EncryptionConfig.CustomerManagedEncryptionOptions", + ) + name: str = proto.Field( proto.STRING, number=1, @@ -243,6 +423,11 @@ class DeleteProtectionState(proto.Enum): number=6, message=timestamp_pb2.Timestamp, ) + delete_time: timestamp_pb2.Timestamp = proto.Field( + proto.MESSAGE, + number=7, + message=timestamp_pb2.Timestamp, + ) location_id: str = proto.Field( proto.STRING, number=9, @@ -286,6 +471,20 @@ class DeleteProtectionState(proto.Enum): number=22, enum=DeleteProtectionState, ) + cmek_config: CmekConfig = proto.Field( + proto.MESSAGE, + number=23, + message=CmekConfig, + ) + previous_id: str = proto.Field( + proto.STRING, + number=25, + ) + source_info: SourceInfo = proto.Field( + proto.MESSAGE, + number=26, + message=SourceInfo, + ) etag: str = proto.Field( proto.STRING, number=99, diff --git a/google/cloud/firestore_admin_v1/types/field.py b/google/cloud/firestore_admin_v1/types/field.py index 31be5fc17a..f878b63313 100644 --- a/google/cloud/firestore_admin_v1/types/field.py +++ b/google/cloud/firestore_admin_v1/types/field.py @@ -34,7 +34,7 @@ class Field(proto.Message): r"""Represents a single field in the database. Fields are grouped by their "Collection Group", which represent - all collections in the database with the same id. + all collections in the database with the same ID. Attributes: name (str): diff --git a/google/cloud/firestore_admin_v1/types/firestore_admin.py b/google/cloud/firestore_admin_v1/types/firestore_admin.py index 17c25a854a..20a105bd61 100644 --- a/google/cloud/firestore_admin_v1/types/firestore_admin.py +++ b/google/cloud/firestore_admin_v1/types/firestore_admin.py @@ -109,7 +109,7 @@ class CreateDatabaseRequest(proto.Message): letter or a number. Must not be UUID-like /[0-9a-f]{8}(-[0-9a-f]{4}){3}-[0-9a-f]{12}/. - "(default)" database id is also valid. + "(default)" database ID is also valid. """ parent: str = proto.Field( @@ -588,8 +588,8 @@ class ExportDocumentsRequest(proto.Message): Required. Database to export. Should be of the form: ``projects/{project_id}/databases/{database_id}``. collection_ids (MutableSequence[str]): - Which collection ids to export. Unspecified - means all collections. Each collection id in + Which collection IDs to export. Unspecified + means all collections. Each collection ID in this list must be unique. output_uri_prefix (str): The output URI. Currently only supports Google Cloud Storage @@ -654,9 +654,9 @@ class ImportDocumentsRequest(proto.Message): Required. Database to import into. Should be of the form: ``projects/{project_id}/databases/{database_id}``. collection_ids (MutableSequence[str]): - Which collection ids to import. Unspecified + Which collection IDs to import. Unspecified means all collections included in the import. - Each collection id in this list must be unique. + Each collection ID in this list must be unique. input_uri_prefix (str): Location of the exported files. This must match the output_uri_prefix of an ExportDocumentsResponse from an @@ -837,7 +837,7 @@ class DeleteBackupRequest(proto.Message): class RestoreDatabaseRequest(proto.Message): r"""The request message for - [FirestoreAdmin.RestoreDatabase][google.firestore.admin.v1.RestoreDatabase]. + [FirestoreAdmin.RestoreDatabase][google.firestore.admin.v1.FirestoreAdmin.RestoreDatabase]. Attributes: parent (str): @@ -846,7 +846,7 @@ class RestoreDatabaseRequest(proto.Message): database_id (str): Required. The ID to use for the database, which will become the final component of the database's resource name. This - database id must not be associated with an existing + database ID must not be associated with an existing database. This value should be 4-63 characters. Valid characters are @@ -854,13 +854,23 @@ class RestoreDatabaseRequest(proto.Message): letter or a number. Must not be UUID-like /[0-9a-f]{8}(-[0-9a-f]{4}){3}-[0-9a-f]{12}/. - "(default)" database id is also valid. + "(default)" database ID is also valid. backup (str): Required. Backup to restore from. Must be from the same project as the parent. + The restored database will be created in the same location + as the source backup. + Format is: ``projects/{project_id}/locations/{location}/backups/{backup}`` + encryption_config (google.cloud.firestore_admin_v1.types.Database.EncryptionConfig): + Optional. Encryption configuration for the restored + database. + + If this field is not specified, the restored database will + use the same encryption configuration as the backup, namely + [use_source_encryption][google.firestore.admin.v1.Database.EncryptionConfig.use_source_encryption]. """ parent: str = proto.Field( @@ -875,6 +885,11 @@ class RestoreDatabaseRequest(proto.Message): proto.STRING, number=3, ) + encryption_config: gfa_database.Database.EncryptionConfig = proto.Field( + proto.MESSAGE, + number=9, + message=gfa_database.Database.EncryptionConfig, + ) __all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/google/cloud/firestore_admin_v1/types/index.py b/google/cloud/firestore_admin_v1/types/index.py index b9739d429e..716213fd22 100644 --- a/google/cloud/firestore_admin_v1/types/index.py +++ b/google/cloud/firestore_admin_v1/types/index.py @@ -43,12 +43,12 @@ class Index(proto.Message): specified allow queries against a collection that is the child of a specific document, specified at query time, and that has the same - collection id. + collection ID. Indexes with a collection group query scope specified allow queries against all collections descended from a specific document, specified at - query time, and that have the same collection id + query time, and that have the same collection ID as this index. api_scope (google.cloud.firestore_admin_v1.types.Index.ApiScope): The API scope supported by this index. @@ -84,11 +84,11 @@ class QueryScope(proto.Enum): specified allow queries against a collection that is the child of a specific document, specified at query time, and that has the - collection id specified by the index. + collection ID specified by the index. COLLECTION_GROUP (2): Indexes with a collection group query scope specified allow queries against all collections - that has the collection id specified by the + that has the collection ID specified by the index. COLLECTION_RECURSIVE (3): Include all the collections's ancestor in the diff --git a/google/cloud/firestore_admin_v1/types/operation.py b/google/cloud/firestore_admin_v1/types/operation.py index bb817e9053..c3e59d10bf 100644 --- a/google/cloud/firestore_admin_v1/types/operation.py +++ b/google/cloud/firestore_admin_v1/types/operation.py @@ -291,11 +291,11 @@ class ExportDocumentsMetadata(proto.Message): progress_bytes (google.cloud.firestore_admin_v1.types.Progress): The progress, in bytes, of this operation. collection_ids (MutableSequence[str]): - Which collection ids are being exported. + Which collection IDs are being exported. output_uri_prefix (str): Where the documents are being exported to. namespace_ids (MutableSequence[str]): - Which namespace ids are being exported. + Which namespace IDs are being exported. snapshot_time (google.protobuf.timestamp_pb2.Timestamp): The timestamp that corresponds to the version of the database that is being exported. If @@ -367,11 +367,11 @@ class ImportDocumentsMetadata(proto.Message): progress_bytes (google.cloud.firestore_admin_v1.types.Progress): The progress, in bytes, of this operation. collection_ids (MutableSequence[str]): - Which collection ids are being imported. + Which collection IDs are being imported. input_uri_prefix (str): The location of the documents being imported. namespace_ids (MutableSequence[str]): - Which namespace ids are being imported. + Which namespace IDs are being imported. """ start_time: timestamp_pb2.Timestamp = proto.Field( @@ -433,10 +433,10 @@ class BulkDeleteDocumentsMetadata(proto.Message): progress_bytes (google.cloud.firestore_admin_v1.types.Progress): The progress, in bytes, of this operation. collection_ids (MutableSequence[str]): - The ids of the collection groups that are + The IDs of the collection groups that are being deleted. namespace_ids (MutableSequence[str]): - Which namespace ids are being deleted. + Which namespace IDs are being deleted. snapshot_time (google.protobuf.timestamp_pb2.Timestamp): The timestamp that corresponds to the version of the database that is being read to get the diff --git a/google/cloud/firestore_admin_v1/types/schedule.py b/google/cloud/firestore_admin_v1/types/schedule.py index 3e6d0dfbad..eb7b138999 100644 --- a/google/cloud/firestore_admin_v1/types/schedule.py +++ b/google/cloud/firestore_admin_v1/types/schedule.py @@ -71,6 +71,9 @@ class BackupSchedule(proto.Message): At what relative time in the future, compared to its creation time, the backup should be deleted, e.g. keep backups for 7 days. + + The maximum supported retention period is 14 + weeks. daily_recurrence (google.cloud.firestore_admin_v1.types.DailyRecurrence): For a schedule that runs daily. @@ -116,8 +119,8 @@ class BackupSchedule(proto.Message): class DailyRecurrence(proto.Message): - r"""Represents a recurring schedule that runs at a specific time - every day. + r"""Represents a recurring schedule that runs every day. + The time zone is UTC. """ diff --git a/google/cloud/firestore_bundle/gapic_version.py b/google/cloud/firestore_bundle/gapic_version.py index f09943f6bd..0f1a446f38 100644 --- a/google/cloud/firestore_bundle/gapic_version.py +++ b/google/cloud/firestore_bundle/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "2.18.0" # {x-release-please-version} +__version__ = "2.19.0" # {x-release-please-version} diff --git a/google/cloud/firestore_v1/__init__.py b/google/cloud/firestore_v1/__init__.py index 1aff5ec740..049eb4183f 100644 --- a/google/cloud/firestore_v1/__init__.py +++ b/google/cloud/firestore_v1/__init__.py @@ -50,6 +50,7 @@ from google.cloud.firestore_v1.collection import CollectionReference from google.cloud.firestore_v1.document import DocumentReference from google.cloud.firestore_v1.query import CollectionGroup, Query +from google.cloud.firestore_v1.query_profile import ExplainOptions from google.cloud.firestore_v1.transaction import Transaction, transactional from google.cloud.firestore_v1.transforms import ( DELETE_FIELD, @@ -131,6 +132,7 @@ "DocumentSnapshot", "DocumentTransform", "ExistsOption", + "ExplainOptions", "FieldFilter", "GeoPoint", "Increment", diff --git a/google/cloud/firestore_v1/aggregation.py b/google/cloud/firestore_v1/aggregation.py index 65106122ab..f0e3f94baf 100644 --- a/google/cloud/firestore_v1/aggregation.py +++ b/google/cloud/firestore_v1/aggregation.py @@ -30,12 +30,14 @@ BaseAggregationQuery, _query_response_to_result, ) -from google.cloud.firestore_v1.base_document import DocumentSnapshot +from google.cloud.firestore_v1.query_results import QueryResultsList from google.cloud.firestore_v1.stream_generator import StreamGenerator # Types needed only for Type Hints -if TYPE_CHECKING: - from google.cloud.firestore_v1 import transaction # pragma: NO COVER +if TYPE_CHECKING: # pragma: NO COVER + from google.cloud.firestore_v1 import transaction + from google.cloud.firestore_v1.query_profile import ExplainMetrics + from google.cloud.firestore_v1.query_profile import ExplainOptions class AggregationQuery(BaseAggregationQuery): @@ -54,10 +56,14 @@ def get( retries.Retry, None, gapic_v1.method._MethodDefault ] = gapic_v1.method.DEFAULT, timeout: float | None = None, - ) -> List[AggregationResult]: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> QueryResultsList[AggregationResult]: """Runs the aggregation query. - This sends a ``RunAggregationQuery`` RPC and returns a list of aggregation results in the stream of ``RunAggregationQueryResponse`` messages. + This sends a ``RunAggregationQuery`` RPC and returns a list of + aggregation results in the stream of ``RunAggregationQueryResponse`` + messages. Args: transaction @@ -70,20 +76,39 @@ def get( should be retried. Defaults to a system-specified policy. timeout (float): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: - list: The aggregation query results + QueryResultsList[AggregationResult]: The aggregation query results. """ - result = self.stream(transaction=transaction, retry=retry, timeout=timeout) - return list(result) # type: ignore + explain_metrics: ExplainMetrics | None = None - def _get_stream_iterator(self, transaction, retry, timeout): + result = self.stream( + transaction=transaction, + retry=retry, + timeout=timeout, + explain_options=explain_options, + ) + result_list = list(result) + + if explain_options is None: + explain_metrics = None + else: + explain_metrics = result.get_explain_metrics() + + return QueryResultsList(result_list, explain_options, explain_metrics) + + def _get_stream_iterator(self, transaction, retry, timeout, explain_options=None): """Helper method for :meth:`stream`.""" request, kwargs = self._prep_stream( transaction, retry, timeout, + explain_options, ) return self._client._firestore_api.run_aggregation_query( @@ -106,9 +131,12 @@ def _retry_query_after_exception(self, exc, retry, transaction): def _make_stream( self, transaction: Optional[transaction.Transaction] = None, - retry: Optional[retries.Retry] = gapic_v1.method.DEFAULT, + retry: Union[ + retries.Retry, None, gapic_v1.method._MethodDefault + ] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> Union[Generator[List[AggregationResult], Any, None]]: + explain_options: Optional[ExplainOptions] = None, + ) -> Generator[List[AggregationResult], Any, Optional[ExplainMetrics]]: """Internal method for stream(). Runs the aggregation query. This sends a ``RunAggregationQuery`` RPC and then returns a generator @@ -127,16 +155,27 @@ def _make_stream( system-specified policy. timeout (Optional[float]): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Yields: - :class:`~google.cloud.firestore_v1.base_aggregation.AggregationResult`: + List[AggregationResult]: The result of aggregations of this query. + + Returns: + (Optional[google.cloud.firestore_v1.types.query_profile.ExplainMetrtics]): + The results of query profiling, if received from the service. + """ + metrics: ExplainMetrics | None = None response_iterator = self._get_stream_iterator( transaction, retry, timeout, + explain_options, ) while True: try: @@ -154,15 +193,26 @@ def _make_stream( if response is None: # EOI break + + if metrics is None and response.explain_metrics: + metrics = response.explain_metrics + result = _query_response_to_result(response) - yield result + if result: + yield result + + return metrics def stream( self, transaction: Optional["transaction.Transaction"] = None, - retry: Optional[retries.Retry] = gapic_v1.method.DEFAULT, + retry: Union[ + retries.Retry, None, gapic_v1.method._MethodDefault + ] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> "StreamGenerator[DocumentSnapshot]": + *, + explain_options: Optional[ExplainOptions] = None, + ) -> StreamGenerator[List[AggregationResult]]: """Runs the aggregation query. This sends a ``RunAggregationQuery`` RPC and then returns a generator @@ -181,13 +231,19 @@ def stream( system-specified policy. timeout (Optinal[float]): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: - `StreamGenerator[DocumentSnapshot]`: A generator of the query results. + `StreamGenerator[List[AggregationResult]]`: + A generator of the query results. """ inner_generator = self._make_stream( transaction=transaction, retry=retry, timeout=timeout, + explain_options=explain_options, ) - return StreamGenerator(inner_generator) + return StreamGenerator(inner_generator, explain_options) diff --git a/google/cloud/firestore_v1/async_aggregation.py b/google/cloud/firestore_v1/async_aggregation.py index 1c75f0cfd8..5855b71614 100644 --- a/google/cloud/firestore_v1/async_aggregation.py +++ b/google/cloud/firestore_v1/async_aggregation.py @@ -20,7 +20,7 @@ """ from __future__ import annotations -from typing import TYPE_CHECKING, AsyncGenerator, List, Optional, Union +from typing import TYPE_CHECKING, Any, AsyncGenerator, List, Optional, Union from google.api_core import gapic_v1 from google.api_core import retry_async as retries @@ -28,13 +28,15 @@ from google.cloud.firestore_v1 import transaction from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator from google.cloud.firestore_v1.base_aggregation import ( - AggregationResult, BaseAggregationQuery, _query_response_to_result, ) +from google.cloud.firestore_v1.query_results import QueryResultsList if TYPE_CHECKING: # pragma: NO COVER - from google.cloud.firestore_v1.base_document import DocumentSnapshot + from google.cloud.firestore_v1.base_aggregation import AggregationResult + from google.cloud.firestore_v1.query_profile import ExplainMetrics, ExplainOptions + import google.cloud.firestore_v1.types.query_profile as query_profile_pb class AsyncAggregationQuery(BaseAggregationQuery): @@ -53,7 +55,9 @@ async def get( retries.AsyncRetry, None, gapic_v1.method._MethodDefault ] = gapic_v1.method.DEFAULT, timeout: float | None = None, - ) -> List[AggregationResult]: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> QueryResultsList[List[AggregationResult]]: """Runs the aggregation query. This sends a ``RunAggregationQuery`` RPC and returns a list of aggregation results in the stream of ``RunAggregationQueryResponse`` messages. @@ -69,23 +73,39 @@ async def get( should be retried. Defaults to a system-specified policy. timeout (float): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: - list: The aggregation query results + QueryResultsList[List[AggregationResult]]: The aggregation query results. """ + explain_metrics: ExplainMetrics | None = None + stream_result = self.stream( - transaction=transaction, retry=retry, timeout=timeout + transaction=transaction, + retry=retry, + timeout=timeout, + explain_options=explain_options, ) result = [aggregation async for aggregation in stream_result] - return result # type: ignore + + if explain_options is None: + explain_metrics = None + else: + explain_metrics = await stream_result.get_explain_metrics() + + return QueryResultsList(result, explain_options, explain_metrics) async def _make_stream( self, transaction: Optional[transaction.Transaction] = None, retry: Optional[retries.AsyncRetry] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> Union[AsyncGenerator[List[AggregationResult], None]]: + explain_options: Optional[ExplainOptions] = None, + ) -> AsyncGenerator[List[AggregationResult] | query_profile_pb.ExplainMetrics, Any]: """Internal method for stream(). Runs the aggregation query. This sends a ``RunAggregationQuery`` RPC and then returns a generator which @@ -105,15 +125,23 @@ async def _make_stream( system-specified policy. timeout (Optional[float]): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Yields: - :class:`~google.cloud.firestore_v1.base_aggregation.AggregationResult`: - The result of aggregations of this query + List[AggregationResult] | query_profile_pb.ExplainMetrics: + The result of aggregations of this query. Query results will be + yielded as `List[AggregationResult]`. When the result contains + returned explain metrics, yield `query_profile_pb.ExplainMetrics` + individually. """ request, kwargs = self._prep_stream( transaction, retry, timeout, + explain_options, ) response_iterator = await self._client._firestore_api.run_aggregation_query( @@ -124,14 +152,21 @@ async def _make_stream( async for response in response_iterator: result = _query_response_to_result(response) - yield result + if result: + yield result + + if response.explain_metrics: + metrics = response.explain_metrics + yield metrics def stream( self, transaction: Optional[transaction.Transaction] = None, retry: Optional[retries.AsyncRetry] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> "AsyncStreamGenerator[DocumentSnapshot]": + *, + explain_options: Optional[ExplainOptions] = None, + ) -> AsyncStreamGenerator[List[AggregationResult]]: """Runs the aggregation query. This sends a ``RunAggregationQuery`` RPC and then returns a generator @@ -150,9 +185,13 @@ def stream( system-specified policy. timeout (Optional[float]): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: - `AsyncStreamGenerator[DocumentSnapshot]`: + `AsyncStreamGenerator[List[AggregationResult]]`: A generator of the query results. """ @@ -160,5 +199,6 @@ def stream( transaction=transaction, retry=retry, timeout=timeout, + explain_options=explain_options, ) - return AsyncStreamGenerator(inner_generator) + return AsyncStreamGenerator(inner_generator, explain_options) diff --git a/google/cloud/firestore_v1/async_collection.py b/google/cloud/firestore_v1/async_collection.py index 77761f2ad1..ec15de65f4 100644 --- a/google/cloud/firestore_v1/async_collection.py +++ b/google/cloud/firestore_v1/async_collection.py @@ -13,6 +13,7 @@ # limitations under the License. """Classes for representing collections for the Google Cloud Firestore API.""" +from __future__ import annotations from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, Tuple @@ -35,6 +36,8 @@ if TYPE_CHECKING: # pragma: NO COVER from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator from google.cloud.firestore_v1.base_document import DocumentSnapshot + from google.cloud.firestore_v1.query_profile import ExplainOptions + from google.cloud.firestore_v1.query_results import QueryResultsList class AsyncCollectionReference(BaseCollectionReference[async_query.AsyncQuery]): @@ -192,7 +195,9 @@ async def get( transaction: Optional[transaction.Transaction] = None, retry: Optional[retries.AsyncRetry] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> list: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> QueryResultsList[DocumentSnapshot]: """Read the documents in this collection. This sends a ``RunQuery`` RPC and returns a list of documents @@ -207,14 +212,21 @@ async def get( system-specified policy. timeout (Otional[float]): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. If a ``transaction`` is used and it already has write operations added, this method cannot be used (i.e. read-after-write is not allowed). Returns: - list: The documents in this collection that match the query. + QueryResultsList[DocumentSnapshot]: + The documents in this collection that match the query. """ query, kwargs = self._prep_get_or_stream(retry, timeout) + if explain_options is not None: + kwargs["explain_options"] = explain_options return await query.get(transaction=transaction, **kwargs) @@ -223,7 +235,9 @@ def stream( transaction: Optional[transaction.Transaction] = None, retry: Optional[retries.AsyncRetry] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> "AsyncStreamGenerator[DocumentSnapshot]": + *, + explain_options: Optional[ExplainOptions] = None, + ) -> AsyncStreamGenerator[DocumentSnapshot]: """Read the documents in this collection. This sends a ``RunQuery`` RPC and then returns a generator which @@ -250,11 +264,17 @@ def stream( system-specified policy. timeout (Optional[float]): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: `AsyncStreamGenerator[DocumentSnapshot]`: A generator of the query results. """ query, kwargs = self._prep_get_or_stream(retry, timeout) + if explain_options: + kwargs["explain_options"] = explain_options return query.stream(transaction=transaction, **kwargs) diff --git a/google/cloud/firestore_v1/async_query.py b/google/cloud/firestore_v1/async_query.py index ca83c26306..76559d7897 100644 --- a/google/cloud/firestore_v1/async_query.py +++ b/google/cloud/firestore_v1/async_query.py @@ -20,13 +20,13 @@ """ from __future__ import annotations -from typing import TYPE_CHECKING, AsyncGenerator, List, Optional, Type +from typing import TYPE_CHECKING, Any, AsyncGenerator, List, Optional, Type from google.api_core import gapic_v1 from google.api_core import retry_async as retries from google.cloud import firestore_v1 -from google.cloud.firestore_v1 import async_document, transaction +from google.cloud.firestore_v1 import transaction from google.cloud.firestore_v1.async_aggregation import AsyncAggregationQuery from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator from google.cloud.firestore_v1.async_vector_query import AsyncVectorQuery @@ -38,12 +38,15 @@ _enum_from_direction, _query_response_to_snapshot, ) +from google.cloud.firestore_v1.query_results import QueryResultsList if TYPE_CHECKING: # pragma: NO COVER # Types needed only for Type Hints from google.cloud.firestore_v1.base_document import DocumentSnapshot from google.cloud.firestore_v1.base_vector_query import DistanceMeasure from google.cloud.firestore_v1.field_path import FieldPath + from google.cloud.firestore_v1.query_profile import ExplainMetrics, ExplainOptions + import google.cloud.firestore_v1.types.query_profile as query_profile_pb from google.cloud.firestore_v1.vector import Vector @@ -177,7 +180,9 @@ async def get( transaction: Optional[transaction.Transaction] = None, retry: Optional[retries.AsyncRetry] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> list: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> QueryResultsList[DocumentSnapshot]: """Read the documents in the collection that match this query. This sends a ``RunQuery`` RPC and returns a list of documents @@ -192,14 +197,21 @@ async def get( system-specified policy. timeout (Otional[float]): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. If a ``transaction`` is used and it already has write operations added, this method cannot be used (i.e. read-after-write is not allowed). Returns: - list: The documents in the collection that match this query. + QueryResultsList[DocumentSnapshot]: The documents in the collection + that match this query. """ + explain_metrics: ExplainMetrics | None = None + is_limited_to_last = self._limit_to_last if self._limit_to_last: @@ -217,12 +229,18 @@ async def get( transaction=transaction, retry=retry, timeout=timeout, + explain_options=explain_options, ) - result = [d async for d in result] + result_list = [d async for d in result] if is_limited_to_last: - result = list(reversed(result)) + result_list = list(reversed(result_list)) + + if explain_options is None: + explain_metrics = None + else: + explain_metrics = await result.get_explain_metrics() - return result + return QueryResultsList(result_list, explain_options, explain_metrics) def find_nearest( self, @@ -314,7 +332,8 @@ async def _make_stream( transaction: Optional[transaction.Transaction] = None, retry: Optional[retries.AsyncRetry] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> AsyncGenerator[async_document.DocumentSnapshot, None]: + explain_options: Optional[ExplainOptions] = None, + ) -> AsyncGenerator[DocumentSnapshot | query_profile_pb.ExplainMetrics, Any]: """Internal method for stream(). Read the documents in the collection that match this query. @@ -342,15 +361,23 @@ async def _make_stream( system-specified policy. timeout (Optional[float]): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Yields: - :class:`~google.cloud.firestore_v1.async_document.DocumentSnapshot`: - The next document that fulfills the query. + [:class:`~google.cloud.firestore_v1.base_document.DocumentSnapshot` \ + | google.cloud.firestore_v1.types.query_profile.ExplainMetrtics]: + The next document that fulfills the query. Query results will be + yielded as `DocumentSnapshot`. When the result contains returned + explain metrics, yield `query_profile_pb.ExplainMetrics` individually. """ request, expected_prefix, kwargs = self._prep_stream( transaction, retry, timeout, + explain_options, ) response_iterator = await self._client._firestore_api.run_query( @@ -371,12 +398,18 @@ async def _make_stream( if snapshot is not None: yield snapshot + if response.explain_metrics: + metrics = response.explain_metrics + yield metrics + def stream( self, transaction: Optional[transaction.Transaction] = None, retry: Optional[retries.AsyncRetry] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> "AsyncStreamGenerator[DocumentSnapshot]": + *, + explain_options: Optional[ExplainOptions] = None, + ) -> AsyncStreamGenerator[DocumentSnapshot]: """Read the documents in the collection that match this query. This sends a ``RunQuery`` RPC and then returns a generator which @@ -403,17 +436,22 @@ def stream( system-specified policy. timeout (Optional[float]): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: - `AsyncStreamGenerator[DocumentSnapshot]`: A generator of the query - results. + `AsyncStreamGenerator[DocumentSnapshot]`: + An asynchronous generator of the queryresults. """ inner_generator = self._make_stream( transaction=transaction, retry=retry, timeout=timeout, + explain_options=explain_options, ) - return AsyncStreamGenerator(inner_generator) + return AsyncStreamGenerator(inner_generator, explain_options) @staticmethod def _get_collection_reference_class() -> ( diff --git a/google/cloud/firestore_v1/async_stream_generator.py b/google/cloud/firestore_v1/async_stream_generator.py index ca0481c0d1..c38e6eea1b 100644 --- a/google/cloud/firestore_v1/async_stream_generator.py +++ b/google/cloud/firestore_v1/async_stream_generator.py @@ -15,27 +15,101 @@ """Classes for iterating over stream results async for the Google Cloud Firestore API. """ +from __future__ import annotations -from collections import abc +from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Optional, TypeVar +from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + QueryExplainError, +) +import google.cloud.firestore_v1.types.query_profile as query_profile_pb -class AsyncStreamGenerator(abc.AsyncGenerator): - """Asynchronous generator for the streamed results.""" +if TYPE_CHECKING: # pragma: NO COVER + from google.cloud.firestore_v1.query_profile import ExplainOptions - def __init__(self, response_generator): + +T = TypeVar("T") + + +class AsyncStreamGenerator(AsyncGenerator[T, Any]): + """Asynchronous Generator for the streamed results. + + Args: + response_generator (AsyncGenerator): + The inner generator that yields the returned results in the stream. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Query profiling options for this stream request. + """ + + def __init__( + self, + response_generator: AsyncGenerator[T | query_profile_pb.ExplainMetrics, Any], + explain_options: Optional[ExplainOptions] = None, + ): self._generator = response_generator + self._explain_options = explain_options + self._explain_metrics = None - def __aiter__(self): - return self._generator + def __aiter__(self) -> AsyncGenerator[T, Any]: + return self - def __anext__(self): - return self._generator.__anext__() + async def __anext__(self) -> T: + try: + next_value = await self._generator.__anext__() + if type(next_value) is query_profile_pb.ExplainMetrics: + self._explain_metrics = ExplainMetrics._from_pb(next_value) + raise StopAsyncIteration + else: + return next_value + except StopAsyncIteration: + raise - def asend(self, value=None): + def asend(self, value: Any = None) -> Awaitable[T]: return self._generator.asend(value) - def athrow(self, exp=None): - return self._generator.athrow(exp) + def athrow(self, *args, **kwargs) -> Awaitable[T]: + return self._generator.athrow(*args, **kwargs) def aclose(self): return self._generator.aclose() + + @property + def explain_options(self) -> ExplainOptions | None: + """Query profiling options for this stream request.""" + return self._explain_options + + async def get_explain_metrics(self) -> ExplainMetrics: + """ + Get the metrics associated with the query execution. + Metrics are only available when explain_options is set on the query. If + ExplainOptions.analyze is False, only plan_summary is available. If it is + True, execution_stats is also available. + :rtype: :class:`~google.cloud.firestore_v1.query_profile.ExplainMetrics` + :returns: The metrics associated with the query execution. + :raises: :class:`~google.cloud.firestore_v1.query_profile.QueryExplainError` + if explain_metrics is not available on the query. + """ + if self._explain_metrics is not None: + return self._explain_metrics + elif self._explain_options is None: + raise QueryExplainError("explain_options not set on query.") + elif self._explain_options.analyze is False: + # We need to run the query to get the explain_metrics. Since no + # query results are returned, it's ok to discard the returned value. + try: + await self.__anext__() + except StopAsyncIteration: + pass + + if self._explain_metrics is None: + raise QueryExplainError( + "Did not receive explain_metrics for this query, despite " + "explain_options is set and analyze = False." + ) + else: + return self._explain_metrics + raise QueryExplainError( + "explain_metrics not available until query is complete." + ) diff --git a/google/cloud/firestore_v1/async_transaction.py b/google/cloud/firestore_v1/async_transaction.py index 7281a68e56..559bea96f4 100644 --- a/google/cloud/firestore_v1/async_transaction.py +++ b/google/cloud/firestore_v1/async_transaction.py @@ -13,18 +13,15 @@ # limitations under the License. """Helpers for applying Google Cloud Firestore changes in a transaction.""" +from __future__ import annotations - -from typing import Any, AsyncGenerator, Callable, Coroutine +from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Coroutine, Optional from google.api_core import exceptions, gapic_v1 from google.api_core import retry_async as retries from google.cloud.firestore_v1 import _helpers, async_batch -from google.cloud.firestore_v1.async_document import ( - AsyncDocumentReference, - DocumentSnapshot, -) +from google.cloud.firestore_v1.async_document import AsyncDocumentReference from google.cloud.firestore_v1.async_query import AsyncQuery from google.cloud.firestore_v1.base_transaction import ( _CANT_BEGIN, @@ -37,6 +34,12 @@ _BaseTransactional, ) +# Types needed only for Type Hints +if TYPE_CHECKING: # pragma: NO COVER + from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator + from google.cloud.firestore_v1.base_document import DocumentSnapshot + from google.cloud.firestore_v1.query_profile import ExplainOptions + class AsyncTransaction(async_batch.AsyncWriteBatch, BaseTransaction): """Accumulate read-and-write operations to be sent in a transaction. @@ -169,31 +172,51 @@ async def get_all( async def get( self, - ref_or_query, + ref_or_query: AsyncDocumentReference | AsyncQuery, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, - ) -> AsyncGenerator[DocumentSnapshot, Any]: + timeout: Optional[float] = None, + *, + explain_options: Optional[ExplainOptions] = None, + ) -> AsyncGenerator[DocumentSnapshot, Any] | AsyncStreamGenerator[DocumentSnapshot]: """ Retrieve a document or a query result from the database. Args: - ref_or_query The document references or query object to return. + ref_or_query (AsyncDocumentReference | AsyncQuery): + The document references or query object to return. retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. Defaults to a system-specified policy. timeout (float): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. + Can only be used when running a query, not a document reference. Yields: - .DocumentSnapshot: The next document snapshot that fulfills the - query, or :data:`None` if the document does not exist. + DocumentSnapshot: The next document snapshot that fulfills the query, + or :data:`None` if the document does not exist. + + Raises: + ValueError: if `ref_or_query` is not one of the supported types, or + explain_options is provided when `ref_or_query` is a document + reference. """ kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) if isinstance(ref_or_query, AsyncDocumentReference): + if explain_options is not None: + raise ValueError( + "When type of `ref_or_query` is `AsyncDocumentReference`, " + "`explain_options` cannot be provided." + ) return await self._client.get_all( [ref_or_query], transaction=self, **kwargs ) elif isinstance(ref_or_query, AsyncQuery): - return await ref_or_query.stream(transaction=self, **kwargs) + if explain_options is not None: + kwargs["explain_options"] = explain_options + return ref_or_query.stream(transaction=self, **kwargs) else: raise ValueError( 'Value for argument "ref_or_query" must be a AsyncDocumentReference or a AsyncQuery.' diff --git a/google/cloud/firestore_v1/async_vector_query.py b/google/cloud/firestore_v1/async_vector_query.py index a77bc4343f..97ea3d0aa9 100644 --- a/google/cloud/firestore_v1/async_vector_query.py +++ b/google/cloud/firestore_v1/async_vector_query.py @@ -14,19 +14,26 @@ from __future__ import annotations -from typing import AsyncGenerator, List, Optional, TypeVar, Union +from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, TypeVar, Union from google.api_core import gapic_v1 from google.api_core import retry_async as retries -from google.cloud.firestore_v1 import async_document -from google.cloud.firestore_v1.base_document import DocumentSnapshot +from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator from google.cloud.firestore_v1.base_query import ( BaseQuery, _collection_group_query_response_to_snapshot, _query_response_to_snapshot, ) from google.cloud.firestore_v1.base_vector_query import BaseVectorQuery +from google.cloud.firestore_v1.query_results import QueryResultsList + +# Types needed only for Type Hints +if TYPE_CHECKING: # pragma: NO COVER + from google.cloud.firestore_v1.base_document import DocumentSnapshot + from google.cloud.firestore_v1.query_profile import ExplainMetrics, ExplainOptions + from google.cloud.firestore_v1 import transaction + import google.cloud.firestore_v1.types.query_profile as query_profile_pb TAsyncVectorQuery = TypeVar("TAsyncVectorQuery", bound="AsyncVectorQuery") @@ -49,7 +56,9 @@ async def get( transaction=None, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> List[DocumentSnapshot]: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> QueryResultsList[DocumentSnapshot]: """Runs the vector query. This sends a ``RunQuery`` RPC and returns a list of document messages. @@ -65,25 +74,43 @@ async def get( should be retried. Defaults to a system-specified policy. timeout (float): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: - list: The vector query results. + QueryResultsList[DocumentSnapshot]: The documents in the collection + that match this query. """ + explain_metrics: ExplainMetrics | None = None + stream_result = self.stream( - transaction=transaction, retry=retry, timeout=timeout + transaction=transaction, + retry=retry, + timeout=timeout, + explain_options=explain_options, ) result = [snapshot async for snapshot in stream_result] - return result # type: ignore - async def stream( + if explain_options is None: + explain_metrics = None + else: + explain_metrics = await stream_result.get_explain_metrics() + + return QueryResultsList(result, explain_options, explain_metrics) + + async def _make_stream( self, - transaction=None, - retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, + transaction: Optional[transaction.Transaction] = None, + retry: Optional[retries.Retry] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> AsyncGenerator[async_document.DocumentSnapshot, None]: - """Reads the documents in the collection that match this query. + explain_options: Optional[ExplainOptions] = None, + ) -> AsyncGenerator[[DocumentSnapshot | query_profile_pb.ExplainMetrics], Any]: + """Internal method for stream(). Read the documents in the collection + that match this query. - This sends a ``RunQuery`` RPC and then returns an iterator which + This sends a ``RunQuery`` RPC and then returns a generator which consumes each document returned in the stream of ``RunQueryResponse`` messages. @@ -92,22 +119,31 @@ async def stream( allowed). Args: - transaction - (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]): - An existing transaction that this query will run in. - retry (google.api_core.retry.Retry): Designation of what errors, if any, - should be retried. Defaults to a system-specified policy. - timeout (float): The timeout for this request. Defaults to a - system-specified value. + transaction (Optional[:class:`~google.cloud.firestore_v1.transaction.\ + Transaction`]): + An existing transaction that the query will run in. + retry (Optional[google.api_core.retry.Retry]): Designation of what + errors, if any, should be retried. Defaults to a + system-specified policy. + timeout (Optional[float]): The timeout for this request. Defaults + to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Yields: - :class:`~google.cloud.firestore_v1.document.DocumentSnapshot`: - The next document that fulfills the query. + [:class:`~google.cloud.firestore_v1.base_document.DocumentSnapshot` \ + | google.cloud.firestore_v1.types.query_profile.ExplainMetrtics]: + The next document that fulfills the query. Query results will be + yielded as `DocumentSnapshot`. When the result contains returned + explain metrics, yield `query_profile_pb.ExplainMetrics` individually. """ request, expected_prefix, kwargs = self._prep_stream( transaction, retry, timeout, + explain_options, ) response_iterator = await self._client._firestore_api.run_query( @@ -127,3 +163,51 @@ async def stream( ) if snapshot is not None: yield snapshot + + if response.explain_metrics: + metrics = response.explain_metrics + yield metrics + + def stream( + self, + transaction=None, + retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + *, + explain_options: Optional[ExplainOptions] = None, + ) -> AsyncStreamGenerator[DocumentSnapshot]: + """Reads the documents in the collection that match this query. + + This sends a ``RunQuery`` RPC and then returns an iterator which + consumes each document returned in the stream of ``RunQueryResponse`` + messages. + + If a ``transaction`` is used and it already has write operations + added, this method cannot be used (i.e. read-after-write is not + allowed). + + Args: + transaction + (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]): + An existing transaction that this query will run in. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. Defaults to a system-specified policy. + timeout (float): The timeout for this request. Defaults to a + system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. + + Returns: + `AsyncStreamGenerator[DocumentSnapshot]`: + An asynchronous generator of the queryresults. + """ + + inner_generator = self._make_stream( + transaction=transaction, + retry=retry, + timeout=timeout, + explain_options=explain_options, + ) + return AsyncStreamGenerator(inner_generator, explain_options) diff --git a/google/cloud/firestore_v1/base_aggregation.py b/google/cloud/firestore_v1/base_aggregation.py index f922663791..807c753f1f 100644 --- a/google/cloud/firestore_v1/base_aggregation.py +++ b/google/cloud/firestore_v1/base_aggregation.py @@ -24,17 +24,7 @@ import abc from abc import ABC -from typing import ( - TYPE_CHECKING, - Any, - AsyncGenerator, - Coroutine, - Generator, - List, - Optional, - Tuple, - Union, -) +from typing import TYPE_CHECKING, Any, Coroutine, List, Optional, Tuple, Union from google.api_core import gapic_v1 from google.api_core import retry as retries @@ -47,8 +37,14 @@ ) # Types needed only for Type Hints -if TYPE_CHECKING: - from google.cloud.firestore_v1 import transaction # pragma: NO COVER +if TYPE_CHECKING: # pragma: NO COVER + from google.cloud.firestore_v1 import transaction + from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator + from google.cloud.firestore_v1.query_profile import ExplainOptions + from google.cloud.firestore_v1.query_results import QueryResultsList + from google.cloud.firestore_v1.stream_generator import ( + StreamGenerator, + ) class AggregationResult(object): @@ -62,7 +58,7 @@ class AggregationResult(object): :param value: The resulting read_time """ - def __init__(self, alias: str, value: int, read_time=None): + def __init__(self, alias: str, value: float, read_time=None): self.alias = alias self.value = value self.read_time = read_time @@ -211,6 +207,7 @@ def _prep_stream( transaction=None, retry: Union[retries.Retry, None, gapic_v1.method._MethodDefault] = None, timeout: float | None = None, + explain_options: Optional[ExplainOptions] = None, ) -> Tuple[dict, dict]: parent_path, expected_prefix = self._collection_ref._parent_info() request = { @@ -218,6 +215,8 @@ def _prep_stream( "structured_aggregation_query": self._to_protobuf(), "transaction": _helpers.get_transaction_id(transaction), } + if explain_options: + request["explain_options"] = explain_options._to_dict() kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return request, kwargs @@ -230,10 +229,17 @@ def get( retries.Retry, None, gapic_v1.method._MethodDefault ] = gapic_v1.method.DEFAULT, timeout: float | None = None, - ) -> List[AggregationResult] | Coroutine[Any, Any, List[AggregationResult]]: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> ( + QueryResultsList[AggregationResult] + | Coroutine[Any, Any, List[List[AggregationResult]]] + ): """Runs the aggregation query. - This sends a ``RunAggregationQuery`` RPC and returns a list of aggregation results in the stream of ``RunAggregationQueryResponse`` messages. + This sends a ``RunAggregationQuery`` RPC and returns a list of + aggregation results in the stream of ``RunAggregationQueryResponse`` + messages. Args: transaction @@ -246,21 +252,29 @@ def get( should be retried. Defaults to a system-specified policy. timeout (float): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: - list: The aggregation query results - + (QueryResultsList[List[AggregationResult]] | Coroutine[Any, Any, List[List[AggregationResult]]]): + The aggregation query results. """ @abc.abstractmethod def stream( self, transaction: Optional[transaction.Transaction] = None, - retry: Optional[retries.Retry] = gapic_v1.method.DEFAULT, + retry: Union[ + retries.Retry, None, gapic_v1.method._MethodDefault + ] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, + *, + explain_options: Optional[ExplainOptions] = None, ) -> ( - Generator[List[AggregationResult], Any, None] - | AsyncGenerator[List[AggregationResult], None] + StreamGenerator[List[AggregationResult]] + | AsyncStreamGenerator[List[AggregationResult]] ): """Runs the aggregation query. @@ -274,8 +288,13 @@ def stream( errors, if any, should be retried. Defaults to a system-specified policy. timeout (Optinal[float]): The timeout for this request. Defaults - to a system-specified value. + to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: + StreamGenerator[List[AggregationResult]] | AsyncStreamGenerator[List[AggregationResult]]: A generator of the query results. """ diff --git a/google/cloud/firestore_v1/base_collection.py b/google/cloud/firestore_v1/base_collection.py index 18c62aa33b..1ac1ba3184 100644 --- a/google/cloud/firestore_v1/base_collection.py +++ b/google/cloud/firestore_v1/base_collection.py @@ -25,7 +25,6 @@ Generator, Generic, Iterable, - Iterator, NoReturn, Optional, Tuple, @@ -35,19 +34,24 @@ from google.api_core import retry as retries from google.cloud.firestore_v1 import _helpers -from google.cloud.firestore_v1.base_aggregation import BaseAggregationQuery from google.cloud.firestore_v1.base_query import QueryType -from google.cloud.firestore_v1.base_vector_query import BaseVectorQuery, DistanceMeasure -from google.cloud.firestore_v1.document import DocumentReference -from google.cloud.firestore_v1.vector import Vector if TYPE_CHECKING: # pragma: NO COVER # Types needed only for Type Hints - from firestore_v1.vector_query import VectorQuery - + from google.cloud.firestore_v1.base_aggregation import BaseAggregationQuery from google.cloud.firestore_v1.base_document import DocumentSnapshot + from google.cloud.firestore_v1.base_vector_query import ( + BaseVectorQuery, + DistanceMeasure, + ) + from google.cloud.firestore_v1.document import DocumentReference from google.cloud.firestore_v1.field_path import FieldPath + from google.cloud.firestore_v1.query_profile import ExplainOptions + from google.cloud.firestore_v1.query_results import QueryResultsList + from google.cloud.firestore_v1.stream_generator import StreamGenerator from google.cloud.firestore_v1.transaction import Transaction + from google.cloud.firestore_v1.vector import Vector + from google.cloud.firestore_v1.vector_query import VectorQuery _AUTO_ID_CHARS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" @@ -492,9 +496,12 @@ def get( transaction: Optional[Transaction] = None, retry: Optional[retries.Retry] = None, timeout: Optional[float] = None, - ) -> Union[ - Generator[DocumentSnapshot, Any, Any], AsyncGenerator[DocumentSnapshot, Any] - ]: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> ( + QueryResultsList[DocumentSnapshot] + | Coroutine[Any, Any, QueryResultsList[DocumentSnapshot]] + ): raise NotImplementedError def stream( @@ -502,7 +509,9 @@ def stream( transaction: Optional[Transaction] = None, retry: Optional[retries.Retry] = None, timeout: Optional[float] = None, - ) -> Union[Iterator[DocumentSnapshot], AsyncIterator[DocumentSnapshot]]: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> StreamGenerator[DocumentSnapshot] | AsyncIterator[DocumentSnapshot]: raise NotImplementedError def on_snapshot(self, callback) -> NoReturn: diff --git a/google/cloud/firestore_v1/base_document.py b/google/cloud/firestore_v1/base_document.py index 1418ea34d0..ada42acb3e 100644 --- a/google/cloud/firestore_v1/base_document.py +++ b/google/cloud/firestore_v1/base_document.py @@ -13,17 +13,29 @@ # limitations under the License. """Classes for representing documents for the Google Cloud Firestore API.""" +from __future__ import annotations import copy -from typing import Any, Dict, Iterable, NoReturn, Optional, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Iterable, + NoReturn, + Optional, + Tuple, + Union, +) from google.api_core import retry as retries from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1 import field_path as field_path_module +from google.cloud.firestore_v1.types import common # Types needed only for Type Hints -from google.cloud.firestore_v1.types import Document, common, firestore, write +if TYPE_CHECKING: # pragma: NO COVER + from google.cloud.firestore_v1.types import Document, firestore, write class BaseDocumentReference(object): diff --git a/google/cloud/firestore_v1/base_query.py b/google/cloud/firestore_v1/base_query.py index cfed454b93..5cdbf4c50a 100644 --- a/google/cloud/firestore_v1/base_query.py +++ b/google/cloud/firestore_v1/base_query.py @@ -27,8 +27,8 @@ from typing import ( TYPE_CHECKING, Any, + Coroutine, Dict, - Generator, Iterable, NoReturn, Optional, @@ -59,8 +59,13 @@ from google.cloud.firestore_v1.vector import Vector if TYPE_CHECKING: # pragma: NO COVER + from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator from google.cloud.firestore_v1.base_vector_query import BaseVectorQuery from google.cloud.firestore_v1.field_path import FieldPath + from google.cloud.firestore_v1.query_profile import ExplainOptions + from google.cloud.firestore_v1.query_results import QueryResultsList + from google.cloud.firestore_v1.stream_generator import StreamGenerator + _BAD_DIR_STRING: str _BAD_OP_NAN_NULL: str @@ -1008,7 +1013,12 @@ def get( transaction=None, retry: Optional[retries.Retry] = None, timeout: Optional[float] = None, - ) -> Iterable[DocumentSnapshot]: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> ( + QueryResultsList[DocumentSnapshot] + | Coroutine[Any, Any, QueryResultsList[DocumentSnapshot]] + ): raise NotImplementedError def _prep_stream( @@ -1016,6 +1026,7 @@ def _prep_stream( transaction=None, retry: Optional[retries.Retry] = None, timeout: Optional[float] = None, + explain_options: Optional[ExplainOptions] = None, ) -> Tuple[dict, str, dict]: """Shared setup for async / sync :meth:`stream`""" if self._limit_to_last: @@ -1030,6 +1041,8 @@ def _prep_stream( "structured_query": self._to_protobuf(), "transaction": _helpers.get_transaction_id(transaction), } + if explain_options is not None: + request["explain_options"] = explain_options._to_dict() kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return request, expected_prefix, kwargs @@ -1039,7 +1052,12 @@ def stream( transaction=None, retry: Optional[retries.Retry] = None, timeout: Optional[float] = None, - ) -> Generator[document.DocumentSnapshot, Any, None]: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> ( + StreamGenerator[document.DocumentSnapshot] + | AsyncStreamGenerator[DocumentSnapshot] + ): raise NotImplementedError def on_snapshot(self, callback) -> NoReturn: diff --git a/google/cloud/firestore_v1/base_transaction.py b/google/cloud/firestore_v1/base_transaction.py index 09f0c1fb9a..752c83169d 100644 --- a/google/cloud/firestore_v1/base_transaction.py +++ b/google/cloud/firestore_v1/base_transaction.py @@ -13,13 +13,31 @@ # limitations under the License. """Helpers for applying Google Cloud Firestore changes in a transaction.""" - -from typing import Any, Coroutine, NoReturn, Optional, Union +from __future__ import annotations + +from typing import ( + TYPE_CHECKING, + Any, + AsyncGenerator, + Coroutine, + Generator, + NoReturn, + Optional, + Union, +) from google.api_core import retry as retries from google.cloud.firestore_v1 import types +# Types needed only for Type Hints +if TYPE_CHECKING: # pragma: NO COVER + from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator + from google.cloud.firestore_v1.document import DocumentSnapshot + from google.cloud.firestore_v1.query_profile import ExplainOptions + from google.cloud.firestore_v1.stream_generator import StreamGenerator + + _CANT_BEGIN: str _CANT_COMMIT: str _CANT_RETRY_READ_ONLY: str @@ -142,7 +160,10 @@ def get_all( references: list, retry: retries.Retry = None, timeout: float = None, - ) -> NoReturn: + ) -> ( + Generator[DocumentSnapshot, Any, None] + | Coroutine[Any, Any, AsyncGenerator[DocumentSnapshot, Any]] + ): raise NotImplementedError def get( @@ -150,7 +171,14 @@ def get( ref_or_query, retry: retries.Retry = None, timeout: float = None, - ) -> NoReturn: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> ( + StreamGenerator[DocumentSnapshot] + | Generator[DocumentSnapshot, Any, None] + | Coroutine[Any, Any, AsyncGenerator[DocumentSnapshot, Any]] + | Coroutine[Any, Any, AsyncStreamGenerator[DocumentSnapshot]] + ): raise NotImplementedError diff --git a/google/cloud/firestore_v1/base_vector_query.py b/google/cloud/firestore_v1/base_vector_query.py index 26cd5b1997..30c79bc7e2 100644 --- a/google/cloud/firestore_v1/base_vector_query.py +++ b/google/cloud/firestore_v1/base_vector_query.py @@ -14,19 +14,26 @@ """Classes for representing vector queries for the Google Cloud Firestore API. """ +from __future__ import annotations import abc from abc import ABC from enum import Enum -from typing import Iterable, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Coroutine, Optional, Tuple, Union from google.api_core import gapic_v1 from google.api_core import retry as retries -from google.cloud.firestore_v1 import _helpers, document -from google.cloud.firestore_v1.base_document import DocumentSnapshot +from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1.types import query -from google.cloud.firestore_v1.vector import Vector + +if TYPE_CHECKING: # pragma: NO COVER + from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator + from google.cloud.firestore_v1.base_document import DocumentSnapshot + from google.cloud.firestore_v1.query_profile import ExplainOptions + from google.cloud.firestore_v1.query_results import QueryResultsList + from google.cloud.firestore_v1.stream_generator import StreamGenerator + from google.cloud.firestore_v1.vector import Vector class DistanceMeasure(Enum): @@ -94,6 +101,7 @@ def _prep_stream( transaction=None, retry: Union[retries.Retry, None, gapic_v1.method._MethodDefault] = None, timeout: Optional[float] = None, + explain_options: Optional[ExplainOptions] = None, ) -> Tuple[dict, str, dict]: parent_path, expected_prefix = self._collection_ref._parent_info() request = { @@ -103,6 +111,9 @@ def _prep_stream( } kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) + if explain_options is not None: + request["explain_options"] = explain_options._to_dict() + return request, expected_prefix, kwargs @abc.abstractmethod @@ -111,8 +122,14 @@ def get( transaction=None, retry: retries.Retry = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> Iterable[DocumentSnapshot]: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> ( + QueryResultsList[DocumentSnapshot] + | Coroutine[Any, Any, QueryResultsList[DocumentSnapshot]] + ): """Runs the vector query.""" + raise NotImplementedError def find_nearest( self, @@ -137,6 +154,9 @@ def stream( self, transaction=None, retry: retries.Retry = gapic_v1.method.DEFAULT, - timeout: float = None, - ) -> Iterable[document.DocumentSnapshot]: + timeout: Optional[float] = None, + *, + explain_options: Optional[ExplainOptions] = None, + ) -> StreamGenerator[DocumentSnapshot] | AsyncStreamGenerator[DocumentSnapshot]: """Reads the documents in the collection that match this query.""" + raise NotImplementedError diff --git a/google/cloud/firestore_v1/collection.py b/google/cloud/firestore_v1/collection.py index 96dadf2e70..372dacd7b1 100644 --- a/google/cloud/firestore_v1/collection.py +++ b/google/cloud/firestore_v1/collection.py @@ -13,6 +13,7 @@ # limitations under the License. """Classes for representing collections for the Google Cloud Firestore API.""" +from __future__ import annotations from typing import TYPE_CHECKING, Any, Callable, Generator, Optional, Tuple, Union @@ -26,10 +27,12 @@ BaseCollectionReference, _item_to_document_ref, ) +from google.cloud.firestore_v1.query_results import QueryResultsList from google.cloud.firestore_v1.watch import Watch if TYPE_CHECKING: # pragma: NO COVER from google.cloud.firestore_v1.base_document import DocumentSnapshot + from google.cloud.firestore_v1.query_profile import ExplainOptions from google.cloud.firestore_v1.stream_generator import StreamGenerator @@ -169,7 +172,9 @@ def get( transaction: Union[transaction.Transaction, None] = None, retry: retries.Retry = gapic_v1.method.DEFAULT, timeout: Union[float, None] = None, - ) -> list: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> QueryResultsList[DocumentSnapshot]: """Read the documents in this collection. This sends a ``RunQuery`` RPC and returns a list of documents @@ -183,15 +188,22 @@ def get( should be retried. Defaults to a system-specified policy. timeout (float): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. If a ``transaction`` is used and it already has write operations added, this method cannot be used (i.e. read-after-write is not allowed). Returns: - list: The documents in this collection that match the query. + QueryResultsList[DocumentSnapshot]: The documents in this collection + that match the query. """ query, kwargs = self._prep_get_or_stream(retry, timeout) + if explain_options is not None: + kwargs["explain_options"] = explain_options return query.get(transaction=transaction, **kwargs) @@ -200,7 +212,9 @@ def stream( transaction: Optional[transaction.Transaction] = None, retry: Optional[retries.Retry] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> "StreamGenerator[DocumentSnapshot]": + *, + explain_options: Optional[ExplainOptions] = None, + ) -> StreamGenerator[DocumentSnapshot]: """Read the documents in this collection. This sends a ``RunQuery`` RPC and then returns an iterator which @@ -227,11 +241,17 @@ def stream( system-specified policy. timeout (Optional[float]): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: `StreamGenerator[DocumentSnapshot]`: A generator of the query results. """ query, kwargs = self._prep_get_or_stream(retry, timeout) + if explain_options: + kwargs["explain_options"] = explain_options return query.stream(transaction=transaction, **kwargs) diff --git a/google/cloud/firestore_v1/gapic_version.py b/google/cloud/firestore_v1/gapic_version.py index f09943f6bd..0f1a446f38 100644 --- a/google/cloud/firestore_v1/gapic_version.py +++ b/google/cloud/firestore_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "2.18.0" # {x-release-please-version} +__version__ = "2.19.0" # {x-release-please-version} diff --git a/google/cloud/firestore_v1/query.py b/google/cloud/firestore_v1/query.py index eb8f51dc8d..818a713c5b 100644 --- a/google/cloud/firestore_v1/query.py +++ b/google/cloud/firestore_v1/query.py @@ -27,7 +27,10 @@ from google.cloud import firestore_v1 from google.cloud.firestore_v1 import aggregation, transaction -from google.cloud.firestore_v1.base_document import DocumentSnapshot +from google.cloud.firestore_v1.query_results import QueryResultsList +from google.cloud.firestore_v1.base_document import ( + DocumentSnapshot, +) from google.cloud.firestore_v1.base_query import ( BaseCollectionGroup, BaseQuery, @@ -36,14 +39,15 @@ _enum_from_direction, _query_response_to_snapshot, ) -from google.cloud.firestore_v1.base_vector_query import DistanceMeasure from google.cloud.firestore_v1.stream_generator import StreamGenerator from google.cloud.firestore_v1.vector import Vector from google.cloud.firestore_v1.vector_query import VectorQuery from google.cloud.firestore_v1.watch import Watch if TYPE_CHECKING: # pragma: NO COVER + from google.cloud.firestore_v1.base_vector_query import DistanceMeasure from google.cloud.firestore_v1.field_path import FieldPath + from google.cloud.firestore_v1.query_profile import ExplainMetrics, ExplainOptions class Query(BaseQuery): @@ -134,8 +138,10 @@ def get( self, transaction=None, retry: retries.Retry = gapic_v1.method.DEFAULT, - timeout: float = None, - ) -> List[DocumentSnapshot]: + timeout: Optional[float] = None, + *, + explain_options: Optional[ExplainOptions] = None, + ) -> QueryResultsList[DocumentSnapshot]: """Read the documents in the collection that match this query. This sends a ``RunQuery`` RPC and returns a list of documents @@ -152,10 +158,17 @@ def get( should be retried. Defaults to a system-specified policy. timeout (float): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: - list: The documents in the collection that match this query. + QueryResultsList[DocumentSnapshot]: The documents in the collection + that match this query. """ + explain_metrics: ExplainMetrics | None = None + is_limited_to_last = self._limit_to_last if self._limit_to_last: @@ -174,11 +187,18 @@ def get( transaction=transaction, retry=retry, timeout=timeout, + explain_options=explain_options, ) + result_list = list(result) if is_limited_to_last: - result = reversed(list(result)) + result_list = list(reversed(result_list)) + + if explain_options is None: + explain_metrics = None + else: + explain_metrics = result.get_explain_metrics() - return list(result) + return QueryResultsList(result_list, explain_options, explain_metrics) def _chunkify( self, chunk_size: int @@ -218,12 +238,13 @@ def _chunkify( ): return - def _get_stream_iterator(self, transaction, retry, timeout): + def _get_stream_iterator(self, transaction, retry, timeout, explain_options=None): """Helper method for :meth:`stream`.""" request, expected_prefix, kwargs = self._prep_stream( transaction, retry, timeout, + explain_options, ) response_iterator = self._client._firestore_api.run_query( @@ -331,7 +352,8 @@ def _make_stream( transaction: Optional[transaction.Transaction] = None, retry: Optional[retries.Retry] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> Generator[DocumentSnapshot, Any, None]: + explain_options: Optional[ExplainOptions] = None, + ) -> Generator[DocumentSnapshot, Any, Optional[ExplainMetrics]]: """Internal method for stream(). Read the documents in the collection that match this query. @@ -360,15 +382,26 @@ def _make_stream( system-specified policy. timeout (Optional[float]): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Yields: - :class:`~google.cloud.firestore_v1.document.DocumentSnapshot`: + DocumentSnapshot: The next document that fulfills the query. + + Returns: + ([google.cloud.firestore_v1.types.query_profile.ExplainMetrtics | None]): + The results of query profiling, if received from the service. """ + metrics: ExplainMetrics | None = None + response_iterator, expected_prefix = self._get_stream_iterator( transaction, retry, timeout, + explain_options, ) last_snapshot = None @@ -391,6 +424,9 @@ def _make_stream( if response is None: # EOI break + if metrics is None and response.explain_metrics: + metrics = response.explain_metrics + if self._all_descendants: snapshot = _collection_group_query_response_to_snapshot( response, self._parent @@ -403,12 +439,16 @@ def _make_stream( last_snapshot = snapshot yield snapshot + return metrics + def stream( self, transaction: Optional[transaction.Transaction] = None, retry: Optional[retries.Retry] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> "StreamGenerator[DocumentSnapshot]": + *, + explain_options: Optional[ExplainOptions] = None, + ) -> StreamGenerator[DocumentSnapshot]: """Read the documents in the collection that match this query. This sends a ``RunQuery`` RPC and then returns a generator which @@ -434,7 +474,11 @@ def stream( errors, if any, should be retried. Defaults to a system-specified policy. timeout (Optinal[float]): The timeout for this request. Defaults - to a system-specified value. + to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: `StreamGenerator[DocumentSnapshot]`: A generator of the query results. @@ -443,8 +487,9 @@ def stream( transaction=transaction, retry=retry, timeout=timeout, + explain_options=explain_options, ) - return StreamGenerator(inner_generator) + return StreamGenerator(inner_generator, explain_options) def on_snapshot(self, callback: Callable) -> Watch: """Monitor the documents in this collection that match this query. diff --git a/google/cloud/firestore_v1/query_profile.py b/google/cloud/firestore_v1/query_profile.py new file mode 100644 index 0000000000..6925f83ffa --- /dev/null +++ b/google/cloud/firestore_v1/query_profile.py @@ -0,0 +1,145 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from typing import Any + +import datetime + +from dataclasses import dataclass +from google.protobuf.json_format import MessageToDict + + +@dataclass(frozen=True) +class ExplainOptions: + """ + Explain options for the query. + Set on a query object using the explain_options attribute at query + construction time. + + :type analyze: bool + :param analyze: Optional. Whether to execute this query. When false + (the default), the query will be planned, returning only metrics from the + planning stages. When true, the query will be planned and executed, + returning the full query results along with both planning and execution + stage metrics. + """ + + analyze: bool = False + + def _to_dict(self): + return {"analyze": self.analyze} + + +@dataclass(frozen=True) +class PlanSummary: + """ + Contains planning phase information about a query.` + + :type indexes_used: list[dict[str, Any]] + :param indexes_used: The indexes selected for this query. + """ + + indexes_used: list[dict[str, Any]] + + +@dataclass(frozen=True) +class ExecutionStats: + """ + Execution phase information about a query. + + Only available when explain_options.analyze is True. + + :type results_returned: int + :param results_returned: Total number of results returned, including + documents, projections, aggregation results, keys. + :type execution_duration: datetime.timedelta + :param execution_duration: Total time to execute the query in the backend. + :type read_operations: int + :param read_operations: Total billable read operations. + :type debug_stats: dict[str, Any] + :param debug_stats: Debugging statistics from the execution of the query. + Note that the debugging stats are subject to change as Firestore evolves + """ + + results_returned: int + execution_duration: datetime.timedelta + read_operations: int + debug_stats: dict[str, Any] + + +@dataclass(frozen=True) +class ExplainMetrics: + """ + ExplainMetrics contains information about the planning and execution of a query. + + When explain_options.analyze is false, only plan_summary is available. + When explain_options.analyze is true, execution_stats is also available. + + :type plan_summary: PlanSummary + :param plan_summary: Planning phase information about the query. + :type execution_stats: ExecutionStats + :param execution_stats: Execution phase information about the query. + """ + + plan_summary: PlanSummary + + @staticmethod + def _from_pb(metrics_pb): + dict_repr = MessageToDict(metrics_pb._pb, preserving_proto_field_name=True) + plan_summary = PlanSummary( + indexes_used=dict_repr.get("plan_summary", {}).get("indexes_used", []) + ) + if "execution_stats" in dict_repr: + stats_dict = dict_repr.get("execution_stats", {}) + execution_stats = ExecutionStats( + results_returned=int(stats_dict.get("results_returned", 0)), + execution_duration=metrics_pb.execution_stats.execution_duration, + read_operations=int(stats_dict.get("read_operations", 0)), + debug_stats=stats_dict.get("debug_stats", {}), + ) + return _ExplainAnalyzeMetrics( + plan_summary=plan_summary, _execution_stats=execution_stats + ) + else: + return ExplainMetrics(plan_summary=plan_summary) + + @property + def execution_stats(self) -> ExecutionStats: + raise QueryExplainError( + "execution_stats not available when explain_options.analyze=False." + ) + + +@dataclass(frozen=True) +class _ExplainAnalyzeMetrics(ExplainMetrics): + """ + Subclass of ExplainMetrics that includes execution_stats. + Only available when explain_options.analyze is True. + """ + + plan_summary: PlanSummary + _execution_stats: ExecutionStats + + @property + def execution_stats(self) -> ExecutionStats: + return self._execution_stats + + +class QueryExplainError(Exception): + """ + Error returned when there is a problem accessing query profiling information. + """ + + pass diff --git a/google/cloud/firestore_v1/query_results.py b/google/cloud/firestore_v1/query_results.py new file mode 100644 index 0000000000..47dddf9de7 --- /dev/null +++ b/google/cloud/firestore_v1/query_results.py @@ -0,0 +1,87 @@ +# Copyright 2024 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExplainOptions, + QueryExplainError, +) + +from typing import List, Optional, TypeVar + +T = TypeVar("T") + + +class QueryResultsList(List[T]): + """A list of received query results from the query call. + + This is a subclass of the built-in list. A new property `explain_metrics` + is added to return the query profile results. + + Args: + docs (list): + The list of query results. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. + explain_metrics (Optional[ExplainMetrics]): + Query profile results. + """ + + def __init__( + self, + docs: List, + explain_options: Optional[ExplainOptions] = None, + explain_metrics: Optional[ExplainMetrics] = None, + ): + super().__init__(docs) + + # When explain_options is set, explain_metrics should be non-empty too. + if explain_options is not None and explain_metrics is None: + raise ValueError( + "If explain_options is set, explain_metrics must be non-empty." + ) + elif explain_options is None and explain_metrics is not None: + raise ValueError( + "If explain_options is empty, explain_metrics must be empty." + ) + + self._explain_options = explain_options + self._explain_metrics = explain_metrics + + @property + def explain_options(self) -> Optional[ExplainOptions]: + """Query profiling options for getting these query results.""" + return self._explain_options + + def get_explain_metrics(self) -> ExplainMetrics: + """ + Get the metrics associated with the query execution. + Metrics are only available when explain_options is set on the query. If + ExplainOptions.analyze is False, only plan_summary is available. If it is + True, execution_stats is also available. + :rtype: :class:`~google.cloud.firestore_v1.query_profile.ExplainMetrics` + :returns: The metrics associated with the query execution. + :raises: :class:`~google.cloud.firestore_v1.query_profile.QueryExplainError` + if explain_metrics is not available on the query. + """ + if self._explain_options is None: + raise QueryExplainError("explain_options not set on query.") + elif self._explain_metrics is None: + raise QueryExplainError( + "explain_metrics is empty despite explain_options is set." + ) + else: + return self._explain_metrics diff --git a/google/cloud/firestore_v1/stream_generator.py b/google/cloud/firestore_v1/stream_generator.py index 0a95af8d1f..7e39a3fbab 100644 --- a/google/cloud/firestore_v1/stream_generator.py +++ b/google/cloud/firestore_v1/stream_generator.py @@ -14,27 +14,98 @@ """Classes for iterating over stream results for the Google Cloud Firestore API. """ +from __future__ import annotations -from collections import abc +from typing import TYPE_CHECKING, Any, Generator, Optional, TypeVar +from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + QueryExplainError, +) -class StreamGenerator(abc.Generator): - """Generator for the streamed results.""" +if TYPE_CHECKING: # pragma: NO COVER + from google.cloud.firestore_v1.query_profile import ExplainOptions - def __init__(self, response_generator): + +T = TypeVar("T") + + +class StreamGenerator(Generator[T, Any, Optional[ExplainMetrics]]): + """Generator for the streamed results. + + Args: + response_generator (Generator[T, Any, Optional[ExplainMetrics]]): + The inner generator that yields the returned document in the stream. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Query profiling options for this stream request. + """ + + def __init__( + self, + response_generator: Generator[T, Any, Optional[ExplainMetrics]], + explain_options: Optional[ExplainOptions] = None, + ): self._generator = response_generator + self._explain_options = explain_options + self._explain_metrics = None - def __iter__(self): - return self._generator + def __iter__(self) -> StreamGenerator: + return self - def __next__(self): - return self._generator.__next__() + def __next__(self) -> T: + try: + return self._generator.__next__() + except StopIteration as e: + # If explain_metrics is available, it would be returned. + if e.value: + self._explain_metrics = ExplainMetrics._from_pb(e.value) + raise - def send(self, value=None): + def send(self, value: Any = None) -> T: return self._generator.send(value) - def throw(self, exp=None): - return self._generator.throw(exp) + def throw(self, *args, **kwargs) -> T: + return self._generator.throw(*args, **kwargs) def close(self): return self._generator.close() + + @property + def explain_options(self) -> ExplainOptions | None: + """Query profiling options for this stream request.""" + return self._explain_options + + def get_explain_metrics(self) -> ExplainMetrics: + """ + Get the metrics associated with the query execution. + Metrics are only available when explain_options is set on the query. If + ExplainOptions.analyze is False, only plan_summary is available. If it is + True, execution_stats is also available. + :rtype: :class:`~google.cloud.firestore_v1.query_profile.ExplainMetrics` + :returns: The metrics associated with the query execution. + :raises: :class:`~google.cloud.firestore_v1.query_profile.QueryExplainError` + if explain_metrics is not available on the query. + """ + if self._explain_metrics is not None: + return self._explain_metrics + elif self._explain_options is None: + raise QueryExplainError("explain_options not set on query.") + elif self._explain_options.analyze is False: + # We need to run the query to get the explain_metrics. Since no + # query results are returned, it's ok to discard the returned value. + try: + next(self) + except StopIteration: + pass + + if self._explain_metrics is None: + raise QueryExplainError( + "Did not receive explain_metrics for this query, despite " + "explain_options is set and analyze = False." + ) + else: + return self._explain_metrics + raise QueryExplainError( + "explain_metrics not available until query is complete." + ) diff --git a/google/cloud/firestore_v1/transaction.py b/google/cloud/firestore_v1/transaction.py index 8f92ddaf0d..a01c1ed53e 100644 --- a/google/cloud/firestore_v1/transaction.py +++ b/google/cloud/firestore_v1/transaction.py @@ -13,17 +13,14 @@ # limitations under the License. """Helpers for applying Google Cloud Firestore changes in a transaction.""" +from __future__ import annotations - -from typing import Any, Callable, Generator +from typing import TYPE_CHECKING, Any, Callable, Generator, Optional from google.api_core import exceptions, gapic_v1 from google.api_core import retry as retries from google.cloud.firestore_v1 import _helpers, batch - -# Types needed only for Type Hints -from google.cloud.firestore_v1.base_document import DocumentSnapshot from google.cloud.firestore_v1.base_transaction import ( _CANT_BEGIN, _CANT_COMMIT, @@ -37,6 +34,12 @@ from google.cloud.firestore_v1.document import DocumentReference from google.cloud.firestore_v1.query import Query +# Types needed only for Type Hints +if TYPE_CHECKING: # pragma: NO COVER + from google.cloud.firestore_v1.base_document import DocumentSnapshot + from google.cloud.firestore_v1.query_profile import ExplainOptions + from google.cloud.firestore_v1.stream_generator import StreamGenerator + class Transaction(batch.WriteBatch, BaseTransaction): """Accumulate read-and-write operations to be sent in a transaction. @@ -169,27 +172,47 @@ def get_all( def get( self, - ref_or_query, + ref_or_query: DocumentReference | Query, retry: retries.Retry = gapic_v1.method.DEFAULT, - timeout: float = None, - ) -> Generator[DocumentSnapshot, Any, None]: + timeout: Optional[float] = None, + *, + explain_options: Optional[ExplainOptions] = None, + ) -> StreamGenerator[DocumentSnapshot] | Generator[DocumentSnapshot, Any, None]: """Retrieve a document or a query result from the database. Args: - ref_or_query: The document references or query object to return. + ref_or_query (DocumentReference | Query): + The document references or query object to return. retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. Defaults to a system-specified policy. timeout (float): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. + Can only be used when running a query, not a document reference. Yields: .DocumentSnapshot: The next document snapshot that fulfills the query, or :data:`None` if the document does not exist. + + Raises: + ValueError: if `ref_or_query` is not one of the supported types, or + explain_options is provided when `ref_or_query` is a document + reference. """ kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) if isinstance(ref_or_query, DocumentReference): + if explain_options is not None: + raise ValueError( + "When type of `ref_or_query` is `AsyncDocumentReference`, " + "`explain_options` cannot be provided." + ) return self._client.get_all([ref_or_query], transaction=self, **kwargs) elif isinstance(ref_or_query, Query): + if explain_options is not None: + kwargs["explain_options"] = explain_options return ref_or_query.stream(transaction=self, **kwargs) else: raise ValueError( diff --git a/google/cloud/firestore_v1/transforms.py b/google/cloud/firestore_v1/transforms.py index ae061f6b30..5ec15b3dc2 100644 --- a/google/cloud/firestore_v1/transforms.py +++ b/google/cloud/firestore_v1/transforms.py @@ -102,7 +102,7 @@ class _NumericValue(object): """Hold a single integer / float value. Args: - value (int | float): value held in the helper. + value (float): value held in the helper. """ def __init__(self, value) -> None: @@ -116,7 +116,7 @@ def value(self): """Value used by the transform. Returns: - (Integer | Float) value passed in the constructor. + (Lloat) value passed in the constructor. """ return self._value @@ -133,7 +133,7 @@ class Increment(_NumericValue): https://cloud.google.com/firestore/docs/reference/rpc/google.firestore.v1#google.firestore.v1.DocumentTransform.FieldTransform.FIELDS.google.firestore.v1.ArrayValue.google.firestore.v1.DocumentTransform.FieldTransform.increment Args: - value (int | float): value used to increment the field. + value (float): value used to increment the field. """ @@ -144,7 +144,7 @@ class Maximum(_NumericValue): https://cloud.google.com/firestore/docs/reference/rpc/google.firestore.v1#google.firestore.v1.DocumentTransform.FieldTransform.FIELDS.google.firestore.v1.ArrayValue.google.firestore.v1.DocumentTransform.FieldTransform.maximum Args: - value (int | float): value used to bound the field. + value (float): value used to bound the field. """ @@ -155,5 +155,5 @@ class Minimum(_NumericValue): https://cloud.google.com/firestore/docs/reference/rpc/google.firestore.v1#google.firestore.v1.DocumentTransform.FieldTransform.FIELDS.google.firestore.v1.ArrayValue.google.firestore.v1.DocumentTransform.FieldTransform.minimum Args: - value (int | float): value used to bound the field. + value (float): value used to bound the field. """ diff --git a/google/cloud/firestore_v1/vector_query.py b/google/cloud/firestore_v1/vector_query.py index a419dba63a..9e2d4ad0f0 100644 --- a/google/cloud/firestore_v1/vector_query.py +++ b/google/cloud/firestore_v1/vector_query.py @@ -14,12 +14,14 @@ """Classes for representing vector queries for the Google Cloud Firestore API. """ +from __future__ import annotations -from typing import TYPE_CHECKING, Any, Generator, Iterable, Optional, TypeVar, Union +from typing import TYPE_CHECKING, Any, Generator, Optional, TypeVar, Union from google.api_core import gapic_v1 from google.api_core import retry as retries +from google.cloud.firestore_v1.query_results import QueryResultsList from google.cloud.firestore_v1.base_query import ( BaseQuery, _collection_group_query_response_to_snapshot, @@ -32,6 +34,8 @@ if TYPE_CHECKING: # pragma: NO COVER from google.cloud.firestore_v1 import transaction from google.cloud.firestore_v1.base_document import DocumentSnapshot + from google.cloud.firestore_v1.query_profile import ExplainMetrics + from google.cloud.firestore_v1.query_profile import ExplainOptions TVectorQuery = TypeVar("TVectorQuery", bound="VectorQuery") @@ -55,7 +59,9 @@ def get( transaction=None, retry: retries.Retry = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> Iterable["DocumentSnapshot"]: + *, + explain_options: Optional[ExplainOptions] = None, + ) -> QueryResultsList[DocumentSnapshot]: """Runs the vector query. This sends a ``RunQuery`` RPC and returns a list of document messages. @@ -71,20 +77,38 @@ def get( should be retried. Defaults to a system-specified policy. timeout (float): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: - list: The vector query results. + QueryResultsList[DocumentSnapshot]: The vector query results. """ - result = self.stream(transaction=transaction, retry=retry, timeout=timeout) + explain_metrics: ExplainMetrics | None = None - return list(result) + result = self.stream( + transaction=transaction, + retry=retry, + timeout=timeout, + explain_options=explain_options, + ) + result_list = list(result) + + if explain_options is None: + explain_metrics = None + else: + explain_metrics = result.get_explain_metrics() + + return QueryResultsList(result_list, explain_options, explain_metrics) - def _get_stream_iterator(self, transaction, retry, timeout): + def _get_stream_iterator(self, transaction, retry, timeout, explain_options=None): """Helper method for :meth:`stream`.""" request, expected_prefix, kwargs = self._prep_stream( transaction, retry, timeout, + explain_options, ) response_iterator = self._client._firestore_api.run_query( @@ -100,7 +124,8 @@ def _make_stream( transaction: Optional["transaction.Transaction"] = None, retry: Optional[retries.Retry] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> Generator["DocumentSnapshot", Any, None]: + explain_options: Optional[ExplainOptions] = None, + ) -> Generator[DocumentSnapshot, Any, Optional[ExplainMetrics]]: """Reads the documents in the collection that match this query. This sends a ``RunQuery`` RPC and then returns a generator which @@ -120,15 +145,26 @@ def _make_stream( system-specified policy. timeout (Optional[float]): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Yields: - :class:`~google.cloud.firestore_v1.document.DocumentSnapshot`: + DocumentSnapshot: The next document that fulfills the query. + + Returns: + ([google.cloud.firestore_v1.types.query_profile.ExplainMetrtics | None]): + The results of query profiling, if received from the service. """ + metrics: ExplainMetrics | None = None + response_iterator, expected_prefix = self._get_stream_iterator( transaction, retry, timeout, + explain_options, ) while True: @@ -137,6 +173,9 @@ def _make_stream( if response is None: # EOI break + if metrics is None and response.explain_metrics: + metrics = response.explain_metrics + if self._nested_query._all_descendants: snapshot = _collection_group_query_response_to_snapshot( response, self._nested_query._parent @@ -148,12 +187,16 @@ def _make_stream( if snapshot is not None: yield snapshot + return metrics + def stream( self, transaction: Optional["transaction.Transaction"] = None, retry: Optional[retries.Retry] = gapic_v1.method.DEFAULT, timeout: Optional[float] = None, - ) -> "StreamGenerator[DocumentSnapshot]": + *, + explain_options: Optional[ExplainOptions] = None, + ) -> StreamGenerator[DocumentSnapshot]: """Reads the documents in the collection that match this query. This sends a ``RunQuery`` RPC and then returns a generator which @@ -173,6 +216,10 @@ def stream( system-specified policy. timeout (Optinal[float]): The timeout for this request. Defaults to a system-specified value. + explain_options + (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): + Options to enable query profiling for this query. When set, + explain_metrics will be available on the returned generator. Returns: `StreamGenerator[DocumentSnapshot]`: A generator of the query results. @@ -181,5 +228,6 @@ def stream( transaction=transaction, retry=retry, timeout=timeout, + explain_options=explain_options, ) - return StreamGenerator(inner_generator) + return StreamGenerator(inner_generator, explain_options) diff --git a/scripts/fixup_firestore_admin_v1_keywords.py b/scripts/fixup_firestore_admin_v1_keywords.py index 9bc2afbcd0..1c2d4ec8d8 100644 --- a/scripts/fixup_firestore_admin_v1_keywords.py +++ b/scripts/fixup_firestore_admin_v1_keywords.py @@ -65,7 +65,7 @@ class firestore_adminCallTransformer(cst.CSTTransformer): 'list_databases': ('parent', 'show_deleted', ), 'list_fields': ('parent', 'filter', 'page_size', 'page_token', ), 'list_indexes': ('parent', 'filter', 'page_size', 'page_token', ), - 'restore_database': ('parent', 'database_id', 'backup', ), + 'restore_database': ('parent', 'database_id', 'backup', 'encryption_config', ), 'update_backup_schedule': ('backup_schedule', 'update_mask', ), 'update_database': ('database', 'update_mask', ), 'update_field': ('field', 'update_mask', ), diff --git a/tests/system/test_system.py b/tests/system/test_system.py index b67b8aecca..ed525db576 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -99,6 +99,124 @@ def test_collections_w_import(database): assert isinstance(collections, list) +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +def test_collection_stream_or_get_w_no_explain_options(database, query_docs, method): + from google.cloud.firestore_v1.query_profile import QueryExplainError + + collection, _, _ = query_docs + + # Tests either `stream()` or `get()`. + method_under_test = getattr(collection, method) + results = method_under_test() + + # verify explain_metrics isn't available + with pytest.raises( + QueryExplainError, + match="explain_options not set on query.", + ): + results.get_explain_metrics() + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["get", "stream"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +def test_collection_stream_or_get_w_explain_options_analyze_false( + database, method, query_docs +): + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExplainOptions, + PlanSummary, + QueryExplainError, + ) + + collection, _, _ = query_docs + + # Tests either `stream()` or `get()`. + method_under_test = getattr(collection, method) + results = method_under_test(explain_options=ExplainOptions(analyze=False)) + + # Verify explain_metrics and plan_summary. + explain_metrics = results.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert plan_summary.indexes_used[0]["properties"] == "(__name__ ASC)" + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats isn't available. + with pytest.raises( + QueryExplainError, + match="execution_stats not available when explain_options.analyze=False", + ): + explain_metrics.execution_stats + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["get", "stream"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +def test_collection_stream_or_get_w_explain_options_analyze_true( + database, method, query_docs +): + from google.cloud.firestore_v1.query_profile import ( + ExecutionStats, + ExplainMetrics, + ExplainOptions, + PlanSummary, + QueryExplainError, + ) + + collection, _, _ = query_docs + + # Tests either `stream()` or `get()`. + method_under_test = getattr(collection, method) + results = method_under_test(explain_options=ExplainOptions(analyze=True)) + + # In the case of `stream()`, an exception should be raised when accessing + # explain_metrics before query finishes. + if method == "stream": + with pytest.raises( + QueryExplainError, + match="explain_metrics not available until query is complete", + ): + results.get_explain_metrics() + + # Finish iterating results, and explain_metrics should be available. + num_results = len(list(results)) + + # Verify explain_metrics and plan_summary. + explain_metrics = results.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert plan_summary.indexes_used[0]["properties"] == "(__name__ ASC)" + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats. + execution_stats = explain_metrics.execution_stats + assert isinstance(execution_stats, ExecutionStats) + assert execution_stats.results_returned == num_results + assert execution_stats.read_operations == num_results + duration = execution_stats.execution_duration.total_seconds() + assert duration > 0 + assert duration < 1 # we expect a number closer to 0.05 + assert isinstance(execution_stats.debug_stats, dict) + assert "billing_details" in execution_stats.debug_stats + assert "documents_scanned" in execution_stats.debug_stats + assert "index_entries_scanned" in execution_stats.debug_stats + assert len(execution_stats.debug_stats) > 0 + + @pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) def test_create_document(client, cleanup, database): now = datetime.datetime.now(tz=datetime.timezone.utc) @@ -414,6 +532,156 @@ def test_vector_search_collection_group_with_distance_parameters_cosine( } +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +def test_vector_query_stream_or_get_w_no_explain_options(client, database, method): + from google.cloud.firestore_v1.query_profile import QueryExplainError + + collection_id = "vector_search" + collection_group = client.collection_group(collection_id) + + vector_query = collection_group.where("color", "==", "red").find_nearest( + vector_field="embedding", + query_vector=Vector([1.0, 2.0, 3.0]), + distance_measure=DistanceMeasure.EUCLIDEAN, + limit=1, + ) + + # Tests either `stream()` or `get()`. + method_under_test = getattr(vector_query, method) + results = method_under_test() + + # verify explain_metrics isn't available + with pytest.raises( + QueryExplainError, + match="explain_options not set on query.", + ): + results.get_explain_metrics() + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +def test_vector_query_stream_or_get_w_explain_options_analyze_true( + client, database, method +): + from google.cloud.firestore_v1.query_profile import ( + ExecutionStats, + ExplainMetrics, + ExplainOptions, + PlanSummary, + QueryExplainError, + ) + + collection_id = "vector_search" + collection_group = client.collection_group(collection_id) + + vector_query = collection_group.where("color", "==", "red").find_nearest( + vector_field="embedding", + query_vector=Vector([1.0, 2.0, 3.0]), + distance_measure=DistanceMeasure.EUCLIDEAN, + limit=1, + ) + + # Tests either `stream()` or `get()`. + method_under_test = getattr(vector_query, method) + results = method_under_test(explain_options=ExplainOptions(analyze=True)) + + # With `stream()`, an exception should be raised when accessing + # explain_metrics before query finishes. + if method == "stream": + with pytest.raises( + QueryExplainError, + match="explain_metrics not available until query is complete", + ): + results.get_explain_metrics() + + # Finish iterating results, and explain_metrics should be available. + num_results = len(list(results)) + + # Verify explain_metrics and plan_summary. + explain_metrics = results.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert ( + plan_summary.indexes_used[0]["properties"] + == "(color ASC, __name__ ASC, embedding VECTOR<3>)" + ) + assert plan_summary.indexes_used[0]["query_scope"] == "Collection group" + + # Verify execution_stats. + execution_stats = explain_metrics.execution_stats + assert isinstance(execution_stats, ExecutionStats) + assert execution_stats.results_returned == num_results + assert execution_stats.read_operations > 0 + duration = execution_stats.execution_duration.total_seconds() + assert duration > 0 + assert duration < 1 # we expect a number closer to 0.05 + assert isinstance(execution_stats.debug_stats, dict) + assert "billing_details" in execution_stats.debug_stats + assert "documents_scanned" in execution_stats.debug_stats + assert "index_entries_scanned" in execution_stats.debug_stats + assert len(execution_stats.debug_stats) > 0 + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +def test_vector_query_stream_or_get_w_explain_options_analyze_false( + client, database, method +): + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExplainOptions, + PlanSummary, + QueryExplainError, + ) + + collection_id = "vector_search" + collection_group = client.collection_group(collection_id) + + vector_query = collection_group.where("color", "==", "red").find_nearest( + vector_field="embedding", + query_vector=Vector([1.0, 2.0, 3.0]), + distance_measure=DistanceMeasure.EUCLIDEAN, + limit=1, + ) + # Tests either `stream()` or `get()`. + method_under_test = getattr(vector_query, method) + results = method_under_test(explain_options=ExplainOptions(analyze=False)) + + results_list = list(results) + assert len(results_list) == 0 + + # Verify explain_metrics and plan_summary. + explain_metrics = results.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert ( + plan_summary.indexes_used[0]["properties"] + == "(color ASC, __name__ ASC, embedding VECTOR<3>)" + ) + assert plan_summary.indexes_used[0]["query_scope"] == "Collection group" + + # Verify execution_stats isn't available. + with pytest.raises( + QueryExplainError, + match="execution_stats not available when explain_options.analyze=False", + ): + explain_metrics.execution_stats + + @pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) def test_create_document_w_subcollection(client, cleanup, database): collection_id = "doc-create-sub" + UNIQUE_RESOURCE_ID @@ -1056,6 +1324,132 @@ def test_query_stream_w_offset(query_docs, database): assert value["b"] == 2 +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +def test_query_stream_or_get_w_no_explain_options(query_docs, database, method): + from google.cloud.firestore_v1.query_profile import QueryExplainError + + collection, _, allowed_vals = query_docs + num_vals = len(allowed_vals) + query = collection.where(filter=FieldFilter("a", "in", [1, num_vals + 100])) + + # Tests either `stream()` or `get()`. + method_under_test = getattr(query, method) + results = method_under_test() + + # If no explain_option is passed, raise an exception if explain_metrics + # is called + with pytest.raises(QueryExplainError, match="explain_options not set on query"): + results.get_explain_metrics() + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +def test_query_stream_or_get_w_explain_options_analyze_true( + query_docs, database, method +): + from google.cloud.firestore_v1.query_profile import ( + ExecutionStats, + ExplainMetrics, + ExplainOptions, + PlanSummary, + QueryExplainError, + ) + + collection, _, allowed_vals = query_docs + num_vals = len(allowed_vals) + query = collection.where(filter=FieldFilter("a", "in", [1, num_vals + 100])) + + # Tests either `stream()` or `get()`. + method_under_test = getattr(query, method) + results = method_under_test(explain_options=ExplainOptions(analyze=True)) + + # With `stream()`, an exception should be raised when accessing + # explain_metrics before query finishes. + if method == "stream": + with pytest.raises( + QueryExplainError, + match="explain_metrics not available until query is complete", + ): + results.get_explain_metrics() + + # Finish iterating results, and explain_metrics should be available. + num_results = len(list(results)) + + # Verify explain_metrics and plan_summary. + explain_metrics = results.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert plan_summary.indexes_used[0]["properties"] == "(a ASC, __name__ ASC)" + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats. + execution_stats = explain_metrics.execution_stats + assert isinstance(execution_stats, ExecutionStats) + assert execution_stats.results_returned == num_results + assert execution_stats.read_operations == num_results + duration = execution_stats.execution_duration.total_seconds() + assert duration > 0 + assert duration < 1 # we expect a number closer to 0.05 + assert isinstance(execution_stats.debug_stats, dict) + assert "billing_details" in execution_stats.debug_stats + assert "documents_scanned" in execution_stats.debug_stats + assert "index_entries_scanned" in execution_stats.debug_stats + assert len(execution_stats.debug_stats) > 0 + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +def test_query_stream_or_get_w_explain_options_analyze_false( + query_docs, database, method +): + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExplainOptions, + PlanSummary, + QueryExplainError, + ) + + collection, _, allowed_vals = query_docs + num_vals = len(allowed_vals) + query = collection.where(filter=FieldFilter("a", "in", [1, num_vals + 100])) + + # Tests either `stream()` or `get()`. + method_under_test = getattr(query, method) + results = method_under_test(explain_options=ExplainOptions(analyze=False)) + + # Verify that no results are returned. + results_list = list(results) + assert len(results_list) == 0 + + # Verify explain_metrics and plan_summary. + explain_metrics = results.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert plan_summary.indexes_used[0]["properties"] == "(a ASC, __name__ ASC)" + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats isn't available. + with pytest.raises( + QueryExplainError, + match="execution_stats not available when explain_options.analyze=False", + ): + explain_metrics.execution_stats + + @pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) def test_query_with_order_dot_key(client, cleanup, database): db = client @@ -2428,6 +2822,140 @@ def test_avg_query_with_start_at(query, database): assert avg_result[0].value == expected_avg +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +def test_aggregation_query_stream_or_get_w_no_explain_options(query, database, method): + # Because all aggregation methods end up calling AggregationQuery.get() or + # AggregationQuery.stream(), only use count() for testing here. + from google.cloud.firestore_v1.query_profile import QueryExplainError + + result = query.get() + start_doc = result[1] + + # start new query that starts at the second result + count_query = query.start_at(start_doc).count("a") + + # Tests either `stream()` or `get()`. + method_under_test = getattr(count_query, method) + results = method_under_test() + + # If no explain_option is passed, raise an exception if explain_metrics + # is called + with pytest.raises(QueryExplainError, match="explain_options not set on query"): + results.get_explain_metrics() + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +def test_aggregation_query_stream_or_get_w_explain_options_analyze_true( + query, database, method +): + # Because all aggregation methods end up calling AggregationQuery.get() or + # AggregationQuery.stream(), only use count() for testing here. + from google.cloud.firestore_v1.query_profile import ( + ExecutionStats, + ExplainMetrics, + ExplainOptions, + PlanSummary, + QueryExplainError, + ) + + result = query.get() + start_doc = result[1] + + # start new query that starts at the second result + count_query = query.start_at(start_doc).count("a") + + # Tests either `stream()` or `get()`. + method_under_test = getattr(count_query, method) + results = method_under_test(explain_options=ExplainOptions(analyze=True)) + + # With `stream()`, an exception should be raised when accessing + # explain_metrics before query finishes. + if method == "stream": + with pytest.raises( + QueryExplainError, + match="explain_metrics not available until query is complete", + ): + results.get_explain_metrics() + + # Finish iterating results, and explain_metrics should be available. + num_results = len(list(results)) + + # Verify explain_metrics and plan_summary. + explain_metrics = results.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert plan_summary.indexes_used[0]["properties"] == "(a ASC, __name__ ASC)" + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats. + execution_stats = explain_metrics.execution_stats + assert isinstance(execution_stats, ExecutionStats) + assert execution_stats.results_returned == num_results + assert execution_stats.read_operations == num_results + duration = execution_stats.execution_duration.total_seconds() + assert duration > 0 + assert duration < 1 # we expect a number closer to 0.05 + assert isinstance(execution_stats.debug_stats, dict) + assert "billing_details" in execution_stats.debug_stats + assert "documents_scanned" in execution_stats.debug_stats + assert "index_entries_scanned" in execution_stats.debug_stats + assert len(execution_stats.debug_stats) > 0 + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +def test_aggregation_query_stream_or_get_w_explain_options_analyze_false( + query, database, method +): + # Because all aggregation methods end up calling AggregationQuery.get() or + # AggregationQuery.stream(), only use count() for testing here. + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExplainOptions, + PlanSummary, + QueryExplainError, + ) + + result = query.get() + start_doc = result[1] + + # start new query that starts at the second result + count_query = query.start_at(start_doc).count("a") + + # Tests either `stream()` or `get()`. + method_under_test = getattr(count_query, method) + results = method_under_test(explain_options=ExplainOptions(analyze=False)) + + # Verify explain_metrics and plan_summary. + explain_metrics = results.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert plan_summary.indexes_used[0]["properties"] == "(a ASC, __name__ ASC)" + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats isn't available. + with pytest.raises( + QueryExplainError, + match="execution_stats not available when explain_options.analyze=False", + ): + explain_metrics.execution_stats + + @pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) def test_query_with_and_composite_filter(collection, database): and_filter = And( @@ -2602,6 +3130,61 @@ def in_transaction(transaction): assert inner_fn_ran is True +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +def test_query_in_transaction_with_explain_options(client, cleanup, database): + """ + Test query profiling in transactions. + """ + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExplainOptions, + QueryExplainError, + ) + + collection_id = "doc-create" + UNIQUE_RESOURCE_ID + doc_ids = [f"doc{i}" + UNIQUE_RESOURCE_ID for i in range(5)] + doc_refs = [client.document(collection_id, doc_id) for doc_id in doc_ids] + for doc_ref in doc_refs: + cleanup(doc_ref.delete) + doc_refs[0].create({"a": 1, "b": 2}) + doc_refs[1].create({"a": 1, "b": 1}) + + collection = client.collection(collection_id) + query = collection.where(filter=FieldFilter("a", "==", 1)) + + with client.transaction() as transaction: + # should work when transaction is initiated through transactional decorator + @firestore.transactional + def in_transaction(transaction): + global inner_fn_ran + + # When no explain_options value is passed, an exception shoud be + # raised when accessing explain_metrics. + result_1 = query.get(transaction=transaction) + with pytest.raises( + QueryExplainError, match="explain_options not set on query." + ): + result_1.get_explain_metrics() + + result_2 = query.get( + transaction=transaction, + explain_options=ExplainOptions(analyze=True), + ) + explain_metrics = result_2.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + assert explain_metrics.plan_summary is not None + assert explain_metrics.execution_stats is not None + + inner_fn_ran = True + + in_transaction(transaction) + # make sure we didn't skip assertions in inner function + assert inner_fn_ran is True + + @pytest.mark.parametrize("with_rollback,expected", [(True, 2), (False, 3)]) @pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) def test_transaction_rollback(client, cleanup, database, with_rollback, expected): diff --git a/tests/system/test_system_async.py b/tests/system/test_system_async.py index 78bd64c5c5..675b23a98a 100644 --- a/tests/system/test_system_async.py +++ b/tests/system/test_system_async.py @@ -36,6 +36,14 @@ from google.cloud import firestore_v1 as firestore from google.cloud.firestore_v1.base_query import And, FieldFilter, Or from google.cloud.firestore_v1.base_vector_query import DistanceMeasure +from google.cloud.firestore_v1.query_profile import ( + ExecutionStats, + ExplainMetrics, + ExplainOptions, + PlanSummary, + QueryExplainError, +) +from google.cloud.firestore_v1.query_results import QueryResultsList from google.cloud.firestore_v1.vector import Vector from test__helpers import ( EMULATOR_CREDS, @@ -78,6 +86,58 @@ def _get_credentials_and_project(): return credentials, project +def _verify_explain_metrics_analyze_true(explain_metrics, num_results): + from google.cloud.firestore_v1.query_profile import ( + ExecutionStats, + ExplainMetrics, + PlanSummary, + ) + + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert plan_summary.indexes_used[0]["properties"] == "(a ASC, __name__ ASC)" + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats. + execution_stats = explain_metrics.execution_stats + assert isinstance(execution_stats, ExecutionStats) + assert execution_stats.results_returned == num_results + assert execution_stats.read_operations == num_results + duration = execution_stats.execution_duration.total_seconds() + assert duration > 0 + assert duration < 1 # we expect a number closer to 0.05 + assert isinstance(execution_stats.debug_stats, dict) + assert "billing_details" in execution_stats.debug_stats + assert "documents_scanned" in execution_stats.debug_stats + assert "index_entries_scanned" in execution_stats.debug_stats + assert len(execution_stats.debug_stats) > 0 + + +def _verify_explain_metrics_analyze_false(explain_metrics): + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + PlanSummary, + QueryExplainError, + ) + + # Verify explain_metrics and plan_summary. + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert plan_summary.indexes_used[0]["properties"] == "(a ASC, __name__ ASC)" + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats isn't available. + with pytest.raises( + QueryExplainError, + match="execution_stats not available when explain_options.analyze=False", + ): + explain_metrics.execution_stats + + @pytest.fixture(scope="session") def database(request): return request.param @@ -359,7 +419,7 @@ async def test_vector_search_collection(client, database, distance_measure): distance_measure=distance_measure, ) returned = await vector_query.get() - assert isinstance(returned, list) + assert isinstance(returned, QueryResultsList) assert len(returned) == 1 assert returned[0].to_dict() == { "embedding": Vector([1.0, 2.0, 3.0]), @@ -387,7 +447,7 @@ async def test_vector_search_collection_with_filter(client, database, distance_m distance_measure=distance_measure, ) returned = await vector_query.get() - assert isinstance(returned, list) + assert isinstance(returned, QueryResultsList) assert len(returned) == 1 assert returned[0].to_dict() == { "embedding": Vector([1.0, 2.0, 3.0]), @@ -537,7 +597,7 @@ async def test_vector_search_collection_group_with_distance_parameters_euclid( distance_threshold=1.0, ) returned = await vector_query.get() - assert isinstance(returned, list) + assert isinstance(returned, QueryResultsList) assert len(returned) == 2 assert returned[0].to_dict() == { "embedding": Vector([1.0, 2.0, 3.0]), @@ -569,7 +629,7 @@ async def test_vector_search_collection_group_with_distance_parameters_cosine( distance_threshold=0.02, ) returned = await vector_query.get() - assert isinstance(returned, list) + assert isinstance(returned, QueryResultsList) assert len(returned) == 2 assert returned[0].to_dict() == { "embedding": Vector([1.0, 2.0, 3.0]), @@ -583,6 +643,186 @@ async def test_vector_search_collection_group_with_distance_parameters_cosine( } +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_vector_query_stream_or_get_w_no_explain_options( + client, database, method +): + from google.cloud.firestore_v1.query_profile import QueryExplainError + + collection_id = "vector_search" + collection_group = client.collection_group(collection_id) + + vector_query = collection_group.where("color", "==", "red").find_nearest( + vector_field="embedding", + query_vector=Vector([1.0, 2.0, 3.0]), + distance_measure=DistanceMeasure.EUCLIDEAN, + limit=1, + ) + + # Tests either `stream()` or `get()`. + method_under_test = getattr(vector_query, method) + if method == "get": + results = await method_under_test() + else: + results = method_under_test() + + # verify explain_metrics isn't available + with pytest.raises( + QueryExplainError, + match="explain_options not set on query.", + ): + await results.get_explain_metrics() + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_vector_query_stream_or_get_w_explain_options_analyze_true( + client, query_docs, database, method +): + from google.cloud.firestore_v1.query_profile import ( + ExecutionStats, + ExplainMetrics, + ExplainOptions, + PlanSummary, + QueryExplainError, + ) + + collection_id = "vector_search" + collection_group = client.collection_group(collection_id) + + vector_query = collection_group.where("color", "==", "red").find_nearest( + vector_field="embedding", + query_vector=Vector([1.0, 2.0, 3.0]), + distance_measure=DistanceMeasure.EUCLIDEAN, + limit=1, + ) + + # Tests either `stream()` or `get()`. + method_under_test = getattr(vector_query, method) + if method == "stream": + results = method_under_test(explain_options=ExplainOptions(analyze=True)) + else: + results = await method_under_test(explain_options=ExplainOptions(analyze=True)) + + # With `stream()`, an exception should be raised when accessing + # explain_metrics before query finishes. + if method == "stream": + with pytest.raises( + QueryExplainError, + match="explain_metrics not available until query is complete", + ): + await results.get_explain_metrics() + + # Finish iterating results, and explain_metrics should be available. + if method == "stream": + results_list = [item async for item in results] + explain_metrics = await results.get_explain_metrics() + else: + results_list = list(results) + explain_metrics = results.get_explain_metrics() + + # Finish iterating results, and explain_metrics should be available. + num_results = len(results_list) + + # Verify explain_metrics and plan_summary. + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert ( + plan_summary.indexes_used[0]["properties"] + == "(color ASC, __name__ ASC, embedding VECTOR<3>)" + ) + assert plan_summary.indexes_used[0]["query_scope"] == "Collection group" + + # Verify execution_stats. + execution_stats = explain_metrics.execution_stats + assert isinstance(execution_stats, ExecutionStats) + assert execution_stats.results_returned == num_results + assert execution_stats.read_operations > 0 + duration = execution_stats.execution_duration.total_seconds() + assert duration > 0 + assert duration < 1 # we expect a number closer to 0.05 + assert isinstance(execution_stats.debug_stats, dict) + assert "billing_details" in execution_stats.debug_stats + assert "documents_scanned" in execution_stats.debug_stats + assert "index_entries_scanned" in execution_stats.debug_stats + assert len(execution_stats.debug_stats) > 0 + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_vector_query_stream_or_get_w_explain_options_analyze_false( + client, query_docs, database, method +): + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExplainOptions, + PlanSummary, + QueryExplainError, + ) + + collection_id = "vector_search" + collection_group = client.collection_group(collection_id) + + vector_query = collection_group.where("color", "==", "red").find_nearest( + vector_field="embedding", + query_vector=Vector([1.0, 2.0, 3.0]), + distance_measure=DistanceMeasure.EUCLIDEAN, + limit=1, + ) + + # Tests either `stream()` or `get()`. + method_under_test = getattr(vector_query, method) + if method == "get": + results = await method_under_test(explain_options=ExplainOptions(analyze=False)) + else: + results = method_under_test(explain_options=ExplainOptions(analyze=False)) + + # Verify that no results are returned. + if method == "stream": + results_list = [item async for item in results] + explain_metrics = await results.get_explain_metrics() + else: + results_list = list(results) + explain_metrics = results.get_explain_metrics() + assert len(results_list) == 0 + + # Finish iterating results, and explain_metrics should be available. + if method == "stream": + explain_metrics = await results.get_explain_metrics() + else: + explain_metrics = results.get_explain_metrics() + + # Verify explain_metrics and plan_summary. + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert ( + plan_summary.indexes_used[0]["properties"] + == "(color ASC, __name__ ASC, embedding VECTOR<3>)" + ) + assert plan_summary.indexes_used[0]["query_scope"] == "Collection group" + + # Verify execution_stats isn't available. + with pytest.raises( + QueryExplainError, + match="execution_stats not available when explain_options.analyze=False", + ): + explain_metrics.execution_stats + + @pytest.mark.skipif(FIRESTORE_EMULATOR, reason="Internal Issue b/137867104") @pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) async def test_update_document(client, cleanup, database): @@ -1040,6 +1280,115 @@ async def test_query_stream_w_offset(query_docs, database): assert value["b"] == 2 +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_query_stream_or_get_w_no_explain_options(query_docs, database, method): + from google.cloud.firestore_v1.query_profile import QueryExplainError + + collection, _, allowed_vals = query_docs + num_vals = len(allowed_vals) + query = collection.where(filter=FieldFilter("a", "in", [1, num_vals + 100])) + + # Tests either `stream()` or `get()`. + method_under_test = getattr(query, method) + if method == "get": + results = await method_under_test() + else: + results = method_under_test() + + # If no explain_option is passed, raise an exception if explain_metrics + # is called + with pytest.raises(QueryExplainError, match="explain_options not set on query"): + await results.get_explain_metrics() + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_query_stream_or_get_w_explain_options_analyze_true( + query_docs, database, method +): + from google.cloud.firestore_v1.query_profile import ( + ExplainOptions, + QueryExplainError, + ) + + collection, _, allowed_vals = query_docs + num_vals = len(allowed_vals) + query = collection.where(filter=FieldFilter("a", "in", [1, num_vals + 100])) + + # Tests either `stream()` or `get()`. + method_under_test = getattr(query, method) + if method == "get": + results = await method_under_test(explain_options=ExplainOptions(analyze=True)) + else: + results = method_under_test(explain_options=ExplainOptions(analyze=True)) + + # With `stream()`, an exception should be raised when accessing + # explain_metrics before query finishes. + if method == "stream": + with pytest.raises( + QueryExplainError, + match="explain_metrics not available until query is complete", + ): + await results.get_explain_metrics() + + # Finish iterating results, and explain_metrics should be available. + if method == "stream": + results_list = [item async for item in results] + explain_metrics = await results.get_explain_metrics() + else: + results_list = list(results) + explain_metrics = results.get_explain_metrics() + + num_results = len(results_list) + _verify_explain_metrics_analyze_true(explain_metrics, num_results) + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_query_stream_or_get_w_explain_options_analyze_false( + query_docs, database, method +): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + collection, _, allowed_vals = query_docs + num_vals = len(allowed_vals) + query = collection.where(filter=FieldFilter("a", "in", [1, num_vals + 100])) + + # Tests either `stream()` or `get()`. + method_under_test = getattr(query, method) + if method == "get": + results = await method_under_test(explain_options=ExplainOptions(analyze=False)) + else: + results = method_under_test(explain_options=ExplainOptions(analyze=False)) + + # Verify that no results are returned. + if method == "stream": + results_list = [item async for item in results] + explain_metrics = await results.get_explain_metrics() + else: + results_list = list(results) + explain_metrics = results.get_explain_metrics() + assert len(results_list) == 0 + + # Finish iterating results, and explain_metrics should be available. + if method == "stream": + explain_metrics = await results.get_explain_metrics() + else: + explain_metrics = results.get_explain_metrics() + + _verify_explain_metrics_analyze_false(explain_metrics) + + @pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) async def test_query_with_order_dot_key(client, cleanup, database): db = client @@ -1205,10 +1554,9 @@ async def test_collection_group_queries_filters(client, cleanup, database): ] batch = client.batch() - - for index, doc_path in enumerate(doc_paths): + for doc_path in doc_paths: doc_ref = client.document(doc_path) - batch.set(doc_ref, {"x": index}) + batch.set(doc_ref, {"x": doc_path}) cleanup(doc_ref.delete) await batch.commit() @@ -1256,6 +1604,154 @@ async def test_collection_group_queries_filters(client, cleanup, database): assert found == set(["cg-doc2"]) +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_collection_stream_or_get_w_no_explain_options( + query_docs, database, method +): + from google.cloud.firestore_v1.query_profile import QueryExplainError + + collection, _, _ = query_docs + + # Tests either `stream()` or `get()`. + method_under_test = getattr(collection, method) + if method == "get": + results = await method_under_test() + else: + results = method_under_test() + + # If no explain_option is passed, raise an exception if explain_metrics + # is called + with pytest.raises(QueryExplainError, match="explain_options not set on query"): + await results.get_explain_metrics() + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_collection_stream_or_get_w_explain_options_analyze_true( + query_docs, database, method +): + from google.cloud.firestore_v1.query_profile import ( + ExplainOptions, + QueryExplainError, + ) + + collection, _, _ = query_docs + + # Tests either `stream()` or `get()`. + method_under_test = getattr(collection, method) + if method == "get": + results = await method_under_test(explain_options=ExplainOptions(analyze=True)) + else: + results = method_under_test(explain_options=ExplainOptions(analyze=True)) + + # With `stream()`, an exception should be raised when accessing + # explain_metrics before query finishes. + if method == "stream": + with pytest.raises( + QueryExplainError, + match="explain_metrics not available until query is complete", + ): + await results.get_explain_metrics() + + # Finish iterating results, and explain_metrics should be available. + if method == "stream": + results_list = [item async for item in results] + explain_metrics = await results.get_explain_metrics() + else: + results_list = list(results) + explain_metrics = results.get_explain_metrics() + + num_results = len(results_list) + from google.cloud.firestore_v1.query_profile import ( + ExecutionStats, + ExplainMetrics, + PlanSummary, + ) + + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert plan_summary.indexes_used[0]["properties"] == "(__name__ ASC)" + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats. + execution_stats = explain_metrics.execution_stats + assert isinstance(execution_stats, ExecutionStats) + assert execution_stats.results_returned == num_results + assert execution_stats.read_operations == num_results + duration = execution_stats.execution_duration.total_seconds() + assert duration > 0 + assert duration < 1 # we expect a number closer to 0.05 + assert isinstance(execution_stats.debug_stats, dict) + assert "billing_details" in execution_stats.debug_stats + assert "documents_scanned" in execution_stats.debug_stats + assert "index_entries_scanned" in execution_stats.debug_stats + assert len(execution_stats.debug_stats) > 0 + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("method", ["stream", "get"]) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_collection_stream_or_get_w_explain_options_analyze_false( + query_docs, database, method +): + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExplainOptions, + PlanSummary, + QueryExplainError, + ) + + collection, _, _ = query_docs + + # Tests either `stream()` or `get()`. + method_under_test = getattr(collection, method) + if method == "get": + results = await method_under_test(explain_options=ExplainOptions(analyze=False)) + else: + results = method_under_test(explain_options=ExplainOptions(analyze=False)) + + # Verify that no results are returned. + if method == "stream": + results_list = [item async for item in results] + explain_metrics = await results.get_explain_metrics() + else: + results_list = list(results) + explain_metrics = results.get_explain_metrics() + assert len(results_list) == 0 + + # Finish iterating results, and explain_metrics should be available. + if method == "stream": + explain_metrics = await results.get_explain_metrics() + else: + explain_metrics = results.get_explain_metrics() + + # Verify explain_metrics and plan_summary. + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert plan_summary.indexes_used[0]["properties"] == "(__name__ ASC)" + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats isn't available. + with pytest.raises( + QueryExplainError, + match="execution_stats not available when explain_options.analyze=False", + ): + explain_metrics.execution_stats + + @pytest.mark.skipif( FIRESTORE_EMULATOR, reason="PartitionQuery not implemented in emulator" ) @@ -2036,6 +2532,83 @@ async def test_async_avg_query_get_multiple_aggregations(collection, database): assert found_alias == set(expected_aliases) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_async_avg_query_get_w_no_explain_options(collection, database): + avg_query = collection.avg("stats.product", alias="total") + results = await avg_query.get() + with pytest.raises(QueryExplainError, match="explain_options not set"): + results.get_explain_metrics() + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_async_avg_query_get_w_explain_options_analyze_true(collection, database): + avg_query = collection.avg("stats.product", alias="total") + results = await avg_query.get(explain_options=ExplainOptions(analyze=True)) + + num_results = len(results) + explain_metrics = results.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert ( + plan_summary.indexes_used[0]["properties"] + == "(stats.product ASC, __name__ ASC)" + ) + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats. + execution_stats = explain_metrics.execution_stats + assert isinstance(execution_stats, ExecutionStats) + assert execution_stats.results_returned == num_results + assert execution_stats.read_operations == num_results + duration = execution_stats.execution_duration.total_seconds() + assert duration > 0 + assert duration < 1 # we expect a number closer to 0.05 + assert isinstance(execution_stats.debug_stats, dict) + assert "billing_details" in execution_stats.debug_stats + assert "documents_scanned" in execution_stats.debug_stats + assert "index_entries_scanned" in execution_stats.debug_stats + assert len(execution_stats.debug_stats) > 0 + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_async_avg_query_get_w_explain_options_analyze_false( + collection, database +): + avg_query = collection.avg("stats.product", alias="total") + results = await avg_query.get(explain_options=ExplainOptions(analyze=False)) + + # Verify that no results are returned. + assert len(results) == 0 + + explain_metrics = results.get_explain_metrics() + + # Verify explain_metrics and plan_summary. + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert ( + plan_summary.indexes_used[0]["properties"] + == "(stats.product ASC, __name__ ASC)" + ) + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats isn't available. + with pytest.raises( + QueryExplainError, + match="execution_stats not available when explain_options.analyze=False", + ): + explain_metrics.execution_stats + + @pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) async def test_async_avg_query_stream_default_alias(collection, database): avg_query = collection.avg("stats.product") @@ -2083,6 +2656,94 @@ async def test_async_avg_query_stream_multiple_aggregations(collection, database assert aggregation_result.alias in ["total", "all"] +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_async_avg_query_stream_w_no_explain_options(collection, database): + avg_query = collection.avg("stats.product", alias="total") + results = avg_query.stream() + with pytest.raises(QueryExplainError, match="explain_options not set"): + await results.get_explain_metrics() + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_async_avg_query_stream_w_explain_options_analyze_true( + collection, database +): + avg_query = collection.avg("stats.product", alias="total") + results = avg_query.stream(explain_options=ExplainOptions(analyze=True)) + with pytest.raises( + QueryExplainError, + match="explain_metrics not available until query is complete", + ): + await results.get_explain_metrics() + + results_list = [item async for item in results] + num_results = len(results_list) + + explain_metrics = await results.get_explain_metrics() + + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert ( + plan_summary.indexes_used[0]["properties"] + == "(stats.product ASC, __name__ ASC)" + ) + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats. + execution_stats = explain_metrics.execution_stats + assert isinstance(execution_stats, ExecutionStats) + assert execution_stats.results_returned == num_results + assert execution_stats.read_operations == num_results + duration = execution_stats.execution_duration.total_seconds() + assert duration > 0 + assert duration < 1 # we expect a number closer to 0.05 + assert isinstance(execution_stats.debug_stats, dict) + assert "billing_details" in execution_stats.debug_stats + assert "documents_scanned" in execution_stats.debug_stats + assert "index_entries_scanned" in execution_stats.debug_stats + assert len(execution_stats.debug_stats) > 0 + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_async_avg_query_stream_w_explain_options_analyze_false( + collection, database +): + avg_query = collection.avg("stats.product", alias="total") + results = avg_query.stream(explain_options=ExplainOptions(analyze=False)) + + # Verify that no results are returned. + results_list = [item async for item in results] + assert len(results_list) == 0 + + explain_metrics = await results.get_explain_metrics() + + # Verify explain_metrics and plan_summary. + assert isinstance(explain_metrics, ExplainMetrics) + plan_summary = explain_metrics.plan_summary + assert isinstance(plan_summary, PlanSummary) + assert len(plan_summary.indexes_used) > 0 + assert ( + plan_summary.indexes_used[0]["properties"] + == "(stats.product ASC, __name__ ASC)" + ) + assert plan_summary.indexes_used[0]["query_scope"] == "Collection" + + # Verify execution_stats isn't available. + with pytest.raises( + QueryExplainError, + match="execution_stats not available when explain_options.analyze=False", + ): + explain_metrics.execution_stats + + @firestore.async_transactional async def create_in_transaction_helper( transaction, client, collection_id, cleanup, database @@ -2265,3 +2926,232 @@ async def test_or_query_in_transaction(client, cleanup, database): assert ( count == 2 ) # assert only 2 results, the third one was rolledback and not created + + +async def _make_transaction_query(client, cleanup): + collection_id = "doc-create" + UNIQUE_RESOURCE_ID + doc_ids = [f"doc{i}" + UNIQUE_RESOURCE_ID for i in range(5)] + doc_refs = [client.document(collection_id, doc_id) for doc_id in doc_ids] + for doc_ref in doc_refs: + cleanup(doc_ref.delete) + await doc_refs[0].create({"a": 1, "b": 2}) + await doc_refs[1].create({"a": 1, "b": 1}) + + collection = client.collection(collection_id) + query = collection.where(filter=FieldFilter("a", "==", 1)) + return query + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_transaction_w_query_w_no_explain_options(client, cleanup, database): + from google.cloud.firestore_v1.query_profile import QueryExplainError + + inner_fn_ran = False + query = await _make_transaction_query(client, cleanup) + transaction = client.transaction() + + # should work when transaction is initiated through transactional decorator + @firestore.async_transactional + async def in_transaction(transaction): + nonlocal inner_fn_ran + + # When no explain_options value is passed, an exception shoud be raised + # when accessing explain_metrics. + returned_generator = await transaction.get(query) + + with pytest.raises( + QueryExplainError, match="explain_options not set on query." + ): + await returned_generator.get_explain_metrics() + + inner_fn_ran = True + + await in_transaction(transaction) + + # make sure we didn't skip assertions in inner function + assert inner_fn_ran is True + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_transaction_w_query_w_explain_options_analyze_true( + client, cleanup, database +): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + inner_fn_ran = False + query = await _make_transaction_query(client, cleanup) + transaction = client.transaction() + + # should work when transaction is initiated through transactional decorator + @firestore.async_transactional + async def in_transaction(transaction): + nonlocal inner_fn_ran + + returned_generator = await transaction.get( + query, + explain_options=ExplainOptions(analyze=True), + ) + + # explain_metrics should not be available before reading all results. + with pytest.raises( + QueryExplainError, + match="explain_metrics not available until query is complete", + ): + await returned_generator.get_explain_metrics() + + result = [x async for x in returned_generator] + explain_metrics = await returned_generator.get_explain_metrics() + _verify_explain_metrics_analyze_true(explain_metrics, len(result)) + + inner_fn_ran = True + + await in_transaction(transaction) + + # make sure we didn't skip assertions in inner function + assert inner_fn_ran is True + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_transaction_w_query_w_explain_options_analyze_false( + client, cleanup, database +): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + inner_fn_ran = False + query = await _make_transaction_query(client, cleanup) + transaction = client.transaction() + + # should work when transaction is initiated through transactional decorator + @firestore.async_transactional + async def in_transaction(transaction): + nonlocal inner_fn_ran + + returned_generator = await transaction.get( + query, + explain_options=ExplainOptions(analyze=False), + ) + explain_metrics = await returned_generator.get_explain_metrics() + _verify_explain_metrics_analyze_false(explain_metrics) + + # When analyze == False, result should be empty. + result = [x async for x in returned_generator] + assert not result + + inner_fn_ran = True + + await in_transaction(transaction) + + # make sure we didn't skip assertions in inner function + assert inner_fn_ran is True + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_query_in_transaction_w_no_explain_options(client, cleanup, database): + from google.cloud.firestore_v1.query_profile import QueryExplainError + + inner_fn_ran = False + query = await _make_transaction_query(client, cleanup) + transaction = client.transaction() + + # should work when transaction is initiated through transactional decorator + @firestore.async_transactional + async def in_transaction(transaction): + nonlocal inner_fn_ran + + # When no explain_options value is passed, an exception shoud be raised + # when accessing explain_metrics. + result = await query.get(transaction=transaction) + + with pytest.raises( + QueryExplainError, match="explain_options not set on query." + ): + result.get_explain_metrics() + + inner_fn_ran = True + + await in_transaction(transaction) + + # make sure we didn't skip assertions in inner function + assert inner_fn_ran is True + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_query_in_transaction_w_explain_options_analyze_true( + client, cleanup, database +): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + inner_fn_ran = False + query = await _make_transaction_query(client, cleanup) + transaction = client.transaction() + + # should work when transaction is initiated through transactional decorator + @firestore.async_transactional + async def in_transaction(transaction): + nonlocal inner_fn_ran + + result = await query.get( + transaction=transaction, + explain_options=ExplainOptions(analyze=True), + ) + + explain_metrics = result.get_explain_metrics() + _verify_explain_metrics_analyze_true(explain_metrics, len(result)) + + inner_fn_ran = True + + await in_transaction(transaction) + + # make sure we didn't skip assertions in inner function + assert inner_fn_ran is True + + +@pytest.mark.skipif( + FIRESTORE_EMULATOR, reason="Query profile not supported in emulator." +) +@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True) +async def test_query_in_transaction_w_explain_options_analyze_false( + client, cleanup, database +): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + inner_fn_ran = False + query = await _make_transaction_query(client, cleanup) + transaction = client.transaction() + + # should work when transaction is initiated through transactional decorator + @firestore.async_transactional + async def in_transaction(transaction): + nonlocal inner_fn_ran + + result = await query.get( + transaction=transaction, + explain_options=ExplainOptions(analyze=False), + ) + explain_metrics = result.get_explain_metrics() + _verify_explain_metrics_analyze_false(explain_metrics) + + # When analyze == False, result should be empty. + assert not result + + inner_fn_ran = True + + await in_transaction(transaction) + + # make sure we didn't skip assertions in inner function + assert inner_fn_ran is True diff --git a/tests/unit/gapic/firestore_admin_v1/test_firestore_admin.py b/tests/unit/gapic/firestore_admin_v1/test_firestore_admin.py index 07d4c09d52..8353d5b180 100644 --- a/tests/unit/gapic/firestore_admin_v1/test_firestore_admin.py +++ b/tests/unit/gapic/firestore_admin_v1/test_firestore_admin.py @@ -5647,6 +5647,7 @@ def test_get_database(request_type, transport: str = "grpc"): app_engine_integration_mode=database.Database.AppEngineIntegrationMode.ENABLED, key_prefix="key_prefix_value", delete_protection_state=database.Database.DeleteProtectionState.DELETE_PROTECTION_DISABLED, + previous_id="previous_id_value", etag="etag_value", ) response = client.get_database(request) @@ -5677,6 +5678,7 @@ def test_get_database(request_type, transport: str = "grpc"): response.delete_protection_state == database.Database.DeleteProtectionState.DELETE_PROTECTION_DISABLED ) + assert response.previous_id == "previous_id_value" assert response.etag == "etag_value" @@ -5785,6 +5787,7 @@ async def test_get_database_empty_call_async(): app_engine_integration_mode=database.Database.AppEngineIntegrationMode.ENABLED, key_prefix="key_prefix_value", delete_protection_state=database.Database.DeleteProtectionState.DELETE_PROTECTION_DISABLED, + previous_id="previous_id_value", etag="etag_value", ) ) @@ -5863,6 +5866,7 @@ async def test_get_database_async( app_engine_integration_mode=database.Database.AppEngineIntegrationMode.ENABLED, key_prefix="key_prefix_value", delete_protection_state=database.Database.DeleteProtectionState.DELETE_PROTECTION_DISABLED, + previous_id="previous_id_value", etag="etag_value", ) ) @@ -5894,6 +5898,7 @@ async def test_get_database_async( response.delete_protection_state == database.Database.DeleteProtectionState.DELETE_PROTECTION_DISABLED ) + assert response.previous_id == "previous_id_value" assert response.etag == "etag_value" @@ -13867,6 +13872,7 @@ def test_create_database_rest(request_type): "uid": "uid_value", "create_time": {"seconds": 751, "nanos": 543}, "update_time": {}, + "delete_time": {}, "location_id": "location_id_value", "type_": 1, "concurrency_mode": 1, @@ -13876,6 +13882,18 @@ def test_create_database_rest(request_type): "app_engine_integration_mode": 1, "key_prefix": "key_prefix_value", "delete_protection_state": 1, + "cmek_config": { + "kms_key_name": "kms_key_name_value", + "active_key_version": [ + "active_key_version_value1", + "active_key_version_value2", + ], + }, + "previous_id": "previous_id_value", + "source_info": { + "backup": {"backup": "backup_value"}, + "operation": "operation_value", + }, "etag": "etag_value", } # The version of a generated dependency at test runtime may differ from the version used during generation. @@ -14286,6 +14304,7 @@ def test_get_database_rest(request_type): app_engine_integration_mode=database.Database.AppEngineIntegrationMode.ENABLED, key_prefix="key_prefix_value", delete_protection_state=database.Database.DeleteProtectionState.DELETE_PROTECTION_DISABLED, + previous_id="previous_id_value", etag="etag_value", ) @@ -14320,6 +14339,7 @@ def test_get_database_rest(request_type): response.delete_protection_state == database.Database.DeleteProtectionState.DELETE_PROTECTION_DISABLED ) + assert response.previous_id == "previous_id_value" assert response.etag == "etag_value" @@ -14905,6 +14925,7 @@ def test_update_database_rest(request_type): "uid": "uid_value", "create_time": {"seconds": 751, "nanos": 543}, "update_time": {}, + "delete_time": {}, "location_id": "location_id_value", "type_": 1, "concurrency_mode": 1, @@ -14914,6 +14935,18 @@ def test_update_database_rest(request_type): "app_engine_integration_mode": 1, "key_prefix": "key_prefix_value", "delete_protection_state": 1, + "cmek_config": { + "kms_key_name": "kms_key_name_value", + "active_key_version": [ + "active_key_version_value1", + "active_key_version_value2", + ], + }, + "previous_id": "previous_id_value", + "source_info": { + "backup": {"backup": "backup_value"}, + "operation": "operation_value", + }, "etag": "etag_value", } # The version of a generated dependency at test runtime may differ from the version used during generation. @@ -19305,8 +19338,34 @@ def test_parse_location_path(): assert expected == actual +def test_operation_path(): + project = "cuttlefish" + database = "mussel" + operation = "winkle" + expected = "projects/{project}/databases/{database}/operations/{operation}".format( + project=project, + database=database, + operation=operation, + ) + actual = FirestoreAdminClient.operation_path(project, database, operation) + assert expected == actual + + +def test_parse_operation_path(): + expected = { + "project": "nautilus", + "database": "scallop", + "operation": "abalone", + } + path = FirestoreAdminClient.operation_path(**expected) + + # Check that the path construction is reversible. + actual = FirestoreAdminClient.parse_operation_path(path) + assert expected == actual + + def test_common_billing_account_path(): - billing_account = "cuttlefish" + billing_account = "squid" expected = "billingAccounts/{billing_account}".format( billing_account=billing_account, ) @@ -19316,7 +19375,7 @@ def test_common_billing_account_path(): def test_parse_common_billing_account_path(): expected = { - "billing_account": "mussel", + "billing_account": "clam", } path = FirestoreAdminClient.common_billing_account_path(**expected) @@ -19326,7 +19385,7 @@ def test_parse_common_billing_account_path(): def test_common_folder_path(): - folder = "winkle" + folder = "whelk" expected = "folders/{folder}".format( folder=folder, ) @@ -19336,7 +19395,7 @@ def test_common_folder_path(): def test_parse_common_folder_path(): expected = { - "folder": "nautilus", + "folder": "octopus", } path = FirestoreAdminClient.common_folder_path(**expected) @@ -19346,7 +19405,7 @@ def test_parse_common_folder_path(): def test_common_organization_path(): - organization = "scallop" + organization = "oyster" expected = "organizations/{organization}".format( organization=organization, ) @@ -19356,7 +19415,7 @@ def test_common_organization_path(): def test_parse_common_organization_path(): expected = { - "organization": "abalone", + "organization": "nudibranch", } path = FirestoreAdminClient.common_organization_path(**expected) @@ -19366,7 +19425,7 @@ def test_parse_common_organization_path(): def test_common_project_path(): - project = "squid" + project = "cuttlefish" expected = "projects/{project}".format( project=project, ) @@ -19376,7 +19435,7 @@ def test_common_project_path(): def test_parse_common_project_path(): expected = { - "project": "clam", + "project": "mussel", } path = FirestoreAdminClient.common_project_path(**expected) @@ -19386,8 +19445,8 @@ def test_parse_common_project_path(): def test_common_location_path(): - project = "whelk" - location = "octopus" + project = "winkle" + location = "nautilus" expected = "projects/{project}/locations/{location}".format( project=project, location=location, @@ -19398,8 +19457,8 @@ def test_common_location_path(): def test_parse_common_location_path(): expected = { - "project": "oyster", - "location": "nudibranch", + "project": "scallop", + "location": "abalone", } path = FirestoreAdminClient.common_location_path(**expected) diff --git a/tests/unit/v1/_test_helpers.py b/tests/unit/v1/_test_helpers.py index 564ec32bc3..39f27ee8c2 100644 --- a/tests/unit/v1/_test_helpers.py +++ b/tests/unit/v1/_test_helpers.py @@ -76,11 +76,20 @@ def make_async_aggregation_query(*args, **kw): return AsyncAggregationQuery(*args, **kw) -def make_aggregation_query_response(aggregations, read_time=None, transaction=None): +def make_aggregation_query_response( + aggregations, + read_time=None, + transaction=None, + explain_metrics=None, +): from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.firestore_v1 import _helpers - from google.cloud.firestore_v1.types import aggregation_result, firestore + from google.cloud.firestore_v1.types import ( + aggregation_result, + firestore, + query_profile, + ) if read_time is None: now = datetime.datetime.now(tz=datetime.timezone.utc) @@ -99,6 +108,9 @@ def make_aggregation_query_response(aggregations, read_time=None, transaction=No if transaction is not None: kwargs["transaction"] = transaction + if explain_metrics is not None: + kwargs["explain_metrics"] = query_profile.ExplainMetrics(explain_metrics) + return firestore.RunAggregationQueryResponse(**kwargs) diff --git a/tests/unit/v1/test_aggregation.py b/tests/unit/v1/test_aggregation.py index 59fe5378c8..4d1eed1980 100644 --- a/tests/unit/v1/test_aggregation.py +++ b/tests/unit/v1/test_aggregation.py @@ -23,6 +23,9 @@ CountAggregation, SumAggregation, ) +from google.cloud.firestore_v1.query_profile import ExplainMetrics, QueryExplainError +from google.cloud.firestore_v1.query_results import QueryResultsList +from google.cloud.firestore_v1.stream_generator import StreamGenerator from tests.unit.v1._test_helpers import ( make_aggregation_query, make_aggregation_query_response, @@ -355,10 +358,45 @@ def test_aggregation_query_prep_stream_with_transaction(): assert kwargs == {"retry": None} -def _aggregation_query_get_helper(retry=None, timeout=None, read_time=None): +def test_aggregation_query_prep_stream_with_explain_options(): + from google.cloud.firestore_v1 import query_profile + + client = make_client() + parent = client.collection("dee") + query = make_query(parent) + aggregation_query = make_aggregation_query(query) + + aggregation_query.count(alias="all") + aggregation_query.sum("someref", alias="sumall") + aggregation_query.avg("anotherref", alias="avgall") + + explain_options = query_profile.ExplainOptions(analyze=True) + request, kwargs = aggregation_query._prep_stream(explain_options=explain_options) + + parent_path, _ = parent._parent_info() + expected_request = { + "parent": parent_path, + "structured_aggregation_query": aggregation_query._to_protobuf(), + "transaction": None, + "explain_options": explain_options._to_dict(), + } + assert request == expected_request + assert kwargs == {"retry": None} + + +def _aggregation_query_get_helper( + retry=None, + timeout=None, + read_time=None, + explain_options=None, +): from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.firestore_v1 import _helpers + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + QueryExplainError, + ) # Create a minimal fake GAPIC. firestore_api = mock.Mock(spec=["run_aggregation_query"]) @@ -375,15 +413,21 @@ def _aggregation_query_get_helper(retry=None, timeout=None, read_time=None): aggregation_result = AggregationResult(alias="total", value=5, read_time=read_time) + if explain_options is not None: + explain_metrics = {"execution_stats": {"results_returned": 1}} + else: + explain_metrics = None response_pb = make_aggregation_query_response( - [aggregation_result], read_time=read_time + [aggregation_result], + read_time=read_time, + explain_metrics=explain_metrics, ) firestore_api.run_aggregation_query.return_value = iter([response_pb]) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) # Execute the query and check the response. - returned = aggregation_query.get(**kwargs) - assert isinstance(returned, list) + returned = aggregation_query.get(**kwargs, explain_options=explain_options) + assert isinstance(returned, QueryResultsList) assert len(returned) == 1 for result in returned: @@ -394,14 +438,29 @@ def _aggregation_query_get_helper(retry=None, timeout=None, read_time=None): result_datetime = _datetime_to_pb_timestamp(r.read_time) assert result_datetime == read_time - # Verify the mock call. + assert returned._explain_options == explain_options + assert returned.explain_options == explain_options + + if explain_options is None: + with pytest.raises(QueryExplainError, match="explain_options not set"): + returned.get_explain_metrics() + else: + actual_explain_metrics = returned.get_explain_metrics() + assert isinstance(actual_explain_metrics, ExplainMetrics) + assert actual_explain_metrics.execution_stats.results_returned == 1 + parent_path, _ = parent._parent_info() + expected_request = { + "parent": parent_path, + "structured_aggregation_query": aggregation_query._to_protobuf(), + "transaction": None, + } + if explain_options is not None: + expected_request["explain_options"] = explain_options._to_dict() + + # Verify the mock call. firestore_api.run_aggregation_query.assert_called_once_with( - request={ - "parent": parent_path, - "structured_aggregation_query": aggregation_query._to_protobuf(), - "transaction": None, - }, + request=expected_request, metadata=client._rpc_metadata, **kwargs, ) @@ -482,6 +541,12 @@ def test_aggregation_query_get_transaction(): ) +def test_aggregation_query_get_w_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + _aggregation_query_get_helper(explain_options=ExplainOptions(analyze=True)) + + _not_passed = object() @@ -604,6 +669,113 @@ def test_aggregation_query_stream_w_retriable_exc_w_transaction(): _aggregation_query_stream_w_retriable_exc_helper(transaction=txn) +def _aggregation_query_stream_helper( + retry=None, + timeout=None, + read_time=None, + explain_options=None, +): + from google.cloud._helpers import _datetime_to_pb_timestamp + + from google.cloud.firestore_v1 import _helpers + + # Create a minimal fake GAPIC. + firestore_api = mock.Mock(spec=["run_aggregation_query"]) + + # Attach the fake GAPIC to a real client. + client = make_client() + client._firestore_api_internal = firestore_api + + # Make a **real** collection reference as parent. + parent = client.collection("dee") + query = make_query(parent) + aggregation_query = make_aggregation_query(query) + aggregation_query.count(alias="all") + + if explain_options is not None and explain_options.analyze is False: + results_list = [] + else: + aggregation_result = AggregationResult( + alias="total", value=5, read_time=read_time + ) + results_list = [aggregation_result] + + if explain_options is not None: + explain_metrics = {"execution_stats": {"results_returned": 1}} + else: + explain_metrics = None + response_pb = make_aggregation_query_response( + results_list, + read_time=read_time, + explain_metrics=explain_metrics, + ) + firestore_api.run_aggregation_query.return_value = iter([response_pb]) + kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) + + # Execute the query and check the response. + returned = aggregation_query.stream(**kwargs, explain_options=explain_options) + assert isinstance(returned, StreamGenerator) + + results = [] + for result in returned: + for r in result: + assert r.alias == aggregation_result.alias + assert r.value == aggregation_result.value + if read_time is not None: + result_datetime = _datetime_to_pb_timestamp(r.read_time) + assert result_datetime == read_time + results.append(result) + assert len(results) == len(results_list) + + if explain_options is None: + with pytest.raises(QueryExplainError, match="explain_options not set"): + returned.get_explain_metrics() + else: + explain_metrics = returned.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + assert explain_metrics.execution_stats.results_returned == 1 + + parent_path, _ = parent._parent_info() + expected_request = { + "parent": parent_path, + "structured_aggregation_query": aggregation_query._to_protobuf(), + "transaction": None, + } + if explain_options is not None: + expected_request["explain_options"] = explain_options._to_dict() + + # Verify the mock call. + firestore_api.run_aggregation_query.assert_called_once_with( + request=expected_request, + metadata=client._rpc_metadata, + **kwargs, + ) + + +def test_aggregation_query_stream(): + _aggregation_query_stream_helper() + + +def test_aggregation_query_stream_with_readtime(): + from google.cloud._helpers import _datetime_to_pb_timestamp + + one_hour_ago = datetime.now(tz=timezone.utc) - timedelta(hours=1) + read_time = _datetime_to_pb_timestamp(one_hour_ago) + _aggregation_query_stream_helper(read_time=read_time) + + +def test_aggregation_query_stream_w_explain_options_analyze_true(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + _aggregation_query_stream_helper(explain_options=ExplainOptions(analyze=True)) + + +def test_aggregation_query_stream_w_explain_options_analyze_false(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + _aggregation_query_stream_helper(explain_options=ExplainOptions(analyze=False)) + + def test_aggregation_from_query(): from google.cloud.firestore_v1 import _helpers diff --git a/tests/unit/v1/test_async_aggregation.py b/tests/unit/v1/test_async_aggregation.py index e51592ae3a..8977d3468b 100644 --- a/tests/unit/v1/test_async_aggregation.py +++ b/tests/unit/v1/test_async_aggregation.py @@ -13,15 +13,7 @@ # limitations under the License. from datetime import datetime, timedelta, timezone - import pytest - -from google.cloud.firestore_v1.base_aggregation import ( - AggregationResult, - AvgAggregation, - CountAggregation, - SumAggregation, -) from tests.unit.v1._test_helpers import ( make_aggregation_query_response, make_async_aggregation_query, @@ -30,6 +22,17 @@ ) from tests.unit.v1.test__helpers import AsyncIter, AsyncMock +from google.cloud.firestore_v1.base_aggregation import ( + AggregationResult, + AvgAggregation, + CountAggregation, + SumAggregation, +) +from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator +from google.cloud.firestore_v1.query_profile import ExplainMetrics, QueryExplainError +from google.cloud.firestore_v1.query_results import QueryResultsList + + _PROJECT = "PROJECT" @@ -292,8 +295,36 @@ def test_async_aggregation_query_prep_stream_with_transaction(): assert kwargs == {"retry": None} +def test_async_aggregation_query_prep_stream_with_explain_options(): + from google.cloud.firestore_v1 import query_profile + + client = make_async_client() + parent = client.collection("dee") + query = make_async_query(parent) + aggregation_query = make_async_aggregation_query(query) + + aggregation_query.count(alias="all") + aggregation_query.sum("someref", alias="sumall") + aggregation_query.avg("anotherref", alias="avgall") + + explain_options = query_profile.ExplainOptions(analyze=True) + request, kwargs = aggregation_query._prep_stream(explain_options=explain_options) + + parent_path, _ = parent._parent_info() + expected_request = { + "parent": parent_path, + "structured_aggregation_query": aggregation_query._to_protobuf(), + "transaction": None, + "explain_options": explain_options._to_dict(), + } + assert request == expected_request + assert kwargs == {"retry": None} + + @pytest.mark.asyncio -async def _async_aggregation_query_get_helper(retry=None, timeout=None, read_time=None): +async def _async_aggregation_query_get_helper( + retry=None, timeout=None, read_time=None, explain_options=None +): from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.firestore_v1 import _helpers @@ -312,15 +343,23 @@ async def _async_aggregation_query_get_helper(retry=None, timeout=None, read_tim aggregation_query.count(alias="all") aggregation_result = AggregationResult(alias="total", value=5, read_time=read_time) + + if explain_options is not None: + explain_metrics = {"execution_stats": {"results_returned": 1}} + else: + explain_metrics = None + response_pb = make_aggregation_query_response( - [aggregation_result], read_time=read_time + [aggregation_result], + read_time=read_time, + explain_metrics=explain_metrics, ) firestore_api.run_aggregation_query.return_value = AsyncIter([response_pb]) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) # Execute the query and check the response. - returned = await aggregation_query.get(**kwargs) - assert isinstance(returned, list) + returned = await aggregation_query.get(**kwargs, explain_options=explain_options) + assert isinstance(returned, QueryResultsList) assert len(returned) == 1 for result in returned: @@ -331,14 +370,25 @@ async def _async_aggregation_query_get_helper(retry=None, timeout=None, read_tim result_datetime = _datetime_to_pb_timestamp(r.read_time) assert result_datetime == read_time + if explain_options is None: + with pytest.raises(QueryExplainError, match="explain_options not set"): + returned.get_explain_metrics() + else: + explain_metrics = returned.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + assert explain_metrics.execution_stats.results_returned == 1 + # Verify the mock call. parent_path, _ = parent._parent_info() + expected_request = { + "parent": parent_path, + "structured_aggregation_query": aggregation_query._to_protobuf(), + "transaction": None, + } + if explain_options is not None: + expected_request["explain_options"] = explain_options._to_dict() firestore_api.run_aggregation_query.assert_called_once_with( - request={ - "parent": parent_path, - "structured_aggregation_query": aggregation_query._to_protobuf(), - "transaction": None, - }, + request=expected_request, metadata=client._rpc_metadata, **kwargs, ) @@ -358,6 +408,14 @@ async def test_async_aggregation_query_get_with_readtime(): await _async_aggregation_query_get_helper(read_time=read_time) +@pytest.mark.asyncio +async def test_async_aggregation_query_get_with_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + await _async_aggregation_query_get_helper(explain_options=explain_options) + + @pytest.mark.asyncio async def test_async_aggregation_query_get_retry_timeout(): from google.api_core.retry import Retry @@ -481,3 +539,102 @@ async def test_async_aggregation_from_query(): metadata=client._rpc_metadata, **kwargs, ) + + +async def _async_aggregation_query_stream_helper( + retry=None, + timeout=None, + read_time=None, + explain_options=None, +): + from google.cloud.firestore_v1 import _helpers + + # Create a minimal fake GAPIC. + firestore_api = AsyncMock(spec=["run_aggregation_query"]) + + # Attach the fake GAPIC to a real client. + client = make_async_client() + client._firestore_api_internal = firestore_api + + # Make a **real** collection reference as parent. + parent = client.collection("dee") + query = make_async_query(parent) + aggregation_query = make_async_aggregation_query(query) + aggregation_query.count(alias="all") + + if explain_options and explain_options.analyze is True: + aggregation_result = AggregationResult( + alias="total", value=5, read_time=read_time + ) + results_list = [aggregation_result] + else: + results_list = [] + + if explain_options is not None: + explain_metrics = {"execution_stats": {"results_returned": 1}} + else: + explain_metrics = None + response_pb = make_aggregation_query_response( + results_list, + read_time=read_time, + explain_metrics=explain_metrics, + ) + firestore_api.run_aggregation_query.return_value = AsyncIter([response_pb]) + kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) + + # Execute the query and check the response. + returned = aggregation_query.stream(**kwargs, explain_options=explain_options) + assert isinstance(returned, AsyncStreamGenerator) + + results = [] + async for result in returned: + for r in result: + assert r.alias == aggregation_result.alias + assert r.value == aggregation_result.value + results.append(result) + assert len(results) == len(results_list) + + if explain_options is None: + with pytest.raises(QueryExplainError, match="explain_options not set"): + await returned.get_explain_metrics() + else: + explain_metrics = await returned.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + assert explain_metrics.execution_stats.results_returned == 1 + + parent_path, _ = parent._parent_info() + expected_request = { + "parent": parent_path, + "structured_aggregation_query": aggregation_query._to_protobuf(), + "transaction": None, + } + if explain_options is not None: + expected_request["explain_options"] = explain_options._to_dict() + + # Verify the mock call. + firestore_api.run_aggregation_query.assert_called_once_with( + request=expected_request, + metadata=client._rpc_metadata, + **kwargs, + ) + + +@pytest.mark.asyncio +async def test_aggregation_query_stream(): + await _async_aggregation_query_stream_helper() + + +@pytest.mark.asyncio +async def test_aggregation_query_stream_w_explain_options_analyze_true(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + await _async_aggregation_query_stream_helper(explain_options=explain_options) + + +@pytest.mark.asyncio +async def test_aggregation_query_stream_w_explain_options_analyze_false(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=False) + await _async_aggregation_query_stream_helper(explain_options=explain_options) diff --git a/tests/unit/v1/test_async_collection.py b/tests/unit/v1/test_async_collection.py index 43884911b4..497fc455fa 100644 --- a/tests/unit/v1/test_async_collection.py +++ b/tests/unit/v1/test_async_collection.py @@ -433,6 +433,23 @@ async def test_asynccollectionreference_get_with_transaction(query_class): query_instance.get.assert_called_once_with(transaction=transaction) +@mock.patch("google.cloud.firestore_v1.async_query.AsyncQuery", autospec=True) +@pytest.mark.asyncio +async def test_asynccollectionreference_get_w_explain_options(query_class): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + + collection = _make_async_collection_reference("collection") + await collection.get(explain_options=ExplainOptions(analyze=True)) + + query_class.assert_called_once_with(collection) + query_instance = query_class.return_value + query_instance.get.assert_called_once_with( + transaction=None, explain_options=explain_options + ) + + @mock.patch("google.cloud.firestore_v1.async_query.AsyncQuery", autospec=True) @pytest.mark.asyncio async def test_asynccollectionreference_stream(query_class): @@ -490,6 +507,51 @@ async def test_asynccollectionreference_stream_with_transaction(query_class): query_instance.stream.assert_called_once_with(transaction=transaction) +@mock.patch("google.cloud.firestore_v1.async_query.AsyncQuery", autospec=True) +@pytest.mark.asyncio +async def test_asynccollectionreference_stream_w_explain_options(query_class): + from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExplainOptions, + QueryExplainError, + ) + import google.cloud.firestore_v1.types.query_profile as query_profile_pb2 + + explain_options = ExplainOptions(analyze=True) + explain_metrics = query_profile_pb2.ExplainMetrics( + {"execution_stats": {"results_returned": 1}} + ) + + async def response_generator(): + yield 1 + yield explain_metrics + + query_class.return_value.stream.return_value = AsyncStreamGenerator( + response_generator(), explain_options + ) + + collection = _make_async_collection_reference("collection") + stream_response = collection.stream(explain_options=ExplainOptions(analyze=True)) + assert isinstance(stream_response, AsyncStreamGenerator) + + with pytest.raises(QueryExplainError, match="explain_metrics not available"): + await stream_response.get_explain_metrics() + + async for _ in stream_response: + pass + + query_class.assert_called_once_with(collection) + query_instance = query_class.return_value + query_instance.stream.assert_called_once_with( + transaction=None, explain_options=explain_options + ) + + explain_metrics = await stream_response.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + assert explain_metrics.execution_stats.results_returned == 1 + + def test_asynccollectionreference_recursive(): from google.cloud.firestore_v1.async_query import AsyncQuery diff --git a/tests/unit/v1/test_async_query.py b/tests/unit/v1/test_async_query.py index cacf0220b1..6af09ec13e 100644 --- a/tests/unit/v1/test_async_query.py +++ b/tests/unit/v1/test_async_query.py @@ -17,6 +17,8 @@ import mock import pytest +from google.cloud.firestore_v1.query_profile import ExplainMetrics, QueryExplainError +from google.cloud.firestore_v1.query_results import QueryResultsList from tests.unit.v1._test_helpers import ( DEFAULT_TEST_PROJECT, make_async_client, @@ -39,7 +41,7 @@ def test_asyncquery_constructor(): assert not query._all_descendants -async def _get_helper(retry=None, timeout=None): +async def _get_helper(retry=None, timeout=None, explain_options=None): from google.cloud.firestore_v1 import _helpers # Create a minimal fake GAPIC. @@ -56,30 +58,46 @@ async def _get_helper(retry=None, timeout=None): _, expected_prefix = parent._parent_info() name = "{}/sleep".format(expected_prefix) data = {"snooze": 10} + explain_metrics = {"execution_stats": {"results_returned": 1}} - response_pb = _make_query_response(name=name, data=data) + response_pb = _make_query_response( + name=name, data=data, explain_metrics=explain_metrics + ) firestore_api.run_query.return_value = AsyncIter([response_pb]) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) # Execute the query and check the response. query = make_async_query(parent) - returned = await query.get(**kwargs) + returned = await query.get(**kwargs, explain_options=explain_options) - assert isinstance(returned, list) + assert isinstance(returned, QueryResultsList) assert len(returned) == 1 snapshot = returned[0] assert snapshot.reference._path == ("dee", "sleep") assert snapshot.to_dict() == data - # Verify the mock call. + if explain_options is None: + with pytest.raises(QueryExplainError, match="explain_options not set"): + returned.get_explain_metrics() + else: + actual_explain_metrics = returned.get_explain_metrics() + assert isinstance(actual_explain_metrics, ExplainMetrics) + assert actual_explain_metrics.execution_stats.results_returned == 1 + + # Create expected request body. parent_path, _ = parent._parent_info() + request = { + "parent": parent_path, + "structured_query": query._to_protobuf(), + "transaction": None, + } + if explain_options: + request["explain_options"] = explain_options._to_dict() + + # Verify the mock call. firestore_api.run_query.assert_called_once_with( - request={ - "parent": parent_path, - "structured_query": query._to_protobuf(), - "transaction": None, - }, + request=request, metadata=client._rpc_metadata, **kwargs, ) @@ -158,6 +176,14 @@ async def test_asyncquery_get_limit_to_last(): ) +@pytest.mark.asyncio +async def test_asyncquery_get_w_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + await _get_helper(explain_options=explain_options) + + def test_asyncquery_sum(): from google.cloud.firestore_v1.base_aggregation import SumAggregation from google.cloud.firestore_v1.field_path import FieldPath @@ -310,7 +336,7 @@ async def test_asyncquery_chunkify_w_chunksize_gt_limit(): assert [snapshot.id for snapshot in chunks[0]] == expected_ids -async def _stream_helper(retry=None, timeout=None): +async def _stream_helper(retry=None, timeout=None, explain_options=None): from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator @@ -328,30 +354,50 @@ async def _stream_helper(retry=None, timeout=None): _, expected_prefix = parent._parent_info() name = "{}/sleep".format(expected_prefix) data = {"snooze": 10} - response_pb = _make_query_response(name=name, data=data) + if explain_options is not None: + explain_metrics = {"execution_stats": {"results_returned": 1}} + else: + explain_metrics = None + response_pb = _make_query_response( + name=name, data=data, explain_metrics=explain_metrics + ) firestore_api.run_query.return_value = AsyncIter([response_pb]) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) # Execute the query and check the response. query = make_async_query(parent) - get_response = query.stream(**kwargs) + stream_response = query.stream(**kwargs, explain_options=explain_options) + assert isinstance(stream_response, AsyncStreamGenerator) - assert isinstance(get_response, AsyncStreamGenerator) - returned = [x async for x in get_response] + returned = [x async for x in stream_response] assert len(returned) == 1 snapshot = returned[0] assert snapshot.reference._path == ("dee", "sleep") assert snapshot.to_dict() == data - # Verify the mock call. + # Verify explain_metrics. + if explain_options is None: + with pytest.raises(QueryExplainError, match="explain_options not set"): + await stream_response.get_explain_metrics() + else: + explain_metrics = await stream_response.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + assert explain_metrics.execution_stats.results_returned == 1 + + # Create expected request body. parent_path, _ = parent._parent_info() + request = { + "parent": parent_path, + "structured_query": query._to_protobuf(), + "transaction": None, + } + if explain_options is not None: + request["explain_options"] = explain_options._to_dict() + + # Verify the mock call. firestore_api.run_query.assert_called_once_with( - request={ - "parent": parent_path, - "structured_query": query._to_protobuf(), - "transaction": None, - }, + request=request, metadata=client._rpc_metadata, **kwargs, ) @@ -638,6 +684,14 @@ async def test_asyncquery_stream_w_collection_group(): ) +@pytest.mark.asyncio +async def test_asyncquery_stream_w_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + await _stream_helper(explain_options=explain_options) + + def _make_async_collection_group(*args, **kwargs): from google.cloud.firestore_v1.async_query import AsyncCollectionGroup diff --git a/tests/unit/v1/test_async_stream_generator.py b/tests/unit/v1/test_async_stream_generator.py index c2e7507b5d..5aa51bc4d1 100644 --- a/tests/unit/v1/test_async_stream_generator.py +++ b/tests/unit/v1/test_async_stream_generator.py @@ -14,8 +14,10 @@ import pytest +from google.protobuf import struct_pb2 -def _make_async_stream_generator(iterable): + +def _make_async_stream_generator(iterable, explain_options=None): from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator async def _inner_generator(): @@ -23,8 +25,9 @@ async def _inner_generator(): X = yield i if X: yield X + # return explain_metrics - return AsyncStreamGenerator(_inner_generator()) + return AsyncStreamGenerator(_inner_generator(), explain_options) @pytest.mark.asyncio @@ -84,7 +87,7 @@ async def test_async_stream_generator_athrow(): @pytest.mark.asyncio -async def test_stream_generator_aclose(): +async def test_async_stream_generator_aclose(): expected_results = [0, 1] inst = _make_async_stream_generator(expected_results) @@ -93,3 +96,162 @@ async def test_stream_generator_aclose(): # Verifies that generator is closed. with pytest.raises(StopAsyncIteration): await inst.__anext__() + + +def test_async_stream_generator_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + inst = _make_async_stream_generator([], explain_options) + assert inst.explain_options == explain_options + + +@pytest.mark.asyncio +async def test_async_stream_generator_explain_metrics_explain_options_analyze_true(): + from google.protobuf import duration_pb2 + from google.protobuf import struct_pb2 + + import google.cloud.firestore_v1.query_profile as query_profile + import google.cloud.firestore_v1.types.query_profile as query_profile_pb2 + + indexes_used_dict = { + "indexes_used": struct_pb2.Value( + struct_value=struct_pb2.Struct( + fields={ + "query_scope": struct_pb2.Value(string_value="Collection"), + "properties": struct_pb2.Value( + string_value="(foo ASC, **name** ASC)" + ), + } + ) + ) + } + plan_summary = query_profile_pb2.PlanSummary() + plan_summary.indexes_used.append(indexes_used_dict) + execution_stats = query_profile_pb2.ExecutionStats( + { + "results_returned": 1, + "execution_duration": duration_pb2.Duration(seconds=2), + "read_operations": 3, + "debug_stats": struct_pb2.Struct( + fields={ + "billing_details": struct_pb2.Value( + string_value="billing_details_results" + ), + "documents_scanned": struct_pb2.Value( + string_value="documents_scanned_results" + ), + "index_entries_scanned": struct_pb2.Value( + string_value="index_entries_scanned" + ), + } + ), + } + ) + + explain_options = query_profile.ExplainOptions(analyze=True) + expected_explain_metrics = query_profile_pb2.ExplainMetrics( + plan_summary=plan_summary, + execution_stats=execution_stats, + ) + iterator = [1, 2, expected_explain_metrics] + + inst = _make_async_stream_generator(iterator, explain_options) + + # Raise an exception if query isn't complete when explain_metrics is called. + with pytest.raises( + query_profile.QueryExplainError, + match="explain_metrics not available until query is complete.", + ): + await inst.get_explain_metrics() + + results = [doc async for doc in inst] + assert len(results) == 2 + + actual_explain_metrics = await inst.get_explain_metrics() + assert isinstance(actual_explain_metrics, query_profile._ExplainAnalyzeMetrics) + assert actual_explain_metrics == query_profile.ExplainMetrics._from_pb( + expected_explain_metrics + ) + assert actual_explain_metrics.plan_summary.indexes_used == [ + { + "indexes_used": { + "query_scope": "Collection", + "properties": "(foo ASC, **name** ASC)", + } + } + ] + assert actual_explain_metrics.execution_stats.results_returned == 1 + duration = actual_explain_metrics.execution_stats.execution_duration.total_seconds() + assert duration == 2 + assert actual_explain_metrics.execution_stats.read_operations == 3 + + expected_debug_stats = { + "billing_details": "billing_details_results", + "documents_scanned": "documents_scanned_results", + "index_entries_scanned": "index_entries_scanned", + } + assert actual_explain_metrics.execution_stats.debug_stats == expected_debug_stats + + +@pytest.mark.asyncio +async def test_async_stream_generator_explain_metrics_explain_options_analyze_false(): + import google.cloud.firestore_v1.query_profile as query_profile + import google.cloud.firestore_v1.types.query_profile as query_profile_pb2 + + explain_options = query_profile.ExplainOptions(analyze=False) + indexes_used_dict = { + "indexes_used": struct_pb2.Value( + struct_value=struct_pb2.Struct( + fields={ + "query_scope": struct_pb2.Value(string_value="Collection"), + "properties": struct_pb2.Value( + string_value="(foo ASC, **name** ASC)" + ), + } + ) + ) + } + plan_summary = query_profile_pb2.PlanSummary() + plan_summary.indexes_used.append(indexes_used_dict) + expected_explain_metrics = query_profile_pb2.ExplainMetrics( + plan_summary=plan_summary + ) + iterator = [expected_explain_metrics] + + inst = _make_async_stream_generator(iterator, explain_options) + actual_explain_metrics = await inst.get_explain_metrics() + assert isinstance(actual_explain_metrics, query_profile.ExplainMetrics) + assert actual_explain_metrics.plan_summary.indexes_used == [ + { + "indexes_used": { + "query_scope": "Collection", + "properties": "(foo ASC, **name** ASC)", + } + } + ] + + +@pytest.mark.asyncio +async def test_async_stream_generator_explain_metrics_missing_explain_options_analyze_false(): + import google.cloud.firestore_v1.query_profile as query_profile + + explain_options = query_profile.ExplainOptions(analyze=False) + inst = _make_async_stream_generator([("1", None)], explain_options) + with pytest.raises( + query_profile.QueryExplainError, match="Did not receive explain_metrics" + ): + await inst.get_explain_metrics() + + +@pytest.mark.asyncio +async def test_stream_generator_explain_metrics_no_explain_options(): + from google.cloud.firestore_v1.query_profile import QueryExplainError + + inst = _make_async_stream_generator([]) + + with pytest.raises( + QueryExplainError, + match="explain_options not set on query.", + ): + await inst.get_explain_metrics() diff --git a/tests/unit/v1/test_async_transaction.py b/tests/unit/v1/test_async_transaction.py index 85d693950e..766c0637e4 100644 --- a/tests/unit/v1/test_async_transaction.py +++ b/tests/unit/v1/test_async_transaction.py @@ -15,7 +15,9 @@ import mock import pytest -from tests.unit.v1.test__helpers import AsyncMock +from tests.unit.v1._test_helpers import make_async_client +from tests.unit.v1.test__helpers import AsyncIter, AsyncMock +from tests.unit.v1.test_base_query import _make_query_response def _make_async_transaction(*args, **kwargs): @@ -314,7 +316,7 @@ async def test_asynctransaction_get_all_w_retry_timeout(): await _get_all_helper(retry=retry, timeout=timeout) -async def _get_w_document_ref_helper(retry=None, timeout=None): +async def _get_w_document_ref_helper(retry=None, timeout=None, explain_options=None): from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1.async_document import AsyncDocumentReference @@ -323,7 +325,7 @@ async def _get_w_document_ref_helper(retry=None, timeout=None): ref = AsyncDocumentReference("documents", "doc-id") kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) - result = await transaction.get(ref, **kwargs) + result = await transaction.get(ref, **kwargs, explain_options=explain_options) client.get_all.assert_called_once_with([ref], transaction=transaction, **kwargs) assert result is client.get_all.return_value @@ -343,26 +345,93 @@ async def test_asynctransaction_get_w_document_ref_w_retry_timeout(): await _get_w_document_ref_helper(retry=retry, timeout=timeout) -async def _get_w_query_helper(retry=None, timeout=None): +@pytest.mark.asyncio +async def test_transaction_get_w_document_ref_w_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + with pytest.raises(ValueError, match="`explain_options` cannot be provided."): + await _get_w_document_ref_helper( + explain_options=ExplainOptions(analyze=True), + ) + + +async def _get_w_query_helper(retry=None, timeout=None, explain_options=None): from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1.async_query import AsyncQuery + from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + QueryExplainError, + ) - client = AsyncMock(spec=[]) - transaction = _make_async_transaction(client) - query = AsyncQuery(parent=AsyncMock(spec=[])) - query.stream = AsyncMock() + # Create a minimal fake GAPIC. + firestore_api = AsyncMock(spec=["run_query"]) + + # Attach the fake GAPIC to a real client. + client = make_async_client() + client._firestore_api_internal = firestore_api + + # Make a **real** collection reference as parent. + parent = client.collection("dee") + + # Add a dummy response to the minimal fake GAPIC. + _, expected_prefix = parent._parent_info() + name = "{}/sleep".format(expected_prefix) + data = {"snooze": 10} + if explain_options is not None: + explain_metrics = {"execution_stats": {"results_returned": 1}} + else: + explain_metrics = None + response_pb = _make_query_response( + name=name, data=data, explain_metrics=explain_metrics + ) + firestore_api.run_query.return_value = AsyncIter([response_pb]) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) - result = await transaction.get( + # Run the transaction with query. + transaction = _make_async_transaction(client) + txn_id = b"beep-fail-commit" + transaction._id = txn_id + query = AsyncQuery(parent) + returned_generator = await transaction.get( query, **kwargs, + explain_options=explain_options, ) - query.stream.assert_called_once_with( - transaction=transaction, + # Verify the response. + assert isinstance(returned_generator, AsyncStreamGenerator) + results = [x async for x in returned_generator] + assert len(results) == 1 + snapshot = results[0] + assert snapshot.reference._path == ("dee", "sleep") + assert snapshot.to_dict() == data + + # Verify explain_metrics. + if explain_options is None: + with pytest.raises(QueryExplainError, match="explain_options not set"): + await returned_generator.get_explain_metrics() + else: + explain_metrics = await returned_generator.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + assert explain_metrics.execution_stats.results_returned == 1 + + # Create expected request body. + parent_path, _ = parent._parent_info() + request = { + "parent": parent_path, + "structured_query": query._to_protobuf(), + "transaction": b"beep-fail-commit", + } + if explain_options is not None: + request["explain_options"] = explain_options._to_dict() + + # Verify the mock call. + firestore_api.run_query.assert_called_once_with( + request=request, + metadata=client._rpc_metadata, **kwargs, ) - assert result is query.stream.return_value @pytest.mark.asyncio @@ -375,6 +444,13 @@ async def test_asynctransaction_get_w_query_w_retry_timeout(): await _get_w_query_helper() +@pytest.mark.asyncio +async def test_transaction_get_w_query_w_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + await _get_w_query_helper(explain_options=ExplainOptions(analyze=True)) + + @pytest.mark.asyncio async def test_asynctransaction_get_failure(): client = _make_client() diff --git a/tests/unit/v1/test_async_vector_query.py b/tests/unit/v1/test_async_vector_query.py index 390190b534..01cded2cc6 100644 --- a/tests/unit/v1/test_async_vector_query.py +++ b/tests/unit/v1/test_async_vector_query.py @@ -15,7 +15,14 @@ import pytest from google.cloud.firestore_v1._helpers import encode_value, make_retry_timeout_kwargs +from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator from google.cloud.firestore_v1.base_vector_query import DistanceMeasure +from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExplainOptions, + QueryExplainError, +) +from google.cloud.firestore_v1.query_results import QueryResultsList from google.cloud.firestore_v1.types.query import StructuredQuery from google.cloud.firestore_v1.vector import Vector from tests.unit.v1._test_helpers import ( @@ -60,6 +67,81 @@ def _expected_pb( return expected_pb +async def _async_vector_query_get_helper( + distance_measure, + expected_distance, + explain_options=None, +): + # Create a minimal fake GAPIC. + firestore_api = AsyncMock(spec=["run_query"]) + client = make_async_client() + client._firestore_api_internal = firestore_api + + # Make a **real** collection reference as parent. + parent = client.collection("dee") + parent_path, expected_prefix = parent._parent_info() + + data = {"snooze": 10, "embedding": Vector([1.0, 2.0, 3.0])} + if explain_options: + explain_metrics = {"execution_stats": {"results_returned": 1}} + else: + explain_metrics = None + response_pb1 = _make_query_response( + name="{}/test_doc".format(expected_prefix), + data=data, + explain_metrics=explain_metrics, + ) + + kwargs = make_retry_timeout_kwargs(retry=None, timeout=None) + + # Execute the vector query and check the response. + firestore_api.run_query.return_value = AsyncIter([response_pb1]) + + vector_async_query = parent.find_nearest( + vector_field="embedding", + query_vector=Vector([1.0, 2.0, 3.0]), + distance_measure=distance_measure, + limit=5, + ) + + returned = await vector_async_query.get( + transaction=_transaction(client), explain_options=explain_options, **kwargs + ) + assert isinstance(returned, QueryResultsList) + assert len(returned) == 1 + assert returned[0].to_dict() == data + + if explain_options is None: + with pytest.raises(QueryExplainError, match="explain_options not set"): + returned.get_explain_metrics() + else: + actual_explain_metrics = returned.get_explain_metrics() + assert isinstance(actual_explain_metrics, ExplainMetrics) + assert actual_explain_metrics.execution_stats.results_returned == 1 + + # Create expected request body. + expected_pb = _expected_pb( + parent=parent, + vector_field="embedding", + vector=Vector([1.0, 2.0, 3.0]), + distance_type=expected_distance, + limit=5, + ) + request = { + "parent": parent_path, + "structured_query": expected_pb, + "transaction": _TXN_ID, + } + if explain_options: + request["explain_options"] = explain_options._to_dict() + + firestore_api.run_query.assert_called_once_with( + request=request, + metadata=client._rpc_metadata, + **kwargs, + ) + + def test_async_vector_query_int_threshold_constructor_to_pb(): client = make_async_client() parent = client.collection("dee") @@ -103,54 +185,31 @@ def test_async_vector_query_int_threshold_constructor_to_pb(): ], ) @pytest.mark.asyncio -async def test_async_vector_query(distance_measure, expected_distance): - # Create a minimal fake GAPIC. - firestore_api = AsyncMock(spec=["run_query"]) - client = make_async_client() - client._firestore_api_internal = firestore_api - - # Make a **real** collection reference as parent. - parent = client.collection("dee") - parent_path, expected_prefix = parent._parent_info() - - data = {"snooze": 10, "embedding": Vector([1.0, 2.0, 3.0])} - response_pb1 = _make_query_response( - name="{}/test_doc".format(expected_prefix), data=data - ) - - kwargs = make_retry_timeout_kwargs(retry=None, timeout=None) - - # Execute the vector query and check the response. - firestore_api.run_query.return_value = AsyncIter([response_pb1]) +async def test_async_vector_query_get(distance_measure, expected_distance): + await _async_vector_query_get_helper(distance_measure, expected_distance) - vector_async_query = parent.find_nearest( - vector_field="embedding", - query_vector=Vector([1.0, 2.0, 3.0]), - distance_measure=distance_measure, - limit=5, - ) - returned = await vector_async_query.get(transaction=_transaction(client), **kwargs) - assert isinstance(returned, list) - assert len(returned) == 1 - assert returned[0].to_dict() == data - - expected_pb = _expected_pb( - parent=parent, - vector_field="embedding", - vector=Vector([1.0, 2.0, 3.0]), - distance_type=expected_distance, - limit=5, - ) - - firestore_api.run_query.assert_called_once_with( - request={ - "parent": parent_path, - "structured_query": expected_pb, - "transaction": _TXN_ID, - }, - metadata=client._rpc_metadata, - **kwargs, +@pytest.mark.parametrize( + "distance_measure, expected_distance", + [ + ( + DistanceMeasure.EUCLIDEAN, + StructuredQuery.FindNearest.DistanceMeasure.EUCLIDEAN, + ), + (DistanceMeasure.COSINE, StructuredQuery.FindNearest.DistanceMeasure.COSINE), + ( + DistanceMeasure.DOT_PRODUCT, + StructuredQuery.FindNearest.DistanceMeasure.DOT_PRODUCT, + ), + ], +) +@pytest.mark.asyncio +async def test_async_vector_query_get_w_explain_options( + distance_measure, expected_distance +): + explain_options = ExplainOptions(analyze=True) + await _async_vector_query_get_helper( + distance_measure, expected_distance, explain_options ) @@ -491,3 +550,123 @@ async def test_async_query_stream_multiple_empty_response_in_stream(): }, metadata=client._rpc_metadata, ) + + +async def _async_vector_query_stream_helper( + distance_measure, + expected_distance, + explain_options=None, +): + # Create a minimal fake GAPIC. + firestore_api = AsyncMock(spec=["run_query"]) + client = make_async_client() + client._firestore_api_internal = firestore_api + + # Make a **real** collection reference as parent. + parent = client.collection("dee") + parent_path, expected_prefix = parent._parent_info() + + data = {"snooze": 10, "embedding": Vector([1.0, 2.0, 3.0])} + if explain_options: + explain_metrics = {"execution_stats": {"results_returned": 1}} + else: + explain_metrics = None + response_pb1 = _make_query_response( + name="{}/test_doc".format(expected_prefix), + data=data, + explain_metrics=explain_metrics, + ) + + kwargs = make_retry_timeout_kwargs(retry=None, timeout=None) + + # Execute the vector query and check the response. + firestore_api.run_query.return_value = AsyncIter([response_pb1]) + + vector_async_query = parent.find_nearest( + vector_field="embedding", + query_vector=Vector([1.0, 2.0, 3.0]), + distance_measure=distance_measure, + limit=5, + ) + + returned = vector_async_query.stream( + transaction=_transaction(client), explain_options=explain_options, **kwargs + ) + assert isinstance(returned, AsyncStreamGenerator) + + results_list = [item async for item in returned] + assert len(results_list) == 1 + assert results_list[0].to_dict() == data + + if explain_options is None: + with pytest.raises(QueryExplainError, match="explain_options not set"): + await returned.get_explain_metrics() + else: + actual_explain_metrics = await returned.get_explain_metrics() + assert isinstance(actual_explain_metrics, ExplainMetrics) + assert actual_explain_metrics.execution_stats.results_returned == 1 + + # Create expected request body. + expected_pb = _expected_pb( + parent=parent, + vector_field="embedding", + vector=Vector([1.0, 2.0, 3.0]), + distance_type=expected_distance, + limit=5, + ) + request = { + "parent": parent_path, + "structured_query": expected_pb, + "transaction": _TXN_ID, + } + if explain_options: + request["explain_options"] = explain_options._to_dict() + + firestore_api.run_query.assert_called_once_with( + request=request, + metadata=client._rpc_metadata, + **kwargs, + ) + + +@pytest.mark.parametrize( + "distance_measure, expected_distance", + [ + ( + DistanceMeasure.EUCLIDEAN, + StructuredQuery.FindNearest.DistanceMeasure.EUCLIDEAN, + ), + (DistanceMeasure.COSINE, StructuredQuery.FindNearest.DistanceMeasure.COSINE), + ( + DistanceMeasure.DOT_PRODUCT, + StructuredQuery.FindNearest.DistanceMeasure.DOT_PRODUCT, + ), + ], +) +@pytest.mark.asyncio +async def test_async_vector_query_stream(distance_measure, expected_distance): + await _async_vector_query_stream_helper(distance_measure, expected_distance) + + +@pytest.mark.parametrize( + "distance_measure, expected_distance", + [ + ( + DistanceMeasure.EUCLIDEAN, + StructuredQuery.FindNearest.DistanceMeasure.EUCLIDEAN, + ), + (DistanceMeasure.COSINE, StructuredQuery.FindNearest.DistanceMeasure.COSINE), + ( + DistanceMeasure.DOT_PRODUCT, + StructuredQuery.FindNearest.DistanceMeasure.DOT_PRODUCT, + ), + ], +) +@pytest.mark.asyncio +async def test_async_vector_query_stream_w_explain_options( + distance_measure, expected_distance +): + explain_options = ExplainOptions(analyze=True) + await _async_vector_query_stream_helper( + distance_measure, expected_distance, explain_options + ) diff --git a/tests/unit/v1/test_base_document.py b/tests/unit/v1/test_base_document.py index 8098afd76a..b2dff117cd 100644 --- a/tests/unit/v1/test_base_document.py +++ b/tests/unit/v1/test_base_document.py @@ -362,6 +362,92 @@ def test_documentsnapshot_non_existent(): assert as_dict is None +def _make_query_results_list(*args, **kwargs): + from google.cloud.firestore_v1.query_results import QueryResultsList + + return QueryResultsList(*args, **kwargs) + + +def _make_explain_metrics(): + from google.cloud.firestore_v1.query_profile import ExplainMetrics, PlanSummary + + plan_summary = PlanSummary( + indexes_used=[{"properties": "(__name__ ASC)", "query_scope": "Collection"}], + ) + return ExplainMetrics(plan_summary=plan_summary) + + +def test_query_results_list_constructor(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + client = mock.sentinel.client + reference = _make_base_document_reference("hi", "bye", client=client) + data_1 = {"zoop": 83} + data_2 = {"zoop": 30} + snapshot_1 = _make_document_snapshot( + reference, + data_1, + True, + mock.sentinel.read_time, + mock.sentinel.create_time, + mock.sentinel.update_time, + ) + snapshot_2 = _make_document_snapshot( + reference, + data_2, + True, + mock.sentinel.read_time, + mock.sentinel.create_time, + mock.sentinel.update_time, + ) + explain_metrics = _make_explain_metrics() + explain_options = ExplainOptions(analyze=True) + snapshot_list = _make_query_results_list( + [snapshot_1, snapshot_2], + explain_options=explain_options, + explain_metrics=explain_metrics, + ) + assert len(snapshot_list) == 2 + assert snapshot_list[0] == snapshot_1 + assert snapshot_list[1] == snapshot_2 + assert snapshot_list._explain_options == explain_options + assert snapshot_list._explain_metrics == explain_metrics + + +def test_query_results_list_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + explain_metrics = _make_explain_metrics() + snapshot_list = _make_query_results_list( + [], explain_options=explain_options, explain_metrics=explain_metrics + ) + + assert snapshot_list.explain_options == explain_options + + +def test_query_results_list_explain_metrics_w_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_metrics = _make_explain_metrics() + snapshot_list = _make_query_results_list( + [], + explain_options=ExplainOptions(analyze=True), + explain_metrics=explain_metrics, + ) + + assert snapshot_list.get_explain_metrics() == explain_metrics + + +def test_query_results_list_explain_metrics_wo_explain_options(): + from google.cloud.firestore_v1.query_profile import QueryExplainError + + snapshot_list = _make_query_results_list([]) + + with pytest.raises(QueryExplainError): + snapshot_list.get_explain_metrics() + + def test__get_document_path(): from google.cloud.firestore_v1.base_document import _get_document_path diff --git a/tests/unit/v1/test_base_query.py b/tests/unit/v1/test_base_query.py index 227b46933f..24caa5e40c 100644 --- a/tests/unit/v1/test_base_query.py +++ b/tests/unit/v1/test_base_query.py @@ -1962,11 +1962,12 @@ def _make_order_pb(field_path, direction): def _make_query_response(**kwargs): - # kwargs supported are ``skipped_results``, ``name`` and ``data`` + # kwargs supported are ``skipped_results``, ``name``, ``data`` + # and ``explain_metrics`` from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.firestore_v1 import _helpers - from google.cloud.firestore_v1.types import document, firestore + from google.cloud.firestore_v1.types import document, firestore, query_profile now = datetime.datetime.now(tz=datetime.timezone.utc) read_time = _datetime_to_pb_timestamp(now) @@ -1984,6 +1985,10 @@ def _make_query_response(**kwargs): kwargs["document"] = document_pb + explain_metrics = kwargs.pop("explain_metrics", None) + if explain_metrics is not None: + kwargs["explain_metrics"] = query_profile.ExplainMetrics(explain_metrics) + return firestore.RunQueryResponse(**kwargs) diff --git a/tests/unit/v1/test_collection.py b/tests/unit/v1/test_collection.py index 98c83664e1..29f76108d1 100644 --- a/tests/unit/v1/test_collection.py +++ b/tests/unit/v1/test_collection.py @@ -385,6 +385,24 @@ def test_get_with_transaction(query_class): query_instance.get.assert_called_once_with(transaction=transaction) +@mock.patch("google.cloud.firestore_v1.query.Query", autospec=True) +def test_get_w_explain_options(query_class): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + collection = _make_collection_reference("collection") + get_response = collection.get(explain_options=explain_options) + + query_class.assert_called_once_with(collection) + query_instance = query_class.return_value + + assert get_response is query_instance.get.return_value + query_instance.get.assert_called_once_with( + transaction=None, + explain_options=explain_options, + ) + + @mock.patch("google.cloud.firestore_v1.query.Query", autospec=True) def test_stream(query_class): collection = _make_collection_reference("collection") @@ -427,6 +445,24 @@ def test_stream_with_transaction(query_class): query_instance.stream.assert_called_once_with(transaction=transaction) +@mock.patch("google.cloud.firestore_v1.query.Query", autospec=True) +def test_stream_w_explain_options(query_class): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + collection = _make_collection_reference("collection") + get_response = collection.stream(explain_options=explain_options) + + query_class.assert_called_once_with(collection) + query_instance = query_class.return_value + + assert get_response is query_instance.stream.return_value + query_instance.stream.assert_called_once_with( + transaction=None, + explain_options=explain_options, + ) + + @mock.patch("google.cloud.firestore_v1.collection.Watch", autospec=True) def test_on_snapshot(watch): collection = _make_collection_reference("collection") diff --git a/tests/unit/v1/test_query.py b/tests/unit/v1/test_query.py index b7add63f36..f30a4fcdff 100644 --- a/tests/unit/v1/test_query.py +++ b/tests/unit/v1/test_query.py @@ -18,6 +18,8 @@ import pytest from google.cloud.firestore_v1.base_client import DEFAULT_DATABASE +from google.cloud.firestore_v1.query_profile import ExplainMetrics, QueryExplainError +from google.cloud.firestore_v1.query_results import QueryResultsList from tests.unit.v1._test_helpers import DEFAULT_TEST_PROJECT, make_client, make_query from tests.unit.v1.test_base_query import _make_cursor_pb, _make_query_response @@ -35,7 +37,12 @@ def test_query_constructor(): assert not query._all_descendants -def _query_get_helper(retry=None, timeout=None, database=None): +def _query_get_helper( + retry=None, + timeout=None, + database=None, + explain_options=None, +): from google.cloud.firestore_v1 import _helpers # Create a minimal fake GAPIC. @@ -52,30 +59,48 @@ def _query_get_helper(retry=None, timeout=None, database=None): _, expected_prefix = parent._parent_info() name = "{}/sleep".format(expected_prefix) data = {"snooze": 10} + explain_metrics = {"execution_stats": {"results_returned": 1}} - response_pb = _make_query_response(name=name, data=data) + response_pb = _make_query_response( + name=name, + data=data, + explain_metrics=explain_metrics, + ) firestore_api.run_query.return_value = iter([response_pb]) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) # Execute the query and check the response. query = make_query(parent) - returned = query.get(**kwargs) + returned = query.get(**kwargs, explain_options=explain_options) - assert isinstance(returned, list) + assert isinstance(returned, QueryResultsList) assert len(returned) == 1 snapshot = returned[0] assert snapshot.reference._path, "dee" == "sleep" assert snapshot.to_dict() == data - # Verify the mock call. + if explain_options is None: + with pytest.raises(QueryExplainError, match="explain_options not set"): + returned.get_explain_metrics() + else: + actual_explain_metrics = returned.get_explain_metrics() + assert isinstance(actual_explain_metrics, ExplainMetrics) + assert actual_explain_metrics.execution_stats.results_returned == 1 + + # Create expected request body. parent_path, _ = parent._parent_info() + request = { + "parent": parent_path, + "structured_query": query._to_protobuf(), + "transaction": None, + } + if explain_options: + request["explain_options"] = explain_options._to_dict() + + # Verify the mock call. firestore_api.run_query.assert_called_once_with( - request={ - "parent": parent_path, - "structured_query": query._to_protobuf(), - "transaction": None, - }, + request=request, metadata=client._rpc_metadata, **kwargs, ) @@ -149,6 +174,13 @@ def test_query_get_limit_to_last(database): ) +def test_query_get_w_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + _query_get_helper(explain_options=explain_options) + + @pytest.mark.parametrize("database", [None, "somedb"]) def test_query_sum(database): from google.cloud.firestore_v1.base_aggregation import SumAggregation @@ -301,7 +333,12 @@ def test_query_chunkify_w_chunksize_gt_limit(database, expected): assert chunk_ids == expected_ids -def _query_stream_helper(retry=None, timeout=None, database=None): +def _query_stream_helper( + retry=None, + timeout=None, + database=None, + explain_options=None, +): from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1.stream_generator import StreamGenerator @@ -319,14 +356,20 @@ def _query_stream_helper(retry=None, timeout=None, database=None): _, expected_prefix = parent._parent_info() name = "{}/sleep".format(expected_prefix) data = {"snooze": 10} - response_pb = _make_query_response(name=name, data=data) + if explain_options is not None: + explain_metrics = {"execution_stats": {"results_returned": 1}} + else: + explain_metrics = None + response_pb = _make_query_response( + name=name, data=data, explain_metrics=explain_metrics + ) firestore_api.run_query.return_value = iter([response_pb]) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) # Execute the query and check the response. query = make_query(parent) - get_response = query.stream(**kwargs) + get_response = query.stream(**kwargs, explain_options=explain_options) assert isinstance(get_response, StreamGenerator) returned = list(get_response) @@ -335,14 +378,28 @@ def _query_stream_helper(retry=None, timeout=None, database=None): assert snapshot.reference._path == ("dee", "sleep") assert snapshot.to_dict() == data - # Verify the mock call. + # Verify explain_metrics. + if explain_options is None: + with pytest.raises(QueryExplainError, match="explain_options not set"): + get_response.get_explain_metrics() + else: + explain_metrics = get_response.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + assert explain_metrics.execution_stats.results_returned == 1 + + # Create expected request body. parent_path, _ = parent._parent_info() + request = { + "parent": parent_path, + "structured_query": query._to_protobuf(), + "transaction": None, + } + if explain_options is not None: + request["explain_options"] = explain_options._to_dict() + + # Verify the mock call. firestore_api.run_query.assert_called_once_with( - request={ - "parent": parent_path, - "structured_query": query._to_protobuf(), - "transaction": None, - }, + request=request, metadata=client._rpc_metadata, **kwargs, ) @@ -747,6 +804,13 @@ def test_query_stream_w_retriable_exc_w_transaction(): _query_stream_w_retriable_exc_helper(transaction=txn) +def test_query_stream_w_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + _query_stream_helper(explain_options=explain_options) + + @mock.patch("google.cloud.firestore_v1.query.Watch", autospec=True) def test_query_on_snapshot(watch): query = make_query(mock.sentinel.parent) diff --git a/tests/unit/v1/test_query_profile.py b/tests/unit/v1/test_query_profile.py new file mode 100644 index 0000000000..a3b0390c61 --- /dev/null +++ b/tests/unit/v1/test_query_profile.py @@ -0,0 +1,126 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + + +def test_explain_metrics__from_pb(): + """ + Test creating an instance of ExplainMetrics from a protobuf. + """ + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + _ExplainAnalyzeMetrics, + QueryExplainError, + PlanSummary, + ) + from google.cloud.firestore_v1.types import query_profile as query_profile_pb2 + from google.protobuf import struct_pb2, duration_pb2 + + # test without execution_stats field + expected_metrics = query_profile_pb2.ExplainMetrics( + plan_summary=query_profile_pb2.PlanSummary( + indexes_used=struct_pb2.ListValue(values=[]) + ) + ) + metrics = ExplainMetrics._from_pb(expected_metrics) + assert isinstance(metrics, ExplainMetrics) + assert isinstance(metrics.plan_summary, PlanSummary) + assert metrics.plan_summary.indexes_used == [] + with pytest.raises(QueryExplainError) as exc: + metrics.execution_stats + assert "execution_stats not available when explain_options.analyze=False" in str( + exc.value + ) + # test with execution_stats field + expected_metrics.execution_stats = query_profile_pb2.ExecutionStats( + results_returned=1, + execution_duration=duration_pb2.Duration(seconds=2), + read_operations=3, + debug_stats=struct_pb2.Struct( + fields={"foo": struct_pb2.Value(string_value="bar")} + ), + ) + metrics = ExplainMetrics._from_pb(expected_metrics) + assert isinstance(metrics, ExplainMetrics) + assert isinstance(metrics, _ExplainAnalyzeMetrics) + assert metrics.execution_stats.results_returned == 1 + assert metrics.execution_stats.execution_duration.total_seconds() == 2 + assert metrics.execution_stats.read_operations == 3 + assert metrics.execution_stats.debug_stats == {"foo": "bar"} + + +def test_explain_metrics__from_pb_empty(): + """ + Test with empty ExplainMetrics protobuf. + """ + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExecutionStats, + _ExplainAnalyzeMetrics, + PlanSummary, + ) + from google.cloud.firestore_v1.types import query_profile as query_profile_pb2 + from google.protobuf import struct_pb2 + + expected_metrics = query_profile_pb2.ExplainMetrics( + plan_summary=query_profile_pb2.PlanSummary( + indexes_used=struct_pb2.ListValue(values=[]) + ), + execution_stats=query_profile_pb2.ExecutionStats(), + ) + metrics = ExplainMetrics._from_pb(expected_metrics) + assert isinstance(metrics, ExplainMetrics) + assert isinstance(metrics, _ExplainAnalyzeMetrics) + assert isinstance(metrics.plan_summary, PlanSummary) + assert isinstance(metrics.execution_stats, ExecutionStats) + assert metrics.plan_summary.indexes_used == [] + assert metrics.execution_stats.results_returned == 0 + assert metrics.execution_stats.execution_duration.total_seconds() == 0 + assert metrics.execution_stats.read_operations == 0 + assert metrics.execution_stats.debug_stats == {} + + +def test_explain_metrics_execution_stats(): + """ + Standard ExplainMetrics class should raise exception when execution_stats is accessed. + _ExplainAnalyzeMetrics should include the field + """ + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + QueryExplainError, + _ExplainAnalyzeMetrics, + ) + + metrics = ExplainMetrics(plan_summary=object()) + with pytest.raises(QueryExplainError) as exc: + metrics.execution_stats + assert "execution_stats not available when explain_options.analyze=False" in str( + exc.value + ) + expected_stats = object() + metrics = _ExplainAnalyzeMetrics( + plan_summary=object(), _execution_stats=expected_stats + ) + assert metrics.execution_stats is expected_stats + + +def test_explain_options__to_dict(): + """ + Should be able to create a dict representation of ExplainOptions + """ + from google.cloud.firestore_v1.query_profile import ExplainOptions + + assert ExplainOptions(analyze=True)._to_dict() == {"analyze": True} + assert ExplainOptions(analyze=False)._to_dict() == {"analyze": False} diff --git a/tests/unit/v1/test_query_results.py b/tests/unit/v1/test_query_results.py new file mode 100644 index 0000000000..59e7878de7 --- /dev/null +++ b/tests/unit/v1/test_query_results.py @@ -0,0 +1,158 @@ +# Copyright 2020 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import mock +import pytest + +from google.cloud.firestore_v1.query_profile import QueryExplainError + + +def _make_base_document_reference(*args, **kwargs): + from google.cloud.firestore_v1.base_document import BaseDocumentReference + + return BaseDocumentReference(*args, **kwargs) + + +def _make_document_snapshot(*args, **kwargs): + from google.cloud.firestore_v1.document import DocumentSnapshot + + return DocumentSnapshot(*args, **kwargs) + + +def _make_query_results_list(*args, **kwargs): + from google.cloud.firestore_v1.query_results import QueryResultsList + + return QueryResultsList(*args, **kwargs) + + +def _make_explain_metrics(): + from google.cloud.firestore_v1.query_profile import ExplainMetrics, PlanSummary + + plan_summary = PlanSummary( + indexes_used=[{"properties": "(__name__ ASC)", "query_scope": "Collection"}], + ) + return ExplainMetrics(plan_summary=plan_summary) + + +def test_query_results_list_constructor(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + client = mock.sentinel.client + reference = _make_base_document_reference("hi", "bye", client=client) + data_1 = {"zoop": 83} + data_2 = {"zoop": 30} + snapshot_1 = _make_document_snapshot( + reference, + data_1, + True, + mock.sentinel.read_time, + mock.sentinel.create_time, + mock.sentinel.update_time, + ) + snapshot_2 = _make_document_snapshot( + reference, + data_2, + True, + mock.sentinel.read_time, + mock.sentinel.create_time, + mock.sentinel.update_time, + ) + explain_metrics = _make_explain_metrics() + explain_options = ExplainOptions(analyze=True) + snapshot_list = _make_query_results_list( + [snapshot_1, snapshot_2], + explain_options=explain_options, + explain_metrics=explain_metrics, + ) + assert len(snapshot_list) == 2 + assert snapshot_list[0] == snapshot_1 + assert snapshot_list[1] == snapshot_2 + assert snapshot_list._explain_options == explain_options + assert snapshot_list._explain_metrics == explain_metrics + + +def test_query_results_list_constructor_w_explain_options_and_wo_explain_metrics(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + with pytest.raises( + ValueError, + match="If explain_options is set, explain_metrics must be non-empty.", + ): + _make_query_results_list( + [], + explain_options=ExplainOptions(analyze=True), + explain_metrics=None, + ) + + +def test_query_results_list_constructor_wo_explain_options_and_w_explain_metrics(): + with pytest.raises( + ValueError, match="If explain_options is empty, explain_metrics must be empty." + ): + _make_query_results_list( + [], + explain_options=None, + explain_metrics=_make_explain_metrics(), + ) + + +def test_query_results_list_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + explain_metrics = _make_explain_metrics() + snapshot_list = _make_query_results_list( + [], explain_options=explain_options, explain_metrics=explain_metrics + ) + + assert snapshot_list.explain_options == explain_options + + +def test_query_results_list_explain_metrics_w_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_metrics = _make_explain_metrics() + snapshot_list = _make_query_results_list( + [], + explain_options=ExplainOptions(analyze=True), + explain_metrics=explain_metrics, + ) + + assert snapshot_list.get_explain_metrics() == explain_metrics + + +def test_query_results_list_explain_metrics_wo_explain_options(): + snapshot_list = _make_query_results_list([]) + + with pytest.raises(QueryExplainError, match="explain_options not set on query."): + snapshot_list.get_explain_metrics() + + +def test_query_results_list_explain_metrics_empty(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_metrics = _make_explain_metrics() + snapshot_list = _make_query_results_list( + [], + explain_options=ExplainOptions(analyze=True), + explain_metrics=explain_metrics, + ) + snapshot_list._explain_metrics = None + + with pytest.raises( + QueryExplainError, + match="explain_metrics is empty despite explain_options is set.", + ): + snapshot_list.get_explain_metrics() diff --git a/tests/unit/v1/test_stream_generator.py b/tests/unit/v1/test_stream_generator.py index bfc11cf6f6..0e8a552607 100644 --- a/tests/unit/v1/test_stream_generator.py +++ b/tests/unit/v1/test_stream_generator.py @@ -14,8 +14,10 @@ import pytest +from google.protobuf import struct_pb2 -def _make_stream_generator(iterable): + +def _make_stream_generator(iterable, explain_options=None, explain_metrics=None): from google.cloud.firestore_v1.stream_generator import StreamGenerator def _inner_generator(): @@ -23,14 +25,27 @@ def _inner_generator(): X = yield i if X: yield X + return explain_metrics + + return StreamGenerator(_inner_generator(), explain_options) + + +def test_stream_generator_constructor(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + from google.cloud.firestore_v1.stream_generator import StreamGenerator + + explain_options = ExplainOptions(analyze=True) + inner_generator = object() + inst = StreamGenerator(inner_generator, explain_options) - return StreamGenerator(_inner_generator()) + assert inst._generator == inner_generator + assert inst._explain_options == explain_options + assert inst._explain_metrics is None def test_stream_generator_iter(): expected_results = [0, 1, 2] inst = _make_stream_generator(expected_results) - actual_results = [] for result in inst: actual_results.append(result) @@ -82,3 +97,159 @@ def test_stream_generator_close(): # Verifies that generator is closed. with pytest.raises(StopIteration): next(inst) + + +def test_stream_generator_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + inst = _make_stream_generator([], explain_options) + assert inst.explain_options == explain_options + + +def test_stream_generator_explain_metrics_explain_options_analyze_true(): + from google.protobuf import duration_pb2 + from google.protobuf import struct_pb2 + + import google.cloud.firestore_v1.query_profile as query_profile + import google.cloud.firestore_v1.types.query_profile as query_profile_pb2 + + iterator = [1, 2] + + indexes_used_dict = { + "indexes_used": struct_pb2.Value( + struct_value=struct_pb2.Struct( + fields={ + "query_scope": struct_pb2.Value(string_value="Collection"), + "properties": struct_pb2.Value( + string_value="(foo ASC, **name** ASC)" + ), + } + ) + ) + } + plan_summary = query_profile_pb2.PlanSummary() + plan_summary.indexes_used.append(indexes_used_dict) + execution_stats = query_profile_pb2.ExecutionStats( + { + "results_returned": 1, + "execution_duration": duration_pb2.Duration(seconds=2), + "read_operations": 3, + "debug_stats": struct_pb2.Struct( + fields={ + "billing_details": struct_pb2.Value( + string_value="billing_details_results" + ), + "documents_scanned": struct_pb2.Value( + string_value="documents_scanned_results" + ), + "index_entries_scanned": struct_pb2.Value( + string_value="index_entries_scanned" + ), + } + ), + } + ) + + explain_options = query_profile.ExplainOptions(analyze=True) + expected_explain_metrics = query_profile_pb2.ExplainMetrics( + plan_summary=plan_summary, + execution_stats=execution_stats, + ) + + inst = _make_stream_generator(iterator, explain_options, expected_explain_metrics) + + # Raise an exception if query isn't complete when explain_metrics is called. + with pytest.raises( + query_profile.QueryExplainError, + match="explain_metrics not available until query is complete.", + ): + inst.get_explain_metrics() + + list(inst) + + actual_explain_metrics = inst.get_explain_metrics() + assert isinstance(actual_explain_metrics, query_profile._ExplainAnalyzeMetrics) + assert actual_explain_metrics == query_profile.ExplainMetrics._from_pb( + expected_explain_metrics + ) + assert actual_explain_metrics.plan_summary.indexes_used == [ + { + "indexes_used": { + "query_scope": "Collection", + "properties": "(foo ASC, **name** ASC)", + } + } + ] + assert actual_explain_metrics.execution_stats.results_returned == 1 + duration = actual_explain_metrics.execution_stats.execution_duration.total_seconds() + assert duration == 2 + assert actual_explain_metrics.execution_stats.read_operations == 3 + + expected_debug_stats = { + "billing_details": "billing_details_results", + "documents_scanned": "documents_scanned_results", + "index_entries_scanned": "index_entries_scanned", + } + assert actual_explain_metrics.execution_stats.debug_stats == expected_debug_stats + + +def test_stream_generator_explain_metrics_explain_options_analyze_false(): + import google.cloud.firestore_v1.query_profile as query_profile + import google.cloud.firestore_v1.types.query_profile as query_profile_pb2 + + iterator = [] + + explain_options = query_profile.ExplainOptions(analyze=False) + indexes_used_dict = { + "indexes_used": struct_pb2.Value( + struct_value=struct_pb2.Struct( + fields={ + "query_scope": struct_pb2.Value(string_value="Collection"), + "properties": struct_pb2.Value( + string_value="(foo ASC, **name** ASC)" + ), + } + ) + ) + } + plan_summary = query_profile_pb2.PlanSummary() + plan_summary.indexes_used.append(indexes_used_dict) + expected_explain_metrics = query_profile_pb2.ExplainMetrics( + plan_summary=plan_summary + ) + + inst = _make_stream_generator(iterator, explain_options, expected_explain_metrics) + actual_explain_metrics = inst.get_explain_metrics() + assert isinstance(actual_explain_metrics, query_profile.ExplainMetrics) + assert actual_explain_metrics.plan_summary.indexes_used == [ + { + "indexes_used": { + "query_scope": "Collection", + "properties": "(foo ASC, **name** ASC)", + } + } + ] + + +def test_stream_generator_explain_metrics_missing_explain_options_analyze_false(): + import google.cloud.firestore_v1.query_profile as query_profile + + explain_options = query_profile.ExplainOptions(analyze=False) + inst = _make_stream_generator([("1", None)], explain_options) + with pytest.raises( + query_profile.QueryExplainError, match="Did not receive explain_metrics" + ): + inst.get_explain_metrics() + + +def test_stream_generator_explain_metrics_no_explain_options(): + from google.cloud.firestore_v1.query_profile import QueryExplainError + + inst = _make_stream_generator([]) + + with pytest.raises( + QueryExplainError, + match="explain_options not set on query.", + ): + inst.get_explain_metrics() diff --git a/tests/unit/v1/test_transaction.py b/tests/unit/v1/test_transaction.py index d37be34ea0..b5beef6c2d 100644 --- a/tests/unit/v1/test_transaction.py +++ b/tests/unit/v1/test_transaction.py @@ -15,6 +15,8 @@ import mock import pytest +from tests.unit.v1.test_base_query import _make_query_response + def _make_transaction(*args, **kwargs): from google.cloud.firestore_v1.transaction import Transaction @@ -328,7 +330,11 @@ def test_transaction_get_all_w_retry_timeout(): _transaction_get_all_helper(retry=retry, timeout=timeout) -def _transaction_get_w_document_ref_helper(retry=None, timeout=None): +def _transaction_get_w_document_ref_helper( + retry=None, + timeout=None, + explain_options=None, +): from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1.document import DocumentReference @@ -337,8 +343,14 @@ def _transaction_get_w_document_ref_helper(retry=None, timeout=None): ref = DocumentReference("documents", "doc-id") kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) + if explain_options is not None: + kwargs["explain_options"] = explain_options + result = transaction.get(ref, **kwargs) + # explain_options should not be in the request even if it's provided. + kwargs.pop("explain_options", None) + assert result is client.get_all.return_value client.get_all.assert_called_once_with([ref], transaction=transaction, **kwargs) @@ -355,20 +367,96 @@ def test_transaction_get_w_document_ref_w_retry_timeout(): _transaction_get_w_document_ref_helper(retry=retry, timeout=timeout) -def _transaction_get_w_query_helper(retry=None, timeout=None): +def test_transaction_get_w_document_ref_w_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + with pytest.raises(ValueError, match="`ref_or_query` is `AsyncDocumentReference`"): + _transaction_get_w_document_ref_helper( + explain_options=ExplainOptions(analyze=True), + ) + + +def _transaction_get_w_query_helper( + retry=None, + timeout=None, + explain_options=None, +): from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1.query import Query + from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + QueryExplainError, + ) + from google.cloud.firestore_v1.stream_generator import StreamGenerator - client = mock.Mock(spec=[]) - transaction = _make_transaction(client) - query = Query(parent=mock.Mock(spec=[])) - query.stream = mock.MagicMock() - kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) + # Create a minimal fake GAPIC. + firestore_api = mock.Mock(spec=["run_query"]) + + # Attach the fake GAPIC to a real client. + client = _make_client() + client._firestore_api_internal = firestore_api - result = transaction.get(query, **kwargs) + # Make a **real** collection reference as parent. + parent = client.collection("dee") + + # Add a dummy response to the minimal fake GAPIC. + _, expected_prefix = parent._parent_info() + name = "{}/sleep".format(expected_prefix) + data = {"snooze": 10} + if explain_options is not None: + explain_metrics = {"execution_stats": {"results_returned": 1}} + else: + explain_metrics = None + response_pb = _make_query_response( + name=name, data=data, explain_metrics=explain_metrics + ) + firestore_api.run_query.return_value = iter([response_pb]) + kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) - assert result is query.stream.return_value - query.stream.assert_called_once_with(transaction=transaction, **kwargs) + # Run the transaction with query. + transaction = _make_transaction(client) + txn_id = b"beep-fail-commit" + transaction._id = txn_id + query = Query(parent) + returned_generator = transaction.get( + query, + **kwargs, + explain_options=explain_options, + ) + + # Verify the response. + assert isinstance(returned_generator, StreamGenerator) + results = list(returned_generator) + assert len(results) == 1 + snapshot = results[0] + assert snapshot.reference._path == ("dee", "sleep") + assert snapshot.to_dict() == data + + # Verify explain_metrics. + if explain_options is None: + with pytest.raises(QueryExplainError, match="explain_options not set"): + returned_generator.get_explain_metrics() + else: + explain_metrics = returned_generator.get_explain_metrics() + assert isinstance(explain_metrics, ExplainMetrics) + assert explain_metrics.execution_stats.results_returned == 1 + + # Create expected request body. + parent_path, _ = parent._parent_info() + request = { + "parent": parent_path, + "structured_query": query._to_protobuf(), + "transaction": b"beep-fail-commit", + } + if explain_options is not None: + request["explain_options"] = explain_options._to_dict() + + # Verify the mock call. + firestore_api.run_query.assert_called_once_with( + request=request, + metadata=client._rpc_metadata, + **kwargs, + ) def test_transaction_get_w_query(): @@ -383,6 +471,12 @@ def test_transaction_get_w_query_w_retry_timeout(): _transaction_get_w_query_helper(retry=retry, timeout=timeout) +def test_transaction_get_w_query_w_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + _transaction_get_w_query_helper(explain_options=ExplainOptions(analyze=True)) + + @pytest.mark.parametrize("database", [None, "somedb"]) def test_transaction_get_failure(database): client = _make_client(database=database) diff --git a/tests/unit/v1/test_vector_query.py b/tests/unit/v1/test_vector_query.py index a5b1d342bd..eb5328ace6 100644 --- a/tests/unit/v1/test_vector_query.py +++ b/tests/unit/v1/test_vector_query.py @@ -17,6 +17,12 @@ from google.cloud.firestore_v1._helpers import encode_value, make_retry_timeout_kwargs from google.cloud.firestore_v1.base_vector_query import DistanceMeasure +from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExplainOptions, + QueryExplainError, +) +from google.cloud.firestore_v1.query_results import QueryResultsList from google.cloud.firestore_v1.types.query import StructuredQuery from google.cloud.firestore_v1.vector import Vector from tests.unit.v1._test_helpers import make_client, make_query, make_vector_query @@ -146,21 +152,7 @@ def _expected_pb( return expected_pb -@pytest.mark.parametrize( - "distance_measure, expected_distance", - [ - ( - DistanceMeasure.EUCLIDEAN, - StructuredQuery.FindNearest.DistanceMeasure.EUCLIDEAN, - ), - (DistanceMeasure.COSINE, StructuredQuery.FindNearest.DistanceMeasure.COSINE), - ( - DistanceMeasure.DOT_PRODUCT, - StructuredQuery.FindNearest.DistanceMeasure.DOT_PRODUCT, - ), - ], -) -def test_vector_query(distance_measure, expected_distance): +def _vector_query_get_helper(distance_measure, expected_distance, explain_options=None): # Create a minimal fake GAPIC. firestore_api = mock.Mock(spec=["run_query"]) client = make_client() @@ -171,8 +163,14 @@ def test_vector_query(distance_measure, expected_distance): parent_path, expected_prefix = parent._parent_info() data = {"snooze": 10, "embedding": Vector([1.0, 2.0, 3.0])} + if explain_options is not None: + explain_metrics = {"execution_stats": {"results_returned": 1}} + else: + explain_metrics = None response_pb = _make_query_response( - name="{}/test_doc".format(expected_prefix), data=data + name="{}/test_doc".format(expected_prefix), + data=data, + explain_metrics=explain_metrics, ) kwargs = make_retry_timeout_kwargs(retry=None, timeout=None) @@ -187,11 +185,21 @@ def test_vector_query(distance_measure, expected_distance): limit=5, ) - returned = vector_query.get(transaction=_transaction(client), **kwargs) - assert isinstance(returned, list) + returned = vector_query.get( + transaction=_transaction(client), **kwargs, explain_options=explain_options + ) + assert isinstance(returned, QueryResultsList) assert len(returned) == 1 assert returned[0].to_dict() == data + if explain_options is None: + with pytest.raises(QueryExplainError, match="explain_options not set"): + returned.get_explain_metrics() + else: + actual_explain_metrics = returned.get_explain_metrics() + assert isinstance(actual_explain_metrics, ExplainMetrics) + assert actual_explain_metrics.execution_stats.results_returned == 1 + expected_pb = _expected_pb( parent=parent, vector_field="embedding", @@ -199,17 +207,49 @@ def test_vector_query(distance_measure, expected_distance): distance_type=expected_distance, limit=5, ) + expected_request = { + "parent": parent_path, + "structured_query": expected_pb, + "transaction": _TXN_ID, + } + if explain_options is not None: + expected_request["explain_options"] = explain_options._to_dict() firestore_api.run_query.assert_called_once_with( - request={ - "parent": parent_path, - "structured_query": expected_pb, - "transaction": _TXN_ID, - }, + request=expected_request, metadata=client._rpc_metadata, **kwargs, ) +@pytest.mark.parametrize( + "distance_measure, expected_distance", + [ + ( + DistanceMeasure.EUCLIDEAN, + StructuredQuery.FindNearest.DistanceMeasure.EUCLIDEAN, + ), + (DistanceMeasure.COSINE, StructuredQuery.FindNearest.DistanceMeasure.COSINE), + ( + DistanceMeasure.DOT_PRODUCT, + StructuredQuery.FindNearest.DistanceMeasure.DOT_PRODUCT, + ), + ], +) +def test_vector_query(distance_measure, expected_distance): + _vector_query_get_helper( + distance_measure=distance_measure, expected_distance=expected_distance + ) + + +def test_vector_query_w_explain_options(): + explain_options = ExplainOptions(analyze=True) + _vector_query_get_helper( + distance_measure=DistanceMeasure.EUCLIDEAN, + expected_distance=StructuredQuery.FindNearest.DistanceMeasure.EUCLIDEAN, + explain_options=explain_options, + ) + + @pytest.mark.parametrize( "distance_measure, expected_distance", [