forked from googleapis/python-bigtable
/
row_set.py
210 lines (158 loc) · 7.12 KB
/
row_set.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""User-friendly container for Google Cloud Bigtable RowSet """
from google.cloud._helpers import _to_bytes
class RowSet(object):
""" Convenience wrapper of google.bigtable.v2.RowSet
Useful for creating a set of row keys and row ranges, which can
be passed to yield_rows method of class:`.Table.yield_rows`.
"""
def __init__(self):
self.row_keys = []
self.row_ranges = []
def __eq__(self, other):
if not isinstance(other, self.__class__):
return NotImplemented
if len(other.row_keys) != len(self.row_keys):
return False
if len(other.row_ranges) != len(self.row_ranges):
return False
if not set(other.row_keys) == set(self.row_keys):
return False
if not set(other.row_ranges) == set(self.row_ranges):
return False
return True
def __ne__(self, other):
return not self == other
def add_row_key(self, row_key):
"""Add row key to row_keys list.
For example:
.. literalinclude:: snippets_table.py
:start-after: [START bigtable_add_row_key]
:end-before: [END bigtable_add_row_key]
:type row_key: bytes
:param row_key: The key of a row to read
"""
self.row_keys.append(row_key)
def add_row_range(self, row_range):
"""Add row_range to row_ranges list.
For example:
.. literalinclude:: snippets_table.py
:start-after: [START bigtable_add_row_range]
:end-before: [END bigtable_add_row_range]
:type row_range: class:`RowRange`
:param row_range: The row range object having start and end key
"""
self.row_ranges.append(row_range)
def add_row_range_from_keys(
self, start_key=None, end_key=None, start_inclusive=True, end_inclusive=False
):
"""Add row range to row_ranges list from the row keys
For example:
.. literalinclude:: snippets_table.py
:start-after: [START bigtable_row_range_from_keys]
:end-before: [END bigtable_row_range_from_keys]
:type start_key: bytes
:param start_key: (Optional) Start key of the row range. If left empty,
will be interpreted as the empty string.
:type end_key: bytes
:param end_key: (Optional) End key of the row range. If left empty,
will be interpreted as the empty string and range will
be unbounded on the high end.
:type start_inclusive: bool
:param start_inclusive: (Optional) Whether the ``start_key`` should be
considered inclusive. The default is True (inclusive).
:type end_inclusive: bool
:param end_inclusive: (Optional) Whether the ``end_key`` should be
considered inclusive. The default is False (exclusive).
"""
row_range = RowRange(start_key, end_key, start_inclusive, end_inclusive)
self.row_ranges.append(row_range)
def add_row_range_by_prefix_from_keys(self, row_key_prefix):
"""Add row range to row_ranges list that start with the row_key_prefix from the row keys
For example:
.. literalinclude:: snippets_table.py
:start-after: [START bigtable_row_range_by_prefix_from_keys]
:end-before: [END bigtable_row_range_by_prefix_from_keys]
:type row_key_prefix: str
:param row_key_prefix: To retrieve all rows that start with this row key prefix.
Prefix cannot be zero length."""
end_key = row_key_prefix[:-1] + chr(ord(row_key_prefix[-1]) + 1)
self.add_row_range_from_keys(
row_key_prefix.encode("utf-8"), end_key.encode("utf-8")
)
def _update_message_request(self, message):
"""Add row keys and row range to given request message
:type message: class:`data_messages_v2_pb2.ReadRowsRequest`
:param message: The ``ReadRowsRequest`` protobuf
"""
for each in self.row_keys:
message.rows.row_keys.append(_to_bytes(each))
for each in self.row_ranges:
r_kwrags = each.get_range_kwargs()
message.rows.row_ranges.add(**r_kwrags)
class RowRange(object):
""" Convenience wrapper of google.bigtable.v2.RowRange
:type start_key: bytes
:param start_key: (Optional) Start key of the row range. If left empty,
will be interpreted as the empty string.
:type end_key: bytes
:param end_key: (Optional) End key of the row range. If left empty,
will be interpreted as the empty string and range will
be unbounded on the high end.
:type start_inclusive: bool
:param start_inclusive: (Optional) Whether the ``start_key`` should be
considered inclusive. The default is True (inclusive).
:type end_inclusive: bool
:param end_inclusive: (Optional) Whether the ``end_key`` should be
considered inclusive. The default is False (exclusive).
"""
def __init__(
self, start_key=None, end_key=None, start_inclusive=True, end_inclusive=False
):
self.start_key = start_key
self.start_inclusive = start_inclusive
self.end_key = end_key
self.end_inclusive = end_inclusive
def _key(self):
"""A tuple key that uniquely describes this field.
Used to compute this instance's hashcode and evaluate equality.
Returns:
Tuple[str]: The contents of this :class:`.RowRange`.
"""
return (self.start_key, self.start_inclusive, self.end_key, self.end_inclusive)
def __hash__(self):
return hash(self._key())
def __eq__(self, other):
if not isinstance(other, self.__class__):
return NotImplemented
return self._key() == other._key()
def __ne__(self, other):
return not self == other
def get_range_kwargs(self):
""" Convert row range object to dict which can be passed to
google.bigtable.v2.RowRange add method.
"""
range_kwargs = {}
if self.start_key is not None:
start_key_key = "start_key_open"
if self.start_inclusive:
start_key_key = "start_key_closed"
range_kwargs[start_key_key] = _to_bytes(self.start_key)
if self.end_key is not None:
end_key_key = "end_key_open"
if self.end_inclusive:
end_key_key = "end_key_closed"
range_kwargs[end_key_key] = _to_bytes(self.end_key)
return range_kwargs