8000 improved proto conversion · googleapis/python-bigtable@59946f7 · GitHub
[go: up one dir, main page]
Skip to content

Commit 59946f7

Browse files
daniel-sanchedaniel-sanche
committed
improved proto conversion
1 parent 98510c3 commit 59946f7

File tree

2 files changed

+25
-18
lines changed

2 files changed

+25
-18
lines changed

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -68,10 +68,10 @@ def __init__(
6868
attempt_timeout, operation_timeout
6969
)
7070
self.operation_timeout = operation_timeout
71-
query_dict = query._to_dict() if not isinstance(query, dict) else query
72-
self.request = ReadRowsRequestPB(**query_dict, table_name=table.table_name)
73-
if table.app_profile_id:
74-
self.request.app_profile_id = table.app_profile_id
71+
if isinstance(query, dict):
72+
self.request = ReadRowsRequestPB(**query, table_name=table.table_name, app_profile_id=table.app_profile_id)
73+
else:
74+
self.request = query._to_pb(table)
7575
self.table = table
7676
self._predicate = retries.if_exception_type(
7777
core_exceptions.DeadlineExceeded,

‎google/cloud/bigtable/data/read_rows_query.py

Lines changed: 21 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,10 @@
2020
from dataclasses import dataclass
2121
from google.cloud.bigtable.data.row_filters import RowFilter
2222

23+
from google.cloud.bigtable_v2.types import RowRange as RowRangePB
24+
from google.cloud.bigtable_v2.types import RowSet as RowSetPB
25+
from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB
26+
2327
if TYPE_CHECKING:
2428
from google.cloud.bigtable.data import RowKeySamples
2529
from google.cloud.bigtable.data import ShardedQuery
@@ -279,16 +283,9 @@ def filter(self, row_filter: RowFilter | None):
279283
280284
Args:
281285
- row_filter: a RowFilter to apply to this query
282-
Can be a RowFilter object or a dict representation
283286
Returns:
284287
- a reference to this query for chaining
285288
"""
286-
if not (
287-
isinstance(row_filter, dict)
288-
or isinstance(row_filter, RowFilter)
289-
or row_filter is None
290-
):
291-
raise ValueError("row_filter must be a RowFilter or dict")
292289
self._filter = row_filter
293290

294291
def add_key(self, row_key: str | bytes):
@@ -312,20 +309,14 @@ def add_key(self, row_key: str | bytes):
312309

313310
def add_range(
314311
self,
315-
row_range: RowRange | dict[str, bytes],
312+
row_range: RowRange,
316313
):
317314
"""
318315
Add a range of row keys to this query.
319316
320317
Args:
321318
- row_range: a range of row keys to add to this query
322-
Can be a RowRange object or a dict representation in
323-
RowRange proto format
324319
"""
325-
if not (isinstance(row_range, dict) or isinstance(row_range, RowRange)):
326-
raise ValueError("row_range must be a RowRange or dict")
327-
if isinstance(row_range, dict):
328-
row_range = RowRange._from_dict(row_range)
329320
self.row_ranges.add(row_range)
330321

331322
def shard(self, shard_keys: RowKeySamples) -> ShardedQuery:
@@ -478,6 +469,22 @@ def _to_dict(self) -> dict[str, Any]:
478469
final_dict["rows_limit"] = self.limit
479470
return final_dict
480471

472+
def _to_pb(self, table) -> ReadRowsRequestPB:
473+
"""
474+
Convert this query into a dictionary that can be used to construct a
475+
ReadRowsRequest protobuf
476+
"""
477+
return ReadRowsRequestPB(
478+
table_name=table.table_name,
479+
app_profile_id=table.app_profile_id,
480+
rows=RowSetPB(
481+
row_keys=list(self.row_keys),
482+
row_ranges=[RowRangePB(*r._to_dict()) for r in self.row_ranges],
483+
),
484+
_filter=self._filter._to_pb() if self._filter else None,
485+
rows_limit=self.limit,
486+
)
487+
481488
def __eq__(self, other):
482489
"""
483490
RowRanges are equal if they have the same row keys, row ranges,

0 commit comments

Comments
 (0)
0