diff --git a/pymongo/pool.py b/pymongo/pool.py index a7b2fcb8c..a9271490f 100644 --- a/pymongo/pool.py +++ b/pymongo/pool.py @@ -25,6 +25,7 @@ import sys import threading import time import weakref +from pathlib import Path from typing import ( TYPE_CHECKING, Any, @@ -268,6 +269,25 @@ else: (platform.python_implementation(), ".".join(map(str, sys.version_info))) ) +DOCKER_ENV_PATH = "/.dockerenv" +ENV_VAR_K8S = "KUBERNETES_SERVICE_HOST" + +RUNTIME_NAME_DOCKER = "docker" +ORCHESTRATOR_NAME_K8S = "kubernetes" + + +def get_container_env_info() -> dict[str, str]: + """Returns the runtime and orchestrator of a container. + If neither value is present, the metadata client.env.container field will be omitted.""" + container = {} + + if Path(DOCKER_ENV_PATH).exists(): + container["runtime"] = RUNTIME_NAME_DOCKER + if os.getenv(ENV_VAR_K8S): + container["orchestrator"] = ORCHESTRATOR_NAME_K8S + + return container + def _is_lambda() -> bool: if os.getenv("AWS_LAMBDA_RUNTIME_API"): @@ -307,6 +327,9 @@ def _getenv_int(key: str) -> Optional[int]: def _metadata_env() -> dict[str, Any]: env: dict[str, Any] = {} + container = get_container_env_info() + if container: + env["container"] = container # Skip if multiple (or no) envs are matched. if (_is_lambda(), _is_azure_func(), _is_gcp_func(), _is_vercel()).count(True) != 1: return env diff --git a/test/test_client.py b/test/test_client.py index 8b2716a26..aceb15312 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -29,6 +29,7 @@ import sys import threading import time from typing import Iterable, Type, no_type_check +from unittest import mock from unittest.mock import patch sys.path[0:0] = [""] @@ -97,7 +98,7 @@ from pymongo.errors import ( ) from pymongo.mongo_client import MongoClient from pymongo.monitoring import ServerHeartbeatListener, ServerHeartbeatStartedEvent -from pymongo.pool import _METADATA, Connection, PoolOptions +from pymongo.pool import _METADATA, DOCKER_ENV_PATH, ENV_VAR_K8S, Connection, PoolOptions from pymongo.read_preferences import ReadPreference from pymongo.server_description import ServerDescription from pymongo.server_selectors import readable_server_selector, writable_server_selector @@ -347,6 +348,15 @@ class ClientUnitTest(unittest.TestCase): options = client._MongoClient__options self.assertEqual(options.pool_options.metadata, metadata) + @mock.patch.dict("os.environ", {ENV_VAR_K8S: "1"}) + def test_container_metadata(self): + metadata = copy.deepcopy(_METADATA) + metadata["env"] = {} + metadata["env"]["container"] = {"orchestrator": "kubernetes"} + client = MongoClient("mongodb://foo:27017/?appname=foobar&connect=false") + options = client._MongoClient__options + self.assertEqual(options.pool_options.metadata["env"], metadata["env"]) + def test_kwargs_codec_options(self): class MyFloatType: def __init__(self, x):