PYTHON-3837 Driver Container and Kubernetes Awareness (#1418)

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

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):