centralizing

Author: HansVRP

openeo/extra/job_management/__init__.py

@@ -30,9 +30,7 @@
     _JobStartTask,
 )
 from openeo.extra.job_management.process_based_job_creator import ProcessBasedJobCreator
-from openeo.extra.job_management._job_database import FullDataFrameJobDatabase, JobDatabaseInterface, ParquetJobDatabase, CsvJobDatabase, create_job_db, get_job_db
-from openeo.extra.job_management._dataframe_utils import normalize_dataframe
-
+from openeo.extra.job_management._job_database import FullDataFrameJobDatabase, JobDatabaseInterface, ParquetJobDatabase, CsvJobDatabase, create_job_db, get_job_db, normalize_dataframe
 
 from openeo.rest import OpenEoApiError
 from openeo.rest.auth.auth import BearerAuth

openeo/extra/job_management/_dataframe_utils.py

@@ -12,12 +12,19 @@
 import shapely.errors
 import shapely.wkt
 
-from openeo.extra.job_management._dataframe_utils import normalize_dataframe, COLUMN_REQUIREMENTS
-
 
 
 _log = logging.getLogger(__name__)
 
+import pandas as pd
+
+class _ColumnProperties:
+    def __init__(self, dtype: str, default=None):
+        self.dtype = dtype
+        self.default = default
+
+
+
 class JobDatabaseInterface(metaclass=abc.ABCMeta):
     """
     Interface for a database of job metadata to use with the :py:class:`MultiBackendJobManager`,
@@ -78,7 +85,20 @@ def get_by_indices(self, indices: Iterable[Union[int, str]]) -> pd.DataFrame:
         """
         ...
 
-        
+# Expected columns in the job DB dataframes.
+# TODO: make this part of public API when settled?
+# TODO: move non official statuses to seperate column (not_started, queued_for_start)
+COLUMN_REQUIREMENTS = {
+    "id": _ColumnProperties(dtype="str"),
+    "backend_name": _ColumnProperties(dtype="str"),
+    "status": _ColumnProperties(dtype="str", default="not_started"),
+    "start_time": _ColumnProperties(dtype="str"),
+    "running_start_time": _ColumnProperties(dtype="str"),
+    "cpu": _ColumnProperties(dtype="str"),
+    "memory": _ColumnProperties(dtype="str"),
+    "duration": _ColumnProperties(dtype="str"),
+    "costs": _ColumnProperties(dtype="float64"),
+}   
 
 class FullDataFrameJobDatabase(JobDatabaseInterface):
     def __init__(self):
@@ -271,6 +291,18 @@ def persist(self, df: pd.DataFrame):
         self.path.parent.mkdir(parents=True, exist_ok=True)
         self.df.to_parquet(self.path, index=False)    
 
+def normalize_dataframe(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    Normalize given pandas dataframe (creating a new one):
+    ensure we have the required columns.
+
+    :param df: The dataframe to normalize.
+    :return: a new dataframe that is normalized.
+    """
+    new_columns = {col: req.default for (col, req) in COLUMN_REQUIREMENTS.items() if col not in df.columns}
+    df = df.assign(**new_columns)
+    return df 
+
 def create_job_db(path: Union[str, Path], df: pd.DataFrame, *, on_exists: str = "error"):
     """
     Factory to create a job database at given path,