Add logic to create calculated columns.

Author: fabienDaou

dss-plugin-visual-edit/python-lib/aggrid_utils.py

@@ -1,16 +1,19 @@
-"""
-AG Grid utility functions to generate column configuration for Dash AG Grid.
-"""
-
-import logging
-from typing import Union
+from typing import List, Union
 
 from pandas import DataFrame
 
 from DataEditor import DataEditor
+from webapp.config.models import EditSchema
+
 
+def _get_user_column_def_override(col_name: str, edit_schema: List[EditSchema]) -> dict:
+    for col_def in edit_schema:
+        if col_def.name == col_name:
+            return col_def.aggrid_col_def
+    return {}
 
-def _get_column_aggrid_data_type(dku_col_type: str, dku_col_meaning: str) -> dict:
+
+def _to_aggrid_data_type(dku_col_type: str, dku_col_meaning: str) -> dict:
     if dku_col_type in ["int", "bigint", "double", "float"]:
         return {"cellDataType": "number"}
     elif dku_col_type == "boolean":
@@ -26,7 +29,7 @@ def _get_column_aggrid_data_type(dku_col_type: str, dku_col_meaning: str) -> dic
         return {"cellDataType": "text"}
 
 
-def _get_column_aggrid_cell_editor(dku_col_type: str, dku_col_meaning: str) -> dict:
+def _to_aggrid_cell_editor(dku_col_type: str, dku_col_meaning: str) -> dict:
     if dku_col_type in ["int", "bigint", "double", "float"]:
         return {"cellEditor": "agNumericCellEditor"}
     elif dku_col_type == "boolean":
@@ -37,8 +40,10 @@ def _get_column_aggrid_cell_editor(dku_col_type: str, dku_col_meaning: str) -> d
         return {"cellEditor": "agLargeTextCellEditor", "cellEditorPopup": True}
 
 
-def _get_pk_column_def_aggrid(col_name: str, dku_col_type: str, dku_col_meaning: str, show_header_filter: bool) -> dict:
-    ag_data_type = _get_column_aggrid_data_type(dku_col_type, dku_col_meaning)
+def _get_pk_column_def_aggrid(
+    col_name: str, dku_col_type: str, dku_col_meaning: str, show_header_filter: bool, edit_schema: List[EditSchema]
+) -> dict:
+    ag_data_type = _to_aggrid_data_type(dku_col_type, dku_col_meaning)
     col_def = {
         "field": col_name,
         "headerName": col_name,
@@ -47,13 +52,13 @@ def _get_pk_column_def_aggrid(col_name: str, dku_col_type: str, dku_col_meaning:
         "filter": show_header_filter,
         "filterParams": {"buttons": ["reset", "apply"], "closeOnApply": True},
     }
-    return col_def | ag_data_type
+    return col_def | ag_data_type | _get_user_column_def_override(col_name, edit_schema)
 
 
 def _get_display_column_def_aggrid(
-    col_name: str, dku_col_type: str, dku_col_meaning: str, show_header_filter: bool
+    col_name: str, dku_col_type: str, dku_col_meaning: str, show_header_filter: bool, edit_schema: List[EditSchema]
 ) -> dict:
-    ag_data_type = _get_column_aggrid_data_type(dku_col_type, dku_col_meaning)
+    ag_data_type = _to_aggrid_data_type(dku_col_type, dku_col_meaning)
     col_def = {
         "field": col_name,
         "headerName": col_name,
@@ -62,14 +67,19 @@ def _get_display_column_def_aggrid(
         "filter": show_header_filter,
         "filterParams": {"buttons": ["reset", "apply"], "closeOnApply": True},
     }
-    return col_def | ag_data_type
+    return col_def | ag_data_type | _get_user_column_def_override(col_name, edit_schema)
 
 
 def _get_editable_column_def_aggrid(
-    col_name: str, dku_col_type: str, dku_col_meaning: str, show_header_filter: bool, freeze_editable_columns: bool
+    col_name: str,
+    dku_col_type: str,
+    dku_col_meaning: str,
+    show_header_filter: bool,
+    freeze_editable_columns: bool,
+    edit_schema: List[EditSchema],
 ) -> dict:
-    ag_data_type = _get_column_aggrid_data_type(dku_col_type, dku_col_meaning)
-    ag_cell_editor = _get_column_aggrid_cell_editor(dku_col_type, dku_col_meaning)
+    ag_data_type = _to_aggrid_data_type(dku_col_type, dku_col_meaning)
+    ag_cell_editor = _to_aggrid_cell_editor(dku_col_type, dku_col_meaning)
     # Alternate red shades for editable columns (even: deep red, odd: light red)
     col_def = {
         "field": col_name,
@@ -79,84 +89,125 @@ def _get_editable_column_def_aggrid(
         "filter": show_header_filter,
         "filterParams": {"buttons": ["reset", "apply"], "closeOnApply": True},
     }
-    return col_def | ag_data_type | ag_cell_editor
+    return col_def | ag_data_type | ag_cell_editor | _get_user_column_def_override(col_name, edit_schema)
 
 
 def _get_linked_record_column_def_aggrid(
     col_name: str, dku_col_type: str, dku_col_meaning: str, linked_ds_name: str, show_header_filter: bool
 ) -> dict:
-    ag_data_type = _get_column_aggrid_data_type(dku_col_type, dku_col_meaning)
+    ag_data_type = _to_aggrid_data_type(dku_col_type, dku_col_meaning)
     return {
         "field": col_name,
         "headerName": col_name,
         "editable": True,
         "cellEditor": "agRichSelectCellEditor",
         "cellEditorParams": {
-            # "values": [0, 1, 2, 3, 4, 5],  # Placeholder values; in practice, this should be dynamically fetched
-            # "values": { "function": "dkuLinkedRecordsCellEditor(params)" },
             "function": "dkuLinkedRecordsCellEditor(params)",
-            # "allowTyping": True,
-            # "filterList": True,
-            # "highlightMatch": True
         },
         "linkedDatasetName": linked_ds_name,
         "filter": show_header_filter,
         "resizable": True,
     } | ag_data_type
 
 
-def get_columns_aggrid(de: DataEditor, show_header_filter=True, freeze_editable_columns=False):
-    # Similar logic to Tabulator, but adapted for AG Grid
-    linked_record_names = []
-    if de.linked_records:
-        try:
-            linked_records_df = de.linked_records_df
-            linked_record_names = linked_records_df.index.values.tolist()
-        except Exception:
-            logging.exception("Failed to get linked record names.")
-
-    ag_cols = []
-
+def _get_pk_column_defs(de: DataEditor, edit_schema: List[EditSchema], show_header_filter: bool) -> List[dict]:
     ag_pk_cols = []
-    for pk_col in de.primary_keys:
-        col_type: str = de.schema_columns_df.loc[pk_col, "type"]  # type: ignore
+    for col_name in de.primary_keys:
+        col_type: str = de.schema_columns_df.loc[col_name, "type"]  # type: ignore
         col_meaning: str = (
-            de.schema_columns_df.loc[pk_col, "meaning_id"] if "meaning_id" in de.schema_columns_df.columns else ""
+            de.schema_columns_df.loc[col_name, "meaning_id"] if "meaning_id" in de.schema_columns_df.columns else ""
         )  # type: ignore
-        col_def = _get_pk_column_def_aggrid(pk_col, col_type, col_meaning, show_header_filter)
+        col_def = _get_pk_column_def_aggrid(col_name, col_type, col_meaning, show_header_filter, edit_schema)
         ag_pk_cols.append(col_def)
+    return ag_pk_cols
 
-    ag_cols.append({"headerName": "Row identification", "children": ag_pk_cols, "marryChildren": True})
 
+def _get_display_column_defs(de: DataEditor, edit_schema: List[EditSchema], show_header_filter: bool) -> List[dict]:
     ag_display_cols = []
-    for display_col in de.display_column_names:
-        col_type: str = de.schema_columns_df.loc[display_col, "type"]  # type: ignore
+    for col_name in de.display_column_names:
+        col_type: str = de.schema_columns_df.loc[col_name, "type"]  # type: ignore
         col_meaning: str = (
-            de.schema_columns_df.loc[display_col, "meaning"] if "meaning" in de.schema_columns_df.columns else ""
+            de.schema_columns_df.loc[col_name, "meaning"] if "meaning" in de.schema_columns_df.columns else ""
         )  # type: ignore
-        col_def = _get_display_column_def_aggrid(display_col, col_type, col_meaning, show_header_filter)
+        col_def = _get_display_column_def_aggrid(col_name, col_type, col_meaning, show_header_filter, edit_schema)
         ag_display_cols.append(col_def)
+    return ag_display_cols
 
-    ag_cols.append({"headerName": "Contextual information", "children": ag_display_cols, "marryChildren": True})
 
-    schema_df = de.schema_columns_df
-    for col_name in de.editable_column_names:
-        col_type = schema_df.loc[col_name, "type"]  # type: ignore
+def _get_editable_column_defs(
+    de: DataEditor, edit_schema: List[EditSchema], show_header_filter: bool, freeze_editable_columns: bool
+) -> List[dict]:
+    ag_editable_cols = []
+    linked_record_names = de.linked_records_df.index.values.tolist()
+    editable_column_names = set(de.editable_column_names) - set(linked_record_names)
+    for col_name in editable_column_names:
+        col_type: str = de.schema_columns_df.loc[col_name, "type"]  # type: ignore
         col_meaning: str = (
             de.schema_columns_df.loc[col_name, "meaning"] if "meaning" in de.schema_columns_df.columns else ""
         )  # type: ignore
-        if col_name in linked_record_names:
-            # Linked record column
+        col_def = _get_editable_column_def_aggrid(
+            col_name,
+            col_type,
+            col_meaning,
+            show_header_filter,
+            freeze_editable_columns,
+            edit_schema,
+        )
+        ag_editable_cols.append(col_def)
+    return ag_editable_cols
+
+
+def _get_linked_record_column_defs(
+    de: DataEditor, edit_schema: List[EditSchema], show_header_filter: bool
+) -> List[dict]:
+    ag_linked_record_cols = []
+    if de.linked_records:
+        linked_records_df = de.linked_records_df
+        linked_record_names = linked_records_df.index.values.tolist()
+        for col_name in linked_record_names:
+            col_type: str = de.schema_columns_df.loc[col_name, "type"]  # type: ignore
+            col_meaning: str = (
+                de.schema_columns_df.loc[col_name, "meaning"] if "meaning" in de.schema_columns_df.columns else ""
+            )  # type: ignore
             linked_ds_name: str = linked_records_df.loc[col_name, "ds_name"]  # type: ignore
             col_def = _get_linked_record_column_def_aggrid(
                 col_name, col_type, col_meaning, linked_ds_name, show_header_filter
             )
-            ag_cols.append(col_def)
-            continue
-        col_def = _get_editable_column_def_aggrid(
-            col_name, col_type, col_meaning, show_header_filter, freeze_editable_columns
-        )
-        ag_cols.append(col_def)
+            ag_linked_record_cols.append(col_def)
+    return ag_linked_record_cols
+
+
+def _get_extra_column_defs(de: DataEditor, edit_schema: List[EditSchema]) -> List[dict]:
+    ag_extra_cols = []
+    all_col_names = de.primary_keys + de.display_column_names + de.editable_column_names
+    extra_col_defs = set([col.name for col in edit_schema]) - set(all_col_names)
+    for extra_col_name in extra_col_defs:
+        col_def = next((col.aggrid_col_def for col in edit_schema if col.name == extra_col_name), None)
+        if col_def:
+            ag_extra_cols.append(col_def)
+    return ag_extra_cols
+
+
+def get_columns_aggrid(
+    de: DataEditor, edit_schema: List[EditSchema], show_header_filter=True, freeze_editable_columns=False
+) -> List[dict]:
+    ag_cols = []
+
+    pk_col_defs = _get_pk_column_defs(de, edit_schema, show_header_filter)
+    ag_cols.append({"headerName": "Row identification", "children": pk_col_defs, "marryChildren": True})
+
+    display_col_defs = _get_display_column_defs(de, edit_schema, show_header_filter)
+    ag_cols.append({"headerName": "Contextual information", "children": display_col_defs, "marryChildren": True})
+
+    editable_col_defs = _get_editable_column_defs(de, edit_schema, show_header_filter, freeze_editable_columns)
+    ag_cols.extend(editable_col_defs)
+
+    linked_record_col_defs = _get_linked_record_column_defs(de, edit_schema, show_header_filter)
+    ag_cols.extend(linked_record_col_defs)
+
+    extra_col_defs = _get_extra_column_defs(de, edit_schema)
+    ag_cols.extend(extra_col_defs)
+
     return ag_cols
 
 

dss-plugin-visual-edit/python-lib/webapp/config/loader.py

@@ -1,16 +1,19 @@
 from __future__ import annotations
+
 import json
 import logging
 import os
-from typing import Any, List
-from dataiku.customwebapp import get_webapp_config
 from json import load
 from os import getenv
+from typing import Any, List
+
+from dataiku.customwebapp import get_webapp_config
+
 from webapp.config.models import (
     Config,
     EditSchema,
-    LinkedRecordInfo,
     LinkedRecord,
+    LinkedRecordInfo,
 )
 
 
@@ -24,9 +27,7 @@ def __init__(self) -> None:
             self.__validate_original_dataset_name__(typed_config.original_dataset)
 
         else:
-            original_ds_name = self.__validate_original_dataset_name__(
-                getenv("ORIGINAL_DATASET")
-            )
+            original_ds_name = self.__validate_original_dataset_name__(getenv("ORIGINAL_DATASET"))
             with open(f"../../webapp-settings/{original_ds_name}.json") as fp:
                 dic_config = load(fp)
             typed_config = Config(**dic_config)
@@ -39,9 +40,7 @@ def __init__(self) -> None:
         self.freeze_editable_columns = typed_config.freeze_editable_columns
         self.group_column_names = typed_config.group_column_names
         self.linked_records_count = typed_config.linked_records_count
-        self.linked_records = self.__get_linked_records__(
-            dic_config, self.linked_records_count
-        )
+        self.linked_records = self.__get_linked_records__(dic_config, self.linked_records_count)
 
         self.editschema_manual = self.__cast_editschema__(typed_config)
 
@@ -82,9 +81,7 @@ def __get_project_key__(self) -> str:
             )
         return key
 
-    def __get_linked_records__(
-        self, params, linked_records_count: int
-    ) -> List[LinkedRecord]:
+    def __get_linked_records__(self, params, linked_records_count: int) -> List[LinkedRecord]:
         """
         Get linked records from webapp parameters, as a list of dictionaries
         """
@@ -102,10 +99,6 @@ def __get_linked_records__(
                 if not ds_lookup_columns:
                     ds_lookup_columns = []
                 linked_records.append(
-                    LinkedRecord(
-                        LinkedRecordInfo(
-                            name, ds_name, ds_key, ds_label, ds_lookup_columns
-                        )
-                    )
+                    LinkedRecord(LinkedRecordInfo(name, ds_name, ds_key, ds_label, ds_lookup_columns))
                 )
         return linked_records

dss-plugin-visual-edit/python-lib/webapp/config/models.py

@@ -1,14 +1,17 @@
 from __future__ import annotations
-from pydantic import BaseModel
+
 from dataclasses import dataclass
 from typing import Any, List
-from DatasetSQL import DatasetSQL
+
 from pandas import DataFrame
+from pydantic import BaseModel
+
+from DatasetSQL import DatasetSQL
 
 
 class EditSchema(BaseModel):
     name: str
-    type: str
+    aggrid_col_def: dict
 
 
 @dataclass(frozen=True)

dss-plugin-visual-edit/webapps/visual-edit/backend.py

@@ -81,7 +81,9 @@
 )
 
 
-columns = get_columns_aggrid(de, webapp_config.show_header_filter, webapp_config.freeze_editable_columns)
+columns = get_columns_aggrid(
+    de, webapp_config.editschema_manual, webapp_config.show_header_filter, webapp_config.freeze_editable_columns
+)
 
 last_build_date_initial = ""
 last_build_date_ok = False