[dy] Show the load balancer ingress host if there is no host in the ingress rules (mage-ai#4575)

* [dy] Don't return url when host is None

* [dy] Check the status field

* [dy] Add unit test

Author: dy46

mage_ai/cluster_manager/kubernetes/workload_manager.py

@@ -43,6 +43,7 @@
 )
 from mage_ai.settings import MAGE_SETTINGS_ENVIRONMENT_VARIABLES
 from mage_ai.shared.array import find
+from mage_ai.shared.hash import safe_dig
 
 
 class WorkloadManager:
@@ -183,7 +184,6 @@ def format_workloads(
                             pass
                         if ip:
                             workload['ip'] = f'{ip}:{port}'
-
                 elif stateful_set and stateful_set.spec.replicas == 0:
                     workload['status'] = 'STOPPED'
                 else:
@@ -515,6 +515,18 @@ def get_url_from_ingress(self, ingress_name: str, workload_name: str) -> str:
         rule = ingress.spec.rules[0]
         host = rule.host
 
+        if host is None:
+            ingress_dict = ingress.to_dict()
+            lb_ingress = safe_dig(ingress_dict, ['status', 'load_balancer', 'ingress[0]']) or {}
+            if 'hostname' in lb_ingress:
+                host = lb_ingress['hostname']
+            # Alternatively, check for `ip` if `hostname` is not available
+            elif 'ip' in lb_ingress:
+                host = lb_ingress['ip']
+
+        if host is None:
+            return None
+
         tls_enabled = False
         try:
             tls = ingress.spec.tls[0]

mage_ai/shared/hash.py

@@ -38,6 +38,53 @@ def _build(obj, key):
             return obj.get(key)
         else:
             return obj
+
+    return reduce(_build, arr, obj_arg)
+
+
+def safe_dig(obj_arg, arr_or_string):
+    """
+    Safely retrieves nested values from a dictionary or list using a dot-separated path.
+
+    Args:
+        obj_arg: The object (dictionary or list) to navigate.
+        arr_or_string (str or list): A dot-separated path string or a list of
+            keys/indexes.
+
+    Returns:
+        The value retrieved from the nested structure, or None if any intermediate
+            key/index is missing or the object is None.
+    """
+    if isinstance(arr_or_string, str):
+        arr_or_string = arr_or_string.split('.')
+    arr = list(map(str.strip, arr_or_string))
+
+    def _build(obj, key):
+        # Return None if the object is None or not a dictionary
+        if obj is None or not isinstance(obj, dict) and not isinstance(obj, list):
+            return None
+
+        tup = re.split(r'\[(\d+)\]$', key)
+        if len(tup) >= 2:
+            key, index = filter(lambda x: x, tup)
+            index = int(index) if index else None
+            if key and index is not None:
+                if key not in obj:
+                    return None  # Return None if the key is not present
+                return (
+                    obj[key][index]
+                    if isinstance(obj[key], list) and len(obj[key]) > index
+                    else None
+                )
+            elif index is not None:
+                return (
+                    obj[index] if isinstance(obj, list) and len(obj) > index else None
+                )
+        elif isinstance(obj, dict):
+            return obj.get(key)
+        else:
+            return None
+
     return reduce(_build, arr, obj_arg)
 
 
@@ -93,6 +140,7 @@ def _build(obj, key):
         if include_blank_values or val is not None:
             obj[key] = val
         return obj
+
     return reduce(_build, keys, {})
 
 
@@ -111,6 +159,7 @@ def _build(obj, item):
             obj[val] = []
         obj[val].append(item)
         return obj
+
     return reduce(_build, arr, {})
 
 
@@ -141,6 +190,7 @@ def _replace_nan_value(v):
         if isinstance(v, float) and math.isnan(v):
             return None
         return v
+
     return {k: _replace_nan_value(v) for k, v in d.items()}
 
 

mage_ai/tests/shared/test_hash.py

@@ -2,6 +2,7 @@
     camel_case_keys_to_snake_case,
     combine_into,
     get_json_value,
+    safe_dig,
     set_value,
 )
 from mage_ai.tests.base_test import TestCase
@@ -107,3 +108,48 @@ def test_combine_into(self):
             ),
             ice=3,
         ))
+
+    def test_safe_dig_dict_navigation(self):
+        data = {
+            'foo': {
+                'bar': {
+                    'baz': 42
+                }
+            }
+        }
+        self.assertEqual(safe_dig(data, 'foo.bar.baz'), 42)
+
+    def test_safe_dig_list_navigation(self):
+        data = {
+            'foo': [
+                {'bar': 42},
+                {'bar': 43}
+            ]
+        }
+        self.assertEqual(safe_dig(data, 'foo[0].bar'), 42)
+        self.assertEqual(safe_dig(data, 'foo[1].bar'), 43)
+        self.assertEqual(safe_dig(data, ['foo[0]', 'bar']), 42)
+        self.assertEqual(safe_dig(data, ['foo[1]', 'bar']), 43)
+
+    def test_safe_dig_missing_key(self):
+        data = {
+            'foo': {
+                'bar': {
+                    'baz': 42
+                }
+            }
+        }
+        self.assertIsNone(safe_dig(data, 'foo.bar.baz.qux'))
+        self.assertIsNone(safe_dig(data, ['foo', 'bar', 'baz', 'qux']))
+
+    def test_safe_dig_none_value(self):
+        data = None
+        self.assertIsNone(safe_dig(data, 'foo.bar.baz'))
+        self.assertIsNone(safe_dig(data, ['foo', 'bar', 'baz']))
+
+    def test_safe_dig_invalid_index(self):
+        data = {
+            'foo': [1, 2, 3]
+        }
+        self.assertIsNone(safe_dig(data, 'foo[3]'))
+        self.assertIsNone(safe_dig(data, ['foo[3]']))