Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions src/google/adk/integrations/bigquery/query_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,15 @@

from __future__ import annotations

import asyncio
import collections
import functools
import json
import re
import threading
import types
from typing import Any
from typing import Awaitable
from typing import Callable
from typing import Optional
import uuid
Expand Down Expand Up @@ -338,7 +340,7 @@ def _execute_sql(
}


def execute_sql(
async def execute_sql(
project_id: str,
query: str,
credentials: Credentials,
Expand Down Expand Up @@ -424,7 +426,8 @@ def execute_sql(
}
}
"""
return _execute_sql(
return await asyncio.to_thread(
_execute_sql,
project_id=project_id,
query=query,
credentials=credentials,
Expand Down Expand Up @@ -870,7 +873,7 @@ def _execute_sql_protected_write_mode(

def _execute_sql_with_docstring(
docstring: str | None,
) -> Callable[..., dict[str, Any]]:
) -> Callable[..., Awaitable[dict[str, Any]]]:
"""Clone execute_sql, keeping its signature but replacing its docstring."""
# Create a new function object using the original function's code and globals.
# We pass the original code, globals, name, defaults, and closure.
Expand Down Expand Up @@ -907,16 +910,16 @@ def _execute_sql_with_docstring(

def get_execute_sql(
settings: BigQueryToolConfig,
) -> Callable[..., dict[str, Any]]:
) -> Callable[..., Awaitable[dict[str, Any]]]:
"""Get the execute_sql tool customized as per the given tool settings.

Args:
settings: BigQuery tool settings indicating the behavior of the
execute_sql tool.

Returns:
callable[..., dict]: A version of the execute_sql tool respecting the tool
settings.
callable[..., Awaitable[dict]]: A version of the execute_sql tool
respecting the tool settings.
"""

if not settings or settings.write_mode == WriteMode.BLOCKED:
Expand Down
Loading