GG-30304 Fix partition awareness.

Author: SammyVimes

pygridgain/api/affinity.py

@@ -16,7 +16,7 @@
 from typing import Iterable, Union
 
 from pygridgain.datatypes import Bool, Int, Long, UUIDObject
-from pygridgain.datatypes.internal import StructArray
+from pygridgain.datatypes.internal import StructArray, Conditional, Struct
 from pygridgain.queries import Query
 from pygridgain.queries.op_codes import OP_CACHE_PARTITIONS
 from pygridgain.utils import is_iterable
@@ -46,10 +46,22 @@
     ('cache_config', cache_config),
 ])
 
+empty_cache_mapping = StructArray([
+    ('cache_id', Int)
+])
+
+empty_node_mapping = Struct([])
+
 partition_mapping = StructArray([
     ('is_applicable', Bool),
-    ('cache_mapping', cache_mapping),
-    ('node_mapping', node_mapping),
+
+    ('cache_mapping', Conditional(lambda ctx: ctx['is_applicable'] == b'\x01',
+                                  lambda ctx: ctx['is_applicable'] is True,
+                                  cache_mapping, empty_cache_mapping)),
+
+    ('node_mapping', Conditional(lambda ctx: ctx['is_applicable'] == b'\x01',
+                                 lambda ctx: ctx['is_applicable'] is True,
+                                 node_mapping, empty_node_mapping)),
 ])
 
 

pygridgain/cache.py

@@ -282,7 +282,8 @@ def get_best_node(
                 # calculate the number of partitions
                 parts = sum(
                     [len(p) for _, p in self.affinity['node_mapping'].items()]
-                )
+                ) if 'node_mapping' in self.affinity else 0
+
                 self.affinity['number_of_partitions'] = parts
             else:
                 # get number of partitions

pygridgain/datatypes/internal.py

@@ -17,7 +17,7 @@
 import ctypes
 import decimal
 from datetime import date, datetime, timedelta
-from typing import Any, Tuple, Union
+from typing import Any, Tuple, Union, Callable
 import uuid
 
 import attr
@@ -111,6 +111,20 @@ def tc_map(key: bytes, _memo_map: dict = {}):
     return _memo_map[key]
 
 
+class Conditional:
+
+    def __init__(self, predicate1: Callable[[any], bool], predicate2: Callable[[any], bool], var1, var2):
+        self.predicate1 = predicate1
+        self.predicate2 = predicate2
+        self.var1 = var1
+        self.var2 = var2
+
+    def parse(self, client: 'Client', context):
+        return self.var1.parse(client) if self.predicate1(context) else self.var2.parse(client)
+
+    def to_python(self, ctype_object, context, *args, **kwargs):
+        return self.var1.to_python(ctype_object, *args, **kwargs) if self.predicate2(context) else self.var2.to_python(ctype_object, *args, **kwargs)
+
 @attr.s
 class StructArray:
     """ `counter_type` counter, followed by count*following structure. """
@@ -193,13 +207,17 @@ def parse(
     ) -> Tuple[ctypes.LittleEndianStructure, bytes]:
         buffer = b''
         fields = []
+        values = {}
 
         for name, c_type in self.fields:
-            c_type, buffer_fragment = c_type.parse(client)
+            is_cond = isinstance(c_type, Conditional)
+            c_type, buffer_fragment = c_type.parse(client, values) if is_cond else c_type.parse(client)
             buffer += buffer_fragment
 
             fields.append((name, c_type))
 
+            values[name] = buffer_fragment
+
         data_class = type(
             'StructLE',
             (ctypes.LittleEndianStructure,),
@@ -216,7 +234,12 @@ def to_python(
     ) -> Union[dict, OrderedDict]:
         result = self.dict_type()
         for name, c_type in self.fields:
+            is_cond = isinstance(c_type, Conditional)
             result[name] = c_type.to_python(
+                getattr(ctype_object, name),
+                result,
+                *args, **kwargs
+            ) if is_cond else c_type.to_python(
                 getattr(ctype_object, name),
                 *args, **kwargs
             )