|
| 1 | +import tarfile |
| 2 | +import tempfile |
| 3 | + |
| 4 | +from kubernetes import client, config |
| 5 | +from kubernetes.stream import stream |
| 6 | +from tests import test_logger |
| 7 | + |
| 8 | +logger = test_logger.get_test_logger(__name__) |
| 9 | + |
| 10 | +TOOLS_POD_NAME = "mongodb-tools-pod" |
| 11 | +TOOLS_POD_IMAGE = "mongodb/mongodb-community-server:8.0-ubi9" |
| 12 | + |
| 13 | + |
| 14 | +class ToolsPod: |
| 15 | + """A pod running MongoDB tools for executing commands like mongorestore inside the cluster.""" |
| 16 | + |
| 17 | + def __init__(self, namespace: str): |
| 18 | + self.namespace = namespace |
| 19 | + self.pod_name = TOOLS_POD_NAME |
| 20 | + config.load_incluster_config() |
| 21 | + self.core_v1 = client.CoreV1Api() |
| 22 | + |
| 23 | + def run_command(self, cmd: list[str]): |
| 24 | + """Execute a command in the tools pod and return the output.""" |
| 25 | + logger.debug(f"Running command in {self.pod_name}: {' '.join(cmd)}") |
| 26 | + resp = stream( |
| 27 | + self.core_v1.connect_get_namespaced_pod_exec, |
| 28 | + self.pod_name, |
| 29 | + self.namespace, |
| 30 | + command=cmd, |
| 31 | + stderr=True, |
| 32 | + stdin=False, |
| 33 | + stdout=True, |
| 34 | + tty=False, |
| 35 | + ) |
| 36 | + logger.debug(f"Command output: {resp}") |
| 37 | + return resp |
| 38 | + |
| 39 | + def copy_file_to_pod(self, src_path: str, dest_path: str): |
| 40 | + """Copy a file from the local filesystem to the tools pod.""" |
| 41 | + logger.debug(f"Copying {src_path} to {self.pod_name}:{dest_path}") |
| 42 | + |
| 43 | + # Create a tar archive containing the file |
| 44 | + with tempfile.NamedTemporaryFile(suffix=".tar") as tar_file: |
| 45 | + with tarfile.open(tar_file.name, "w") as tar: |
| 46 | + tar.add(src_path, arcname=dest_path.split("/")[-1]) |
| 47 | + |
| 48 | + tar_file.seek(0) |
| 49 | + tar_data = tar_file.read() |
| 50 | + |
| 51 | + # Extract the tar archive in the pod |
| 52 | + exec_command = ["tar", "xf", "-", "-C", "/".join(dest_path.split("/")[:-1]) or "/"] |
| 53 | + resp = stream( |
| 54 | + self.core_v1.connect_get_namespaced_pod_exec, |
| 55 | + self.pod_name, |
| 56 | + self.namespace, |
| 57 | + command=exec_command, |
| 58 | + stderr=True, |
| 59 | + stdin=True, |
| 60 | + stdout=True, |
| 61 | + tty=False, |
| 62 | + _preload_content=False, |
| 63 | + ) |
| 64 | + |
| 65 | + # Send the tar data |
| 66 | + resp.write_stdin(tar_data) |
| 67 | + resp.close() |
| 68 | + logger.debug(f"File copied to {self.pod_name}:{dest_path}") |
| 69 | + |
| 70 | + def run_pod_and_wait(self): |
| 71 | + """Create the tools pod and wait for it to be ready.""" |
| 72 | + pod_body = client.V1Pod( |
| 73 | + api_version="v1", |
| 74 | + kind="Pod", |
| 75 | + metadata=client.V1ObjectMeta(name=self.pod_name, labels={"app": "mongodb-tools"}), |
| 76 | + spec=client.V1PodSpec( |
| 77 | + containers=[ |
| 78 | + client.V1Container( |
| 79 | + name="mongodb-tools", |
| 80 | + image=TOOLS_POD_IMAGE, |
| 81 | + command=["/bin/bash", "-c"], |
| 82 | + args=["sleep infinity"], |
| 83 | + ) |
| 84 | + ], |
| 85 | + restart_policy="Never", |
| 86 | + ), |
| 87 | + ) |
| 88 | + |
| 89 | + try: |
| 90 | + self.core_v1.create_namespaced_pod(namespace=self.namespace, body=pod_body) |
| 91 | + logger.info(f"Created {self.pod_name} in namespace {self.namespace}") |
| 92 | + except client.exceptions.ApiException as e: |
| 93 | + if e.status == 409: |
| 94 | + logger.info(f"Pod {self.pod_name} already exists") |
| 95 | + else: |
| 96 | + raise |
| 97 | + |
| 98 | + # Wait for pod to be ready |
| 99 | + from kubernetes.watch import Watch |
| 100 | + |
| 101 | + w = Watch() |
| 102 | + for event in w.stream( |
| 103 | + self.core_v1.list_namespaced_pod, |
| 104 | + namespace=self.namespace, |
| 105 | + label_selector="app=mongodb-tools", |
| 106 | + timeout_seconds=120, |
| 107 | + ): |
| 108 | + pod = event["object"] |
| 109 | + if pod.status.phase == "Running": |
| 110 | + # Check if container is ready |
| 111 | + if pod.status.container_statuses: |
| 112 | + for container_status in pod.status.container_statuses: |
| 113 | + if container_status.ready: |
| 114 | + logger.info(f"{self.pod_name} is ready") |
| 115 | + w.stop() |
| 116 | + return |
| 117 | + raise TimeoutError(f"Timed out waiting for {self.pod_name} to be ready") |
| 118 | + |
| 119 | + |
| 120 | +def get_tools_pod(namespace: str) -> ToolsPod: |
| 121 | + """Create and return a ready tools pod in the given namespace.""" |
| 122 | + tools_pod = ToolsPod(namespace) |
| 123 | + tools_pod.run_pod_and_wait() |
| 124 | + return tools_pod |
0 commit comments