• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2023 gRPC authors.
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"""Helper to read observability config."""
15
16from dataclasses import dataclass
17from dataclasses import field
18import json
19import os
20from typing import Mapping, Optional
21
22GRPC_GCP_OBSERVABILITY_CONFIG_FILE_ENV = "GRPC_GCP_OBSERVABILITY_CONFIG_FILE"
23GRPC_GCP_OBSERVABILITY_CONFIG_ENV = "GRPC_GCP_OBSERVABILITY_CONFIG"
24
25
26@dataclass
27class GcpObservabilityConfig:
28    project_id: str = ""
29    stats_enabled: bool = False
30    tracing_enabled: bool = False
31    labels: Optional[Mapping[str, str]] = field(default_factory=dict)
32    sampling_rate: Optional[float] = 0.0
33
34    def load_from_string_content(self, config_contents: str) -> None:
35        """Loads the configuration from a string.
36
37        Args:
38            config_contents: The configuration string.
39
40        Raises:
41            ValueError: If the configuration is invalid.
42        """
43        try:
44            config_json = json.loads(config_contents)
45        except json.decoder.JSONDecodeError:
46            raise ValueError("Failed to load Json configuration.")
47
48        if config_json and not isinstance(config_json, dict):
49            raise ValueError("Found invalid configuration.")
50
51        self.project_id = config_json.get("project_id", "")
52        self.labels = config_json.get("labels", {})
53        self.stats_enabled = "cloud_monitoring" in config_json.keys()
54        self.tracing_enabled = "cloud_trace" in config_json.keys()
55        tracing_config = config_json.get("cloud_trace", {})
56        self.sampling_rate = tracing_config.get("sampling_rate", 0.0)
57
58
59def read_config() -> GcpObservabilityConfig:
60    """Reads the GCP observability config from the environment variables.
61
62    Returns:
63        The GCP observability config.
64
65    Raises:
66        ValueError: If the configuration is invalid.
67    """
68    config_contents = _get_gcp_observability_config_contents()
69    config = GcpObservabilityConfig()
70    config.load_from_string_content(config_contents)
71
72    if not config.project_id:
73        # Get project ID from GCP environment variables since project ID was not
74        # set it in the GCP observability config.
75        config.project_id = _get_gcp_project_id_from_env_var()
76        if not config.project_id:
77            # Could not find project ID from GCP environment variables either.
78            raise ValueError("GCP Project ID not found.")
79    return config
80
81
82def _get_gcp_project_id_from_env_var() -> Optional[str]:
83    """Gets the project ID from the GCP environment variables.
84
85    Returns:
86        The project ID, or an empty string if the project ID could not be found.
87    """
88
89    project_id = ""
90    project_id = os.getenv("GCP_PROJECT")
91    if project_id:
92        return project_id
93
94    project_id = os.getenv("GCLOUD_PROJECT")
95    if project_id:
96        return project_id
97
98    project_id = os.getenv("GOOGLE_CLOUD_PROJECT")
99    if project_id:
100        return project_id
101
102    return project_id
103
104
105def _get_gcp_observability_config_contents() -> str:
106    """Get the contents of the observability config from environment variable or file.
107
108    Returns:
109        The content from environment variable.
110
111    Raises:
112        ValueError: If no configuration content was found.
113    """
114
115    contents_str = ""
116    # First try get config from GRPC_GCP_OBSERVABILITY_CONFIG_FILE_ENV.
117    config_path = os.getenv(GRPC_GCP_OBSERVABILITY_CONFIG_FILE_ENV)
118    if config_path:
119        with open(config_path, "r") as f:
120            contents_str = f.read()
121
122    # Next, try GRPC_GCP_OBSERVABILITY_CONFIG_ENV env var.
123    if not contents_str:
124        contents_str = os.getenv(GRPC_GCP_OBSERVABILITY_CONFIG_ENV)
125
126    if not contents_str:
127        raise ValueError("Configuration content not found.")
128
129    return contents_str
130