Fix bug with cursor state naemspace

Flared · May 7, 2024 · bf18f8e · bf18f8e
1 parent d4b7727
commit bf18f8e
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 6 deletions.
diff --git a/example/definitions/simple.yaml b/example/definitions/simple.yaml
@@ -210,6 +210,7 @@ spec:
     config:
       job_state:
         cursors_states_enabled: true
+        cursors_states_namespace: example-namespace
       job:
         buffer_size: 5
 
diff --git a/src/saturn_engine/stores/jobs_store.py b/src/saturn_engine/stores/jobs_store.py
@@ -4,6 +4,7 @@
 import time
 from datetime import datetime
 
+import sqlalchemy as sa
 from sqlalchemy import select
 from sqlalchemy import union_all
 from sqlalchemy import update
@@ -190,13 +191,19 @@ def fetch_cursors_states(
     fetch_stmts = []
     for job, cursors in query.items():
         fetch_stmts.append(
-            select(Job.name, JobCursorState)
+            select(JobCursorState, sa.func.coalesce(Job.name, job).label("name"))
             .join(
-                JobCursorState,
+                Job,
                 Job.job_definition_name == JobCursorState.job_definition_name,
+                isouter=True,
             )
             .where(
-                Job.name == job,
+                sa.or_(
+                    Job.name == job,
+                    sa.and_(
+                        Job.name.is_(None), JobCursorState.job_definition_name == job
+                    ),
+                ),
                 JobCursorState.cursor.in_(cursors),
             )
         )

diff --git a/tests/worker_manager/api/test_jobs.py b/tests/worker_manager/api/test_jobs.py
@@ -630,7 +630,7 @@ def test_sync_states(
                             "completed_at": "2020-01-01T01:02:03+04:00",
                             "error": "ValueError: boo",
                         },
-                    }
+                    },
                 }
             }
         },
@@ -649,7 +649,7 @@ def test_sync_states(
                         },
                     },
                     orphan_job.name: {"cursors_states": {"a": {"x": 1}}},
-                }
+                },
             }
         },
     )
@@ -677,6 +677,7 @@ def test_fetch_cursors_states(
     new_job: Job,
     orphan_job: Job,
 ) -> None:
+    custom_ns = "custom-namespace"
     resp = client.post(
         "/api/jobs/_states",
         json={
@@ -687,7 +688,12 @@ def test_fetch_cursors_states(
                             "a": {"x": 1},
                             "b": {"x": 2},
                         },
-                    }
+                    },
+                    custom_ns: {
+                        "cursors_states": {
+                            "c": {"x": 1},
+                        }
+                    },
                 }
             }
         },
@@ -702,6 +708,7 @@ def test_fetch_cursors_states(
                 new_job.name: ["a", "c"],
                 orphan_job.name: ["a"],
                 "do-not-exist": ["a"],
+                custom_ns: ["c", "d"],
             }
         },
     )
@@ -714,6 +721,10 @@ def test_fetch_cursors_states(
             },
             orphan_job.name: {"a": None},
             "do-not-exist": {"a": None},
+            custom_ns: {
+                "c": {"x": 1},
+                "d": None,
+            },
         }
     }