Coverage for /builds/BuildGrid/buildgrid/buildgrid/browser/utils.py: 57.58%

1# Copyright (C) 2021 Bloomberg LP 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# <http://www.apache.org/licenses/LICENSE-2.0> 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15from typing import Optional 

16 

17from aiohttp_middlewares.cors import ACCESS_CONTROL_ALLOW_CREDENTIALS, ACCESS_CONTROL_ALLOW_ORIGIN 

18 

19from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import ActionResult 

20from buildgrid._protos.google.longrunning.operations_pb2 import Operation 

21from buildgrid.server.utils.async_lru_cache import LruCache 

22from buildgrid.settings import ( 

23 BROWSER_BLOB_CACHE_MAX_LENGTH, 

24 BROWSER_BLOB_CACHE_TTL, 

25 BROWSER_OPERATION_CACHE_MAX_LENGTH, 

26 BROWSER_OPERATION_CACHE_TTL, 

27 BROWSER_RESULT_CACHE_MAX_LENGTH, 

28 BROWSER_RESULT_CACHE_TTL 

29) 

30 

31 

32TARBALL_DIRECTORY_PREFIX = 'bgd-browser-tarball-' 

33 

34 

35def get_cors_headers(origin, origins, allow_all): 

36 headers = {} 

37 if origin: 

38 headers[ACCESS_CONTROL_ALLOW_CREDENTIALS] = "true" 

39 if origin in origins: 

40 headers[ACCESS_CONTROL_ALLOW_ORIGIN] = origin 

41 elif allow_all: 

42 headers[ACCESS_CONTROL_ALLOW_ORIGIN] = "*" 

43 return headers 

44 

45 

46class ResponseCache: 

47 

48 def __init__(self): 

49 self._action_results: LruCache[ActionResult] = LruCache(BROWSER_RESULT_CACHE_MAX_LENGTH) 

50 self._blobs: LruCache[bytes] = LruCache(BROWSER_BLOB_CACHE_MAX_LENGTH) 

51 self._operations: LruCache[Operation] = LruCache(BROWSER_OPERATION_CACHE_MAX_LENGTH) 

52 

53 async def store_action_result(self, key: str, result: ActionResult) -> None: 

54 await self._action_results.set(key, result, BROWSER_RESULT_CACHE_TTL) 

55 

56 async def get_action_result(self, key: str) -> Optional[ActionResult]: 

57 return await self._action_results.get(key) 

58 

59 async def store_blob(self, key: str, blob: bytes) -> None: 

60 await self._blobs.set(key, blob, BROWSER_BLOB_CACHE_TTL) 

61 

62 async def get_blob(self, key: str) -> Optional[bytes]: 

63 return await self._blobs.get(key) 

64 

65 async def store_operation(self, key: str, operation: Operation) -> None: 

66 await self._operations.set(key, operation, BROWSER_OPERATION_CACHE_TTL) 

67 

68 async def get_operation(self, key: str) -> Optional[Operation]: 

69 return await self._operations.get(key)