Merge pull request #54 from Small-Bodies-Node/queue-position-in-catch-route

Queue position in catch route

Author: mkelley

live-tests/test_alive.py

@@ -10,34 +10,32 @@
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
-    parser.add_argument('--url', default='http://127.0.0.1:5000')
-    parser.add_argument('--target', default='65P')
-    parser.add_argument('--source', '-s', default=None, action='append')
-    parser.add_argument('--no-cached', '--force', '-f',
-                        action='store_false', dest='cached')
-    parser.add_argument('-v', dest='verbose', action='store_true',
-                        help='verbose mode')
+    parser.add_argument("--url", default="http://127.0.0.1:5000")
+    parser.add_argument("--target", default="65P")
+    parser.add_argument("--source", "-s", default=None, action="append")
+    parser.add_argument(
+        "--no-cached", "--force", "-f", action="store_false", dest="cached"
+    )
+    parser.add_argument("-v", dest="verbose", action="store_true", help="verbose mode")
     return parser.parse_args()
 
 
 def query(args: argparse.Namespace) -> Tuple[str, float, Any]:
     start = time.monotonic()
 
-    params = {
-        'target': args.target,
-        'cached': args.cached
-    }
+    params = {"target": args.target, "cached": args.cached}
     if args.source is not None:
-        params['sources'] = args.source
+        params["sources"] = args.source
 
-    res = requests.get(f'{args.url}/catch', params=params)
+    res = requests.get(f"{args.url}/catch", params=params)
     data = res.json()
     if args.verbose:
         print(data)
         print()
 
-    if data['queued']:
-        messages = SSEClient(data['message_stream'], chunk_size=128)
+    if data["queued"]:
+        print(f"Queue position: {data['queue_position']}")
+        messages = SSEClient(data["message_stream"], chunk_size=128)
         for message in messages:
             if len(message.data) == 0:
                 continue
@@ -51,58 +49,61 @@ def query(args: argparse.Namespace) -> Tuple[str, float, Any]:
             if not isinstance(message_data, dict):
                 continue
 
-            if message_data['job_prefix'] != data['job_id'][:8]:
+            if message_data["job_prefix"] != data["job_id"][:8]:
                 # this mesage is not for us
                 continue
 
             # this message is for us, print the text
-            print(message_data['text'], file=sys.stderr)
+            print(message_data["text"], file=sys.stderr)
 
             # Message status may be 'success', 'error', 'running', 'queued'.
-            if message_data['status'] == 'error':
-                raise Exception(message_data['text'])
+            if message_data["status"] == "error":
+                raise Exception(message_data["text"])
 
-            if message_data['status'] == 'success':
+            if message_data["status"] == "success":
                 break
 
-    elif data['results'] is None:
-        raise Exception(data['message'])
+    elif data["results"] is None:
+        raise Exception(data["message"])
 
     # 'results' is the URL to the search results
-    res = requests.get(data['results'])
+    res = requests.get(data["results"])
     dt = time.monotonic() - start
 
     # response is JSON formatted
-    return data['results'], dt, res.json()
+    return data["results"], dt, res.json()
 
 
 def summarize(results_url: str, dt: float, data: Any) -> None:
     count_by_survey = []
-    for survey in set([row['source'] for row in data['data']]):
-        n = len([row for row in data['data']
-                 if row['source'] == survey])
-        count_by_survey.append(f'  - {n} {survey}\n')
+    for survey in set([row["source"] for row in data["data"]]):
+        n = len([row for row in data["data"] if row["source"] == survey])
+        count_by_survey.append(f"  - {n} {survey}\n")
 
-    print(f'''
+    print(
+        f"""
 Job ID: {data['job_id']}
 Results: {results_url}
 Elapsed time: {dt:.1f} s
-Count: {data['count']}''')
+Count: {data['count']}"""
+    )
 
     if len(count_by_survey) > 0:
-        print(''.join(count_by_survey))
+        print("".join(count_by_survey))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     args = parse_args()
 
-    print(f'''
+    print(
+        f"""
 CATCH APIs Query
 
 Base URL: {args.url}
 Target: {args.target}
 Cached results allowed? {args.cached}
-''')
+"""
+    )
 
     results_url, dt, data = query(args)
     summarize(results_url, dt, data)

live-tests/test_query_fixed.py

@@ -14,10 +14,12 @@
 import numpy as np
 from catch_apis.app import app
 
+
 @pytest.fixture()
 def test_client() -> TestClient:
     return app.test_client()
 
+
 def test_point_full_search(test_client: TestClient):
     # Crab nebula
     parameters = {
@@ -31,13 +33,36 @@ def test_point_full_search(test_client: TestClient):
     assert results["message"] == ""
     assert np.isclose(results["query"]["ra"], (5 + (34 + 32 / 60) / 60) * 15)
     assert np.isclose(results["query"]["dec"], 22 + 48 / 60 / 60)
-    assert "neat_palomar_tricam" in results["query"]["sources"]
-    assert len(results["data"]) == 428
+    assert set([row["source"] for row in results["data"]]) == {
+        "atlas_haleakela",
+        "atlas_mauna_loa",
+        "catalina_bigelow",
+        "catalina_lemmon",
+        "neat_maui_geodss",
+        "ps1dr2",
+        "spacewatch",
+    }
+
+    # only count static data sets
+    sources = [row["source"] for row in results["data"]]
+    expected = {
+        "neat_maui_geodss": 1,
+        "neat_palomar_tricam": 0,
+        "ps1dr2": 73,
+        "spacewatch": 166,
+    }
+    for source, count in expected.items():
+        assert sources.count(source) == count
+
     product_ids = [obs["product_id"] for obs in results["data"]]
 
-    # these are verified to be the Crab:
+    # verified to be the Crab:
     assert "rings.v3.skycell.1784.059.wrp.g.55560_46188.fits" in product_ids
-    assert "urn:nasa:pds:gbo.ast.catalina.survey:data_calibrated:703_20211108_2b_n24018_01_0002.arch" in product_ids
+    assert (
+        "urn:nasa:pds:gbo.ast.catalina.survey:data_calibrated:703_20211108_2b_n24018_01_0002.arch"
+        in product_ids
+    )
+
 
 def test_point_date_range(test_client: TestClient):
     # Crab nebula
@@ -51,30 +76,31 @@ def test_point_date_range(test_client: TestClient):
     results = response.json()
     assert results["query"]["start_date"] is None
     assert results["query"]["stop_date"] is None
-    assert len(results["data"]) == 167
+    assert len(results["data"]) == 166
 
     parameters["start_date"] = "2007-01-01"
     response = test_client.get("/fixed", params=parameters)
     response.raise_for_status()
     results = response.json()
     assert results["query"]["start_date"] == "2007-01-01 00:00:00.000"
-    assert len(results["data"]) == 118
+    assert len(results["data"]) == 117
 
     parameters["stop_date"] = "2008-01-01"
     response = test_client.get("/fixed", params=parameters)
     response.raise_for_status()
     results = response.json()
     assert results["query"]["stop_date"] == "2008-01-01 00:00:00.000"
     assert len(results["data"]) == 24
-    
+
     del parameters["start_date"]
     response = test_client.get("/fixed", params=parameters)
     response.raise_for_status()
     results = response.json()
     assert results["query"]["start_date"] is None
     assert len(results["data"]) == 73
 
-    # 118 + 73 - 24 = 167
+    # 117 + 73 - 24 = 166
+
 
 def test_areal_search(test_client: TestClient):
     # Crab nebula
@@ -107,4 +133,4 @@ def test_areal_search(test_client: TestClient):
     response = test_client.get("/fixed", params=parameters)
     response.raise_for_status()
     results = response.json()
-    assert len(results["data"]) == 73
+    assert len(results["data"]) == 73

live-tests/test_query_moving.py

@@ -26,7 +26,7 @@
 ENV.REDIS_TASK_MESSAGES = "TASK_MESSAGES_TESTING"
 
 from catch_apis.app import app
-from catch_apis.services.stream import messages as stream_messages
+from catch_apis.services.stream import messages_service
 from catch_apis import woRQer
 
 
@@ -39,8 +39,8 @@ def test_client() -> TestClient:
 def mock_stream_messages(timeout):
     """Patch message stream to timeout in an absolute amount of time."""
     with mock.patch(
-        "catch_apis.services.stream.messages",
-        partial(stream_messages, timeout),
+        "catch_apis.services.stream.messages_service",
+        partial(messages_service, timeout),
     ):
         yield
 
@@ -124,7 +124,7 @@ def _query(
         # hang until the timeout is reached, run a worker in "burst" mode, then
         # directly read messages from the message function
         woRQer.run(True)
-        messages = [message for message in stream_messages(1)]
+        messages = [message for message in messages_service(1)]
         for message in messages:
             if len(message) == 0 or not message.startswith("data:"):
                 continue
@@ -160,11 +160,11 @@ def _query(
 def test_equivalencies(test_client: TestClient, targets: List[str]) -> None:
     source = "neat_maui_geodss"
     catch0, caught0, queued0 = _query(test_client, targets[0], True, source=source)
-    data0 = caught0["data"]
+    data0 = sorted(caught0["data"], key=lambda row: row["product_id"])
     assert len(data0) > 0
     for target in targets[1:]:
         catch, caught, queued = _query(test_client, target, True, source=source)
-        data = caught["data"]
+        data = sorted(caught["data"], key=lambda row: row["product_id"])
         for a, b in zip(data0, data):
             for k in COMPARE_KEYS:
                 # np.isclose using rtol = 1% in case of ephemeris updates and

src/catch_apis/api/__init__.py

@@ -1,4 +0,0 @@
-from .catch import catch  # noqa: F401
-from .caught import caught  # noqa: F401
-from .fixed import fixed_target_query  # noqa: F401
- 

src/catch_apis/api/catch.py

@@ -9,7 +9,9 @@
 from astropy.time import Time
 
 from ..config import allowed_sources, get_logger, QueryStatus
-from .. import services
+from ..services.target_name import parse_target_name
+from ..services.catch import catch_service
+from ..services.queue import JobsQueue
 from ..services.message import (
     Message,
     listen_for_task_messages,
@@ -32,7 +34,7 @@ def _format_date(date):
     return date if date is None else date.iso
 
 
-def catch(
+def catch_controller(
     target: str,
     sources: Optional[List[str]] = None,
     start_date: Optional[str] = None,
@@ -75,7 +77,7 @@ def catch(
 
     target_type: str
     sanitized_target: str
-    target_type, sanitized_target = services.parse_target_name(target)
+    target_type, sanitized_target = parse_target_name(target)
     if sanitized_target == "":
         messages.append("Invalid target: empty string")
         valid_query = False
@@ -117,14 +119,15 @@ def catch(
         "job_id": job_id.hex,
         "queued": False,
         "queue_full": False,
+        "queue_position": None,
         "message": None,
         "version": version,
     }
 
     Message.reset_t0()
     listen_for_task_messages(job_id)
 
-    status: QueryStatus = services.catch(
+    status: QueryStatus = catch_service(
         job_id,
         sanitized_target,
         sources=_sources,
@@ -142,8 +145,13 @@ def catch(
     message_stream_url: str = urllib.parse.urlunsplit(
         (parsed[0], parsed[1], os.path.join(parsed[2], "stream"), "", "")
     )
-
     if status == QueryStatus.QUEUED:
+        queue = JobsQueue()
+        for job in queue.jobs:
+            if job.args[0].hex == job_id.hex:
+                result["queue_position"] = job.get_position()
+                break
+
         result["queued"] = True
         result["message_stream"] = message_stream_url
         result["results"] = results_url

src/catch_apis/api/caught.py

@@ -1,20 +1,20 @@
 import uuid
-from typing import Any, Dict, List, Tuple, Union
 
-from .. import services
+from ..services.caught import caught_service
+from ..services.status.job_id import job_id_service
 from .. import __version__ as version
 
 
-def caught(job_id: str) -> Union[dict, Tuple[str, int]]:
+def caught_controller(job_id: str) -> dict | tuple[str, int]:
     """Controller for returning caught data."""
 
     try:
         _job_id: uuid.UUID = uuid.UUID(job_id, version=4)
     except ValueError:
         return "Invalid job ID", 400
 
-    parameters, status = services.status.job_id(_job_id)
-    data: List[Dict[str, Any]] = services.caught(_job_id)
+    parameters, status = job_id_service(_job_id)
+    data = caught_service(_job_id)
     return {
         "parameters": parameters,
         "status": status,

src/catch_apis/api/fixed.py

@@ -7,7 +7,7 @@
 from astropy.coordinates import Angle
 import astropy.units as u
 
-from .. import services
+from ..services.fixed import fixed_target_query_service
 from ..config import CatchApisException, get_logger, allowed_sources
 from .. import __version__ as version
 
@@ -60,7 +60,7 @@ def _format_date(date):
     return date if date is None else date.iso
 
 
-def fixed_target_query(
+def fixed_target_query_controller(
     ra: str,
     dec: str,
     sources: Optional[List[str]] = None,
@@ -128,7 +128,7 @@ def fixed_target_query(
     data: List[dict] = []
     if valid_query:
         try:
-            data = services.fixed_target_query(
+            data = fixed_target_query_service(
                 job_id,
                 sanitized_ra,
                 sanitized_dec,