Arber Xhindoli pushed to branch arber/91-get-tree at BuildGrid / buildgrid
Commits:
-
e5b09463
by Arber Xhindoli at 2018-11-19T16:01:08Z
2 changed files:
Changes:
| ... | ... | @@ -61,23 +61,35 @@ class ContentAddressableStorageInstance: |
| 61 | 61 |
def get_tree(self, request):
|
| 62 | 62 |
storage = self._storage
|
| 63 | 63 |
|
| 64 |
- # Create getTreeResponse message
|
|
| 65 |
- # TODO: handle page_token, page_size
|
|
| 66 | 64 |
response = re_pb2.GetTreeResponse()
|
| 65 |
+ page_size = request.page_size
|
|
| 66 |
+ next_page_token = ""
|
|
| 67 |
+ directories = []
|
|
| 67 | 68 |
|
| 68 |
- # Set to MAX_REQUEST_COUNT, will use in the future to limit size of GRPC response
|
|
| 69 |
- if not request.page_size:
|
|
| 70 |
- request.page_size = 500
|
|
| 69 |
+ if not page_size:
|
|
| 70 |
+ # MAX_REQUEST_COUNT
|
|
| 71 |
+ page_size = 500
|
|
| 71 | 72 |
|
| 72 | 73 |
def _get_tree(node_digest):
|
| 74 |
+ nonlocal directories, page_size, next_page_token
|
|
| 75 |
+ if next_page_token:
|
|
| 76 |
+ # next page token has been set unwind the stack.
|
|
| 77 |
+ return
|
|
| 78 |
+ if page_size <= 0:
|
|
| 79 |
+ # save the next digest hash in order to continue later
|
|
| 80 |
+ next_page_token = str(node_digest.hash)
|
|
| 81 |
+ return
|
|
| 73 | 82 |
directory_from_digest = storage.get_message(node_digest, re_pb2.Directory)
|
| 74 |
- directories = [directory_from_digest]
|
|
| 83 |
+ directories.append(directory_from_digest)
|
|
| 75 | 84 |
for directory in directory_from_digest.directories:
|
| 76 |
- directories.extend(_get_tree(directory.digest))
|
|
| 77 |
- return directories
|
|
| 85 |
+ page_size -= 1
|
|
| 86 |
+ _get_tree(directory.digest)
|
|
| 87 |
+ return
|
|
| 78 | 88 |
|
| 79 |
- response.directories.extend(_get_tree(request.root_digest))
|
|
| 80 |
- yield response
|
|
| 89 |
+ _get_tree(request.root_digest)
|
|
| 90 |
+ response.directories.extend(directories)
|
|
| 91 |
+ response.next_page_token = next_page_token
|
|
| 92 |
+ return response
|
|
| 81 | 93 |
|
| 82 | 94 |
|
| 83 | 95 |
class ByteStreamInstance:
|
| ... | ... | @@ -89,9 +89,17 @@ class ContentAddressableStorageService(remote_execution_pb2_grpc.ContentAddressa |
| 89 | 89 |
|
| 90 | 90 |
try:
|
| 91 | 91 |
instance = self._get_instance(request.instance_name)
|
| 92 |
- response = instance.get_tree(request)
|
|
| 93 |
- |
|
| 94 |
- return response
|
|
| 92 |
+ while True:
|
|
| 93 |
+ response = instance.get_tree(request)
|
|
| 94 |
+ if not response.next_page_token:
|
|
| 95 |
+ # This is ugly, but handles the case in which there is only one iteration
|
|
| 96 |
+ # of the while loop. We need to manually raise stopiteration from this generator.
|
|
| 97 |
+ # If return is called, before a yield is ever called, we get an empty response.
|
|
| 98 |
+ yield response
|
|
| 99 |
+ raise StopIteration # pylint: disable=stop-iteration-return
|
|
| 100 |
+ |
|
| 101 |
+ yield response
|
|
| 102 |
+ request.root_digest.hash = response.next_page_token
|
|
| 95 | 103 |
|
| 96 | 104 |
except InvalidArgumentError as e:
|
| 97 | 105 |
self.__logger.error(e)
|