Skip to content

Commit

Permalink
rewrite
Browse files Browse the repository at this point in the history
  • Loading branch information
Rachel Chen authored and Rachel Chen committed Sep 16, 2024
1 parent 935a29c commit ef9c4fa
Show file tree
Hide file tree
Showing 5 changed files with 73 additions and 117 deletions.
47 changes: 17 additions & 30 deletions snuba/cli/jobs.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,25 +2,24 @@

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
params[k] = v
return params


def _insufficient_job_specification(
Expand All @@ -37,43 +36,31 @@ def _redundant_job_specification(

@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):
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()
19 changes: 14 additions & 5 deletions snuba/manual_jobs/__init__.py
Original file line number Diff line number Diff line change
@@ -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:
Expand Down
33 changes: 6 additions & 27 deletions snuba/manual_jobs/job_loader.py
Original file line number Diff line number Diff line change
@@ -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


Expand All @@ -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))
29 changes: 25 additions & 4 deletions snuba/manual_jobs/manifest_reader.py
Original file line number Diff line number Diff line change
@@ -1,18 +1,39 @@
import os
import typing
from typing import Any, Sequence

import simplejson

from snuba.manual_jobs import JobSpec


def convert_to_dict(contents: list) -> dict[Any, Any]:
contents_dict = {}
for i in range(0, len(contents), 2):
k = contents[i]
v = contents[i + 1]
contents_dict[k] = v
return contents_dict


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")
62 changes: 11 additions & 51 deletions tests/cli/test_jobs.py
Original file line number Diff line number Diff line change
@@ -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"]
Expand All @@ -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"]
Expand All @@ -88,16 +57,6 @@ 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(
Expand All @@ -123,7 +82,8 @@ def test_toy():
[
"--json_manifest",
"run_manifest.json",
"--job_id",
"abc1234",
],
)
assert result.exit_code == 0
assert False

0 comments on commit ef9c4fa

Please sign in to comment.