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

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 Dict, Optional 

16 

17from aiohttp_middlewares.annotations import UrlCollection 

18from aiohttp_middlewares.cors import ACCESS_CONTROL_ALLOW_CREDENTIALS, ACCESS_CONTROL_ALLOW_ORIGIN 

19 

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

21from buildgrid.server.settings import ( 

22 BROWSER_BLOB_CACHE_MAX_LENGTH, 

23 BROWSER_BLOB_CACHE_TTL, 

24 BROWSER_RESULT_CACHE_MAX_LENGTH, 

25 BROWSER_RESULT_CACHE_TTL, 

26) 

27from buildgrid.server.utils.async_lru_cache import LruCache 

28 

29TARBALL_DIRECTORY_PREFIX = "bgd-browser-tarball-" 

30 

31 

32def get_cors_headers(origin: Optional[str], origins: UrlCollection, allow_all: bool) -> Dict[str, str]: 

33 headers = {} 

34 if origin: 

35 headers[ACCESS_CONTROL_ALLOW_CREDENTIALS] = "true" 

36 if origin in origins: 

37 headers[ACCESS_CONTROL_ALLOW_ORIGIN] = origin 

38 elif allow_all: 

39 headers[ACCESS_CONTROL_ALLOW_ORIGIN] = "*" 

40 return headers 

41 

42 

43class ResponseCache: 

44 def __init__(self) -> None: 

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

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

47 

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

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

50 

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

52 return await self._action_results.get(key) 

53 

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

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

56 

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

58 return await self._blobs.get(key)