PR #736 addressing some review notes (exclude bearer from repr)

and some code style tweaks

Author: soxofaan

openeo/extra/job_management/__init__.py

@@ -553,6 +553,7 @@ def _job_update_loop(
         not_started = job_db.get_by_status(statuses=["not_started"], max=200).copy()
         if len(not_started) > 0:
             # Check number of jobs running at each backend
+            # TODO: should "created" be included in here? Calling this "running" is quite misleading then.
             running = job_db.get_by_status(statuses=["created", "queued", "queued_for_start", "running"])
             stats["job_db get_by_status"] += 1
             per_backend = running.groupby("backend_name").size().to_dict()
@@ -570,7 +571,7 @@ def _job_update_loop(
                         stats["job_db persist"] += 1
                         total_added += 1
 
-        self._process_threadworker_updates(self._worker_pool, job_db, stats)
+        self._process_threadworker_updates(self._worker_pool, job_db=job_db, stats=stats)
 
         # TODO: move this back closer to the `_track_statuses` call above, once job done/error handling is also handled in threads?
         for job, row in jobs_done:
@@ -641,7 +642,7 @@ def _launch_job(self, start_job, df, i, backend_name, stats: Optional[dict] = No
                                 root_url=job_con.root_url,
                                 bearer_token=job_con.auth.bearer if isinstance(job_con.auth, BearerAuth) else None,
                                 job_id=job.job_id,
-                                df_idx = i
+                                df_idx=i,
                             )
                             _log.info(f"Submitting task {task} to thread pool")
                             self._worker_pool.submit_task(task)
@@ -659,8 +660,9 @@ def _launch_job(self, start_job, df, i, backend_name, stats: Optional[dict] = No
 
     def _process_threadworker_updates(
         self,
-        worker_pool: '_JobManagerWorkerThreadPool',
-        job_db: 'JobDatabaseInterface',
+        worker_pool: _JobManagerWorkerThreadPool,
+        *,
+        job_db: JobDatabaseInterface,
         stats: Dict[str, int],
     ) -> None:
         """
@@ -669,8 +671,6 @@ def _process_threadworker_updates(
         (matched by df_idx) are upserted via job_db.persist(). Any results
         targeting unknown df_idx indices are logged as errors but not persisted.
 
-
-
         :param worker_pool: Thread-pool managing asynchronous Task executes
         :param job_db:      Interface to append/upsert to the job database
         :param stats:       Dictionary accumulating statistic counters
@@ -684,44 +684,43 @@ def _process_threadworker_updates(
             # Process database updates
             if res.db_update:
                 try:
-                    updates.append({
-                        'id': res.job_id,
-                        'df_idx': res.df_idx,
-                        **res.db_update,
-                    })
+                    updates.append(
+                        {
+                            "id": res.job_id,
+                            "df_idx": res.df_idx,
+                            **res.db_update,
+                        }
+                    )
                 except Exception as e:
-                    _log.error(f"Skipping invalid db_update '{res.db_update}' for job '{res.job_id}': {e}", )
-                    
+                    _log.error(f"Skipping invalid db_update {res.db_update!r} for job {res.job_id!r}: {e}")
+
             # Process stats updates
             if res.stats_update:
                 try:
                     for key, val in res.stats_update.items():
                         count = int(val)
                         stats[key] = stats.get(key, 0) + count
                 except Exception as e:
-                    _log.error(
-                        f"Skipping invalid stats_update {res.stats_update} for job '{res.job_id}': {e}"
-                    )
+                    _log.error(f"Skipping invalid stats_update {res.stats_update!r} for job {res.job_id!r}: {e}")
 
         # No valid updates: nothing to persist
         if not updates:
             return
 
         # Build DataFrame of updates indexed by df_idx
-        df_updates = pd.DataFrame(updates).set_index('df_idx', drop=True)
+        df_updates = pd.DataFrame(updates).set_index("df_idx", drop=True)
 
         # Determine which rows to upsert
         existing_indices = set(df_updates.index).intersection(job_db.read().index)
         if existing_indices:
             df_upsert = df_updates.loc[sorted(existing_indices)]
             job_db.persist(df_upsert)
-            stats['job_db persist'] = stats.get('job_db persist', 0) + 1
+            stats["job_db persist"] = stats.get("job_db persist", 0) + 1
 
         # Any df_idx not in original index are errors
         missing = set(df_updates.index) - existing_indices
         if missing:
-            _log.error(f"Skipping non-existing dataframe indiches: {sorted(missing)}")
-
+            _log.error(f"Skipping non-existing dataframe indices: {sorted(missing)}")
 
     def on_job_done(self, job: BatchJob, row):
         """
@@ -977,12 +976,11 @@ def get_by_status(self, statuses, max=None) -> pd.DataFrame:
 
     def _merge_into_df(self, df: pd.DataFrame):
         if self._df is not None:
-            self._df.update(df, overwrite=True) 
+            self._df.update(df, overwrite=True)
         else:
             self._df = df
 
 
-
 class CsvJobDatabase(FullDataFrameJobDatabase):
     """
     Persist/load job metadata with a CSV file.

openeo/extra/job_management/_thread_worker.py

@@ -35,7 +35,7 @@ class _TaskResult:
     """
 
     job_id: str  # Mandatory
-    df_idx: int # Mandatory
+    df_idx: int  # Mandatory
     db_update: Dict[str, Any] = field(default_factory=dict)  # Optional
     stats_update: Dict[str, int] = field(default_factory=dict)  # Optional
 
@@ -56,14 +56,14 @@ class Task(ABC):
 
     :param job_id:
         Identifier of the job to start on the backend.
-    
+
     :param df_idx:
         Index of the row of the job in the dataframe.
 
     """
 
     job_id: str
-    df_idx: int 
+    df_idx: int
 
     @abstractmethod
     def execute(self) -> _TaskResult:
@@ -88,7 +88,7 @@ class ConnectedTask(Task):
     """
 
     root_url: str
-    bearer_token: Optional[str]
+    bearer_token: Optional[str] = field(default=None, repr=False)
 
     def get_connection(self) -> openeo.Connection:
         connection = openeo.connect(self.root_url)
@@ -118,7 +118,7 @@ def execute(self) -> _TaskResult:
             _log.info(f"Job {self.job_id!r} started successfully")
             return _TaskResult(
                 job_id=self.job_id,
-                df_idx = self.df_idx,
+                df_idx=self.df_idx,
                 db_update={"status": "queued"},
                 stats_update={"job start": 1},
             )
@@ -127,9 +127,9 @@ def execute(self) -> _TaskResult:
             # TODO: more insights about the failure (e.g. the exception) are just logged, but lost from the result
             return _TaskResult(
                 job_id=self.job_id,
-                df_idx = self.df_idx,
+                df_idx=self.df_idx,
                 db_update={"status": "start_failed"},
-                stats_update={"start_job error": 1}
+                stats_update={"start_job error": 1},
             )
 
 
@@ -188,7 +188,7 @@ def process_futures(self, timeout: Union[float, None] = 0) -> Tuple[List[_TaskRe
                     _log.exception(f"Threaded task {task!r} failed: {e!r}")
                     result = _TaskResult(
                         job_id=task.job_id,
-                        df_idx = task.df_idx,
+                        df_idx=task.df_idx,
                         db_update={"status": "threaded task failed"},
                         stats_update={"threaded task failed": 1},
                     )

tests/extra/job_management/test_job_management.py

@@ -94,15 +94,14 @@ class DummyTask(Task):
     """
 
     def __init__(self, job_id, df_idx, db_update, stats_update):
-        super().__init__(job_id=job_id, df_idx = df_idx)
+        super().__init__(job_id=job_id, df_idx=df_idx)
         self._db_update = db_update or {}
         self._stats_update = stats_update or {}
 
     def execute(self) -> _TaskResult:
-
         return _TaskResult(
             job_id=self.job_id,
-            df_idx = self.df_idx,
+            df_idx=self.df_idx,
             db_update=self._db_update,
             stats_update=self._stats_update,
         )
@@ -739,7 +738,7 @@ def get_status(job_id, current_status):
         # Mock sleep() to skip one hour at a time instead of actually sleeping
         with mock.patch.object(openeo.extra.job_management.time, "sleep", new=lambda s: time_machine.shift(60 * 60)):
             job_manager.run_jobs(df=df, start_job=self._create_year_job, job_db=job_db_path)
-    
+
         final_df = CsvJobDatabase(job_db_path).read()
 
         # Validate running_start_time is a valid datetime object
@@ -759,10 +758,12 @@ def test_process_threadworker_updates(self, tmp_path, caplog):
         # Invalid index (not in DB)
         pool.submit_task(DummyTask("j-missing", df_idx=4, db_update={"status": "created"}, stats_update=None))
 
-        df_initial = pd.DataFrame({
-            "id": ["j-0", "j-1", "j-2", "j-3"],
-            "status": ["created", "created", "created", "created"],
-        })
+        df_initial = pd.DataFrame(
+            {
+                "id": ["j-0", "j-1", "j-2", "j-3"],
+                "status": ["created", "created", "created", "created"],
+            }
+        )
         job_db = CsvJobDatabase(tmp_path / "jobs.csv").initialize_from_df(df_initial)
 
         mgr = MultiBackendJobManager(root_dir=tmp_path / "jobs")
@@ -786,7 +787,7 @@ def test_process_threadworker_updates(self, tmp_path, caplog):
         assert stats["job_db persist"] == 1
 
         # Assert error log for invalid index
-        assert any("Skipping non-existing dataframe indiches" in msg for msg in caplog.messages)
+        assert any("Skipping non-existing dataframe indices" in msg for msg in caplog.messages)
 
     def test_no_results_leaves_db_and_stats_untouched(self, tmp_path, caplog):
         pool = _JobManagerWorkerThreadPool(max_workers=2)
@@ -796,13 +797,12 @@ def test_no_results_leaves_db_and_stats_untouched(self, tmp_path, caplog):
         job_db = CsvJobDatabase(tmp_path / "jobs.csv").initialize_from_df(df_initial)
         mgr = MultiBackendJobManager(root_dir=tmp_path / "jobs")
 
-        mgr._process_threadworker_updates(pool, job_db, stats)
+        mgr._process_threadworker_updates(pool, job_db=job_db, stats=stats)
 
         df_final = job_db.read()
         assert df_final.loc[0, "status"] == "created"
         assert stats == {}
 
-
     def test_logs_on_invalid_update(self, tmp_path, caplog):
         pool = _JobManagerWorkerThreadPool(max_workers=2)
         stats = collections.defaultdict(int)
@@ -824,7 +824,7 @@ def execute(self):
         mgr = MultiBackendJobManager(root_dir=tmp_path / "jobs")
 
         with caplog.at_level(logging.ERROR):
-            mgr._process_threadworker_updates(pool, job_db, stats)
+            mgr._process_threadworker_updates(pool, job_db=job_db, stats=stats)
 
         # DB should remain unchanged
         df_final = job_db.read()

tests/extra/job_management/test_thread_worker.py

@@ -25,9 +25,9 @@ def dummy_backend(requests_mock) -> DummyBackend:
 
 class TestTaskResult:
     def test_default(self):
-        result = _TaskResult(job_id="j-123", df_idx = 0)
+        result = _TaskResult(job_id="j-123", df_idx=0)
         assert result.job_id == "j-123"
-        assert result.df_idx ==0
+        assert result.df_idx == 0
         assert result.db_update == {}
         assert result.stats_update == {}
 
@@ -37,12 +37,14 @@ def test_start_success(self, dummy_backend, caplog):
         caplog.set_level(logging.WARNING)
         job = dummy_backend.connection.create_job(process_graph={})
 
-        task = _JobStartTask(job_id=job.job_id, df_idx=0, root_url=dummy_backend.connection.root_url, bearer_token="h4ll0")
+        task = _JobStartTask(
+            job_id=job.job_id, df_idx=0, root_url=dummy_backend.connection.root_url, bearer_token="h4ll0"
+        )
         result = task.execute()
 
         assert result == _TaskResult(
             job_id="job-000",
-            df_idx = 0,
+            df_idx=0,
             db_update={"status": "queued"},
             stats_update={"job start": 1},
         )
@@ -54,7 +56,9 @@ def test_start_failure(self, dummy_backend, caplog):
         job = dummy_backend.connection.create_job(process_graph={})
         dummy_backend.setup_job_start_failure()
 
-        task = _JobStartTask(job_id=job.job_id, df_idx=0, root_url=dummy_backend.connection.root_url, bearer_token="h4ll0")
+        task = _JobStartTask(
+            job_id=job.job_id, df_idx=0, root_url=dummy_backend.connection.root_url, bearer_token="h4ll0"
+        )
         result = task.execute()
 
         assert result == _TaskResult(
@@ -68,7 +72,13 @@ def test_start_failure(self, dummy_backend, caplog):
             "Failed to start job 'job-000': OpenEoApiError('[500] Internal: No job starting " "for you, buddy')"
         ]
 
-
+    @pytest.mark.parametrize("serializer", [repr, str])
+    def test_hide_token(self, serializer):
+        secret = "Secret!"
+        task = _JobStartTask(job_id="job-123", df_idx=0, root_url="https://example.com", bearer_token=secret)
+        serialized = serializer(task)
+        assert "job-123" in serialized
+        assert secret not in serialized
 
 
 class NopTask(Task):
@@ -107,8 +117,6 @@ def execute(self) -> _TaskResult:
         return _TaskResult(job_id=self.job_id, df_idx=self.df_idx, db_update={"status": "all fine"})
 
 
-
-
 class TestJobManagerWorkerThreadPool:
     @pytest.fixture
     def worker_pool(self) -> Iterator[_JobManagerWorkerThreadPool]:
@@ -146,7 +154,7 @@ def test_submit_and_process_with_error(self, worker_pool):
         assert results == [
             _TaskResult(
                 job_id="j-666",
-                df_idx = 0,
+                df_idx=0,
                 db_update={"status": "threaded task failed"},
                 stats_update={"threaded task failed": 1},
             ),
@@ -163,7 +171,7 @@ def test_submit_and_process_iterative(self, worker_pool):
         worker_pool.submit_task(NopTask(job_id="j-22", df_idx=22))
         worker_pool.submit_task(NopTask(job_id="j-222", df_idx=222))
         results, remaining = worker_pool.process_futures(timeout=1)
-        assert results == [_TaskResult(job_id="j-22", df_idx=22), _TaskResult(job_id="j-222",  df_idx=222)]
+        assert results == [_TaskResult(job_id="j-22", df_idx=22), _TaskResult(job_id="j-222", df_idx=222)]
         assert remaining == 0
 
     def test_submit_multiple_simple(self, worker_pool):
@@ -204,7 +212,7 @@ def test_submit_multiple_blocking_and_failing(self, worker_pool):
         events[0].set()
         results, remaining = worker_pool.process_futures(timeout=0.1)
         assert results == [
-            _TaskResult(job_id="j-0", df_idx = 0, db_update={"status": "all fine"}),
+            _TaskResult(job_id="j-0", df_idx=0, db_update={"status": "all fine"}),
         ]
         assert remaining == n - 1
 
@@ -213,10 +221,13 @@ def test_submit_multiple_blocking_and_failing(self, worker_pool):
             events[j].set()
         results, remaining = worker_pool.process_futures(timeout=0.1)
         assert results == [
-            _TaskResult(job_id="j-1", df_idx = 1, db_update={"status": "all fine"}),
-            _TaskResult(job_id="j-2", df_idx = 2, db_update={"status": "all fine"}),
+            _TaskResult(job_id="j-1", df_idx=1, db_update={"status": "all fine"}),
+            _TaskResult(job_id="j-2", df_idx=2, db_update={"status": "all fine"}),
             _TaskResult(
-                job_id="j-3", df_idx = 3, db_update={"status": "threaded task failed"}, stats_update={"threaded task failed": 1}
+                job_id="j-3",
+                df_idx=3,
+                db_update={"status": "threaded task failed"},
+                stats_update={"threaded task failed": 1},
             ),
         ]
         assert remaining == 1
@@ -226,7 +237,7 @@ def test_submit_multiple_blocking_and_failing(self, worker_pool):
             events[j].set()
         results, remaining = worker_pool.process_futures(timeout=0.1)
         assert results == [
-            _TaskResult(job_id="j-4", df_idx = 4, db_update={"status": "all fine"}),
+            _TaskResult(job_id="j-4", df_idx=4, db_update={"status": "all fine"}),
         ]
         assert remaining == 0
 
@@ -252,7 +263,7 @@ def test_job_start_task(self, worker_pool, dummy_backend, caplog):
         assert results == [
             _TaskResult(
                 job_id="job-000",
-                df_idx = 0,
+                df_idx=0,
                 db_update={"status": "queued"},
                 stats_update={"job start": 1},
             )
@@ -270,7 +281,9 @@ def test_job_start_task_failure(self, worker_pool, dummy_backend, caplog):
 
         results, remaining = worker_pool.process_futures(timeout=1)
         assert results == [
-            _TaskResult(job_id="job-000", df_idx=0, db_update={"status": "start_failed"}, stats_update={"start_job error": 1})
+            _TaskResult(
+                job_id="job-000", df_idx=0, db_update={"status": "start_failed"}, stats_update={"start_job error": 1}
+            )
         ]
         assert remaining == 0
         assert caplog.messages == [