|
3 | 3 |
|
4 | 4 | @author: wf |
5 | 5 | """ |
6 | | - |
| 6 | +from tqdm import tqdm |
7 | 7 | import logging |
8 | 8 | import sys |
9 | 9 | from argparse import ArgumentParser |
@@ -87,6 +87,12 @@ def getArgParser(self, description: str, version_msg) -> ArgumentParser: |
87 | 87 | action=StoreDictKeyPair, |
88 | 88 | help="query parameters as Key-value pairs in the format key1=value1,key2=value2", |
89 | 89 | ) |
| 90 | + parser.add_argument( |
| 91 | + "--progress", |
| 92 | + action="store_true", |
| 93 | + help="show progress bars when testing queries (--testQueries)", |
| 94 | + ) |
| 95 | + |
90 | 96 | parser.add_argument( |
91 | 97 | "--domain", |
92 | 98 | type=str, |
@@ -144,6 +150,56 @@ def cmd_parse(self, argv: Optional[list] = None): |
144 | 150 | self.args.func(self.args) |
145 | 151 | return self.args |
146 | 152 |
|
| 153 | + def handle_test_queries(self): |
| 154 | + """ |
| 155 | + Handle the --testQueries option by executing queries against endpoints. |
| 156 | + The endpoint is the outer loop, queries are the inner loop. |
| 157 | + """ |
| 158 | + # Determine which endpoints to use |
| 159 | + if self.args.endpointName: |
| 160 | + endpoint_names = [self.args.endpointName] |
| 161 | + else: |
| 162 | + endpoint_names = list(self.nqm.endpoints.keys()) |
| 163 | + |
| 164 | + # Get all queries to test |
| 165 | + queries = self.nqm.get_all_queries(domain=self.args.domain, namespace=self.args.namespace) |
| 166 | + |
| 167 | + # Create execution instance |
| 168 | + execution = Execution(self.nqm, debug=self.args.debug) |
| 169 | + |
| 170 | + # Outer loop: endpoints |
| 171 | + endpoint_iter = tqdm(endpoint_names, desc="Testing endpoints") if self.args.progress else endpoint_names |
| 172 | + for endpoint_name in endpoint_iter: |
| 173 | + # Inner loop: queries |
| 174 | + query_iter = tqdm(queries, desc=f"Queries for {endpoint_name}", leave=False) if self.args.progress else queries |
| 175 | + for i, nq in enumerate(query_iter, start=1): |
| 176 | + execution.execute( |
| 177 | + nq, |
| 178 | + endpoint_name=endpoint_name, |
| 179 | + context=self.args.context, |
| 180 | + title=f"{endpoint_name}::query {i:3}/{len(queries)}", |
| 181 | + prefix_merger=QueryPrefixMerger.get_by_name(self.args.prefix_merger), |
| 182 | + ) |
| 183 | + |
| 184 | + def handle_test_queries_no_progress_version(self): |
| 185 | + if self.args.endpointName: |
| 186 | + endpoint_names = [self.args.endpointName] |
| 187 | + else: |
| 188 | + endpoint_names = list(self.nqm.endpoints.keys()) |
| 189 | + queries = self.nqm.get_all_queries(domain=self.args.domain, namespace=self.args.namespace) |
| 190 | + execution = Execution(self.nqm, debug=self.args.debug) |
| 191 | + query_iter = tqdm(queries, desc="Testing queries") if self.args.progress else queries |
| 192 | + for i, nq in enumerate(query_iter, start=1): |
| 193 | + for endpoint_name in endpoint_names: |
| 194 | + execution.execute( |
| 195 | + nq, |
| 196 | + endpoint_name=endpoint_name, |
| 197 | + context=self.args.context, |
| 198 | + title=f"query {i:3}/{len(queries)}::{endpoint_name}", |
| 199 | + prefix_merger=QueryPrefixMerger.get_by_name(self.args.prefix_merger), |
| 200 | + ) |
| 201 | + |
| 202 | + |
147 | 203 | def handle_args(self, args) -> bool: |
148 | 204 | """ |
149 | 205 | handle the command line args |
@@ -174,21 +230,8 @@ def handle_args(self, args) -> bool: |
174 | 230 | print(f"{namespace}:{count}") |
175 | 231 | handled = True |
176 | 232 | elif self.args.testQueries: |
177 | | - if self.args.endpointName: |
178 | | - endpoint_names = [self.args.endpointName] |
179 | | - else: |
180 | | - endpoint_names = list(nqm.endpoints.keys()) |
181 | | - queries = self.nqm.get_all_queries(domain=self.args.domain, namespace=self.args.namespace) |
182 | | - execution = Execution(self.nqm, debug=self.args.debug) |
183 | | - for i, nq in enumerate(queries, start=1): |
184 | | - for endpoint_name in endpoint_names: |
185 | | - execution.execute( |
186 | | - nq, |
187 | | - endpoint_name=endpoint_name, |
188 | | - context=self.args.context, |
189 | | - title=f"query {i:3}/{len(queries)}::{endpoint_name}", |
190 | | - prefix_merger=QueryPrefixMerger.get_by_name(self.args.prefix_merger), |
191 | | - ) |
| 233 | + self.handle_test_queries() |
| 234 | + handled = True |
192 | 235 | elif self.args.queryName is not None or self.args.query_id is not None: |
193 | 236 | if self.args.query_id is not None: |
194 | 237 | query_name = QueryName.from_query_id(self.args.query_id) |
|
0 commit comments