1# Copyright 2024 Google LLC
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14from __future__ import annotations
15
16from typing import Any
17
18import datetime
19
20from dataclasses import dataclass
21from google.protobuf.json_format import MessageToDict
22from google.cloud.firestore_v1.types.document import MapValue
23from google.cloud.firestore_v1.types.document import Value
24from google.cloud.firestore_v1.types.explain_stats import (
25 ExplainStats as ExplainStats_pb,
26)
27from google.protobuf.wrappers_pb2 import StringValue
28
29
30@dataclass(frozen=True)
31class ExplainOptions:
32 """
33 Explain options for the query.
34 Set on a query object using the explain_options attribute at query
35 construction time.
36
37 :type analyze: bool
38 :param analyze: Optional. Whether to execute this query. When false
39 (the default), the query will be planned, returning only metrics from the
40 planning stages. When true, the query will be planned and executed,
41 returning the full query results along with both planning and execution
42 stage metrics.
43 """
44
45 analyze: bool = False
46
47 def _to_dict(self):
48 return {"analyze": self.analyze}
49
50
51@dataclass(frozen=True)
52class PipelineExplainOptions:
53 """
54 Explain options for pipeline queries.
55
56 Set on a pipeline.execution() or pipeline.stream() call, to provide
57 explain_stats in the pipeline output
58
59 :type mode: str
60 :param mode: Optional. The mode of operation for this explain query.
61 When set to 'analyze', the query will be executed and return the full
62 query results along with execution statistics.
63
64 :type output_format: str | None
65 :param output_format: Optional. The format in which to return the explain
66 stats.
67 """
68
69 mode: str = "analyze"
70
71 def _to_value(self):
72 out_dict = {"mode": Value(string_value=self.mode)}
73 value_pb = MapValue(fields=out_dict)
74 return Value(map_value=value_pb)
75
76
77@dataclass(frozen=True)
78class PlanSummary:
79 """
80 Contains planning phase information about a query.`
81
82 :type indexes_used: list[dict[str, Any]]
83 :param indexes_used: The indexes selected for this query.
84 """
85
86 indexes_used: list[dict[str, Any]]
87
88
89@dataclass(frozen=True)
90class ExecutionStats:
91 """
92 Execution phase information about a query.
93
94 Only available when explain_options.analyze is True.
95
96 :type results_returned: int
97 :param results_returned: Total number of results returned, including
98 documents, projections, aggregation results, keys.
99 :type execution_duration: datetime.timedelta
100 :param execution_duration: Total time to execute the query in the backend.
101 :type read_operations: int
102 :param read_operations: Total billable read operations.
103 :type debug_stats: dict[str, Any]
104 :param debug_stats: Debugging statistics from the execution of the query.
105 Note that the debugging stats are subject to change as Firestore evolves
106 """
107
108 results_returned: int
109 execution_duration: datetime.timedelta
110 read_operations: int
111 debug_stats: dict[str, Any]
112
113
114@dataclass(frozen=True)
115class ExplainMetrics:
116 """
117 ExplainMetrics contains information about the planning and execution of a query.
118
119 When explain_options.analyze is false, only plan_summary is available.
120 When explain_options.analyze is true, execution_stats is also available.
121
122 :type plan_summary: PlanSummary
123 :param plan_summary: Planning phase information about the query.
124 :type execution_stats: ExecutionStats
125 :param execution_stats: Execution phase information about the query.
126 """
127
128 plan_summary: PlanSummary
129
130 @staticmethod
131 def _from_pb(metrics_pb):
132 dict_repr = MessageToDict(metrics_pb._pb, preserving_proto_field_name=True)
133 plan_summary = PlanSummary(
134 indexes_used=dict_repr.get("plan_summary", {}).get("indexes_used", [])
135 )
136 if "execution_stats" in dict_repr:
137 stats_dict = dict_repr.get("execution_stats", {})
138 execution_stats = ExecutionStats(
139 results_returned=int(stats_dict.get("results_returned", 0)),
140 execution_duration=metrics_pb.execution_stats.execution_duration,
141 read_operations=int(stats_dict.get("read_operations", 0)),
142 debug_stats=stats_dict.get("debug_stats", {}),
143 )
144 return _ExplainAnalyzeMetrics(
145 plan_summary=plan_summary, _execution_stats=execution_stats
146 )
147 else:
148 return ExplainMetrics(plan_summary=plan_summary)
149
150 @property
151 def execution_stats(self) -> ExecutionStats:
152 raise QueryExplainError(
153 "execution_stats not available when explain_options.analyze=False."
154 )
155
156
157@dataclass(frozen=True)
158class _ExplainAnalyzeMetrics(ExplainMetrics):
159 """
160 Subclass of ExplainMetrics that includes execution_stats.
161 Only available when explain_options.analyze is True.
162 """
163
164 plan_summary: PlanSummary
165 _execution_stats: ExecutionStats
166
167 @property
168 def execution_stats(self) -> ExecutionStats:
169 return self._execution_stats
170
171
172class QueryExplainError(Exception):
173 """
174 Error returned when there is a problem accessing query profiling information.
175 """
176
177 pass
178
179
180class ExplainStats:
181 """
182 Contains query profiling statistics for a pipeline query.
183
184 This class is not meant to be instantiated directly by the user. Instead, an
185 instance of `ExplainStats` may be returned by pipeline execution methods
186 when `explain_options` are provided.
187
188 It provides methods to access the explain statistics in different formats.
189 """
190
191 def __init__(self, stats_pb: ExplainStats_pb):
192 """
193 Args:
194 stats_pb (ExplainStats_pb): The raw protobuf message for explain stats.
195 """
196 self._stats_pb = stats_pb
197
198 def get_text(self) -> str:
199 """
200 Returns the explain stats as a string.
201
202 This method is suitable for explain formats that have a text-based output,
203 such as 'text' or 'json'.
204
205 Returns:
206 str: The string representation of the explain stats.
207
208 Raises:
209 QueryExplainError: If the explain stats payload from the backend is not
210 a string. This can happen if a non-text output format was requested.
211 """
212 pb_data = self._stats_pb._pb.data
213 content = StringValue()
214 if pb_data.Unpack(content):
215 return content.value
216 raise QueryExplainError(
217 "Unable to decode explain stats. Did you request an output format that returns a string value, such as 'text' or 'json'?"
218 )
219
220 def get_raw(self) -> ExplainStats_pb:
221 """
222 Returns the explain stats in an encoded proto format, as returned from the Firestore backend.
223 The caller is responsible for unpacking this proto message.
224
225 Returns:
226 google.cloud.firestore_v1.types.explain_stats.ExplainStats: the proto from the backend
227 """
228 return self._stats_pb