Skip to content
Merged
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
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,7 @@
)
from google.cloud.bigtable.data.execute_query._parameters_formatting import (
_format_execute_query_params,
_format_execute_query_view_params,
_to_param_types,
)
from google.cloud.bigtable.data.execute_query.metadata import (
Expand Down Expand Up @@ -717,6 +718,7 @@ async def execute_query(
*,
parameters: dict[str, ExecuteQueryValueType] | None = None,
parameter_types: dict[str, SqlType.Type] | None = None,
view_parameters: dict[str, str] | None = None,
app_profile_id: str | None = None,
operation_timeout: float = 600,
attempt_timeout: float | None = 20,
Expand Down Expand Up @@ -758,6 +760,8 @@ async def execute_query(
Required to contain entries only for parameters whose type cannot be
detected automatically (i.e. the value can be None, an empty list or
an empty dict).
view_parameters: Dictionary with values for all view parameters. Currently only
string values are supported.
app_profile_id: The app profile to associate with requests.
https://cloud.google.com/bigtable/docs/app-profiles
operation_timeout: the time budget for the entire executeQuery operation, in seconds.
Expand Down Expand Up @@ -883,12 +887,14 @@ async def execute_query(
retryable_excs = [_get_error_type(e) for e in retryable_errors]

pb_params = _format_execute_query_params(parameters, parameter_types)
pb_view_params = _format_execute_query_view_params(view_parameters)

request_body = {
"instance_name": instance_name,
"app_profile_id": app_profile_id,
"prepared_query": prepare_result.prepared_query,
"params": pb_params,
"view_parameters": pb_view_params,
}
operation_timeout, attempt_timeout = _align_timeouts(
operation_timeout, attempt_timeout
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@
)
from google.cloud.bigtable.data.execute_query._parameters_formatting import (
_format_execute_query_params,
_format_execute_query_view_params,
_to_param_types,
)
from google.cloud.bigtable.data.execute_query.metadata import (
Expand Down Expand Up @@ -532,6 +533,7 @@ def execute_query(
*,
parameters: dict[str, ExecuteQueryValueType] | None = None,
parameter_types: dict[str, SqlType.Type] | None = None,
view_parameters: dict[str, str] | None = None,
app_profile_id: str | None = None,
operation_timeout: float = 600,
attempt_timeout: float | None = 20,
Expand Down Expand Up @@ -572,6 +574,8 @@ def execute_query(
Required to contain entries only for parameters whose type cannot be
detected automatically (i.e. the value can be None, an empty list or
an empty dict).
view_parameters: Dictionary with values for all view parameters. Currently only
string values are supported.
app_profile_id: The app profile to associate with requests.
https://cloud.google.com/bigtable/docs/app-profiles
operation_timeout: the time budget for the entire executeQuery operation, in seconds.
Expand Down Expand Up @@ -692,11 +696,13 @@ def execute_query(
prepare_metadata = _pb_metadata_to_metadata_types(prepare_result.metadata)
retryable_excs = [_get_error_type(e) for e in retryable_errors]
pb_params = _format_execute_query_params(parameters, parameter_types)
pb_view_params = _format_execute_query_view_params(view_parameters)
request_body = {
"instance_name": instance_name,
"app_profile_id": app_profile_id,
"prepared_query": prepare_result.prepared_query,
"params": pb_params,
"view_parameters": pb_view_params,
}
operation_timeout, attempt_timeout = _align_timeouts(
operation_timeout, attempt_timeout
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,27 @@
from google.cloud.bigtable_v2.types.data import Value


def _format_execute_query_view_params(
view_parameters: Optional[Dict[str, str]],
) -> Dict[str, Value]:
"""
Takes a dictionary of view_param_name -> view_param_value (string) and formats
them into a dictionary of string-typed Value objects.
"""
if not view_parameters:
return {}

result_values = {}
for key, value in view_parameters.items():
if not isinstance(value, str):
raise TypeError(
f"View parameter {key} must be a string, got {type(value).__name__}"
)
result_values[key] = Value(string_value=value)

return result_values


def _format_execute_query_params(
params: Optional[Dict[str, ExecuteQueryValueType]],
parameter_types: Optional[Dict[str, SqlType.Type]],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3462,6 +3462,44 @@ async def test_execute_query_with_params(
assert execute_query_mock.call_count == 1
assert prepare_mock.call_count == 1

@CrossSync.pytest
async def test_execute_query_with_view_parameters(
self, client, execute_query_mock, prepare_mock
):
values = [
*chunked_responses(2, str_val("test2"), int_val(9), token=b"r2"),
]
execute_query_mock.return_value = self._make_gapic_stream(values)
query_str = f"SELECT a, b FROM {self.TABLE_NAME} WHERE user_id = VIEW_PARAMETERS('user_id')"
result = await client.execute_query(
query_str,
self.INSTANCE_NAME,
view_parameters={"user_id": "alice"},
)
results = [r async for r in result]
assert len(results) == 1
assert results[0]["a"] == "test2"
assert results[0]["b"] == 9
assert execute_query_mock.call_count == 1
assert prepare_mock.call_count == 1
assert prepare_mock.call_args[1]["request"]["query"] == query_str

request = execute_query_mock.call_args[0][0]
assert "user_id" in request.view_parameters
assert request.view_parameters["user_id"].string_value == "alice"

@CrossSync.pytest
async def test_execute_query_with_view_parameters_invalid_type(
self, client, execute_query_mock, prepare_mock
):
with pytest.raises(TypeError) as e:
await client.execute_query(
f"SELECT a, b FROM {self.TABLE_NAME}",
self.INSTANCE_NAME,
view_parameters={"user_id": 123},
)
assert "View parameter user_id must be a string, got int" in str(e.value)

@CrossSync.pytest
async def test_execute_query_error_before_metadata(
self, client, execute_query_mock, prepare_mock
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2959,6 +2959,42 @@ def test_execute_query_with_params(self, client, execute_query_mock, prepare_moc
assert execute_query_mock.call_count == 1
assert prepare_mock.call_count == 1

def test_execute_query_with_view_parameters(
self, client, execute_query_mock, prepare_mock
):
values = [
*chunked_responses(2, str_val("test2"), int_val(9), token=b"r2"),
]
execute_query_mock.return_value = self._make_gapic_stream(values)
query_str = f"SELECT a, b FROM {self.TABLE_NAME} WHERE user_id = VIEW_PARAMETERS('user_id')"
result = client.execute_query(
query_str,
self.INSTANCE_NAME,
view_parameters={"user_id": "alice"},
)
results = [r for r in result]
assert len(results) == 1
assert results[0]["a"] == "test2"
assert results[0]["b"] == 9
assert execute_query_mock.call_count == 1
assert prepare_mock.call_count == 1
assert prepare_mock.call_args[1]["request"]["query"] == query_str

request = execute_query_mock.call_args[0][0]
assert "user_id" in request.view_parameters
assert request.view_parameters["user_id"].string_value == "alice"

def test_execute_query_with_view_parameters_invalid_type(
self, client, execute_query_mock, prepare_mock
):
with pytest.raises(TypeError) as e:
client.execute_query(
f"SELECT a, b FROM {self.TABLE_NAME}",
self.INSTANCE_NAME,
view_parameters={"user_id": 123},
)
assert "View parameter user_id must be a string, got int" in str(e.value)

def test_execute_query_error_before_metadata(
self, client, execute_query_mock, prepare_mock
):
Expand Down
Loading