Skip to content

Commit d1c0dbb

Browse files
lsierantlsierant
committed
Run mongorestore in a tools pod
1 parent 540a9e6 commit d1c0dbb

11 files changed

+244
-95
lines changed

‎docker/mongodb-kubernetes-tests/tests/common/mongodb_tools_pod/__init__.py‎

Whitespace-only changes.

‎docker/mongodb-kubernetes-tests/tests/common/mongodb_tools_pod/mongodb_tools_pod.py‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
import tarfile
2+
import tempfile
3+
from typing import Optional
4+
5+
from kubernetes import client
6+
from kubernetes.client.exceptions import ApiException
7+
from kubernetes.stream import stream
8+
from tests import test_logger
9+
10+
logger = test_logger.get_test_logger(__name__)
11+
12+
13+
class ToolsPod:
14+
namespace: str
15+
pod_name: str
16+
container_name: str
17+
api_client: Optional[client.ApiClient] = None
18+
19+
def __init__(self, namespace: str, api_client: Optional[client.ApiClient] = None):
20+
self.namespace = namespace
21+
self.pod_name = "mongodb-tools-pod"
22+
self.container_name = "mongodb-tools"
23+
self.api_client = api_client
24+
25+
def run_command(self, cmd: list[str]):
26+
api_client = client.CoreV1Api(api_client=self.api_client)
27+
resp = stream(
28+
api_client.connect_get_namespaced_pod_exec,
29+
self.pod_name,
30+
self.namespace,
31+
container=self.container_name,
32+
command=cmd,
33+
stdout=True,
34+
stderr=True,
35+
_preload_content=False,
36+
)
37+
resp.run_forever(timeout=60)
38+
exit_code = resp.returncode
39+
output = resp.read_all()
40+
41+
if exit_code != 0:
42+
raise RuntimeError(f"Command failed with exit code {exit_code}: {output}")
43+
44+
return output
45+
46+
def copy_file_to_pod(self, src_path: str, dest_path: str):
47+
api_client = client.CoreV1Api(api_client=self.api_client)
48+
try:
49+
exec_command = ["tar", "xvf", "-", "-C", "/"]
50+
resp = stream(
51+
api_client.connect_get_namespaced_pod_exec,
52+
self.pod_name,
53+
self.namespace,
54+
container=self.container_name,
55+
command=exec_command,
56+
stderr=True,
57+
stdin=True,
58+
stdout=True,
59+
tty=False,
60+
_preload_content=False,
61+
)
62+
63+
stdout_output = []
64+
stderr_output = []
65+
66+
with tempfile.TemporaryFile() as tar_buffer:
67+
with tarfile.open(fileobj=tar_buffer, mode="w") as tar:
68+
tar.add(src_path, dest_path)
69+
70+
tar_buffer.seek(0)
71+
commands = [tar_buffer.read()]
72+
73+
while resp.is_open():
74+
resp.update(timeout=1)
75+
if resp.peek_stdout():
76+
stdout = resp.read_stdout()
77+
stdout_output.append(stdout)
78+
if resp.peek_stderr():
79+
stderr = resp.read_stderr()
80+
stderr_output.append(stderr)
81+
if commands:
82+
c = commands.pop(0)
83+
resp.write_stdin(c.decode())
84+
else:
85+
break
86+
87+
resp.run_forever(timeout=5)
88+
exit_code = resp.returncode
89+
if exit_code is not None and exit_code != 0:
90+
output = "".join(stdout_output + stderr_output)
91+
raise RuntimeError(
92+
f"Failed to copy file to pod: tar command exited with code {exit_code}\nOutput: {output}"
93+
)
94+
95+
except ApiException as e:
96+
raise Exception(f"Failed to copy file to the pod: {e}")
97+
98+
def run_pod_and_wait(self):
99+
from kubetester import get_pod_when_ready
100+
101+
pod_exists = False
102+
try:
103+
client.CoreV1Api(api_client=self.api_client).read_namespaced_pod(self.pod_name, self.namespace)
104+
pod_exists = True
105+
logger.debug(f"{self.pod_name} already exists in namespace {self.namespace}")
106+
except Exception:
107+
pass
108+
109+
if not pod_exists:
110+
pod_body = {
111+
"apiVersion": "v1",
112+
"kind": "Pod",
113+
"metadata": {
114+
"name": self.pod_name,
115+
"labels": {"app": self.pod_name},
116+
},
117+
"spec": {
118+
"containers": [
119+
{
120+
"name": self.container_name,
121+
"image": "mongodb/mongodb-community-server:8.0-ubi9",
122+
"command": ["/bin/bash", "-c"],
123+
"args": ["trap 'exit 0' SIGTERM; while true; do sleep 1; done"],
124+
}
125+
],
126+
"restartPolicy": "Never",
127+
},
128+
}
129+
client.CoreV1Api(api_client=self.api_client).create_namespaced_pod(self.namespace, pod_body)
130+
logger.info(f"Created {self.pod_name} in namespace {self.namespace}")
131+
132+
logger.debug(f"Waiting for tools pod ({self.namespace}/{self.pod_name}) to become ready")
133+
get_pod_when_ready(self.namespace, f"app={self.pod_name}", default_retry=60)
134+
135+
136+
def get_tools_pod(namespace: str) -> ToolsPod:
137+
tools_pod = ToolsPod(namespace)
138+
tools_pod.run_pod_and_wait()
139+
140+
return tools_pod

‎docker/mongodb-kubernetes-tests/tests/common/search/movies_search_helper.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import pymongo.errors
44
from kubetester import kubetester
55
from tests import test_logger
6+
from tests.common.mongodb_tools_pod import mongodb_tools_pod
67
from tests.common.search.search_tester import SearchTester
78

89
logger = test_logger.get_test_logger(__name__)
@@ -13,15 +14,19 @@ class SampleMoviesSearchHelper:
1314
db_name: str
1415
col_name: str
1516
archive_url: str
17+
tools_pod: mongodb_tools_pod.ToolsPod
1618

17-
def __init__(self, search_tester: SearchTester):
19+
def __init__(self, search_tester: SearchTester, tools_pod: mongodb_tools_pod.ToolsPod):
1820
self.search_tester = search_tester
1921
self.db_name = "sample_mflix"
2022
self.col_name = "movies"
23+
self.tools_pod = tools_pod
2124

2225
def restore_sample_database(self):
2326
self.search_tester.mongorestore_from_url(
24-
"https://atlas-education.s3.amazonaws.com/sample_mflix.archive", f"{self.db_name}.*"
27+
"https://atlas-education.s3.amazonaws.com/sample_mflix.archive",
28+
f"{self.db_name}.*",
29+
tools_pod=self.tools_pod,
2530
)
2631

2732
def create_search_index(self):

‎docker/mongodb-kubernetes-tests/tests/common/search/search_tester.py‎

Lines changed: 36 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
import os
2-
import tempfile
2+
from dataclasses import dataclass
33

44
import kubetester
5-
import requests
6-
from kubetester.helm import process_run_and_check
5+
from kubetester.kubetester import KubernetesTester
76
from kubetester.mongotester import MongoTester
87
from pymongo.operations import SearchIndexModel
98
from tests import test_logger
9+
from tests.common.mongodb_tools_pod import mongodb_tools_pod
1010

1111
logger = test_logger.get_test_logger(__name__)
1212

@@ -20,19 +20,39 @@ def __init__(
2020
):
2121
super().__init__(connection_string, use_ssl, ca_path)
2222

23-
def mongorestore_from_url(self, archive_url: str, ns_include: str, mongodb_tools_dir: str = ""):
24-
logger.debug(f"running mongorestore from {archive_url}")
25-
with tempfile.NamedTemporaryFile(delete=False) as sample_file:
26-
resp = requests.get(archive_url)
27-
size = sample_file.write(resp.content)
28-
logger.debug(f"Downloaded sample file from {archive_url} to {sample_file.name} (size: {size})")
29-
mongorestore_path = os.path.join(mongodb_tools_dir, "mongorestore")
30-
mongorestore_cmd = f"{mongorestore_path} --archive={sample_file.name} --verbose=1 --drop --nsInclude {ns_include} --uri={self.cnx_string}"
31-
if self.default_opts.get("tls", False):
32-
mongorestore_cmd += " --ssl"
33-
if ca_path := self.default_opts.get("tlsCAFile"):
34-
mongorestore_cmd += " --sslCAFile=" + ca_path
35-
process_run_and_check(mongorestore_cmd.split(), capture_output=True)
23+
def mongorestore_from_url(self, archive_url: str, ns_include: str, tools_pod: mongodb_tools_pod.ToolsPod):
24+
logger.debug(f"running mongorestore from {archive_url} in pod {tools_pod.pod_name}")
25+
26+
archive_file = f"/tmp/mongodb_archive_{os.urandom(4).hex()}.archive"
27+
28+
logger.debug(f"Downloading archive to {archive_file}")
29+
download_cmd = ["curl", "-L", "-o", archive_file, archive_url]
30+
tools_pod.run_command(cmd=download_cmd)
31+
32+
logger.debug("Running mongorestore")
33+
mongorestore_cmd = [
34+
"mongorestore",
35+
f"--archive={archive_file}",
36+
"--verbose=1",
37+
"--drop",
38+
f"--nsInclude={ns_include}",
39+
f"--uri={self.cnx_string}",
40+
]
41+
42+
if self.default_opts.get("tls", False):
43+
mongorestore_cmd.append("--ssl")
44+
if ca_path := self.default_opts.get("tlsCAFile"):
45+
tools_pod.copy_file_to_pod(ca_path, "/tmp/ca.crt")
46+
mongorestore_cmd.append(f"--sslCAFile=/tmp/ca.crt")
47+
48+
result = tools_pod.run_command(cmd=mongorestore_cmd)
49+
logger.debug(f"mongorestore completed: {result}")
50+
51+
logger.debug("Cleaning up archive file")
52+
cleanup_cmd = ["rm", "-f", archive_file]
53+
tools_pod.run_command(cmd=cleanup_cmd)
54+
55+
return result
3656

3757
def create_search_index(self, database_name: str, collection_name: str):
3858
database = self.client[database_name]

‎docker/mongodb-kubernetes-tests/tests/search/search_community_basic.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from kubetester.phase import Phase
66
from pytest import fixture, mark
77
from tests import test_logger
8+
from tests.common.mongodb_tools_pod import mongodb_tools_pod
89
from tests.common.search import movies_search_helper
910
from tests.common.search.movies_search_helper import SampleMoviesSearchHelper
1011
from tests.common.search.search_tester import SearchTester
@@ -86,9 +87,10 @@ def test_wait_for_community_resource_ready(mdbc: MongoDBCommunity):
8687

8788

8889
@fixture(scope="function")
89-
def sample_movies_helper(mdbc: MongoDBCommunity) -> SampleMoviesSearchHelper:
90+
def sample_movies_helper(mdbc: MongoDBCommunity, tools_pod: mongodb_tools_pod.ToolsPod) -> SampleMoviesSearchHelper:
9091
return movies_search_helper.SampleMoviesSearchHelper(
91-
SearchTester(get_connection_string(mdbc, USER_NAME, USER_PASSWORD))
92+
SearchTester(get_connection_string(mdbc, USER_NAME, USER_PASSWORD)),
93+
tools_pod=tools_pod,
9294
)
9395

9496

‎docker/mongodb-kubernetes-tests/tests/search/search_community_external_mongod_basic.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from kubetester.phase import Phase
66
from pytest import fixture, mark
77
from tests import test_logger
8+
from tests.common.mongodb_tools_pod import mongodb_tools_pod
89
from tests.common.search import movies_search_helper
910
from tests.common.search.movies_search_helper import SampleMoviesSearchHelper
1011
from tests.common.search.search_tester import SearchTester
@@ -118,9 +119,10 @@ def test_wait_for_community_resource_ready(mdbc: MongoDBCommunity):
118119

119120

120121
@fixture(scope="function")
121-
def sample_movies_helper(mdbc: MongoDBCommunity) -> SampleMoviesSearchHelper:
122+
def sample_movies_helper(mdbc: MongoDBCommunity, tools_pod: mongodb_tools_pod.ToolsPod) -> SampleMoviesSearchHelper:
122123
return movies_search_helper.SampleMoviesSearchHelper(
123-
SearchTester(get_connection_string(mdbc, USER_NAME, USER_PASSWORD))
124+
SearchTester(get_connection_string(mdbc, USER_NAME, USER_PASSWORD)),
125+
tools_pod=tools_pod,
124126
)
125127

126128

‎docker/mongodb-kubernetes-tests/tests/search/search_community_external_mongod_tls.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from kubetester.phase import Phase
77
from pytest import fixture, mark
88
from tests import test_logger
9+
from tests.common.mongodb_tools_pod import mongodb_tools_pod
910
from tests.common.search import movies_search_helper
1011
from tests.common.search.movies_search_helper import SampleMoviesSearchHelper
1112
from tests.common.search.search_tester import SearchTester
@@ -157,13 +158,16 @@ def test_wait_for_community_resource_ready(mdbc: MongoDBCommunity):
157158

158159

159160
@fixture(scope="function")
160-
def sample_movies_helper(mdbc: MongoDBCommunity, issuer_ca_filepath: str) -> SampleMoviesSearchHelper:
161+
def sample_movies_helper(
162+
mdbc: MongoDBCommunity, issuer_ca_filepath: str, tools_pod: mongodb_tools_pod.ToolsPod
163+
) -> SampleMoviesSearchHelper:
161164
return movies_search_helper.SampleMoviesSearchHelper(
162165
SearchTester(
163166
get_connection_string(mdbc, USER_NAME, USER_PASSWORD),
164167
use_ssl=True,
165168
ca_path=issuer_ca_filepath,
166-
)
169+
),
170+
tools_pod=tools_pod,
167171
)
168172

169173

‎docker/mongodb-kubernetes-tests/tests/search/search_community_tls.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from kubetester.phase import Phase
77
from pytest import fixture, mark
88
from tests import test_logger
9+
from tests.common.mongodb_tools_pod import mongodb_tools_pod
910
from tests.common.search import movies_search_helper
1011
from tests.common.search.movies_search_helper import SampleMoviesSearchHelper
1112
from tests.common.search.search_tester import SearchTester
@@ -124,9 +125,12 @@ def test_wait_for_community_resource_ready(mdbc: MongoDBCommunity):
124125

125126

126127
@fixture(scope="function")
127-
def sample_movies_helper(mdbc: MongoDBCommunity, issuer_ca_filepath: str) -> SampleMoviesSearchHelper:
128+
def sample_movies_helper(
129+
mdbc: MongoDBCommunity, issuer_ca_filepath: str, tools_pod: mongodb_tools_pod.ToolsPod
130+
) -> SampleMoviesSearchHelper:
128131
return movies_search_helper.SampleMoviesSearchHelper(
129132
SearchTester(get_connection_string(mdbc, USER_NAME, USER_PASSWORD), use_ssl=True, ca_path=issuer_ca_filepath),
133+
tools_pod=tools_pod,
130134
)
131135

132136

‎docker/mongodb-kubernetes-tests/tests/search/search_enterprise_basic.py‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
from kubetester.phase import Phase
99
from pytest import fixture, mark
1010
from tests import test_logger
11+
from tests.common.mongodb_tools_pod import mongodb_tools_pod
1112
from tests.common.search import movies_search_helper
1213
from tests.common.search.movies_search_helper import SampleMoviesSearchHelper
1314
from tests.common.search.search_tester import SearchTester
@@ -155,27 +156,28 @@ def test_wait_for_database_resource_ready(mdb: MongoDB):
155156
), "mongot parameters not found in mongod config"
156157

157158

158-
@mark.e2e_search_enterprise_basic
159-
def test_search_restore_sample_database(mdb: MongoDB):
160-
sample_movies_helper = movies_search_helper.SampleMoviesSearchHelper(
161-
SearchTester(get_connection_string(mdb, ADMIN_USER_NAME, ADMIN_USER_PASSWORD))
159+
@fixture(scope="function")
160+
def sample_movies_helper(
161+
mdb: MongoDB, issuer_ca_filepath: str, tools_pod: mongodb_tools_pod.ToolsPod
162+
) -> SampleMoviesSearchHelper:
163+
return movies_search_helper.SampleMoviesSearchHelper(
164+
SearchTester(get_connection_string(mdb, USER_NAME, USER_PASSWORD), use_ssl=True, ca_path=issuer_ca_filepath),
165+
tools_pod=tools_pod,
162166
)
167+
168+
169+
@mark.e2e_search_enterprise_basic
170+
def test_search_restore_sample_database(sample_movies_helper: SampleMoviesSearchHelper):
163171
sample_movies_helper.restore_sample_database()
164172

165173

166174
@mark.e2e_search_enterprise_basic
167-
def test_search_create_search_index(mdb: MongoDB):
168-
sample_movies_helper = movies_search_helper.SampleMoviesSearchHelper(
169-
SearchTester(get_connection_string(mdb, USER_NAME, USER_PASSWORD))
170-
)
175+
def test_search_create_search_index(sample_movies_helper: SampleMoviesSearchHelper):
171176
sample_movies_helper.create_search_index()
172177

173178

174179
@mark.e2e_search_enterprise_basic
175-
def test_search_assert_search_query(mdb: MongoDB):
176-
sample_movies_helper = movies_search_helper.SampleMoviesSearchHelper(
177-
SearchTester(get_connection_string(mdb, USER_NAME, USER_PASSWORD))
178-
)
180+
def test_search_assert_search_query(sample_movies_helper: SampleMoviesSearchHelper):
179181
sample_movies_helper.assert_search_query(retry_timeout=60)
180182

181183

0 commit comments

Comments
 (0)