Skip to content

Commit

Permalink
c
Browse files Browse the repository at this point in the history
  • Loading branch information
Rachel Chen authored and Rachel Chen committed Sep 11, 2024
1 parent d02cff4 commit 1160ad8
Show file tree
Hide file tree
Showing 5 changed files with 190 additions and 22 deletions.
54 changes: 47 additions & 7 deletions snuba/cli/jobs.py
Original file line number Diff line number Diff line change
@@ -1,28 +1,68 @@
from typing import Tuple
from collections.abc import Mapping
from typing import Any, MutableMapping, Tuple

import click

from snuba.manual_jobs.job_loader import JobLoader

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, job_id, json_manifest) -> bool:
return not json_manifest and not (job_name and job_id)


def _redundant_job_specification(job_name, job_id, json_manifest) -> bool:
return json_manifest and (job_name or job_id)


@jobs.command()
@click.argument("job_name")
@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_name: str, dry_run: bool, pairs: Tuple[str, ...]) -> None:
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:
JobLoader.parse_json_manifest(json_manifest, kwargs)
else:
kwargs["job_type"] = job_type
kwargs["job_id"] = job_id
print(pairs)
_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(
kwargs["job_type"], kwargs["job_id"], dry_run, **kwargs
)
job_to_run.execute()
5 changes: 4 additions & 1 deletion snuba/manual_jobs/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -23,6 +24,8 @@ def config_key(cls) -> str:
def get_from_name(cls, name: str) -> "Job":
return cast("Job", cls.class_from_name(name))

# from kwargs vs from api endpoing


import_submodules_in_directory(
os.path.dirname(os.path.realpath(__file__)), "snuba.manual_jobs"
Expand Down
34 changes: 29 additions & 5 deletions snuba/manual_jobs/job_loader.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,39 @@
from typing import Any, cast
import json
from json import JSONDecodeError
from typing import Any, Mapping, MutableMapping, 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]) -> None:
try:
with open(filepath, "r") as json_manifest_file:
json_manifest = json.load(json_manifest_file)
kwargs["job_id"] = json_manifest["id"]
kwargs["job_type"] = json_manifest["job_type"]
if "params" in json_manifest:
for k, v in json_manifest["params"]:
kwargs[k] = v
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))
8 changes: 5 additions & 3 deletions snuba/manual_jobs/toy_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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() + "`"
)
111 changes: 105 additions & 6 deletions tests/cli/test_jobs.py
Original file line number Diff line number Diff line change
@@ -1,17 +1,116 @@
import json
import os

from click.testing import CliRunner

from snuba.cli.jobs import run
from snuba.cli.jobs import JOB_SPECIFICATION_ERROR_MSG, run


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_name",
"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_name",
"ToyJob",
"--json_manifest",
"doesntmatter.json",
"--dry_run",
"True",
"k1=v1",
"k2=v2",
],
)
assert result.exit_code == 1
assert result.output == JOB_SPECIFICATION_ERROR_MSG


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 == JOB_SPECIFICATION_ERROR_MSG


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 == JOB_SPECIFICATION_ERROR_MSG


def test_provides_only_job_name_errors() -> None:
runner = CliRunner()
result = runner.invoke(
run, ["--job_name", "ToyJob", "--dry_run", "True", "k1=v1", "k2=v2"]
)
assert result.exit_code == 1
assert result.output == JOB_SPECIFICATION_ERROR_MSG

def test_invalid_job() -> None:

def test_provides_only_job_id_errors() -> None:
runner = CliRunner()
result = runner.invoke(
run, ["SomeJobThatDoesntExist", "--dry_run", "True", "k1=v1", "k2=v2"]
run, ["--job_id", "0001", "--dry_run", "True", "k1=v1", "k2=v2"]
)
assert result.exit_code == 1
assert result.output == JOB_SPECIFICATION_ERROR_MSG


def test_provides_json_manifest_is_valid(tmp_path) -> 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_name",
"ToyJob",
"--job_id",
"0001",
"--dry_run",
"True",
"k1=v1",
"k2=v2",
],
)
assert result.exit_code == 0

0 comments on commit 1160ad8

Please sign in to comment.