diff --git a/snuba/cli/jobs.py b/snuba/cli/jobs.py index dcb6ebbb04..71d84ccf8b 100644 --- a/snuba/cli/jobs.py +++ b/snuba/cli/jobs.py @@ -2,12 +2,11 @@ import click +from snuba.manual_jobs import JobSpec 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" -) +JOB_SPECIFICATION_ERROR_MSG = "Missing job type and/or job id" @click.group() @@ -15,65 +14,41 @@ def jobs() -> None: pass -def _override_with_command_line_args( - kwargs: MutableMapping[Any, Any], pairs: Tuple[str, ...] -) -> None: +def _parse_params(pairs: Tuple[str, ...]) -> MutableMapping[Any, Any]: + params = {} 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) + params[k] = v + return params @jobs.command() @click.option("--json_manifest", required=True) +@click.option("--job_id") @click.option( "--dry_run", default=True, ) -@click.argument("pairs", nargs=-1) -def run_from_manifest(*, json_manifest: str, dry_run: bool, pairs: Tuple[str, ...]): - contents = ManifestReader.read(json_manifest)[0] - print(contents) +def run_from_manifest(*, json_manifest: str, job_id: str, dry_run: bool) -> None: + job_specs = ManifestReader.read(json_manifest) + for job_spec in job_specs: + if job_spec.job_id == job_id: + job_to_run = JobLoader.get_job_instance(job_spec, dry_run) + job_to_run.execute() @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): +def run(*, job_type: str, job_id: str, dry_run: bool, pairs: Tuple[str, ...]) -> None: + if not job_type or not job_id: raise click.ClickException(JOB_SPECIFICATION_ERROR_MSG) + job_spec = JobSpec(job_id=job_id, job_type=job_type, params=_parse_params(pairs)) - kwargs = {} - 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_type, job_id, dry_run, **kwargs) + job_to_run = JobLoader.get_job_instance(job_spec, dry_run) job_to_run.execute() diff --git a/snuba/manual_jobs/__init__.py b/snuba/manual_jobs/__init__.py index 2dd1dbb7fd..fe82969d7b 100644 --- a/snuba/manual_jobs/__init__.py +++ b/snuba/manual_jobs/__init__.py @@ -1,16 +1,25 @@ import os from abc import ABC, abstractmethod -from typing import Any, cast +from dataclasses import dataclass +from typing import Any, MutableMapping, Optional, cast from snuba.utils.registered_class import RegisteredClass, import_submodules_in_directory +@dataclass +class JobSpec: + job_id: str + job_type: str + params: Optional[MutableMapping[Any, Any]] + + class Job(ABC, metaclass=RegisteredClass): - def __init__(self, job_id: str, dry_run: bool, **kwargs: Any) -> None: - self.id = job_id + def __init__(self, job_spec: JobSpec, dry_run: bool) -> None: + self.job_spec = job_spec self.dry_run = dry_run - for k, v in kwargs.items(): - setattr(self, k, v) + if job_spec.params: + for k, v in job_spec.params.items(): + setattr(self, k, v) @abstractmethod def execute(self) -> None: diff --git a/snuba/manual_jobs/job_loader.py b/snuba/manual_jobs/job_loader.py index c3584a6a87..1e787a4647 100644 --- a/snuba/manual_jobs/job_loader.py +++ b/snuba/manual_jobs/job_loader.py @@ -1,7 +1,6 @@ -import json -from typing import Any, MutableMapping, Tuple, cast +from typing import cast -from snuba.manual_jobs import Job +from snuba.manual_jobs import Job, JobSpec from snuba.utils.serializable_exception import SerializableException @@ -11,31 +10,11 @@ class JsonManifestParsingException(SerializableException): class JobLoader: @staticmethod - 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) + def get_job_instance(job_spec: JobSpec, dry_run: bool) -> "Job": + job_type_class = Job.class_from_name(job_spec.job_type) if job_type_class is None: raise Exception( - f"Job does not exist. Did you make a file {job_type}.py yet?" + f"Job does not exist. Did you make a file {job_spec.job_type}.py yet?" ) - return cast("Job", job_type_class(job_id, dry_run=dry_run, **kwargs)) + return cast("Job", job_type_class(job_spec, dry_run=dry_run)) diff --git a/snuba/manual_jobs/manifest_reader.py b/snuba/manual_jobs/manifest_reader.py index 530de4c5f4..a31091c842 100644 --- a/snuba/manual_jobs/manifest_reader.py +++ b/snuba/manual_jobs/manifest_reader.py @@ -1,18 +1,30 @@ import os -from typing import Any, Sequence +import typing +from typing import Sequence import simplejson +from snuba.manual_jobs import JobSpec + class ManifestReader: @staticmethod - def read(filename: str) -> Sequence[Any]: + def read(filename: str) -> Sequence[JobSpec]: local_root = os.path.dirname(__file__) with open(os.path.join(local_root, filename)) as stream: contents = simplejson.loads(stream.read()) - assert isinstance(contents, Sequence) - return contents + contents_dict = dict(zip(contents[::2], contents[1::2])) + job_specs = [] + for content in contents: + assert isinstance(content, dict) + job_spec = JobSpec( + job_id=typing.cast(str, contents_dict.get("id")), + job_type=typing.cast(str, contents_dict.get("job_type")), + params=contents_dict.get("params"), + ) + job_specs.append(job_spec) + return job_specs -def read_jobs_manifest() -> Sequence[Any]: +def read_jobs_manifest() -> Sequence[JobSpec]: return ManifestReader.read("run_manifest.json") diff --git a/snuba/manual_jobs/toy_job.py b/snuba/manual_jobs/toy_job.py index c85d0b89f6..c9aec34595 100644 --- a/snuba/manual_jobs/toy_job.py +++ b/snuba/manual_jobs/toy_job.py @@ -1,13 +1,15 @@ -from typing import Any - import click -from snuba.manual_jobs import Job +from snuba.manual_jobs import Job, JobSpec class ToyJob(Job): - def __init__(self, job_id: str, dry_run: bool, **kwargs: Any): - super().__init__(job_id, dry_run, **kwargs) + def __init__( + self, + job_spec: JobSpec, + dry_run: bool, + ): + super().__init__(job_spec, dry_run) def _build_query(self) -> str: if self.dry_run: @@ -17,5 +19,9 @@ def _build_query(self) -> str: def execute(self) -> None: click.echo( - "executing job " + self.id + " with query `" + self._build_query() + "`" + "executing job " + + self.job_spec.job_id + + " with query `" + + self._build_query() + + "`" ) diff --git a/tests/cli/test_jobs.py b/tests/cli/test_jobs.py index 616f74d6f9..25f48e479e 100644 --- a/tests/cli/test_jobs.py +++ b/tests/cli/test_jobs.py @@ -1,76 +1,45 @@ -import json -from typing import Any - from click.testing import CliRunner from snuba.cli.jobs import JOB_SPECIFICATION_ERROR_MSG, run, run_from_manifest -def test_invalid_job_errors() -> None: +def test_cmd_line_valid() -> None: runner = CliRunner() result = runner.invoke( run, - [ - "--job_type", - "NonexistentJob", - "--job_id", - "0001", - "--dry_run", - "True", - "k1=v1", - "k2=v2", - ], + ["--job_type", "ToyJob", "--job_id", "0001"], ) - assert result.exit_code == 1 + assert result.exit_code == 0 -def test_provides_json_manifest_and_job_name_errors() -> None: +def test_invalid_job_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, - [ + "NonexistentJob", "--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: +def test_cmd_line_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: +def test_cmd_line_no_job_id_errors() -> None: runner = CliRunner() result = runner.invoke( run, ["--job_type", "ToyJob", "--dry_run", "True", "k1=v1", "k2=v2"] @@ -79,7 +48,7 @@ def test_provides_only_job_name_errors() -> None: assert result.output == "Error: " + JOB_SPECIFICATION_ERROR_MSG + "\n" -def test_provides_only_job_id_errors() -> None: +def test_cmd_line_no_job_type_errors() -> None: runner = CliRunner() result = runner.invoke( run, ["--job_id", "0001", "--dry_run", "True", "k1=v1", "k2=v2"] @@ -88,42 +57,15 @@ def test_provides_only_job_id_errors() -> None: 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(): +def test_json_valid() -> None: runner = CliRunner() result = runner.invoke( run_from_manifest, [ "--json_manifest", "run_manifest.json", + "--job_id", + "abc1234", ], ) assert result.exit_code == 0 - assert False