Merge branch '1847-DAPS-refactor-python-client-to-use-proto3' into 1848-DAPS-refactor-authz-library

JoshuaSBrown · JoshuaSBrown · commit c2aa1bad92a6 · 2026-02-07T00:00:18.000-05:00
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -199,7 +199,7 @@ endif()
 
 if( BUILD_PYTHON_CLIENT )
   # make target = pydatafed
-  file(COPY ${PROJECT_SOURCE_DIR}/external/DataFedDependencies/python/datafed_pkg/requirements.txt DESTINATION ${PROJECT_SOURCE_DIR}/python/datafed_pkg/requirements.txt)
+  file(COPY ${PROJECT_SOURCE_DIR}/external/DataFedDependencies/python/datafed_pkg/requirements.txt DESTINATION ${PROJECT_SOURCE_DIR}/python/datafed_pkg)
   add_subdirectory( python EXCLUDE_FROM_ALL )
 endif()
 
diff --git a/python/datafed_pkg/CMakeLists.txt b/python/datafed_pkg/CMakeLists.txt
@@ -35,4 +35,4 @@ endforeach()
 add_subdirectory( datafed )
 
 add_custom_target( pydatafed )
-add_dependencies( pydatafed pydatafed_src)
+add_dependencies( pydatafed pydatafed_proto_src)
diff --git a/python/datafed_pkg/datafed/CMakeLists.txt b/python/datafed_pkg/datafed/CMakeLists.txt
@@ -35,30 +35,13 @@ set(PROTO_SUBDIRS anon auth enums messages)
 #   from .anon import ack_reply_pb2 as ...
 #   from .enums import error_code_pb2 as ...
 #   from . import envelope_pb2 as ...
+
+# Create the import fixup script
 add_custom_command(TARGET pydatafed_src POST_BUILD
-    COMMAND find ${CMAKE_CURRENT_BINARY_DIR} -name "*_pb2.py" -exec
-        sed -i -r
-        -e "s:^from anon import:from .anon import:g"
-        -e "s:^from anon\\.:from .anon.:g"
-        -e "s:^from auth import:from .auth import:g"
-        -e "s:^from auth\\.:from .auth.:g"
-        -e "s:^from enums import:from .enums import:g"
-        -e "s:^from enums\\.:from .enums.:g"
-        -e "s:^from messages import:from .messages import:g"
-        -e "s:^from messages\\.:from .messages.:g"
-        -e "s:^import ([a-zA-Z0-9_]*_pb2):from . import \\1:g"
-        {} +
+	COMMAND sh ${DataFed_SOURCE_DIR}/python/datafed_pkg/scripts/fix_proto_imports.sh  ${CMAKE_CURRENT_BINARY_DIR}
     COMMENT "Rewriting protobuf imports to relative"
 )
 
-# Ensure __init__.py exists in each generated subdirectory
-foreach(subdir ${PROTO_SUBDIRS})
-    add_custom_command(TARGET pydatafed_src POST_BUILD
-        COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_CURRENT_BINARY_DIR}/${subdir}
-        COMMAND ${CMAKE_COMMAND} -E touch ${CMAKE_CURRENT_BINARY_DIR}/${subdir}/__init__.py
-    )
-endforeach()
-
 # Copy generated files back to source tree for testing
 add_custom_target(pydatafed_proto_src DEPENDS pydatafed_src)
 add_custom_command(TARGET pydatafed_proto_src POST_BUILD
diff --git a/python/datafed_pkg/datafed/CommandLib.py b/python/datafed_pkg/datafed/CommandLib.py
@@ -14,10 +14,9 @@
 import time
 import pathlib
 import requests
-from . import envelope_pb2 as sdms
 from . import MessageLib
 from . import Config
-
+from . import envelope_pb2 as sdms
 
 class API:
     """
diff --git a/python/datafed_pkg/datafed/Connection.py b/python/datafed_pkg/datafed/Connection.py
@@ -67,6 +67,10 @@ def __init__(
         self._msg_desc_by_type = {}
         self._msg_desc_by_name = {}
         self._msg_type_by_desc = {}
+        self._field_by_msg_desc = {}
+
+        self._envelope_class = None
+        self._envelope_desc = None
 
         self._address = "tcp://{0}:{1}".format(server_host, server_port)
         # init zeromq
@@ -130,6 +134,10 @@ def registerEnvelope(self, envelope_module, envelope_class_name="Envelope"):
         envelope_class = getattr(envelope_module, envelope_class_name)
         envelope_desc = envelope_class.DESCRIPTOR
 
+        # Store for envelope wrapping/unwrapping
+        self._envelope_class = envelope_class
+        self._envelope_desc = envelope_desc
+
         for field in envelope_desc.fields:
             if field.message_type is None:
                 # Skip non-message fields (e.g. scalars) if any exist
@@ -141,6 +149,7 @@ def registerEnvelope(self, envelope_module, envelope_class_name="Envelope"):
             self._msg_desc_by_type[msg_type] = desc
             self._msg_desc_by_name[desc.name] = desc
             self._msg_type_by_desc[desc] = msg_type
+            self._field_by_msg_desc[desc] = field
 
         self._logger.debug(
             "Registered %d message types from %s",
@@ -177,15 +186,15 @@ def registerProtocol(self, msg_module):
     ##
     # @brief Receive a message
     #
-    # Receive a protobuf message with timeout. This method automatically
-    # parses and creates a new protobuf message class based on received
-    # framing. The new message object, the message name (defined in the
-    # associated proto file), and re-association context are returned as
-    # a tuple. On timeout, (None,None,None) is returned.
+    # Receive a protobuf message with timeout. The wire payload is an
+    # Envelope message; this method deserializes the Envelope and extracts
+    # the inner message via the oneof payload field. The inner message
+    # object, its name, and re-association context are returned as a tuple.
+    # On timeout, (None, None, None) is returned.
     #
     # @param timeout - Timeout in milliseconds
-    # @return Tuple of message, message type, and re-association context
-    # @retval (object,str,int) or (None,None,None) on timeout
+    # @return Tuple of message, message name, and re-association context
+    # @retval (object, str, int) or (None, None, None) on timeout
     # @exception Exception: if unregistered message type is received.
     #
     def recv(self, a_timeout=1000):
@@ -219,48 +228,62 @@ def recv(self, a_timeout=1000):
             # client
             self._socket.recv_string(0)
 
-            # receive custom frame header and unpack
+            # Receive frame: 8 bytes = uint32 size + uint16 msg_type + uint16 context
             frame_data = self._socket.recv(0)
-            frame_values = struct.unpack(">LBBH", frame_data)
-            msg_type = (frame_values[1] << 8) | frame_values[2]
+            frame_values = struct.unpack(">LHH", frame_data)
+            body_size = frame_values[0]
+            msg_type = frame_values[1]
+            ctxt = frame_values[2]
 
-            # find message descriptor based on type (descriptor index)
-
-            if not (msg_type in self._msg_desc_by_type):
+            if msg_type not in self._msg_desc_by_type:
                 raise Exception(
                     "received unregistered message type: {}".format(msg_type)
                 )
 
-            desc = self._msg_desc_by_type[msg_type]
+            data = self._socket.recv(0)
 
-            if frame_values[0] > 0:
-                # Create message by parsing content
-                data = self._socket.recv(0)
-                reply = GetMessageClass(desc)()
-                reply.ParseFromString(data)
+            if body_size > 0:
+                # Deserialize as Envelope
+                envelope = self._envelope_class()
+                envelope.ParseFromString(data)
+
+                # Extract inner message from the oneof
+                payload_field = envelope.WhichOneof("payload")
+                if payload_field is None:
+                    raise Exception("Received Envelope with no payload set")
+                reply = getattr(envelope, payload_field)
             else:
-                # No content, just create message instance
-                data = self._socket.recv(0)
+                # Zero-size body: create empty message instance from type
+                desc = self._msg_desc_by_type[msg_type]
                 reply = GetMessageClass(desc)()
 
-            return reply, desc.name, frame_values[3]
+            return reply, reply.DESCRIPTOR.name, ctxt
         else:
             return None, None, None
 
     ##
     # @brief Send a message
     #
-    # Serializes and sends framing and message payload over connection.
+    # Wraps the inner message in an Envelope, serializes it, and sends
+    # framing and payload over the connection. The frame header carries the
+    # message type (Envelope field number) for efficient routing on the
+    # server side.
     #
     # @param message - The protobuf message object to be sent
     # @param ctxt - Reply re-association value (int)
     # @exception Exception: if unregistered message type is sent.
     #
     def send(self, message, ctxt):
         # Find msg type by descriptor look-up
-        if not (message.DESCRIPTOR in self._msg_type_by_desc):
+        if message.DESCRIPTOR not in self._msg_type_by_desc:
             raise Exception("Attempt to send unregistered message type.")
+
         msg_type = self._msg_type_by_desc[message.DESCRIPTOR]
+        field = self._field_by_msg_desc[message.DESCRIPTOR]
+
+        # Wrap inner message in Envelope
+        envelope = self._envelope_class()
+        getattr(envelope, field.name).CopyFrom(message)
 
         # Initial Null frame
         self._socket.send_string("BEGIN_DATAFED", zmq.SNDMORE)
@@ -274,12 +297,12 @@ def send(self, message, ctxt):
         self._socket.send_string(self._pub_key, zmq.SNDMORE)
         self._socket.send_string("no_user", zmq.SNDMORE)
 
-        # Serialize
-        data = message.SerializeToString()
+        # Serialize the Envelope (not the inner message)
+        data = envelope.SerializeToString()
         data_sz = len(data)
 
-        # Build the message frame, to match C-struct MessageFrame
-        frame = struct.pack(">LBBH", data_sz, msg_type >> 8, msg_type & 0xFF, ctxt)
+        # Build the message frame: uint32 size + uint16 msg_type + uint16 context
+        frame = struct.pack(">LHH", data_sz, msg_type, ctxt)
 
         if data_sz > 0:
             # Send frame and payload
diff --git a/python/datafed_pkg/scripts/fix_proto_imports.sh b/python/datafed_pkg/scripts/fix_proto_imports.sh
@@ -0,0 +1,65 @@
+#!/bin/sh
+set -e
+
+PROTO_DIR="$1"
+ROOT_DIR="${2:-$1}"
+
+if [ -z "$PROTO_DIR" ]; then
+    echo "Usage: fix_proto_imports.sh <proto_output_dir> [root_dir]"
+    echo "  proto_output_dir: directory to find and fix _pb2.py files"
+    echo "  root_dir: package root for computing relative depth (defaults to proto_output_dir)"
+    exit 1
+fi
+
+find "$PROTO_DIR" -name '*_pb2.py' | while read f; do
+    relpath=$(realpath --relative-to="$ROOT_DIR" "$f")
+    case "$relpath" in
+        */*)
+            sed -i \
+                -e 's:^from anon import:from ..anon import:g' \
+                -e 's:^from anon\.:from ..anon.:g' \
+                -e 's:^from auth import:from ..auth import:g' \
+                -e 's:^from auth\.:from ..auth.:g' \
+                -e 's:^from enums import:from ..enums import:g' \
+                -e 's:^from enums\.:from ..enums.:g' \
+                -e 's:^from messages import:from ..messages import:g' \
+                -e 's:^from messages\.:from ..messages.:g' \
+                -e 's:^import \(.*_pb2\):from . import \1:g' \
+                "$f"
+            ;;
+        *)
+            sed -i \
+                -e 's:^from anon import:from .anon import:g' \
+                -e 's:^from anon\.:from .anon.:g' \
+                -e 's:^from auth import:from .auth import:g' \
+                -e 's:^from auth\.:from .auth.:g' \
+                -e 's:^from enums import:from .enums import:g' \
+                -e 's:^from enums\.:from .enums.:g' \
+                -e 's:^from messages import:from .messages import:g' \
+                -e 's:^from messages\.:from .messages.:g' \
+                -e 's:^import \(.*_pb2\):from . import \1:g' \
+                "$f"
+            ;;
+    esac
+done
+
+for subdir in anon auth enums messages; do
+    if [ -d "$ROOT_DIR/$subdir" ]; then
+        touch "$ROOT_DIR/$subdir/__init__.py"
+    fi
+done
+
+# Append re-exports to envelope_pb2.py for backward compatibility
+# Connection.py uses getattr(envelope_module, class_name) for dynamic dispatch
+echo "" >> "$ROOT_DIR/envelope_pb2.py"
+echo "# Re-export all message and enum classes for dynamic lookup" >> "$ROOT_DIR/envelope_pb2.py"
+
+for subdir in anon auth enums messages; do
+    if [ -d "$ROOT_DIR/$subdir" ]; then
+        for f in "$ROOT_DIR/$subdir"/*_pb2.py; do
+            [ -f "$f" ] || continue
+            module=$(basename "$f" .py)
+            echo "from .$subdir.$module import *" >> "$ROOT_DIR/envelope_pb2.py"
+        done
+    fi
+done
diff --git a/python/datafed_pkg/setup.py b/python/datafed_pkg/setup.py
@@ -22,6 +22,13 @@
     long_description_content_type="text/markdown",
     url="https://github.com/ORNL/DataFed",
     packages=setuptools.find_packages(),
+    package_data={
+        "datafed": ["*.py"],
+        "datafed.anon": ["*.py"],
+        "datafed.auth": ["*.py"],
+        "datafed.enums": ["*.py"],
+        "datafed.messages": ["*.py"],
+    },
     setup_requires=["setuptools"],
     install_requires=install_requires,
     entry_points={"console_scripts": ["datafed = datafed.CLI:run"]},
@@ -31,3 +38,5 @@
         "Operating System :: OS Independent",
     ],
 )
+
+
