diff --git a/snuba/cli/jobs.py b/snuba/cli/jobs.py index 5ae0e824aad..b2ff51388ba 100644 --- a/snuba/cli/jobs.py +++ b/snuba/cli/jobs.py @@ -1,28 +1,66 @@ -from typing import Tuple +from typing import Any, MutableMapping, Tuple import click from snuba.manual_jobs.job_loader import JobLoader +JOB_SPECIFICATION_ERROR_MSG = ( + "Either a json manifest or job name & job id must be provided, but not both" +) + @click.group() 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("--job_type") +@click.option("--job_id") +@click.option("--json_manifest") @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( + *, + 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 d7b9e6e34e1..2dd1dbb7fdc 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 cce933f89fe..c3584a6a87f 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/toy_job.py b/snuba/manual_jobs/toy_job.py index e48733c678e..c85d0b89f67 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 dc7d9f0ffa6..eba0a933965 100644 --- a/tests/cli/test_jobs.py +++ b/tests/cli/test_jobs.py @@ -1,17 +1,116 @@ +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 -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_provides_only_job_name_errors() -> None: + runner = CliRunner() + result = runner.invoke( + 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_invalid_job() -> None: + +def test_provides_only_job_id_errors() -> None: runner = CliRunner() result = runner.invoke( - run, ["SomeJobThatDoesntExist", "--dry_run", "True", "k1=v1", "k2=v2"] + 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