From c5bdf4c64a6d469d6b81d88ed1805d9de337bd12 Mon Sep 17 00:00:00 2001
From: Sampath Ganta <neigant1274@gmail.com>
Date: Thu, 22 Feb 2024 15:32:30 +0530
Subject: [PATCH] Added Master Arbitration support for Get and Set

Created Extension for Master Arbitration
---
 pygnmi/client.py                | 19 ++++++++++++++-----
 pygnmi/create_gnmi_extension.py |  8 ++++++++
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/pygnmi/client.py b/pygnmi/client.py
index 66b2a89..25db545 100644
--- a/pygnmi/client.py
+++ b/pygnmi/client.py
@@ -388,7 +388,7 @@ def convert_encoding(self, requested_encoding: str, is_encoding_explicitly_set:
         encoding = requested_encoding or self.__encoding
         return Encoding.Value(encoding.upper())  # may raise ValueError
 
-    def get(self, prefix: str = "", path: list = None, target: str = None, datatype: str = "all", encoding: str = None):
+    def get(self, prefix: str = "", path: list = None, target: str = None, datatype: str = "all", encoding: str = None, extension: dict = None):
         """
         Collecting the information about the resources from defined paths.
 
@@ -426,6 +426,7 @@ def get(self, prefix: str = "", path: list = None, target: str = None, datatype:
 
         # Set Protobuf value for encoding
         pb_encoding = self.convert_encoding(encoding)
+        gnmi_extension = get_gnmi_extension(ext=extension)
 
         # Gnmi PREFIX
         try:
@@ -448,7 +449,10 @@ def get(self, prefix: str = "", path: list = None, target: str = None, datatype:
             raise gNMIException("Conversion of gNMI paths to the Protobuf format failed", e)
 
         try:
-            gnmi_message_request = GetRequest(prefix=protobuf_prefix, path=protobuf_paths, type=pb_datatype, encoding=pb_encoding)
+            if gnmi_extension:
+                gnmi_message_request = GetRequest(prefix=protobuf_prefix, path=protobuf_paths, type=pb_datatype, encoding=pb_encoding, extension = [gnmi_extension])
+            else:
+                gnmi_message_request = GetRequest(prefix=protobuf_prefix, path=protobuf_paths, type=pb_datatype, encoding=pb_encoding)
             debug_gnmi_msg(self.__debug, gnmi_message_request, "gNMI request")
 
             gnmi_message_response = self.__stub.Get(gnmi_message_request, metadata=self.__metadata)
@@ -553,6 +557,7 @@ def set(
         encoding: str = None,
         prefix: str = "",
         target: str = None,
+        extension: dict = None,
     ):
         """
         Changing the configuration on the destination network elements.
@@ -583,6 +588,7 @@ def set(
 
         # Set the encoding to auto-discovered value, unless overridden
         encoding = encoding or self.__encoding
+        gnmi_extension = get_gnmi_extension(ext=extension)
 
         # Gnmi PREFIX
         try:
@@ -630,9 +636,12 @@ def set(
                                            encoding=encoding,
                                            datatype='config')
 
-            gnmi_message_request = SetRequest(
-                prefix=protobuf_prefix, delete=del_protobuf_paths, update=update_msg, replace=replace_msg
-            )
+            if gnmi_extension:
+                gnmi_message_request = SetRequest(prefix=protobuf_prefix, delete=del_protobuf_paths, update=update_msg, replace=replace_msg, extension=[gnmi_extension])
+            else:
+                gnmi_message_request = SetRequest(
+                    prefix=protobuf_prefix, delete=del_protobuf_paths, update=update_msg, replace=replace_msg
+                )
             debug_gnmi_msg(self.__debug, gnmi_message_request, "gNMI request")
 
             gnmi_message_response = self.__stub.Set(gnmi_message_request, metadata=self.__metadata)
diff --git a/pygnmi/create_gnmi_extension.py b/pygnmi/create_gnmi_extension.py
index 1d7640d..09ed72c 100644
--- a/pygnmi/create_gnmi_extension.py
+++ b/pygnmi/create_gnmi_extension.py
@@ -36,6 +36,14 @@ def get_gnmi_extension(ext: dict = None) -> list:
                 result.history.range.start = _get_time_ns_epoch(ext["history"]["range"]["start"])
                 result.history.range.end = _get_time_ns_epoch(ext["history"]["range"]["end"])
 
+    if "master_arbitration" in ext:
+
+        result.master_arbitration.election_id.high = ext["master_arbitration"]["election_id"]["high"]
+        result.master_arbitration.election_id.low = ext["master_arbitration"]["election_id"]["low"]
+
+        if "role" in ext["master_arbitration"]:
+            result.master_arbitration.role.id = ext["master_arbitration"]["role"]["id"]
+
     return result