feat(job-runner): track job status

snuba/cli/jobs.py

@@ -2,7 +2,7 @@
 
 import click
 
-from snuba.manual_jobs import JobSpec
+from snuba.manual_jobs import JobSpec, get_job_status
 from snuba.manual_jobs.job_loader import JobLoader
 from snuba.manual_jobs.manifest_reader import _ManifestReader
 
@@ -49,3 +49,9 @@ def run(*, job_type: str, job_id: str, dry_run: bool, pairs: Tuple[str, ...]) ->
 
     job_to_run = JobLoader.get_job_instance(job_spec, dry_run)
     job_to_run.execute()
+
+
+@jobs.command()
+@click.option("--job_id")
+def status(*, job_id: str) -> str:
+    return get_job_status(job_id)

snuba/manual_jobs/__init__.py

@@ -1,19 +1,27 @@
 import logging
 import os
+import typing
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Any, MutableMapping, Optional, cast
 
+from snuba.redis import RedisClientKey, get_redis_client
 from snuba.utils.registered_class import RegisteredClass, import_submodules_in_directory
 
 logger = logging.getLogger("snuba_init")
+redis_client = get_redis_client(RedisClientKey.JOB)
+job_status_hash = "snuba-job-status-hash"
+
+RUNNING = "running"
+FINISHED = "finished"
+NOT_STARTED = "not started"
 
 
 @dataclass
 class JobSpec:
     job_id: str
     job_type: str
-    params: Optional[MutableMapping[Any, Any]]
+    params: Optional[MutableMapping[Any, Any]] = None
 
 
 class Job(ABC, metaclass=RegisteredClass):
@@ -23,9 +31,15 @@ def __init__(self, job_spec: JobSpec, dry_run: bool) -> None:
         if job_spec.params:
             for k, v in job_spec.params.items():
                 setattr(self, k, v)
+        self._set_status(NOT_STARTED)
 
-    @abstractmethod
     def execute(self) -> None:
+        self._set_status(RUNNING)
+        self._execute()
+        self._set_status(FINISHED)
+
+    @abstractmethod
+    def _execute(self) -> None:
         pass
 
     @classmethod
@@ -36,6 +50,17 @@ def config_key(cls) -> str:
     def get_from_name(cls, name: str) -> "Job":
         return cast("Job", cls.class_from_name(name))
 
+    def _set_status(self, status: str) -> None:
+        redis_client.hset(job_status_hash, self.job_spec.job_id, status)
+
+
+def get_job_status(job_id: str) -> str:
+    status = redis_client.hget(job_status_hash, job_id)
+    if status is None:
+        raise KeyError
+    else:
+        return typing.cast(str, status.decode("utf-8"))
+
 
 import_submodules_in_directory(
     os.path.dirname(os.path.realpath(__file__)), "snuba.manual_jobs"

snuba/manual_jobs/toy_job.py

@@ -15,7 +15,7 @@ def _build_query(self) -> str:
         else:
             return "not dry run query"
 
-    def execute(self) -> None:
+    def _execute(self) -> None:
         logger.info(
             "executing job "
             + self.job_spec.job_id

snuba/redis.py

@@ -114,6 +114,7 @@ class RedisClientKey(Enum):
     DLQ = "dlq"
     OPTIMIZE = "optimize"
     ADMIN_AUTH = "admin_auth"
+    JOB = "job"
 
 
 _redis_clients: Mapping[RedisClientKey, RedisClientType] = {
@@ -141,6 +142,9 @@ class RedisClientKey(Enum):
     RedisClientKey.ADMIN_AUTH: _initialize_specialized_redis_cluster(
         settings.REDIS_CLUSTERS["admin_auth"]
     ),
+    RedisClientKey.JOB: _initialize_specialized_redis_cluster(
+        settings.REDIS_CLUSTERS["job"]
+    ),
 }
 
 

snuba/settings/__init__.py

@@ -175,6 +175,7 @@ class RedisClusters(TypedDict):
     dlq: RedisClusterConfig | None
     optimize: RedisClusterConfig | None
     admin_auth: RedisClusterConfig | None
+    job: RedisClusterConfig | None
 
 
 REDIS_CLUSTERS: RedisClusters = {
@@ -186,6 +187,7 @@ class RedisClusters(TypedDict):
     "dlq": None,
     "optimize": None,
     "admin_auth": None,
+    "job": None,
 }
 
 # Query Recording Options

snuba/settings/settings_test.py

@@ -60,6 +60,7 @@
         (7, "dlq"),
         (8, "optimize"),
         (9, "admin_auth"),
+        (10, "job"),
     ]
 }
 VALIDATE_DATASET_YAMLS_ON_STARTUP = True

tests/cli/test_jobs.py

@@ -1,6 +1,7 @@
+import pytest
 from click.testing import CliRunner
 
-from snuba.cli.jobs import JOB_SPECIFICATION_ERROR_MSG, run, run_from_manifest
+from snuba.cli.jobs import JOB_SPECIFICATION_ERROR_MSG, run, run_from_manifest, status
 
 
 def test_cmd_line_valid() -> None:
@@ -69,3 +70,19 @@ def test_json_valid() -> None:
         ],
     )
     assert result.exit_code == 0
+
+
+@pytest.mark.redis_db
+def test_jobs_status() -> None:
+    runner = CliRunner()
+    runner.invoke(
+        run_from_manifest,
+        [
+            "--json_manifest",
+            "job_manifest.json",
+            "--job_id",
+            "abc1234",
+        ],
+    )
+    result = runner.invoke(status, ["--job_id", "abc1234"])
+    assert result.exit_code == 0

tests/jobs/test_job.py

@@ -0,0 +1,27 @@
+import pytest
+
+from snuba.manual_jobs import FINISHED, NOT_STARTED, JobSpec, get_job_status
+from snuba.manual_jobs.toy_job import ToyJob
+
+JOB_ID = "abc1234"
+test_job_spec = JobSpec(job_id=JOB_ID, job_type="ToyJob")
+
+# @pytest.fixture(autouse=True)
+# def setup_test_job() -> Job:
+#     return ToyJob(test_job_spec, dry_run=True)
+
+
+@pytest.mark.redis_db
+def test_job_status_changes_to_finished() -> None:
+    # test_job = setup_test_job
+    test_job = ToyJob(test_job_spec, dry_run=True)
+    assert get_job_status(JOB_ID) == NOT_STARTED
+    test_job.execute()
+    assert get_job_status(JOB_ID) == FINISHED
+
+
+@pytest.mark.redis_db
+def test_job_status_with_invalid_job_id() -> None:
+    assert get_job_status(JOB_ID) == NOT_STARTED
+    with pytest.raises(KeyError):
+        get_job_status("invalid_job_id")