c

snuba/cli/jobs.py

@@ -1,28 +1,79 @@
-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()
 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()

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)

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

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

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

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