Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/cleanup/janitor/config.py: 98.11%

53 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2025-02-11 15:07 +0000

1# Copyright (C) 2024 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 

15 

16import os 

17from typing import Any 

18 

19import yaml 

20from pydantic import BaseModel, ConfigDict, Field 

21 

22from buildgrid.server.app.settings.config import populate_monitoring_config 

23from buildgrid.server.app.settings.parser import load_sql_connection, string_definitions 

24from buildgrid.server.monitoring import set_monitoring_bus 

25from buildgrid.server.sql.provider import SqlProvider 

26 

27 

28class S3Config(BaseModel): 

29 access_key: str 

30 bucket_regex: str 

31 endpoint: str 

32 path_prefix: str 

33 hash_prefix_size: int = Field(default=0) 

34 secret_key: str 

35 sleep_interval: int | None = Field(default=None) 

36 

37 

38class SQLStorageConfig(BaseModel): 

39 sql: SqlProvider 

40 sql_ro: SqlProvider 

41 sleep_interval: int | None = Field(default=None) 

42 batch_size: int 

43 

44 model_config = ConfigDict(arbitrary_types_allowed=True) 

45 

46 

47class RedisConfig(BaseModel): 

48 db: int | None = Field(default=None) 

49 dns_srv_record: str | None = Field(default=None) 

50 index_prefix: str 

51 key_batch_size: int 

52 password: str | None = Field(default=None) 

53 host: str | None = Field(default=None) 

54 port: int | None = Field(default=None) 

55 sentinel_master_name: str | None = Field(default=None) 

56 

57 

58class JanitorConfig(BaseModel): 

59 redis: RedisConfig | None = Field(default=None) 

60 sleep_interval: int | None = Field(default=None) 

61 s3: S3Config | None = Field(default=None) 

62 sql_storage_config: SQLStorageConfig | None = Field(default=None) 

63 sql_connection_string: str | None = Field(default=None) 

64 

65 

66def parse_janitor_config(path: str | bytes | os.PathLike[str]) -> JanitorConfig: 

67 class Loader(yaml.SafeLoader): 

68 def string_loader(self, node: yaml.MappingNode) -> Any: 

69 return string_definitions[node.tag](node.value) 

70 

71 for kind in string_definitions: 

72 Loader.add_constructor(kind, Loader.string_loader) 

73 

74 with open(path) as config_file: 

75 config = yaml.load(config_file, Loader=Loader) 

76 if "sql_storage_config" in config: 

77 config["sql_storage_config"]["sql"] = load_sql_connection(**config["sql_storage_config"]["sql"]) 

78 config["sql_storage_config"]["sql_ro"] = load_sql_connection(**config["sql_storage_config"]["sql_ro"]) 

79 if "monitoring" in config: 

80 monitoring_config = populate_monitoring_config(config["monitoring"]) 

81 if monitoring_config is not None: 

82 set_monitoring_bus(monitoring_config) 

83 return JanitorConfig(**config)