rewrite

snuba/cli/jobs.py

@@ -2,78 +2,54 @@
 
 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)
+ if job_id not in job_specs.keys():
+ raise click.ClickException("Provide a valid job id")
+
+ job_to_run = JobLoader.get_job_instance(job_specs[job_id], 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,41 +1,15 @@
-import json
-from typing import Any, MutableMapping, Tuple, cast
+from typing import cast
 
-from snuba.manual_jobs import Job
-from snuba.utils.serializable_exception import SerializableException
-
-
-class JsonManifestParsingException(SerializableException):
- pass
+from snuba.manual_jobs import Job, JobSpec
 
 
 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,35 @@
 import os
-from typing import Any, Sequence
+from typing import Mapping
 
 import simplejson
 
+from snuba.manual_jobs import JobSpec
+
 
 class ManifestReader:
  @staticmethod
- def read(filename: str) -> Sequence[Any]:
+ def read(filename: str) -> Mapping[str, 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
+
+ job_specs = {}
+ for content in contents:
+ content_dict = dict(zip(content[::2], content[1::2]))
+
+ job_id = content_dict["id"]
+ assert isinstance(job_id, str)
+ job_type = content_dict["job_type"]
+ assert isinstance(job_type, str)
+
+ job_spec = JobSpec(
+ job_id=job_id,
+ job_type=job_type,
+ params=content_dict.get("params"),
+ )
+ job_specs[job_id] = job_spec
+ return job_specs
 
 
-def read_jobs_manifest() -> Sequence[Any]:
+def read_jobs_manifest() -> Mapping[str, 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