From 935a29c60ea34500704689f6b283a541cdfd9a1e Mon Sep 17 00:00:00 2001 From: Rachel Chen Date: Wed, 11 Sep 2024 03:19:57 -0700 Subject: [PATCH] c --- snuba/cli/jobs.py | 65 ++++++++++++-- snuba/manual_jobs/__init__.py | 3 +- snuba/manual_jobs/job_loader.py | 36 ++++++-- snuba/manual_jobs/manifest_reader.py | 4 +- snuba/manual_jobs/toy_job.py | 8 +- tests/cli/test_jobs.py | 124 +++++++++++++++++++++++++-- 6 files changed, 216 insertions(+), 24 deletions(-) diff --git a/snuba/cli/jobs.py b/snuba/cli/jobs.py index 5ae0e824aa..dcb6ebbb04 100644 --- a/snuba/cli/jobs.py +++ b/snuba/cli/jobs.py @@ -1,8 +1,13 @@ -from typing import Tuple +from typing import Any, MutableMapping, Tuple import click from snuba.manual_jobs.job_loader import JobLoader +from snuba.manual_jobs.manifest_reader import ManifestReader + +JOB_SPECIFICATION_ERROR_MSG = ( + "Either a json manifest or job name & job id must be provided, but not both" +) @click.group() @@ -10,19 +15,65 @@ def jobs() -> None: pass +def _override_with_command_line_args( + kwargs: MutableMapping[Any, Any], pairs: Tuple[str, ...] +) -> None: + for pair in pairs: + k, v = pair.split("=") + kwargs[k] = v + + +def _insufficient_job_specification( + job_name: str, job_id: str, json_manifest: str +) -> bool: + return not json_manifest and not (job_name and job_id) + + +def _redundant_job_specification( + job_name: str, job_id: str, json_manifest: str +) -> bool: + return json_manifest and (job_name or job_id) + + @jobs.command() -@click.argument("job_name") +@click.option("--json_manifest", required=True) @click.option( "--dry_run", default=True, ) @click.argument("pairs", nargs=-1) -def run(*, job_name: str, dry_run: bool, pairs: Tuple[str, ...]) -> None: +def run_from_manifest(*, json_manifest: str, dry_run: bool, pairs: Tuple[str, ...]): + contents = ManifestReader.read(json_manifest)[0] + print(contents) + + +@jobs.command() +@click.option("--job_type") +@click.option("--job_id") +@click.option("--json_manifest") +@click.option( + "--dry_run", + default=True, +) +@click.argument("pairs", nargs=-1) +def run( + *, + job_type: str, + job_id: str, + json_manifest: str, + dry_run: bool, + pairs: Tuple[str, ...] +) -> None: + if _insufficient_job_specification( + job_type, job_id, json_manifest + ) or _redundant_job_specification(job_type, job_id, json_manifest): + raise click.ClickException(JOB_SPECIFICATION_ERROR_MSG) kwargs = {} - for pair in pairs: - k, v = pair.split("=") - kwargs[k] = v + if json_manifest: + job_type, job_id = JobLoader.parse_json_manifest(json_manifest, kwargs) + + _override_with_command_line_args(kwargs, pairs) - job_to_run = JobLoader.get_job_instance(job_name, dry_run, **kwargs) + job_to_run = JobLoader.get_job_instance(job_type, job_id, dry_run, **kwargs) job_to_run.execute() diff --git a/snuba/manual_jobs/__init__.py b/snuba/manual_jobs/__init__.py index d7b9e6e34e..2dd1dbb7fd 100644 --- a/snuba/manual_jobs/__init__.py +++ b/snuba/manual_jobs/__init__.py @@ -6,7 +6,8 @@ class Job(ABC, metaclass=RegisteredClass): - def __init__(self, dry_run: bool, **kwargs: Any) -> None: + def __init__(self, job_id: str, dry_run: bool, **kwargs: Any) -> None: + self.id = job_id self.dry_run = dry_run for k, v in kwargs.items(): setattr(self, k, v) diff --git a/snuba/manual_jobs/job_loader.py b/snuba/manual_jobs/job_loader.py index cce933f89f..c3584a6a87 100644 --- a/snuba/manual_jobs/job_loader.py +++ b/snuba/manual_jobs/job_loader.py @@ -1,15 +1,41 @@ -from typing import Any, cast +import json +from typing import Any, MutableMapping, Tuple, cast from snuba.manual_jobs import Job +from snuba.utils.serializable_exception import SerializableException + + +class JsonManifestParsingException(SerializableException): + pass class JobLoader: @staticmethod - def get_job_instance(class_name: str, dry_run: bool, **kwargs: Any) -> "Job": - job_type_class = Job.class_from_name(class_name) + def parse_json_manifest( + filepath: str, kwargs: MutableMapping[Any, Any] + ) -> Tuple[str, str]: + try: + with open(filepath, "r") as json_manifest_file: + json_manifest = json.load(json_manifest_file) + job_id = json_manifest["id"] + job_type = json_manifest["job_type"] + for k, v in json_manifest["params"].items(): + kwargs[k] = v + return job_type, job_id + + except KeyError: + raise JsonManifestParsingException( + "The JSON manifest file should contain the `id` field and the `job_type` field" + ) + + @staticmethod + def get_job_instance( + job_type: str, job_id: str, dry_run: bool, **kwargs: Any + ) -> "Job": + job_type_class = Job.class_from_name(job_type) if job_type_class is None: raise Exception( - f"Job does not exist. Did you make a file {class_name}.py yet?" + f"Job does not exist. Did you make a file {job_type}.py yet?" ) - return cast("Job", job_type_class(dry_run=dry_run, **kwargs)) + return cast("Job", job_type_class(job_id, dry_run=dry_run, **kwargs)) diff --git a/snuba/manual_jobs/manifest_reader.py b/snuba/manual_jobs/manifest_reader.py index 9756f0c5d2..530de4c5f4 100644 --- a/snuba/manual_jobs/manifest_reader.py +++ b/snuba/manual_jobs/manifest_reader.py @@ -4,7 +4,7 @@ import simplejson -class _ManifestReader: +class ManifestReader: @staticmethod def read(filename: str) -> Sequence[Any]: local_root = os.path.dirname(__file__) @@ -15,4 +15,4 @@ def read(filename: str) -> Sequence[Any]: def read_jobs_manifest() -> Sequence[Any]: - return _ManifestReader.read("run_manifest.json") + return ManifestReader.read("run_manifest.json") diff --git a/snuba/manual_jobs/toy_job.py b/snuba/manual_jobs/toy_job.py index e48733c678..c85d0b89f6 100644 --- a/snuba/manual_jobs/toy_job.py +++ b/snuba/manual_jobs/toy_job.py @@ -6,8 +6,8 @@ class ToyJob(Job): - def __init__(self, dry_run: bool, **kwargs: Any): - super().__init__(dry_run, **kwargs) + def __init__(self, job_id: str, dry_run: bool, **kwargs: Any): + super().__init__(job_id, dry_run, **kwargs) def _build_query(self) -> str: if self.dry_run: @@ -16,4 +16,6 @@ def _build_query(self) -> str: return "not dry run query" def execute(self) -> None: - click.echo("executing query `" + self._build_query() + "`") + click.echo( + "executing job " + self.id + " with query `" + self._build_query() + "`" + ) diff --git a/tests/cli/test_jobs.py b/tests/cli/test_jobs.py index dc7d9f0ffa..616f74d6f9 100644 --- a/tests/cli/test_jobs.py +++ b/tests/cli/test_jobs.py @@ -1,17 +1,129 @@ +import json +from typing import Any + from click.testing import CliRunner -from snuba.cli.jobs import run +from snuba.cli.jobs import JOB_SPECIFICATION_ERROR_MSG, run, run_from_manifest -def test_valid_job() -> None: +def test_invalid_job_errors() -> None: runner = CliRunner() - result = runner.invoke(run, ["ToyJob", "--dry_run", "True", "k1=v1", "k2=v2"]) - assert result.exit_code == 0 + result = runner.invoke( + run, + [ + "--job_type", + "NonexistentJob", + "--job_id", + "0001", + "--dry_run", + "True", + "k1=v1", + "k2=v2", + ], + ) + + assert result.exit_code == 1 + + +def test_provides_json_manifest_and_job_name_errors() -> None: + runner = CliRunner() + result = runner.invoke( + run, + [ + "--job_type", + "ToyJob", + "--json_manifest", + "doesntmatter.json", + "--dry_run", + "True", + "k1=v1", + "k2=v2", + ], + ) + assert result.exit_code == 1 + assert result.output == "Error: " + JOB_SPECIFICATION_ERROR_MSG + "\n" + + +def test_provides_json_manifest_and_job_id_errors() -> None: + runner = CliRunner() + result = runner.invoke( + run, + [ + "--job_id", + "0001", + "--json_manifest", + "doesntmatter.json", + "--dry_run", + "True", + "k1=v1", + "k2=v2", + ], + ) + assert result.exit_code == 1 + assert result.output == "Error: " + JOB_SPECIFICATION_ERROR_MSG + "\n" + + +def test_provides_no_job_specification_errors() -> None: + runner = CliRunner() + result = runner.invoke(run, ["--dry_run", "True", "k1=v1", "k2=v2"]) + assert result.exit_code == 1 + assert result.output == "Error: " + JOB_SPECIFICATION_ERROR_MSG + "\n" -def test_invalid_job() -> None: +def test_provides_only_job_name_errors() -> None: runner = CliRunner() result = runner.invoke( - run, ["SomeJobThatDoesntExist", "--dry_run", "True", "k1=v1", "k2=v2"] + run, ["--job_type", "ToyJob", "--dry_run", "True", "k1=v1", "k2=v2"] ) assert result.exit_code == 1 + assert result.output == "Error: " + JOB_SPECIFICATION_ERROR_MSG + "\n" + + +def test_provides_only_job_id_errors() -> None: + runner = CliRunner() + result = runner.invoke( + run, ["--job_id", "0001", "--dry_run", "True", "k1=v1", "k2=v2"] + ) + assert result.exit_code == 1 + assert result.output == "Error: " + JOB_SPECIFICATION_ERROR_MSG + "\n" + + +def test_provides_json_manifest_is_valid(tmp_path: Any) -> None: + runner = CliRunner() + json_manifest_file = tmp_path / "manifest_file.json" + data = {"id": "abc1234", "job_type": "ToyJob", "params": {"p1": "value1"}} + with json_manifest_file.open("w") as f: + json.dump(data, f) + result = runner.invoke(run, ["--json_manifest", str(json_manifest_file)]) + assert result.exit_code == 0 + + +def test_provides_job_name_and_job_id_is_valid() -> None: + runner = CliRunner() + result = runner.invoke( + run, + [ + "--job_type", + "ToyJob", + "--job_id", + "0001", + "--dry_run", + "True", + "k1=v1", + "k2=v2", + ], + ) + assert result.exit_code == 0 + + +def test_toy(): + runner = CliRunner() + result = runner.invoke( + run_from_manifest, + [ + "--json_manifest", + "run_manifest.json", + ], + ) + assert result.exit_code == 0 + assert False