Align DoubleJobRegistry.*dependency* methods with JobRegistryInterface

soxofaan · soxofaan · commit e2678e73a0b3 · 2025-04-16T18:15:07.000+02:00
refs: #863, #1123
diff --git a/openeogeotrellis/_version.py b/openeogeotrellis/_version.py
@@ -1 +1 @@
-__version__ = "0.65.0a10"
+__version__ = "0.65.0a11"
diff --git a/openeogeotrellis/backend.py b/openeogeotrellis/backend.py
@@ -1595,7 +1595,9 @@ def job_results_status(job_results_dependency: dict) -> (str, Optional[str]):
 
         def fail_job():
             with self._double_job_registry as registry:
-                registry.set_dependency_status(job_id, user_id, DEPENDENCY_STATUS.ERROR)
+                registry.set_dependency_status(
+                    job_id=job_id, user_id=user_id, dependency_status=DEPENDENCY_STATUS.ERROR
+                )
                 registry.set_status(job_id, user_id, JOB_STATUS.ERROR)
 
             job_info["status"] = JOB_STATUS.ERROR  # TODO: avoid mutation
@@ -1633,7 +1635,9 @@ def fail_job():
                 job_info.get("dependency_status") != DEPENDENCY_STATUS.AWAITING_RETRY
             ):  # haven't retried yet: retry
                 with self._double_job_registry as registry:
-                    registry.set_dependency_status(job_id, user_id, DEPENDENCY_STATUS.AWAITING_RETRY)
+                    registry.set_dependency_status(
+                        job_id=job_id, user_id=user_id, dependency_status=DEPENDENCY_STATUS.AWAITING_RETRY
+                    )
 
                 retries = [retry for details, _, retry in batch_processes.values() if details.status() == "PARTIAL"]
 
@@ -1742,11 +1746,15 @@ def processing_units_spent(value_estimate: Decimal, temporal_step: Optional[str]
                 logger.debug(f"Total cost of Sentinel Hub batch processes: {batch_process_processing_units} PU",
                              extra={'job_id': job_id, 'user_id': user_id})
 
-                registry.set_dependencies(job_id, user_id, dependencies)
-                registry.set_dependency_status(job_id, user_id, DEPENDENCY_STATUS.AVAILABLE)
+                registry.set_dependencies(job_id=job_id, user_id=user_id, dependencies=dependencies)
+                registry.set_dependency_status(
+                    job_id=job_id, user_id=user_id, dependency_status=DEPENDENCY_STATUS.AVAILABLE
+                )
 
                 if batch_process_processing_units:
-                    registry.set_dependency_usage(job_id, user_id, batch_process_processing_units)
+                    registry.set_dependency_usage(
+                        job_id=job_id, user_id=user_id, dependency_usage=batch_process_processing_units
+                    )
 
             self._start_job(job_id, User(user_id=user_id), lambda _: vault_token, dependencies)
         else:  # still some running: continue polling
@@ -1897,7 +1905,7 @@ def _start_job(self, job_id: str, user: User, get_vault_token: Callable[[str], s
                         else get_vault_token(sentinel_hub_client_alias),
                     )
                     dbl_registry.set_dependency_status(
-                        job_id, user_id, DEPENDENCY_STATUS.AWAITING
+                        job_id=job_id, user_id=user_id, dependency_status=DEPENDENCY_STATUS.AWAITING
                     )
                     dbl_registry.set_status(job_id, user_id, JOB_STATUS.QUEUED)
 
@@ -3091,7 +3099,7 @@ def cancel_job(self, job_id: str, user_id: str):
                         registry.set_status(job_id, user_id, JOB_STATUS.CANCELED)
         else:
             with self._double_job_registry as registry:
-                registry.remove_dependencies(job_id, user_id)
+                registry.remove_dependencies(job_id=job_id, user_id=user_id)
                 registry.set_status(job_id, user_id, JOB_STATUS.CANCELED)
 
     def delete_job(self, job_id: str, user_id: str):
diff --git a/openeogeotrellis/job_registry.py b/openeogeotrellis/job_registry.py
@@ -663,18 +663,18 @@ def set_status(
         return self.db[job_id]
 
     def set_dependencies(
-        self, job_id: str, dependencies: List[Dict[str, str]]
-    ) -> JobDict:
-        return self._update(job_id=job_id, dependencies=dependencies)
+        self, job_id: str, *, user_id: Optional[str] = None, dependencies: List[Dict[str, str]]
+    ) -> None:
+        self._update(job_id=job_id, dependencies=dependencies)
 
-    def remove_dependencies(self, job_id: str) -> JobDict:
-        return self._update(job_id=job_id, dependencies=None, dependency_status=None)
+    def remove_dependencies(self, job_id: str, *, user_id: Optional[str] = None) -> None:
+        self._update(job_id=job_id, dependencies=None, dependency_status=None)
 
-    def set_dependency_status(self, job_id: str, dependency_status: str) -> JobDict:
-        return self._update(job_id=job_id, dependency_status=dependency_status)
+    def set_dependency_status(self, job_id: str, *, user_id: Optional[str] = None, dependency_status: str) -> None:
+        self._update(job_id=job_id, dependency_status=dependency_status)
 
-    def set_dependency_usage(self, job_id: str, dependency_usage: Decimal) -> JobDict:
-        return self._update(job_id, dependency_usage=str(dependency_usage))
+    def set_dependency_usage(self, job_id: str, *, user_id: Optional[str] = None, dependency_usage: Decimal) -> None:
+        self._update(job_id, dependency_usage=str(dependency_usage))
 
     def set_proxy_user(self, job_id: str, proxy_user: str) -> JobDict:
         return self._update(job_id=job_id, proxy_user=proxy_user)
@@ -829,7 +829,7 @@ def get_job(self, job_id: str, *, user_id: Optional[str] = None) -> dict:
         # TODO: eliminate get_job/get_job_metadata duplication?
         zk_job = ejr_job = None
         if self.zk_job_registry:
-            assert user_id
+            assert user_id, "user_id is required in ZkJobRegistry"
             with contextlib.suppress(JobNotFoundException, ZkStrippedSpecification):
                 zk_job = self.zk_job_registry.get_job(
                     job_id=job_id, user_id=user_id, parse_specification=True, omit_raw_specification=True
@@ -881,47 +881,45 @@ def set_status(self, job_id: str, user_id: str, status: str,
 
     def delete_job(self, job_id: str, *, user_id: Optional[str] = None) -> None:
         if self.zk_job_registry:
-            assert user_id
+            assert user_id, "user_id is required in ZkJobRegistry"
             self.zk_job_registry.delete(job_id=job_id, user_id=user_id)
         if self.elastic_job_registry:
             self.elastic_job_registry.delete_job(job_id=job_id, user_id=user_id)
 
     def set_dependencies(
-        self, job_id: str, user_id: str, dependencies: List[Dict[str, str]]
-    ):
+        self, job_id: str, *, user_id: Optional[str] = None, dependencies: List[Dict[str, str]]
+    ) -> None:
         if self.zk_job_registry:
+            assert user_id, "user_id is required in ZkJobRegistry"
             self.zk_job_registry.set_dependencies(job_id=job_id, user_id=user_id, dependencies=dependencies)
         if self.elastic_job_registry:
-            self.elastic_job_registry.set_dependencies(
-                job_id=job_id, dependencies=dependencies
-            )
+            self.elastic_job_registry.set_dependencies(job_id=job_id, user_id=user_id, dependencies=dependencies)
 
-    def remove_dependencies(self, job_id: str, user_id: str):
+    def remove_dependencies(self, job_id: str, *, user_id: Optional[str] = None) -> None:
         if self.zk_job_registry:
+            assert user_id, "user_id is required in ZkJobRegistry"
             self.zk_job_registry.remove_dependencies(job_id=job_id, user_id=user_id)
         if self.elastic_job_registry:
-            self.elastic_job_registry.remove_dependencies(job_id=job_id)
+            self.elastic_job_registry.remove_dependencies(job_id=job_id, user_id=user_id)
 
-    def set_dependency_status(
-        self, job_id: str, user_id: str, dependency_status: str
-    ) -> None:
+    def set_dependency_status(self, job_id: str, *, user_id: Optional[str] = None, dependency_status: str) -> None:
         if self.zk_job_registry:
+            assert user_id, "user_id is required in ZkJobRegistry"
             self.zk_job_registry.set_dependency_status(
                 job_id=job_id, user_id=user_id, dependency_status=dependency_status
             )
         if self.elastic_job_registry:
             self.elastic_job_registry.set_dependency_status(
-                job_id=job_id, dependency_status=dependency_status
+                job_id=job_id, user_id=user_id, dependency_status=dependency_status
             )
 
-    def set_dependency_usage(
-        self, job_id: str, user_id: str, dependency_usage: Decimal
-    ):
+    def set_dependency_usage(self, job_id: str, *, user_id: Optional[str] = None, dependency_usage: Decimal) -> None:
         if self.zk_job_registry:
+            assert user_id, "user_id is required in ZkJobRegistry"
             self.zk_job_registry.set_dependency_usage(job_id=job_id, user_id=user_id, processing_units=dependency_usage)
         if self.elastic_job_registry:
             self.elastic_job_registry.set_dependency_usage(
-                job_id=job_id, dependency_usage=dependency_usage
+                job_id=job_id, user_id=user_id, dependency_usage=dependency_usage
             )
 
     def set_proxy_user(self, job_id: str, user_id: str, proxy_user: str):
diff --git a/openeogeotrellis/job_tracker_v2.py b/openeogeotrellis/job_tracker_v2.py
@@ -511,7 +511,7 @@ def _sync_job_status(
             stats[f"reached final status {job_metadata.status}"] += 1
             result_metadata = self._batch_jobs.load_results_metadata(job_id, user_id)
 
-            double_job_registry.remove_dependencies(job_id, user_id)
+            double_job_registry.remove_dependencies(job_id=job_id, user_id=user_id)
 
             # there can be duplicates if batch processes are recycled
             dependency_sources = list(set(get_deletable_dependency_sources(job_info)))
diff --git a/setup.py b/setup.py
@@ -57,7 +57,7 @@
     tests_require=tests_require,
     install_requires=[
         "openeo>=0.33.0",
-        "openeo_driver>=0.133.0.a3.dev",
+        "openeo_driver>=0.133.0.a5.dev",
         'pyspark==3.5.3; python_version>="3.8"',
         'pyspark>=2.3.1,<2.4.0; python_version<"3.8"',
         'geopyspark==0.4.9+openeo',

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-__version__ = "0.65.0a10"`
	`1`	`+__version__ = "0.65.0a11"`