Mirror of https://github.com/roostorg/osprey
github.com/roostorg/osprey
use `read_rows` for `select_all` in Bigtable result store (#40)
+21
-6
+21
-6
osprey_worker/src/osprey/worker/lib/storage/stored_execution_result.py
+21
-6
osprey_worker/src/osprey/worker/lib/storage/stored_execution_result.py
···
11
11
import google.cloud.storage as storage
12
12
import pytz
13
13
from google.api_core import retry
14
-
from google.cloud.bigtable import row_filters
14
+
from google.cloud.bigtable import row_filters, row_set
15
15
from google.cloud.bigtable.row import Row
16
16
from minio import Minio
17
17
from minio.error import S3Error
···
193
193
# TODO: Add `select_*_minimal` methods
194
194
195
195
def select_many(self, action_ids: List[int]) -> List[Dict[str, Any]]:
196
-
return [
197
-
row
198
-
for row in gevent.pool.Pool(BIGTABLE_CONCURRENCY_LIMIT).imap(self.select_one, action_ids)
199
-
if row is not None
200
-
]
196
+
if not action_ids:
197
+
return []
198
+
199
+
row_set_obj = row_set.RowSet()
200
+
for action_id in action_ids:
201
+
row_set_obj.add_row_key(StoredExecutionResultBigTable._encode_action_id(action_id))
202
+
203
+
rows = osprey_bigtable.table('stored_execution_result').read_rows(
204
+
row_set=row_set,
205
+
filter_=row_filters.CellsColumnLimitFilter(1),
206
+
retry=self.retry_policy,
207
+
)
208
+
209
+
results: List[Dict[str, Any]] = []
210
+
for row in rows:
211
+
if not row:
212
+
continue
213
+
results.append(StoredExecutionResultBigTable._execution_result_dict_from_row(row))
214
+
215
+
return results
201
216
202
217
def insert(
203
218
self,