diff --git a/tools-tests/upload_test.py b/tools-tests/upload_test.py index ce73499296..96c6c39f9d 100644 --- a/tools-tests/upload_test.py +++ b/tools-tests/upload_test.py @@ -40,6 +40,17 @@ def main_norun_redacted_zip(main_norun): yield +@pytest.fixture +def taskrunner_workdir(tmp_path: pathlib.Path): + """ + Creates a temporary workdir for the TaskRunner to use. The directory is required to exist before TaskRunner is + instantiated and closed when TaskRunner goes out of scope. + """ + workdir = tmp_path / "workdir" + workdir.mkdir() + yield workdir + + class FakeResponse: def __init__(self, status_code): self.status_code = status_code @@ -47,6 +58,9 @@ def __init__(self, status_code): def getcode(self): return self.status_code + def read(self) -> bytes: + return b"{}" + class FakeSuccessUrlOpener: def __init__(self, *args, **kwargs): @@ -66,28 +80,40 @@ def open(self, *args, **kwargs): @pytest.mark.usefixtures("main_norun") @pytest.mark.parametrize("args", [[], ["--log-redaction-level", "none"]]) -def test_main_output_exists(args): +def test_main_output_exists(args, taskrunner_workdir): with pytest.raises(SystemExit, check=lambda e: e.code == 0): - with unittest.mock.patch("sys.argv", ["sg_collect", *args, ZIP_NAME]): + with unittest.mock.patch( + "sys.argv", ["sg_collect", *args, "--tmp-dir", taskrunner_workdir, ZIP_NAME] + ): sgcollect.main() assert pathlib.Path(ZIP_NAME).exists() assert not pathlib.Path(REDACTED_ZIP_NAME).exists() + assert not [x for x in taskrunner_workdir.iterdir()] @pytest.mark.usefixtures("main_norun_redacted_zip") -def test_main_output_exists_with_redacted(): +def test_main_output_exists_with_redacted(taskrunner_workdir): with pytest.raises(SystemExit, check=lambda e: e.code == 0): with unittest.mock.patch( - "sys.argv", ["sg_collect", "--log-redaction-level", "partial", ZIP_NAME] + "sys.argv", + [ + "sg_collect", + "--log-redaction-level", + "partial", + "--tmp-dir", + taskrunner_workdir, + ZIP_NAME, + ], ): sgcollect.main() assert pathlib.Path(ZIP_NAME).exists() assert pathlib.Path(REDACTED_ZIP_NAME).exists() + assert not [x for x in taskrunner_workdir.iterdir()] @pytest.mark.usefixtures("main_norun") @pytest.mark.parametrize("args", [[], ["--log-redaction-level", "none"]]) -def test_main_zip_deleted_on_upload_success(args): +def test_main_zip_deleted_on_upload_success(args, taskrunner_workdir): with unittest.mock.patch("tasks.urllib.request.build_opener", FakeSuccessUrlOpener): with unittest.mock.patch( "sys.argv", @@ -98,6 +124,8 @@ def test_main_zip_deleted_on_upload_success(args): "https://example.com", "--customer", "fakeCustomer", + "--tmp-dir", + taskrunner_workdir, ZIP_NAME, ], ): @@ -106,11 +134,12 @@ def test_main_zip_deleted_on_upload_success(args): assert exc.value.code == 0 assert not pathlib.Path(ZIP_NAME).exists() assert not pathlib.Path(REDACTED_ZIP_NAME).exists() + assert not [x for x in taskrunner_workdir.iterdir()] @pytest.mark.usefixtures("main_norun") @pytest.mark.parametrize("args", [[], ["--log-redaction-level", "none"]]) -def test_main_zip_deleted_on_upload_failure(args): +def test_main_zip_deleted_on_upload_failure(args, taskrunner_workdir): with unittest.mock.patch("tasks.urllib.request.build_opener", FakeFailureUrlOpener): with unittest.mock.patch( "sys.argv", @@ -129,10 +158,11 @@ def test_main_zip_deleted_on_upload_failure(args): assert exc.value.code == 1 assert not pathlib.Path(ZIP_NAME).exists() assert not pathlib.Path(REDACTED_ZIP_NAME).exists() + assert not [x for x in taskrunner_workdir.iterdir()] @pytest.mark.usefixtures("main_norun_redacted_zip") -def test_main_redacted_zip_deleted_on_upload_success(): +def test_main_redacted_zip_deleted_on_upload_success(taskrunner_workdir): with unittest.mock.patch("tasks.urllib.request.build_opener", FakeSuccessUrlOpener): with unittest.mock.patch( "sys.argv", @@ -144,6 +174,8 @@ def test_main_redacted_zip_deleted_on_upload_success(): "https://example.com", "--customer", "fakeCustomer", + "--tmp-dir", + taskrunner_workdir, ZIP_NAME, ], ): @@ -152,10 +184,11 @@ def test_main_redacted_zip_deleted_on_upload_success(): assert exc.value.code == 0 assert not pathlib.Path(ZIP_NAME).exists() assert not pathlib.Path(REDACTED_ZIP_NAME).exists() + assert not [x for x in taskrunner_workdir.iterdir()] @pytest.mark.usefixtures("main_norun_redacted_zip") -def test_main_redacted_zip_deleted_on_upload_failure(): +def test_main_redacted_zip_deleted_on_upload_failure(taskrunner_workdir): with unittest.mock.patch("tasks.urllib.request.build_opener", FakeFailureUrlOpener): with unittest.mock.patch( "sys.argv", @@ -167,6 +200,8 @@ def test_main_redacted_zip_deleted_on_upload_failure(): "https://example.com", "--customer", "fakeCustomer", + "--tmp-dir", + taskrunner_workdir, ZIP_NAME, ], ): @@ -175,11 +210,12 @@ def test_main_redacted_zip_deleted_on_upload_failure(): assert exc.value.code == 1 assert not pathlib.Path(ZIP_NAME).exists() assert not pathlib.Path(REDACTED_ZIP_NAME).exists() + assert not [x for x in taskrunner_workdir.iterdir()] @pytest.mark.usefixtures("main_norun") @pytest.mark.parametrize("args", [[], ["--log-redaction-level", "none"]]) -def test_main_keep_zip_on_upload_success(args): +def test_main_keep_zip_on_upload_success(args, taskrunner_workdir): with unittest.mock.patch("tasks.urllib.request.build_opener", FakeSuccessUrlOpener): with unittest.mock.patch( "sys.argv", @@ -191,6 +227,8 @@ def test_main_keep_zip_on_upload_success(args): "--customer", "fakeCustomer", "--keep-zip", + "--tmp-dir", + taskrunner_workdir, ZIP_NAME, ], ): @@ -199,11 +237,12 @@ def test_main_keep_zip_on_upload_success(args): assert exc.value.code == 0 assert pathlib.Path(ZIP_NAME).exists() assert not pathlib.Path(REDACTED_ZIP_NAME).exists() + assert not [x for x in taskrunner_workdir.iterdir()] @pytest.mark.usefixtures("main_norun") @pytest.mark.parametrize("args", [[], ["--log-redaction-level", "none"]]) -def test_main_keep_zip_on_upload_failure(args): +def test_main_keep_zip_on_upload_failure(args, taskrunner_workdir): with unittest.mock.patch("tasks.urllib.request.build_opener", FakeFailureUrlOpener): with unittest.mock.patch( "sys.argv", @@ -215,6 +254,8 @@ def test_main_keep_zip_on_upload_failure(args): "--customer", "fakeCustomer", "--keep-zip", + "--tmp-dir", + taskrunner_workdir, ZIP_NAME, ], ): @@ -223,10 +264,11 @@ def test_main_keep_zip_on_upload_failure(args): assert exc.value.code == 1 assert pathlib.Path(ZIP_NAME).exists() assert not pathlib.Path(REDACTED_ZIP_NAME).exists() + assert not [x for x in taskrunner_workdir.iterdir()] @pytest.mark.usefixtures("main_norun_redacted_zip") -def test_main_keep_zip_deleted_on_upload_success(): +def test_main_keep_zip_deleted_on_upload_success(taskrunner_workdir): with unittest.mock.patch("tasks.urllib.request.build_opener", FakeSuccessUrlOpener): with unittest.mock.patch( "sys.argv", @@ -239,6 +281,8 @@ def test_main_keep_zip_deleted_on_upload_success(): "--customer", "fakeCustomer", "--keep-zip", + "--tmp-dir", + taskrunner_workdir, ZIP_NAME, ], ): @@ -247,10 +291,11 @@ def test_main_keep_zip_deleted_on_upload_success(): assert exc.value.code == 0 assert pathlib.Path(ZIP_NAME).exists() assert pathlib.Path(REDACTED_ZIP_NAME).exists() + assert not [x for x in taskrunner_workdir.iterdir()] @pytest.mark.usefixtures("main_norun_redacted_zip") -def test_main_keep_zip_deleted_on_upload_failure(): +def test_main_keep_zip_deleted_on_upload_failure(taskrunner_workdir): with unittest.mock.patch("tasks.urllib.request.build_opener", FakeFailureUrlOpener): with unittest.mock.patch( "sys.argv", @@ -263,6 +308,8 @@ def test_main_keep_zip_deleted_on_upload_failure(): "--customer", "fakeCustomer", "--keep-zip", + "--tmp-dir", + taskrunner_workdir, ZIP_NAME, ], ): @@ -271,6 +318,7 @@ def test_main_keep_zip_deleted_on_upload_failure(): assert exc.value.code == 1 assert pathlib.Path(ZIP_NAME).exists() assert pathlib.Path(REDACTED_ZIP_NAME).exists() + assert not [x for x in taskrunner_workdir.iterdir()] @pytest.fixture(scope="session") diff --git a/tools/sgcollect.py b/tools/sgcollect.py index d079530917..90595b8323 100755 --- a/tools/sgcollect.py +++ b/tools/sgcollect.py @@ -927,77 +927,75 @@ def main() -> NoReturn: do_upload(args[0], options.just_upload_into, options.upload_proxy) # Create a TaskRunner and run all of the OS tasks (collect top, netstat, etc) - # The output of the tasks will go directly into couchbase.log - runner = TaskRunner( + with TaskRunner( verbosity=options.verbosity, default_name="sync_gateway.log", tmp_dir=options.tmp_dir, - ) - - if not options.product_only: - for task in make_os_tasks(["sync_gateway"]): - runner.run(task) - - # Output the Python version if verbosity was enabled - if options.verbosity: - log("Python version: %s" % sys.version) - - # Find path to sg binary - sg_binary_path = discover_sg_binary_path(options, sg_url, auth_headers) - - # Run SG specific tasks - for task in make_sg_tasks( - sg_url=sg_url, - auth_headers=auth_headers, - sync_gateway_config_path_option=options.sync_gateway_config, - sync_gateway_executable_path=options.sync_gateway_executable, - should_redact=should_redact, - ): - runner.run(task) - - if ( - sg_binary_path is not None - and sg_binary_path != "" - and os.path.exists(sg_binary_path) - ): - runner.collect_file(sg_binary_path) - else: - print( - "WARNING: unable to find Sync Gateway executable, omitting from result. Go pprofs will not be accurate." - ) + ) as runner: + if not options.product_only: + for task in make_os_tasks(["sync_gateway"]): + runner.run(task) - runner.run(get_sgcollect_info_options_task(options, args)) + # Output the Python version if verbosity was enabled + if options.verbosity: + log("Python version: %s" % sys.version) - runner.close_all_files() + # Find path to sg binary + sg_binary_path = discover_sg_binary_path(options, sg_url, auth_headers) - # Build redacted zip file - if should_redact: - log("Redacting log files to level: %s" % options.redact_level) - runner.redact_and_zip( - redact_zip_file, "sgcollect_info", options.salt_value, platform.node() - ) + # Run SG specific tasks + for task in make_sg_tasks( + sg_url=sg_url, + auth_headers=auth_headers, + sync_gateway_config_path_option=options.sync_gateway_config, + sync_gateway_executable_path=options.sync_gateway_executable, + should_redact=should_redact, + ): + runner.run(task) - # Build the actual zip file - runner.zip(zip_filename, "sgcollect_info", platform.node()) + if ( + sg_binary_path is not None + and sg_binary_path != "" + and os.path.exists(sg_binary_path) + ): + runner.collect_file(sg_binary_path) + else: + print( + "WARNING: unable to find Sync Gateway executable, omitting from result. Go pprofs will not be accurate." + ) - if options.redact_level != "none": - print("Zipfile built: {0}".format(redact_zip_file)) + runner.run(get_sgcollect_info_options_task(options, args)) - print("Zipfile built: {0}".format(zip_filename)) + runner.close_all_files() - if not upload_url: - sys.exit(0) - # Upload the zip to the URL to S3 if required - try: + # Build redacted zip file if should_redact: - exit_code = do_upload(redact_zip_file, upload_url, options.upload_proxy) - else: - exit_code = do_upload(zip_filename, upload_url, options.upload_proxy) - finally: - if not options.keep_zip: - delete_zip(zip_filename) - delete_zip(redact_zip_file) - sys.exit(exit_code) + log("Redacting log files to level: %s" % options.redact_level) + runner.redact_and_zip( + redact_zip_file, "sgcollect_info", options.salt_value, platform.node() + ) + + try: + # Build the actual zip file + runner.zip(zip_filename, "sgcollect_info", platform.node()) + + if options.redact_level != "none": + print("Zipfile built: {0}".format(redact_zip_file)) + + print("Zipfile built: {0}".format(zip_filename)) + + if not upload_url: + sys.exit(0) + # Upload the zip to the URL to S3 if required + if should_redact: + exit_code = do_upload(redact_zip_file, upload_url, options.upload_proxy) + else: + exit_code = do_upload(zip_filename, upload_url, options.upload_proxy) + sys.exit(exit_code) + finally: + if not options.keep_zip and upload_url: + delete_zip(zip_filename) + delete_zip(redact_zip_file) def ud(value, should_redact=True): diff --git a/tools/tasks.py b/tools/tasks.py index e36d10939c..aeb2d7ca8a 100644 --- a/tools/tasks.py +++ b/tools/tasks.py @@ -18,11 +18,13 @@ import os import pathlib import re +import shutil import sys import tempfile import threading import time import traceback +import types import urllib.error import urllib.parse import urllib.request @@ -243,6 +245,11 @@ def will_run(self): class TaskRunner(object): + """ + TaskRunner manages running tasks and collecting their output into files. Use as a context manager to cleanup the + temporary files when finished. + """ + def __init__(self, verbosity=0, default_name="couchbase.log", tmp_dir=None): self.files = {} self.tasks = {} @@ -256,7 +263,7 @@ def __init__(self, verbosity=0, default_name="couchbase.log", tmp_dir=None): tmp_dir = os.path.abspath(os.path.expanduser(tmp_dir)) try: - self.tmpdir = tempfile.mkdtemp(dir=tmp_dir) + self.tmpdir = tempfile.mkdtemp(dir=tmp_dir, prefix="sgcollect") except OSError as e: print("Could not use temporary dir {0}: {1}".format(tmp_dir, e)) sys.exit(1) @@ -270,6 +277,18 @@ def __init__(self, verbosity=0, default_name="couchbase.log", tmp_dir=None): log("Could not use TMPDIR {0}".format(os.getenv("TMPDIR"))) log("Using temporary dir {0}".format(os.path.split(self.tmpdir)[0])) + def __enter__(self): + return self + + def __exit__( + self, + exc_type: Optional[type[BaseException]], + exc_value: Optional[type[BaseException]], + traceback: Optional[types.TracebackType], + ): + self.close_all_files() + shutil.rmtree(self.tmpdir, ignore_errors=True) + def collect_file(self, filename): """Add a file to the list of files collected. Used to capture the exact file (including timestamps) from the Couchbase instance.