rewrite

snuba/cli/jobs.py

@@ -2,78 +2,53 @@
 
 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()
 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()

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:

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))

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")

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()
+            + "`"
         )

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