|
16 | 16 |
|
17 | 17 | import functools |
18 | 18 | import json |
19 | | -import sys |
20 | | -import textwrap |
21 | 19 | import types |
22 | 20 | from typing import Callable |
23 | 21 |
|
@@ -79,7 +77,6 @@ def execute_sql( |
79 | 77 | ] |
80 | 78 | } |
81 | 79 | """ |
82 | | - |
83 | 80 | try: |
84 | 81 | # Get BigQuery client |
85 | 82 | bq_client = client.get_bigquery_client( |
@@ -189,7 +186,47 @@ def execute_sql( |
189 | 186 | } |
190 | 187 |
|
191 | 188 |
|
192 | | -_execute_sql_write_examples = """ |
| 189 | +def _execute_sql_write_mode(*args, **kwargs) -> dict: |
| 190 | + """Run a BigQuery or BigQuery ML SQL query in the project and return the result. |
| 191 | +
|
| 192 | + Args: |
| 193 | + project_id (str): The GCP project id in which the query should be |
| 194 | + executed. |
| 195 | + query (str): The BigQuery SQL query to be executed. |
| 196 | + credentials (Credentials): The credentials to use for the request. |
| 197 | + settings (BigQueryToolConfig): The settings for the tool. |
| 198 | + tool_context (ToolContext): The context for the tool. |
| 199 | +
|
| 200 | + Returns: |
| 201 | + dict: Dictionary representing the result of the query. |
| 202 | + If the result contains the key "result_is_likely_truncated" with |
| 203 | + value True, it means that there may be additional rows matching the |
| 204 | + query not returned in the result. |
| 205 | +
|
| 206 | + Examples: |
| 207 | + Fetch data or insights from a table: |
| 208 | +
|
| 209 | + >>> execute_sql("my_project", |
| 210 | + ... "SELECT island, COUNT(*) AS population " |
| 211 | + ... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island") |
| 212 | + { |
| 213 | + "status": "SUCCESS", |
| 214 | + "rows": [ |
| 215 | + { |
| 216 | + "island": "Dream", |
| 217 | + "population": 124 |
| 218 | + }, |
| 219 | + { |
| 220 | + "island": "Biscoe", |
| 221 | + "population": 168 |
| 222 | + }, |
| 223 | + { |
| 224 | + "island": "Torgersen", |
| 225 | + "population": 52 |
| 226 | + } |
| 227 | + ] |
| 228 | + } |
| 229 | +
|
193 | 230 | Create a table with schema prescribed: |
194 | 231 |
|
195 | 232 | >>> execute_sql("my_project", |
@@ -328,9 +365,50 @@ def execute_sql( |
328 | 365 | - Use "CREATE OR REPLACE MODEL" instead of "CREATE MODEL". |
329 | 366 | - First run "DROP MODEL", followed by "CREATE MODEL". |
330 | 367 | """ |
| 368 | + return execute_sql(*args, **kwargs) |
331 | 369 |
|
332 | 370 |
|
333 | | -_execute_sql_protecetd_write_examples = """ |
| 371 | +def _execute_sql_protected_write_mode(*args, **kwargs) -> dict: |
| 372 | + """Run a BigQuery or BigQuery ML SQL query in the project and return the result. |
| 373 | +
|
| 374 | + Args: |
| 375 | + project_id (str): The GCP project id in which the query should be |
| 376 | + executed. |
| 377 | + query (str): The BigQuery SQL query to be executed. |
| 378 | + credentials (Credentials): The credentials to use for the request. |
| 379 | + settings (BigQueryToolConfig): The settings for the tool. |
| 380 | + tool_context (ToolContext): The context for the tool. |
| 381 | +
|
| 382 | + Returns: |
| 383 | + dict: Dictionary representing the result of the query. |
| 384 | + If the result contains the key "result_is_likely_truncated" with |
| 385 | + value True, it means that there may be additional rows matching the |
| 386 | + query not returned in the result. |
| 387 | +
|
| 388 | + Examples: |
| 389 | + Fetch data or insights from a table: |
| 390 | +
|
| 391 | + >>> execute_sql("my_project", |
| 392 | + ... "SELECT island, COUNT(*) AS population " |
| 393 | + ... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island") |
| 394 | + { |
| 395 | + "status": "SUCCESS", |
| 396 | + "rows": [ |
| 397 | + { |
| 398 | + "island": "Dream", |
| 399 | + "population": 124 |
| 400 | + }, |
| 401 | + { |
| 402 | + "island": "Biscoe", |
| 403 | + "population": 168 |
| 404 | + }, |
| 405 | + { |
| 406 | + "island": "Torgersen", |
| 407 | + "population": 52 |
| 408 | + } |
| 409 | + ] |
| 410 | + } |
| 411 | +
|
334 | 412 | Create a temporary table with schema prescribed: |
335 | 413 |
|
336 | 414 | >>> execute_sql("my_project", |
@@ -460,6 +538,7 @@ def execute_sql( |
460 | 538 | - Only temporary models can be created or deleted. Please do not try |
461 | 539 | creating a permanent model (non-TEMP model) or deleting one. |
462 | 540 | """ |
| 541 | + return execute_sql(*args, **kwargs) |
463 | 542 |
|
464 | 543 |
|
465 | 544 | def get_execute_sql(settings: BigQueryToolConfig) -> Callable[..., dict]: |
@@ -496,17 +575,8 @@ def get_execute_sql(settings: BigQueryToolConfig) -> Callable[..., dict]: |
496 | 575 |
|
497 | 576 | # Now, set the new docstring |
498 | 577 | if settings.write_mode == WriteMode.PROTECTED: |
499 | | - examples = _execute_sql_protecetd_write_examples |
| 578 | + execute_sql_wrapper.__doc__ = _execute_sql_protected_write_mode.__doc__ |
500 | 579 | else: |
501 | | - examples = _execute_sql_write_examples |
502 | | - |
503 | | - # Handle Python 3.13+ inspect.cleandoc behavior change |
504 | | - # Python 3.13 changed inspect.cleandoc from lstrip() to lstrip(' '), making it |
505 | | - # more conservative. The appended examples have inconsistent indentation that |
506 | | - # Python 3.11/3.12's aggressive cleandoc would fix, but 3.13+ needs help. |
507 | | - if sys.version_info >= (3, 13): |
508 | | - examples = textwrap.dedent(examples) |
509 | | - |
510 | | - execute_sql_wrapper.__doc__ += examples |
| 580 | + execute_sql_wrapper.__doc__ = _execute_sql_write_mode.__doc__ |
511 | 581 |
|
512 | 582 | return execute_sql_wrapper |
0 commit comments