Skip to content

Commit 89b8da8

Browse files
jackdingilianjackdingilian
authored
feat: Add support for array and float32 SQL query params (#1078)
1 parent c94ef3f commit 89b8da8

File tree

9 files changed

+381
-36
lines changed

9 files changed

+381
-36
lines changed

‎google/cloud/bigtable/data/_async/client.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -569,6 +569,8 @@ async def execute_query(
569569
will be chained with a RetryExceptionGroup containing GoogleAPIError exceptions
570570
from any retries that failed
571571
google.api_core.exceptions.GoogleAPIError: raised if the request encounters an unrecoverable error
572+
google.cloud.bigtable.data.exceptions.ParameterTypeInferenceFailed: Raised if
573+
a parameter is passed without an explicit type, and the type cannot be infered
572574
"""
573575
warnings.warn(
574576
"ExecuteQuery is in preview and may change in the future.",

‎google/cloud/bigtable/data/_sync_autogen/client.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -439,6 +439,8 @@ def execute_query(
439439
will be chained with a RetryExceptionGroup containing GoogleAPIError exceptions
440440
from any retries that failed
441441
google.api_core.exceptions.GoogleAPIError: raised if the request encounters an unrecoverable error
442+
google.cloud.bigtable.data.exceptions.ParameterTypeInferenceFailed: Raised if
443+
a parameter is passed without an explicit type, and the type cannot be infered
442444
"""
443445
warnings.warn(
444446
"ExecuteQuery is in preview and may change in the future.",

‎google/cloud/bigtable/data/execute_query/_async/execute_query_iterator.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222
Tuple,
2323
TYPE_CHECKING,
2424
)
25-
2625
from google.api_core import retry as retries
2726

2827
from google.cloud.bigtable.data.execute_query._byte_cursor import _ByteCursor
@@ -116,7 +115,6 @@ def __init__(
116115
exception_factory=_retry_exception_factory,
117116
)
118117
self._req_metadata = req_metadata
119-
120118
try:
121119
self._register_instance_task = CrossSync.create_task(
122120
self._client._register_instance,

‎google/cloud/bigtable/data/execute_query/_parameters_formatting.py

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,14 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
from typing import Any, Dict, Optional
1615
import datetime
16+
from typing import Any, Dict, Optional
17+
1718
from google.api_core.datetime_helpers import DatetimeWithNanoseconds
19+
1820
from google.cloud.bigtable.data.exceptions import ParameterTypeInferenceFailed
19-
from google.cloud.bigtable.data.execute_query.values import ExecuteQueryValueType
2021
from google.cloud.bigtable.data.execute_query.metadata import SqlType
22+
from google.cloud.bigtable.data.execute_query.values import ExecuteQueryValueType
2123

2224

2325
def _format_execute_query_params(
@@ -48,7 +50,6 @@ def _format_execute_query_params(
4850
parameter_types = parameter_types or {}
4951

5052
result_values = {}
51-
5253
for key, value in params.items():
5354
user_provided_type = parameter_types.get(key)
5455
try:
@@ -109,6 +110,16 @@ def _detect_type(value: ExecuteQueryValueType) -> SqlType.Type:
109110
"Cannot infer type of None, please provide the type manually."
110111
)
111112

113+
if isinstance(value, list):
114+
raise ParameterTypeInferenceFailed(
115+
"Cannot infer type of ARRAY parameters, please provide the type manually."
116+
)
117+
118+
if isinstance(value, float):
119+
raise ParameterTypeInferenceFailed(
120+
"Cannot infer type of float, must specify either FLOAT32 or FLOAT64 type manually."
121+
)
122+
112123
for field_type, type_dict in _TYPES_TO_TYPE_DICTS:
113124
if isinstance(value, field_type):
114125
return type_dict

‎google/cloud/bigtable/data/execute_query/_query_result_parsing_utils.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
SqlType.Bytes: "bytes_value",
2323
SqlType.String: "string_value",
2424
SqlType.Int64: "int_value",
25+
SqlType.Float32: "float_value",
2526
SqlType.Float64: "float_value",
2627
SqlType.Bool: "bool_value",
2728
SqlType.Timestamp: "timestamp_value",

‎google/cloud/bigtable/data/execute_query/metadata.py

Lines changed: 30 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -21,23 +21,16 @@
2121
"""
2222

2323
from collections import defaultdict
24-
from typing import (
25-
Optional,
26-
List,
27-
Dict,
28-
Set,
29-
Type,
30-
Union,
31-
Tuple,
32-
Any,
33-
)
24+
import datetime
25+
from typing import Any, Dict, List, Optional, Set, Tuple, Type, Union
26+
27+
from google.api_core.datetime_helpers import DatetimeWithNanoseconds
28+
from google.protobuf import timestamp_pb2 # type: ignore
29+
from google.type import date_pb2 # type: ignore
30+
3431
from google.cloud.bigtable.data.execute_query.values import _NamedList
3532
from google.cloud.bigtable_v2 import ResultSetMetadata
3633
from google.cloud.bigtable_v2 import Type as PBType
37-
from google.type import date_pb2 # type: ignore
38-
from google.protobuf import timestamp_pb2 # type: ignore
39-
from google.api_core.datetime_helpers import DatetimeWithNanoseconds
40-
import datetime
4134

4235

4336
class SqlType:
@@ -127,6 +120,8 @@ class Array(Type):
127120
def __init__(self, element_type: "SqlType.Type"):
128121
if isinstance(element_type, SqlType.Array):
129122
raise ValueError("Arrays of arrays are not supported.")
123+
if isinstance(element_type, SqlType.Map):
124+
raise ValueError("Arrays of Maps are not supported.")
130125
self._element_type = element_type
131126

132127
@property
@@ -140,10 +135,21 @@ def from_pb_type(cls, type_pb: Optional[PBType] = None) -> "SqlType.Array":
140135
return cls(_pb_type_to_metadata_type(type_pb.array_type.element_type))
141136

142137
def _to_value_pb_dict(self, value: Any):
143-
raise NotImplementedError("Array is not supported as a query parameter")
138+
if value is None:
139+
return {}
140+
141+
return {
142+
"array_value": {
143+
"values": [
144+
self.element_type._to_value_pb_dict(entry) for entry in value
145+
]
146+
}
147+
}
144148

145149
def _to_type_pb_dict(self) -> Dict[str, Any]:
146-
raise NotImplementedError("Array is not supported as a query parameter")
150+
return {
151+
"array_type": {"element_type": self.element_type._to_type_pb_dict()}
152+
}
147153

148154
def __eq__(self, other):
149155
return super().__eq__(other) and self.element_type == other.element_type
@@ -222,6 +228,13 @@ class Float64(Type):
222228
value_pb_dict_field_name = "float_value"
223229
type_field_name = "float64_type"
224230

231+
class Float32(Type):
232+
"""Float32 SQL type."""
233+
234+
expected_type = float
235+
value_pb_dict_field_name = "float_value"
236+
type_field_name = "float32_type"
237+
225238
class Bool(Type):
226239
"""Bool SQL type."""
227240

@@ -376,6 +389,7 @@ def _pb_metadata_to_metadata_types(
376389
"bytes_type": SqlType.Bytes,
377390
"string_type": SqlType.String,
378391
"int64_type": SqlType.Int64,
392+
"float32_type": SqlType.Float32,
379393
"float64_type": SqlType.Float64,
380394
"bool_type": SqlType.Bool,
381395
"timestamp_type": SqlType.Timestamp,

‎tests/system/data/test_system_async.py

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,16 @@
1414

1515
import pytest
1616
import asyncio
17+
import datetime
1718
import uuid
1819
import os
1920
from google.api_core import retry
2021
from google.api_core.exceptions import ClientError
2122

23+
from google.cloud.bigtable.data.execute_query.metadata import SqlType
2224
from google.cloud.bigtable.data.read_modify_write_rules import _MAX_INCREMENT_VALUE
2325
from google.cloud.environment_vars import BIGTABLE_EMULATOR
26+
from google.type import date_pb2
2427

2528
from google.cloud.bigtable.data._cross_sync import CrossSync
2629

@@ -1027,3 +1030,83 @@ async def test_execute_query_simple(self, client, table_id, instance_id):
10271030
row = rows[0]
10281031
assert row["a"] == 1
10291032
assert row["b"] == "foo"
1033+
1034+
@CrossSync.pytest
1035+
@pytest.mark.usefixtures("client")
1036+
@CrossSync.Retry(
1037+
predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
1038+
)
1039+
async def test_execute_query_params(self, client, table_id, instance_id):
1040+
query = (
1041+
"SELECT @stringParam AS strCol, @bytesParam as bytesCol, @int64Param AS intCol, "
1042+
"@float32Param AS float32Col, @float64Param AS float64Col, @boolParam AS boolCol, "
1043+
"@tsParam AS tsCol, @dateParam AS dateCol, @byteArrayParam AS byteArrayCol, "
1044+
"@stringArrayParam AS stringArrayCol, @intArrayParam AS intArrayCol, "
1045+
"@float32ArrayParam AS float32ArrayCol, @float64ArrayParam AS float64ArrayCol, "
1046+
"@boolArrayParam AS boolArrayCol, @tsArrayParam AS tsArrayCol, "
1047+
"@dateArrayParam AS dateArrayCol"
1048+
)
1049+
parameters = {
1050+
"stringParam": "foo",
1051+
"bytesParam": b"bar",
1052+
"int64Param": 12,
1053+
"float32Param": 1.1,
1054+
"float64Param": 1.2,
1055+
"boolParam": True,
1056+
"tsParam": datetime.datetime.fromtimestamp(1000, tz=datetime.timezone.utc),
1057+
"dateParam": datetime.date(2025, 1, 16),
1058+
"byteArrayParam": [b"foo", b"bar", None],
1059+
"stringArrayParam": ["foo", "bar", None],
1060+
"intArrayParam": [1, None, 2],
1061+
"float32ArrayParam": [1.2, None, 1.3],
1062+
"float64ArrayParam": [1.4, None, 1.5],
1063+
"boolArrayParam": [None, False, True],
1064+
"tsArrayParam": [
1065+
datetime.datetime.fromtimestamp(1000, tz=datetime.timezone.utc),
1066+
datetime.datetime.fromtimestamp(2000, tz=datetime.timezone.utc),
1067+
None,
1068+
],
1069+
"dateArrayParam": [
1070+
datetime.date(2025, 1, 16),
1071+
datetime.date(2025, 1, 17),
1072+
None,
1073+
],
1074+
}
1075+
param_types = {
1076+
"float32Param": SqlType.Float32(),
1077+
"float64Param": SqlType.Float64(),
1078+
"byteArrayParam": SqlType.Array(SqlType.Bytes()),
1079+
"stringArrayParam": SqlType.Array(SqlType.String()),
1080+
"intArrayParam": SqlType.Array(SqlType.Int64()),
1081+
"float32ArrayParam": SqlType.Array(SqlType.Float32()),
1082+
"float64ArrayParam": SqlType.Array(SqlType.Float64()),
1083+
"boolArrayParam": SqlType.Array(SqlType.Bool()),
1084+
"tsArrayParam": SqlType.Array(SqlType.Timestamp()),
1085+
"dateArrayParam": SqlType.Array(SqlType.Date()),
1086+
}
1087+
result = await client.execute_query(
1088+
query, instance_id, parameters=parameters, parameter_types=param_types
1089+
)
1090+
rows = [r async for r in result]
1091+
assert len(rows) == 1
1092+
row = rows[0]
1093+
assert row["strCol"] == parameters["stringParam"]
1094+
assert row["bytesCol"] == parameters["bytesParam"]
1095+
assert row["intCol"] == parameters["int64Param"]
1096+
assert row["float32Col"] == pytest.approx(parameters["float32Param"])
1097+
assert row["float64Col"] == pytest.approx(parameters["float64Param"])
1098+
assert row["boolCol"] == parameters["boolParam"]
1099+
assert row["tsCol"] == parameters["tsParam"]
1100+
assert row["dateCol"] == date_pb2.Date(year=2025, month=1, day=16)
1101+
assert row["stringArrayCol"] == parameters["stringArrayParam"]
1102+
assert row["byteArrayCol"] == parameters["byteArrayParam"]
1103+
assert row["intArrayCol"] == parameters["intArrayParam"]
1104+
assert row["float32ArrayCol"] == pytest.approx(parameters["float32ArrayParam"])
1105+
assert row["float64ArrayCol"] == pytest.approx(parameters["float64ArrayParam"])
1106+
assert row["boolArrayCol"] == parameters["boolArrayParam"]
1107+
assert row["tsArrayCol"] == parameters["tsArrayParam"]
1108+
assert row["dateArrayCol"] == [
1109+
date_pb2.Date(year=2025, month=1, day=16),
1110+
date_pb2.Date(year=2025, month=1, day=17),
1111+
None,
1112+
]

‎tests/system/data/test_system_autogen.py

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,15 @@
1616
# This file is automatically generated by CrossSync. Do not edit manually.
1717

1818
import pytest
19+
import datetime
1920
import uuid
2021
import os
2122
from google.api_core import retry
2223
from google.api_core.exceptions import ClientError
24+
from google.cloud.bigtable.data.execute_query.metadata import SqlType
2325
from google.cloud.bigtable.data.read_modify_write_rules import _MAX_INCREMENT_VALUE
2426
from google.cloud.environment_vars import BIGTABLE_EMULATOR
27+
from google.type import date_pb2
2528
from google.cloud.bigtable.data._cross_sync import CrossSync
2629
from . import TEST_FAMILY, TEST_FAMILY_2
2730

@@ -838,3 +841,74 @@ def test_execute_query_simple(self, client, table_id, instance_id):
838841
row = rows[0]
839842
assert row["a"] == 1
840843
assert row["b"] == "foo"
844+
845+
@pytest.mark.usefixtures("client")
846+
@CrossSync._Sync_Impl.Retry(
847+
predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
848+
)
849+
def test_execute_query_params(self, client, table_id, instance_id):
850+
query = "SELECT @stringParam AS strCol, @bytesParam as bytesCol, @int64Param AS intCol, @float32Param AS float32Col, @float64Param AS float64Col, @boolParam AS boolCol, @tsParam AS tsCol, @dateParam AS dateCol, @byteArrayParam AS byteArrayCol, @stringArrayParam AS stringArrayCol, @intArrayParam AS intArrayCol, @float32ArrayParam AS float32ArrayCol, @float64ArrayParam AS float64ArrayCol, @boolArrayParam AS boolArrayCol, @tsArrayParam AS tsArrayCol, @dateArrayParam AS dateArrayCol"
851+
parameters = {
852+
"stringParam": "foo",
853+
"bytesParam": b"bar",
854+
"int64Param": 12,
855+
"float32Param": 1.1,
856+
"float64Param": 1.2,
857+
"boolParam": True,
858+
"tsParam": datetime.datetime.fromtimestamp(1000, tz=datetime.timezone.utc),
859+
"dateParam": datetime.date(2025, 1, 16),
860+
"byteArrayParam": [b"foo", b"bar", None],
861+
"stringArrayParam": ["foo", "bar", None],
862+
"intArrayParam": [1, None, 2],
863+
"float32ArrayParam": [1.2, None, 1.3],
864+
"float64ArrayParam": [1.4, None, 1.5],
865+
"boolArrayParam": [None, False, True],
866+
"tsArrayParam": [
867+
datetime.datetime.fromtimestamp(1000, tz=datetime.timezone.utc),
868+
datetime.datetime.fromtimestamp(2000, tz=datetime.timezone.utc),
869+
None,
870+
],
871+
"dateArrayParam": [
872+
datetime.date(2025, 1, 16),
873+
datetime.date(2025, 1, 17),
874+
None,
875+
],
876+
}
877+
param_types = {
878+
"float32Param": SqlType.Float32(),
879+
"float64Param": SqlType.Float64(),
880+
"byteArrayParam": SqlType.Array(SqlType.Bytes()),
881+
"stringArrayParam": SqlType.Array(SqlType.String()),
882+
"intArrayParam": SqlType.Array(SqlType.Int64()),
883+
"float32ArrayParam": SqlType.Array(SqlType.Float32()),
884+
"float64ArrayParam": SqlType.Array(SqlType.Float64()),
885+
"boolArrayParam": SqlType.Array(SqlType.Bool()),
886+
"tsArrayParam": SqlType.Array(SqlType.Timestamp()),
887+
"dateArrayParam": SqlType.Array(SqlType.Date()),
888+
}
889+
result = client.execute_query(
890+
query, instance_id, parameters=parameters, parameter_types=param_types
891+
)
892+
rows = [r for r in result]
893+
assert len(rows) == 1
894+
row = rows[0]
895+
assert row["strCol"] == parameters["stringParam"]
896+
assert row["bytesCol"] == parameters["bytesParam"]
897+
assert row["intCol"] == parameters["int64Param"]
898+
assert row["float32Col"] == pytest.approx(parameters["float32Param"])
899+
assert row["float64Col"] == pytest.approx(parameters["float64Param"])
900+
assert row["boolCol"] == parameters["boolParam"]
901+
assert row["tsCol"] == parameters["tsParam"]
902+
assert row["dateCol"] == date_pb2.Date(year=2025, month=1, day=16)
903+
assert row["stringArrayCol"] == parameters["stringArrayParam"]
904+
assert row["byteArrayCol"] == parameters["byteArrayParam"]
905+
assert row["intArrayCol"] == parameters["intArrayParam"]
906+
assert row["float32ArrayCol"] == pytest.approx(parameters["float32ArrayParam"])
907+
assert row["float64ArrayCol"] == pytest.approx(parameters["float64ArrayParam"])
908+
assert row["boolArrayCol"] == parameters["boolArrayParam"]
909+
assert row["tsArrayCol"] == parameters["tsArrayParam"]
910+
assert row["dateArrayCol"] == [
911+
date_pb2.Date(year=2025, month=1, day=16),
912+
date_pb2.Date(year=2025, month=1, day=17),
913+
None,
914+
]

0 commit comments

Comments
 (0)