Update src/ukg/schedule.py

src/ukg/schedule.py

@@ -1,172 +1,262 @@
+"""Helper functions for interacting with UKG (Kronos) APIs.
+
+This module centralizes the logic for fetching open shifts, locations and
+employee data from the UKG demo API. It also ensures that return values are
+safe to consume within a Streamlit application by handling missing keys
+gracefully and avoiding common exceptions such as ``KeyError`` when
+expected columns are missing from returned data frames.
+"""
+
 import os
-import requests
-import pandas as pd
 import json
-import streamlit as st  # ✅ إضافة هذا السطر
+from typing import List, Optional, Iterable, Dict, Any
+
+import pandas as pd
+import requests
+import streamlit as st
 
-def fetch_open_shifts(start_date="2000-01-01", end_date="3000-01-01", location_ids=None):
-    if location_ids is None:
-        location_ids = ["2401","2402","2953","2955","2927","2928","2401","2955"]
 
-    url = "https://partnerdemo-019.cfn.mykronos.com/api/v1/scheduling/schedule/multi_read"
-    headers = {
+def _get_auth_header() -> Dict[str, str]:
+    """Construct a common authorization header for UKG API calls.
+
+    The UKG API uses two headers for authentication: ``appkey`` and
+    ``Authorization``. The latter expects the string ``Bearer `` followed by
+    the token. Both values are read from environment variables ``UKG_APP_KEY``
+    and ``UKG_AUTH_TOKEN``. If either variable is missing, a warning is
+    emitted via Streamlit.
+
+    Returns
+    -------
+    dict
+        Header dictionary suitable for passing to ``requests`` calls.
+    """
+    app_key = os.environ.get("UKG_APP_KEY")
+    token = os.environ.get("UKG_AUTH_TOKEN")
+    if not app_key or not token:
+        st.warning(
+            "UKG authentication variables (UKG_APP_KEY and/or UKG_AUTH_TOKEN) are"
+            " not set. API calls may fail."
+        )
+    return {
         "Content-Type": "application/json",
-        "appkey": os.environ.get("UKG_APP_KEY"),
-        "Authorization": os.environ.get("UKG_AUTH_TOKEN")
+        "appkey": app_key or "",
+        "Authorization": f"Bearer {token}" if token else "",
     }
 
+
+def fetch_open_shifts(
+    start_date: str = "2000-01-01",
+    end_date: str = "3000-01-01",
+    location_ids: Optional[Iterable[str]] = None,
+) -> pd.DataFrame:
+    """Fetch open shift instances from the UKG demo API.
+
+    Parameters
+    ----------
+    start_date : str
+        ISO date (YYYY-MM-DD) for the beginning of the date range.
+    end_date : str
+        ISO date (YYYY-MM-DD) for the end of the date range.
+    location_ids : iterable of str, optional
+        A collection of location identifiers to filter the search. If not
+        provided, a default list of IDs is used.
+
+    Returns
+    -------
+    pandas.DataFrame
+        A DataFrame containing information about open shifts. If the API
+        returns no shifts or encounters an error, an empty DataFrame is
+        returned.
+    """
+    if location_ids is None:
+        location_ids = ["2401", "2402", "2953", "2955", "2927", "2928", "2401", "2955"]
+
+    url = (
+        "https://partnerdemo-019.cfn.mykronos.com/api/v1/"
+        "scheduling/schedule/multi_read"
+    )
+    headers = _get_auth_header()
+
     payload = {
         "select": ["OPENSHIFTS"],
         "where": {
             "locations": {
                 "dateRange": {
                     "startDate": start_date,
-                    "endDate": end_date
+                    "endDate": end_date,
                 },
                 "includeEmployeeTransfer": False,
-                "locations": {
-                    "ids": location_ids
-                }
+                "locations": {"ids": list(location_ids)},
             }
-        }
+        },
     }
 
     try:
-        print("Payload being sent:")
-        print(json.dumps(payload, indent=2))
-        print("Headers:")
-        print(headers)
-
         response = requests.post(url, headers=headers, json=payload)
-        print("Status Code:", response.status_code)
-        print("Raw Response:", response.text)
-
         response.raise_for_status()
         data = response.json()
-
         open_shifts = data.get("openShifts", [])
-        print(f"Found {len(open_shifts)} open shifts")
-
-        rows = []
+        rows: List[Dict[str, Any]] = []
         for shift in open_shifts:
-            rows.append({
-                "ID": shift.get("id"),
-                "Start": shift.get("startDateTime"),
-                "End": shift.get("endDateTime"),
-                "Label": shift.get("label"),
-                "Org Job": shift.get("segments", [{}])[0].get("orgJobRef", {}).get("qualifier", "") if shift.get("segments") else "",
-                "Posted": shift.get("posted"),
-                "Self Serviced": shift.get("selfServiced"),
-                "Locked": shift.get("locked")
-            })
-
+            rows.append(
+                {
+                    "ID": shift.get("id"),
+                    "Start": shift.get("startDateTime"),
+                    "End": shift.get("endDateTime"),
+                    "Label": shift.get("label"),
+                    "Org Job": shift.get("segments", [{}])[0]
+                    .get("orgJobRef", {})
+                    .get("qualifier", "")
+                    if shift.get("segments")
+                    else "",
+                    "Posted": shift.get("posted"),
+                    "Self Serviced": shift.get("selfServiced"),
+                    "Locked": shift.get("locked"),
+                }
+            )
         return pd.DataFrame(rows)
-
     except Exception as e:
-        print(f"❌ UKG Open Shifts API call failed: {e}")
+        st.error(f"❌ UKG Open Shifts API call failed: {e}")
         return pd.DataFrame()
-        
 
-# ---- UKG Location API FETCH FUNCTION ----
-def fetch_location_data():
-    url = "https://partnerdemo-019.cfn.mykronos.com/api/v1/commons/locations/multi_read"
-    headers = {
-        "Content-Type": "application/json",
-        "appkey": os.environ.get("UKG_APP_KEY"),
-        "Authorization": os.environ.get("UKG_AUTH_TOKEN")
-    }
 
+def fetch_location_data(date: str = "2025-07-13", query: str = "Medsurg") -> pd.DataFrame:
+    """Fetch location information from the UKG demo API.
+
+    Parameters
+    ----------
+    date : str
+        The effective date for the location context in ISO format.
+    query : str
+        A search string to filter locations.
+
+    Returns
+    -------
+    pandas.DataFrame
+        A DataFrame containing location attributes, or an empty DataFrame
+        if the call fails.
+    """
+    url = (
+        "https://partnerdemo-019.cfn.mykronos.com/api/v1/"
+        "commons/locations/multi_read"
+    )
+    headers = _get_auth_header()
     payload = {
-        "multiReadOptions": {
-            "includeOrgPathDetails": True
-        },
+        "multiReadOptions": {"includeOrgPathDetails": True},
         "where": {
-            "query": {
-                "context": "ORG",
-                "date": "2025-07-13",
-                "q": "Medsurg"
-            }
-        }
+            "query": {"context": "ORG", "date": date, "q": query}
+        },
     }
-
     try:
         response = requests.post(url, headers=headers, json=payload)
         response.raise_for_status()
         data = response.json()
-
-        # استخراج المواقع
+        # The API returns a list of location objects
         rows = []
-        for item in data:
-            rows.append({
-                "Node ID": item.get("nodeId", ""),
-                "Name": item.get("name", ""),
-                "Full Name": item.get("fullName", ""),
-                "Description": item.get("description", ""),
-                "Org Path": item.get("orgPath", ""),
-                "Persistent ID": item.get("persistentId", "")
-            })
-
+        for item in data if isinstance(data, list) else data.get("locations", []):
+            rows.append(
+                {
+                    "Node ID": item.get("nodeId", ""),
+                    "Name": item.get("name", ""),
+                    "Full Name": item.get("fullName", ""),
+                    "Description": item.get("description", ""),
+                    "Org Path": item.get("orgPath", ""),
+                    "Persistent ID": item.get("persistentId", ""),
+                }
+            )
         return pd.DataFrame(rows)
-
     except Exception as e:
         st.error(f"❌ UKG Location API call failed: {e}")
         return pd.DataFrame()
 
 
-import os
-import requests
-import pandas as pd
+def fetch_employees(employee_ids: Iterable[int]) -> pd.DataFrame:
+    """Fetch employee information from the UKG demo API.
 
-# ---- Fetch employee data from UKG API ----
-def fetch_employees(employee_ids):
-    UKG_BASE_URL = "https://partnerdemo-019.cfn.mykronos.com/api/v1/commons/persons/"
-    headers = {
-        "Content-Type": "application/json",
-        "appkey": os.environ.get("UKG_APP_KEY"),
-        "Authorization": f"Bearer {os.environ.get('UKG_AUTH_TOKEN')}"
-    }
+    For each employee ID provided, this function queries the person information
+    endpoint and extracts basic details such as the person number, full name,
+    phone number and organizational path. It also derives a ``JobRole`` from
+    the last segment of the ``organizationPath``. When no records are returned
+    or a particular key is missing, the returned DataFrame still contains
+    the expected columns to avoid ``KeyError`` exceptions downstream.
 
-    def fetch_employee_data(emp_id):
+    Parameters
+    ----------
+    employee_ids : iterable of int
+        A list of employee identifiers.
+
+    Returns
+    -------
+    pandas.DataFrame
+        A DataFrame with one row per successfully fetched employee. If no
+        data is available, a DataFrame with the appropriate columns but no
+        rows is returned.
+    """
+    base_url = (
+        "https://partnerdemo-019.cfn.mykronos.com/api/v1/commons/persons/"
+    )
+    headers = _get_auth_header()
+
+    def fetch_employee_data(emp_id: int) -> Optional[Dict[str, Any]]:
         try:
-            response = requests.get(f"{UKG_BASE_URL}{emp_id}", headers=headers)
-            if response.status_code == 200:
-                data = response.json()
-                person_number = data.get("personInformation", {}).get("person", {}).get("personNumber", None)
-                org_path = ""
-                phone = ""
-                full_name = ""
+            resp = requests.get(f"{base_url}{emp_id}", headers=headers)
+            if resp.status_code == 200:
+                data = resp.json()
+                person_info = data.get("personInformation", {}).get("person", {})
+                person_number = person_info.get("personNumber")
+                full_name = person_info.get("fullName", "")
 
-                primary_accounts = data.get("jobAssignment", {}).get("primaryLaborAccounts", [])
+                # Extract organization path from primary labor accounts
+                org_path = ""
+                primary_accounts = (
+                    data.get("jobAssignment", {}).get("primaryLaborAccounts", [])
+                )
                 if primary_accounts:
                     org_path = primary_accounts[0].get("organizationPath", "")
 
+                # Extract first phone number if present
+                phone = ""
                 phones = data.get("personInformation", {}).get("telephoneNumbers", [])
                 if phones:
                     phone = phones[0].get("phoneNumber", "")
 
-                full_name = data.get("personInformation", {}).get("person", {}).get("fullName", "")
-
                 return {
                     "personNumber": person_number,
                     "organizationPath": org_path,
                     "phoneNumber": phone,
-                    "fullName": full_name
+                    "fullName": full_name,
                 }
             else:
-                print(f"⚠️ Could not fetch ID {emp_id}: {response.status_code}")
-                return None
+                st.warning(
+                    f"⚠️ Could not fetch employee {emp_id}: {resp.status_code}"
+                )
         except Exception as e:
-            print(f"❌ Error for ID {emp_id}: {e}")
-            return None
+            st.error(f"❌ Error fetching employee {emp_id}: {e}")
+        return None
 
-    employee_records = [emp for emp_id in employee_ids if (emp := fetch_employee_data(emp_id))]
-    df_employees = pd.DataFrame(employee_records)
+    records: List[Dict[str, Any]] = []
+    for emp_id in employee_ids:
+        data = fetch_employee_data(emp_id)
+        if data:
+            records.append(data)
 
-    # Extract job role from organizationPath
-    df_employees["JobRole"] = df_employees["organizationPath"].apply(
-        lambda x: x.split("/")[-1] if isinstance(x, str) else ""
-    )
-
-    return df_employees
+    # Build DataFrame with guaranteed columns
+    df_employees = pd.DataFrame(records)
+    # Ensure mandatory columns exist even if DataFrame is empty
+    for col in [
+        "personNumber",
+        "organizationPath",
+        "phoneNumber",
+        "fullName",
+    ]:
+        if col not in df_employees.columns:
+            df_employees[col] = []
 
+    # Derive JobRole from organizationPath safely
+    def derive_role(path: Any) -> str:
+        if isinstance(path, str) and path:
+            return path.split("/")[-1]
+        return ""
 
+    df_employees["JobRole"] = df_employees["organizationPath"].apply(derive_role)
+    return df_employees